From ce2a735ef8f0a2a485956e5999a5a46bcf5534e4 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 24 Oct 2023 10:07:56 +0300 Subject: [PATCH 1/2] simplify code, change requirement from PhysicalSortExpr to PhysicalSortRequirement --- .../enforce_distribution.rs | 11 +++----- .../src/physical_optimizer/enforce_sorting.rs | 16 +++++------ .../replace_with_order_preserving_variants.rs | 18 ++++++------- .../src/physical_optimizer/sort_pushdown.rs | 27 +++++++++---------- .../core/src/physical_optimizer/utils.rs | 11 ++++---- 5 files changed, 38 insertions(+), 45 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index 89036e9f8ccc3..16ef4813866bf 100644 --- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs @@ -55,9 +55,7 @@ use datafusion_physical_expr::expressions::{Column, NoOp}; use datafusion_physical_expr::utils::{ map_columns_before_projection, ordering_satisfy_requirement_concrete, }; -use datafusion_physical_expr::{ - expr_list_eq_strict_order, PhysicalExpr, PhysicalSortRequirement, -}; +use datafusion_physical_expr::{expr_list_eq_strict_order, PhysicalExpr}; use datafusion_physical_plan::unbounded_output; use datafusion_physical_plan::windows::{get_best_fitting_window, BoundedWindowAggExec}; @@ -1374,10 +1372,7 @@ fn ensure_distribution( // make sure ordering requirements are still satisfied after. if ordering_satisfied { // Make sure to satisfy ordering requirement: - let sort_expr = PhysicalSortRequirement::to_sort_exprs( - required_input_ordering.clone(), - ); - add_sort_above(&mut child, sort_expr, None)?; + add_sort_above(&mut child, required_input_ordering, None)?; } } // Stop tracking distribution changing operators @@ -1715,7 +1710,7 @@ mod tests { use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; use datafusion_physical_expr::{ expressions, expressions::binary, expressions::lit, expressions::Column, - LexOrdering, PhysicalExpr, PhysicalSortExpr, + LexOrdering, PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement, }; /// Models operators like BoundedWindowExec that require an input diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index 92db3bbd053ea..e7384bf60639f 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -399,7 +399,11 @@ fn parallelize_sorts( let mut prev_layer = plan.clone(); update_child_to_remove_coalesce(&mut prev_layer, &mut coalesce_onwards[0])?; let (sort_exprs, fetch) = get_sort_exprs(&plan)?; - add_sort_above(&mut prev_layer, sort_exprs.to_vec(), fetch)?; + add_sort_above( + &mut prev_layer, + &PhysicalSortRequirement::from_sort_exprs(sort_exprs), + fetch, + )?; let spm = SortPreservingMergeExec::new(sort_exprs.to_vec(), prev_layer) .with_fetch(fetch); return Ok(Transformed::Yes(PlanWithCorrespondingCoalescePartitions { @@ -456,9 +460,7 @@ fn ensure_sorting( ) { // Make sure we preserve the ordering requirements: update_child_to_remove_unnecessary_sort(child, sort_onwards, &plan)?; - let sort_expr = - PhysicalSortRequirement::to_sort_exprs(required_ordering); - add_sort_above(child, sort_expr, None)?; + add_sort_above(child, &required_ordering, None)?; if is_sort(child) { *sort_onwards = Some(ExecTree::new(child.clone(), idx, vec![])); } else { @@ -468,8 +470,7 @@ fn ensure_sorting( } (Some(required), None) => { // Ordering requirement is not met, we should add a `SortExec` to the plan. - let sort_expr = PhysicalSortRequirement::to_sort_exprs(required); - add_sort_above(child, sort_expr, None)?; + add_sort_above(child, &required, None)?; *sort_onwards = Some(ExecTree::new(child.clone(), idx, vec![])); } (None, Some(_)) => { @@ -603,9 +604,8 @@ fn analyze_window_sort_removal( .required_input_ordering() .swap_remove(0) .unwrap_or_default(); - let sort_expr = PhysicalSortRequirement::to_sort_exprs(reqs); // Satisfy the ordering requirement so that the window can run: - add_sort_above(&mut window_child, sort_expr, None)?; + add_sort_above(&mut window_child, &reqs, None)?; let uses_bounded_memory = window_expr.iter().all(|e| e.uses_bounded_memory()); let new_window = if uses_bounded_memory { diff --git a/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs index fb1a50e18d6fe..fb75c083a70ae 100644 --- a/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs +++ b/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs @@ -71,14 +71,15 @@ impl OrderPreservationContext { // ordering, (or that can maintain ordering with the replacement of // its variant) let plan = item.plan; + let children = plan.children(); let ordering_onwards = item.ordering_onwards; - if plan.children().is_empty() { + if children.is_empty() { // Plan has no children, there is nothing to propagate. None } else if ordering_onwards[0].is_none() && ((is_repartition(&plan) && !plan.maintains_input_order()[0]) || (is_coalesce_partitions(&plan) - && plan.children()[0].output_ordering().is_some())) + && children[0].output_ordering().is_some())) { Some(ExecTree::new(plan, idx, vec![])) } else { @@ -175,19 +176,18 @@ fn get_updated_plan( // When a `RepartitionExec` doesn't preserve ordering, replace it with // a `SortPreservingRepartitionExec` if appropriate: if is_repartition(&plan) && !plan.maintains_input_order()[0] && is_spr_better { - let child = plan.children()[0].clone(); - plan = Arc::new( - RepartitionExec::try_new(child, plan.output_partitioning())? - .with_preserve_order(true), - ) as _ + let child = plan.children().swap_remove(0); + let repartition = RepartitionExec::try_new(child, plan.output_partitioning())?; + plan = Arc::new(repartition.with_preserve_order(true)) as _ } // When the input of a `CoalescePartitionsExec` has an ordering, replace it // with a `SortPreservingMergeExec` if appropriate: + let mut children = plan.children(); if is_coalesce_partitions(&plan) - && plan.children()[0].output_ordering().is_some() + && children[0].output_ordering().is_some() && is_spm_better { - let child = plan.children()[0].clone(); + let child = children.swap_remove(0); plan = Arc::new(SortPreservingMergeExec::new( child.output_ordering().unwrap_or(&[]).to_vec(), child, diff --git a/datafusion/core/src/physical_optimizer/sort_pushdown.rs b/datafusion/core/src/physical_optimizer/sort_pushdown.rs index 9b81ad3efb50d..68671280f1798 100644 --- a/datafusion/core/src/physical_optimizer/sort_pushdown.rs +++ b/datafusion/core/src/physical_optimizer/sort_pushdown.rs @@ -139,9 +139,7 @@ pub(crate) fn pushdown_sorts( || plan.ordering_equivalence_properties(), ) { // If the current plan is a SortExec, modify it to satisfy parent requirements: - let parent_required_expr = PhysicalSortRequirement::to_sort_exprs( - parent_required.ok_or_else(err)?.iter().cloned(), - ); + let parent_required_expr = parent_required.ok_or_else(err)?; new_plan = sort_exec.input().clone(); add_sort_above(&mut new_plan, parent_required_expr, sort_exec.fetch())?; }; @@ -149,13 +147,13 @@ pub(crate) fn pushdown_sorts( .output_ordering() .map(PhysicalSortRequirement::from_sort_exprs); // Since new_plan is a SortExec, we can safely get the 0th index. - let child = &new_plan.children()[0]; + let child = new_plan.children().swap_remove(0); if let Some(adjusted) = - pushdown_requirement_to_children(child, required_ordering.as_deref())? + pushdown_requirement_to_children(&child, required_ordering.as_deref())? { // Can push down requirements Ok(Transformed::Yes(SortPushDown { - plan: child.clone(), + plan: child, required_ordering: None, adjusted_request_ordering: adjusted, })) @@ -180,16 +178,14 @@ pub(crate) fn pushdown_sorts( // Can not satisfy the parent requirements, check whether the requirements can be pushed down: if let Some(adjusted) = pushdown_requirement_to_children(plan, parent_required)? { Ok(Transformed::Yes(SortPushDown { - plan: plan.clone(), + plan: requirements.plan, required_ordering: None, adjusted_request_ordering: adjusted, })) } else { // Can not push down requirements, add new SortExec: - let parent_required_expr = PhysicalSortRequirement::to_sort_exprs( - parent_required.ok_or_else(err)?.iter().cloned(), - ); - let mut new_plan = plan.clone(); + let parent_required_expr = parent_required.ok_or_else(err)?; + let mut new_plan = requirements.plan; add_sort_above(&mut new_plan, parent_required_expr, None)?; Ok(Transformed::Yes(SortPushDown::init(new_plan))) } @@ -206,7 +202,7 @@ fn pushdown_requirement_to_children( if is_window(plan) { let required_input_ordering = plan.required_input_ordering(); let request_child = required_input_ordering[0].as_deref(); - let child_plan = plan.children()[0].clone(); + let child_plan = plan.children().swap_remove(0); match determine_children_requirement(parent_required, request_child, child_plan) { RequirementsCompatibility::Satisfy => { Ok(Some(vec![request_child.map(|r| r.to_vec())])) @@ -355,16 +351,17 @@ fn try_pushdown_requirements_to_join( || smj.ordering_equivalence_properties(), ) .then(|| { - let required_input_ordering = smj.required_input_ordering(); + let mut required_input_ordering = smj.required_input_ordering(); let new_req = Some(PhysicalSortRequirement::from_sort_exprs(&sort_expr)); match push_side { JoinSide::Left => { - vec![new_req, required_input_ordering[1].clone()] + required_input_ordering[0] = new_req; } JoinSide::Right => { - vec![required_input_ordering[0].clone(), new_req] + required_input_ordering[1] = new_req; } } + required_input_ordering })) } diff --git a/datafusion/core/src/physical_optimizer/utils.rs b/datafusion/core/src/physical_optimizer/utils.rs index 0d6c85f9f22b7..dade35b8373b2 100644 --- a/datafusion/core/src/physical_optimizer/utils.rs +++ b/datafusion/core/src/physical_optimizer/utils.rs @@ -31,8 +31,8 @@ use crate::physical_plan::union::UnionExec; use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use crate::physical_plan::{displayable, ExecutionPlan}; -use datafusion_physical_expr::utils::ordering_satisfy; -use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::utils::ordering_satisfy_requirement; +use datafusion_physical_expr::PhysicalSortRequirement; /// This object implements a tree that we use while keeping track of paths /// leading to [`SortExec`]s. @@ -101,16 +101,17 @@ pub(crate) fn get_children_exectrees( /// given ordering requirements while preserving the original partitioning. pub fn add_sort_above( node: &mut Arc, - sort_expr: Vec, + sort_requirement: &[PhysicalSortRequirement], fetch: Option, ) -> Result<()> { // If the ordering requirement is already satisfied, do not add a sort. - if !ordering_satisfy( + if !ordering_satisfy_requirement( node.output_ordering(), - Some(&sort_expr), + Some(sort_requirement), || node.equivalence_properties(), || node.ordering_equivalence_properties(), ) { + let sort_expr = PhysicalSortRequirement::to_sort_exprs(sort_requirement.to_vec()); let new_sort = SortExec::new(sort_expr, node.clone()).with_fetch(fetch); *node = Arc::new(if node.output_partitioning().partition_count() > 1 { From ffdeba31f757074e041f8d268419e2619ad2921f Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 24 Oct 2023 13:26:22 +0300 Subject: [PATCH 2/2] Remove unnecessary result --- .../core/src/physical_optimizer/enforce_distribution.rs | 2 +- datafusion/core/src/physical_optimizer/enforce_sorting.rs | 8 ++++---- datafusion/core/src/physical_optimizer/sort_pushdown.rs | 4 ++-- datafusion/core/src/physical_optimizer/utils.rs | 4 +--- 4 files changed, 8 insertions(+), 10 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index 16ef4813866bf..9cd7eff4722b1 100644 --- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs @@ -1372,7 +1372,7 @@ fn ensure_distribution( // make sure ordering requirements are still satisfied after. if ordering_satisfied { // Make sure to satisfy ordering requirement: - add_sort_above(&mut child, required_input_ordering, None)?; + add_sort_above(&mut child, required_input_ordering, None); } } // Stop tracking distribution changing operators diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index e7384bf60639f..913dae07faa13 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -403,7 +403,7 @@ fn parallelize_sorts( &mut prev_layer, &PhysicalSortRequirement::from_sort_exprs(sort_exprs), fetch, - )?; + ); let spm = SortPreservingMergeExec::new(sort_exprs.to_vec(), prev_layer) .with_fetch(fetch); return Ok(Transformed::Yes(PlanWithCorrespondingCoalescePartitions { @@ -460,7 +460,7 @@ fn ensure_sorting( ) { // Make sure we preserve the ordering requirements: update_child_to_remove_unnecessary_sort(child, sort_onwards, &plan)?; - add_sort_above(child, &required_ordering, None)?; + add_sort_above(child, &required_ordering, None); if is_sort(child) { *sort_onwards = Some(ExecTree::new(child.clone(), idx, vec![])); } else { @@ -470,7 +470,7 @@ fn ensure_sorting( } (Some(required), None) => { // Ordering requirement is not met, we should add a `SortExec` to the plan. - add_sort_above(child, &required, None)?; + add_sort_above(child, &required, None); *sort_onwards = Some(ExecTree::new(child.clone(), idx, vec![])); } (None, Some(_)) => { @@ -605,7 +605,7 @@ fn analyze_window_sort_removal( .swap_remove(0) .unwrap_or_default(); // Satisfy the ordering requirement so that the window can run: - add_sort_above(&mut window_child, &reqs, None)?; + add_sort_above(&mut window_child, &reqs, None); let uses_bounded_memory = window_expr.iter().all(|e| e.uses_bounded_memory()); let new_window = if uses_bounded_memory { diff --git a/datafusion/core/src/physical_optimizer/sort_pushdown.rs b/datafusion/core/src/physical_optimizer/sort_pushdown.rs index 68671280f1798..808c4a3dadae7 100644 --- a/datafusion/core/src/physical_optimizer/sort_pushdown.rs +++ b/datafusion/core/src/physical_optimizer/sort_pushdown.rs @@ -141,7 +141,7 @@ pub(crate) fn pushdown_sorts( // If the current plan is a SortExec, modify it to satisfy parent requirements: let parent_required_expr = parent_required.ok_or_else(err)?; new_plan = sort_exec.input().clone(); - add_sort_above(&mut new_plan, parent_required_expr, sort_exec.fetch())?; + add_sort_above(&mut new_plan, parent_required_expr, sort_exec.fetch()); }; let required_ordering = new_plan .output_ordering() @@ -186,7 +186,7 @@ pub(crate) fn pushdown_sorts( // Can not push down requirements, add new SortExec: let parent_required_expr = parent_required.ok_or_else(err)?; let mut new_plan = requirements.plan; - add_sort_above(&mut new_plan, parent_required_expr, None)?; + add_sort_above(&mut new_plan, parent_required_expr, None); Ok(Transformed::Yes(SortPushDown::init(new_plan))) } } diff --git a/datafusion/core/src/physical_optimizer/utils.rs b/datafusion/core/src/physical_optimizer/utils.rs index dade35b8373b2..403af4b16ec71 100644 --- a/datafusion/core/src/physical_optimizer/utils.rs +++ b/datafusion/core/src/physical_optimizer/utils.rs @@ -21,7 +21,6 @@ use std::fmt; use std::fmt::Formatter; use std::sync::Arc; -use crate::error::Result; use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::repartition::RepartitionExec; @@ -103,7 +102,7 @@ pub fn add_sort_above( node: &mut Arc, sort_requirement: &[PhysicalSortRequirement], fetch: Option, -) -> Result<()> { +) { // If the ordering requirement is already satisfied, do not add a sort. if !ordering_satisfy_requirement( node.output_ordering(), @@ -120,7 +119,6 @@ pub fn add_sort_above( new_sort }) as _ } - Ok(()) } /// Checks whether the given operator is a limit;