Skip to content

Commit 6636d6b

Browse files
adriangbclaude
andauthored
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/ensure_requirements.rs‎

Lines changed: 144 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,15 @@
2121
//! so the tests live alongside the rest of the `physical_optimizer/` integration
2222
//! suite and can use real `ExecutionPlan`s where convenient.
2323
24+
use insta::assert_snapshot;
25+
2426
use datafusion_common::config::ConfigOptions;
27+
use datafusion_common::tree_node::{TransformedResult, TreeNode};
2528
use datafusion_physical_optimizer::PhysicalOptimizerRule;
2629
use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
30+
use datafusion_physical_optimizer::ensure_requirements::enforce_sorting::{
31+
PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
32+
};
2733

2834
use std::sync::Arc;
2935

@@ -65,6 +71,19 @@ struct MockMultiPartitionExec {
6571

6672
impl MockMultiPartitionExec {
6773
fn new(partition_count: usize) -> Self {
74+
Self::with_partitioning(Partitioning::UnknownPartitioning(partition_count))
75+
}
76+
77+
/// A source that is already partitioned on `a`, as an aggregate or a partitioned
78+
/// join below the node under test would be.
79+
fn hash_partitioned_on_a(partition_count: usize) -> Self {
80+
Self::with_partitioning(Partitioning::Hash(
81+
vec![Arc::new(Column::new("a", 0))],
82+
partition_count,
83+
))
84+
}
85+
86+
fn with_partitioning(partitioning: Partitioning) -> Self {
6887
let schema = Arc::new(Schema::new(vec![
6988
Field::new("a", DataType::Int64, false),
7089
Field::new("b", DataType::Int64, false),
@@ -81,7 +100,7 @@ impl MockMultiPartitionExec {
81100
}
82101
let properties = PlanProperties::new(
83102
eq,
84-
Partitioning::UnknownPartitioning(partition_count),
103+
partitioning,
85104
EmissionType::Incremental,
86105
Boundedness::Bounded,
87106
);
@@ -1252,3 +1271,127 @@ fn test_idempotent_union_projection_sort() {
12521271

12531272
assert_idempotent(plan);
12541273
}
1274+
1275+
/// Builds the plan shape that phase 3a (`parallelize_sorts`) sees in the reproducer,
1276+
/// i.e. the output of the distribution + sorting phases, not a freshly planned tree:
1277+
///
1278+
/// ```text
1279+
/// CoalescePartitionsExec <- the node `parallelize_sorts` rewrites
1280+
/// HashJoinExec: mode=CollectLeft
1281+
/// CoalescePartitionsExec <- satisfies `SinglePartition` on the build side
1282+
/// <build>
1283+
/// RepartitionExec: RoundRobinBatch
1284+
/// CoalescePartitionsExec <- links the join into the coalesce cascade
1285+
/// MockMultiPartitionExec
1286+
/// ```
1287+
///
1288+
/// Both coalesces below the join matter. The probe-side one is what makes
1289+
/// `update_coalesce_ctx_children` mark the join as connected — it only skips children that
1290+
/// require `SinglePartition`, and the probe side does not — so the walk descends into the
1291+
/// join. The build-side one is the one that must survive.
1292+
fn collect_left_plan_before_parallelize_sorts(
1293+
build: Arc<dyn ExecutionPlan>,
1294+
join_type: JoinType,
1295+
) -> Result<Arc<dyn ExecutionPlan>> {
1296+
let build: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(build));
1297+
let probe: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new(
1298+
Arc::new(CoalescePartitionsExec::new(Arc::new(
1299+
MockMultiPartitionExec::new(4),
1300+
))),
1301+
Partitioning::RoundRobinBatch(TEST_TARGET_PARTITIONS),
1302+
)?);
1303+
1304+
let on = vec![(
1305+
Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
1306+
Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
1307+
)];
1308+
let join: Arc<dyn ExecutionPlan> = Arc::new(HashJoinExec::try_new(
1309+
build,
1310+
probe,
1311+
on,
1312+
None,
1313+
&join_type,
1314+
None,
1315+
PartitionMode::CollectLeft,
1316+
NullEquality::NullEqualsNothing,
1317+
false,
1318+
)?);
1319+
1320+
Ok(Arc::new(CoalescePartitionsExec::new(join)))
1321+
}
1322+
1323+
/// Runs phase 3a of `EnsureRequirements` (`parallelize_sorts`) on its own, the same way
1324+
/// the rule drives it, and checks the result with `SanityCheckPlan`.
1325+
///
1326+
/// The phase is driven directly rather than through `EnsureRequirements::optimize` because
1327+
/// the earlier phases would rebuild the plan shape above into something that never reaches
1328+
/// the code path under test.
1329+
fn parallelize_sorts_and_sanity_check(
1330+
plan: Arc<dyn ExecutionPlan>,
1331+
) -> Result<Arc<dyn ExecutionPlan>> {
1332+
let ctx = PlanWithCorrespondingCoalescePartitions::new_default(plan);
1333+
let rewritten = ctx.transform_up(parallelize_sorts).data()?.plan;
1334+
SanityCheckPlan::new().optimize(Arc::clone(&rewritten), &test_config())?;
1335+
Ok(rewritten)
1336+
}
1337+
1338+
/// A `CollectLeft` `HashJoinExec` requires `Distribution::SinglePartition` on its build
1339+
/// (left) child, so the distribution phase puts a `CoalescePartitionsExec` on top of a
1340+
/// multi-partition build side. The sort-parallelization phase must not take that coalesce
1341+
/// back out again.
1342+
///
1343+
/// It used to, because `remove_bottleneck_in_subplan` removed a coalesce found at
1344+
/// `children[0]` positionally, without consulting the parent's distribution requirement for
1345+
/// that child. The result was a build side left multi-partition with nothing to re-enforce
1346+
/// distribution afterwards, which `SanityCheckPlan` rejected with "does not satisfy
1347+
/// distribution requirements: SinglePartition".
1348+
#[test]
1349+
fn test_collect_left_join_keeps_build_side_coalesce() -> Result<()> {
1350+
let plan = collect_left_plan_before_parallelize_sorts(
1351+
Arc::new(MockMultiPartitionExec::new(4)),
1352+
JoinType::Left,
1353+
)?;
1354+
1355+
let rewritten = parallelize_sorts_and_sanity_check(plan)?;
1356+
1357+
// The build-side coalesce is retained; the probe-side one is still removed, which is
1358+
// the parallelization this phase exists for.
1359+
assert_snapshot!(plan_string(&rewritten), @r"
1360+
CoalescePartitionsExec
1361+
HashJoinExec: mode=CollectLeft, join_type=Left, on=[(a@0, a@0)]
1362+
CoalescePartitionsExec
1363+
MockMultiPartitionExec
1364+
RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4
1365+
MockMultiPartitionExec
1366+
");
1367+
1368+
Ok(())
1369+
}
1370+
1371+
/// The same removal, with a build side that is already hash-partitioned on the join key
1372+
/// rather than `UnknownPartitioning`. This is the shape a `JoinSelection` input swap leaves
1373+
/// behind (a `CollectLeft` join reported as `join_type=Right`) when the build subtree is the
1374+
/// output of an aggregate or a partitioned join: the build side satisfies the join's *hash*
1375+
/// requirement but still not `SinglePartition`, so the coalesce is just as load-bearing.
1376+
#[test]
1377+
fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result<()> {
1378+
let plan = collect_left_plan_before_parallelize_sorts(
1379+
Arc::new(MockMultiPartitionExec::hash_partitioned_on_a(
1380+
TEST_TARGET_PARTITIONS,
1381+
)),
1382+
JoinType::Right,
1383+
)?;
1384+
1385+
let rewritten = parallelize_sorts_and_sanity_check(plan)?;
1386+
1387+
assert_snapshot!(plan_string(&rewritten), @r"
1388+
CoalescePartitionsExec
1389+
HashJoinExec: mode=CollectLeft, join_type=Right, on=[(a@0, a@0)]
1390+
CoalescePartitionsExec
1391+
MockMultiPartitionExec
1392+
RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4
1393+
MockMultiPartitionExec
1394+
");
1395+
1396+
Ok(())
1397+
}

‎datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs‎

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -609,11 +609,49 @@ fn adjust_window_sort_removal(
609609
/// the plan, some of the remaining `RepartitionExec`s might become unnecessary.
610610
/// Removes such `RepartitionExec`s from the plan as well.
611611
fn remove_bottleneck_in_subplan(
612+
requirements: PlanWithCorrespondingCoalescePartitions,
613+
) -> Result<PlanWithCorrespondingCoalescePartitions> {
614+
// The root is the node `parallelize_sorts` is rewriting (a `SortExec`,
615+
// `SortPreservingMergeExec` or `CoalescePartitionsExec`). Its own distribution
616+
// requirement does not constrain the removal, because the caller drops the node and
617+
// rebuilds the cascade around the result.
618+
remove_bottleneck_in_subplan_impl(requirements, true)
619+
}
620+
621+
fn remove_bottleneck_in_subplan_impl(
612622
mut requirements: PlanWithCorrespondingCoalescePartitions,
623+
is_root: bool,
613624
) -> Result<PlanWithCorrespondingCoalescePartitions> {
614625
let plan = &requirements.plan;
626+
// Below the root, a `CoalescePartitionsExec` feeding a child that requires
627+
// `Distribution::SinglePartition` is not an avoidable bottleneck: it is what satisfies
628+
// that requirement. Removing it leaves the parent with a multi-partition input it cannot
629+
// accept, and nothing re-runs distribution enforcement afterwards, so the plan reaches
630+
// `SanityCheckPlan` invalid. The traversal reaches such a node because
631+
// `update_coalesce_ctx_children` marks a node as connected when *any* child qualifies:
632+
// a `CollectLeft` `HashJoinExec` whose probe side is connected is descended into even
633+
// though its build side must stay single-partition.
634+
//
635+
// Only `SinglePartition` is protected. A `HashPartitioned` child is in principle in the
636+
// same position — a single-partition input trivially satisfies a hash requirement, so a
637+
// coalesce below one is also load-bearing — but nothing puts a coalesce there:
638+
// `ensure_distribution` satisfies a hash requirement with a `RepartitionExec`, never a
639+
// `CoalescePartitionsExec`. Widening the check would be dead code today.
640+
let dist_reqs = plan.input_distribution_requirements();
641+
let removable = |idx: usize| {
642+
is_root
643+
|| !matches!(
644+
dist_reqs.child_distribution(idx),
645+
Some(Distribution::SinglePartition)
646+
)
647+
};
648+
let remove_from_first_child = requirements
649+
.children
650+
.first()
651+
.is_some_and(|child| is_coalesce_partitions(&child.plan))
652+
&& removable(0);
615653
let children = &mut requirements.children;
616-
if is_coalesce_partitions(&children[0].plan) {
654+
if remove_from_first_child {
617655
// We can safely use the 0th index since we have a `CoalescePartitionsExec`.
618656
let mut new_child_node = children[0].children.swap_remove(0);
619657
while new_child_node.plan.output_partitioning() == plan.output_partitioning()
@@ -627,9 +665,14 @@ fn remove_bottleneck_in_subplan(
627665
requirements.children = requirements
628666
.children
629667
.into_iter()
630-
.map(|node| {
631-
if node.data {
632-
remove_bottleneck_in_subplan(node)
668+
.enumerate()
669+
.map(|(idx, node)| {
670+
// Deliberately conservative: not descending at all also skips legitimate
671+
// cleanups *below* a protected child (a redundant second coalesce under the
672+
// load-bearing one, say). This could later be narrowed to "descend, but
673+
// protect only the topmost coalesce" if that turns out to matter.
674+
if node.data && removable(idx) {
675+
remove_bottleneck_in_subplan_impl(node, false)
633676
} else {
634677
Ok(node)
635678
}

‎datafusion/sqllogictest/test_files/joins.slt‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5608,3 +5608,77 @@ set datafusion.execution.target_partitions = 4;
56085608

56095609
statement ok
56105610
reset datafusion.execution.batch_size;
5611+
5612+
# Regression test: a `CollectLeft` `HashJoinExec` requires `SinglePartition` on its build
5613+
# (left) child, and the `CoalescePartitionsExec` that satisfies it must survive the
5614+
# sort-parallelization phase of `EnsureRequirements`. It used to be removed positionally
5615+
# (the traversal descends into the join because the *probe* side is linked to a coalesce),
5616+
# leaving a multi-partition build side that `SanityCheckPlan` rejects with
5617+
# "does not satisfy distribution requirements: SinglePartition".
5618+
5619+
statement ok
5620+
set datafusion.execution.target_partitions = 8;
5621+
5622+
# Keep the scan multi-partition as written, i.e. one partition per file.
5623+
statement ok
5624+
set datafusion.optimizer.repartition_file_scans = false;
5625+
5626+
statement ok
5627+
CREATE TABLE collect_left_src (id INT, ts INT) AS VALUES (1, 10), (2, 20), (3, 30);
5628+
5629+
query I
5630+
COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/0.parquet' STORED AS PARQUET;
5631+
----
5632+
3
5633+
5634+
query I
5635+
COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/1.parquet' STORED AS PARQUET;
5636+
----
5637+
3
5638+
5639+
query I
5640+
COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/2.parquet' STORED AS PARQUET;
5641+
----
5642+
3
5643+
5644+
query I
5645+
COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/3.parquet' STORED AS PARQUET;
5646+
----
5647+
3
5648+
5649+
statement ok
5650+
CREATE EXTERNAL TABLE collect_left STORED AS PARQUET LOCATION 'test_files/scratch/joins/collect_left/';
5651+
5652+
# The build side is the 4-partition scan; the probe side is the `DISTINCT ON` aggregate,
5653+
# whose `CoalescePartitionsExec` is what makes the traversal reach the join.
5654+
query I
5655+
SELECT a.id
5656+
FROM collect_left a
5657+
LEFT JOIN (SELECT DISTINCT ON (id) id, ts FROM collect_left ORDER BY id, ts) f
5658+
ON a.id = f.id
5659+
ORDER BY a.id;
5660+
----
5661+
1
5662+
1
5663+
1
5664+
1
5665+
2
5666+
2
5667+
2
5668+
2
5669+
3
5670+
3
5671+
3
5672+
3
5673+
5674+
statement ok
5675+
DROP TABLE collect_left;
5676+
5677+
statement ok
5678+
DROP TABLE collect_left_src;
5679+
5680+
statement ok
5681+
reset datafusion.optimizer.repartition_file_scans;
5682+
5683+
statement ok
5684+
set datafusion.execution.target_partitions = 4;

0 commit comments

Comments
 (0)