From 847a4c366dff1bc1832160a77098eddd7e97b87b Mon Sep 17 00:00:00 2001 From: Ivan Chesnov Date: Wed, 28 Jun 2023 13:42:11 +0300 Subject: [PATCH 1/4] Added creating MapWriter in ComplexWriter. --- .../main/codegen/templates/BaseWriter.java | 1 + .../complex/impl/ComplexWriterImpl.java | 39 ++- .../complex/writer/TestComplexWriter.java | 275 +++++++----------- 3 files changed, 147 insertions(+), 168 deletions(-) diff --git a/java/vector/src/main/codegen/templates/BaseWriter.java b/java/vector/src/main/codegen/templates/BaseWriter.java index 4d63fb73e98..3b35d22692e 100644 --- a/java/vector/src/main/codegen/templates/BaseWriter.java +++ b/java/vector/src/main/codegen/templates/BaseWriter.java @@ -106,6 +106,7 @@ public interface ComplexWriter { void copyReader(FieldReader reader); StructWriter rootAsStruct(); ListWriter rootAsList(); + MapWriter rootAsMap(boolean keysSorted); void setPosition(int index); void setValueCount(int count); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java index 13b26bb67da..4d732a04226 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java @@ -19,6 +19,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.StateTool; import org.apache.arrow.vector.complex.StructVector; @@ -32,6 +33,7 @@ public class ComplexWriterImpl extends AbstractFieldWriter implements ComplexWri private NullableStructWriter structRoot; private UnionListWriter listRoot; + private UnionMapWriter mapRoot; private final NonNullableStructVector container; Mode mode = Mode.INIT; @@ -39,7 +41,7 @@ public class ComplexWriterImpl extends AbstractFieldWriter implements ComplexWri private final boolean unionEnabled; private final NullableStructWriterFactory nullableStructWriterFactory; - private enum Mode { INIT, STRUCT, LIST } + private enum Mode { INIT, STRUCT, LIST, MAP } /** * Constructs a new instance. @@ -107,6 +109,9 @@ public void clear() { case LIST: listRoot.clear(); break; + case MAP: + mapRoot.clear(); + break; default: break; } @@ -121,6 +126,9 @@ public void setValueCount(int count) { case LIST: listRoot.setValueCount(count); break; + case MAP: + mapRoot.setValueCount(count); + break; default: break; } @@ -136,6 +144,9 @@ public void setPosition(int index) { case LIST: listRoot.setPosition(index); break; + case MAP: + mapRoot.setPosition(index); + break; default: break; } @@ -223,5 +234,29 @@ public ListWriter rootAsList() { return listRoot; } + @Override + public MapWriter rootAsMap(boolean keysSorted) { + switch (mode) { -} + case INIT: + int vectorCount = container.size(); + // TODO allow dictionaries in complex types + MapVector mapVector = container.addOrGetMap(name, keysSorted); + if (container.size() > vectorCount) { + mapVector.allocateNew(); + } + mapRoot = new UnionMapWriter(mapVector); + mapRoot.setPosition(idx()); + mode = Mode.MAP; + break; + + case MAP: + break; + + default: + check(Mode.INIT, Mode.STRUCT); + } + + return mapRoot; + } +} \ No newline at end of file diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java index 9f7f66083ac..3ca6880acbf 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java @@ -21,7 +21,6 @@ import java.math.BigDecimal; import java.time.LocalDateTime; -import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -49,6 +48,7 @@ import org.apache.arrow.vector.complex.impl.UnionListReader; 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.impl.UnionReader; import org.apache.arrow.vector.complex.impl.UnionWriter; import org.apache.arrow.vector.complex.reader.BaseReader.StructReader; @@ -62,15 +62,8 @@ import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.holders.DecimalHolder; -import org.apache.arrow.vector.holders.DurationHolder; -import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.holders.IntHolder; -import org.apache.arrow.vector.holders.NullableDurationHolder; -import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder; -import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; -import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; -import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.ArrowType.Int; @@ -362,125 +355,6 @@ public void listDecimalType() { } } - @Test - public void listTimeStampMilliTZType() { - try (ListVector listVector = ListVector.empty("list", allocator)) { - listVector.allocateNew(); - UnionListWriter listWriter = new UnionListWriter(listVector); - for (int i = 0; i < COUNT; i++) { - listWriter.startList(); - for (int j = 0; j < i % 7; j++) { - if (j % 2 == 0) { - listWriter.writeNull(); - } else { - TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); - holder.timezone = "FakeTimeZone"; - holder.value = j; - listWriter.timeStampMilliTZ().write(holder); - } - } - listWriter.endList(); - } - listWriter.setValueCount(COUNT); - UnionListReader listReader = new UnionListReader(listVector); - for (int i = 0; i < COUNT; i++) { - listReader.setPosition(i); - for (int j = 0; j < i % 7; j++) { - listReader.next(); - if (j % 2 == 0) { - assertFalse("index is set: " + j, listReader.reader().isSet()); - } else { - NullableTimeStampMilliTZHolder actual = new NullableTimeStampMilliTZHolder(); - listReader.reader().read(actual); - assertEquals(j, actual.value); - assertEquals("FakeTimeZone", actual.timezone); - } - } - } - } - } - - @Test - public void listDurationType() { - try (ListVector listVector = ListVector.empty("list", allocator)) { - listVector.allocateNew(); - UnionListWriter listWriter = new UnionListWriter(listVector); - for (int i = 0; i < COUNT; i++) { - listWriter.startList(); - for (int j = 0; j < i % 7; j++) { - if (j % 2 == 0) { - listWriter.writeNull(); - } else { - DurationHolder holder = new DurationHolder(); - holder.unit = TimeUnit.MICROSECOND; - holder.value = j; - listWriter.duration().write(holder); - } - } - listWriter.endList(); - } - listWriter.setValueCount(COUNT); - UnionListReader listReader = new UnionListReader(listVector); - for (int i = 0; i < COUNT; i++) { - listReader.setPosition(i); - for (int j = 0; j < i % 7; j++) { - listReader.next(); - if (j % 2 == 0) { - assertFalse("index is set: " + j, listReader.reader().isSet()); - } else { - NullableDurationHolder actual = new NullableDurationHolder(); - listReader.reader().read(actual); - assertEquals(TimeUnit.MICROSECOND, actual.unit); - assertEquals(j, actual.value); - } - } - } - } - } - - @Test - public void listFixedSizeBinaryType() throws Exception { - List bufs = new ArrayList(); - try (ListVector listVector = ListVector.empty("list", allocator)) { - listVector.allocateNew(); - UnionListWriter listWriter = new UnionListWriter(listVector); - for (int i = 0; i < COUNT; i++) { - listWriter.startList(); - for (int j = 0; j < i % 7; j++) { - if (j % 2 == 0) { - listWriter.writeNull(); - } else { - ArrowBuf buf = allocator.buffer(4); - buf.setInt(0, j); - FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); - holder.byteWidth = 4; - holder.buffer = buf; - listWriter.fixedSizeBinary().write(holder); - bufs.add(buf); - } - } - listWriter.endList(); - } - listWriter.setValueCount(COUNT); - UnionListReader listReader = new UnionListReader(listVector); - for (int i = 0; i < COUNT; i++) { - listReader.setPosition(i); - for (int j = 0; j < i % 7; j++) { - listReader.next(); - if (j % 2 == 0) { - assertFalse("index is set: " + j, listReader.reader().isSet()); - } else { - NullableFixedSizeBinaryHolder actual = new NullableFixedSizeBinaryHolder(); - listReader.reader().read(actual); - assertEquals(j, actual.buffer.getInt(0)); - assertEquals(4, actual.byteWidth); - } - } - } - } - AutoCloseables.close(bufs); - } - @Test public void listScalarTypeNullable() { try (ListVector listVector = ListVector.empty("list", allocator)) { @@ -732,33 +606,14 @@ private void checkListMap(ListVector listVector) { } @Test - public void simpleUnion() throws Exception { - List bufs = new ArrayList(); + public void simpleUnion() { UnionVector vector = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); UnionWriter unionWriter = new UnionWriter(vector); unionWriter.allocate(); for (int i = 0; i < COUNT; i++) { unionWriter.setPosition(i); - if (i % 5 == 0) { + if (i % 2 == 0) { unionWriter.writeInt(i); - } else if (i % 5 == 1) { - TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); - holder.value = (long) i; - holder.timezone = "AsdfTimeZone"; - unionWriter.write(holder); - } else if (i % 5 == 2) { - DurationHolder holder = new DurationHolder(); - holder.value = (long) i; - holder.unit = TimeUnit.NANOSECOND; - unionWriter.write(holder); - } else if (i % 5 == 3) { - FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); - ArrowBuf buf = allocator.buffer(4); - buf.setInt(0, i); - holder.byteWidth = 4; - holder.buffer = buf; - unionWriter.write(holder); - bufs.add(buf); } else { unionWriter.writeFloat4((float) i); } @@ -767,29 +622,13 @@ public void simpleUnion() throws Exception { UnionReader unionReader = new UnionReader(vector); for (int i = 0; i < COUNT; i++) { unionReader.setPosition(i); - if (i % 5 == 0) { + if (i % 2 == 0) { Assert.assertEquals(i, i, unionReader.readInteger()); - } else if (i % 5 == 1) { - NullableTimeStampMilliTZHolder holder = new NullableTimeStampMilliTZHolder(); - unionReader.read(holder); - Assert.assertEquals(i, holder.value); - Assert.assertEquals("AsdfTimeZone", holder.timezone); - } else if (i % 5 == 2) { - NullableDurationHolder holder = new NullableDurationHolder(); - unionReader.read(holder); - Assert.assertEquals(i, holder.value); - Assert.assertEquals(TimeUnit.NANOSECOND, holder.unit); - } else if (i % 5 == 3) { - NullableFixedSizeBinaryHolder holder = new NullableFixedSizeBinaryHolder(); - unionReader.read(holder); - assertEquals(i, holder.buffer.getInt(0)); - assertEquals(4, holder.byteWidth); } else { Assert.assertEquals((float) i, unionReader.readFloat(), 1e-12); } } vector.close(); - AutoCloseables.close(bufs); } @Test @@ -1563,4 +1402,108 @@ public void testStructOfList() { Assert.assertEquals(0, size); } } -} + + @Test + public void testMap() { + try (MapVector mapVector = MapVector.empty("map", allocator, false)) { + mapVector.allocateNew(); + UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + for (int i = 0; i < COUNT; i++) { + mapWriter.startMap(); + for (int j = 0; j < i % 7; j++) { + mapWriter.startEntry(); + if (j % 2 == 0) { + mapWriter.key().integer().writeInt(j); + mapWriter.value().integer().writeInt(j + 1); + } else { + IntHolder keyHolder = new IntHolder(); + keyHolder.value = j; + IntHolder valueHolder = new IntHolder(); + valueHolder.value = j + 1; + mapWriter.key().integer().write(keyHolder); + mapWriter.value().integer().write(valueHolder); + } + mapWriter.endEntry(); + } + mapWriter.endMap(); + } + mapWriter.setValueCount(COUNT); + UnionMapReader mapReader = new UnionMapReader(mapVector); + for (int i = 0; i < COUNT; i++) { + mapReader.setPosition(i); + for (int j = 0; j < i % 7; j++) { + mapReader.next(); + assertEquals(j, mapReader.key().readInteger().intValue()); + assertEquals(j + 1, mapReader.value().readInteger().intValue()); + } + } + } + } + + @Test + public void testMapWithNulls() { + try (MapVector mapVector = MapVector.empty("map", allocator, false)) { + mapVector.allocateNew(); + UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().integer().writeNull(); + mapWriter.value().integer().writeInt(1); + mapWriter.endEntry(); + mapWriter.endMap(); + mapWriter.setValueCount(1); + UnionMapReader mapReader = new UnionMapReader(mapVector); + Assert.assertNull(mapReader.key().readInteger()); + assertEquals(1, mapReader.value().readInteger().intValue()); + } + } + + @Test + public void testMapWithListKey() { + try (MapVector mapVector = MapVector.empty("map", allocator, false)) { + mapVector.allocateNew(); + UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().list().startList(); + for (int i = 0; i < 3; i++) { + mapWriter.key().list().integer().writeInt(i); + } + mapWriter.key().list().endList(); + mapWriter.value().integer().writeInt(1); + mapWriter.endEntry(); + mapWriter.endMap(); + mapWriter.setValueCount(1); + UnionMapReader mapReader = new UnionMapReader(mapVector); + mapReader.key().next(); + assertEquals(0, mapReader.key().reader().readInteger().intValue()); + mapReader.key().next(); + assertEquals(1, mapReader.key().reader().readInteger().intValue()); + mapReader.key().next(); + assertEquals(2, mapReader.key().reader().readInteger().intValue()); + assertEquals(1, mapReader.value().readInteger().intValue()); + } + } + + @Test + public void testMapWithStruckKey() { + try (MapVector mapVector = MapVector.empty("map", allocator, false)) { + mapVector.allocateNew(); + UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().struct().start(); + mapWriter.key().struct().integer("value1").writeInt(1); + mapWriter.key().struct().integer("value2").writeInt(2); + mapWriter.key().struct().end(); + mapWriter.value().integer().writeInt(1); + mapWriter.endEntry(); + mapWriter.endMap(); + mapWriter.setValueCount(1); + UnionMapReader mapReader = new UnionMapReader(mapVector); + assertEquals(1, mapReader.key().reader("value1").readInteger().intValue()); + assertEquals(2, mapReader.key().reader("value2").readInteger().intValue()); + assertEquals(1, mapReader.value().readInteger().intValue()); + } + } +} \ No newline at end of file From 7176860310113f08006e6681719dcec2c7f1a557 Mon Sep 17 00:00:00 2001 From: Ivan Chesnov Date: Thu, 29 Jun 2023 15:54:30 +0300 Subject: [PATCH 2/4] Revert accidentally removed tests --- .../complex/writer/TestComplexWriter.java | 168 +++++++++++++++++- 1 file changed, 165 insertions(+), 3 deletions(-) diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java index 3ca6880acbf..c0694f743dd 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java @@ -21,6 +21,7 @@ import java.math.BigDecimal; import java.time.LocalDateTime; +import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -62,8 +63,15 @@ import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.holders.DecimalHolder; +import org.apache.arrow.vector.holders.DurationHolder; +import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.holders.IntHolder; +import org.apache.arrow.vector.holders.NullableDurationHolder; +import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder; +import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; +import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; +import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.ArrowType.Int; @@ -355,6 +363,125 @@ public void listDecimalType() { } } + @Test + public void listTimeStampMilliTZType() { + try (ListVector listVector = ListVector.empty("list", allocator)) { + listVector.allocateNew(); + UnionListWriter listWriter = new UnionListWriter(listVector); + for (int i = 0; i < COUNT; i++) { + listWriter.startList(); + for (int j = 0; j < i % 7; j++) { + if (j % 2 == 0) { + listWriter.writeNull(); + } else { + TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); + holder.timezone = "FakeTimeZone"; + holder.value = j; + listWriter.timeStampMilliTZ().write(holder); + } + } + listWriter.endList(); + } + listWriter.setValueCount(COUNT); + UnionListReader listReader = new UnionListReader(listVector); + for (int i = 0; i < COUNT; i++) { + listReader.setPosition(i); + for (int j = 0; j < i % 7; j++) { + listReader.next(); + if (j % 2 == 0) { + assertFalse("index is set: " + j, listReader.reader().isSet()); + } else { + NullableTimeStampMilliTZHolder actual = new NullableTimeStampMilliTZHolder(); + listReader.reader().read(actual); + assertEquals(j, actual.value); + assertEquals("FakeTimeZone", actual.timezone); + } + } + } + } + } + + @Test + public void listDurationType() { + try (ListVector listVector = ListVector.empty("list", allocator)) { + listVector.allocateNew(); + UnionListWriter listWriter = new UnionListWriter(listVector); + for (int i = 0; i < COUNT; i++) { + listWriter.startList(); + for (int j = 0; j < i % 7; j++) { + if (j % 2 == 0) { + listWriter.writeNull(); + } else { + DurationHolder holder = new DurationHolder(); + holder.unit = TimeUnit.MICROSECOND; + holder.value = j; + listWriter.duration().write(holder); + } + } + listWriter.endList(); + } + listWriter.setValueCount(COUNT); + UnionListReader listReader = new UnionListReader(listVector); + for (int i = 0; i < COUNT; i++) { + listReader.setPosition(i); + for (int j = 0; j < i % 7; j++) { + listReader.next(); + if (j % 2 == 0) { + assertFalse("index is set: " + j, listReader.reader().isSet()); + } else { + NullableDurationHolder actual = new NullableDurationHolder(); + listReader.reader().read(actual); + assertEquals(TimeUnit.MICROSECOND, actual.unit); + assertEquals(j, actual.value); + } + } + } + } + } + + @Test + public void listFixedSizeBinaryType() throws Exception { + List bufs = new ArrayList(); + try (ListVector listVector = ListVector.empty("list", allocator)) { + listVector.allocateNew(); + UnionListWriter listWriter = new UnionListWriter(listVector); + for (int i = 0; i < COUNT; i++) { + listWriter.startList(); + for (int j = 0; j < i % 7; j++) { + if (j % 2 == 0) { + listWriter.writeNull(); + } else { + ArrowBuf buf = allocator.buffer(4); + buf.setInt(0, j); + FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); + holder.byteWidth = 4; + holder.buffer = buf; + listWriter.fixedSizeBinary().write(holder); + bufs.add(buf); + } + } + listWriter.endList(); + } + listWriter.setValueCount(COUNT); + UnionListReader listReader = new UnionListReader(listVector); + for (int i = 0; i < COUNT; i++) { + listReader.setPosition(i); + for (int j = 0; j < i % 7; j++) { + listReader.next(); + if (j % 2 == 0) { + assertFalse("index is set: " + j, listReader.reader().isSet()); + } else { + NullableFixedSizeBinaryHolder actual = new NullableFixedSizeBinaryHolder(); + listReader.reader().read(actual); + assertEquals(j, actual.buffer.getInt(0)); + assertEquals(4, actual.byteWidth); + } + } + } + } + AutoCloseables.close(bufs); + } + @Test public void listScalarTypeNullable() { try (ListVector listVector = ListVector.empty("list", allocator)) { @@ -606,14 +733,33 @@ private void checkListMap(ListVector listVector) { } @Test - public void simpleUnion() { + public void simpleUnion() throws Exception { + List bufs = new ArrayList(); UnionVector vector = new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); UnionWriter unionWriter = new UnionWriter(vector); unionWriter.allocate(); for (int i = 0; i < COUNT; i++) { unionWriter.setPosition(i); - if (i % 2 == 0) { + if (i % 5 == 0) { unionWriter.writeInt(i); + } else if (i % 5 == 1) { + TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); + holder.value = (long) i; + holder.timezone = "AsdfTimeZone"; + unionWriter.write(holder); + } else if (i % 5 == 2) { + DurationHolder holder = new DurationHolder(); + holder.value = (long) i; + holder.unit = TimeUnit.NANOSECOND; + unionWriter.write(holder); + } else if (i % 5 == 3) { + FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); + ArrowBuf buf = allocator.buffer(4); + buf.setInt(0, i); + holder.byteWidth = 4; + holder.buffer = buf; + unionWriter.write(holder); + bufs.add(buf); } else { unionWriter.writeFloat4((float) i); } @@ -622,13 +768,29 @@ public void simpleUnion() { UnionReader unionReader = new UnionReader(vector); for (int i = 0; i < COUNT; i++) { unionReader.setPosition(i); - if (i % 2 == 0) { + if (i % 5 == 0) { Assert.assertEquals(i, i, unionReader.readInteger()); + } else if (i % 5 == 1) { + NullableTimeStampMilliTZHolder holder = new NullableTimeStampMilliTZHolder(); + unionReader.read(holder); + Assert.assertEquals(i, holder.value); + Assert.assertEquals("AsdfTimeZone", holder.timezone); + } else if (i % 5 == 2) { + NullableDurationHolder holder = new NullableDurationHolder(); + unionReader.read(holder); + Assert.assertEquals(i, holder.value); + Assert.assertEquals(TimeUnit.NANOSECOND, holder.unit); + } else if (i % 5 == 3) { + NullableFixedSizeBinaryHolder holder = new NullableFixedSizeBinaryHolder(); + unionReader.read(holder); + assertEquals(i, holder.buffer.getInt(0)); + assertEquals(4, holder.byteWidth); } else { Assert.assertEquals((float) i, unionReader.readFloat(), 1e-12); } } vector.close(); + AutoCloseables.close(bufs); } @Test From b5f1333c69f843a34f16f037218a2df9e3d233c5 Mon Sep 17 00:00:00 2001 From: Ivan Chesnov Date: Thu, 29 Jun 2023 20:15:05 +0300 Subject: [PATCH 3/4] fixed formating --- .../org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java | 2 +- .../apache/arrow/vector/complex/writer/TestComplexWriter.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java index 4d732a04226..8d2694b6df8 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java @@ -259,4 +259,4 @@ public MapWriter rootAsMap(boolean keysSorted) { return mapRoot; } -} \ No newline at end of file +} diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java index c0694f743dd..1e675a46997 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java @@ -1668,4 +1668,4 @@ public void testMapWithStruckKey() { assertEquals(1, mapReader.value().readInteger().intValue()); } } -} \ No newline at end of file +} From 88cd2c26c40c2ecb224b2548297a593cfbd9a466 Mon Sep 17 00:00:00 2001 From: Ivan Chesnov Date: Tue, 4 Jul 2023 18:39:30 +0300 Subject: [PATCH 4/4] Updated test for using ComplexWriterImpl --- .../complex/writer/TestComplexWriter.java | 43 +++++++++---------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java index 1e675a46997..6f7f5abd30a 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java @@ -49,7 +49,6 @@ import org.apache.arrow.vector.complex.impl.UnionListReader; 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.impl.UnionReader; import org.apache.arrow.vector.complex.impl.UnionWriter; import org.apache.arrow.vector.complex.reader.BaseReader.StructReader; @@ -1567,9 +1566,9 @@ public void testStructOfList() { @Test public void testMap() { - try (MapVector mapVector = MapVector.empty("map", allocator, false)) { - mapVector.allocateNew(); - UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { + ComplexWriter writer = new ComplexWriterImpl("root", parent); + MapWriter mapWriter = writer.rootAsMap(false); for (int i = 0; i < COUNT; i++) { mapWriter.startMap(); for (int j = 0; j < i % 7; j++) { @@ -1589,8 +1588,8 @@ public void testMap() { } mapWriter.endMap(); } - mapWriter.setValueCount(COUNT); - UnionMapReader mapReader = new UnionMapReader(mapVector); + writer.setValueCount(COUNT); + UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); for (int i = 0; i < COUNT; i++) { mapReader.setPosition(i); for (int j = 0; j < i % 7; j++) { @@ -1604,17 +1603,17 @@ public void testMap() { @Test public void testMapWithNulls() { - try (MapVector mapVector = MapVector.empty("map", allocator, false)) { - mapVector.allocateNew(); - UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { + ComplexWriter writer = new ComplexWriterImpl("root", parent); + MapWriter mapWriter = writer.rootAsMap(false); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().integer().writeNull(); mapWriter.value().integer().writeInt(1); mapWriter.endEntry(); mapWriter.endMap(); - mapWriter.setValueCount(1); - UnionMapReader mapReader = new UnionMapReader(mapVector); + writer.setValueCount(1); + UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); Assert.assertNull(mapReader.key().readInteger()); assertEquals(1, mapReader.value().readInteger().intValue()); } @@ -1622,9 +1621,9 @@ public void testMapWithNulls() { @Test public void testMapWithListKey() { - try (MapVector mapVector = MapVector.empty("map", allocator, false)) { - mapVector.allocateNew(); - UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { + ComplexWriter writer = new ComplexWriterImpl("root", parent); + MapWriter mapWriter = writer.rootAsMap(false); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().list().startList(); @@ -1635,8 +1634,8 @@ public void testMapWithListKey() { mapWriter.value().integer().writeInt(1); mapWriter.endEntry(); mapWriter.endMap(); - mapWriter.setValueCount(1); - UnionMapReader mapReader = new UnionMapReader(mapVector); + writer.setValueCount(1); + UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); mapReader.key().next(); assertEquals(0, mapReader.key().reader().readInteger().intValue()); mapReader.key().next(); @@ -1648,10 +1647,10 @@ public void testMapWithListKey() { } @Test - public void testMapWithStruckKey() { - try (MapVector mapVector = MapVector.empty("map", allocator, false)) { - mapVector.allocateNew(); - UnionMapWriter mapWriter = new UnionMapWriter(mapVector); + public void testMapWithStructKey() { + try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { + ComplexWriter writer = new ComplexWriterImpl("root", parent); + MapWriter mapWriter = writer.rootAsMap(false); mapWriter.startMap(); mapWriter.startEntry(); mapWriter.key().struct().start(); @@ -1661,8 +1660,8 @@ public void testMapWithStruckKey() { mapWriter.value().integer().writeInt(1); mapWriter.endEntry(); mapWriter.endMap(); - mapWriter.setValueCount(1); - UnionMapReader mapReader = new UnionMapReader(mapVector); + writer.setValueCount(1); + UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root"); assertEquals(1, mapReader.key().reader("value1").readInteger().intValue()); assertEquals(2, mapReader.key().reader("value2").readInteger().intValue()); assertEquals(1, mapReader.value().readInteger().intValue());