GH-88: Support importing C Data Interface arrays with a non-zero offset

c/src/main/java/org/apache/arrow/c/ArrayImporter.java

@@ -25,29 +25,46 @@
 import java.util.List;
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.util.MemoryUtil;
 import org.apache.arrow.util.Preconditions;
+import org.apache.arrow.vector.ExtensionTypeVector;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
+import org.apache.arrow.vector.types.UnionMode;
 import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;
 import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
+import org.apache.arrow.vector.types.pojo.Field;
 
 /** Importer for {@link ArrowArray}. */
 final class ArrayImporter {
   private static final int MAX_IMPORT_RECURSION_LEVEL = 64;
 
   private final BufferAllocator allocator;
   private final FieldVector vector;
+  // The field the vector was created from, or null to use the vector's own field. Unlike the
+  // field of an ExtensionTypeVector that does not override getField(), it has the extension type.
+  private final Field field;
   private final DictionaryProvider dictionaryProvider;
 
   private ReferenceCountedArrowArray underlyingAllocation;
   private int recursionLevel;
 
   ArrayImporter(
       BufferAllocator allocator, FieldVector vector, DictionaryProvider dictionaryProvider) {
+    this(allocator, vector, /* field= */ null, dictionaryProvider);
+  }
+
+  ArrayImporter(
+      BufferAllocator allocator,
+      FieldVector vector,
+      Field field,
+      DictionaryProvider dictionaryProvider) {
     this.allocator = Preconditions.checkNotNull(allocator);
     this.vector = Preconditions.checkNotNull(vector);
+    this.field = field;
     this.dictionaryProvider = dictionaryProvider;
   }
 
@@ -71,8 +88,7 @@ void importArray(ArrowArray src) {
     }
   }
 
-  private void importChild(ArrayImporter parent, ArrowArray src) {
-    ArrowArray.Snapshot snapshot = src.snapshot();
+  private void importChild(ArrayImporter parent, ArrowArray.Snapshot snapshot) {
     checkState(snapshot.release != NULL, "Cannot import released ArrowArray");
     recursionLevel = parent.recursionLevel + 1;
     checkState(
@@ -84,15 +100,78 @@ private void importChild(ArrayImporter parent, ArrowArray src) {
   }
 
   private void doImport(ArrowArray.Snapshot snapshot) {
-    // A non-zero offset is only meaningful for arrays that have buffers to offset into.
-    // Null arrays carry no buffers, so their offset is inert and safe to ignore.
-    checkState(
-        snapshot.offset == 0
-            || snapshot.length == 0
-            || vector.getField().getType().getTypeID() == ArrowType.ArrowTypeID.Null,
-        "ArrowArray struct has non-zero offset (%s), which is not supported",
-        snapshot.offset);
+    checkState(snapshot.offset >= 0, "ArrowArray struct has negative offset (%s)", snapshot.offset);
+    if (snapshot.offset != 0 && snapshot.length != 0) {
+      importSliced(snapshot);
+    } else {
+      importUnsliced(snapshot, /* skippedRows= */ 0, /* trimChildren= */ false);
+    }
+  }
+
+  /**
+   * Java vectors cannot represent an offset, so the rows up to the end of the slice are imported
+   * into a temporary vector and the slice is then transferred into the target. Leading rows that
+   * fill whole bytes of a bitmap are skipped by advancing the buffer pointers, so that the
+   * temporary vector, and any validity bitmap synthesized for it, has fewer than 8 rows before the
+   * slice.
+   */
+  private void importSliced(ArrowArray.Snapshot snapshot) {
+    // Run-end encoded arrays have no buffers, and their children are not indexed by row
+    long skippedRows =
+        storageType().getTypeID() == ArrowTypeID.RunEndEncoded ? 0 : snapshot.offset & ~7L;
+    int offset = checkedCastToInt(snapshot.offset - skippedRows);
+    int length = checkedCastToInt(snapshot.length);
+
+    snapshot.offset = 0;
+    snapshot.length = Math.addExact(offset, length);
+    snapshot.null_count = unknownNullCount(snapshot);
+
+    Field unslicedField = field != null ? field : vector.getField();
+    try (FieldVector unslicedVector = unslicedField.createVector(allocator)) {
+      // When the target is an extension vector whose getField() reports its storage type, the
+      // temporary vector is a storage vector, so transfer between the storage vectors
+      FieldVector from = ExtensionTypeVector.getStorageVector(unslicedVector);
+      FieldVector to = ExtensionTypeVector.getStorageVector(vector);
+      checkTransferable(from, to);
+      ArrayImporter unslicedImporter =
+          new ArrayImporter(allocator, unslicedVector, unslicedField, dictionaryProvider);
+      unslicedImporter.recursionLevel = recursionLevel;
+      unslicedImporter.underlyingAllocation = underlyingAllocation;
+      unslicedImporter.importUnsliced(snapshot, skippedRows, /* trimChildren= */ true);
+      from.makeTransferPair(to).splitAndTransfer(offset, length);
+    }
+  }
 
+  /**
+   * Transfer pairs of nested vectors create or reuse the children of the target based on the
+   * classes of the source's children, so every vector of the target must be an instance of the
+   * class of its counterpart in the source.
+   */
+  private static void checkTransferable(FieldVector from, FieldVector to) {
+    if (!from.getClass().isInstance(to)) {
+      String reason =
+          to instanceof ExtensionTypeVector && !(from instanceof ExtensionTypeVector)
+              ? "its getField() does not report its extension type"
+              : "its field creates a " + from.getClass().getSimpleName();
+      throw new UnsupportedOperationException(
+          String.format(
+              "Cannot import an ArrowArray with a non-zero offset into %s %s, because %s",
+              to.getClass().getSimpleName(), to.getName(), reason));
+    }
+    List<FieldVector> fromChildren = from.getChildrenFromFields();
+    List<FieldVector> toChildren = to.getChildrenFromFields();
+    for (int i = 0; i < Math.min(fromChildren.size(), toChildren.size()); i++) {
+      checkTransferable(fromChildren.get(i), toChildren.get(i));
+    }
+  }
+
+  /**
+   * Imports an array whose offset is 0, skipping the first {@code skippedRows} rows of its buffers,
+   * which must be a multiple of 8. With {@code trimChildren}, children whose rows line up with the
+   * rows of the array are imported only for the rows of the array, rather than in full.
+   */
+  private void importUnsliced(
+      ArrowArray.Snapshot snapshot, long skippedRows, boolean trimChildren) {
     // First import children (required for reconstituting parent array data)
     long[] children =
         NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children));
@@ -103,11 +182,31 @@ private void doImport(ArrowArray.Snapshot snapshot) {
           "ArrowArray struct has %s children (expected %s)",
           children.length,
           childVectors.size());
+      List<Field> childFields = field != null ? field.getChildren() : null;
+      long childRowsPerRow = trimChildren ? childRowsPerRow(storageType()) : 0;
       for (int i = 0; i < children.length; i++) {
         checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i);
+        ArrowArray.Snapshot child = ArrowArray.wrap(children[i]).snapshot();
+        if (childRowsPerRow != 0) {
+          long firstRow = Math.multiplyExact(skippedRows, childRowsPerRow);
+          long rowCount = Math.multiplyExact(snapshot.length, childRowsPerRow);
+          checkState(
+              child.length - firstRow >= rowCount,
+              "ArrowArray struct child at position %s has length %s, but its parent needs %s rows",
+              i,
+              child.length,
+              firstRow + rowCount);
+          child.offset = Math.addExact(child.offset, firstRow);
+          child.length = rowCount;
+          child.null_count = unknownNullCount(child);
+        }
+        Field childField =
+            childFields != null && childFields.size() == children.length
+                ? childFields.get(i)
+                : null;
         ArrayImporter childImporter =
-            new ArrayImporter(allocator, childVectors.get(i), dictionaryProvider);
-        childImporter.importChild(this, ArrowArray.wrap(children[i]));
+            new ArrayImporter(allocator, childVectors.get(i), childField, dictionaryProvider);
+        childImporter.importChild(this, child);
       }
     }
 
@@ -124,7 +223,7 @@ private void doImport(ArrowArray.Snapshot snapshot) {
 
       ArrayImporter dictionaryImporter =
           new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider);
-      dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary));
+      dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary).snapshot());
     }
 
     // Import main data
@@ -133,7 +232,8 @@ private void doImport(ArrowArray.Snapshot snapshot) {
         NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers));
 
     try (final BufferImportTypeVisitor visitor =
-        new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) {
+        new BufferImportTypeVisitor(
+            allocator, underlyingAllocation, fieldNode, bufferPointers, skippedRows)) {
       final List<ArrowBuf> buffers;
       if (bufferPointers == null || bufferPointers.length == 0) {
         buffers = Collections.emptyList();
@@ -150,4 +250,42 @@ private void doImport(ArrowArray.Snapshot snapshot) {
           e);
     }
   }
+
+  /**
+   * Returns the null count to import a different range of the rows of {@code snapshot} with than
+   * the range its null count covers. Without a validity bitmap it must be 0 for the buffers to
+   * load, otherwise it is unknown.
+   */
+  private static long unknownNullCount(ArrowArray.Snapshot snapshot) {
+    boolean hasBitmap =
+        snapshot.n_buffers > 0
+            && snapshot.buffers != NULL
+            && MemoryUtil.getLong(snapshot.buffers) != NULL;
+    return hasBitmap ? -1 : 0;
+  }
+
+  /** The type of the vector, or the storage type of an extension type. */
+  private ArrowType storageType() {
+    ArrowType type = vector.getField().getType();
+    return type instanceof ArrowType.ExtensionType
+        ? ((ArrowType.ExtensionType) type).storageType()
+        : type;
+  }
+
+  /**
+   * Returns how many rows of each child a row of an array of {@code type} spans, or 0 if its
+   * children are not indexed by its rows.
+   */
+  private static long childRowsPerRow(ArrowType type) {
+    switch (type.getTypeID()) {
+      case Struct:
+        return 1;
+      case Union:
+        return ((ArrowType.Union) type).getMode() == UnionMode.Sparse ? 1 : 0;
+      case FixedSizeList:
+        return ((ArrowType.FixedSizeList) type).getListSize();
+      default:
+        return 0;
+    }
+  }
 }

c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java

@@ -62,17 +62,34 @@ class BufferImportTypeVisitor implements ArrowType.ArrowTypeVisitor<List<ArrowBu
   private final ReferenceCountedArrowArray underlyingAllocation;
   private final ArrowFieldNode fieldNode;
   private final long[] buffers;
+  private final long skippedRows;
   private final List<ArrowBuf> imported;
 
   BufferImportTypeVisitor(
       BufferAllocator allocator,
       ReferenceCountedArrowArray underlyingAllocation,
       ArrowFieldNode fieldNode,
       long[] buffers) {
+    this(allocator, underlyingAllocation, fieldNode, buffers, /* skippedRows= */ 0);
+  }
+
+  /**
+   * Imports the buffers of the {@code fieldNode.getLength()} rows that follow the first {@code
+   * skippedRows} rows of the C array. The skipped rows must fill whole bytes of a bitmap, which
+   * makes skipping them a matter of advancing the buffer pointers of each row.
+   */
+  BufferImportTypeVisitor(
+      BufferAllocator allocator,
+      ReferenceCountedArrowArray underlyingAllocation,
+      ArrowFieldNode fieldNode,
+      long[] buffers,
+      long skippedRows) {
+    checkState(skippedRows % 8 == 0, "Skipped rows (%s) must be a multiple of 8", skippedRows);
     this.allocator = allocator;
     this.underlyingAllocation = underlyingAllocation;
     this.fieldNode = fieldNode;
     this.buffers = buffers;
+    this.skippedRows = skippedRows;
     this.imported = new ArrayList<>();
   }
 
@@ -83,6 +100,10 @@ public void close() throws Exception {
 
   @VisibleForTesting
   ArrowBuf importBuffer(ArrowType type, int index, long capacity) {
+    return importBuffer(type, index, /* skippedBytes= */ 0, capacity);
+  }
+
+  private ArrowBuf importBuffer(ArrowType type, int index, long skippedBytes, long capacity) {
     checkState(
         buffers.length > index,
         "Expected at least %s buffers for type %s, but found %s",
@@ -102,24 +123,28 @@ ArrowBuf importBuffer(ArrowType type, int index, long capacity) {
       }
     }
 
-    ArrowBuf buf = underlyingAllocation.unsafeAssociateAllocation(allocator, capacity, bufferPtr);
+    ArrowBuf buf =
+        underlyingAllocation.unsafeAssociateAllocation(
+            allocator, capacity, Math.addExact(bufferPtr, skippedBytes));
     imported.add(buf);
     return buf;
   }
 
+  /** Imports a buffer that holds {@code bitsPerSlot} bits for each row. */
   private ArrowBuf importFixedBits(ArrowType type, int index, long bitsPerSlot) {
     final long capacity = DataSizeRoundingUtil.divideBy8Ceil(bitsPerSlot * fieldNode.getLength());
-    return importBuffer(type, index, capacity);
+    return importBuffer(type, index, Math.multiplyExact(skippedRows, bitsPerSlot) / 8, capacity);
   }
 
+  /** Imports a buffer that holds {@code bytesPerSlot} bytes for each row. */
   private ArrowBuf importFixedBytes(ArrowType type, int index, long bytesPerSlot) {
     final long capacity = bytesPerSlot * fieldNode.getLength();
-    return importBuffer(type, index, capacity);
+    return importBuffer(type, index, Math.multiplyExact(skippedRows, bytesPerSlot), capacity);
   }
 
   private ArrowBuf importOffsets(ArrowType type, long bytesPerSlot) {
     final long capacity = bytesPerSlot * (fieldNode.getLength() + 1);
-    return importBuffer(type, 1, capacity);
+    return importBuffer(type, 1, Math.multiplyExact(skippedRows, bytesPerSlot), capacity);
   }
 
   private ArrowBuf importData(ArrowType type, long capacity) {

c/src/main/java/org/apache/arrow/c/Data.java

@@ -16,6 +16,7 @@
  */
 package org.apache.arrow.c;
 
+import java.util.List;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorLoader;
@@ -35,8 +36,12 @@
 /**
  * Functions for working with the C data interface.
  *
- * <p>This API is EXPERIMENTAL. Note that currently only 64bit systems are supported. Importing
- * {@link ArrowArray ArrowArrays} with a non-zero offset is not supported.
+ * <p>This API is EXPERIMENTAL. Note that currently only 64bit systems are supported.
+ *
+ * <p>Importing {@link ArrowArray ArrowArrays} with a non-zero offset is not always zero-copy:
+ * buffers that cannot be sliced at the offset, such as offset buffers or view data, are copied.
+ * Importing such an array into an existing vector fails if an extension vector nested in it does
+ * not report its extension type in {@link FieldVector#getField()}.
  */
 public final class Data {
 
@@ -358,7 +363,22 @@ public static void importIntoVector(
       FieldVector vector,
       DictionaryProvider provider,
       boolean closeImportedStructs) {
-    ArrayImporter importer = new ArrayImporter(allocator, vector, provider);
+    importIntoVector(allocator, array, vector, /* field= */ null, provider, closeImportedStructs);
+  }
+
+  /**
+   * Import Java vector from the C data interface, given the field that the vector was created from,
+   * if any. The field reports the type of extension vectors whose getField() only reports their
+   * storage type, which an import with a non-zero offset needs to create a temporary vector.
+   */
+  private static void importIntoVector(
+      BufferAllocator allocator,
+      ArrowArray array,
+      FieldVector vector,
+      Field field,
+      DictionaryProvider provider,
+      boolean closeImportedStructs) {
+    ArrayImporter importer = new ArrayImporter(allocator, vector, field, provider);
     importer.importArray(array);
     if (closeImportedStructs) {
       array.close();
@@ -407,7 +427,7 @@ public static FieldVector importVector(
       boolean closeImportedStructs) {
     Field field = importField(allocator, schema, provider, closeImportedStructs);
     FieldVector vector = field.createVector(allocator);
-    importIntoVector(allocator, array, vector, provider, closeImportedStructs);
+    importIntoVector(allocator, array, vector, field, provider, closeImportedStructs);
     return vector;
   }
 
@@ -454,8 +474,10 @@ public static void importIntoVectorSchemaRoot(
       DictionaryProvider provider,
       boolean closeImportedStructs) {
     try (StructVector structVector = StructVector.emptyWithDuplicates("", allocator)) {
-      structVector.initializeChildrenFromFields(root.getSchema().getFields());
-      importIntoVector(allocator, array, structVector, provider, closeImportedStructs);
+      List<Field> fields = root.getSchema().getFields();
+      structVector.initializeChildrenFromFields(fields);
+      Field field = new Field("", structVector.getField().getFieldType(), fields);
+      importIntoVector(allocator, array, structVector, field, provider, closeImportedStructs);
       StructVectorUnloader unloader = new StructVectorUnloader(structVector);
       VectorLoader loader = new VectorLoader(root);
       try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) {