diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index 1528b28294..c152f86f3a 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -19,10 +19,11 @@ use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::limit::GlobalLimitExec; +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::OrderingRequirements; +use datafusion_physical_expr::{Distribution, OrderingRequirements}; #[derive(Debug)] pub struct EnsureGlobalLimitForFetch; @@ -54,16 +55,22 @@ impl EnsureGlobalLimitForFetch { let plan = if children.is_empty() { plan } else { + let required_input_distribution = plan.required_input_distribution(); let required_input_ordering = plan.required_input_ordering(); let maintains_input_order = plan.maintains_input_order(); let child_parent = ParentContext { global_fetch: provided_global_fetch(&plan), required_ordering: None, + required_distribution: Distribution::UnspecifiedDistribution, }; let children = children .into_iter() .enumerate() .map(|(idx, child)| { + let required_distribution = required_input_distribution + .get(idx) + .cloned() + .unwrap_or(Distribution::UnspecifiedDistribution); let required_ordering = required_input_ordering .get(idx) .cloned() @@ -78,6 +85,7 @@ impl EnsureGlobalLimitForFetch { }); let parent = ParentContext { required_ordering, + required_distribution, ..child_parent.clone() }; Self::optimize_plan(Arc::clone(child), parent) @@ -99,14 +107,30 @@ impl EnsureGlobalLimitForFetch { return Ok(plan); } - Ok(add_global_fetch(plan, fetch, parent.required_ordering)) + add_global_fetch( + plan, + fetch, + parent.required_ordering, + parent.required_distribution, + ) } } -#[derive(Clone, Default)] +#[derive(Clone)] struct ParentContext { global_fetch: Option, required_ordering: Option, + required_distribution: Distribution, +} + +impl Default for ParentContext { + fn default() -> Self { + Self { + global_fetch: None, + required_ordering: None, + required_distribution: Distribution::UnspecifiedDistribution, + } + } } fn provided_global_fetch(plan: &Arc) -> Option { @@ -121,14 +145,35 @@ fn add_global_fetch( plan: Arc, fetch: usize, required_ordering: Option, -) -> Arc { - if required_ordering.is_some() + required_distribution: Distribution, +) -> DfResult> { + let original_partition_count = plan.output_partitioning().partition_count(); + let plan = if required_ordering.is_some() && let Some(ordering) = plan.output_ordering().cloned() { Arc::new(SortPreservingMergeExec::new(ordering, plan).with_fetch(Some(fetch))) + as Arc } else { Arc::new(CoalescePartitionsExec::new(plan).with_fetch(Some(fetch))) + as Arc + }; + + restore_required_distribution(plan, required_distribution, original_partition_count) +} + +fn restore_required_distribution( + plan: Arc, + required_distribution: Distribution, + partition_count: usize, +) -> DfResult> { + if partition_count <= 1 || !matches!(&required_distribution, Distribution::HashPartitioned(_)) { + return Ok(plan); } + + let partitioning = required_distribution.create_partitioning(partition_count); + Ok(Arc::new( + RepartitionExec::try_new(plan, partitioning)?.with_preserve_order(), + )) } #[cfg(test)] @@ -139,10 +184,13 @@ mod tests { use datafusion::arrow::record_batch::RecordBatch; use datafusion::physical_expr::expressions::{col, lit}; use datafusion::physical_plan::filter::FilterExecBuilder; + use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion::physical_plan::limit::GlobalLimitExec; use datafusion::physical_plan::projection::ProjectionExec; + use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::test::TestMemoryExec; - use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion_common::{JoinType, NullEquality}; + use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr}; use super::*; @@ -244,7 +292,9 @@ mod tests { filter, 1, Some(OrderingRequirements::from(required_ordering)), - ); + Distribution::UnspecifiedDistribution, + ) + .unwrap(); let merge = optimized .as_any() .downcast_ref::() @@ -278,6 +328,42 @@ mod tests { assert_eq!(child.fetch(), Some(1)); } + #[test] + fn restores_parent_hash_distribution_after_global_fetch() { + let left = filter_fetch(hash_repartition(unordered_input()), 1); + let right = hash_repartition(unordered_input()); + let on = vec![( + col("a", left.schema().as_ref()).unwrap(), + col("a", right.schema().as_ref()).unwrap(), + )]; + let join = Arc::new( + HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) as Arc; + + let optimized = + EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap(); + let left = optimized.children()[0]; + let repartition = left.as_any().downcast_ref::().unwrap(); + + 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()); @@ -307,6 +393,11 @@ mod tests { ) } + fn hash_repartition(input: Arc) -> Arc { + let partitioning = Partitioning::Hash(vec![col("a", input.schema().as_ref()).unwrap()], 3); + Arc::new(RepartitionExec::try_new(input, partitioning).unwrap()) + } + fn schema() -> Arc { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) }