diff --git a/vector/src/main/codegen/templates/ComplexCopier.java b/vector/src/main/codegen/templates/ComplexCopier.java index 6655f6c2a7..8152cc646f 100644 --- a/vector/src/main/codegen/templates/ComplexCopier.java +++ b/vector/src/main/codegen/templates/ComplexCopier.java @@ -30,6 +30,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. */ @@ -121,7 +125,7 @@ public static void copy(FieldReader reader, FieldWriter writer) { <#assign fields = minor.fields!type.fields /> <#assign uncappedName = name?uncap_first/> - <#if !minor.typeParams?? || minor.class?starts_with("Decimal") > + <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) > case ${name?upper_case}: if (reader.isSet()) { @@ -158,6 +162,15 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr return (FieldWriter) writer.${uncappedName}(name); } + <#if is_timestamp_tz(minor.class)> + case ${name?upper_case}: + if (reader.getField().getType() instanceof ArrowType.Timestamp) { + ArrowType.Timestamp type = (ArrowType.Timestamp) reader.getField().getType(); + return (FieldWriter) writer.${uncappedName}(name, type.getTimezone()); + } else { + return (FieldWriter) writer.${uncappedName}(name); + } + case STRUCT: @@ -182,7 +195,7 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter <#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.typeParams?? || minor.class?starts_with("Decimal") > + <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) > case ${name?upper_case}: return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}(); @@ -209,7 +222,7 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w <#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.typeParams?? || minor.class?starts_with("Decimal") > + <#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) > case ${name?upper_case}: return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}(); diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java index b2a8cf9ba4..998aa07183 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java +++ b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java @@ -16,10 +16,13 @@ */ package org.apache.arrow.vector.complex.impl; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; +import java.util.Arrays; import java.util.UUID; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -36,8 +39,10 @@ import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.DecimalHolder; +import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types; 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.DecimalUtility; import org.junit.jupiter.api.AfterEach; @@ -954,4 +959,125 @@ public void testCopyStructVectorWithExtensionType() { assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); } } + + @Test + public void testCopyListOfTimeStampNanoTZVector() { + try (ListVector from = ListVector.empty("v", allocator); + ListVector to = ListVector.empty("v", allocator)) { + from.addOrGetVector(FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"))); + to.addOrGetVector(FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"))); + + UnionListWriter listWriter = from.getWriter(); + listWriter.allocate(); + for (int i = 0; i < COUNT; i++) { + listWriter.setPosition(i); + listWriter.startList(); + listWriter.timeStampNanoTZ().writeTimeStampNanoTZ(i * 1_000_000L); + listWriter.timeStampNanoTZ().writeTimeStampNanoTZ(i * 2_000_000L); + listWriter.endList(); + } + from.setValueCount(COUNT); + + // copy values + FieldReader in = from.getReader(); + FieldWriter out = to.getWriter(); + for (int i = 0; i < COUNT; i++) { + in.setPosition(i); + out.setPosition(i); + ComplexCopier.copy(in, out); + } + to.setValueCount(COUNT); + + // validate equals + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } + + @Test + public void testCopyStructOfTimeStampNanoTZVector() { + try (final StructVector from = StructVector.empty("v", allocator); + final StructVector to = StructVector.empty("v", allocator)) { + from.allocateNewSafe(); + + NullableStructWriter structWriter = from.getWriter(); + for (int i = 0; i < COUNT; i++) { + structWriter.setPosition(i); + structWriter.start(); + structWriter.timeStampNanoTZ("ts", "UTC").writeTimeStampNanoTZ(i * 1_000_000L); + structWriter.end(); + } + from.setValueCount(COUNT); + + // copy values + FieldReader in = from.getReader(); + FieldWriter out = to.getWriter(); + for (int i = 0; i < COUNT; i++) { + in.setPosition(i); + out.setPosition(i); + ComplexCopier.copy(in, out); + } + to.setValueCount(COUNT); + + // validate equals + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } + + @Test + public void testCopyMapOfTimeStampNanoTZVector() { + // Before the fix ComplexCopier.getMapWriterForReader threw + // UnsupportedOperationException for a TimeStamp*TZ map value. The map value + // child is pre-typed so the reader reports TIMESTAMPNANOTZ and the map case + // is actually exercised. vectorEquals is not asserted because the MapWriter + // scalar API has no timezone overload (unlike StructWriter), so the copied + // value child does not round-trip to the same timezone; this test locks in + // that copying the map no longer throws. + Field key = new Field("key", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field value = + new Field( + "value", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), + null); + Field entries = + new Field( + MapVector.DATA_VECTOR_NAME, + FieldType.notNullable(ArrowType.Struct.INSTANCE), + Arrays.asList(key, value)); + Field mapField = + new Field("v", FieldType.nullable(new ArrowType.Map(false)), Arrays.asList(entries)); + + try (final MapVector from = (MapVector) mapField.createVector(allocator); + final MapVector to = MapVector.empty("v", allocator, false)) { + + from.allocateNew(); + + UnionMapWriter mapWriter = from.getWriter(); + for (int i = 0; i < COUNT; i++) { + mapWriter.setPosition(i); + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().integer().writeInt(i); + mapWriter.value().timeStampNanoTZ().writeTimeStampNanoTZ(i * 1_000_000L); + mapWriter.endEntry(); + mapWriter.endMap(); + } + + from.setValueCount(COUNT); + + // copy values + FieldReader in = from.getReader(); + FieldWriter out = to.getWriter(); + assertDoesNotThrow( + () -> { + for (int i = 0; i < COUNT; i++) { + in.setPosition(i); + out.setPosition(i); + ComplexCopier.copy(in, out); + } + }); + to.setValueCount(COUNT); + + assertEquals(from.getValueCount(), to.getValueCount()); + } + } }