GH-551: [Java] Add TimeStamp*TZ support to ComplexCopier

ComplexCopier skipped types with typeParams except Decimal.
Generate TimeStamp*TZ cases and pass timezone to the struct writer.

Author: AzazelSensei

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")>
+</#function>
+
 /*
  * 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>
+    <#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);
+      }
+    </#if>
 
     </#list></#list>
     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}</#if>();
     </#if>
@@ -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}</#if>();
     </#if>

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));
+    }
+  }
 }