What is the problem the feature request solves?
Comet's opt-in native Arrow UDF path (spark.comet.exec.pyarrowUDF.enabled) covers mapInArrow and mapInPandas (#4234). Scalar Python UDFs still fall back entirely — nothing under spark/src/main matches ArrowEvalPythonExec — so a Comet scan feeding a scalar UDF pays a ColumnarToRow transition and Spark then re-encodes those rows back into Arrow to reach the Python worker.
One operator covers three user-facing UDF families, because ExtractPythonUDFs routes all of them to ArrowEvalPythonExec:
- regular
udf() when spark.sql.execution.pythonUDF.arrow.enabled=true (eval type SQL_ARROW_BATCHED_UDF, available since Spark 3.4, still off by default in 4.1)
@pandas_udf scalar (SQL_SCALAR_PANDAS_UDF)
- Spark 4.1's
@arrow_udf (SQL_SCALAR_ARROW_UDF)
The first of those is what makes ordinary PySpark UDF code benefit from Comet. With the Spark conf enabled a plain udf() already speaks the same Arrow IPC protocol as mapInArrow, so this is the operator that turns "Comet accelerates PyArrow UDFs" into "Comet accelerates Python UDFs".
The win should be larger than it was for mapInArrow, because EvalPythonEvaluatorFactory is considerably more row-bound than MapInBatchExec was. Per row, the vanilla path does:
HybridRowQueue.add(UnsafeRow), which copies the whole input row and can spill to disk
- a
MutableProjection to materialize the UDF arguments
- on the way back,
queue.remove(), a JoinedRow, and an UnsafeProjection over the full output schema
batch.rowIterator in ArrowEvalPythonEvaluatorFactory, shredding Python's Arrow output back into rows
all of it behind the ColumnarToRow that a Comet child forces.
Describe the potential solution
Add CometArrowEvalPythonExec, matched in EliminateRedundantTransitions the way EligibleMapInBatch matches today. Rather than a row queue, retain the input ColumnarBatch, send only the UDF-argument columns to the Python worker, and append the returned vectors to the input batch's vectors to form the output batch. That removes the row queue and its spill path, the joined row and both projections, and keeps the result columnar for downstream native operators.
CometArrowPythonRunnerBase should be reusable close to as-is: the input side already copies ColumnarBatch vectors into a struct root, and the output side already decodes the worker's IPC into CometVectors.
Things to scope deliberately:
- Batch alignment. Appending Python output onto a retained input batch assumes a 1:1 batch relationship. That holds for the non-iterator eval types.
SQL_SCALAR_PANDAS_ITER_UDF and SQL_SCALAR_ARROW_ITER_UDF only guarantee equal total row counts, so they need row-count-based re-alignment, or should be excluded from the first version.
- Non-attribute arguments. Spark does not project UDF arguments below the operator;
allInputs in EvalPythonEvaluatorFactory can hold arbitrary expressions. Either evaluate them in a native projection beneath the operator, or restrict the first version to bare attribute references.
- Several UDFs in one operator, including the result-struct shape Spark uses when
udfs.length > 1.
- The existing pyarrowUDF limitations carry over: Arrow is required on the input side (so native shuffle when an exchange intervenes),
spark.sql.execution.arrow.useLargeVarTypes is unsupported, and the UTC timestamp labelling difference applies.
Comet should not flip spark.sql.execution.pythonUDF.arrow.enabled on the user's behalf, since that conf changes Spark's own type coercion and error semantics rather than just the transport. A fallback reason on BatchEvalPythonExec pointing users at it would be a good companion change.
Additional context
Follows #4234. Pickled SQL_BATCHED_UDF execution stays out of scope: per-row boxing, pickle and interpreter cost dominate, and the unpickled Java-object results would have to be converted into Arrow vectors to keep batches all-Comet. #5123 covers the remaining Arrow-based Python operators (grouped-aggregate, window, applyInArrow, cogroup, Arrow UDTF) and builds on the pattern established here. #4384 tracks a fuzz harness for the vector-copy path.
What is the problem the feature request solves?
Comet's opt-in native Arrow UDF path (
spark.comet.exec.pyarrowUDF.enabled) coversmapInArrowandmapInPandas(#4234). Scalar Python UDFs still fall back entirely — nothing underspark/src/mainmatchesArrowEvalPythonExec— so a Comet scan feeding a scalar UDF pays aColumnarToRowtransition and Spark then re-encodes those rows back into Arrow to reach the Python worker.One operator covers three user-facing UDF families, because
ExtractPythonUDFsroutes all of them toArrowEvalPythonExec:udf()whenspark.sql.execution.pythonUDF.arrow.enabled=true(eval typeSQL_ARROW_BATCHED_UDF, available since Spark 3.4, still off by default in 4.1)@pandas_udfscalar (SQL_SCALAR_PANDAS_UDF)@arrow_udf(SQL_SCALAR_ARROW_UDF)The first of those is what makes ordinary PySpark UDF code benefit from Comet. With the Spark conf enabled a plain
udf()already speaks the same Arrow IPC protocol asmapInArrow, so this is the operator that turns "Comet accelerates PyArrow UDFs" into "Comet accelerates Python UDFs".The win should be larger than it was for
mapInArrow, becauseEvalPythonEvaluatorFactoryis considerably more row-bound thanMapInBatchExecwas. Per row, the vanilla path does:HybridRowQueue.add(UnsafeRow), which copies the whole input row and can spill to diskMutableProjectionto materialize the UDF argumentsqueue.remove(), aJoinedRow, and anUnsafeProjectionover the full output schemabatch.rowIteratorinArrowEvalPythonEvaluatorFactory, shredding Python's Arrow output back into rowsall of it behind the
ColumnarToRowthat a Comet child forces.Describe the potential solution
Add
CometArrowEvalPythonExec, matched inEliminateRedundantTransitionsthe wayEligibleMapInBatchmatches today. Rather than a row queue, retain the inputColumnarBatch, send only the UDF-argument columns to the Python worker, and append the returned vectors to the input batch's vectors to form the output batch. That removes the row queue and its spill path, the joined row and both projections, and keeps the result columnar for downstream native operators.CometArrowPythonRunnerBaseshould be reusable close to as-is: the input side already copiesColumnarBatchvectors into a struct root, and the output side already decodes the worker's IPC intoCometVectors.Things to scope deliberately:
SQL_SCALAR_PANDAS_ITER_UDFandSQL_SCALAR_ARROW_ITER_UDFonly guarantee equal total row counts, so they need row-count-based re-alignment, or should be excluded from the first version.allInputsinEvalPythonEvaluatorFactorycan hold arbitrary expressions. Either evaluate them in a native projection beneath the operator, or restrict the first version to bare attribute references.udfs.length > 1.spark.sql.execution.arrow.useLargeVarTypesis unsupported, and the UTC timestamp labelling difference applies.Comet should not flip
spark.sql.execution.pythonUDF.arrow.enabledon the user's behalf, since that conf changes Spark's own type coercion and error semantics rather than just the transport. A fallback reason onBatchEvalPythonExecpointing users at it would be a good companion change.Additional context
Follows #4234. Pickled
SQL_BATCHED_UDFexecution stays out of scope: per-row boxing, pickle and interpreter cost dominate, and the unpickled Java-object results would have to be converted into Arrow vectors to keep batches all-Comet. #5123 covers the remaining Arrow-based Python operators (grouped-aggregate, window,applyInArrow, cogroup, Arrow UDTF) and builds on the pattern established here. #4384 tracks a fuzz harness for the vector-copy path.