fix: keep a CoalescePartitionsExec required by a SinglePartition child - #170
Open
LiaCastaneda wants to merge 1 commit into
Open
fix: keep a CoalescePartitionsExec required by a SinglePartition child#170LiaCastaneda wants to merge 1 commit into
LiaCastaneda wants to merge 1 commit into
Conversation
apache#23948) - None filed; happy to open one if preferred. 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): ```sql 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. `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. 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. 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. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit 6636d6b)
LiaCastaneda
force-pushed
the
lia/cherry-pick-CoalescePartitionsExec-fix
branch
from
August 14, 2026 09:18
960df74 to
7b13c4c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-picks apache#23948