fix: export large_string and large_binary vectors with 32-bit offsets

.ai/skills/review-comet-ffi-pr/SKILL.md

@@ -92,9 +92,12 @@ ownership. A new input shape needs a reader, not a special case elsewhere.
 
 `NativeUtil.exportBatch()` matches on the column type and has exactly two cases. A `CometVector`
 exports its underlying `FieldVector` through `Data.exportVector`, passing the dictionary provider
-only when the field carries a dictionary. Spark's own `ConstantColumnVector` is materialized into a
-fresh Arrow vector first, because native takes Arrow arrays only. Anything else throws
-`"Comet execution only takes Arrow Arrays"`.
+only when the field carries a dictionary. A large-offset string or binary column
+(`LargeVarCharVector` or `LargeVarBinaryVector`, which a PyArrow UDF can return) is exported as it
+is, although `getFieldVector` rejects it: native columnar-to-row, the only consumer, casts it to the
+planned `Utf8` or `Binary` in `maybe_cast_to_schema_type`, so the JVM does not copy it. Spark's own
+`ConstantColumnVector` is materialized into a fresh Arrow vector first, because native takes Arrow
+arrays only. Anything else throws `"Comet execution only takes Arrow Arrays"`.
 
 The `CometVector` hierarchy is the abstract `CometVector`, the abstract `CometDecodedVector` under
 it, and the concrete `CometPlainVector`, `CometDictionaryVector`, `CometListVector`,

spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala

@@ -24,7 +24,7 @@ import scala.jdk.CollectionConverters._
 
 import org.apache.arrow.c.{ArrowArray, ArrowImporter, ArrowSchema, CDataDictionaryProvider, Data}
 import org.apache.arrow.util.AutoCloseables
-import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.{LargeVarBinaryVector, LargeVarCharVector, VectorSchemaRoot}
 import org.apache.arrow.vector.dictionary.DictionaryProvider
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field}
 import org.apache.spark.SparkException
@@ -142,12 +142,15 @@ class NativeUtil extends AutoCloseable {
           // Don't need to deallocate them here.
           val arrowSchema = ArrowSchema.wrap(schemaAddrs(index))
           val arrowArray = ArrowArray.wrap(arrayAddrs(index))
-          Data.exportVector(
-            allocator,
-            getFieldVector(valueVector, "export"),
-            provider,
-            arrowArray,
-            arrowSchema)
+          // A PyArrow UDF can return large_string or large_binary, which getFieldVector rejects.
+          // Export those as they are: native columnar-to-row casts each column to its planned
+          // type in maybe_cast_to_schema_type, which rebuilds the offsets and keeps the data.
+          val fieldVector = valueVector match {
+            case v: LargeVarCharVector => v
+            case v: LargeVarBinaryVector => v
+            case v => getFieldVector(v, "export")
+          }
+          Data.exportVector(allocator, fieldVector, provider, arrowArray, arrowSchema)
         case cv: ConstantColumnVector =>
           // Spark uses ConstantColumnVector for partition columns / per-batch constants (e.g.
           // partition values, synthetic columns). Materialise to a fresh Arrow vector so Comet's

spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala

@@ -106,8 +106,9 @@ case class CometMapInBatchExec(
       columnarBatchIter.map { batch =>
         // Python returns a single struct column; flatten to the user's output columns. The runner
         // produces Comet vectors, so the struct's children are already CometVectors that downstream
-        // consumers (a stacked CometMapInBatchExec, or NativeUtil.exportBatch for a native Comet
-        // operator) can use directly.
+        // consumers (a stacked CometMapInBatchExec, the Arrow C Stream in CometNativeArrowSource
+        // for a native Comet operator, or NativeUtil.exportBatch for native columnar-to-row) can
+        // use directly.
         val structVector = batch.column(0).asInstanceOf[CometStructVector]
         val outputVectors: Array[ColumnVector] =
           outputAttrs.indices.map(i => structVector.getChild(i)).toArray

spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala

@@ -536,6 +536,51 @@ class CometNativeColumnarToRowSuite extends CometTestBase with AdaptiveSparkPlan
     }
   }
 
+  test("large-offset string and binary columns") {
+    import java.nio.charset.StandardCharsets
+
+    import org.apache.arrow.vector.{LargeVarBinaryVector, LargeVarCharVector}
+    import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+    import org.apache.comet.CometArrowAllocator
+    import org.apache.comet.vector.CometPlainVector
+
+    // A PyArrow UDF can return large_string or large_binary. exportBatch hands them over with
+    // 64-bit offsets, and the converter casts them to the Utf8 and Binary it plans for.
+    val values = Seq("hello", null, "", "w\u00f6rld")
+    val schema = new StructType().add("s", StringType).add("b", BinaryType)
+    val strings = new LargeVarCharVector("s", CometArrowAllocator)
+    val binaries = new LargeVarBinaryVector("b", CometArrowAllocator)
+    val converter = new NativeColumnarToRowConverter(schema, values.length)
+    try {
+      Seq(strings, binaries).foreach { vector =>
+        vector.allocateNew(values.length)
+        values.zipWithIndex.foreach { case (value, i) =>
+          if (value != null) vector.setSafe(i, value.getBytes(StandardCharsets.UTF_8))
+        }
+        vector.setValueCount(values.length)
+      }
+      val batch = new ColumnarBatch(
+        Array[ColumnVector](new CometPlainVector(strings), new CometPlainVector(binaries)),
+        values.length)
+
+      val rows = converter
+        .convert(batch)
+        .map { row =>
+          val s = if (row.isNullAt(0)) null else row.getUTF8String(0).toString
+          val b =
+            if (row.isNullAt(1)) null else new String(row.getBinary(1), StandardCharsets.UTF_8)
+          (s, b)
+        }
+        .toList
+      assert(rows == values.map(v => (v, v)))
+    } finally {
+      converter.close()
+      strings.close()
+      binaries.close()
+    }
+  }
+
   /**
    * Helper to create a parquet table from a DataFrame and run a function with it.
    */

spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala

@@ -27,7 +27,7 @@ import scala.util.Using
 
 import org.apache.arrow.c.{ArrowArray, ArrowSchema, Data}
 import org.apache.arrow.memory.RootAllocator
-import org.apache.arrow.vector.{IntVector, TimeStampMicroTZVector, UInt4Vector, VarCharVector}
+import org.apache.arrow.vector.{IntVector, LargeVarBinaryVector, LargeVarCharVector, TimeStampMicroTZVector, UInt4Vector, VarCharVector}
 import org.apache.arrow.vector.complex.StructVector
 import org.apache.arrow.vector.dictionary.Dictionary
 import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider
@@ -373,6 +373,63 @@ class NativeUtilSuite extends CometTestBase {
     }
   }
 
+  test("exportBatch exports large-offset string and binary vectors as they are") {
+    // Native columnar-to-row narrows them to its planned type; see CometNativeColumnarToRowSuite.
+    val values = Seq("hello", null, "", "wörld")
+    val before = CometArrowAllocator.getAllocatedMemory
+    val strings = new LargeVarCharVector("s", CometArrowAllocator)
+    val binaries = new LargeVarBinaryVector("b", CometArrowAllocator)
+    Seq(strings, binaries).foreach { vector =>
+      vector.allocateNew(values.length)
+      values.zipWithIndex.foreach { case (value, i) =>
+        if (value != null) vector.setSafe(i, value.getBytes(StandardCharsets.UTF_8))
+      }
+      vector.setValueCount(values.length)
+    }
+    val batch = new ColumnarBatch(
+      Array[ColumnVector](new CometPlainVector(strings), new CometPlainVector(binaries)),
+      values.length)
+
+    val structAllocator = new RootAllocator(Long.MaxValue)
+    val arrays = Array.fill(2)(ArrowArray.allocateNew(structAllocator))
+    val schemas = Array.fill(2)(ArrowSchema.allocateNew(structAllocator))
+    val nativeUtil = new NativeUtil
+    var imported: ColumnarBatch = null
+    try {
+      val exportedRows = nativeUtil.exportBatch(
+        arrays.map(_.memoryAddress()),
+        schemas.map(_.memoryAddress()),
+        batch)
+      assert(exportedRows == values.length)
+      strings.close()
+      binaries.close()
+
+      imported =
+        new ColumnarBatch(nativeUtil.importVector(arrays, schemas).toArray, values.length)
+
+      val stringColumn = imported.column(0).asInstanceOf[CometVector]
+      val binaryColumn = imported.column(1).asInstanceOf[CometVector]
+      assert(stringColumn.getValueVector.isInstanceOf[LargeVarCharVector])
+      assert(binaryColumn.getValueVector.isInstanceOf[LargeVarBinaryVector])
+      def read(column: CometVector, get: Int => Array[Byte]): Seq[String] =
+        values.indices.map { i =>
+          if (column.isNullAt(i)) null else new String(get(i), StandardCharsets.UTF_8)
+        }
+      assert(read(stringColumn, stringColumn.getUTF8String(_).getBytes) == values)
+      assert(read(binaryColumn, binaryColumn.getBinary) == values)
+    } finally {
+      strings.close()
+      binaries.close()
+      if (imported != null) {
+        imported.close()
+      }
+      nativeUtil.close()
+      (arrays ++ schemas).foreach(_.close())
+      structAllocator.close()
+    }
+    assert(CometArrowAllocator.getAllocatedMemory == before)
+  }
+
   test("imports are charged to the import allocator and roll up into the root") {
     // Arrow charges a buffer to whichever allocator owns it, so the tracing counters can only
     // report the import path apart from the rest of Comet's Arrow memory if imports go to their