PyArrow UDF Acceleration#

Comet can accelerate Python UDFs that use PyArrow-backed batch processing, such as mapInArrow and mapInPandas. These APIs are commonly used for ML inference, feature engineering, and data transformation workloads.

Background#

Spark’s mapInArrow and mapInPandas APIs allow users to apply Python functions that operate on Arrow RecordBatches or Pandas DataFrames. Under the hood, Spark communicates with the Python worker process using the Arrow IPC format.

Without Comet, the execution path for these UDFs involves unnecessary data conversions:

Comet reads data in Arrow columnar format (via CometScan)
Spark inserts a ColumnarToRow transition (converts Arrow to UnsafeRow)
The Python runner converts those rows back to Arrow to send to Python
Python executes the UDF on Arrow batches
Results are returned as Arrow and then converted back to rows

Steps 2 and 3 are redundant since the data starts and ends in Arrow format.

How Comet Optimizes This#

When enabled, Comet detects PythonMapInArrowExec / MapInArrowExec and MapInPandasExec operators in the physical plan and replaces them with CometMapInBatchExec, which:

Reads Arrow columnar batches directly from the upstream Comet operator
Feeds them to the Python runner without the expensive UnsafeProjection copy
Keeps the Python output in columnar format for downstream operators

This eliminates the ColumnarToRow transition and the output row conversion, reducing CPU overhead and memory allocations. The row-to-Arrow re-encoding that Spark’s ArrowPythonRunner performed on the input side is also gone: CometArrowPythonRunner consumes ColumnarBatch directly, so batches are written straight from Comet’s vectors into the IPC root. See Limitations for the copies that remain.

Plan flow#

Without Comet’s optimization:

PythonMapInArrow / MapInArrow / MapInPandas
+- ColumnarToRow         <- Arrow -> Row copy
   +- CometNativeExec    <- Arrow batch
      +- CometScan

With the optimization enabled:

CometMapInBatch          <- Arrow batch in/out, Python runner attached
+- CometNativeExec
   +- CometScan

Configuration#

The optimization is experimental and disabled by default. Enable it with:

spark.comet.exec.pyarrowUDF.enabled=true

The default is false while the feature stabilizes.

Relationship to Spark’s PySpark Arrow conversion conf#

spark.comet.exec.pyarrowUDF.enabled is not the same as PySpark’s spark.sql.execution.arrow.pyspark.enabled. That conf controls whether Spark uses Arrow when materializing a DataFrame to a Pandas DataFrame (toPandas()) or constructing one from Pandas. The Comet conf controls a planner rewrite for mapInArrow / mapInPandas, and only affects how Comet’s columnar batches feed the Python worker. Both confs can be set independently.

Supported APIs#

PySpark API	Spark Plan Node	Supported
`df.mapInArrow(func, schema)`	`PythonMapInArrowExec`	Yes
`df.mapInPandas(func, schema)`	`MapInPandasExec`	Yes
`@pandas_udf` (scalar)	`ArrowEvalPythonExec`	Not yet
`df.applyInPandas(func, schema)`	`FlatMapGroupsInPandasExec`	Not yet

Example#

import pyarrow as pa
from pyspark.sql import SparkSession, types as T

spark = SparkSession.builder \
    .config("spark.plugins", "org.apache.spark.CometPlugin") \
    .config("spark.comet.enabled", "true") \
    .config("spark.comet.exec.enabled", "true") \
    .config("spark.comet.exec.pyarrowUDF.enabled", "true") \
    .config("spark.memory.offHeap.enabled", "true") \
    .config("spark.memory.offHeap.size", "2g") \
    .getOrCreate()

df = spark.read.parquet("data.parquet")

def transform(batch: pa.RecordBatch) -> pa.RecordBatch:
    # Your transformation logic here
    table = batch.to_pandas()
    table["new_col"] = table["value"] * 2
    return pa.RecordBatch.from_pandas(table)

output_schema = T.StructType([
    T.StructField("value", T.DoubleType()),
    T.StructField("new_col", T.DoubleType()),
])

result = df.mapInArrow(transform, output_schema)

Verifying the Optimization#

Use explain() to verify that CometMapInBatch appears in your plan:

result.explain(mode="extended")

You should see:

CometMapInBatch ...
+- CometNativeExec ...
   +- CometScan ...

Instead of the unoptimized plan:

PythonMapInArrow ...
+- ColumnarToRow
   +- CometNativeExec ...
      +- CometScan ...

When AQE is enabled (the Spark default) and the query contains a shuffle, the optimization is applied during stage materialization. Calling explain() before running an action will show the unoptimized plan:

AdaptiveSparkPlan isFinalPlan=false
+- PythonMapInArrow ...
   +- CometExchange ...

To see the optimized plan, run an action first (for example result.collect() or result.cache(); result.count()) and then call explain(). The post-execution plan shows the materialized stages and includes CometMapInBatch if the optimization fired.

Barrier execution#

mapInArrow(..., barrier=True) and mapInPandas(..., barrier=True) are honored: the optimized operator propagates isBarrier through RDD.barrier(), so all tasks are gang-scheduled and BarrierTaskContext.barrier() works inside the UDF the same way it does on the unoptimized path.

Limitations#

The optimization currently applies only to mapInArrow and mapInPandas. Scalar pandas UDFs (@pandas_udf) and grouped operations (applyInPandas) are not yet supported.
The optimization requires Arrow data on the input side. If a shuffle sits between the upstream Comet operator and the Python UDF, use Comet’s columnar shuffle for the optimization to apply. Both the jvm and native shuffle modes can feed CometMapInBatch. Set spark.shuffle.manager to org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager and enable spark.comet.shuffle.enabled=true at session startup. With a vanilla Spark Exchange in the plan the data leaves the shuffle as rows and the optimization cannot fire.
Spark 4.0 or newer is required. On Spark 3.4 and 3.5 the optimization is a no-op even when enabled; vanilla PythonMapInArrowExec / MapInPandasExec handle the operation. The Spark 3.5 PythonArrowInput trait has a different contract than 4.x and a separate implementation has not been written. Track 3.5 support as a future follow-on if there is user demand.
Timestamps are presented to the UDF with a UTC time zone rather than the session time zone. Comet normalizes timestamps to UTC internally, and the accelerated path builds the Arrow schema it sends to Python from Comet’s own vectors, so a TimestampType column reaches the worker labelled Timestamp(MICROSECOND, "UTC"). Vanilla Spark instead labels it with spark.sql.session.timeZone. The stored value is the same absolute instant either way, so a passthrough or value-based UDF round-trips identically. The difference is only observable to a UDF that reads the Arrow field’s time zone or localizes to wall-clock time (for example a mapInPandas UDF that strips the tz and treats the value as naive local time): under a non-UTC 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 supplies large_string and large_binary input columns with 8-byte offsets. Native Comet vectors use 4-byte offsets, and direct serialization advertises their matching string and binary types. This produces a valid IPC stream, but does not preserve the input types requested by the configuration. EliminateRedundantTransitions therefore skips the rewrite and vanilla Spark handles the operation. Comet can read large_string and large_binary columns returned by a Python worker; that output support does not widen the input vectors.
Comet applies spark.sql.execution.arrow.maxRecordsPerBatch to every input batch, including batches with only plain columns. Before decoding dictionary-encoded shuffle columns, Comet also compares their estimated decoded size with spark.sql.execution.arrow.maxBytesPerBatch. When either threshold requires splitting, every column is sliced at the same row boundaries. Temporary slices and decoded dictionary vectors are released after each synchronous write. Comet returns control to Spark after each slice so Spark can drain its Python transport buffer; small slices may share that buffer until Spark reaches its buffering threshold. The source batch remains alive until its last slice has been written.
The byte estimate covers only the logical buffers of decoded dictionary columns: values, offsets, and validity bits. It excludes plain columns and is a soft limit: the row that crosses the threshold stays in the batch, and a single oversized row remains intact. A separate guard prevents combining rows whose estimated decoded dictionary size exceeds Arrow’s signed 32-bit limit (2 GiB minus 1 byte). This guard cannot split an individually oversized row and does not guarantee that Arrow allocations stay below that limit. Arrow rounds buffer capacities up, so an allocation can approach twice its logical size; existing input buffers and other overhead also consume memory. maxBytesPerBatch is therefore not a ceiling on actual memory use.
Dictionary-encoded values nested inside a struct, list, or map are not supported on the optimized input path. Comet rejects them with an error naming the field path. Comet’s current shuffle does not produce these nested dictionaries.
Comet writes input Arrow IPC record batches directly from plain vector buffers. For an unsplit plain batch, the only additional Arrow buffer is the validity bitmap for the non-null struct that wraps the input columns. Slicing may allocate offset or validity buffers. 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. Borrowed buffers are not transferred between Arrow allocators or given new ownership.