From 923d106485edce60a9d9e8eb81a0c104be955960 Mon Sep 17 00:00:00 2001 From: Vibhatha Abeykoon Date: Wed, 26 Jun 2024 16:40:27 +0530 Subject: [PATCH 1/4] fix: temp fix --- .../main/codegen/templates/BaseWriter.java | 7 ++ .../codegen/templates/UnionMapWriter.java | 24 +++++ .../vector/complex/impl/PromotableWriter.java | 3 + .../org/apache/arrow/vector/types/Types.java | 2 +- .../apache/arrow/vector/TestMapVector.java | 89 +++++++++++++++++++ 5 files changed, 124 insertions(+), 1 deletion(-) diff --git a/java/vector/src/main/codegen/templates/BaseWriter.java b/java/vector/src/main/codegen/templates/BaseWriter.java index 35df256b324..5af1f71c4eb 100644 --- a/java/vector/src/main/codegen/templates/BaseWriter.java +++ b/java/vector/src/main/codegen/templates/BaseWriter.java @@ -84,6 +84,13 @@ public interface ListWriter extends BaseWriter { <#assign capName = minor.class?cap_first /> ${capName}Writer ${lowerName}(); + + default Decimal256Writer decimal256(int scale, int precision) { + return null; + } + default DecimalWriter decimal(int scale, int precision) { + return null; + } } public interface MapWriter extends ListWriter { diff --git a/java/vector/src/main/codegen/templates/UnionMapWriter.java b/java/vector/src/main/codegen/templates/UnionMapWriter.java index 606f880377b..c3e509df3d6 100644 --- a/java/vector/src/main/codegen/templates/UnionMapWriter.java +++ b/java/vector/src/main/codegen/templates/UnionMapWriter.java @@ -183,6 +183,30 @@ public Decimal256Writer decimal256() { } } + @Override + public DecimalWriter decimal(int scale, int precision) { + switch (mode) { + case KEY: + return entryWriter.decimal(MapVector.KEY_NAME, scale, precision); + case VALUE: + return entryWriter.decimal(MapVector.VALUE_NAME, scale, precision); + default: + return this; + } + } + + @Override + public Decimal256Writer decimal256(int scale, int precision) { + switch (mode) { + case KEY: + return entryWriter.decimal256(MapVector.KEY_NAME, scale, precision); + case VALUE: + return entryWriter.decimal256(MapVector.VALUE_NAME, scale, precision); + default: + return this; + } + } + @Override public StructWriter struct() { diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java index 7fd0def9673..cc698f04e19 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java @@ -289,6 +289,9 @@ private boolean requiresArrowType(MinorType type) { protected FieldWriter getWriter(MinorType type, ArrowType arrowType) { if (state == State.UNION) { if (requiresArrowType(type)) { + if (arrowType == null) { + arrowType = type.getType(); + } ((UnionWriter) writer).getWriter(type, arrowType); } else { ((UnionWriter) writer).getWriter(type); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/types/Types.java b/java/vector/src/main/java/org/apache/arrow/vector/types/Types.java index ed099890e1b..31a11ce2386 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/types/Types.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/types/Types.java @@ -693,7 +693,7 @@ public FieldWriter getNewFieldWriter(ValueVector vector) { return new DenseUnionWriter((DenseUnionVector) vector); } }, - MAP(null) { + MAP(new Map(false)) { @Override public FieldVector getNewVector( Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 213ffced273..07f317a3d42 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -22,14 +22,18 @@ import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.math.BigDecimal; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.complex.reader.FieldReader; @@ -1204,4 +1208,89 @@ public void testMakeTransferPairPreserveNullability() { assertEquals(intField, vec.getField().getChildren().get(0)); assertEquals(intField, res.getField().getChildren().get(0)); } + + @Test + public void testSimpleMapVectorWithDecimals() { + try (final MapVector vector = MapVector.empty("v", allocator, false)) { + vector.allocateNew(); + UnionMapWriter mapWriter = vector.getWriter(); + + mapWriter.setPosition(0); + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().decimal(1, 2).writeDecimal(BigDecimal.valueOf(10, 1)); + mapWriter.value().decimal(1, 2).writeDecimal(BigDecimal.valueOf(20, 1)); + mapWriter.endEntry(); + mapWriter.endMap(); + + vector.setValueCount(1); + assertEquals(vector.getChildrenFromFields().size(), 1); + StructVector structVector = (StructVector) vector.getChildrenFromFields().get(0); + assertEquals(structVector.getChildrenFromFields().size(), 2); + DecimalVector keyVector = (DecimalVector) structVector.getChildrenFromFields().get(0); + DecimalVector valueVector = (DecimalVector) structVector.getChildrenFromFields().get(1); + assertEquals(keyVector.getObject(0), BigDecimal.valueOf(10, 1)); + assertEquals(valueVector.getObject(0), BigDecimal.valueOf(20, 1)); + } + } + + @Test + public void testWritingDecimals() { + try (ListVector from = ListVector.empty("v", allocator)) { + UnionListWriter listWriter = from.getWriter(); + listWriter.allocate(); + + listWriter.setPosition(0); + listWriter.startList(); + listWriter.map().startMap(); + listWriter.map().startEntry(); + listWriter.map().key().integer().writeInt(10); + listWriter.map().value().integer().writeInt(20); + listWriter.map().endEntry(); + listWriter.map().startEntry(); + listWriter.map().key().decimal(1, 2).writeDecimal(BigDecimal.valueOf(2.0)); + listWriter.map().value().decimal(1, 2).writeDecimal(BigDecimal.valueOf(3.0)); + listWriter.map().endEntry(); + listWriter.map().endMap(); + listWriter.endList(); + + listWriter.setValueCount(1); + from.setValueCount(1); + + System.out.println(from); + } + } + + @Test + public void testListOfStruct() { + try (ListVector from = ListVector.empty("v", new RootAllocator())) { + + UnionListWriter listWriter = from.getWriter(); + listWriter.allocate(); + + // write null, [null,{"f1":1,"f2":2},null, + // {"f1":1,"f2":2},null] alternatively + for (int i = 0; i < 10; i++) { + listWriter.setPosition(i); + if (i % 2 == 0) { + listWriter.writeNull(); + continue; + } + listWriter.startList(); + listWriter.struct().writeNull(); + listWriter.struct().start(); + listWriter.struct().decimal("f1", 1, 2).writeDecimal(BigDecimal.valueOf(2.0)); + listWriter.struct().end(); + listWriter.struct().writeNull(); + listWriter.struct().start(); + listWriter.struct().decimal("f1", 1, 2).writeDecimal(BigDecimal.valueOf(3.0)); + listWriter.struct().end(); + listWriter.struct().writeNull(); + listWriter.endList(); + } + from.setValueCount(10); + + System.out.println(from); + } + } } From dc2f1573a8b5ac3a3c271395e9742f38252a4ff8 Mon Sep 17 00:00:00 2001 From: Vibhatha Abeykoon Date: Thu, 27 Jun 2024 05:18:48 +0530 Subject: [PATCH 2/4] fix: adding map vector test cases --- .../org/apache/arrow/vector/TestMapVector.java | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 07f317a3d42..6e80d6ea332 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -1236,8 +1236,8 @@ public void testSimpleMapVectorWithDecimals() { @Test public void testWritingDecimals() { - try (ListVector from = ListVector.empty("v", allocator)) { - UnionListWriter listWriter = from.getWriter(); + try (ListVector vector = ListVector.empty("v", allocator)) { + UnionListWriter listWriter = vector.getWriter(); listWriter.allocate(); listWriter.setPosition(0); @@ -1255,9 +1255,19 @@ public void testWritingDecimals() { listWriter.endList(); listWriter.setValueCount(1); - from.setValueCount(1); + vector.setValueCount(1); - System.out.println(from); + assertEquals(vector.getChildrenFromFields().size(), 1); + MapVector mapVector = (MapVector) vector.getChildrenFromFields().get(0); + assertEquals(mapVector.getChildrenFromFields().size(), 1); + StructVector structVector = (StructVector) mapVector.getChildrenFromFields().get(0); + assertEquals(structVector.getChildrenFromFields().size(), 2); + assertEquals(structVector.getChildrenFromFields().get(0).getObject(0), 10); + assertEquals( + structVector.getChildrenFromFields().get(0).getObject(1), BigDecimal.valueOf(2.0)); + assertEquals(structVector.getChildrenFromFields().get(1).getObject(0), 20); + assertEquals( + structVector.getChildrenFromFields().get(1).getObject(1), BigDecimal.valueOf(3.0)); } } From 8d750aeeb134bf6015ccb8d86d719bbbc2de037c Mon Sep 17 00:00:00 2001 From: Vibhatha Abeykoon Date: Thu, 27 Jun 2024 07:24:41 +0530 Subject: [PATCH 3/4] fix: adding templates for new decimal functions --- .../templates/AbstractFieldWriter.java | 9 +++++ .../main/codegen/templates/BaseWriter.java | 13 ++++---- .../codegen/templates/UnionListWriter.java | 6 ++++ .../codegen/templates/UnionMapWriter.java | 1 - .../apache/arrow/vector/TestMapVector.java | 33 ------------------- 5 files changed, 21 insertions(+), 41 deletions(-) diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java index 6c2368117f7..01d5f493345 100644 --- a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java +++ b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java @@ -233,6 +233,15 @@ public MapWriter map(String name, boolean keysSorted) { } + <#if lowerName?contains("decimal")> + + @Override + public ${capName}Writer ${lowerName}(<#list minor.typeParams as typeParam>${typeParam.type} ${typeParam.name}<#if typeParam_has_next>, ) { + fail("${capName}(" + <#list minor.typeParams as typeParam>"${typeParam.name}: " + ${typeParam.name} + ", " + ")"); + return null; + } + + @Override public ${capName}Writer ${lowerName}(String name) { fail("${capName}"); diff --git a/java/vector/src/main/codegen/templates/BaseWriter.java b/java/vector/src/main/codegen/templates/BaseWriter.java index 5af1f71c4eb..718ee227f49 100644 --- a/java/vector/src/main/codegen/templates/BaseWriter.java +++ b/java/vector/src/main/codegen/templates/BaseWriter.java @@ -56,6 +56,9 @@ public interface StructWriter extends BaseWriter { <#if minor.typeParams?? > ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}); + <#if lowerName?contains("decimal")> + ${capName}Writer ${lowerName}(<#list minor.typeParams as typeParam>${typeParam.type} ${typeParam.name}<#if typeParam_has_next>, ); + ${capName}Writer ${lowerName}(String name); @@ -83,14 +86,10 @@ public interface ListWriter extends BaseWriter { <#assign upperName = minor.class?upper_case /> <#assign capName = minor.class?cap_first /> ${capName}Writer ${lowerName}(); + <#if lowerName?contains("decimal")> + ${capName}Writer ${lowerName}(<#list minor.typeParams as typeParam>${typeParam.type} ${typeParam.name}<#if typeParam_has_next>, ); + - - default Decimal256Writer decimal256(int scale, int precision) { - return null; - } - default DecimalWriter decimal(int scale, int precision) { - return null; - } } public interface MapWriter extends ListWriter { diff --git a/java/vector/src/main/codegen/templates/UnionListWriter.java b/java/vector/src/main/codegen/templates/UnionListWriter.java index eeb964c055f..af8e19d15b3 100644 --- a/java/vector/src/main/codegen/templates/UnionListWriter.java +++ b/java/vector/src/main/codegen/templates/UnionListWriter.java @@ -127,6 +127,12 @@ public void setPosition(int index) { public ${minor.class}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { return writer.${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name}); } + <#if lowerName?contains("decimal")> + @Override + public ${capName}Writer ${lowerName}(<#list minor.typeParams as typeParam>${typeParam.type} ${typeParam.name}<#if typeParam_has_next>, ) { + return writer.${lowerName}(<#list minor.typeParams as typeParam>${typeParam.name}<#if typeParam_has_next>, ); + } + @Override diff --git a/java/vector/src/main/codegen/templates/UnionMapWriter.java b/java/vector/src/main/codegen/templates/UnionMapWriter.java index c3e509df3d6..938f009047b 100644 --- a/java/vector/src/main/codegen/templates/UnionMapWriter.java +++ b/java/vector/src/main/codegen/templates/UnionMapWriter.java @@ -207,7 +207,6 @@ public Decimal256Writer decimal256(int scale, int precision) { } } - @Override public StructWriter struct() { switch (mode) { diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 6e80d6ea332..1ccb3276c8d 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -1270,37 +1270,4 @@ public void testWritingDecimals() { structVector.getChildrenFromFields().get(1).getObject(1), BigDecimal.valueOf(3.0)); } } - - @Test - public void testListOfStruct() { - try (ListVector from = ListVector.empty("v", new RootAllocator())) { - - UnionListWriter listWriter = from.getWriter(); - listWriter.allocate(); - - // write null, [null,{"f1":1,"f2":2},null, - // {"f1":1,"f2":2},null] alternatively - for (int i = 0; i < 10; i++) { - listWriter.setPosition(i); - if (i % 2 == 0) { - listWriter.writeNull(); - continue; - } - listWriter.startList(); - listWriter.struct().writeNull(); - listWriter.struct().start(); - listWriter.struct().decimal("f1", 1, 2).writeDecimal(BigDecimal.valueOf(2.0)); - listWriter.struct().end(); - listWriter.struct().writeNull(); - listWriter.struct().start(); - listWriter.struct().decimal("f1", 1, 2).writeDecimal(BigDecimal.valueOf(3.0)); - listWriter.struct().end(); - listWriter.struct().writeNull(); - listWriter.endList(); - } - from.setValueCount(10); - - System.out.println(from); - } - } } From 2ac3e11f831ff36c85d9a538d08652b0133db964 Mon Sep 17 00:00:00 2001 From: Vibhatha Abeykoon Date: Thu, 27 Jun 2024 07:31:47 +0530 Subject: [PATCH 4/4] fix: typo --- .../src/test/java/org/apache/arrow/vector/TestMapVector.java | 1 - 1 file changed, 1 deletion(-) diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 1ccb3276c8d..ec56b08d407 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -29,7 +29,6 @@ import java.util.Map; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector;