apache / apache/datafusion-comet
Extend native Arrow UDF path to grouped-aggregate, window, and applyInArrow Python operators
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
### 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).
Contributor guide
Research direction
Start with the existing native path from #4234, the scalar work in #5386, and the PyArrow UDF user guide. Trace the listed operators—starting with FlatMapGroupsInPandasExec—and establish how grouped and window operators should preserve Arrow batches. Done means the in-scope operators use the native path in the proposed order while pickled UDFs and AttachDistributedSequenceExec remain excluded.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, scala
- Domain
- data-engineering, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100