From 5d21653160e5fc20b400d1161cbc13589ac42f4e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Sun, 5 Jul 2026 19:53:19 +0800 Subject: [PATCH] fix: preserve inherited distribution for global limit Signed-off-by: jeremyhi --- src/query/src/optimizer/global_limit.rs | 125 +++++++++++++++++++++--- 1 file changed, 114 insertions(+), 11 deletions(-) diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index c152f86f3a..392f07b8b0 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -23,7 +23,7 @@ use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion_common::Result as DfResult; -use datafusion_physical_expr::{Distribution, OrderingRequirements}; +use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning}; #[derive(Debug)] pub struct EnsureGlobalLimitForFetch; @@ -62,6 +62,7 @@ impl EnsureGlobalLimitForFetch { global_fetch: provided_global_fetch(&plan), required_ordering: None, required_distribution: Distribution::UnspecifiedDistribution, + partitioning_to_restore: None, }; let children = children .into_iter() @@ -71,6 +72,14 @@ impl EnsureGlobalLimitForFetch { .get(idx) .cloned() .unwrap_or(Distribution::UnspecifiedDistribution); + let partitioning_to_restore = + partitioning_to_restore_for(child, &required_distribution).or_else(|| { + inherited_partitioning_to_restore( + &plan, + child, + &parent.required_distribution, + ) + }); let required_ordering = required_input_ordering .get(idx) .cloned() @@ -86,6 +95,7 @@ impl EnsureGlobalLimitForFetch { let parent = ParentContext { required_ordering, required_distribution, + partitioning_to_restore, ..child_parent.clone() }; Self::optimize_plan(Arc::clone(child), parent) @@ -111,7 +121,7 @@ impl EnsureGlobalLimitForFetch { plan, fetch, parent.required_ordering, - parent.required_distribution, + parent.partitioning_to_restore, ) } } @@ -121,6 +131,7 @@ struct ParentContext { global_fetch: Option, required_ordering: Option, required_distribution: Distribution, + partitioning_to_restore: Option, } impl Default for ParentContext { @@ -129,6 +140,7 @@ impl Default for ParentContext { global_fetch: None, required_ordering: None, required_distribution: Distribution::UnspecifiedDistribution, + partitioning_to_restore: None, } } } @@ -145,9 +157,8 @@ fn add_global_fetch( plan: Arc, fetch: usize, required_ordering: Option, - required_distribution: Distribution, + partitioning_to_restore: Option, ) -> DfResult> { - let original_partition_count = plan.output_partitioning().partition_count(); let plan = if required_ordering.is_some() && let Some(ordering) = plan.output_ordering().cloned() { @@ -158,24 +169,69 @@ fn add_global_fetch( as Arc }; - restore_required_distribution(plan, required_distribution, original_partition_count) + restore_required_partitioning(plan, partitioning_to_restore) } -fn restore_required_distribution( +fn restore_required_partitioning( plan: Arc, - required_distribution: Distribution, - partition_count: usize, + partitioning_to_restore: Option, ) -> DfResult> { - if partition_count <= 1 || !matches!(&required_distribution, Distribution::HashPartitioned(_)) { + let Some(partitioning) = partitioning_to_restore else { + return Ok(plan); + }; + + if partitioning.partition_count() <= 1 || !matches!(&partitioning, Partitioning::Hash(_, _)) { return Ok(plan); } - let partitioning = required_distribution.create_partitioning(partition_count); Ok(Arc::new( RepartitionExec::try_new(plan, partitioning)?.with_preserve_order(), )) } +fn partitioning_to_restore_for( + child: &Arc, + required_distribution: &Distribution, +) -> Option { + if !matches!(required_distribution, Distribution::HashPartitioned(_)) + || child.output_partitioning().partition_count() <= 1 + { + return None; + } + + if child + .output_partitioning() + .satisfaction(required_distribution, child.equivalence_properties(), false) + .is_satisfied() + { + Some(child.output_partitioning().clone()) + } else { + Some( + required_distribution + .clone() + .create_partitioning(child.output_partitioning().partition_count()), + ) + } +} + +fn inherited_partitioning_to_restore( + plan: &Arc, + child: &Arc, + required_distribution: &Distribution, +) -> Option { + if !matches!(required_distribution, Distribution::HashPartitioned(_)) + || child.output_partitioning().partition_count() <= 1 + || !matches!(child.output_partitioning(), Partitioning::Hash(_, _)) + { + return None; + } + + plan.output_partitioning() + .satisfaction(required_distribution, plan.equivalence_properties(), false) + .is_satisfied() + .then(|| child.output_partitioning().clone()) +} + #[cfg(test)] mod tests { use datafusion::arrow::array::Int32Array; @@ -292,7 +348,7 @@ mod tests { filter, 1, Some(OrderingRequirements::from(required_ordering)), - Distribution::UnspecifiedDistribution, + None, ) .unwrap(); let merge = optimized @@ -364,6 +420,53 @@ mod tests { assert_eq!(repartition.input().fetch(), Some(1)); } + #[test] + fn restores_inherited_hash_distribution_through_projection() { + let filter = filter_fetch(hash_repartition(unordered_input()), 1); + let projection = Arc::new( + ProjectionExec::try_new( + vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())], + filter, + ) + .unwrap(), + ) as Arc; + let right = hash_repartition(unordered_input()); + let on = vec![( + col("a", projection.schema().as_ref()).unwrap(), + col("a", right.schema().as_ref()).unwrap(), + )]; + let join = Arc::new( + HashJoinExec::try_new( + projection, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) as Arc; + + let optimized = + EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap(); + let projection = optimized.children()[0]; + let repartition = projection.children()[0] + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(projection.as_any().is::()); + assert!(matches!( + repartition.partitioning(), + Partitioning::Hash(_, 3) + )); + assert!(repartition.input().as_any().is::()); + assert_eq!(repartition.input().fetch(), Some(1)); + } + fn unordered_input() -> Arc { let schema = schema(); let batch = batch(schema.clone());