Merge branch 'apache:main' into main

Author: vandop

adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java

@@ -24,7 +24,6 @@
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
-import org.apache.arrow.vector.types.pojo.ArrowType;
 
 /** Composite consumer which hold all consumers. It manages the consume and cleanup process. */
 public class CompositeJdbcConsumer implements JdbcConsumer {
@@ -46,9 +45,9 @@ public void consume(ResultSet rs) throws SQLException, IOException {
           BaseConsumer consumer = (BaseConsumer) consumers[i];
           JdbcFieldInfo fieldInfo =
               new JdbcFieldInfo(rs.getMetaData(), consumer.columnIndexInResultSet);
-          ArrowType arrowType = consumer.vector.getMinorType().getType();
+
           throw new JdbcConsumerException(
-              "Exception while consuming JDBC value", e, fieldInfo, arrowType);
+              "Exception while consuming JDBC value", e, fieldInfo, consumer.vector.getField());
         } else {
           throw e;
         }

adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java

@@ -17,33 +17,33 @@
 package org.apache.arrow.adapter.jdbc.consumer.exceptions;
 
 import org.apache.arrow.adapter.jdbc.JdbcFieldInfo;
-import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
 
 /**
  * Exception while consuming JDBC data. This exception stores the JdbcFieldInfo for the column and
  * the ArrowType for the corresponding vector for easier debugging.
  */
 public class JdbcConsumerException extends RuntimeException {
   final JdbcFieldInfo fieldInfo;
-  final ArrowType arrowType;
+  final Field field;
 
   /**
    * Construct JdbcConsumerException with all fields.
    *
    * @param message error message
    * @param cause original exception
    * @param fieldInfo JdbcFieldInfo for the column
-   * @param arrowType ArrowType for the corresponding vector
+   * @param field ArrowType for the corresponding vector
    */
   public JdbcConsumerException(
-      String message, Throwable cause, JdbcFieldInfo fieldInfo, ArrowType arrowType) {
+      String message, Throwable cause, JdbcFieldInfo fieldInfo, Field field) {
     super(message, cause);
     this.fieldInfo = fieldInfo;
-    this.arrowType = arrowType;
+    this.field = field;
   }
 
-  public ArrowType getArrowType() {
-    return this.arrowType;
+  public Field getField() {
+    return this.field;
   }
 
   public JdbcFieldInfo getFieldInfo() {

vector/src/main/codegen/includes/vv_imports.ftl

@@ -34,6 +34,7 @@ import org.apache.arrow.vector.complex.*;
 import org.apache.arrow.vector.complex.reader.*;
 import org.apache.arrow.vector.complex.impl.*;
 import org.apache.arrow.vector.complex.writer.*;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;

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

@@ -109,6 +109,10 @@ public void copyAsField(String name, ${name}Writer writer) {
 
   </#list></#list>
 
+  public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory) {
+    fail("CopyAsValue StructWriter");
+  }
+
   public void read(ExtensionHolder holder) {
     fail("Extension");
   }

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

@@ -49,6 +49,7 @@ public interface RepeatedStructReader extends StructReader{
     boolean next();
     int size();
     void copyAsValue(StructWriter writer);
+    void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory);
   }
 
   public interface ListReader extends BaseReader{
@@ -59,6 +60,7 @@ public interface RepeatedListReader extends ListReader{
     boolean next();
     int size();
     void copyAsValue(ListWriter writer);
+    void copyAsValue(ListWriter writer, ExtensionTypeWriterFactory writerFactory);
   }
 
   public interface MapReader extends BaseReader{
@@ -69,6 +71,7 @@ public interface RepeatedMapReader extends MapReader{
     boolean next();
     int size();
     void copyAsValue(MapWriter writer);
+    void copyAsValue(MapWriter writer, ExtensionTypeWriterFactory writerFactory);
   }
 
   public interface ScalarReader extends

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

@@ -42,10 +42,14 @@ public class ComplexCopier {
    * @param output field to write to
    */
   public static void copy(FieldReader input, FieldWriter output) {
-    writeValue(input, output);
+    writeValue(input, output, null);
   }
 
-  private static void writeValue(FieldReader reader, FieldWriter writer) {
+  public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
+    writeValue(input, output, extensionTypeWriterFactory);
+  }
+
+  private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
     final MinorType mt = reader.getMinorType();
 
       switch (mt) {
@@ -61,7 +65,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
             FieldReader childReader = reader.reader();
             FieldWriter childWriter = getListWriterForReader(childReader, writer);
             if (childReader.isSet()) {
-              writeValue(childReader, childWriter);
+              writeValue(childReader, childWriter, extensionTypeWriterFactory);
             } else {
               childWriter.writeNull();
             }
@@ -79,8 +83,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
             FieldReader structReader = reader.reader();
             if (structReader.isSet()) {
               writer.startEntry();
-              writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()));
-              writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()));
+              writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()), extensionTypeWriterFactory);
+              writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()), extensionTypeWriterFactory);
               writer.endEntry();
             } else {
               writer.writeNull();
@@ -99,7 +103,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
             if (childReader.getMinorType() != Types.MinorType.NULL) {
               FieldWriter childWriter = getStructWriterForReader(childReader, writer, name);
               if (childReader.isSet()) {
-                writeValue(childReader, childWriter);
+                writeValue(childReader, childWriter, extensionTypeWriterFactory);
               } else {
                 childWriter.writeNull();
               }
@@ -110,6 +114,20 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
           writer.writeNull();
         }
         break;
+      case EXTENSIONTYPE:
+        if (extensionTypeWriterFactory == null) {
+          throw new IllegalArgumentException("Must provide ExtensionTypeWriterFactory");
+        }
+        if (reader.isSet()) {
+          Object value = reader.readObject();
+          if (value != null) {
+            writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory);
+            writer.writeExtension(value);
+          }
+        } else {
+          writer.writeNull();
+        }
+        break;
   <#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/>
@@ -162,6 +180,9 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr
       return (FieldWriter) writer.map(name);
     case LISTVIEW:
       return (FieldWriter) writer.listView(name);
+    case EXTENSIONTYPE:
+      ExtensionWriter extensionWriter = writer.extension(name, reader.getField().getType());
+      return (FieldWriter) extensionWriter;
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());
     }
@@ -186,6 +207,9 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
       return (FieldWriter) writer.list();
     case LISTVIEW:
       return (FieldWriter) writer.listView();
+    case EXTENSIONTYPE:
+      ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+      return (FieldWriter) extensionWriter;
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());
     }
@@ -211,6 +235,9 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w
         return (FieldWriter) writer.listView();
       case MAP:
         return (FieldWriter) writer.map(false);
+      case EXTENSIONTYPE:
+        ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+        return (FieldWriter) extensionWriter;
       default:
         throw new UnsupportedOperationException(reader.getMinorType().toString());
     }

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

@@ -86,6 +86,7 @@ public void read(int arrayIndex, Nullable${name}Holder holder){
   }
   </#list></#list>
 
+  public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory){}
   public void read(ExtensionHolder holder) {
     holder.isSet = 0;
   }

vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java

@@ -22,6 +22,7 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.ReferenceManager;
 import org.apache.arrow.util.Preconditions;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.util.DataSizeRoundingUtil;
 import org.apache.arrow.vector.util.TransferPair;
@@ -260,6 +261,18 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public void copyFrom(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void copyFromSafe(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * Transfer the validity buffer from `validityBuffer` to the target vector's `validityBuffer`.
    * Start at `startIndex` and copy `length` number of elements. If the starting index is 8 byte

vector/src/main/java/org/apache/arrow/vector/NullVector.java

@@ -27,6 +27,7 @@
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.compare.VectorVisitor;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.impl.NullReader;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
@@ -329,6 +330,18 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public void copyFrom(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void copyFromSafe(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public String getName() {
     return this.getField().getName();

vector/src/main/java/org/apache/arrow/vector/ValueVector.java

@@ -22,6 +22,7 @@
 import org.apache.arrow.memory.OutOfMemoryException;
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
 import org.apache.arrow.vector.compare.VectorVisitor;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.Field;
@@ -309,6 +310,30 @@ public interface ValueVector extends Closeable, Iterable<ValueVector> {
    */
   void copyFromSafe(int fromIndex, int thisIndex, ValueVector from);
 
+  /**
+   * Copy a cell value from a particular index in source vector to a particular position in this
+   * vector.
+   *
+   * @param fromIndex position to copy from in source vector
+   * @param thisIndex position to copy to in this vector
+   * @param from source vector
+   * @param writerFactory the extension type writer factory to use for copying extension type values
+   */
+  void copyFrom(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory);
+
+  /**
+   * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the
+   * capacity of the vector needs to be expanded before copy.
+   *
+   * @param fromIndex position to copy from in source vector
+   * @param thisIndex position to copy to in this vector
+   * @param from source vector
+   * @param writerFactory the extension type writer factory to use for copying extension type values
+   */
+  void copyFromSafe(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory);
+
   /**
    * Accept a generic {@link VectorVisitor} and return the result.
    *