hareshkh opened a new pull request, #26961:
URL: https://github.com/apache/datafusion/pull/26961
Credits: @adriangb
## Which issue does this PR close?
- Related to #23948 #23805
## 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):
```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.
## 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
