From 4ba5e10cb216ad3c6922e68a6114e1d9778b1c19 Mon Sep 17 00:00:00 2001 From: axreldable Date: Sat, 28 Feb 2026 16:11:31 +0100 Subject: [PATCH] GH-552: [Vector] Add absent methods to the UnionFixedSizeListWriter --- .../templates/UnionFixedSizeListWriter.java | 151 ++++++---------- .../codegen/templates/UnionListWriter.java | 2 - .../arrow/vector/TestFixedSizeListVector.java | 163 ++++++++++++++++++ .../apache/arrow/vector/TestMapVector.java | 35 ++-- .../org/apache/arrow/vector/TestUtils.java | 13 ++ 5 files changed, 245 insertions(+), 119 deletions(-) diff --git a/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java b/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java index f6e3f63caf..484199ab2a 100644 --- a/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java +++ b/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java @@ -35,6 +35,10 @@ <#include "/@includes/vv_imports.ftl" /> +<#function is_timestamp_tz type> + <#return type?starts_with("TimeStamp") && type?ends_with("TZ")> + + /* * This class is generated using freemarker and the ${.template_name} template. */ @@ -96,55 +100,30 @@ public void close() throws Exception { public void setPosition(int index) { super.setPosition(index); } - <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> - <#assign fields = minor.fields!type.fields /> - <#assign uncappedName = name?uncap_first/> - <#if uncappedName == "int" ><#assign uncappedName = "integer" /> - <#if !minor.typeParams?? > + <#list vv.types as type><#list type.minor as minor> + <#assign lowerName = minor.class?uncap_first /> + <#if lowerName == "int" ><#assign lowerName = "integer" /> + <#assign upperName = minor.class?upper_case /> @Override - public ${name}Writer ${uncappedName}() { + public ${minor.class}Writer ${lowerName}() { return this; } + <#if minor.typeParams?? > @Override - public ${name}Writer ${uncappedName}(String name) { - structName = name; - return writer.${uncappedName}(name); + 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}); } - - - @Override - public DecimalWriter decimal() { - return this; - } - - @Override - public DecimalWriter decimal(String name, int scale, int precision) { - return writer.decimal(name, scale, precision); - } - - @Override - public DecimalWriter decimal(String name) { - return writer.decimal(name); - } - @Override - public Decimal256Writer decimal256() { - return this; - } - - @Override - public Decimal256Writer decimal256(String name, int scale, int precision) { - return writer.decimal256(name, scale, precision); + public ${minor.class}Writer ${lowerName}(String name) { + structName = name; + return writer.${lowerName}(name); } - @Override - public Decimal256Writer decimal256(String name) { - return writer.decimal256(name); - } + @Override public StructWriter struct() { @@ -215,87 +194,86 @@ public void end() { } @Override - public void write(DecimalHolder holder) { - if (writer.idx() >= (idx() + 1) * listSize) { - throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); - } - writer.write(holder); - writer.setPosition(writer.idx() + 1); - } - - @Override - public void write(Decimal256Holder holder) { + public void writeNull() { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.write(holder); - writer.setPosition(writer.idx() + 1); + writer.writeNull(); } + <#list vv.types as type> + <#list type.minor as minor> + <#assign name = minor.class?cap_first /> + <#assign fields = minor.fields!type.fields /> + <#assign uncappedName = name?uncap_first/> @Override - public void writeNull() { + public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeNull(); + writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); + writer.setPosition(writer.idx()+1); } - public void writeDecimal(long start, ArrowBuf buffer, ArrowType arrowType) { + <#if is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> + @Override + public void write(${name}Holder holder) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal(start, buffer, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write(holder); + writer.setPosition(writer.idx()+1); } - public void writeDecimal(BigDecimal value) { + <#elseif minor.class?starts_with("Decimal")> + @Override + public void write${name}(long start, ArrowBuf buffer, ArrowType arrowType) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal(value); - writer.setPosition(writer.idx() + 1); + writer.write${name}(start, buffer, arrowType); + writer.setPosition(writer.idx()+1); } - public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) { + @Override + public void write(${name}Holder holder) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeBigEndianBytesToDecimal(value, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write(holder); + writer.setPosition(writer.idx()+1); } - public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) { + @Override + public void write${name}(BigDecimal value) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal256(start, buffer, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write${name}(value); + writer.setPosition(writer.idx()+1); } - public void writeDecimal256(BigDecimal value) { + @Override + public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType){ if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal256(value); + writer.writeBigEndianBytesTo${name}(value, arrowType); writer.setPosition(writer.idx() + 1); } - - public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { + <#else> + @Override + public void write(${name}Holder holder) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeBigEndianBytesToDecimal256(value, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); + writer.setPosition(writer.idx()+1); } + - - <#list vv.types as type> - <#list type.minor as minor> - <#assign name = minor.class?cap_first /> - <#assign fields = minor.fields!type.fields /> - <#assign uncappedName = name?uncap_first/> - <#if minor.class?ends_with("VarBinary")> + <#if minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { if (writer.idx() >= (idx() + 1) * listSize) { @@ -349,27 +327,8 @@ public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } - - - <#if !minor.typeParams?? > - @Override - public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { - if (writer.idx() >= (idx() + 1) * listSize) { - throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); - } - writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); - writer.setPosition(writer.idx() + 1); - } - - public void write(${name}Holder holder) { - if (writer.idx() >= (idx() + 1) * listSize) { - throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); - } - writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); - writer.setPosition(writer.idx() + 1); - } + - } diff --git a/vector/src/main/codegen/templates/UnionListWriter.java b/vector/src/main/codegen/templates/UnionListWriter.java index 4b54739230..ec634fa4c7 100644 --- a/vector/src/main/codegen/templates/UnionListWriter.java +++ b/vector/src/main/codegen/templates/UnionListWriter.java @@ -123,8 +123,6 @@ public void setPosition(int index) { <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> <#assign upperName = minor.class?upper_case /> - <#assign capName = minor.class?cap_first /> - <#assign vectName = capName /> @Override public ${minor.class}Writer ${lowerName}() { return this; diff --git a/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java index 73a88b3a1e..de7c0c1b80 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java @@ -30,14 +30,20 @@ import java.util.Arrays; import java.util.List; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListReader; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.holders.DurationHolder; +import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; +import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; +import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.TransferPair; @@ -628,6 +634,163 @@ public void testWriteLargeVarBinaryHelpers() throws Exception { } } + @Test + public void testWriterTimeStampNanoTZField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + final int valueCount = 10; + + for (int i = 0; i < valueCount; i++) { + writer.startList(); + writer.timeStampNanoTZ().writeTimeStampNanoTZ(i * 1000L); + writer.timeStampNanoTZ().writeTimeStampNanoTZ((i + 1) * 1000L); + writer.timeStampNanoTZ().writeTimeStampNanoTZ((i + 2) * 1000L); + writer.endList(); + } + vector.setValueCount(valueCount); + + UnionFixedSizeListReader reader = vector.getReader(); + for (int i = 0; i < valueCount; i++) { + reader.setPosition(i); + assertTrue(reader.isSet()); + assertTrue(reader.next()); + assertEquals(i * 1000L, reader.reader().readLong().longValue()); + assertTrue(reader.next()); + assertEquals((i + 1) * 1000L, reader.reader().readLong().longValue()); + assertTrue(reader.next()); + assertEquals((i + 2) * 1000L, reader.reader().readLong().longValue()); + assertFalse(reader.next()); + } + } + } + + @Test + public void testWriterUsingHolderTimestampMilliTZField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); + holder.timezone = "SomeFakeTimeZone"; + writer.startList(); + holder.value = 12341234L; + writer.timeStampMilliTZ().write(holder); + holder.value = 55555L; + writer.timeStampMilliTZ().write(holder); + + // Writing with a different timezone should throw + holder.timezone = "AsdfTimeZone"; + holder.value = 77777; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, () -> writer.timeStampMilliTZ().write(holder)); + assertEquals( + "holder.timezone: AsdfTimeZone not equal to vector timezone: SomeFakeTimeZone", + ex.getMessage()); + + writer.endList(); + vector.setValueCount(1); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "SomeFakeTimeZone")), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(3)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + + @Test + public void testWriterUsingHolderDurationField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + DurationHolder durationHolder = new DurationHolder(); + durationHolder.unit = TimeUnit.MILLISECOND; + + writer.startList(); + durationHolder.value = 812374L; + writer.duration().write(durationHolder); + durationHolder.value = 143451L; + writer.duration().write(durationHolder); + + // Writing with a different unit should throw + durationHolder.unit = TimeUnit.SECOND; + durationHolder.value = 8888888; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, () -> writer.duration().write(durationHolder)); + assertEquals("holder.unit: SECOND not equal to vector unit: MILLISECOND", ex.getMessage()); + + writer.endList(); + vector.setValueCount(1); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.Duration(TimeUnit.MILLISECOND)), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(3)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + + @Test + public void testWriterUsingHolderFixedSizeBinaryField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 2, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + FixedSizeBinaryHolder holder1 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + FixedSizeBinaryHolder holder2 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + + writer.startList(); + writer.fixedSizeBinary().write(holder1); + holder1.buffer.close(); + writer.fixedSizeBinary().write(holder2); + holder2.buffer.close(); + + writer.endList(); + vector.setValueCount(1); + + FieldReader reader = vector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.FixedSizeBinary(2)), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(2)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + private int[] convertListToIntArray(List list) { int[] values = new int[list.size()]; for (int i = 0; i < list.size(); i++) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 274d2973bd..2f520f3882 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -1353,17 +1353,6 @@ public void testCopyFromForExtensionType() throws Exception { } } - private FixedSizeBinaryHolder getFixedSizeBinaryHolder(byte[] array) { - FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); - holder.byteWidth = array.length; - holder.buffer = allocator.buffer(array.length); - for (int i = 0; i < array.length; i++) { - holder.buffer.setByte(i, array[i]); - } - - return holder; - } - /** * Regression test for GH-586: UnionMapWriter.fixedSizeBinary() should properly delegate to the * entry writer for both key and value paths. @@ -1382,8 +1371,10 @@ public void testFixedSizeBinaryWriter() { // {[11, 22] -> null} // {null -> [32, 21]} - wrong "for a given entry, the "key" is non-nullable" - todo: it // shouldn't work. Should it? - FixedSizeBinaryHolder holder1 = getFixedSizeBinaryHolder(new byte[] {11, 22}); - FixedSizeBinaryHolder holder2 = getFixedSizeBinaryHolder(new byte[] {32, 21}); + FixedSizeBinaryHolder holder1 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + FixedSizeBinaryHolder holder2 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); writer.setPosition(0); // optional writer.startMap(); @@ -1399,8 +1390,8 @@ public void testFixedSizeBinaryWriter() { writer.endMap(); // {1 -> [11, 22], 2 -> [32, 21]} - holder1 = getFixedSizeBinaryHolder(new byte[] {11, 22}); - holder2 = getFixedSizeBinaryHolder(new byte[] {32, 21}); + holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); writer.setPosition(1); writer.startMap(); writer.startEntry(); @@ -1416,8 +1407,8 @@ public void testFixedSizeBinaryWriter() { holder2.buffer.close(); // {[11, 22] -> 1, [32, 21] -> 2} - holder1 = getFixedSizeBinaryHolder(new byte[] {11, 22}); - holder2 = getFixedSizeBinaryHolder(new byte[] {32, 21}); + holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); writer.setPosition(3); writer.startMap(); writer.startEntry(); @@ -1433,7 +1424,7 @@ public void testFixedSizeBinaryWriter() { holder2.buffer.close(); // {[11, 22] -> null} - holder1 = getFixedSizeBinaryHolder(new byte[] {11, 22}); + holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); writer.setPosition(4); writer.startMap(); writer.startEntry(); @@ -1443,7 +1434,7 @@ public void testFixedSizeBinaryWriter() { holder1.buffer.close(); // {null -> [32, 21]} - holder2 = getFixedSizeBinaryHolder(new byte[] {32, 21}); + holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); writer.setPosition(5); writer.startMap(); writer.startEntry(); @@ -1536,8 +1527,10 @@ public void testFixedSizeBinaryFirstInitialization() { // populate input vector with the following records // {[11, 22] -> [32, 21]} - FixedSizeBinaryHolder holder1 = getFixedSizeBinaryHolder(new byte[] {11, 22}); - FixedSizeBinaryHolder holder2 = getFixedSizeBinaryHolder(new byte[] {32, 21}); + FixedSizeBinaryHolder holder1 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + FixedSizeBinaryHolder holder2 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); writer.setPosition(0); // optional writer.startMap(); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java index c28751aa58..d91b2004c0 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java @@ -18,6 +18,7 @@ import java.util.Random; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; @@ -73,4 +74,16 @@ public static void ensureRegistered(ArrowType.ExtensionType type) { ExtensionTypeRegistry.register(type); } } + + public static FixedSizeBinaryHolder fixedSizeBinaryHolder( + BufferAllocator allocator, byte[] array) { + FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); + holder.byteWidth = array.length; + holder.buffer = allocator.buffer(array.length); + for (int i = 0; i < array.length; i++) { + holder.buffer.setByte(i, array[i]); + } + + return holder; + } }