feat: Add experimental support for accelerated PyArrow UDFs

[andygrove]

Which issue does this PR close?

Closes #957
Closes #4240

Rationale for this change

Spark's mapInArrow / mapInPandas plans insert ColumnarToRow between a Comet scan and the Python operator, and MapInBatchExec then projects the Python output back to rows. Both projections copy data per row even though the upstream Comet plan and the Python runner both speak Arrow. Inside the Python runner, ArrowPythonRunner then re-encodes those rows into Arrow IPC, so the data round-trips Arrow → Row → Arrow on the JVM side before ever reaching Python.

This PR replaces the row-based plumbing with a custom CometArrowPythonRunner that consumes Iterator[ColumnarBatch] directly. The destination VectorSchemaRoot (in Spark's allocator, used for the IPC stream) is populated per batch by a single Unsafe.copyMemory per buffer per column, copying straight from Comet's vectors into the runner's root. No row materialization, no ArrowWriter.write(InternalRow) loop.

Not zero-copy in the strict sense: Comet's Parquet readers each construct their own RootAllocator, separate from ArrowUtils.rootAllocator, so Arrow's TransferPair cannot share buffers across the boundary. Bulk per-buffer memcpy is the next-best alternative and is materially faster than the per-row write loop, especially on wide rows where the row path is dominated by the InternalRow accessor overhead. A future PR that unifies the allocator parent would unlock true zero-copy via TransferPair.

I/O asymmetry: why memcpy on input but not on output

The input path memcpys; the output path does not, and that is structural rather than an oversight.

• Input (JVM → Python). Two Arrow vector trees already exist in different allocators. The source is Comet's ColumnarBatch, allocated from each Parquet reader's private RootAllocator. The destination is Spark's persistent IPC root, allocated from ArrowUtils.rootAllocator. Because the two roots are unrelated, Arrow's TransferPair / VectorLoader.load cannot rebind buffers across the boundary. CometColumnarPythonInput.copyVector walks the trees in lockstep and setBytes-copies each buffer. This is the bridge between the two allocator trees.
• Output (Python → JVM). There is no second tree to bridge to. Spark's BasicPythonArrowOutput constructs an ArrowStreamReader against a per-iterator BufferAllocator; loadNextBatch decodes the IPC frames coming back from the Python worker directly into a VectorSchemaRoot allocated from that allocator. Each FieldVector in the root is wrapped in an ArrowColumnVector and emitted as a ColumnarBatch. CometMapInBatchExec consumes those vectors as-is - it just unwraps the single struct child and forwards them. Nothing else in the JVM owns those buffers, so there is no boundary to cross.

In short, the input has two pre-existing trees in different allocators (memcpy required); the output has one tree, decoded straight into the right allocator (no copy possible or needed). The future allocator-unification work that would unlock TransferPair on input does not apply to output - the output is already as good as it can get.

Plan (Spark 4.0)

Baseline (spark.comet.exec.pyarrowUdf.enabled=false):

*(1) Filter (isnotnull(id#6L) AND (id#6L < 100))
+- MapInArrow transform(id#3L, value#4)#5, [id#6L, value#7], false
   +- CometNativeColumnarToRow
      +- CometNativeScan parquet [id#3L,value#4] Batched: true, DataFilters: [],
         Format: CometParquet, Location: InMemoryFileIndex(1 paths)[file:/tmp/...],
         PartitionFilters: [], PushedFilters: [],
         ReadSchema: struct<id:bigint,value:double>

Optimized (spark.comet.exec.pyarrowUdf.enabled=true):

*(1) Filter (isnotnull(id#18L) AND (id#18L < 100))
+- CometMapInBatch transform(id#15L, value#16)#17, [id#18L, value#19], false, 207
   +- CometNativeScan parquet [id#15L,value#16] Batched: true, DataFilters: [],
      Format: CometParquet, Location: InMemoryFileIndex(1 paths)[file:/tmp/...],
      PartitionFilters: [], PushedFilters: [],
      ReadSchema: struct<id:bigint,value:double>

What changes are included in this PR?

• New input trait CometColumnarPythonInput under spark/src/main/spark-4.x/.../execution/python/. Extends Spark's private[python] PythonArrowInput[Iterator[ColumnarBatch]]; implements writeNextBatchToArrowStream by walking the destination struct's children, allocating each sized to match the corresponding Comet column, collecting buffer addresses, and issuing one bulk copy for the whole tree.
• New CometArrowPythonRunner per Spark minor (4.0, 4.1, 4.2) under the same package. 4.0 extends BasePythonRunner directly because Spark 4.0's BaseArrowPythonRunner is bound to Iterator[InternalRow]; 4.1/4.2 extend the generic BaseArrowPythonRunner[IN, OUT]. All three mix in CometColumnarPythonInput plus Spark's BasicPythonArrowOutput.
• New helper CometVectorIpcCopier in comet-common crosses the shading boundary using only long primitives. comet-common shades org.apache.arrow.* into org.apache.comet.shaded.arrow.*; comet-spark references the unshaded Spark Arrow classes. The helper does the byte-level walk over the source vector tree inside common, exposing bufferReadableBytes(): long[], valueCounts(): int[], and copyBuffersToAddresses(addresses: long[]): Unit. No shaded type crosses the API.
• CometMapInBatchExec loses all row plumbing (rowIterator, BatchIterator, ContextAwareIterator, InternalRow(_) wrap). Feeds Iterator[ColumnarBatch] straight to the runner; on output, unwraps the single struct column into the user's output columns as before. isBarrier is propagated through RDD.barrier() so mapInArrow(..., barrier=True) keeps its gang-scheduling semantics.
• EliminateRedundantTransitions rewrite unchanged in shape: matches ColumnarToRow + (PythonMapInArrowExec | MapInPandasExec | MapInArrowExec) over a columnar Comet child.
• Per-Spark-version shim for the differing ArrowPythonRunner constructors and the 4.0+ MapInArrowExec rename. Shared 4.x bits live in spark-4.x/.../Spark4xMapInBatchSupport.scala; per-minor shims provide only the runner factory.
• Spark 3.4 and 3.5 are no-ops. 3.5's PythonArrowInput trait has a different contract (writeIteratorToArrowStream one-shot vs 4.x's writeNextBatchToArrowStream batch-at-a-time) and a separate implementation has not been written. The matchers in the 3.4 / 3.5 shims return None; vanilla Spark handles the operation. 3.5 support can be added later if there is user demand.
• New conf spark.comet.exec.pyarrowUdf.enabled, default false while experimental.
• User guide page documenting usage, plan flow, barrier semantics, the buffer-copy boundary, and limitations.

Limitations

• The optimization currently applies only to mapInArrow and mapInPandas. Scalar pandas UDFs (@pandas_udf) and grouped operations (applyInPandas) are not yet supported.
• Spark 4.0 or newer is required. Spark 3.4 and 3.5 are documented no-ops.
• The optimization fires across a shuffle only with Comet's columnar shuffle. A vanilla Spark Exchange outputs rows and breaks the precondition.
• Not zero-copy; bulk per-buffer memcpy. Comet's Parquet readers each construct their own RootAllocator, so cross-root TransferPair cannot be used. A future PR that has the readers allocate from ArrowUtils.rootAllocator would unlock zero-copy.

How are these changes tested?

• CometMapInBatchSuite under spark/src/test/spark-4.x/ covers the JVM-side rule and an end-to-end check: constructs a MapInArrowExec over a stub CometPlan leaf and asserts EliminateRedundantTransitions rewrites it to CometMapInBatchExec; runs a Parquet → mapInArrow query with primitives + nullable varchar and asserts row-equivalence with the un-rewritten output.
• pytest module test_pyarrow_udf.py covers mapInArrow and mapInPandas end-to-end against a real Python worker: nulls, empty input, Python exception propagation, decimal / date / timestamp / array / struct types, post-shuffle correctness, and barrier=True gang scheduling. 24 / 24 cases pass locally on Spark 4.0.
• Dedicated pyarrow_udf_test.yml workflow runs the pytest module on every PR against Spark 4.0 (covers the 4.x shim path; paths allowlist scoped to the feature files).

Wall-clock benchmark (benchmark_pyarrow_udf.py, 1M rows, 5 iters, local[2], Spark 4.0.2 / PySpark 4.0.1):

api	workload	vanilla median	optimized median	rows/s	speedup
mapInArrow	narrow primitives	0.268 s	0.140 s	7.5M	1.92x
mapInPandas	narrow primitives	0.269 s	0.141 s	7.4M	1.91x
mapInArrow	mixed with strings	0.489 s	0.273 s	3.8M	1.79x
mapInPandas	mixed with strings	0.504 s	0.283 s	3.7M	1.78x
mapInArrow	wide rows (50 cols)	2.542 s	1.795 s	584K	1.42x
mapInPandas	wide rows (50 cols)	2.588 s	1.846 s	568K	1.40x

The benchmark UDF is a pure passthrough on local[2], so most of the wall time is Spark's Python fork / IPC overhead. Real UDFs (PyArrow compute, pandas ops, model inference) increase the per-row Python cost and shrink the speedup ratio. The bulk-copy path improves the wide-row case the most because that's where the row path's per-InternalRow overhead is most concentrated.

[comphead]

I assume not only speed is improved, would also be interesting to check memory metrics

[parthchandra]

Some minor comments. Otherwise great PR!

[parthchandra]

+1 !

[parthchandra]

Other Comet operators extend CometPlan. The idea was that all Comet specific common behavior (say some Comet specific metrics, for example) can be in a single place.

[andygrove]

Added

[parthchandra]

Should this be Spark 4 now?
I'm assuming that this is enabled for only one version of Spark because it is experimental?

[andygrove]

With the latest changes, this feature is now only supported for Spark 4.x. If there is user demand, we can create a separate PR to support for 3.x

[comphead]

MAP is confusing IMO, so many things related to map in Spark

[comphead]

maybe COMET_PYARROW_SUPPORT or something like that?

[andygrove]

Thanks, updated docs to clarify.

[andygrove]

Item 7 (CometArrowPythonRunner dedupe across 4.1/4.2). Took a closer look. The shared surface ends up smaller than the diff suggests because Spark's BaseArrowPythonRunner itself changed shape between minors:

• 4.1 constructor: (funcs, evalType, argOffsets, schema, timeZoneId, largeVarTypes, workerConf, pythonMetrics, jobArtifactUUID, sessionUUID)
• 4.2 constructor: (funcs, evalType, argOffsets, schema, timeZoneId, largeVarTypes, pythonMetrics, jobArtifactUUID, sessionUUID) — workerConf dropped, overridden via runnerConf instead

A shared CometArrowPythonRunnerBase would have to live in spark-4.x/ to be visible to both minors, but BaseArrowPythonRunner isn't on the classpath there at a single signature, so the base would still need per-minor constructors and super(...) calls. Net saving: the two trait mix-ins (CometColumnarPythonInput + BasicPythonArrowOutput) factor out, which is roughly 10 lines for one more abstraction layer. I'd rather keep the per-minor pattern that mirrors Spark4xMapInBatchSupport already in the repo. Open to revisiting if 4.3 lands and the two minors converge on a stable signature.

Item 11 (fuzz harness). Borrowing the CometCodegenFuzzSuite shape into test_pyarrow_udf.py is the right next step but I'd like to do it in a follow-up so this PR doesn't grow further. Filed as #4384.

For this round I've added the targeted gaps in commit a520321:

• Decimal precision sweep {1, 9, 17, 18, 19, 28, 38} × scale {0, half, max} — covers the long-backed / FixedSizeBinary boundary at 18/19
• Null density sweep {0, 0.01, 0.5, 0.99, 1.0} for validity-buffer memcpy
• Multi-batch per partition (4000 rows, maxRecordsPerBatch=16) so the persistent IPC root is exercised across ~250 batches
• 50-column mixed schema for the flattened-tree address walk
• Mid-stream zero-row batch (validity sizing + setValueCount(0))
• Transforming array UDF (reverse each list) — catches symmetric encode/decode mistakes a passthrough would invert

Also filed #4383 for item 3 (drop the per-batch Comet→Spark buffer copy via direct ArrowRecordBatch assembly). That one supersedes the bulk-copy framing #4294 was built on; the doc rewrite in b69da2e explains why.

[mbutrovich]

Thanks for continuing to push on this, @andygrove. The latest revision resolves the correctness concern from last round. The output now wraps each vector via CometVector.getVector so the operator emits real CometVectors, extractColumnarChild handles a stacked CometMapInBatchExec directly, and the three downstream shapes I was worried about (chained UDFs, filter on UDF output, aggregate on UDF output) are all pinned by tests. The allocator framing rewrite and the removal of the per-batch VectorUnloader and the cometCodec reimplementation are real improvements, and the decimal, null-density, multi-batch, wide-schema, and nested-type coverage is strong. CI is green across every 4.x matrix plus the dedicated PyArrow workflow.

I left a set of inline comments I would like addressed before merge. None of them changes the approach.

[mbutrovich]

copyVector calls dst.allocateNew(...) on every child for every batch. In Arrow Java BaseFixedWidthVector.allocateNew(int) calls clear() first, which releases both buffers before reallocating (BaseFixedWidthVector.java:292 -> :296 -> releaseBuffer at :240-241), and BaseVariableWidthVector does the same. So the only thing actually persistent across batches is the StructVector container. Every leaf buffer is freed and re-malloc'd per batch. On a 50-column batch that is ~100 free/malloc pairs per batch.

Can we make the root genuinely persistent? Allocate once on the first batch, then for subsequent batches use reset() (BaseFixedWidthVector.java:225, "doesn't release any memory") plus a capacity check, reallocating only when the incoming batch does not fit. That is the setValueCount(0) + grow-on-demand path, and it removes the per-batch churn that the design is otherwise paying for. While you are in here, the test_map_in_arrow_multi_batch_per_partition docstring (test_pyarrow_udf.py, approx lines 437-442) says it exercises "the persistent destination IPC root over multiple batches," which is not true today since the child buffers reallocate. Once the reuse lands the docstring becomes accurate. If for some reason the reuse cannot work, the docstring needs correcting either way.

[andygrove]

Deferred the buffer reuse to #4383, which removes the copyVector bulk-copy path entirely (the per-batch allocateNew goes away with it), so a reset() + grow-on-demand rework here would be throwaway against code that is being deleted. I did correct the test_map_in_arrow_multi_batch_per_partition docstring in 5802e7c: it no longer claims leaf-buffer persistence (only the struct container is reused today) and now points at #4383.

[mbutrovich]

The read path increments pythonNumRowsReceived but nothing increments pythonDataReceived, while the writer does increment pythonDataSent. Stock BasicPythonArrowOutput populates pythonDataReceived as it reads frames. As written the accelerated path reports 0 bytes received, so the metric silently diverges from the fallback path depending on whether the rewrite fired. Please track the bytes read here so the metric matches the fallback. The metrics map in CometMapInBatchExec is where it surfaces to users.

[andygrove]

Fixed in 5802e7c. The reader now meters reader.bytesRead() around loadNextBatch() into pythonDataReceived, matching BasicPythonArrowOutput.

[mbutrovich]

The docstring says BaseFixedWidthVector handles precision <= 18 as "long-backed (8 bytes)" and precision >= 19 as "16-byte FixedSizeBinary." The Arrow DecimalVector that copyVector touches is always 16 bytes wide regardless of precision. The 8-byte long-backed form is Spark's UnsafeRow encoding, a layer this Arrow copy never sees. The sweep is good coverage. Please fix the rationale so it does not point at a buffer-width boundary that does not exist on the Arrow path.

[andygrove]

Fixed. Rewrote the rationale: the Arrow DecimalVector on this path is always 16 bytes wide, and the 8-byte long-backed form is Spark's UnsafeRow encoding the Arrow copy never sees. The docstring now says the sweep guards the precision/scale extremes and the 18/19 point where Spark changes its own representation, rather than an Arrow buffer-width boundary.

[mbutrovich]

The writer constructs ArrowStreamWriter(writeRoot, null, channel) with no compression codec, so spark.sql.execution.arrow.compression.codec is ignored on this path. The fallback path honors it. The Python ArrowStreamReader auto-detects compression from the IPC message metadata, so an uncompressed stream still reads correctly, but a user who sets the codec gets it silently dropped when the rewrite fires. Please resolve the codec from SQLConf and pass it through the VectorUnloader the writer uses, matching the fallback behavior.

[andygrove]

Deferred to #4383. One correctness note on the "fallback honors it" framing: on Spark 4.0 the fallback does not compress input to Python either. PythonArrowInput there builds new ArrowStreamWriter(root, null, dataOut) and writes via arrowWriter.writeBatch(), with no VectorUnloader/codec. The codec-on-input path only exists on 4.1+, where it is applied through VectorUnloader.getRecordBatch() + MessageSerializer. #4383 moves this operator onto that same unloader assembly, which is exactly where the codec drops in and where per-version behavior can be matched cleanly. Applying it in the current copyVector/ArrowStreamWriter writer would be throwaway against that change.

[mbutrovich]

The rule skips the rewrite entirely when useLargeVarTypes is on, because copyVector does a raw setBytes that cannot bridge Comet's 4-byte offsets to the destination's 8-byte offsets. This PR already taught CometPlainVector to read both offset widths, so the boundary is half-solved: the read side handles both, the write side gates the whole optimization off. Rather than gate, can the var-width arm of copyVector widen the offsets when the destination is a BaseLargeVariableWidthVector and the source is a BaseVariableWidthVector (read each 4-byte offset, write the 8-byte equivalent, copy the data buffer as-is)? That lets the rewrite always fire and removes a planner special case that would otherwise become dead code the moment the copy understands the wider offset.

[andygrove]

Deferred to #4383. Agreed the useLargeVarTypes gate becomes dead code once the copy understands 8-byte offsets, but the offset widening belongs in the reworked copy path (#4383 assembles the ArrowRecordBatch directly and is where the offset handling gets re-answered) rather than in the copyVector version that is being removed. The current gate is a safe fallback to vanilla Spark, not a correctness gap, so I left it in place for this PR.

[mbutrovich]

The manual while (b < srcBufs.size) reads more cleanly as srcBufs.asScala.zip(dstBufs.asScala).foreach { case (s, d) => d.setBytes(0, s, 0, s.readableBytes) }, with the buffer-count require moved to a size check beforehand. Drops a nesting level.

[andygrove]

Done, switched to srcBufs.asScala.zip(dstBufs.asScala).foreach { case (s, d) => d.setBytes(0, s, 0, s.readableBytes) } with the buffer-count require kept ahead of it.

[mbutrovich]

schema.head.dataType.asInstanceOf[StructType] is cast in two mutually exclusive arms. Hoist it to a single lazy val on the Writer so the cast lives in one place.

[andygrove]

Done. Hoisted to a single lazy val inputStructType on the Writer; both the empty-stream and first-batch arms use it.

[mbutrovich]

A null-named FFI child gets the positional placeholder _$idx. A real sibling literally named _0 would collide. It only applies to positional FFI fields with null names so it is unlikely, but a one-line comment noting the placeholder assumes no real field uses the _N form would save the next reader the head-scratch.

[andygrove]

Added a comment noting the positional _$idx placeholder is only applied to null-named FFI children and assumes no real sibling uses the _N form.

[mbutrovich]

One more pass, thanks for iterating on this, @andygrove!

[mbutrovich]

The input Arrow schema is built from the source vector's Field (getValueVector.getField), and the empty-input branch hardcodes "UTC". Vanilla Spark builds the schema it sends to Python with ArrowUtils.toArrowSchema(schema, timeZoneId), so a TimestampType reaches the worker as Timestamp(MICROSECOND, sessionLocalTimeZone). Comet normalizes timestamps to UTC internally, so under a non-UTC spark.sql.session.timeZone the worker would see a UTC label where vanilla Spark presents the session zone.

A passthrough UDF round-trips fine, which is why the current tests pass, but a tz-sensitive mapInPandas UDF or a mapInArrow UDF that reads field.type.tz could diverge from vanilla. test_map_in_arrow_date_and_timestamp only runs under the default session timezone. Could we add a non-UTC spark.sql.session.timeZone case with a tz-aware UDF to confirm parity, or document it as a known limitation?

[mbutrovich]

All three per-version runners take timeZoneId: String and largeVarTypes: Boolean, but neither is referenced. The base builds its schema from the source fields, and the rule already falls back to vanilla when useLargeVarTypes is on. Since these are brand-new classes, dropping the dead params (and the corresponding RunnerInputs fields) would keep the constructors honest. This is also the tell for comment #1: sessionLocalTimeZone is resolved on the driver and threaded all the way here, then ignored.

[mbutrovich]

The 4.1 and 4.2 runners differ from 4.0 in the config param handling (workerConf as override val vs pythonRunnerConf + override def) and in writeUDF arity, and those files landed under [skip ci]. The pytest workflow only runs a real Python worker against 4.0, so a wiring divergence in the 4.1/4.2 runners would still pass CI. A single 4.1 or 4.2 smoke run would close that gap. Not a blocker for an experimental drop if out of scope.

[mbutrovich]

Approved pending CI! Thanks for the initial PR and all the subsequent work on this, @andygrove !

[andygrove]

Merged. Thanks @mbutrovich!