apacheGH-38246 added new getTransferPair with Field argument for Comp…

…lex type
bookingnizer · Oct 12, 2023 · c06ac68 · c06ac68
1 parent 14abb2c
commit c06ac68
Show file tree

Hide file tree

Showing 18 changed files with 436 additions and 48 deletions.
diff --git a/java/vector/src/main/codegen/templates/DenseUnionVector.java b/java/vector/src/main/codegen/templates/DenseUnionVector.java
@@ -115,7 +115,7 @@ public class DenseUnionVector extends AbstractContainerVector implements FieldVe
   private long typeBufferAllocationSizeInBytes;
   private long offsetBufferAllocationSizeInBytes;
 
-  private final FieldType fieldType;
+  private final Field field;
 
   public static final byte TYPE_WIDTH = 1;
   public static final byte OFFSET_WIDTH = 4;
@@ -131,7 +131,23 @@ public static DenseUnionVector empty(String name, BufferAllocator allocator) {
 
   public DenseUnionVector(String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) {
     super(name, allocator, callBack);
-    this.fieldType = fieldType;
+    this.field = new Field(name, fieldType, null);
+    this.internalStruct = new NonNullableStructVector(
+        "internal",
+        allocator,
+        INTERNAL_STRUCT_TYPE,
+        callBack,
+        AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE,
+        false);
+    this.typeBuffer = allocator.getEmpty();
+    this.typeBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH;
+    this.offsetBuffer = allocator.getEmpty();
+    this.offsetBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH;
+  }
+
+  public DenseUnionVector(Field field, BufferAllocator allocator, CallBack callBack) {
+    super(field.getName(), allocator, callBack);
+    this.field = field;
     this.internalStruct = new NonNullableStructVector(
         "internal",
         allocator,
@@ -234,8 +250,8 @@ public synchronized byte registerNewTypeId(Field field) {
               typeFields.length + " relative types. Please use union of union instead");
     }
     byte typeId = nextTypeId;
-    if (fieldType != null) {
-      int[] typeIds = ((ArrowType.Union) fieldType.getType()).getTypeIds();
+    if (field.getFieldType() != null) {
+      int[] typeIds = ((ArrowType.Union) field.getFieldType().getType()).getTypeIds();
       if (typeIds != null) {
         int thisTypeId = typeIds[nextTypeId];
         if (thisTypeId > Byte.MAX_VALUE) {
@@ -528,12 +544,12 @@ public Field getField() {
     }
 
     FieldType fieldType;
-    if (this.fieldType == null) {
+    if (field.getFieldType() == null) {
       fieldType = FieldType.nullable(new ArrowType.Union(Dense, typeIds));
     } else {
       final UnionMode mode = UnionMode.Dense;
-      fieldType = new FieldType(this.fieldType.isNullable(), new ArrowType.Union(mode, typeIds),
-              this.fieldType.getDictionary(), this.fieldType.getMetadata());
+      fieldType = new FieldType(field.getFieldType().isNullable(), new ArrowType.Union(mode, typeIds),
+          field.getFieldType().getDictionary(), field.getFieldType().getMetadata());
     }
 
     return new Field(name, fieldType, childFields);
@@ -554,6 +570,16 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallB
     return new org.apache.arrow.vector.complex.DenseUnionVector.TransferImpl(ref, allocator, callBack);
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+    return getTransferPair(field, allocator, null);
+  }
+
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return new org.apache.arrow.vector.complex.DenseUnionVector.TransferImpl(field, allocator, callBack);
+  }
+
   @Override
   public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((DenseUnionVector) target);
@@ -598,6 +624,12 @@ public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) {
       createTransferPairs();
     }
 
+    public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) {
+      to = new DenseUnionVector(field.getName(), allocator, null, callBack);
+      internalStruct.makeTransferPair(to.internalStruct);
+      createTransferPairs();
+    }
+
     public TransferImpl(DenseUnionVector to) {
       this.to = to;
       internalStruct.makeTransferPair(to.internalStruct);

diff --git a/java/vector/src/main/codegen/templates/UnionVector.java b/java/vector/src/main/codegen/templates/UnionVector.java
@@ -103,7 +103,7 @@ public class UnionVector extends AbstractContainerVector implements FieldVector
 
   private int typeBufferAllocationSizeInBytes;
 
-  private final FieldType fieldType;
+  private final Field field;
   private final Field[] typeIds = new Field[Byte.MAX_VALUE + 1];
 
   public static final byte TYPE_WIDTH = 1;
@@ -118,7 +118,21 @@ public static UnionVector empty(String name, BufferAllocator allocator) {
 
   public UnionVector(String name, BufferAllocator allocator, FieldType fieldType, CallBack callBack) {
     super(name, allocator, callBack);
-    this.fieldType = fieldType;
+    this.field = new Field(name, fieldType, null);
+    this.internalStruct = new NonNullableStructVector(
+        "internal",
+        allocator,
+        INTERNAL_STRUCT_TYPE,
+        callBack,
+        AbstractStructVector.ConflictPolicy.CONFLICT_REPLACE,
+        false);
+    this.typeBuffer = allocator.getEmpty();
+    this.typeBufferAllocationSizeInBytes = BaseValueVector.INITIAL_VALUE_ALLOCATION * TYPE_WIDTH;
+  }
+
+  public UnionVector(Field field, BufferAllocator allocator, CallBack callBack) {
+    super(field.getName(), allocator, callBack);
+    this.field = field;
     this.internalStruct = new NonNullableStructVector(
         "internal",
         allocator,
@@ -144,8 +158,8 @@ public void initializeChildrenFromFields(List<Field> children) {
     int count = 0;
     for (Field child: children) {
       int typeId = Types.getMinorTypeForArrowType(child.getType()).ordinal();
-      if (fieldType != null) {
-        int[] typeIds = ((ArrowType.Union)fieldType.getType()).getTypeIds();
+      if (field.getFieldType() != null) {
+        int[] typeIds = ((ArrowType.Union)field.getFieldType().getType()).getTypeIds();
         if (typeIds != null) {
           typeId = typeIds[count++];
         }
@@ -469,12 +483,12 @@ public Field getField() {
     }
 
     FieldType fieldType;
-    if (this.fieldType == null) {
+    if (field.getFieldType() == null) {
       fieldType = FieldType.nullable(new ArrowType.Union(Sparse, typeIds));
     } else {
-      final UnionMode mode = ((ArrowType.Union)this.fieldType.getType()).getMode();
-      fieldType = new FieldType(this.fieldType.isNullable(), new ArrowType.Union(mode, typeIds),
-          this.fieldType.getDictionary(), this.fieldType.getMetadata());
+      final UnionMode mode = ((ArrowType.Union)field.getFieldType().getType()).getMode();
+      fieldType = new FieldType(field.getFieldType().isNullable(), new ArrowType.Union(mode, typeIds),
+          field.getFieldType().getDictionary(), field.getFieldType().getMetadata());
     }
 
     return new Field(name, fieldType, childFields);
@@ -495,6 +509,16 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallB
     return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(ref, allocator, callBack);
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+    return getTransferPair(field, allocator, null);
+  }
+
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(field, allocator, callBack);
+  }
+
   @Override
   public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((UnionVector) target);
@@ -547,6 +571,11 @@ public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) {
       internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct);
     }
 
+    public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) {
+      to = new UnionVector(field.getName(), allocator, null, callBack);
+      internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct);
+    }
+
     public TransferImpl(UnionVector to) {
       this.to = to;
       internalStructVectorTransferPair = internalStruct.makeTransferPair(to.internalStruct);

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java
@@ -569,6 +569,18 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallB
     return getTransferPair(ref, allocator);
   }
 
+  /**
+   * Construct a transfer pair of this vector and another vector of same type.
+   * @param field The field materialized by this vector.
+   * @param allocator allocator for the target vector
+   * @param callBack not used
+   * @return TransferPair
+   */
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return getTransferPair(field, allocator);
+  }
+
   /**
    * Construct a transfer pair of this vector and another vector of same type.
    * @param allocator allocator for the target vector

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java
@@ -662,6 +662,18 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallB
     return getTransferPair(ref, allocator);
   }
 
+  /**
+   * Construct a transfer pair of this vector and another vector of same type.
+   * @param field The field materialized by this vector
+   * @param allocator allocator for the target vector
+   * @param callBack not used
+   * @return TransferPair
+   */
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return getTransferPair(field, allocator);
+  }
+
   /**
    * Construct a transfer pair of this vector and another vector of same type.
    * @param allocator allocator for the target vector
@@ -672,6 +684,7 @@ public TransferPair getTransferPair(BufferAllocator allocator) {
     return getTransferPair(getName(), allocator);
   }
 
+
   /**
    * Construct a transfer pair of this vector and another vector of same type.
    * @param ref name of the target vector
@@ -680,6 +693,14 @@ public TransferPair getTransferPair(BufferAllocator allocator) {
    */
   public abstract TransferPair getTransferPair(String ref, BufferAllocator allocator);
 
+  /**
+   * Construct a transfer pair of this vector and another vector of same type.
+   * @param field The field materialized by this vector
+   * @param allocator allocator for the target vector
+   * @return TransferPair
+   */
+  public abstract TransferPair getTransferPair(Field field, BufferAllocator allocator);
+
   /**
    * Transfer this vector'data to another vector. The memory associated
    * with this vector is transferred to the allocator of target vector

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java
@@ -692,6 +692,18 @@ public void validateScalars() {
     // No validation by default.
   }
 
+  /**
+   * Construct a transfer pair of this vector and another vector of same type.
+   * @param field The field materialized by this vector.
+   * @param allocator allocator for the target vector
+   * @param callBack not used
+   * @return TransferPair
+   */
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return getTransferPair(field, allocator);
+  }
+
   /**
    * Construct a transfer pair of this vector and another vector of same type.
    * @param ref name of the target vector
@@ -722,6 +734,14 @@ public TransferPair getTransferPair(BufferAllocator allocator) {
    */
   public abstract TransferPair getTransferPair(String ref, BufferAllocator allocator);
 
+  /**
+   * Construct a transfer pair of this vector and another vector of same type.
+   * @param field The field materialized by this vector.
+   * @param allocator allocator for the target vector
+   * @return TransferPair
+   */
+  public abstract TransferPair getTransferPair(Field field, BufferAllocator allocator);
+
   /**
    * Transfer this vector'data to another vector. The memory associated
    * with this vector is transferred to the allocator of target vector

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java b/java/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java
@@ -125,6 +125,16 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallB
     return underlyingVector.getTransferPair(ref, allocator, callBack);
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+    return underlyingVector.getTransferPair(field, allocator);
+  }
+
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return underlyingVector.getTransferPair(field, allocator, callBack);
+  }
+
   @Override
   public TransferPair makeTransferPair(ValueVector target) {
     return underlyingVector.makeTransferPair(target);

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/LargeVarBinaryVector.java b/java/vector/src/main/java/org/apache/arrow/vector/LargeVarBinaryVector.java
@@ -253,6 +253,11 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
     return new TransferImpl(ref, allocator);
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+    return new TransferImpl(field, allocator);
+  }
+
   /**
    * Construct a TransferPair with a desired target vector of the same type.
    *
@@ -271,6 +276,10 @@ public TransferImpl(String ref, BufferAllocator allocator) {
       to = new LargeVarBinaryVector(ref, field.getFieldType(), allocator);
     }
 
+    public TransferImpl(Field field, BufferAllocator allocator) {
+      to = new LargeVarBinaryVector(field, allocator);
+    }
+
     public TransferImpl(LargeVarBinaryVector to) {
       this.to = to;
     }

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java b/java/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java
@@ -292,6 +292,11 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
     return new LargeVarCharVector.TransferImpl(ref, allocator);
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+    return new LargeVarCharVector.TransferImpl(field, allocator);
+  }
+
   /**
    * Construct a TransferPair with a desired target vector of the same type.
    *
@@ -310,6 +315,10 @@ public TransferImpl(String ref, BufferAllocator allocator) {
       to = new LargeVarCharVector(ref, field.getFieldType(), allocator);
     }
 
+    public TransferImpl(Field field, BufferAllocator allocator) {
+      to = new LargeVarCharVector(field, allocator);
+    }
+
     public TransferImpl(LargeVarCharVector to) {
       this.to = to;
     }

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java b/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java
@@ -106,7 +106,7 @@ public Types.MinorType getMinorType() {
 
   @Override
   public TransferPair getTransferPair(BufferAllocator allocator) {
-    return getTransferPair(null, allocator);
+    return getTransferPair((String) null, allocator);
   }
 
   @Override
@@ -162,11 +162,21 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
     return new TransferImpl();
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+    return new TransferImpl();
+  }
+
   @Override
   public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
     return getTransferPair(ref, allocator);
   }
 
+  @Override
+  public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+    return getTransferPair(field, allocator);
+  }
+
   @Override
   public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((NullVector) target);

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ValueVector.java b/java/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
@@ -134,8 +134,12 @@ public interface ValueVector extends Closeable, Iterable<ValueVector> {
 
   TransferPair getTransferPair(String ref, BufferAllocator allocator);
 
+  TransferPair getTransferPair(Field field, BufferAllocator allocator);
+
   TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack);
 
+  TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack);
+
   /**
    * Makes a new transfer pair used to transfer underlying buffers.
    *