From 5989789ef05c5eac0000bbe8fccfc7ef12639a63 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 6 Sep 2026 13:56:46 +0800 Subject: [PATCH 1/2] feat(optimizer): enhance InterleaveExec to handle non-interleavable children during rewrites --- .../enforce_distribution.rs | 140 ++++++++++++++++++ datafusion/physical-plan/src/union.rs | 96 +++++++++++- 2 files changed, 228 insertions(+), 8 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index ce025dc3b631c..4aeee422394ca 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -38,8 +38,10 @@ use datafusion::datasource::physical_plan::{CsvSource, ParquetSource}; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::ScalarValue; +use datafusion_common::Statistics; use datafusion_common::config::CsvOptions; use datafusion_common::error::Result; +use datafusion_common::stats::Precision; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; @@ -56,7 +58,9 @@ use datafusion_physical_expr_common::sort_expr::{ use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::enforce_distribution::*; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion_physical_optimizer::join_selection::JoinSelection; use datafusion_physical_optimizer::output_requirements::OutputRequirements; +use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; @@ -2856,6 +2860,142 @@ fn existing_interleave_is_kept_when_children_stay_interleavable() -> Result<()> Ok(()) } +#[test] +fn interleave_broken_by_later_rewrite_is_repaired_on_next_pass() -> Result<()> { + // The chain reported in https://github.com/apache/datafusion/issues/21826: + // a distribution pass builds an interleave, a later rule changes one + // child's partitioning while rebuilding the tree, and another distribution + // pass runs afterwards. + let alias = vec![("a".to_string(), "a1".to_string())]; + let union = union_exec(vec![ + aggregate_exec_with_alias(parquet_exec(), alias.clone()), + aggregate_exec_with_alias(parquet_exec(), alias), + ]); + let config = TestConfig::default().config; + let pass1 = EnsureRequirements::new().optimize(union, &config)?; + assert!(pass1.is::()); + + // Stand-in for the later rewrite: one child loses its hash partitioning. + // Rebuilding the interleave over it used to fail here. + let children = pass1.children(); + let coalesced: Arc = + Arc::new(CoalescePartitionsExec::new(Arc::clone(children[1]))); + let rewritten = Arc::clone(&pass1).replace_children( + vec![Arc::clone(children[0]), coalesced], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + assert!(matches!( + rewritten.output_partitioning(), + Partitioning::UnknownPartitioning(10) + )); + // Without a repair, the plan is rejected rather than executed. + assert!( + SanityCheckPlan::new() + .optimize(Arc::clone(&rewritten), &config) + .is_err() + ); + + // The next distribution pass restores a valid interleave. + let pass2 = EnsureRequirements::new().optimize(rewritten, &config)?; + SanityCheckPlan::new().optimize(Arc::clone(&pass2), &config)?; + assert_plan!(pass2, + @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 + "); + + Ok(()) +} + +/// A single-partition parquet scan with the given (inexact) row and byte +/// statistics, so `JoinSelection` can compare build and probe sizes. +fn parquet_exec_with_size( + num_rows: usize, + total_byte_size: usize, +) -> Arc { + let mut statistics = Statistics::new_unknown(&schema()); + statistics.num_rows = Precision::Inexact(num_rows); + statistics.total_byte_size = Precision::Inexact(total_byte_size); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(ParquetSource::new(schema())), + ) + .with_file(PartitionedFile::new( + "x".to_string(), + total_byte_size as u64, + )) + .with_statistics(statistics) + .build(); + DataSourceExec::from_data_source(config) +} + +#[test] +fn issue_21826_join_selection_after_distribution_pass() -> Result<()> { + // The exact chain from https://github.com/apache/datafusion/issues/21826: + // a distribution pass turns a union of hash joins into an interleave, + // `JoinSelection` then swaps the sides of one join, which changes that + // child's output partitioning, and a second distribution pass follows. + let schema = schema(); + let on: JoinOn = vec![(col("a", &schema)?, col("a", &schema)?)]; + let big = || parquet_exec_with_size(100_000, 8_000_000); + let small = || parquet_exec_with_size(10, 800); + let union = union_exec(vec![ + // Build side is the bigger input: JoinSelection swaps this join. + hash_join_exec(big(), small(), &on, &JoinType::Inner), + // Already the right way around: left as is. + hash_join_exec(small(), big(), &on, &JoinType::Inner), + ]); + let config = TestConfig::default().config; + + let pass1 = EnsureRequirements::new().optimize(union, &config)?; + assert!(pass1.is::()); + + // Rebuilding the interleave over the swapped join used to fail here with + // "Can not create InterleaveExec: new children can not be interleaved". + let reordered = JoinSelection::new().optimize(pass1, &config)?; + assert!(matches!( + reordered.output_partitioning(), + Partitioning::UnknownPartitioning(10) + )); + // Without a repair, the plan is rejected rather than executed. + assert!( + SanityCheckPlan::new() + .optimize(Arc::clone(&reordered), &config) + .is_err() + ); + + // The second distribution pass repairs it. The two joins no longer share + // a partitioning, so the result is a plain union. + let pass2 = EnsureRequirements::new().optimize(reordered, &config)?; + SanityCheckPlan::new().optimize(Arc::clone(&pass2), &config)?; + assert_plan!(pass2, + @r" + UnionExec + ProjectionExec: expr=[a@5 as a, b@6 as b, c@7 as c, d@8 as d, e@9 as e, a@0 as a, b@1 as b, c@2 as c, d@3 as d, e@4 as e] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + RepartitionExec: partitioning=Hash([a@0], 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-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 557d8f21d136b..7ef145757c6af 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -670,6 +670,23 @@ impl InterleaveExec { can_interleave(inputs.iter()), "Not all InterleaveExec children have a consistent hash or range partitioning" ); + Self::try_new_unchecked(inputs) + } + + /// Like [`Self::try_new`], but does not require the inputs to be + /// interleavable. + /// + /// Optimizer rules rebuild every parent from its rewritten children while + /// walking the plan, so a rewrite that changes a child's partitioning + /// (join side swaps, removed repartitions, ...) hands this node children + /// that are no longer interleavable before any rule has had the chance to + /// repair it. Such a node reports [`Partitioning::UnknownPartitioning`], + /// so nothing downstream can rely on a hash layout it does not have, and + /// [`ExecutionPlan::check_invariants`] rejects it at + /// [`InvariantLevel::Executable`] if it is never repaired. This mirrors + /// how unmet distribution requirements are handled for every other + /// operator. + fn try_new_unchecked(inputs: Vec>) -> Result { let schema = union_schema(&inputs)?; let inputs = inputs .into_iter() @@ -694,8 +711,14 @@ impl InterleaveExec { schema: SchemaRef, ) -> Result { let eq_properties = EquivalenceProperties::new(schema); - // Get output partitioning: - let output_partitioning = inputs[0].output_partitioning().clone(); + // Get output partitioning. Only claim the shared hash / range layout + // when every input actually has it (see `try_new_unchecked`). + let first_partitioning = inputs[0].output_partitioning(); + let output_partitioning = if can_interleave(inputs.iter()) { + first_partitioning.clone() + } else { + Partitioning::UnknownPartitioning(first_partitioning.partition_count()) + }; Ok(PlanProperties::new( eq_properties, output_partitioning, @@ -758,16 +781,25 @@ impl ExecutionPlan for InterleaveExec { ..Self::clone(&*self) })), ChildrenPropertiesMode::Recompute => { - // New children are no longer interleavable, which might be a bug of optimization rewrite. - assert_or_internal_err!( - can_interleave(children.iter()), - "Can not create InterleaveExec: new children can not be interleaved" - ); - Ok(Arc::new(InterleaveExec::try_new(children)?)) + // The new children may no longer be interleavable; see + // `try_new_unchecked` for why this is not rejected here. + Ok(Arc::new(InterleaveExec::try_new_unchecked(children)?)) } } } + fn check_invariants(&self, check: InvariantLevel) -> Result<()> { + check_default_invariants(self, check)?; + + if matches!(check, InvariantLevel::Executable) { + assert_or_internal_err!( + can_interleave(self.inputs.iter()), + "Not all InterleaveExec children have a consistent hash or range partitioning" + ); + } + Ok(()) + } + fn with_new_children( self: Arc, children: Vec>, @@ -1751,6 +1783,54 @@ mod tests { )?)) } + #[test] + fn test_interleave_rebuild_defers_partitioning_invariant() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("name", DataType::Int32, true)])); + let hash = || make_hash_exec(&schema, vec!["name"], 3); + let interleave: Arc = + Arc::new(InterleaveExec::try_new(vec![hash()?, hash()?])?); + assert!(matches!( + interleave.output_partitioning(), + Partitioning::Hash(_, 3) + )); + + // A rewrite that takes one child's hash partitioning away. The + // optimizer's tree walk rebuilds the parent from such children before + // any rule can repair it, so the rebuild itself must not fail. + let round_robin: Arc = Arc::new(RepartitionExec::try_new( + hash()?, + Partitioning::RoundRobinBatch(3), + )?); + let rebuilt = interleave.replace_children( + vec![hash()?, round_robin], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + // Explicit construction still requires interleavable inputs. + let children = rebuilt.children().into_iter().cloned().collect(); + assert!(InterleaveExec::try_new(children).is_err()); + + // The rebuilt node no longer claims a hash layout, is structurally + // sound, but is not executable until a distribution pass repairs it. + assert!(matches!( + rebuilt.output_partitioning(), + Partitioning::UnknownPartitioning(3) + )); + rebuilt.check_invariants(InvariantLevel::Always)?; + let err = rebuilt + .check_invariants(InvariantLevel::Executable) + .unwrap_err() + .to_string(); + assert!( + err.contains( + "Not all InterleaveExec children have a consistent hash or range partitioning" + ), + "{err}" + ); + Ok(()) + } + #[test] fn test_can_interleave_matrix() -> Result<()> { let name_column = "name"; From 9abad8d5e5bba21c476e87bc562e1505716522f4 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 6 Sep 2026 22:31:48 +0800 Subject: [PATCH 2/2] fix(interleave): ensure widest partition count is reported for non-interleavable inputs --- datafusion/physical-plan/src/union.rs | 52 +++++++++++++++++++++++++-- 1 file changed, 49 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 7ef145757c6af..da496b99fc6f7 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -713,11 +713,20 @@ impl InterleaveExec { let eq_properties = EquivalenceProperties::new(schema); // Get output partitioning. Only claim the shared hash / range layout // when every input actually has it (see `try_new_unchecked`). - let first_partitioning = inputs[0].output_partitioning(); let output_partitioning = if can_interleave(inputs.iter()) { - first_partitioning.clone() + inputs[0].output_partitioning().clone() } else { - Partitioning::UnknownPartitioning(first_partitioning.partition_count()) + // Non-interleavable inputs need not even agree on a partition + // count. Report the largest one: `execute` errors out for a + // partition that some input lacks, so an unrepaired node fails + // loudly instead of silently dropping the extra partitions of + // the widest input. + let partition_count = inputs + .iter() + .map(|input| input.output_partitioning().partition_count()) + .max() + .unwrap_or(0); + Partitioning::UnknownPartitioning(partition_count) }; Ok(PlanProperties::new( eq_properties, @@ -1831,6 +1840,43 @@ mod tests { Ok(()) } + #[test] + fn test_interleave_rebuild_reports_widest_partition_count() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("name", DataType::Int32, true)])); + let interleave: Arc = + Arc::new(InterleaveExec::try_new(vec![ + make_hash_exec(&schema, vec!["name"], 3)?, + make_hash_exec(&schema, vec!["name"], 3)?, + ])?); + + // A rewrite that leaves the children with different partition counts. + let rebuilt = interleave.replace_children( + vec![ + make_hash_exec(&schema, vec!["name"], 3)?, + make_hash_exec(&schema, vec!["name"], 5)?, + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + // The widest child decides the reported count, so the partitions only + // it has are still visible to callers instead of being dropped. + assert!(matches!( + rebuilt.output_partitioning(), + Partitioning::UnknownPartitioning(5) + )); + // Executing one of them fails loudly rather than returning no rows. + let Err(err) = rebuilt.execute(4, Arc::new(TaskContext::default())) else { + panic!("executing a partition the narrow child lacks must fail"); + }; + let err = err.to_string(); + assert!( + err.contains("Partition 4 not found in InterleaveExec"), + "{err}" + ); + Ok(()) + } + #[test] fn test_can_interleave_matrix() -> Result<()> { let name_column = "name";