andygrove opened a new issue, #5386: URL: https://github.com/apache/datafusion-comet/issues/5386
### 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 `CometVector`s. 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
