fix: keep a CoalescePartitionsExec required by a SinglePartition child

[adriangb]

Which issue does this PR close?

• None filed; happy to open one if preferred.

Rationale for this change

A valid query can be planned into a physical plan that SanityCheckPlan then rejects:

SanityCheckPlan
caused by
Error during planning: Plan: ["HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, id@0)], projection=[id@0]",
  "  DataSourceExec: file_groups={4 groups: [...]}, projection=[id], file_type=parquet",
  "  RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
  "    CoalescePartitionsExec",
  "      ProjectionExec: expr=[first_value(t.id) ORDER BY [...]@1 as id]",
  "        AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[first_value(t.id) ORDER BY [...]]",
  "          RepartitionExec: partitioning=Hash([id@0], 8), input_partitions=4",
  "            AggregateExec: mode=Partial, gby=[id@1 as id], aggr=[first_value(t.id) ORDER BY [...]]",
  "              DataSourceExec: file_groups={4 groups: [...]}, projection=[ts, id], file_type=parquet"]
does not satisfy distribution requirements: SinglePartition. Child-0 output partitioning: UnknownPartitioning(4)

The HashJoinExec is in CollectLeft mode, which requires Distribution::SinglePartition on its build (left) child, but child 0 is a bare 4-partition DataSourceExec with no CoalescePartitionsExec above it.

Self-contained reproducer with datafusion-cli (the four COPY statements are what make the scan multi-partition):

set datafusion.execution.target_partitions = 8;
set datafusion.optimizer.repartition_file_scans = false;

create table src (id int, ts int) as values (1, 10), (2, 20), (3, 30);

copy (select * from src) to 'data/0.parquet' stored as parquet;
copy (select * from src) to 'data/1.parquet' stored as parquet;
copy (select * from src) to 'data/2.parquet' stored as parquet;
copy (select * from src) to 'data/3.parquet' stored as parquet;

create external table t stored as parquet location 'data/';

select a.id
from t a
left join (select distinct on (id) id, ts from t order by id, ts) f on a.id = f.id
order by a.id;

Setting datafusion.optimizer.repartition_sorts = false makes it plan fine, which points at the sort-parallelization phase.

EnsureRequirements does insert the coalesce for the SinglePartition requirement (enforce_distribution.rs, Distribution::SinglePartition => add_merge_on_top(...)). Its own phase 3a (parallelize_sorts) then takes it back out: remove_bottleneck_in_subplan removes a CoalescePartitionsExec found at children[0] positionally, without consulting the parent's distribution requirement for that child.

That parent is reached because update_coalesce_ctx_children marks a node as connected when any child qualifies. It correctly excludes a SinglePartition-requiring child from setting the flag, but the join's other child (UnspecifiedDistribution, connected to a coalesce below) sets it, so the traversal descends into the join and rewrites child 0 anyway. Nothing re-enforces distribution afterwards, so SanityCheckPlan is the first thing to notice. Note the surviving CoalescePartitionsExec on the probe side in the plan above: it is what propagated the flag, and it is untouched because the if returns without recursing into child 1.

The sibling helper on the phase 2b path already does consult the requirement (update_child_to_remove_unnecessary_sort / remove_corresponding_sort_from_sub_plan re-add a merge using the per-child child_distribution(child_idx)); only this path is missing it.

The same failure shows up with a build child that is already hash-partitioned on the join key (Child-0 output partitioning: Hash([k@0], 8)), which is what a JoinSelection input swap leaves behind — a CollectLeft join reported as join_type=Right with an embedded projection.

What changes are included in this PR?

remove_bottleneck_in_subplan now checks the parent's per-child distribution requirement before removing a coalesce, both for children[0] and when recursing into the other children.

The node parallelize_sorts is itself rewriting (the root of the call) is exempt, since the caller drops that node and rebuilds the sort cascade around the result — that is the rule's intended transformation, and gating it too would disable sort parallelization below a global sort. This is threaded through as an is_root flag on a private _impl function; the public entry point keeps its signature.

Are these changes tested?

Yes, at two levels:

• An end-to-end sqllogictest in datafusion/sqllogictest/test_files/joins.slt reproducing it from SQL (the reproducer above, with the data written by COPY inside the test). On main it fails with exactly the distribution error above.
• Two tests in datafusion/core/tests/physical_optimizer/ensure_requirements.rs covering both shapes of the build child (UnknownPartitioning(n) and Hash([k], n)), running the full EnsureRequirements rule and then SanityCheckPlan via the existing optimize_and_sanity_check helper, plus the idempotency check.

cargo test -p datafusion-physical-optimizer, cargo test -p datafusion --test core_integration -- physical_optimizer (530 tests) and the full sqllogictest suite (498 files) pass.

Are there any user-facing changes?

No API changes. Plans that were previously rejected by SanityCheckPlan now plan and execute; a coalesce that is genuinely required is retained where it was previously (incorrectly) removed.

[adriangb]

@zhuqi-lucas @2010YOUY01 would one of you mind taking a look at this change please?

[Copilot]

Pull request overview

Note

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Fixes a regression in EnsureRequirements sort-parallelization where a required CoalescePartitionsExec (needed to satisfy Distribution::SinglePartition for CollectLeft joins) could be removed, leading to SanityCheckPlan failures.

Changes:

• Prevent remove_bottleneck_in_subplan from removing CoalescePartitionsExec when the parent requires SinglePartition for that child.
• Add end-to-end sqllogictest coverage for the CollectLeft join + sort-parallelization reproducer.
• Add physical optimizer regression tests covering both UnknownPartitioning(n) and Hash([k], n) build-side shapes.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs	Adds per-child distribution checks to avoid removing required coalesces during sort parallelization.
datafusion/sqllogictest/test_files/joins.slt	Adds an end-to-end regression test reproducing the CollectLeft + SinglePartition failure.
datafusion/core/tests/physical_optimizer/ensure_requirements.rs	Adds optimizer-level regression tests and snapshots to ensure required coalesces are retained.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.85%. Comparing base (994fc81) to head (8fb3d30).
⚠️ Report is 38 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23948      +/-   ##
==========================================
+ Coverage   80.69%   80.85%   +0.16%     
==========================================
  Files        1095     1099       +4     
  Lines      372626   374345    +1719     
  Branches   372626   374345    +1719     
==========================================
+ Hits       300700   302695    +1995     
+ Misses      53978    53605     -373     
- Partials    17948    18045      +97     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[zhuqi-lucas]

I verified this locally and I'm afraid these two unit tests don't actually cover the regression: with the fix commit reverted (and cargo clean -p datafusion-physical-optimizer to force a rebuild), both tests still pass. I then instrumented the pre-fix children[0]-coalesce-removal branch with an eprintln! probe — it never fires for these mock plans, so they never reach the buggy path at all. The snapshots are identical pre/post fix.

(The sqllogictest is fine — on the reverted code it fails with exactly the reported does not satisfy distribution requirements: SinglePartition error, and passes with the fix.)

My guess at the cause: by the time parallelize_sorts runs, the distribution phase has already restructured the probe side, so the coalesce link that makes the traversal descend into the join in the real reproducer isn't there in this mock shape. Could you either make these construct the pre-phase-3a plan shape directly (so the removal branch demonstrably fires — an eprintln probe on main is a quick way to confirm), or drop them and let the sqllogictest carry the regression? As-is they'd stay green if the bug were reintroduced.

[adriangb]

Yep — confirmed, and your hypothesis was exactly right. Phase 2a drops the probe-side coalesce in the mock shape, so the join never gets marked as connected and the walk never descends into it.

Rebuilt in 8fb3d30: the tests now construct the pre-phase-3a shape directly (probe-side coalesce intact) and drive parallelize_sorts via transform_up, then check with SanityCheckPlan. Both fail with the reported SinglePartition error when the fix is reverted.

[zhuqi-lucas]

Question (non-blocking): should this also protect Distribution::HashPartitioned children? A single-partition input trivially satisfies a hash requirement, so a coalesce sitting under a hash-requiring child is load-bearing in the same way — removing it exposes whatever multi-partitioning is below, which may not be hashed on the right keys. I don't think EnsureRequirements itself ever places a coalesce there (hash requirements get a RepartitionExec), so it's probably unreachable via this rule's own output — but a plan built by other rules or by hand could hit it. Fine to leave as-is if you agree it's unreachable; a one-line comment saying so would help the next reader.

[adriangb]

Agreed on both counts — same hazard in principle, but ensure_distribution satisfies a hash requirement with a RepartitionExec, never a coalesce, so widening the match would be dead code. Left as-is with a comment saying so.

[zhuqi-lucas]

Worth a short comment that this is deliberately conservative: skipping the descent entirely also skips legitimate cleanups below the protected child (e.g. a redundant second coalesce underneath the load-bearing one). Correctness-first is the right call for a fix — just flagging that the gate could later be narrowed to "descend, but protect only the top-level coalesce" if anyone cares about the missed optimization.

[adriangb]

Yes, that's the intended trade — the missed cleanup is a redundant coalesce under a load-bearing one, which I'd rather leave than risk narrowing wrong in a fix. Added a comment recording that, with "descend, but protect only the topmost coalesce" as the follow-up.