GH-891: Add ExtensionTypeWriterFactory to TransferPair

Author: jhrotko

vector/src/main/codegen/templates/ComplexCopier.java

@@ -63,7 +63,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
           writer.startList();
           while (reader.next()) {
             FieldReader childReader = reader.reader();
-            FieldWriter childWriter = getListWriterForReader(childReader, writer);
+            FieldWriter childWriter = getListWriterForReader(childReader, writer, extensionTypeWriterFactory);
             if (childReader.isSet()) {
               writeValue(childReader, childWriter, extensionTypeWriterFactory);
             } else {
@@ -189,6 +189,10 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr
   }
 
   private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter writer) {
+    return getListWriterForReader(reader, writer, null);
+  }
+
+  private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
     switch (reader.getMinorType()) {
     <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
     <#assign fields = minor.fields!type.fields />
@@ -209,6 +213,9 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
       return (FieldWriter) writer.listView();
     case EXTENSIONTYPE:
       ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+      if (extensionTypeWriterFactory != null) {
+        extensionWriter.addExtensionTypeWriterFactory(extensionTypeWriterFactory);
+      }
       return (FieldWriter) extensionWriter;
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());

vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java

@@ -63,6 +63,7 @@
 import org.apache.arrow.vector.util.OversizedAllocationException;
 import org.apache.arrow.vector.util.SchemaChangeRuntimeException;
 import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.util.TransferPairWithExtendedType;
 
 /**
  * A list vector contains lists of a specific type of elements. Its structure contains 3 elements.
@@ -648,7 +649,7 @@ public UnionVector promoteToUnion() {
     return vector;
   }
 
-  private class TransferImpl implements TransferPair {
+  private class TransferImpl implements TransferPairWithExtendedType {
 
     LargeListVector to;
     TransferPair dataTransferPair;
@@ -731,6 +732,12 @@ public ValueVector getTo() {
     public void copyValueSafe(int from, int to) {
       this.to.copyFrom(from, to, LargeListVector.this);
     }
+
+    @Override
+    public void copyValueSafe(
+        int from, int to, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
+      this.to.copyFrom(from, to, LargeListVector.this, extensionTypeWriterFactory);
+    }
   }
 
   @Override

vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java

@@ -56,6 +56,7 @@
 import org.apache.arrow.vector.util.JsonStringArrayList;
 import org.apache.arrow.vector.util.OversizedAllocationException;
 import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.util.TransferPairWithExtendedType;
 
 /**
  * A list vector contains lists of a specific type of elements. Its structure contains 3 elements.
@@ -528,7 +529,7 @@ public <OUT, IN> OUT accept(VectorVisitor<OUT, IN> visitor, IN value) {
     return visitor.visit(this, value);
   }
 
-  private class TransferImpl implements TransferPair {
+  private class TransferImpl implements TransferPairWithExtendedType {
 
     ListVector to;
     TransferPair dataTransferPair;
@@ -612,6 +613,12 @@ public ValueVector getTo() {
     public void copyValueSafe(int from, int to) {
       this.to.copyFrom(from, to, ListVector.this);
     }
+
+    @Override
+    public void copyValueSafe(
+        int from, int to, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
+      this.to.copyFrom(from, to, ListVector.this, extensionTypeWriterFactory);
+    }
   }
 
   @Override

vector/src/main/java/org/apache/arrow/vector/util/TransferPairWithExtendedType.java

@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector.util;
+
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
+
+public interface TransferPairWithExtendedType extends TransferPair {
+  void copyValueSafe(int from, int to, ExtensionTypeWriterFactory extensionTypeWriterFactory);
+}

vector/src/test/java/org/apache/arrow/vector/TestListVector.java

@@ -1314,6 +1314,84 @@ public void testCopyFromForExtensionType() throws Exception {
     }
   }
 
+  @Test
+  public void testCopyValueSafeForExtensionType() throws Exception {
+    try (ListVector inVector = ListVector.empty("input", allocator);
+        ListVector outVector = ListVector.empty("output", allocator)) {
+      UnionListWriter writer = inVector.getWriter();
+      writer.allocate();
+
+      // Create first list with UUIDs
+      writer.setPosition(0);
+      UUID u1 = UUID.randomUUID();
+      UUID u2 = UUID.randomUUID();
+      writer.startList();
+      ExtensionWriter extensionWriter = writer.extension(new UuidType());
+      extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
+      extensionWriter.writeExtension(u1);
+      extensionWriter.writeExtension(u2);
+      writer.endList();
+
+      // Create second list with UUIDs
+      writer.setPosition(1);
+      UUID u3 = UUID.randomUUID();
+      UUID u4 = UUID.randomUUID();
+      writer.startList();
+      extensionWriter = writer.extension(new UuidType());
+      extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory());
+      extensionWriter.writeExtension(u3);
+      extensionWriter.writeExtension(u4);
+      extensionWriter.writeNull();
+
+      writer.endList();
+      writer.setValueCount(2);
+
+      // Use copyFromSafe with ExtensionTypeWriterFactory
+      // This internally calls TransferImpl.copyValueSafe with ExtensionTypeWriterFactory
+      outVector.allocateNew();
+      outVector.copyFromSafe(0, 0, inVector, new UuidWriterFactory());
+      outVector.copyFromSafe(1, 1, inVector, new UuidWriterFactory());
+      outVector.setValueCount(2);
+
+      // Verify first list
+      UnionListReader reader = outVector.getReader();
+      reader.setPosition(0);
+      assertTrue(reader.isSet(), "first list shouldn't be null");
+      reader.next();
+      FieldReader uuidReader = reader.reader();
+      UuidHolder holder = new UuidHolder();
+      uuidReader.read(holder);
+      ByteBuffer bb = ByteBuffer.wrap(holder.value);
+      UUID actualUuid = new UUID(bb.getLong(), bb.getLong());
+      assertEquals(u1, actualUuid);
+      reader.next();
+      uuidReader = reader.reader();
+      uuidReader.read(holder);
+      bb = ByteBuffer.wrap(holder.value);
+      actualUuid = new UUID(bb.getLong(), bb.getLong());
+      assertEquals(u2, actualUuid);
+
+      // Verify second list
+      reader.setPosition(1);
+      assertTrue(reader.isSet(), "second list shouldn't be null");
+      reader.next();
+      uuidReader = reader.reader();
+      uuidReader.read(holder);
+      bb = ByteBuffer.wrap(holder.value);
+      actualUuid = new UUID(bb.getLong(), bb.getLong());
+      assertEquals(u3, actualUuid);
+      reader.next();
+      uuidReader = reader.reader();
+      uuidReader.read(holder);
+      bb = ByteBuffer.wrap(holder.value);
+      actualUuid = new UUID(bb.getLong(), bb.getLong());
+      assertEquals(u4, actualUuid);
+      reader.next();
+      uuidReader = reader.reader();
+      assertFalse(uuidReader.isSet(), "third element should be null");
+    }
+  }
+
   private void writeIntValues(UnionListWriter writer, int[] values) {
     writer.startList();
     for (int v : values) {