From 7f495824c6ace4ddb7bcaad588513670e80a5ecd Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 6 Sep 2026 11:58:36 +0800 Subject: [PATCH] feat(optimizer): normalize InterleaveExec to UnionExec during distribution enforcement --- .../enforce_distribution.rs | 164 +++++++++++++++++- .../enforce_distribution.rs | 35 ++++ .../src/ensure_requirements/mod.rs | 9 + 3 files changed, 207 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 1f742bc38be17..ce025dc3b631c 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -71,8 +71,9 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use datafusion_physical_plan::union::UnionExec; +use datafusion_physical_plan::union::{InterleaveExec, UnionExec}; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, displayable, @@ -731,6 +732,12 @@ impl TestConfig { // After these operations tree nodes should be in a consistent state. // This code block makes sure that these rules doesn't violate tree node integrity. { + // Mirror `EnsureRequirements`: interleaves are normalized to + // unions before distribution is enforced. + let plan = plan + .clone() + .transform_down(replace_interleave_with_union) + .data()?; let adjusted = if self.config.optimizer.top_down_join_key_reordering { // Run adjust_input_keys_ordering rule let plan_requirements = @@ -2694,6 +2701,161 @@ fn union_not_to_interleave() -> Result<()> { Ok(()) } +/// Builds `RepartitionExec(Hash([a], 10))` over a single-partition parquet scan. +fn hash_repartitioned_parquet_exec() -> Result> { + let schema = schema(); + Ok(Arc::new(RepartitionExec::try_new( + parquet_exec(), + Partitioning::Hash(vec![col("a", &schema)?], 10), + )?)) +} + +/// Builds `FinalPartitioned <- RepartitionExec(Hash([alias], 10)) <- Partial <- input`, +/// i.e. an aggregate whose output is natively hash partitioned on its group key. +fn hash_partitioned_aggregate_exec( + input: Arc, + column: &str, + alias: &str, +) -> Result> { + let schema = schema(); + let group_by = PhysicalGroupBy::new_single(vec![( + col(column, &input.schema())?, + alias.to_string(), + )]); + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + vec![], + vec![], + input, + schema.clone(), + )?); + let hash_expr = col(alias, &partial.schema())?; + let repartition = Arc::new(RepartitionExec::try_new( + partial, + Partitioning::Hash(vec![hash_expr], 10), + )?); + let final_grouping = PhysicalGroupBy::new_single(vec![( + Arc::new(Column::new(alias, 0)) as Arc, + alias.to_string(), + )]); + Ok(Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + final_grouping, + vec![], + vec![], + repartition, + schema, + )?)) +} + +#[test] +fn interleave_falls_back_to_union_when_children_lose_partitioning() -> Result<()> { + // An `InterleaveExec` whose children are hash partitioned only by + // `RepartitionExec`s that nothing requires. The rule removes those + // repartitions, after which the children can no longer be interleaved, + // so the node must degrade to a `UnionExec` instead of failing + // (https://github.com/apache/datafusion/issues/21826). + let plan: Arc = Arc::new(InterleaveExec::try_new(vec![ + hash_repartitioned_parquet_exec()?, + hash_repartitioned_parquet_exec()?, + ])?); + + let test_config = TestConfig::default(); + let plan_distrib = test_config.to_plan(plan.clone(), &DISTRIB_DISTRIB_SORT); + assert_plan!(plan_distrib, + @r" + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); + assert_plan!(plan_distrib, plan_sort); + + Ok(()) +} + +#[test] +fn interleave_fallback_still_satisfies_parent_hash_requirement() -> Result<()> { + // Same as above, but a parent requires hash partitioning. After the + // interleave degrades to a union, the parent's requirement is enforced + // with a single repartition above the union. + let interleave: Arc = Arc::new(InterleaveExec::try_new(vec![ + hash_repartitioned_parquet_exec()?, + hash_repartitioned_parquet_exec()?, + ])?); + let plan = + aggregate_exec_with_alias(interleave, vec![("a".to_string(), "a".to_string())]); + + let test_config = TestConfig::default(); + let plan_distrib = test_config.to_plan(plan.clone(), &DISTRIB_DISTRIB_SORT); + assert_plan!(plan_distrib, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); + assert_plan!(plan_distrib, plan_sort); + + Ok(()) +} + +#[test] +fn existing_interleave_is_kept_when_children_stay_interleavable() -> Result<()> { + // An `InterleaveExec` over hash-partitioned aggregates. The aggregates + // require the hash partitioning, so the children stay interleavable and + // the interleave is re-derived. + let plan: Arc = Arc::new(InterleaveExec::try_new(vec![ + hash_partitioned_aggregate_exec(parquet_exec(), "a", "a1")?, + hash_partitioned_aggregate_exec(parquet_exec(), "a", "a1")?, + ])?); + + let test_config = TestConfig::default(); + let plan_distrib = test_config.to_plan(plan.clone(), &DISTRIB_DISTRIB_SORT); + assert_plan!(plan_distrib, + @r" + InterleaveExec + AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[] + RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[] + RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + let plan_sort = test_config.to_plan(plan.clone(), &SORT_DISTRIB_DISTRIB); + assert_plan!(plan_distrib, plan_sort); + + // Interleaves are re-derived from unions, so with `prefer_existing_union` + // the rule keeps the union it normalized the input interleave to. + let test_config = TestConfig::default().with_prefer_existing_union(); + let plan_prefer_union = test_config.to_plan(plan, &DISTRIB_DISTRIB_SORT); + assert_plan!(plan_prefer_union, + @r" + UnionExec + AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[] + RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[] + RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + + Ok(()) +} + #[test] fn added_repartition_to_single_partition() -> Result<()> { let alias = vec![("a".to_string(), "a".to_string())]; diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 7b4c748030654..ae9774c9f8c2d 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -662,6 +662,37 @@ fn new_join_conditions( .collect() } +/// Turns an [`InterleaveExec`] back into the equivalent [`UnionExec`], for +/// use with `transform_down` before distribution is (re-)enforced. +/// +/// An interleave carries no distribution requirement of its own; it exists +/// only because its children happened to share a hash (or range) +/// partitioning when [`ensure_distribution`] created it from a union. Later +/// rewrites can take that partitioning away: this rule removes +/// `RepartitionExec`s nothing requires, join-key reordering changes hash +/// expressions, join side swaps change output partitioning, and so on. +/// `PlanContext` rebuilds every parent from its updated children while +/// walking the tree, and rebuilding an interleave over children that are no +/// longer interleavable fails (see +/// ). +/// +/// So, like the other distribution artifacts this rule strips and re-inserts, +/// interleaves are normalized to unions up front and re-derived by +/// [`ensure_distribution`] where the children still qualify. +/// +/// This must run top-down: demoting a nested interleave first would change +/// its output partitioning and make its parent interleave fail to rebuild. +pub fn replace_interleave_with_union( + plan: Arc, +) -> Result>> { + match plan.downcast_ref::() { + Some(interleave) => { + UnionExec::try_new(interleave.inputs().clone()).map(Transformed::yes) + } + None => Ok(Transformed::no(plan)), + } +} + /// Adds RoundRobin repartition operator to the plan increase parallelism. /// /// # Arguments @@ -1708,6 +1739,10 @@ pub fn ensure_distribution( // - Agg: // Repartition (hash): // Data + // + // An [`InterleaveExec`] present in the input was already turned back + // into a [`UnionExec`] by [`replace_interleave_with_union`], so it is + // re-derived here from the children's actual partitioning. Arc::new(InterleaveExec::try_new(children_plans)?) } else { // Route through `replace_children_if_necessary` so the common diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs index 41a03bb031629..2bc57915b2315 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs @@ -42,6 +42,8 @@ //! ```text //! EnsureRequirements::optimize(plan) //! │ +//! ├─ Phase 0: top-down Interleave → Union (replace_interleave_with_union) +//! │ //! ├─ Phase 1: top-down join-key reorder (adjust_input_keys_ordering) //! │ //! ├─ Phase 2: combined distribution + sorting (single bottom-up pass) @@ -178,6 +180,13 @@ impl PhysicalOptimizerRule for EnsureRequirements { plan: Arc, config: &ConfigOptions, ) -> Result> { + // Phase 0: Normalize `InterleaveExec` back to `UnionExec` (top-down). + // Interleaves are distribution artifacts of Phase 2, which re-derives + // them from the children's final partitioning. Keeping them would + // fail as soon as a child loses the partitioning they depend on. + use super::enforce_distribution::replace_interleave_with_union; + let plan = plan.transform_down(replace_interleave_with_union).data()?; + // Phase 1: Join key reordering (top-down, from EnforceDistribution) use super::enforce_distribution::{ PlanWithKeyRequirements, adjust_input_keys_ordering,