From 6c5a59b28348bb9fa6a0047096edf8b7ac07a307 Mon Sep 17 00:00:00 2001
From: Chao Sun <chao@openai.com>
Date: Fri, 14 Aug 2026 22:20:46 -0700
Subject: [PATCH 1/3] perf: serialize Python input directly from Comet Arrow
 vectors

---
 .github/workflows/pyarrow_udf_test.yml        |   2 +
 docs/source/user-guide/latest/pyarrow-udfs.md |  23 +-
 .../rules/EliminateRedundantTransitions.scala |   7 +-
 .../python/CometArrowPythonRunnerBase.scala   | 196 ++++++-----
 .../pyspark/benchmark_pyarrow_udf.py          |  12 +-
 .../resources/pyspark/test_pyarrow_udf.py     | 171 +++++++--
 .../python/CometArrowPythonRunnerSuite.scala  | 329 ++++++++++++++++++
 7 files changed, 589 insertions(+), 151 deletions(-)
 create mode 100644 spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala

diff --git a/.github/workflows/pyarrow_udf_test.yml b/.github/workflows/pyarrow_udf_test.yml
index 1a03962685d..8fff5d414cd 100644
--- a/.github/workflows/pyarrow_udf_test.yml
+++ b/.github/workflows/pyarrow_udf_test.yml
@@ -39,10 +39,12 @@ on:
       - "spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala"
+      - "spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala"
       - "spark/src/test/resources/pyspark/conftest.py"
       - "spark/src/test/resources/pyspark/test_pyarrow_udf.py"
       - "spark/src/test/spark-3.5/org/apache/spark/sql/comet/CometMapInBatchSuite.scala"
       - "spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala"
+      - "spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala"
       - ".github/workflows/pyarrow_udf_test.yml"
   pull_request:
     paths: *feature-paths
diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md
index ff493de23c0..1c023db83c6 100644
--- a/docs/source/user-guide/latest/pyarrow-udfs.md
+++ b/docs/source/user-guide/latest/pyarrow-udfs.md
@@ -203,18 +203,11 @@ on the unoptimized path.
   session time zone such a UDF can diverge from the unoptimized path. Set
   `spark.comet.exec.pyarrowUDF.enabled=false` for those UDFs.
 - `spark.sql.execution.arrow.useLargeVarTypes=true` is not supported. With this conf enabled,
-  Spark widens `StringType` and `BinaryType` to Arrow's 8-byte-offset variants in the
-  destination IPC root, while Comet's source vectors always use 4-byte offsets. The buffer-copy
-  path cannot bridge that mismatch, so `EliminateRedundantTransitions` skips the rewrite and
-  vanilla Spark handles the operation.
-- Each batch is copied twice on the JVM side: once from Comet's vectors into Spark's
-  destination IPC root (per-buffer `setBytes`), and a second time inside the IPC writer when
-  `VectorUnloader` / `MessageSerializer.serialize` walks the root and writes bytes to the
-  pipe to the Python worker. The pipe write is structural (Spark's transport to Python is
-  fork + pipe + Arrow IPC, so the buffer bytes must reach the pipe at least once); dropping
-  the first copy by serialising directly from Comet's vectors is tracked in
-  [#4294](https://github.com/apache/datafusion-comet/issues/4294). Even after that,
-  true zero-copy at the JVM boundary is blocked because Comet's source `FieldVector`s are
-  imported from native via Arrow C Data Interface (their buffers route `release` through FFI),
-  while Spark's destination IPC root is a child of `ArrowUtils.rootAllocator`. The two
-  reference managers cannot share buffers via `TransferPair`.
+  Spark expects `StringType` and `BinaryType` to use Arrow's 8-byte-offset variants, while
+  Comet's source vectors always use 4-byte offsets. `EliminateRedundantTransitions` skips the
+  rewrite for this incompatible layout and vanilla Spark handles the operation.
+- Comet writes input Arrow IPC record batches directly from its existing vector buffers. The
+  only additional Arrow buffer is the validity bitmap for the non-null struct that wraps the
+  input columns. Writing the IPC bytes to the Python worker's pipe still requires one copy;
+  that copy is inherent to Spark's process-based Python transport. This path does not transfer
+  buffers between Arrow allocators or change their ownership.
diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
index ec277cfc7bc..9f9345bea07 100644
--- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
@@ -112,10 +112,9 @@ case class EliminateRedundantTransitions(session: SparkSession)
       // 4.1+ matches the renamed `MapInArrowExec`.
       //
       // Falls back to vanilla Spark when `spark.sql.execution.arrow.useLargeVarTypes` is enabled:
-      // CometArrowPythonRunnerBase.copyVector does raw `setBytes` on each Arrow buffer, but Comet's
-      // source string/binary vectors always use 4-byte offsets while the destination root is
-      // allocated with 8-byte offsets when this conf is on. The buffer counts match but the
-      // offset width does not, so a direct memcpy would corrupt the offsets.
+      // Comet's source string/binary vectors use 4-byte offsets, while Spark expects 8-byte
+      // offsets when this conf is on. Direct IPC serialization cannot change the source vectors'
+      // physical layout, so forwarding them under Spark's widened schema would corrupt the stream.
       //
       // `EligibleMapInBatch` matches whenever the operator would run natively if the feature were
       // enabled. When it is disabled (the default) we leave the vanilla Spark operator in place
diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
index b22b0a792c3..03f5b18c66a 100644
--- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
@@ -20,14 +20,17 @@
 package org.apache.spark.sql.execution.python
 
 import java.io.{DataInputStream, DataOutputStream}
-import java.nio.channels.Channels
+import java.nio.channels.{Channels, WritableByteChannel}
+import java.util.ArrayList
 import java.util.concurrent.atomic.AtomicBoolean
 
 import scala.jdk.CollectionConverters._
 
-import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
-import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
+import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot, VectorUnloader}
+import org.apache.arrow.vector.complex.StructVector
 import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
@@ -36,7 +39,6 @@ import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
-import org.apache.spark.unsafe.Platform
 
 import org.apache.comet.CometArrowAllocator
 import org.apache.comet.vector.{CometDecodedVector, CometVector}
@@ -53,7 +55,7 @@ import org.apache.comet.vector.{CometDecodedVector, CometVector}
  * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
  * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
  * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
- *   - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ *   - Input: each Comet `ColumnarBatch` is written directly from its shaded Arrow vectors to the
  *     worker with a shaded `ArrowStreamWriter`.
  *   - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
  *     `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
@@ -107,9 +109,8 @@ private[python] trait CometArrowPythonRunnerBase
       private val allocator =
         CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
       private var currentGroup: Iterator[ColumnarBatch] = _
-      private var arrowWriter: ArrowStreamWriter = _
+      private var arrowWriter: CometArrowPythonRunnerBase.DirectArrowStreamWriter = _
       private var writeRoot: VectorSchemaRoot = _
-      private var structVec: StructVector = _
 
       // The runner's input schema is a single struct column ("struct") whose children are the
       // user's input columns (see `schema` above). Cast once here rather than at each use site.
@@ -132,16 +133,18 @@ private[python] trait CometArrowPythonRunnerBase
         writeUDF(dataOut)
       }
 
-      /** Build the destination struct root and start the writer from the given child fields. */
+      /** Build the schema-only struct root and start the writer from the given child fields. */
       private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
         val structField =
           new Field(
             "struct",
             new FieldType(false, ArrowType.Struct.INSTANCE, null),
             childFields.asJava)
-        structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+        val structVec = structField.createVector(allocator).asInstanceOf[StructVector]
         writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
-        arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+        arrowWriter = new CometArrowPythonRunnerBase.DirectArrowStreamWriter(
+          writeRoot,
+          Channels.newChannel(dataOut))
         arrowWriter.start()
       }
 
@@ -167,49 +170,37 @@ private[python] trait CometArrowPythonRunnerBase
 
         val cometBatch = currentGroup.next()
         val startData = dataOut.size()
+        val sourceVectors = (0 until cometBatch.numCols()).map { i =>
+          cometBatch
+            .column(i)
+            .asInstanceOf[CometDecodedVector]
+            .getValueVector
+            .asInstanceOf[FieldVector]
+        }
 
         if (arrowWriter == null) {
-          // Build the destination struct root once, sized to the first batch's child fields.
+          // Build the schema-only struct root once from the first batch's child fields.
           // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
           // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
           // column names from the input schema (the worker reads columns by name, and shaded
-          // Arrow rejects a null field name). The field types and child structure are kept as-is
-          // so copyVector still walks the source and destination trees in lockstep. Keeping the
-          // type as-is also means a TimestampType reaches the worker with Comet's UTC time zone
+          // Arrow rejects a null field name). Keep the field types and child structure as-is so
+          // the advertised schema matches the source buffers. Keeping the type as-is also means
+          // a TimestampType reaches the worker with Comet's UTC time zone
           // rather than the session zone vanilla Spark would label it with; this is a documented
           // limitation (see pyarrow-udfs.md), not a value difference, since the stored instant is
           // identical.
           val childNames = inputStructType.fieldNames
-          val childFields = (0 until cometBatch.numCols()).map { i =>
-            val vecField =
-              cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
-            renamed(vecField, childNames(i), forceNullable = true)
+          val childFields = sourceVectors.zipWithIndex.map { case (vector, i) =>
+            renamed(vector.getField, childNames(i), forceNullable = true)
           }
           startWriter(childFields, dataOut)
         }
 
-        var i = 0
-        while (i < cometBatch.numCols()) {
-          val src = cometBatch
-            .column(i)
-            .asInstanceOf[CometDecodedVector]
-            .getValueVector
-            .asInstanceOf[FieldVector]
-          val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
-          copyVector(src, dst)
-          i += 1
-        }
-        val numRows = cometBatch.numRows()
-        structVec.setValueCount(numRows)
-        // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
-        // allocated and zero-initialised, so without this Python would see an all-null struct.
-        val validityBytes = (numRows + 7) / 8
-        Platform.setMemory(
-          structVec.getValidityBuffer.memoryAddress(),
-          0xff.toByte,
-          validityBytes)
-        writeRoot.setRowCount(numRows)
-        arrowWriter.writeBatch()
+        CometArrowPythonRunnerBase.writeDirectBatch(
+          arrowWriter,
+          sourceVectors,
+          cometBatch.numRows(),
+          allocator)
 
         pythonMetrics("pythonDataSent") += dataOut.size() - startData
         true
@@ -291,8 +282,8 @@ private[python] trait CometArrowPythonRunnerBase
   /**
    * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
    * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
-   * materialize the struct. Keeping the type and structure intact means the destination tree
-   * still mirrors the Comet source tree for [[copyVector]].
+   * materialize the struct. Keeping the type and structure intact means the advertised schema
+   * still mirrors the Comet source vectors serialized directly into each record batch.
    */
   private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
     // A Map's descendants must keep their original nullability: Arrow requires the entries struct
@@ -317,74 +308,81 @@ private[python] trait CometArrowPythonRunnerBase
     // non-nullable Arrow `Field` even for columns that contain nulls (Comet uses positional schema
     // and does not round-trip Spark's nullability), and the worker rejects a null value under a
     // non-nullable field (`from_pandas(pdf, schema=batch.schema)` raises). Marking the field
-    // nullable is a safe superset; `copyVector` fills an all-valid validity buffer when the source
-    // has no nulls.
+    // nullable is a safe superset; Arrow IPC permits an empty validity buffer when its field node
+    // reports no null values.
     val ft = field.getFieldType
     val nullable = forceNullable || ft.isNullable
     val newFt = new FieldType(nullable, ft.getType, ft.getDictionary, ft.getMetadata)
     new Field(name, newFt, newChildren)
   }
+}
+
+private[python] object CometArrowPythonRunnerBase {
 
   /**
-   * Copy a Comet column into the destination FieldVector. Walks both trees in lockstep: sizes
-   * each destination node from the source, copies every buffer with `ArrowBuf.setBytes`, then
-   * sets value counts bottom-up so `setValueCount` does not rewrite the offset bytes we just
-   * copied. Both source and destination are Comet's (shaded) Arrow vectors, so no shaded /
-   * unshaded type crosses.
+   * Expose ArrowWriter's existing record-batch serialization without unloading its schema root.
    */
-  private def copyVector(src: FieldVector, dst: FieldVector): Unit = {
-    val valueCount = src.getValueCount
-
-    dst match {
-      case bfwv: BaseFixedWidthVector =>
-        bfwv.allocateNew(valueCount)
-      case bvwv: BaseVariableWidthVector =>
-        bvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount)
-      case blvwv: BaseLargeVariableWidthVector =>
-        blvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount)
-      case _ =>
-        dst.setInitialCapacity(valueCount)
-        dst.allocateNew()
-    }
-
-    val srcBufs = src.getFieldBuffers
-    val dstBufs = dst.getFieldBuffers
-    require(
-      srcBufs.size == dstBufs.size,
-      s"buffer count mismatch for ${dst.getField}: src=${srcBufs.size}, dst=${dstBufs.size}")
-    srcBufs.asScala.zip(dstBufs.asScala).foreach { case (s, d) =>
-      d.setBytes(0, s, 0, s.readableBytes)
-    }
+  private[python] final class DirectArrowStreamWriter(
+      root: VectorSchemaRoot,
+      channel: WritableByteChannel)
+      extends ArrowStreamWriter(root, null, channel) {
 
-    val srcChildren = src.getChildrenFromFields
-    val dstChildren = dst.getChildrenFromFields
-    require(
-      srcChildren.size == dstChildren.size,
-      s"child count mismatch for ${dst.getField}: src=${srcChildren.size}, dst=${dstChildren.size}")
-    srcChildren.asScala.zip(dstChildren.asScala).foreach { case (sc, dc) =>
-      copyVector(sc.asInstanceOf[FieldVector], dc.asInstanceOf[FieldVector])
+    def writeDirect(batch: ArrowRecordBatch): Unit = {
+      writeRecordBatch(batch)
     }
+  }
 
-    // For vectors that fill offset-buffer "holes" in setValueCount (variable-width and list
-    // types), set lastSet = vc - 1 first so fillHoles is a no-op and the already-copied offset
-    // bytes are preserved.
-    dst match {
-      case v: BaseVariableWidthVector => v.setLastSet(valueCount - 1)
-      case v: BaseLargeVariableWidthVector => v.setLastSet(valueCount - 1)
-      case v: ListVector => v.setLastSet(valueCount - 1)
-      case v: LargeListVector => v.setLastSet(valueCount - 1)
-      case _ =>
-    }
-    dst.setValueCount(valueCount)
-
-    // Every destination field is nullable (see `renamed`), so the worker reads the validity
-    // buffer. When the source has no nulls its validity buffer may be empty (Comet omits it),
-    // which would otherwise leave the freshly-allocated destination validity all-zero and make
-    // the worker see every value as null. Set all-valid in that case. Done after setValueCount,
-    // which can rewrite validity, mirroring the struct-level all-valid fill in writeNextInput.
-    if (valueCount > 0 && dst.getField.isNullable && src.getNullCount == 0) {
-      val validityBytes = (valueCount + 7) / 8
-      Platform.setMemory(dst.getValidityBuffer.memoryAddress(), 0xff.toByte, validityBytes)
+  /**
+   * Write source vectors directly beneath the non-null struct advertised by `writer`.
+   *
+   * VectorUnloader recursively retains the source buffers without moving them between allocators.
+   * The wrapping record batch takes its own references, so closing both temporary batches
+   * restores the original reference counts after the synchronous pipe write. The borrowed
+   * VectorSchemaRoot must never be closed because its vectors are owned by the input
+   * ColumnarBatch.
+   */
+  private[python] def writeDirectBatch(
+      writer: DirectArrowStreamWriter,
+      sourceVectors: Seq[FieldVector],
+      numRows: Int,
+      allocator: BufferAllocator): Unit = {
+    val sourceRoot =
+      new VectorSchemaRoot(sourceVectors.map(_.getField).asJava, sourceVectors.asJava, numRows)
+    val sourceBatch = new VectorUnloader(sourceRoot).getRecordBatch
+    try {
+      val validityBytes = (numRows.toLong + 7L) / 8L
+      val structValidity = allocator.buffer(validityBytes)
+      try {
+        if (validityBytes > 0) {
+          structValidity.setOne(0L, validityBytes)
+        }
+        structValidity.writerIndex(validityBytes)
+
+        val nodes = new ArrayList[ArrowFieldNode](sourceBatch.getNodes.size() + 1)
+        nodes.add(new ArrowFieldNode(numRows, 0))
+        nodes.addAll(sourceBatch.getNodes)
+
+        val buffers = new ArrayList[ArrowBuf](sourceBatch.getBuffers.size() + 1)
+        buffers.add(structValidity)
+        buffers.addAll(sourceBatch.getBuffers)
+
+        val wrappedBatch = new ArrowRecordBatch(
+          numRows,
+          nodes,
+          buffers,
+          sourceBatch.getBodyCompression,
+          sourceBatch.getVariadicBufferCounts,
+          true)
+        try {
+          writer.writeDirect(wrappedBatch)
+        } finally {
+          wrappedBatch.close()
+        }
+      } finally {
+        structValidity.close()
+      }
+    } finally {
+      sourceBatch.close()
     }
   }
 }
diff --git a/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py b/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
index 19d9dac5d32..b179ab76060 100644
--- a/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
+++ b/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
@@ -31,8 +31,8 @@
   * vanilla:   CometScan -> ColumnarToRow + UnsafeProjection -> ArrowPythonRunner
               (per-row InternalRow.getXXX() loop inside ArrowWriter.write)
   * optimized: CometScan -> CometMapInBatchExec -> CometArrowPythonRunner
-              (per-buffer Unsafe.copyMemory from Comet's vectors into the
-              runner's persistent VectorSchemaRoot; no row materialization)
+              (Arrow IPC serialization directly from Comet's source vectors;
+              no row materialization or intermediate vector-buffer copy)
 
 Results are wall-clock seconds, so they include Python interpreter,
 Arrow IPC, and downstream count() costs. That's intentional: the
@@ -49,7 +49,7 @@
     # Build Comet (release for representative numbers):
     make release
 
-    pip install pyspark==3.5.9 pyarrow pandas
+    pip install pyspark==4.0.4 pyarrow pandas
 
     python3 spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
 
@@ -84,6 +84,8 @@ def _build_spark() -> SparkSession:
         .config("spark.plugins", "org.apache.spark.CometPlugin")
         .config("spark.comet.enabled", "true")
         .config("spark.comet.exec.enabled", "true")
+        # Keep Comet's scan and execution rules active with Spark's default shuffle manager.
+        .config("spark.comet.shuffle.enabled", "false")
         .config("spark.memory.offHeap.enabled", "true")
         .config("spark.memory.offHeap.size", "4g")
         .config("spark.driver.memory", "4g")
@@ -157,6 +159,10 @@ def _time_run(spark: SparkSession, parquet_path: str, accelerate: bool, api: str
         df = df.mapInArrow(_passthrough_arrow, schema)
     else:
         df = df.mapInPandas(_passthrough_pandas, schema)
+    plan = df._jdf.queryExecution().executedPlan().toString()
+    if ("CometMapInBatch" in plan) != accelerate:
+        expected = "CometMapInBatch" if accelerate else "vanilla Python execution"
+        raise RuntimeError(f"Expected {expected} for {api}, but found:\n{plan}")
     t0 = time.perf_counter()
     df.count()
     return time.perf_counter() - t0
diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf.py b/spark/src/test/resources/pyspark/test_pyarrow_udf.py
index 84f405e8b2c..60ae72069b5 100644
--- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py
+++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py
@@ -361,11 +361,11 @@ def test_map_in_arrow_decimal_precision_sweep(
     spark, tmp_path, accelerated, precision, scale
 ):
     """
-    The Arrow `DecimalVector` that `copyVector` touches is always 16 bytes wide regardless of
-    precision, so there is no buffer-width boundary on the Arrow path (the 8-byte long-backed form
-    is Spark's `UnsafeRow` encoding, a layer this Arrow buffer copy never sees). This sweep instead
-    guards the precision/scale extremes and the 18/19 point where Spark's own decimal handling
-    changes representation, keeping the round trip value-exact. Scale extremes: 0, half, max.
+    Arrow's `DecimalVector` is always 16 bytes wide regardless of precision, so there is no
+    buffer-width boundary on the direct Arrow path (the 8-byte long-backed form is Spark's
+    `UnsafeRow` encoding, a layer this path never sees). This sweep instead guards the
+    precision/scale extremes and the 18/19 point where Spark's own decimal handling changes
+    representation, keeping the round trip value-exact. Scale extremes: 0, half, max.
     """
     schema_in = T.StructType(
         [
@@ -405,10 +405,9 @@ def test_map_in_arrow_null_density_sweep(
     spark, tmp_path, accelerated, null_fraction
 ):
     """
-    Validity-buffer memcpy is where Arrow Java vector copies historically break. Sweep null
-    density across the corner cases: all-non-null, sparse-null, half-null, sparse-non-null,
-    all-null. Catches off-by-one in validity packing and edge cases where source/destination
-    null counts diverge.
+    Sweep validity-buffer density across the corner cases: all-non-null, sparse-null, half-null,
+    sparse-non-null, all-null. Catches off-by-one errors in validity packing and mismatches
+    between field-node null counts and serialized validity buffers.
     """
     schema_in = T.StructType(
         [
@@ -437,11 +436,10 @@ def passthrough(iterator):
 
 def test_map_in_arrow_multi_batch_per_partition(spark, tmp_path, accelerated):
     """
-    Force many small batches in a single partition so the writer runs its per-batch
-    allocate/copy/write loop hundreds of times against a reused struct container (the leaf
-    buffers are reallocated each batch today; see #4383). Catches errors that only appear across
-    the batch boundary: stale value counts, offset/validity sizing on the second and later
-    batches, and variable-width data-buffer sizing as row content changes batch to batch.
+    Exercise the stream across many rows with a small Spark Arrow batch limit. The accelerated
+    path writes complete Comet source batches, so its source-vector turnover is covered separately
+    by test_map_in_arrow_nested_source_batches below. Catches stale value counts and changing
+    variable-width content on the Spark fallback path.
     """
     schema_in = T.StructType(
         [
@@ -472,10 +470,125 @@ def passthrough(iterator):
         spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", prev_records)
 
 
+def test_map_in_arrow_nested_source_batches(spark, tmp_path, accelerated):
+    """Serialize changing nested/null source vectors directly across actual worker batches."""
+    item_type = T.StructType(
+        [
+            T.StructField("label", T.StringType()),
+            T.StructField("score", T.IntegerType()),
+        ]
+    )
+    payload_type = T.StructType(
+        [
+            T.StructField("items", T.ArrayType(item_type, containsNull=True)),
+            T.StructField(
+                "attrs",
+                T.MapType(T.StringType(), T.IntegerType(), valueContainsNull=True),
+            ),
+        ]
+    )
+    input_schema = T.StructType(
+        [
+            T.StructField("id", T.LongType(), nullable=False),
+            T.StructField("payload", payload_type),
+        ]
+    )
+    output_schema = T.StructType(
+        [
+            *input_schema.fields,
+            T.StructField("batch_index", T.IntegerType(), nullable=False),
+            T.StructField("batch_size", T.IntegerType(), nullable=False),
+        ]
+    )
+
+    rows = []
+    for index in range(37):
+        if index % 11 == 0:
+            payload = None
+        else:
+            if index % 5 == 0:
+                items = None
+            else:
+                items = [
+                    None
+                    if (index + offset) % 6 == 0
+                    else {
+                        "label": None
+                        if (index + offset) % 4 == 0
+                        else f"item_{index}_{'x' * offset}",
+                        "score": None
+                        if (index + offset) % 3 == 0
+                        else index * 10 + offset,
+                    }
+                    for offset in range(index % 4)
+                ]
+
+            if index % 7 == 0:
+                attrs = None
+            elif index % 6 == 0:
+                attrs = {}
+            else:
+                attrs = {
+                    f"key_{index % 3}": None if index % 4 == 0 else index,
+                    "marker": index * 10,
+                }
+            payload = {"items": items, "attrs": attrs}
+        rows.append((index, payload))
+
+    src = str(tmp_path / "nested_source_batches.parquet")
+    spark.createDataFrame(rows, input_schema).coalesce(1).write.parquet(src)
+
+    batch_size = 7
+    previous_comet_batch_size = spark.conf.get("spark.comet.batchSize", "8192")
+    previous_arrow_batch_size = spark.conf.get(
+        "spark.sql.execution.arrow.maxRecordsPerBatch"
+    )
+    spark.conf.set("spark.comet.batchSize", str(batch_size))
+    spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", str(batch_size))
+    try:
+
+        def annotate_batches(iterator):
+            for batch_index, batch in enumerate(iterator):
+                assert batch.schema.names == ["id", "payload"]
+                yield pa.RecordBatch.from_arrays(
+                    [
+                        *batch.columns,
+                        pa.array([batch_index] * batch.num_rows, type=pa.int32()),
+                        pa.array([batch.num_rows] * batch.num_rows, type=pa.int32()),
+                    ],
+                    names=[*batch.schema.names, "batch_index", "batch_size"],
+                )
+
+        result_df = spark.read.parquet(src).mapInArrow(annotate_batches, output_schema)
+        _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+        output = result_df.collect()
+        assert len(output) == len(rows)
+        expected_payloads = dict(rows)
+        observed_batches = {}
+        for row in output:
+            payload = row["payload"]
+            actual_payload = None if payload is None else payload.asDict(recursive=True)
+            assert actual_payload == expected_payloads[row["id"]]
+            observed_batches.setdefault(row["batch_index"], []).append(row)
+
+        assert sorted(observed_batches) == list(range(6))
+        expected_batch_sizes = [batch_size] * 5 + [2]
+        observed_batch_sizes = [len(observed_batches[index]) for index in range(6)]
+        assert observed_batch_sizes == expected_batch_sizes
+        for batch_rows in observed_batches.values():
+            assert {row["batch_size"] for row in batch_rows} == {len(batch_rows)}
+    finally:
+        spark.conf.set("spark.comet.batchSize", previous_comet_batch_size)
+        spark.conf.set(
+            "spark.sql.execution.arrow.maxRecordsPerBatch", previous_arrow_batch_size
+        )
+
+
 def test_map_in_arrow_wide_schema(spark, tmp_path, accelerated):
     """
-    50-column mixed-type schema. The bulk-copy path walks a flattened addresses[] array indexed
-    across the whole vector tree; off-by-one in flattening logic surfaces at depth * width.
+    50-column mixed-type schema. Direct serialization must preserve every source vector and its
+    independent null bitmap when building the IPC field-node and buffer lists.
     """
     fields = [T.StructField("id", T.LongType())]
     for i in range(15):
@@ -664,10 +777,9 @@ def _normalize(row):
 
 def test_map_in_arrow_numeric_scalars(spark, tmp_path, accelerated):
     """
-    Covers the BaseFixedWidthVector branch in CometColumnarPythonInput.copyVector for
-    every fixed-width primitive Comet's scan supports beyond the long/double/int already
-    exercised by other tests: boolean, byte, short, float. Each has a distinct buffer
-    size, and the validity bit handling is independent per column.
+    Covers every fixed-width primitive Comet's scan supports beyond the long/double/int already
+    exercised by other tests: boolean, byte, short, float. Each has a distinct buffer size, and
+    the validity bit handling is independent per directly serialized source vector.
     """
     schema_in = T.StructType(
         [
@@ -781,8 +893,8 @@ def test_map_in_arrow_map_type(spark, tmp_path, accelerated):
     """
     MapType is encoded in Arrow as a List<Struct<key, value>> with extra metadata. The
     buffer layout (offsets + struct child + key/value children) is distinct from a plain
-    list, and CometMapVector is a separate vector class from CometListVector. Without
-    this test the recursive copy path through map-typed columns is unexercised.
+    list, and CometMapVector is a separate vector class from CometListVector. Without this test,
+    direct recursive serialization of map-typed source vectors is unexercised.
     """
     schema_in = T.StructType(
         [
@@ -830,11 +942,10 @@ def _normalize(row):
 
 def test_map_in_arrow_deeply_nested(spark, tmp_path, accelerated):
     """
-    Exercises the recursive descent in CometColumnarPythonInput.copyVector at depth > 1,
-    in every nesting combination: array-of-array, array-of-struct, struct-of-array,
-    struct-of-struct. Single-level nesting is covered by test_map_in_arrow_array_and_struct;
-    the bug surface here is that setLastSet / setValueCount must be applied bottom-up
-    correctly at every level.
+    Exercises direct source-vector serialization at depth > 1, in every nesting combination:
+    array-of-array, array-of-struct, struct-of-array, struct-of-struct. Single-level nesting is
+    covered by test_map_in_arrow_array_and_struct; the bug surface here is that Arrow field
+    nodes and buffers must remain in the expected depth-first order at every level.
     """
     schema_in = T.StructType(
         [
@@ -952,9 +1063,9 @@ def _norm_input_config(c):
 def test_map_in_arrow_falls_back_when_use_large_var_types(spark, tmp_path):
     """
     `spark.sql.execution.arrow.useLargeVarTypes=true` widens StringType / BinaryType to
-    LargeUtf8 / LargeBinary in the destination IPC root (8-byte offsets). Comet's source
-    vectors always use 4-byte offsets; CometColumnarPythonInput.copyVector does a raw
-    setBytes per buffer and would corrupt the offset buffer in this configuration.
+    LargeUtf8 / LargeBinary in Spark's expected IPC schema (8-byte offsets). Comet's source
+    vectors always use 4-byte offsets; serializing those source vectors directly cannot
+    satisfy the 8-byte offset schema expected in this configuration.
     EliminateRedundantTransitions must skip the rewrite in that case so vanilla Spark
     handles the operation. This test does not use the `accelerated` fixture: it sets
     pyarrowUDF.enabled=true AND useLargeVarTypes=true and asserts the plan still falls
diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
new file mode 100644
index 00000000000..682676a840c
--- /dev/null
+++ b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
@@ -0,0 +1,329 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, IOException}
+import java.nio.ByteBuffer
+import java.nio.channels.{Channels, WritableByteChannel}
+
+import scala.jdk.CollectionConverters._
+
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatest.matchers.should.Matchers
+
+import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
+import org.apache.arrow.vector.{FieldVector, IntVector, NullVector, VarCharVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.ipc.ArrowStreamReader
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.{writeDirectBatch, DirectArrowStreamWriter}
+
+class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
+
+  private def withWriter(
+      childFields: Seq[Field],
+      allocator: BufferAllocator,
+      channel: WritableByteChannel)(f: DirectArrowStreamWriter => Unit): Unit = {
+    val structField = new Field(
+      "struct",
+      new FieldType(false, ArrowType.Struct.INSTANCE, null),
+      childFields.asJava)
+    val root = VectorSchemaRoot.create(new Schema(Seq(structField).asJava), allocator)
+    val writer = new DirectArrowStreamWriter(root, channel)
+    try {
+      writer.start()
+      f(writer)
+      writer.end()
+    } finally {
+      writer.close()
+      root.close()
+    }
+  }
+
+  private def withReader(bytes: Array[Byte])(f: ArrowStreamReader => Unit): Unit = {
+    val allocator = new RootAllocator(Long.MaxValue)
+    val reader = new ArrowStreamReader(new ByteArrayInputStream(bytes), allocator)
+    try {
+      f(reader)
+    } finally {
+      reader.close()
+      allocator.close()
+    }
+  }
+
+  test("direct batches retain borrowed buffers without copying them into the writer allocator") {
+    val sourceAllocator = new RootAllocator(Long.MaxValue)
+    val writerAllocator = new RootAllocator(1024)
+    val vector = new VarCharVector("source_name", sourceAllocator)
+    val output = new ByteArrayOutputStream()
+    try {
+      val payload = Array.fill[Byte](16 * 1024)('x'.toByte)
+      vector.allocateNew(payload.length.toLong, 2)
+      vector.setSafe(0, payload)
+      vector.setNull(1)
+      vector.setValueCount(2)
+
+      val field = new Field("payload", vector.getField.getFieldType, vector.getField.getChildren)
+      val buffers = vector.getFieldBuffers.asScala.toSeq
+      val originalReferenceCounts = buffers.map(_.refCnt())
+      val originalLastSet = vector.getLastSet
+
+      withWriter(Seq(field), writerAllocator, Channels.newChannel(output)) { writer =>
+        val originalWriterAllocation = writerAllocator.getAllocatedMemory
+        writeDirectBatch(writer, Seq(vector), 2, writerAllocator)
+
+        writerAllocator.getAllocatedMemory shouldBe originalWriterAllocation
+        buffers.map(_.refCnt()) shouldBe originalReferenceCounts
+        vector.getLastSet shouldBe originalLastSet
+        vector.getValueCount shouldBe 2
+        vector.get(0) shouldBe payload
+        vector.isNull(1) shouldBe true
+      }
+
+      withReader(output.toByteArray) { reader =>
+        reader.loadNextBatch() shouldBe true
+        val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector]
+        struct.getNullCount shouldBe 0
+        val result = struct.getChild("payload").asInstanceOf[VarCharVector]
+        result.get(0) shouldBe payload
+        result.isNull(1) shouldBe true
+        reader.loadNextBatch() shouldBe false
+      }
+    } finally {
+      vector.close()
+      writerAllocator.close()
+      sourceAllocator.close()
+    }
+  }
+
+  test("direct batches preserve nested list, struct, map, and null field layouts") {
+    val sourceAllocator = new RootAllocator(Long.MaxValue)
+    val writerAllocator = new RootAllocator(Long.MaxValue)
+    val list = ListVector.empty("items", sourceAllocator)
+    val struct = StructVector.empty("details", sourceAllocator)
+    val map = MapVector.empty("mapping", sourceAllocator, false)
+    val nulls = new NullVector("nulls", 3)
+    val output = new ByteArrayOutputStream()
+    try {
+      val listWriter = list.getWriter
+      listWriter.setPosition(0)
+      listWriter.startList()
+      listWriter.integer().writeInt(11)
+      listWriter.integer().writeInt(12)
+      listWriter.endList()
+      listWriter.setPosition(1)
+      listWriter.writeNull()
+      listWriter.setPosition(2)
+      listWriter.startList()
+      listWriter.integer().writeInt(13)
+      listWriter.endList()
+      listWriter.setValueCount(3)
+
+      val structWriter = struct.getWriter
+      structWriter.setPosition(0)
+      structWriter.start()
+      structWriter.integer("count").writeInt(21)
+      structWriter.end()
+      structWriter.setPosition(1)
+      structWriter.writeNull()
+      structWriter.setPosition(2)
+      structWriter.start()
+      structWriter.integer("count").writeNull()
+      structWriter.end()
+      structWriter.setValueCount(3)
+
+      val mapWriter = map.getWriter
+      mapWriter.setPosition(0)
+      mapWriter.startMap()
+      mapWriter.startEntry()
+      mapWriter.key().integer().writeInt(31)
+      mapWriter.value().integer().writeInt(32)
+      mapWriter.endEntry()
+      mapWriter.endMap()
+      mapWriter.setPosition(1)
+      mapWriter.writeNull()
+      mapWriter.setPosition(2)
+      mapWriter.startMap()
+      mapWriter.startEntry()
+      mapWriter.key().integer().writeInt(33)
+      mapWriter.value().integer().writeNull()
+      mapWriter.endEntry()
+      mapWriter.endMap()
+      mapWriter.setValueCount(3)
+
+      val vectors = Seq[FieldVector](list, struct, map, nulls)
+      withWriter(vectors.map(_.getField), writerAllocator, Channels.newChannel(output)) {
+        writer =>
+          writeDirectBatch(writer, vectors, 3, writerAllocator)
+      }
+
+      withReader(output.toByteArray) { reader =>
+        reader.loadNextBatch() shouldBe true
+        val result = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector]
+        result.getNullCount shouldBe 0
+
+        val resultList = result.getChild("items").asInstanceOf[ListVector]
+        resultList.getObject(0).asScala.toSeq shouldBe Seq(11, 12)
+        resultList.isNull(1) shouldBe true
+        resultList.getObject(2).asScala.toSeq shouldBe Seq(13)
+
+        val resultStruct = result.getChild("details").asInstanceOf[StructVector]
+        resultStruct.getChild("count").asInstanceOf[IntVector].get(0) shouldBe 21
+        resultStruct.isNull(1) shouldBe true
+        resultStruct.getChild("count").isNull(2) shouldBe true
+
+        val resultMap = result.getChild("mapping").asInstanceOf[MapVector]
+        val entries = resultMap.getDataVector.asInstanceOf[StructVector]
+        entries.getChildByOrdinal(0).getField.getName shouldBe MapVector.KEY_NAME
+        entries.getChildByOrdinal(1).getField.getName shouldBe MapVector.VALUE_NAME
+        entries.getChildByOrdinal(0).asInstanceOf[IntVector].get(0) shouldBe 31
+        entries.getChildByOrdinal(1).asInstanceOf[IntVector].get(0) shouldBe 32
+        resultMap.isNull(1) shouldBe true
+        entries.getChildByOrdinal(1).isNull(1) shouldBe true
+
+        val resultNulls = result.getChild("nulls").asInstanceOf[NullVector]
+        resultNulls.getNullCount shouldBe 3
+        reader.loadNextBatch() shouldBe false
+      }
+    } finally {
+      nulls.close()
+      map.close()
+      struct.close()
+      list.close()
+      writerAllocator.close()
+      sourceAllocator.close()
+    }
+  }
+
+  test("direct batches preserve zero-row batches between populated batches") {
+    val sourceAllocator = new RootAllocator(Long.MaxValue)
+    val writerAllocator = new RootAllocator(Long.MaxValue)
+    val first = new IntVector("value", sourceAllocator)
+    val empty = new IntVector("value", sourceAllocator)
+    val last = new IntVector("value", sourceAllocator)
+    val output = new ByteArrayOutputStream()
+    try {
+      first.allocateNew(2)
+      first.setSafe(0, 41)
+      first.setSafe(1, 42)
+      first.setValueCount(2)
+      empty.setValueCount(0)
+      last.allocateNew(1)
+      last.setSafe(0, 43)
+      last.setValueCount(1)
+
+      withWriter(Seq(first.getField), writerAllocator, Channels.newChannel(output)) { writer =>
+        writeDirectBatch(writer, Seq(first), 2, writerAllocator)
+        writeDirectBatch(writer, Seq(empty), 0, writerAllocator)
+        writeDirectBatch(writer, Seq(last), 1, writerAllocator)
+      }
+
+      withReader(output.toByteArray) { reader =>
+        reader.loadNextBatch() shouldBe true
+        reader.getVectorSchemaRoot.getRowCount shouldBe 2
+        reader.loadNextBatch() shouldBe true
+        reader.getVectorSchemaRoot.getRowCount shouldBe 0
+        reader.loadNextBatch() shouldBe true
+        reader.getVectorSchemaRoot.getRowCount shouldBe 1
+        val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector]
+        struct.getChild("value").asInstanceOf[IntVector].get(0) shouldBe 43
+        reader.loadNextBatch() shouldBe false
+      }
+    } finally {
+      last.close()
+      empty.close()
+      first.close()
+      writerAllocator.close()
+      sourceAllocator.close()
+    }
+  }
+
+  test("direct batches represent non-null structs with no child columns") {
+    val allocator = new RootAllocator(Long.MaxValue)
+    val output = new ByteArrayOutputStream()
+    try {
+      withWriter(Seq.empty, allocator, Channels.newChannel(output)) { writer =>
+        writeDirectBatch(writer, Seq.empty, 3, allocator)
+      }
+
+      withReader(output.toByteArray) { reader =>
+        reader.loadNextBatch() shouldBe true
+        val struct = reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector]
+        struct.getValueCount shouldBe 3
+        struct.getNullCount shouldBe 0
+        struct.getChildrenFromFields.isEmpty shouldBe true
+        reader.loadNextBatch() shouldBe false
+      }
+    } finally {
+      allocator.close()
+    }
+  }
+
+  test("direct batches release temporary references when writing the stream fails") {
+    val sourceAllocator = new RootAllocator(Long.MaxValue)
+    val writerAllocator = new RootAllocator(Long.MaxValue)
+    val source = new IntVector("value", sourceAllocator)
+    val output = new ByteArrayOutputStream()
+    var failWrites = false
+    val channel = new WritableByteChannel {
+      private var open = true
+
+      override def isOpen: Boolean = open
+
+      override def close(): Unit = open = false
+
+      override def write(buffer: ByteBuffer): Int = {
+        if (failWrites) {
+          throw new IOException("injected Arrow IPC write failure")
+        }
+        val bytes = new Array[Byte](buffer.remaining())
+        buffer.get(bytes)
+        output.write(bytes)
+        bytes.length
+      }
+    }
+    try {
+      source.allocateNew(1)
+      source.setSafe(0, 51)
+      source.setValueCount(1)
+
+      withWriter(Seq(source.getField), writerAllocator, channel) { writer =>
+        val originalReferenceCounts = source.getFieldBuffers.asScala.map(_.refCnt()).toSeq
+        val originalWriterAllocation = writerAllocator.getAllocatedMemory
+        failWrites = true
+        try {
+          val error = intercept[IOException] {
+            writeDirectBatch(writer, Seq(source), 1, writerAllocator)
+          }
+          error.getMessage shouldBe "injected Arrow IPC write failure"
+        } finally {
+          failWrites = false
+        }
+        source.getFieldBuffers.asScala.map(_.refCnt()).toSeq shouldBe originalReferenceCounts
+        writerAllocator.getAllocatedMemory shouldBe originalWriterAllocation
+        source.get(0) shouldBe 51
+      }
+    } finally {
+      source.close()
+      writerAllocator.close()
+      sourceAllocator.close()
+    }
+  }
+}

From 72a0db52988e1a13a1460738e3dd0de5d7d2a5b4 Mon Sep 17 00:00:00 2001
From: Chao Sun <chao@openai.com>
Date: Sat, 15 Aug 2026 08:37:22 -0700
Subject: [PATCH 2/3] test: register direct Arrow suite and fix benchmark
 profile

---
 .github/workflows/pr_build_linux.yml                      | 1 +
 .github/workflows/pr_build_macos.yml                      | 1 +
 spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py | 2 +-
 3 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 18dc8e4086c..4a640a03b46 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -342,6 +342,7 @@ jobs:
               org.apache.comet.exec.CometWindowExecSuite
               org.apache.comet.exec.CometJoinSuite
               org.apache.spark.sql.comet.CometMapInBatchSuite
+              org.apache.spark.sql.execution.python.CometArrowPythonRunnerSuite
               org.apache.comet.CometNativeSuite
               org.apache.comet.CometConfSuite
               org.apache.comet.CometPublicApiSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index fa82e4ebca7..8ae0e3f135a 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -158,6 +158,7 @@ jobs:
               org.apache.comet.exec.CometWindowExecSuite
               org.apache.comet.exec.CometJoinSuite
               org.apache.spark.sql.comet.CometMapInBatchSuite
+              org.apache.spark.sql.execution.python.CometArrowPythonRunnerSuite
               org.apache.comet.CometNativeSuite
               org.apache.comet.CometConfSuite
               org.apache.comet.CometPublicApiSuite
diff --git a/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py b/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
index b179ab76060..92d58b489a0 100644
--- a/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
+++ b/spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py
@@ -47,7 +47,7 @@
 
 Usage:
     # Build Comet (release for representative numbers):
-    make release
+    make release PROFILES='-Pspark-4.0 -Pscala-2.13'
 
     pip install pyspark==4.0.4 pyarrow pandas
 

From 64a0114a71d3dfc32a526369d4431b07378b3286 Mon Sep 17 00:00:00 2001
From: Chao Sun <chao@openai.com>
Date: Wed, 19 Aug 2026 14:24:12 -0700
Subject: [PATCH 3/3] refactor: use public Arrow IPC serialization for Python
 input

---
 .../python/CometArrowPythonRunnerBase.scala   | 37 ++++++-------------
 .../python/CometArrowPythonRunnerSuite.scala  | 34 ++++++++---------
 2 files changed, 28 insertions(+), 43 deletions(-)

diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
index 03f5b18c66a..a3d516a85e9 100644
--- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
@@ -20,7 +20,7 @@
 package org.apache.spark.sql.execution.python
 
 import java.io.{DataInputStream, DataOutputStream}
-import java.nio.channels.{Channels, WritableByteChannel}
+import java.nio.channels.Channels
 import java.util.ArrayList
 import java.util.concurrent.atomic.AtomicBoolean
 
@@ -29,8 +29,8 @@ import scala.jdk.CollectionConverters._
 import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
 import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot, VectorUnloader}
 import org.apache.arrow.vector.complex.StructVector
-import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
-import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel}
+import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch, MessageSerializer}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
@@ -109,7 +109,7 @@ private[python] trait CometArrowPythonRunnerBase
       private val allocator =
         CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
       private var currentGroup: Iterator[ColumnarBatch] = _
-      private var arrowWriter: CometArrowPythonRunnerBase.DirectArrowStreamWriter = _
+      private var arrowWriter: ArrowStreamWriter = _
       private var writeRoot: VectorSchemaRoot = _
 
       // The runner's input schema is a single struct column ("struct") whose children are the
@@ -142,9 +142,7 @@ private[python] trait CometArrowPythonRunnerBase
             childFields.asJava)
         val structVec = structField.createVector(allocator).asInstanceOf[StructVector]
         writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
-        arrowWriter = new CometArrowPythonRunnerBase.DirectArrowStreamWriter(
-          writeRoot,
-          Channels.newChannel(dataOut))
+        arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
         arrowWriter.start()
       }
 
@@ -196,8 +194,8 @@ private[python] trait CometArrowPythonRunnerBase
           startWriter(childFields, dataOut)
         }
 
-        CometArrowPythonRunnerBase.writeDirectBatch(
-          arrowWriter,
+        CometArrowPythonRunnerBase.serializeBatch(
+          new WriteChannel(Channels.newChannel(dataOut)),
           sourceVectors,
           cometBatch.numRows(),
           allocator)
@@ -320,20 +318,7 @@ private[python] trait CometArrowPythonRunnerBase
 private[python] object CometArrowPythonRunnerBase {
 
   /**
-   * Expose ArrowWriter's existing record-batch serialization without unloading its schema root.
-   */
-  private[python] final class DirectArrowStreamWriter(
-      root: VectorSchemaRoot,
-      channel: WritableByteChannel)
-      extends ArrowStreamWriter(root, null, channel) {
-
-    def writeDirect(batch: ArrowRecordBatch): Unit = {
-      writeRecordBatch(batch)
-    }
-  }
-
-  /**
-   * Write source vectors directly beneath the non-null struct advertised by `writer`.
+   * Serialize source vectors directly beneath the non-null struct advertised in the IPC stream.
    *
    * VectorUnloader recursively retains the source buffers without moving them between allocators.
    * The wrapping record batch takes its own references, so closing both temporary batches
@@ -341,8 +326,8 @@ private[python] object CometArrowPythonRunnerBase {
    * VectorSchemaRoot must never be closed because its vectors are owned by the input
    * ColumnarBatch.
    */
-  private[python] def writeDirectBatch(
-      writer: DirectArrowStreamWriter,
+  private[python] def serializeBatch(
+      writeChannel: WriteChannel,
       sourceVectors: Seq[FieldVector],
       numRows: Int,
       allocator: BufferAllocator): Unit = {
@@ -374,7 +359,7 @@ private[python] object CometArrowPythonRunnerBase {
           sourceBatch.getVariadicBufferCounts,
           true)
         try {
-          writer.writeDirect(wrappedBatch)
+          MessageSerializer.serialize(writeChannel, wrappedBatch)
         } finally {
           wrappedBatch.close()
         }
diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
index 682676a840c..81f3a76551f 100644
--- a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
+++ b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
@@ -31,25 +31,25 @@ import org.scalatest.matchers.should.Matchers
 import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
 import org.apache.arrow.vector.{FieldVector, IntVector, NullVector, VarCharVector, VectorSchemaRoot}
 import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
-import org.apache.arrow.vector.ipc.ArrowStreamReader
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
-import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.{writeDirectBatch, DirectArrowStreamWriter}
+import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.serializeBatch
 
 class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
 
   private def withWriter(
       childFields: Seq[Field],
       allocator: BufferAllocator,
-      channel: WritableByteChannel)(f: DirectArrowStreamWriter => Unit): Unit = {
+      channel: WritableByteChannel)(f: WritableByteChannel => Unit): Unit = {
     val structField = new Field(
       "struct",
       new FieldType(false, ArrowType.Struct.INSTANCE, null),
       childFields.asJava)
     val root = VectorSchemaRoot.create(new Schema(Seq(structField).asJava), allocator)
-    val writer = new DirectArrowStreamWriter(root, channel)
+    val writer = new ArrowStreamWriter(root, null, channel)
     try {
       writer.start()
-      f(writer)
+      f(channel)
       writer.end()
     } finally {
       writer.close()
@@ -85,9 +85,9 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
       val originalReferenceCounts = buffers.map(_.refCnt())
       val originalLastSet = vector.getLastSet
 
-      withWriter(Seq(field), writerAllocator, Channels.newChannel(output)) { writer =>
+      withWriter(Seq(field), writerAllocator, Channels.newChannel(output)) { channel =>
         val originalWriterAllocation = writerAllocator.getAllocatedMemory
-        writeDirectBatch(writer, Seq(vector), 2, writerAllocator)
+        serializeBatch(new WriteChannel(channel), Seq(vector), 2, writerAllocator)
 
         writerAllocator.getAllocatedMemory shouldBe originalWriterAllocation
         buffers.map(_.refCnt()) shouldBe originalReferenceCounts
@@ -170,8 +170,8 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
 
       val vectors = Seq[FieldVector](list, struct, map, nulls)
       withWriter(vectors.map(_.getField), writerAllocator, Channels.newChannel(output)) {
-        writer =>
-          writeDirectBatch(writer, vectors, 3, writerAllocator)
+        channel =>
+          serializeBatch(new WriteChannel(channel), vectors, 3, writerAllocator)
       }
 
       withReader(output.toByteArray) { reader =>
@@ -229,10 +229,10 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
       last.setSafe(0, 43)
       last.setValueCount(1)
 
-      withWriter(Seq(first.getField), writerAllocator, Channels.newChannel(output)) { writer =>
-        writeDirectBatch(writer, Seq(first), 2, writerAllocator)
-        writeDirectBatch(writer, Seq(empty), 0, writerAllocator)
-        writeDirectBatch(writer, Seq(last), 1, writerAllocator)
+      withWriter(Seq(first.getField), writerAllocator, Channels.newChannel(output)) { channel =>
+        serializeBatch(new WriteChannel(channel), Seq(first), 2, writerAllocator)
+        serializeBatch(new WriteChannel(channel), Seq(empty), 0, writerAllocator)
+        serializeBatch(new WriteChannel(channel), Seq(last), 1, writerAllocator)
       }
 
       withReader(output.toByteArray) { reader =>
@@ -259,8 +259,8 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
     val allocator = new RootAllocator(Long.MaxValue)
     val output = new ByteArrayOutputStream()
     try {
-      withWriter(Seq.empty, allocator, Channels.newChannel(output)) { writer =>
-        writeDirectBatch(writer, Seq.empty, 3, allocator)
+      withWriter(Seq.empty, allocator, Channels.newChannel(output)) { channel =>
+        serializeBatch(new WriteChannel(channel), Seq.empty, 3, allocator)
       }
 
       withReader(output.toByteArray) { reader =>
@@ -304,13 +304,13 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
       source.setSafe(0, 51)
       source.setValueCount(1)
 
-      withWriter(Seq(source.getField), writerAllocator, channel) { writer =>
+      withWriter(Seq(source.getField), writerAllocator, channel) { channel =>
         val originalReferenceCounts = source.getFieldBuffers.asScala.map(_.refCnt()).toSeq
         val originalWriterAllocation = writerAllocator.getAllocatedMemory
         failWrites = true
         try {
           val error = intercept[IOException] {
-            writeDirectBatch(writer, Seq(source), 1, writerAllocator)
+            serializeBatch(new WriteChannel(channel), Seq(source), 1, writerAllocator)
           }
           error.getMessage shouldBe "injected Arrow IPC write failure"
         } finally {