diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index 392f07b8b0..2e5350f779 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -63,6 +63,7 @@ impl EnsureGlobalLimitForFetch { required_ordering: None, required_distribution: Distribution::UnspecifiedDistribution, partitioning_to_restore: None, + preserve_hash_partitioning: false, }; let children = children .into_iter() @@ -73,13 +74,9 @@ impl EnsureGlobalLimitForFetch { .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, - ) - }); + partitioning_to_restore_for(child, &required_distribution) + .or_else(|| inherited_partitioning_to_restore(&plan, child, &parent)); + let preserve_hash_partitioning = partitioning_to_restore.is_some(); let required_ordering = required_input_ordering .get(idx) .cloned() @@ -96,6 +93,7 @@ impl EnsureGlobalLimitForFetch { required_ordering, required_distribution, partitioning_to_restore, + preserve_hash_partitioning, ..child_parent.clone() }; Self::optimize_plan(Arc::clone(child), parent) @@ -132,6 +130,7 @@ struct ParentContext { required_ordering: Option, required_distribution: Distribution, partitioning_to_restore: Option, + preserve_hash_partitioning: bool, } impl Default for ParentContext { @@ -141,6 +140,7 @@ impl Default for ParentContext { required_ordering: None, required_distribution: Distribution::UnspecifiedDistribution, partitioning_to_restore: None, + preserve_hash_partitioning: false, } } } @@ -217,18 +217,30 @@ fn partitioning_to_restore_for( fn inherited_partitioning_to_restore( plan: &Arc, child: &Arc, - required_distribution: &Distribution, + parent: &ParentContext, ) -> Option { - if !matches!(required_distribution, Distribution::HashPartitioned(_)) - || child.output_partitioning().partition_count() <= 1 + if child.output_partitioning().partition_count() <= 1 || !matches!(child.output_partitioning(), Partitioning::Hash(_, _)) + || !matches!(plan.output_partitioning(), Partitioning::Hash(_, _)) + || plan.output_partitioning().partition_count() + != child.output_partitioning().partition_count() { return None; } - plan.output_partitioning() - .satisfaction(required_distribution, plan.equivalence_properties(), false) - .is_satisfied() + let satisfies_parent_distribution = matches!( + parent.required_distribution, + Distribution::HashPartitioned(_) + ) && plan + .output_partitioning() + .satisfaction( + &parent.required_distribution, + plan.equivalence_properties(), + false, + ) + .is_satisfied(); + + (satisfies_parent_distribution || parent.preserve_hash_partitioning) .then(|| child.output_partitioning().clone()) } @@ -467,6 +479,50 @@ mod tests { assert_eq!(repartition.input().fetch(), Some(1)); } + #[test] + fn restores_inherited_hash_distribution_through_multiple_projections() { + let filter = filter_fetch(hash_repartition(unordered_input()), 1); + let projection = project_a(filter); + let projection = project_a(projection); + 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 outer_projection = optimized.children()[0]; + let inner_projection = outer_projection.children()[0]; + let repartition = inner_projection.children()[0] + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(outer_projection.as_any().is::()); + assert!(inner_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()); @@ -501,6 +557,16 @@ mod tests { Arc::new(RepartitionExec::try_new(input, partitioning).unwrap()) } + fn project_a(input: Arc) -> Arc { + Arc::new( + ProjectionExec::try_new( + vec![(col("a", input.schema().as_ref()).unwrap(), "a".to_string())], + input, + ) + .unwrap(), + ) + } + fn schema() -> Arc { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) }