Let custom operators declare that they tolerate a coalesced partition count

Open
#2,728 0 comments 0 reactions 1 assignee View on GitHub

Maintainers usually reply within 1 day

@andygrove is already working on this.

Since Nov 1, 2026.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
12/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
rust, wasm
Domain
backend

Research direction

The coalescing decision lives in CoalescePartitionsRule and the hand-kept tolerates_partition_count_change operator list. Read the Ballista extension trait and the SchedulerConfig resolver that apache/datafusion-ballista#2721 added for partition slicing, since the issue asks for a similar hook. Done would be a design settled with maintainers first, then a trait method that lets a stage opt in without Ballista depending on Comet.

Written by the indexing model from the issue text.

Description

enhancement

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

The adaptive planner's CoalescePartitionsRule coalesces a shuffle's partitions only when every operator between the stage root and the shuffle is on a hand-kept list of operators that produce the same rows over any partitioning of their input (tolerates_partition_count_change). An extension operator from a client that submitted its own physical plan is assumed to depend on the count, which is the safe default. Comet Native Executors (apache/datafusion-comet#7288) runs every stage as a CometFragmentExec, so its queries never get partition coalescing.

Describe the solution you'd like

Let an ExecutionPlan declare this itself, for example through a method on a Ballista extension trait, or a resolver in SchedulerConfig like the one apache/datafusion-ballista#2721 adds for partition slicing. A Comet stage that reads only hash shuffles could then declare it.

Describe alternatives you've considered

Keeping the whitelist and adding Comet's operator to it would make Ballista depend on Comet.

Additional context

Comet's CometFragmentExec also carries a partition descriptor child that would have to coalesce along with its shuffle inputs.

Dominant language
Rust
Stars
2.2k
Forks
337
Avg merge
2d 8h
Merged PRs (30d)
127

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from apache/datafusion-ballista

All issues in apache/datafusion-ballista

Similar issues

More Rust issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.