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..bbdb492fe9 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 @@ -36,6 +36,7 @@ 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.FieldType; @@ -954,4 +955,66 @@ public void testCopyStructVectorWithExtensionType() { assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); } } + + @Test + public void testCopyListOfTimeStampNanoTZ() { + 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); + + 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); + + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } + + @Test + public void testCopyStructOfTimeStampNanoTZ() { + 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); + + 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); + + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } }