What is the problem the feature request solves?
Comet has an opt-in native path (spark.comet.exec.pyarrowUDF.enabled) that keeps data in Arrow format across mapInArrow / mapInPandas, avoiding Spark's Arrow to row to Arrow round trip. That is all #4234 landed — scalar UDFs (ArrowEvalPythonExec) are tracked separately in #5386, and applyInPandas is still outstanding too.
The remaining Arrow-based Python operators still fall back, even though they exchange Arrow batches with the Python worker and the same optimization applies in principle:
FlatMapGroupsInPandasExec (applyInPandas)AggregateInPandasExec / ArrowAggregatePythonExec (grouped-aggregate pandas UDFs, renamed in Spark 4.1)WindowInPandasExec / ArrowWindowPythonExec (window pandas UDFs, renamed in Spark 4.1)FlatMapGroupsInArrowExec (applyInArrow, Spark 4.0+)FlatMapCoGroupsInPandasExec / FlatMapCoGroupsInArrowExec (cogroup(...).applyInPandas / applyInArrow)ArrowEvalPythonUDTFExec (Arrow-based Python UDTFs, Spark 3.5+)
For Python-heavy analytics workloads these operators sit in the middle of otherwise native stages, so each one both pays the row round trip and splits the native stage.
Describe the potential solution
Extend the pyarrowUDF native path to these operators once the scalar work in #5386 lands, in roughly the order listed. applyInPandas comes first because the grouped variants need the grouping and partitioning handled natively and it establishes that pattern; grouped-aggregate and window pandas UDFs follow, since they are the most common in analytics code.
Pickled (non-Arrow) UDF operators are out of scope since there is no columnar boundary to preserve. AttachDistributedSequenceExec is also out of scope here.
Additional context
Related: #4234, #5386 (scalar UDFs), the PyArrow UDF user guide, #4384 (fuzz harness for the pyarrow UDF vector-copy path).
What is the problem the feature request solves?
Comet has an opt-in native path (
spark.comet.exec.pyarrowUDF.enabled) that keeps data in Arrow format acrossmapInArrow/mapInPandas, avoiding Spark's Arrow to row to Arrow round trip. That is all #4234 landed — scalar UDFs (ArrowEvalPythonExec) are tracked separately in #5386, andapplyInPandasis still outstanding too.The remaining Arrow-based Python operators still fall back, even though they exchange Arrow batches with the Python worker and the same optimization applies in principle:
FlatMapGroupsInPandasExec(applyInPandas)AggregateInPandasExec/ArrowAggregatePythonExec(grouped-aggregate pandas UDFs, renamed in Spark 4.1)WindowInPandasExec/ArrowWindowPythonExec(window pandas UDFs, renamed in Spark 4.1)FlatMapGroupsInArrowExec(applyInArrow, Spark 4.0+)FlatMapCoGroupsInPandasExec/FlatMapCoGroupsInArrowExec(cogroup(...).applyInPandas/applyInArrow)ArrowEvalPythonUDTFExec(Arrow-based Python UDTFs, Spark 3.5+)For Python-heavy analytics workloads these operators sit in the middle of otherwise native stages, so each one both pays the row round trip and splits the native stage.
Describe the potential solution
Extend the pyarrowUDF native path to these operators once the scalar work in #5386 lands, in roughly the order listed.
applyInPandascomes first because the grouped variants need the grouping and partitioning handled natively and it establishes that pattern; grouped-aggregate and window pandas UDFs follow, since they are the most common in analytics code.Pickled (non-Arrow) UDF operators are out of scope since there is no columnar boundary to preserve.
AttachDistributedSequenceExecis also out of scope here.Additional context
Related: #4234, #5386 (scalar UDFs), the PyArrow UDF user guide, #4384 (fuzz harness for the pyarrow UDF vector-copy path).