diff --git a/.ai/skills/review-comet-ffi-pr/SKILL.md b/.ai/skills/review-comet-ffi-pr/SKILL.md index 88d47c1683c..c90aa6982cd 100644 --- a/.ai/skills/review-comet-ffi-pr/SKILL.md +++ b/.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`, diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 8dd49d8fb1f..e4c289b1cb9 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/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 diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala index 8e6ff6c2ad3..d85a06ab2b9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala +++ b/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 diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala index 2dc6fcfd325..2a9cf2a3623 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala +++ b/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. */ diff --git a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala index a3046b7f99b..ac13af6f41c 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/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