Repository navigation
Commit 6636d6b
fix: keep a CoalescePartitionsExec required by a SinglePartition child (#23948)
## 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):
```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.
---------
Co-authored-by: Claude Opus 5 <[email protected]>1 parent a589b4b commit 6636d6b
3 files changed
Lines changed: 265 additions & 5 deletions
File tree
- datafusion
- core/tests/physical_optimizer
- physical-optimizer/src/ensure_requirements/enforce_sorting
- sqllogictest/test_files
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
21 | 21 | | |
22 | 22 | | |
23 | 23 | | |
| 24 | + | |
| 25 | + | |
24 | 26 | | |
| 27 | + | |
25 | 28 | | |
26 | 29 | | |
| 30 | + | |
| 31 | + | |
| 32 | + | |
27 | 33 | | |
28 | 34 | | |
29 | 35 | | |
| |||
65 | 71 | | |
66 | 72 | | |
67 | 73 | | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
| 79 | + | |
| 80 | + | |
| 81 | + | |
| 82 | + | |
| 83 | + | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
68 | 87 | | |
69 | 88 | | |
70 | 89 | | |
| |||
81 | 100 | | |
82 | 101 | | |
83 | 102 | | |
84 | | - | |
| 103 | + | |
85 | 104 | | |
86 | 105 | | |
87 | 106 | | |
| |||
1252 | 1271 | | |
1253 | 1272 | | |
1254 | 1273 | | |
| 1274 | + | |
| 1275 | + | |
| 1276 | + | |
| 1277 | + | |
| 1278 | + | |
| 1279 | + | |
| 1280 | + | |
| 1281 | + | |
| 1282 | + | |
| 1283 | + | |
| 1284 | + | |
| 1285 | + | |
| 1286 | + | |
| 1287 | + | |
| 1288 | + | |
| 1289 | + | |
| 1290 | + | |
| 1291 | + | |
| 1292 | + | |
| 1293 | + | |
| 1294 | + | |
| 1295 | + | |
| 1296 | + | |
| 1297 | + | |
| 1298 | + | |
| 1299 | + | |
| 1300 | + | |
| 1301 | + | |
| 1302 | + | |
| 1303 | + | |
| 1304 | + | |
| 1305 | + | |
| 1306 | + | |
| 1307 | + | |
| 1308 | + | |
| 1309 | + | |
| 1310 | + | |
| 1311 | + | |
| 1312 | + | |
| 1313 | + | |
| 1314 | + | |
| 1315 | + | |
| 1316 | + | |
| 1317 | + | |
| 1318 | + | |
| 1319 | + | |
| 1320 | + | |
| 1321 | + | |
| 1322 | + | |
| 1323 | + | |
| 1324 | + | |
| 1325 | + | |
| 1326 | + | |
| 1327 | + | |
| 1328 | + | |
| 1329 | + | |
| 1330 | + | |
| 1331 | + | |
| 1332 | + | |
| 1333 | + | |
| 1334 | + | |
| 1335 | + | |
| 1336 | + | |
| 1337 | + | |
| 1338 | + | |
| 1339 | + | |
| 1340 | + | |
| 1341 | + | |
| 1342 | + | |
| 1343 | + | |
| 1344 | + | |
| 1345 | + | |
| 1346 | + | |
| 1347 | + | |
| 1348 | + | |
| 1349 | + | |
| 1350 | + | |
| 1351 | + | |
| 1352 | + | |
| 1353 | + | |
| 1354 | + | |
| 1355 | + | |
| 1356 | + | |
| 1357 | + | |
| 1358 | + | |
| 1359 | + | |
| 1360 | + | |
| 1361 | + | |
| 1362 | + | |
| 1363 | + | |
| 1364 | + | |
| 1365 | + | |
| 1366 | + | |
| 1367 | + | |
| 1368 | + | |
| 1369 | + | |
| 1370 | + | |
| 1371 | + | |
| 1372 | + | |
| 1373 | + | |
| 1374 | + | |
| 1375 | + | |
| 1376 | + | |
| 1377 | + | |
| 1378 | + | |
| 1379 | + | |
| 1380 | + | |
| 1381 | + | |
| 1382 | + | |
| 1383 | + | |
| 1384 | + | |
| 1385 | + | |
| 1386 | + | |
| 1387 | + | |
| 1388 | + | |
| 1389 | + | |
| 1390 | + | |
| 1391 | + | |
| 1392 | + | |
| 1393 | + | |
| 1394 | + | |
| 1395 | + | |
| 1396 | + | |
| 1397 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
609 | 609 | | |
610 | 610 | | |
611 | 611 | | |
| 612 | + | |
| 613 | + | |
| 614 | + | |
| 615 | + | |
| 616 | + | |
| 617 | + | |
| 618 | + | |
| 619 | + | |
| 620 | + | |
| 621 | + | |
612 | 622 | | |
| 623 | + | |
613 | 624 | | |
614 | 625 | | |
| 626 | + | |
| 627 | + | |
| 628 | + | |
| 629 | + | |
| 630 | + | |
| 631 | + | |
| 632 | + | |
| 633 | + | |
| 634 | + | |
| 635 | + | |
| 636 | + | |
| 637 | + | |
| 638 | + | |
| 639 | + | |
| 640 | + | |
| 641 | + | |
| 642 | + | |
| 643 | + | |
| 644 | + | |
| 645 | + | |
| 646 | + | |
| 647 | + | |
| 648 | + | |
| 649 | + | |
| 650 | + | |
| 651 | + | |
| 652 | + | |
615 | 653 | | |
616 | | - | |
| 654 | + | |
617 | 655 | | |
618 | 656 | | |
619 | 657 | | |
| |||
627 | 665 | | |
628 | 666 | | |
629 | 667 | | |
630 | | - | |
631 | | - | |
632 | | - | |
| 668 | + | |
| 669 | + | |
| 670 | + | |
| 671 | + | |
| 672 | + | |
| 673 | + | |
| 674 | + | |
| 675 | + | |
633 | 676 | | |
634 | 677 | | |
635 | 678 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
5608 | 5608 | | |
5609 | 5609 | | |
5610 | 5610 | | |
| 5611 | + | |
| 5612 | + | |
| 5613 | + | |
| 5614 | + | |
| 5615 | + | |
| 5616 | + | |
| 5617 | + | |
| 5618 | + | |
| 5619 | + | |
| 5620 | + | |
| 5621 | + | |
| 5622 | + | |
| 5623 | + | |
| 5624 | + | |
| 5625 | + | |
| 5626 | + | |
| 5627 | + | |
| 5628 | + | |
| 5629 | + | |
| 5630 | + | |
| 5631 | + | |
| 5632 | + | |
| 5633 | + | |
| 5634 | + | |
| 5635 | + | |
| 5636 | + | |
| 5637 | + | |
| 5638 | + | |
| 5639 | + | |
| 5640 | + | |
| 5641 | + | |
| 5642 | + | |
| 5643 | + | |
| 5644 | + | |
| 5645 | + | |
| 5646 | + | |
| 5647 | + | |
| 5648 | + | |
| 5649 | + | |
| 5650 | + | |
| 5651 | + | |
| 5652 | + | |
| 5653 | + | |
| 5654 | + | |
| 5655 | + | |
| 5656 | + | |
| 5657 | + | |
| 5658 | + | |
| 5659 | + | |
| 5660 | + | |
| 5661 | + | |
| 5662 | + | |
| 5663 | + | |
| 5664 | + | |
| 5665 | + | |
| 5666 | + | |
| 5667 | + | |
| 5668 | + | |
| 5669 | + | |
| 5670 | + | |
| 5671 | + | |
| 5672 | + | |
| 5673 | + | |
| 5674 | + | |
| 5675 | + | |
| 5676 | + | |
| 5677 | + | |
| 5678 | + | |
| 5679 | + | |
| 5680 | + | |
| 5681 | + | |
| 5682 | + | |
| 5683 | + | |
| 5684 | + | |
0 commit comments