Skip to content

Commit ab09536

Browse files
authored
fix: split "union" and "interleave" (#6045)
* fix: split "union" and "interleave" Silently interleaving partitions is confusing and causes bugs. Fixes #5970. * refactor: add interleave during dist enforcement This avoids a top-level repartition in the following query: ```sql SELECT c1, c9 FROM aggregate_test_100 UNION ALL SELECT c1, c3 FROM aggregate_test_100 ORDER BY c9 DESC LIMIT 5 ``` * fix: docs
1 parent 26bf31e commit ab09536

5 files changed

Lines changed: 424 additions & 81 deletions

File tree

‎datafusion/core/src/physical_optimizer/dist_enforcement.rs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ use crate::physical_plan::joins::{
2929
use crate::physical_plan::projection::ProjectionExec;
3030
use crate::physical_plan::repartition::RepartitionExec;
3131
use crate::physical_plan::sorts::sort::SortOptions;
32+
use crate::physical_plan::union::{can_interleave, InterleaveExec, UnionExec};
3233
use crate::physical_plan::windows::WindowAggExec;
3334
use crate::physical_plan::Partitioning;
3435
use crate::physical_plan::{with_new_children_if_necessary, Distribution, ExecutionPlan};
@@ -825,6 +826,35 @@ fn ensure_distribution(
825826
return Ok(Transformed::No(plan));
826827
}
827828

829+
// special case for UnionExec: We want to "bubble up" hash-partitioned data. So instead of:
830+
//
831+
// Agg:
832+
// Repartition (hash):
833+
// Union:
834+
// - Agg:
835+
// Repartition (hash):
836+
// Data
837+
// - Agg:
838+
// Repartition (hash):
839+
// Data
840+
//
841+
// We can use:
842+
//
843+
// Agg:
844+
// Interleave:
845+
// - Agg:
846+
// Repartition (hash):
847+
// Data
848+
// - Agg:
849+
// Repartition (hash):
850+
// Data
851+
if let Some(union_exec) = plan.as_any().downcast_ref::<UnionExec>() {
852+
if can_interleave(union_exec.inputs()) {
853+
let plan = InterleaveExec::try_new(union_exec.inputs().clone())?;
854+
return Ok(Transformed::Yes(Arc::new(plan)));
855+
}
856+
}
857+
828858
let required_input_distributions = plan.required_input_distribution();
829859
let children: Vec<Arc<dyn ExecutionPlan>> = plan.children();
830860
assert_eq!(children.len(), required_input_distributions.len());
@@ -2134,4 +2164,42 @@ mod tests {
21342164
assert_optimized!(expected, exec);
21352165
Ok(())
21362166
}
2167+
2168+
#[test]
2169+
fn union_to_interleave() -> Result<()> {
2170+
// group by (a as a1)
2171+
let left = aggregate_exec_with_alias(
2172+
parquet_exec(),
2173+
vec![("a".to_string(), "a1".to_string())],
2174+
);
2175+
// group by (a as a2)
2176+
let right = aggregate_exec_with_alias(
2177+
parquet_exec(),
2178+
vec![("a".to_string(), "a1".to_string())],
2179+
);
2180+
2181+
// Union
2182+
let plan = Arc::new(UnionExec::new(vec![left, right]));
2183+
2184+
// final agg
2185+
let plan =
2186+
aggregate_exec_with_alias(plan, vec![("a1".to_string(), "a2".to_string())]);
2187+
2188+
// Only two RepartitionExecs added, no final RepartionExec required
2189+
let expected = &[
2190+
"AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]",
2191+
"AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]",
2192+
"InterleaveExec",
2193+
"AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]",
2194+
"RepartitionExec: partitioning=Hash([Column { name: \"a1\", index: 0 }], 10), input_partitions=1",
2195+
"AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]",
2196+
"ParquetExec: limit=None, partitions={1 group: [[x]]}, projection=[a, b, c, d, e]",
2197+
"AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]",
2198+
"RepartitionExec: partitioning=Hash([Column { name: \"a1\", index: 0 }], 10), input_partitions=1",
2199+
"AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]",
2200+
"ParquetExec: limit=None, partitions={1 group: [[x]]}, projection=[a, b, c, d, e]",
2201+
];
2202+
assert_optimized!(expected, plan);
2203+
Ok(())
2204+
}
21372205
}

0 commit comments

Comments
 (0)