fix: defer InterleaveExec partitioning check to InvariantLevel::Executable

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::<InterleaveExec>());
+
+    // 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<dyn ExecutionPlan> =
+        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<dyn ExecutionPlan> {
+    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::<InterleaveExec>());
+
+    // 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())];

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<Arc<dyn ExecutionPlan>>) -> Result<Self> {
         let schema = union_schema(&inputs)?;
         let inputs = inputs
             .into_iter()
@@ -694,8 +711,23 @@ impl InterleaveExec {
         schema: SchemaRef,
     ) -> Result<PlanProperties> {
         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 output_partitioning = if can_interleave(inputs.iter()) {
+            inputs[0].output_partitioning().clone()
+        } else {
+            // 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,
             output_partitioning,
@@ -758,16 +790,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<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
@@ -1751,6 +1792,91 @@ 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<dyn ExecutionPlan> =
+            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<dyn ExecutionPlan> = 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_interleave_rebuild_reports_widest_partition_count() -> Result<()> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("name", DataType::Int32, true)]));
+        let interleave: Arc<dyn ExecutionPlan> =
+            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";