apache / apache/datafusion-comet

Extend native Arrow UDF path to grouped-aggregate, window, and applyInArrow Python operators

Open
#5,123 1 comment 0 reactions 0 assignees View on GitHub
area:udf enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.