You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by oz...@apache.org on 2024/01/27 22:06:03 UTC
(arrow-datafusion) branch main updated: TreeNode refactor code deduplication: Part 3 (#8817)
This is an automated email from the ASF dual-hosted git repository.
ozankabak pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 9c4affe785 TreeNode refactor code deduplication: Part 3 (#8817)
9c4affe785 is described below
commit 9c4affe785a187832927fee41c69721e346a2816
Author: Mehmet Ozan Kabak <oz...@gmail.com>
AuthorDate: Sun Jan 28 01:05:56 2024 +0300
TreeNode refactor code deduplication: Part 3 (#8817)
* Reduce code duplication
* fix doc tests
* addressing ozan's todos
* resolving merge conflicts
* remove test duplication
* Update enforce_sorting.rs
* tmp
* remove coalesce fix
* enforce dist and sort refactors
* Diff decreasing
* Review Part 1
* Review Part 2
* optimize tests
* Review Part 3
* remove new_default
* Review Part 4
* Review Part 5
* Review Part 6
* Review Part 7
* Resolve logical conflicts
* Review Part 8
* Remove clone from PlanContext
* renaming
* Remove deriving Clone on ExprContext
* Review Part 9
* Review Part 10
* Fix failing tests
* Review Part 11
* Stabilize previously unstable tests
* Review Part 12
* Adapt tests to upstream changes
* Review Part 13
* Review Part 14
* Fix import
* Move cross check inside assert optimized
* Move cross check to asset_optimized macro
* Retract aggregate topk tests (these will be solved with another PR)
* better code documentation
---------
Co-authored-by: berkaysynnada <be...@synnada.ai>
Co-authored-by: Mustafa Akur <mu...@synnada.ai>
Co-authored-by: Berkay Şahin <12...@users.noreply.github.com>
---
datafusion/common/src/tree_node.rs | 137 ++--
.../src/physical_optimizer/enforce_distribution.rs | 767 +++++++-----------
.../core/src/physical_optimizer/enforce_sorting.rs | 878 +++++++++------------
.../core/src/physical_optimizer/join_selection.rs | 283 +++----
.../src/physical_optimizer/pipeline_checker.rs | 67 +-
.../replace_with_order_preserving_variants.rs | 322 ++++----
.../core/src/physical_optimizer/sort_pushdown.rs | 201 ++---
.../core/src/physical_optimizer/test_utils.rs | 27 +-
datafusion/core/src/physical_optimizer/utils.rs | 29 +-
datafusion/expr/src/tree_node/expr.rs | 193 ++---
datafusion/expr/src/tree_node/plan.rs | 58 +-
.../physical-expr/src/equivalence/properties.rs | 38 +-
.../physical-expr/src/intervals/cp_solver.rs | 2 +-
datafusion/physical-expr/src/sort_properties.rs | 68 +-
datafusion/physical-expr/src/tree_node.rs | 71 +-
datafusion/physical-expr/src/utils/mod.rs | 53 +-
datafusion/physical-plan/src/tree_node.rs | 70 +-
17 files changed, 1479 insertions(+), 1785 deletions(-)
diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs
index 5f11c8cc1d..c5c4ee824d 100644
--- a/datafusion/common/src/tree_node.rs
+++ b/datafusion/common/src/tree_node.rs
@@ -18,11 +18,29 @@
//! This module provides common traits for visiting or rewriting tree
//! data structures easily.
-use std::borrow::Cow;
use std::sync::Arc;
use crate::Result;
+/// If the function returns [`VisitRecursion::Continue`], the normal execution of the
+/// function continues. If it returns [`VisitRecursion::Skip`], the function returns
+/// with [`VisitRecursion::Continue`] to jump next recursion step, bypassing further
+/// exploration of the current step. In case of [`VisitRecursion::Stop`], the function
+/// return with [`VisitRecursion::Stop`] and recursion halts.
+#[macro_export]
+macro_rules! handle_tree_recursion {
+ ($EXPR:expr) => {
+ match $EXPR {
+ VisitRecursion::Continue => {}
+ // If the recursion should skip, do not apply to its children, let
+ // the recursion continue:
+ VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
+ // If the recursion should stop, do not apply to its children:
+ VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
+ }
+ };
+}
+
/// Defines a visitable and rewriteable a tree node. This trait is
/// implemented for plans ([`ExecutionPlan`] and [`LogicalPlan`]) as
/// well as expression trees ([`PhysicalExpr`], [`Expr`]) in
@@ -33,27 +51,18 @@ use crate::Result;
/// [`PhysicalExpr`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.PhysicalExpr.html
/// [`LogicalPlan`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/enum.LogicalPlan.html
/// [`Expr`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/expr/enum.Expr.html
-pub trait TreeNode: Sized + Clone {
- /// Returns all children of the TreeNode
- fn children_nodes(&self) -> Vec<Cow<Self>>;
-
- /// Use preorder to iterate the node on the tree so that we can
- /// stop fast for some cases.
+pub trait TreeNode: Sized {
+ /// Applies `op` to the node and its children. `op` is applied in a preoder way,
+ /// and it is controlled by [`VisitRecursion`], which means result of the `op`
+ /// on the self node can cause an early return.
///
/// The `op` closure can be used to collect some info from the
/// tree node or do some checking for the tree node.
- fn apply<F>(&self, op: &mut F) -> Result<VisitRecursion>
- where
- F: FnMut(&Self) -> Result<VisitRecursion>,
- {
- match op(self)? {
- VisitRecursion::Continue => {}
- // If the recursion should skip, do not apply to its children. And let the recursion continue
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- // If the recursion should stop, do not apply to its children
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
- };
-
+ fn apply<F: FnMut(&Self) -> Result<VisitRecursion>>(
+ &self,
+ op: &mut F,
+ ) -> Result<VisitRecursion> {
+ handle_tree_recursion!(op(self)?);
self.apply_children(&mut |node| node.apply(op))
}
@@ -89,22 +98,8 @@ pub trait TreeNode: Sized + Clone {
&self,
visitor: &mut V,
) -> Result<VisitRecursion> {
- match visitor.pre_visit(self)? {
- VisitRecursion::Continue => {}
- // If the recursion should skip, do not apply to its children. And let the recursion continue
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- // If the recursion should stop, do not apply to its children
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
- };
-
- match self.apply_children(&mut |node| node.visit(visitor))? {
- VisitRecursion::Continue => {}
- // If the recursion should skip, do not apply to its children. And let the recursion continue
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- // If the recursion should stop, do not apply to its children
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
- }
-
+ handle_tree_recursion!(visitor.pre_visit(self)?);
+ handle_tree_recursion!(self.apply_children(&mut |node| node.visit(visitor))?);
visitor.post_visit(self)
}
@@ -148,7 +143,6 @@ pub trait TreeNode: Sized + Clone {
F: Fn(Self) -> Result<Transformed<Self>>,
{
let after_op_children = self.map_children(|node| node.transform_up(op))?;
-
let new_node = op(after_op_children)?.into();
Ok(new_node)
}
@@ -161,7 +155,6 @@ pub trait TreeNode: Sized + Clone {
F: FnMut(Self) -> Result<Transformed<Self>>,
{
let after_op_children = self.map_children(|node| node.transform_up_mut(op))?;
-
let new_node = op(after_op_children)?.into();
Ok(new_node)
}
@@ -215,17 +208,7 @@ pub trait TreeNode: Sized + Clone {
/// Apply the closure `F` to the node's children
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
where
- F: FnMut(&Self) -> Result<VisitRecursion>,
- {
- for child in self.children_nodes() {
- match op(&child)? {
- VisitRecursion::Continue => {}
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
- }
- }
- Ok(VisitRecursion::Continue)
- }
+ F: FnMut(&Self) -> Result<VisitRecursion>;
/// Apply transform `F` to the node's children, the transform `F` might have a direction(Preorder or Postorder)
fn map_children<F>(self, transform: F) -> Result<Self>
@@ -356,8 +339,15 @@ pub trait DynTreeNode {
/// Blanket implementation for Arc for any tye that implements
/// [`DynTreeNode`] (such as [`Arc<dyn PhysicalExpr>`])
impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.arc_children().into_iter().map(Cow::Owned).collect()
+ /// Apply the closure `F` to the node's children
+ fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
+ where
+ F: FnMut(&Self) -> Result<VisitRecursion>,
+ {
+ for child in self.arc_children() {
+ handle_tree_recursion!(op(&child)?)
+ }
+ Ok(VisitRecursion::Continue)
}
fn map_children<F>(self, transform: F) -> Result<Self>
@@ -366,12 +356,53 @@ impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
{
let children = self.arc_children();
if !children.is_empty() {
- let new_children: Result<Vec<_>> =
- children.into_iter().map(transform).collect();
+ let new_children =
+ children.into_iter().map(transform).collect::<Result<_>>()?;
let arc_self = Arc::clone(&self);
- self.with_new_arc_children(arc_self, new_children?)
+ self.with_new_arc_children(arc_self, new_children)
} else {
Ok(self)
}
}
}
+
+/// Instead of implementing [`TreeNode`], it's recommended to implement a [`ConcreteTreeNode`] for
+/// trees that contain nodes with payloads. This approach ensures safe execution of algorithms
+/// involving payloads, by enforcing rules for detaching and reattaching child nodes.
+pub trait ConcreteTreeNode: Sized {
+ /// Provides read-only access to child nodes.
+ fn children(&self) -> Vec<&Self>;
+
+ /// Detaches the node from its children, returning the node itself and its detached children.
+ fn take_children(self) -> (Self, Vec<Self>);
+
+ /// Reattaches updated child nodes to the node, returning the updated node.
+ fn with_new_children(self, children: Vec<Self>) -> Result<Self>;
+}
+
+impl<T: ConcreteTreeNode> TreeNode for T {
+ /// Apply the closure `F` to the node's children
+ fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
+ where
+ F: FnMut(&Self) -> Result<VisitRecursion>,
+ {
+ for child in self.children() {
+ handle_tree_recursion!(op(child)?)
+ }
+ Ok(VisitRecursion::Continue)
+ }
+
+ fn map_children<F>(self, transform: F) -> Result<Self>
+ where
+ F: FnMut(Self) -> Result<Self>,
+ {
+ let (new_self, children) = self.take_children();
+ if !children.is_empty() {
+ let new_children =
+ children.into_iter().map(transform).collect::<Result<_>>()?;
+ new_self.with_new_children(new_children)
+ } else {
+ Ok(new_self)
+ }
+ }
+}
diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
index a2f530c0e6..0c5c2d78b6 100644
--- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs
+++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
@@ -21,9 +21,7 @@
//! according to the configuration), this rule increases partition counts in
//! the physical plan.
-use std::borrow::Cow;
-use std::fmt;
-use std::fmt::Formatter;
+use std::fmt::Debug;
use std::sync::Arc;
use super::output_requirements::OutputRequirementExec;
@@ -41,11 +39,10 @@ use crate::physical_plan::joins::{
use crate::physical_plan::projection::ProjectionExec;
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use crate::physical_plan::tree_node::PlanContext;
use crate::physical_plan::union::{can_interleave, InterleaveExec, UnionExec};
use crate::physical_plan::windows::WindowAggExec;
-use crate::physical_plan::{
- with_new_children_if_necessary, Distribution, ExecutionPlan, Partitioning,
-};
+use crate::physical_plan::{Distribution, ExecutionPlan, Partitioning};
use arrow::compute::SortOptions;
use datafusion_common::tree_node::{Transformed, TreeNode};
@@ -57,8 +54,8 @@ use datafusion_physical_expr::{
PhysicalSortRequirement,
};
use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::unbounded_output;
use datafusion_physical_plan::windows::{get_best_fitting_window, BoundedWindowAggExec};
-use datafusion_physical_plan::{get_plan_string, unbounded_output};
use itertools::izip;
@@ -199,7 +196,7 @@ impl PhysicalOptimizerRule for EnforceDistribution {
let adjusted = if top_down_join_key_reordering {
// Run a top-down process to adjust input key ordering recursively
- let plan_requirements = PlanWithKeyRequirements::new(plan);
+ let plan_requirements = PlanWithKeyRequirements::new_default(plan);
let adjusted =
plan_requirements.transform_down(&adjust_input_keys_ordering)?;
adjusted.plan
@@ -210,7 +207,7 @@ impl PhysicalOptimizerRule for EnforceDistribution {
})?
};
- let distribution_context = DistributionContext::new(adjusted);
+ let distribution_context = DistributionContext::new_default(adjusted);
// Distribution enforcement needs to be applied bottom-up.
let distribution_context =
distribution_context.transform_up(&|distribution_context| {
@@ -273,8 +270,7 @@ impl PhysicalOptimizerRule for EnforceDistribution {
fn adjust_input_keys_ordering(
mut requirements: PlanWithKeyRequirements,
) -> Result<Transformed<PlanWithKeyRequirements>> {
- let parent_required = requirements.required_key_ordering.clone();
- let plan_any = requirements.plan.as_any();
+ let plan = requirements.plan.clone();
if let Some(HashJoinExec {
left,
@@ -285,7 +281,7 @@ fn adjust_input_keys_ordering(
mode,
null_equals_null,
..
- }) = plan_any.downcast_ref::<HashJoinExec>()
+ }) = plan.as_any().downcast_ref::<HashJoinExec>()
{
match mode {
PartitionMode::Partitioned => {
@@ -302,50 +298,44 @@ fn adjust_input_keys_ordering(
)
.map(|e| Arc::new(e) as _)
};
- reorder_partitioned_join_keys(
- requirements.plan.clone(),
- &parent_required,
+ return reorder_partitioned_join_keys(
+ requirements,
on,
vec![],
&join_constructor,
)
- .map(Transformed::Yes)
+ .map(Transformed::Yes);
}
PartitionMode::CollectLeft => {
- let new_right_request = match join_type {
+ // Push down requirements to the right side
+ requirements.children[1].data = match join_type {
JoinType::Inner | JoinType::Right => shift_right_required(
- &parent_required,
+ &requirements.data,
left.schema().fields().len(),
- ),
+ )
+ .unwrap_or_default(),
JoinType::RightSemi | JoinType::RightAnti => {
- Some(parent_required.clone())
+ requirements.data.clone()
}
JoinType::Left
| JoinType::LeftSemi
| JoinType::LeftAnti
- | JoinType::Full => None,
+ | JoinType::Full => vec![],
};
-
- // Push down requirements to the right side
- requirements.children[1].required_key_ordering =
- new_right_request.unwrap_or(vec![]);
- Ok(Transformed::Yes(requirements))
}
PartitionMode::Auto => {
// Can not satisfy, clear the current requirements and generate new empty requirements
- Ok(Transformed::Yes(PlanWithKeyRequirements::new(
- requirements.plan,
- )))
+ requirements.data.clear();
}
}
} else if let Some(CrossJoinExec { left, .. }) =
- plan_any.downcast_ref::<CrossJoinExec>()
+ plan.as_any().downcast_ref::<CrossJoinExec>()
{
let left_columns_len = left.schema().fields().len();
// Push down requirements to the right side
- requirements.children[1].required_key_ordering =
- shift_right_required(&parent_required, left_columns_len).unwrap_or_default();
- Ok(Transformed::Yes(requirements))
+ requirements.children[1].data =
+ shift_right_required(&requirements.data, left_columns_len)
+ .unwrap_or_default();
} else if let Some(SortMergeJoinExec {
left,
right,
@@ -354,7 +344,7 @@ fn adjust_input_keys_ordering(
sort_options,
null_equals_null,
..
- }) = plan_any.downcast_ref::<SortMergeJoinExec>()
+ }) = plan.as_any().downcast_ref::<SortMergeJoinExec>()
{
let join_constructor =
|new_conditions: (Vec<(Column, Column)>, Vec<SortOptions>)| {
@@ -368,65 +358,56 @@ fn adjust_input_keys_ordering(
)
.map(|e| Arc::new(e) as _)
};
- reorder_partitioned_join_keys(
- requirements.plan.clone(),
- &parent_required,
+ return reorder_partitioned_join_keys(
+ requirements,
on,
sort_options.clone(),
&join_constructor,
)
- .map(Transformed::Yes)
- } else if let Some(aggregate_exec) = plan_any.downcast_ref::<AggregateExec>() {
- if !parent_required.is_empty() {
- match aggregate_exec.mode() {
- AggregateMode::FinalPartitioned => reorder_aggregate_keys(
- requirements.plan.clone(),
- &parent_required,
- aggregate_exec,
- )
- .map(Transformed::Yes),
- _ => Ok(Transformed::Yes(PlanWithKeyRequirements::new(
- requirements.plan,
- ))),
+ .map(Transformed::Yes);
+ } else if let Some(aggregate_exec) = plan.as_any().downcast_ref::<AggregateExec>() {
+ if !requirements.data.is_empty() {
+ if aggregate_exec.mode() == &AggregateMode::FinalPartitioned {
+ return reorder_aggregate_keys(requirements, aggregate_exec)
+ .map(Transformed::Yes);
+ } else {
+ requirements.data.clear();
}
} else {
// Keep everything unchanged
- Ok(Transformed::No(requirements))
+ return Ok(Transformed::No(requirements));
}
- } else if let Some(proj) = plan_any.downcast_ref::<ProjectionExec>() {
+ } else if let Some(proj) = plan.as_any().downcast_ref::<ProjectionExec>() {
let expr = proj.expr();
// For Projection, we need to transform the requirements to the columns before the Projection
// And then to push down the requirements
// Construct a mapping from new name to the the orginal Column
- let new_required = map_columns_before_projection(&parent_required, expr);
- if new_required.len() == parent_required.len() {
- requirements.children[0].required_key_ordering = new_required;
- Ok(Transformed::Yes(requirements))
+ let new_required = map_columns_before_projection(&requirements.data, expr);
+ if new_required.len() == requirements.data.len() {
+ requirements.children[0].data = new_required;
} else {
// Can not satisfy, clear the current requirements and generate new empty requirements
- Ok(Transformed::Yes(PlanWithKeyRequirements::new(
- requirements.plan,
- )))
+ requirements.data.clear();
}
- } else if plan_any.downcast_ref::<RepartitionExec>().is_some()
- || plan_any.downcast_ref::<CoalescePartitionsExec>().is_some()
- || plan_any.downcast_ref::<WindowAggExec>().is_some()
+ } else if plan.as_any().downcast_ref::<RepartitionExec>().is_some()
+ || plan
+ .as_any()
+ .downcast_ref::<CoalescePartitionsExec>()
+ .is_some()
+ || plan.as_any().downcast_ref::<WindowAggExec>().is_some()
{
- Ok(Transformed::Yes(PlanWithKeyRequirements::new(
- requirements.plan,
- )))
+ requirements.data.clear();
} else {
// By default, push down the parent requirements to children
- requirements.children.iter_mut().for_each(|child| {
- child.required_key_ordering = parent_required.clone();
- });
- Ok(Transformed::Yes(requirements))
+ for child in requirements.children.iter_mut() {
+ child.data = requirements.data.clone();
+ }
}
+ Ok(Transformed::Yes(requirements))
}
fn reorder_partitioned_join_keys<F>(
- join_plan: Arc<dyn ExecutionPlan>,
- parent_required: &[Arc<dyn PhysicalExpr>],
+ mut join_plan: PlanWithKeyRequirements,
on: &[(Column, Column)],
sort_options: Vec<SortOptions>,
join_constructor: &F,
@@ -434,56 +415,48 @@ fn reorder_partitioned_join_keys<F>(
where
F: Fn((Vec<(Column, Column)>, Vec<SortOptions>)) -> Result<Arc<dyn ExecutionPlan>>,
{
+ let parent_required = &join_plan.data;
let join_key_pairs = extract_join_keys(on);
+ let eq_properties = join_plan.plan.equivalence_properties();
+
if let Some((
JoinKeyPairs {
left_keys,
right_keys,
},
new_positions,
- )) = try_reorder(
- join_key_pairs.clone(),
- parent_required,
- &join_plan.equivalence_properties(),
- ) {
+ )) = try_reorder(join_key_pairs.clone(), parent_required, &eq_properties)
+ {
if !new_positions.is_empty() {
let new_join_on = new_join_conditions(&left_keys, &right_keys);
- let mut new_sort_options: Vec<SortOptions> = vec![];
- for idx in 0..sort_options.len() {
- new_sort_options.push(sort_options[new_positions[idx]])
- }
- let mut requirement_tree = PlanWithKeyRequirements::new(join_constructor((
- new_join_on,
- new_sort_options,
- ))?);
- requirement_tree.children[0].required_key_ordering = left_keys;
- requirement_tree.children[1].required_key_ordering = right_keys;
- Ok(requirement_tree)
- } else {
- let mut requirement_tree = PlanWithKeyRequirements::new(join_plan);
- requirement_tree.children[0].required_key_ordering = left_keys;
- requirement_tree.children[1].required_key_ordering = right_keys;
- Ok(requirement_tree)
+ let new_sort_options = (0..sort_options.len())
+ .map(|idx| sort_options[new_positions[idx]])
+ .collect();
+ join_plan.plan = join_constructor((new_join_on, new_sort_options))?;
}
+ let mut requirements = join_plan;
+ requirements.children[0].data = left_keys;
+ requirements.children[1].data = right_keys;
+ Ok(requirements)
} else {
- let mut requirement_tree = PlanWithKeyRequirements::new(join_plan);
- requirement_tree.children[0].required_key_ordering = join_key_pairs.left_keys;
- requirement_tree.children[1].required_key_ordering = join_key_pairs.right_keys;
- Ok(requirement_tree)
+ let mut requirements = join_plan;
+ requirements.children[0].data = join_key_pairs.left_keys;
+ requirements.children[1].data = join_key_pairs.right_keys;
+ Ok(requirements)
}
}
fn reorder_aggregate_keys(
- agg_plan: Arc<dyn ExecutionPlan>,
- parent_required: &[Arc<dyn PhysicalExpr>],
+ mut agg_node: PlanWithKeyRequirements,
agg_exec: &AggregateExec,
) -> Result<PlanWithKeyRequirements> {
+ let parent_required = &agg_node.data;
let output_columns = agg_exec
.group_by()
.expr()
.iter()
.enumerate()
- .map(|(index, (_col, name))| Column::new(name, index))
+ .map(|(index, (_, name))| Column::new(name, index))
.collect::<Vec<_>>();
let output_exprs = output_columns
@@ -491,95 +464,82 @@ fn reorder_aggregate_keys(
.map(|c| Arc::new(c.clone()) as _)
.collect::<Vec<_>>();
- if parent_required.len() != output_exprs.len()
- || !agg_exec.group_by().null_expr().is_empty()
- || physical_exprs_equal(&output_exprs, parent_required)
+ if parent_required.len() == output_exprs.len()
+ && agg_exec.group_by().null_expr().is_empty()
+ && !physical_exprs_equal(&output_exprs, parent_required)
{
- Ok(PlanWithKeyRequirements::new(agg_plan))
- } else {
- let new_positions = expected_expr_positions(&output_exprs, parent_required);
- match new_positions {
- None => Ok(PlanWithKeyRequirements::new(agg_plan)),
- Some(positions) => {
- let new_partial_agg = if let Some(agg_exec) =
- agg_exec.input().as_any().downcast_ref::<AggregateExec>()
- {
- if matches!(agg_exec.mode(), &AggregateMode::Partial) {
- let group_exprs = agg_exec.group_by().expr();
- let new_group_exprs = positions
- .into_iter()
- .map(|idx| group_exprs[idx].clone())
- .collect();
- let new_partial_group_by =
- PhysicalGroupBy::new_single(new_group_exprs);
- Some(Arc::new(AggregateExec::try_new(
- AggregateMode::Partial,
- new_partial_group_by,
- agg_exec.aggr_expr().to_vec(),
- agg_exec.filter_expr().to_vec(),
- agg_exec.input().clone(),
- agg_exec.input_schema.clone(),
- )?))
- } else {
- None
- }
- } else {
- None
- };
- if let Some(partial_agg) = new_partial_agg {
- // Build new group expressions that correspond to the output of partial_agg
+ if let Some(positions) = expected_expr_positions(&output_exprs, parent_required) {
+ if let Some(agg_exec) =
+ agg_exec.input().as_any().downcast_ref::<AggregateExec>()
+ {
+ if matches!(agg_exec.mode(), &AggregateMode::Partial) {
+ let group_exprs = agg_exec.group_by().expr();
+ let new_group_exprs = positions
+ .into_iter()
+ .map(|idx| group_exprs[idx].clone())
+ .collect();
+ let partial_agg = Arc::new(AggregateExec::try_new(
+ AggregateMode::Partial,
+ PhysicalGroupBy::new_single(new_group_exprs),
+ agg_exec.aggr_expr().to_vec(),
+ agg_exec.filter_expr().to_vec(),
+ agg_exec.input().clone(),
+ agg_exec.input_schema.clone(),
+ )?);
+ // Build new group expressions that correspond to the output
+ // of the "reordered" aggregator:
let group_exprs = partial_agg.group_expr().expr();
- let new_final_group = partial_agg.output_group_expr();
let new_group_by = PhysicalGroupBy::new_single(
- new_final_group
- .iter()
+ partial_agg
+ .output_group_expr()
+ .into_iter()
.enumerate()
- .map(|(idx, expr)| (expr.clone(), group_exprs[idx].1.clone()))
+ .map(|(idx, expr)| (expr, group_exprs[idx].1.clone()))
.collect(),
);
-
let new_final_agg = Arc::new(AggregateExec::try_new(
AggregateMode::FinalPartitioned,
new_group_by,
agg_exec.aggr_expr().to_vec(),
agg_exec.filter_expr().to_vec(),
- partial_agg,
+ partial_agg.clone(),
agg_exec.input_schema(),
)?);
+ agg_node.plan = new_final_agg.clone();
+ agg_node.data.clear();
+ agg_node.children = vec![PlanWithKeyRequirements::new(
+ partial_agg as _,
+ vec![],
+ agg_node.children.swap_remove(0).children,
+ )];
+
// Need to create a new projection to change the expr ordering back
let agg_schema = new_final_agg.schema();
let mut proj_exprs = output_columns
.iter()
.map(|col| {
let name = col.name();
- (
- Arc::new(Column::new(
- name,
- agg_schema.index_of(name).unwrap(),
- )) as _,
- name.to_owned(),
- )
+ let index = agg_schema.index_of(name)?;
+ Ok((Arc::new(Column::new(name, index)) as _, name.to_owned()))
})
- .collect::<Vec<_>>();
+ .collect::<Result<Vec<_>>>()?;
let agg_fields = agg_schema.fields();
for (idx, field) in
agg_fields.iter().enumerate().skip(output_columns.len())
{
let name = field.name();
- proj_exprs
- .push((Arc::new(Column::new(name, idx)) as _, name.clone()))
+ let plan = Arc::new(Column::new(name, idx)) as _;
+ proj_exprs.push((plan, name.clone()))
}
- // TODO merge adjacent Projections if there are
- Ok(PlanWithKeyRequirements::new(Arc::new(
- ProjectionExec::try_new(proj_exprs, new_final_agg)?,
- )))
- } else {
- Ok(PlanWithKeyRequirements::new(agg_plan))
+ return ProjectionExec::try_new(proj_exprs, new_final_agg).map(|p| {
+ PlanWithKeyRequirements::new(Arc::new(p), vec![], vec![agg_node])
+ });
}
}
}
}
+ Ok(agg_node)
}
fn shift_right_required(
@@ -589,17 +549,11 @@ fn shift_right_required(
let new_right_required = parent_required
.iter()
.filter_map(|r| {
- if let Some(col) = r.as_any().downcast_ref::<Column>() {
- let idx = col.index();
- if idx >= left_columns_len {
- let result = Column::new(col.name(), idx - left_columns_len);
- Some(Arc::new(result) as _)
- } else {
- None
- }
- } else {
- None
- }
+ r.as_any().downcast_ref::<Column>().and_then(|col| {
+ col.index()
+ .checked_sub(left_columns_len)
+ .map(|index| Arc::new(Column::new(col.name(), index)) as _)
+ })
})
.collect::<Vec<_>>();
@@ -700,14 +654,15 @@ pub(crate) fn reorder_join_keys_to_inputs(
let new_sort_options = (0..sort_options.len())
.map(|idx| sort_options[new_positions[idx]])
.collect();
- return Ok(Arc::new(SortMergeJoinExec::try_new(
+ return SortMergeJoinExec::try_new(
left.clone(),
right.clone(),
new_join_on,
*join_type,
new_sort_options,
*null_equals_null,
- )?));
+ )
+ .map(|smj| Arc::new(smj) as _);
}
}
}
@@ -762,21 +717,21 @@ fn try_reorder(
normalized_expected = expected
.iter()
.map(|e| eq_groups.normalize_expr(e.clone()))
- .collect::<Vec<_>>();
+ .collect();
assert_eq!(normalized_expected.len(), expected.len());
normalized_left_keys = join_keys
.left_keys
.iter()
.map(|e| eq_groups.normalize_expr(e.clone()))
- .collect::<Vec<_>>();
+ .collect();
assert_eq!(join_keys.left_keys.len(), normalized_left_keys.len());
normalized_right_keys = join_keys
.right_keys
.iter()
.map(|e| eq_groups.normalize_expr(e.clone()))
- .collect::<Vec<_>>();
+ .collect();
assert_eq!(join_keys.right_keys.len(), normalized_right_keys.len());
if physical_exprs_equal(&normalized_expected, &normalized_left_keys)
@@ -793,23 +748,19 @@ fn try_reorder(
expected_expr_positions(&normalized_right_keys, &normalized_expected)
});
- if let Some(positions) = new_positions {
+ new_positions.map(|positions| {
let mut new_left_keys = vec![];
let mut new_right_keys = vec![];
for pos in positions.iter() {
new_left_keys.push(join_keys.left_keys[*pos].clone());
new_right_keys.push(join_keys.right_keys[*pos].clone());
}
- Some((
- JoinKeyPairs {
- left_keys: new_left_keys,
- right_keys: new_right_keys,
- },
- positions,
- ))
- } else {
- None
- }
+ let pairs = JoinKeyPairs {
+ left_keys: new_left_keys,
+ right_keys: new_right_keys,
+ };
+ (pairs, positions)
+ })
}
/// Return the expected expressions positions.
@@ -894,11 +845,7 @@ fn add_roundrobin_on_top(
let new_plan = Arc::new(repartition) as _;
- Ok(DistributionContext {
- plan: new_plan,
- distribution_connection: true,
- children_nodes: vec![input],
- })
+ Ok(DistributionContext::new(new_plan, true, vec![input]))
} else {
// Partition is not helpful, we already have desired number of partitions.
Ok(input)
@@ -922,7 +869,7 @@ fn add_roundrobin_on_top(
/// A [`Result`] object that contains new execution plan where the desired
/// distribution is satisfied by adding a Hash repartition.
fn add_hash_on_top(
- mut input: DistributionContext,
+ input: DistributionContext,
hash_exprs: Vec<Arc<dyn PhysicalExpr>>,
n_target: usize,
) -> Result<DistributionContext> {
@@ -952,10 +899,9 @@ fn add_hash_on_top(
let partitioning = Partitioning::Hash(hash_exprs, n_target);
let repartition = RepartitionExec::try_new(input.plan.clone(), partitioning)?
.with_preserve_order();
+ let plan = Arc::new(repartition) as _;
- input.children_nodes = vec![input.clone()];
- input.distribution_connection = true;
- input.plan = Arc::new(repartition) as _;
+ return Ok(DistributionContext::new(plan, true, vec![input]));
}
Ok(input)
@@ -992,11 +938,7 @@ fn add_spm_on_top(input: DistributionContext) -> DistributionContext {
Arc::new(CoalescePartitionsExec::new(input.plan.clone())) as _
};
- DistributionContext {
- plan: new_plan,
- distribution_connection: true,
- children_nodes: vec![input],
- }
+ DistributionContext::new(new_plan, true, vec![input])
} else {
input
}
@@ -1027,10 +969,9 @@ fn remove_dist_changing_operators(
|| is_sort_preserving_merge(&distribution_context.plan)
{
// All of above operators have a single child. First child is only child.
- let child = distribution_context.children_nodes.swap_remove(0);
// Remove any distribution changing operators at the beginning:
+ distribution_context = distribution_context.children.swap_remove(0);
// Note that they will be re-inserted later on if necessary or helpful.
- distribution_context = child;
}
Ok(distribution_context)
@@ -1057,42 +998,35 @@ fn remove_dist_changing_operators(
fn replace_order_preserving_variants(
mut context: DistributionContext,
) -> Result<DistributionContext> {
- let mut updated_children = context
- .children_nodes
- .iter()
+ context.children = context
+ .children
+ .into_iter()
.map(|child| {
- if child.distribution_connection {
- replace_order_preserving_variants(child.clone())
+ if child.data {
+ replace_order_preserving_variants(child)
} else {
- Ok(child.clone())
+ Ok(child)
}
})
.collect::<Result<Vec<_>>>()?;
if is_sort_preserving_merge(&context.plan) {
- let child = updated_children.swap_remove(0);
- context.plan = Arc::new(CoalescePartitionsExec::new(child.plan.clone()));
- context.children_nodes = vec![child];
+ let child_plan = context.children[0].plan.clone();
+ context.plan = Arc::new(CoalescePartitionsExec::new(child_plan));
return Ok(context);
} else if let Some(repartition) =
context.plan.as_any().downcast_ref::<RepartitionExec>()
{
if repartition.preserve_order() {
- let child = updated_children.swap_remove(0);
context.plan = Arc::new(RepartitionExec::try_new(
- child.plan.clone(),
+ context.children[0].plan.clone(),
repartition.partitioning().clone(),
)?);
- context.children_nodes = vec![child];
return Ok(context);
}
}
- context.plan = context
- .plan
- .clone()
- .with_new_children(updated_children.into_iter().map(|c| c.plan).collect())?;
- Ok(context)
+ context.update_plan_from_children()
}
/// This utility function adds a [`SortExec`] above an operator according to the
@@ -1109,17 +1043,11 @@ fn add_sort_preserving_partitions(
.ordering_satisfy_requirement(sort_requirement)
{
let sort_expr = PhysicalSortRequirement::to_sort_exprs(sort_requirement.to_vec());
- let new_sort = SortExec::new(sort_expr, node.plan.clone()).with_fetch(fetch);
-
- DistributionContext {
- plan: Arc::new(if node.plan.output_partitioning().partition_count() > 1 {
- new_sort.with_preserve_partitioning(true)
- } else {
- new_sort
- }),
- distribution_connection: false,
- children_nodes: vec![node],
+ let mut new_sort = SortExec::new(sort_expr, node.plan.clone()).with_fetch(fetch);
+ if node.plan.output_partitioning().partition_count() > 1 {
+ new_sort = new_sort.with_preserve_partitioning(true);
}
+ DistributionContext::new(Arc::new(new_sort), false, vec![node])
} else {
node
}
@@ -1133,7 +1061,7 @@ fn ensure_distribution(
dist_context: DistributionContext,
config: &ConfigOptions,
) -> Result<Transformed<DistributionContext>> {
- let dist_context = dist_context.update_children()?;
+ let dist_context = update_children(dist_context)?;
if dist_context.plan.children().is_empty() {
return Ok(Transformed::No(dist_context));
@@ -1154,8 +1082,8 @@ fn ensure_distribution(
// Remove unnecessary repartition from the physical plan if any
let DistributionContext {
mut plan,
- distribution_connection,
- children_nodes,
+ data,
+ children,
} = remove_dist_changing_operators(dist_context)?;
if let Some(exec) = plan.as_any().downcast_ref::<WindowAggExec>() {
@@ -1181,8 +1109,8 @@ fn ensure_distribution(
// - Satisfy the distribution requirements of every child, if it is not
// already satisfied.
// We store the updated children in `new_children`.
- let children_nodes = izip!(
- children_nodes.into_iter(),
+ let children = izip!(
+ children.into_iter(),
plan.required_input_distribution().iter(),
plan.required_input_ordering().iter(),
plan.benefits_from_input_partitioning(),
@@ -1249,7 +1177,7 @@ fn ensure_distribution(
.equivalence_properties()
.ordering_satisfy_requirement(required_input_ordering);
if (!ordering_satisfied || !order_preserving_variants_desirable)
- && child.distribution_connection
+ && child.data
{
child = replace_order_preserving_variants(child)?;
// If ordering requirements were satisfied before repartitioning,
@@ -1264,7 +1192,7 @@ fn ensure_distribution(
}
}
// Stop tracking distribution changing operators
- child.distribution_connection = false;
+ child.data = false;
} else {
// no ordering requirement
match requirement {
@@ -1286,151 +1214,78 @@ fn ensure_distribution(
)
.collect::<Result<Vec<_>>>()?;
- let new_distribution_context = DistributionContext {
- plan: if plan.as_any().is::<UnionExec>()
- && can_interleave(children_nodes.iter().map(|c| c.plan.clone()))
- {
- // Add a special case for [`UnionExec`] since we want to "bubble up"
- // hash-partitioned data. So instead of
- //
- // Agg:
- // Repartition (hash):
- // Union:
- // - Agg:
- // Repartition (hash):
- // Data
- // - Agg:
- // Repartition (hash):
- // Data
- //
- // we can use:
- //
- // Agg:
- // Interleave:
- // - Agg:
- // Repartition (hash):
- // Data
- // - Agg:
- // Repartition (hash):
- // Data
- Arc::new(InterleaveExec::try_new(
- children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?)
- } else {
- plan.with_new_children(
- children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- },
- distribution_connection,
- children_nodes,
+ let children_plans = children.iter().map(|c| c.plan.clone()).collect::<Vec<_>>();
+ plan = if plan.as_any().is::<UnionExec>() && can_interleave(children_plans.iter()) {
+ // Add a special case for [`UnionExec`] since we want to "bubble up"
+ // hash-partitioned data. So instead of
+ //
+ // Agg:
+ // Repartition (hash):
+ // Union:
+ // - Agg:
+ // Repartition (hash):
+ // Data
+ // - Agg:
+ // Repartition (hash):
+ // Data
+ //
+ // we can use:
+ //
+ // Agg:
+ // Interleave:
+ // - Agg:
+ // Repartition (hash):
+ // Data
+ // - Agg:
+ // Repartition (hash):
+ // Data
+ Arc::new(InterleaveExec::try_new(children_plans)?)
+ } else {
+ plan.with_new_children(children_plans)?
};
- Ok(Transformed::Yes(new_distribution_context))
-}
-
-/// A struct to keep track of distribution changing operators
-/// (`RepartitionExec`, `SortPreservingMergeExec`, `CoalescePartitionsExec`),
-/// and their associated parents inside `plan`. Using this information,
-/// we can optimize distribution of the plan if/when necessary.
-#[derive(Debug, Clone)]
-struct DistributionContext {
- plan: Arc<dyn ExecutionPlan>,
- /// Indicates whether this plan is connected to a distribution-changing
- /// operator.
- distribution_connection: bool,
- children_nodes: Vec<Self>,
-}
-
-impl DistributionContext {
- /// Creates a tree according to the plan with empty states.
- fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- distribution_connection: false,
- children_nodes: children.into_iter().map(Self::new).collect(),
- }
- }
-
- fn update_children(mut self) -> Result<Self> {
- for child_context in self.children_nodes.iter_mut() {
- child_context.distribution_connection = match &child_context.plan {
- plan if is_repartition(plan)
- || is_coalesce_partitions(plan)
- || is_sort_preserving_merge(plan) =>
- {
- true
- }
- _ => {
- child_context.plan.children().is_empty()
- || child_context.children_nodes[0].distribution_connection
- || child_context
- .plan
- .required_input_distribution()
- .iter()
- .zip(child_context.children_nodes.iter())
- .any(|(required_dist, child_context)| {
- child_context.distribution_connection
- && matches!(
- required_dist,
- Distribution::UnspecifiedDistribution
- )
- })
- }
- };
- }
-
- let children_plans = self
- .children_nodes
- .iter()
- .map(|context| context.plan.clone())
- .collect::<Vec<_>>();
-
- Ok(Self {
- plan: with_new_children_if_necessary(self.plan, children_plans)?.into(),
- distribution_connection: false,
- children_nodes: self.children_nodes,
- })
- }
+ Ok(Transformed::Yes(DistributionContext::new(
+ plan, data, children,
+ )))
}
-impl TreeNode for DistributionContext {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children_nodes.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children_nodes.is_empty() {
- self.children_nodes = self
- .children_nodes
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- }
- Ok(self)
+/// Keeps track of distribution changing operators (like `RepartitionExec`,
+/// `SortPreservingMergeExec`, `CoalescePartitionsExec`) and their ancestors.
+/// Using this information, we can optimize distribution of the plan if/when
+/// necessary.
+type DistributionContext = PlanContext<bool>;
+
+fn update_children(mut dist_context: DistributionContext) -> Result<DistributionContext> {
+ for child_context in dist_context.children.iter_mut() {
+ let child_plan_any = child_context.plan.as_any();
+ child_context.data =
+ if let Some(repartition) = child_plan_any.downcast_ref::<RepartitionExec>() {
+ !matches!(
+ repartition.partitioning(),
+ Partitioning::UnknownPartitioning(_)
+ )
+ } else {
+ child_plan_any.is::<SortPreservingMergeExec>()
+ || child_plan_any.is::<CoalescePartitionsExec>()
+ || child_context.plan.children().is_empty()
+ || child_context.children[0].data
+ || child_context
+ .plan
+ .required_input_distribution()
+ .iter()
+ .zip(child_context.children.iter())
+ .any(|(required_dist, child_context)| {
+ child_context.data
+ && matches!(
+ required_dist,
+ Distribution::UnspecifiedDistribution
+ )
+ })
+ }
}
-}
-/// implement Display method for `DistributionContext` struct.
-impl fmt::Display for DistributionContext {
- fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
- let plan_string = get_plan_string(&self.plan);
- write!(f, "plan: {:?}", plan_string)?;
- write!(
- f,
- "distribution_connection:{}",
- self.distribution_connection,
- )?;
- write!(f, "")
- }
+ dist_context.data = false;
+ Ok(dist_context)
}
#[derive(Debug, Clone)]
@@ -1439,49 +1294,8 @@ struct JoinKeyPairs {
right_keys: Vec<Arc<dyn PhysicalExpr>>,
}
-#[derive(Debug, Clone)]
-struct PlanWithKeyRequirements {
- plan: Arc<dyn ExecutionPlan>,
- /// Parent required key ordering
- required_key_ordering: Vec<Arc<dyn PhysicalExpr>>,
- children: Vec<Self>,
-}
-
-impl PlanWithKeyRequirements {
- fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- required_key_ordering: vec![],
- children: children.into_iter().map(Self::new).collect(),
- }
- }
-}
-
-impl TreeNode for PlanWithKeyRequirements {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children.is_empty() {
- self.children = self
- .children
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- }
- Ok(self)
- }
-}
+/// Keeps track of parent required key orderings.
+type PlanWithKeyRequirements = PlanContext<Vec<Arc<dyn PhysicalExpr>>>;
/// Since almost all of these tests explicitly use `ParquetExec` they only run with the parquet feature flag on
#[cfg(feature = "parquet")]
@@ -1497,6 +1311,9 @@ pub(crate) mod tests {
use crate::datasource::physical_plan::{CsvExec, FileScanConfig};
use crate::physical_optimizer::enforce_sorting::EnforceSorting;
use crate::physical_optimizer::output_requirements::OutputRequirements;
+ use crate::physical_optimizer::test_utils::{
+ check_integrity, coalesce_partitions_exec, repartition_exec,
+ };
use crate::physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
@@ -1506,16 +1323,12 @@ pub(crate) mod tests {
use crate::physical_plan::joins::{
utils::JoinOn, HashJoinExec, PartitionMode, SortMergeJoinExec,
};
+ use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use crate::physical_plan::projection::ProjectionExec;
+ use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use crate::physical_plan::{displayable, DisplayAs, DisplayFormatType, Statistics};
- use crate::physical_optimizer::test_utils::{
- coalesce_partitions_exec, repartition_exec,
- };
- use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
- use crate::physical_plan::sorts::sort::SortExec;
-
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion_common::ScalarValue;
@@ -1897,7 +1710,7 @@ pub(crate) mod tests {
target_partitions: usize,
prefer_existing_sort: bool,
) -> Result<Arc<dyn ExecutionPlan>> {
- let distribution_context = DistributionContext::new(plan);
+ let distribution_context = DistributionContext::new_default(plan);
let mut config = ConfigOptions::new();
config.execution.target_partitions = target_partitions;
config.optimizer.enable_round_robin_repartition = true;
@@ -1963,6 +1776,35 @@ pub(crate) mod tests {
let optimizer = OutputRequirements::new_add_mode();
let optimized = optimizer.optimize($PLAN.clone(), &config)?;
+ // This file has 2 rules that use tree node, apply these rules to original plan consecutively
+ // After these operations tree nodes should be in a consistent state.
+ // This code block makes sure that these rules doesn't violate tree node integrity.
+ {
+ let adjusted = if config.optimizer.top_down_join_key_reordering {
+ // Run adjust_input_keys_ordering rule
+ let plan_requirements =
+ PlanWithKeyRequirements::new_default($PLAN.clone());
+ let adjusted = plan_requirements
+ .transform_down(&adjust_input_keys_ordering)
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+ adjusted.plan
+ } else {
+ // Run reorder_join_keys_to_inputs rule
+ $PLAN.clone().transform_up(&|plan| {
+ Ok(Transformed::Yes(reorder_join_keys_to_inputs(plan)?))
+ })?
+ };
+
+ // Then run ensure_distribution rule
+ DistributionContext::new_default(adjusted)
+ .transform_up(&|distribution_context| {
+ ensure_distribution(distribution_context, &config)
+ })
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+ }
+
let optimized = if $FIRST_ENFORCE_DIST {
// Run enforce distribution rule first:
let optimizer = EnforceDistribution::new();
@@ -2255,9 +2097,9 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, top_join.clone(), true);
assert_optimized!(expected, top_join, false);
+
Ok(())
}
@@ -2313,6 +2155,7 @@ pub(crate) mod tests {
assert_optimized!(expected, top_join.clone(), true);
assert_optimized!(expected, top_join, false);
+
Ok(())
}
@@ -2352,6 +2195,7 @@ pub(crate) mod tests {
];
assert_optimized!(expected, join.clone(), true);
assert_optimized!(expected, join, false);
+
Ok(())
}
@@ -2404,6 +2248,7 @@ pub(crate) mod tests {
];
assert_optimized!(expected, join.clone(), true);
assert_optimized!(expected, join, false);
+
Ok(())
}
@@ -2520,6 +2365,7 @@ pub(crate) mod tests {
];
assert_optimized!(expected, filter_top_join.clone(), true);
assert_optimized!(expected, filter_top_join, false);
+
Ok(())
}
@@ -3143,6 +2989,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
assert_optimized!(expected_first_sort_enforcement, join, false, true);
+
Ok(())
}
@@ -3162,7 +3009,8 @@ pub(crate) mod tests {
let exec = Arc::new(CoalesceBatchesExec::new(exec, 4096));
// Merge from multiple parquet files and keep the data sorted
- let exec = Arc::new(SortPreservingMergeExec::new(sort_key, exec));
+ let exec: Arc<dyn ExecutionPlan> =
+ Arc::new(SortPreservingMergeExec::new(sort_key, exec));
// The optimizer should not add an additional SortExec as the
// data is already sorted
@@ -3172,6 +3020,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]",
];
assert_optimized!(expected, exec, true);
+
// In this case preserving ordering through order preserving operators is not desirable
// (according to flag: PREFER_EXISTING_SORT)
// hence in this case ordering lost during CoalescePartitionsExec and re-introduced with
@@ -3183,6 +3032,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]",
];
assert_optimized!(expected, exec, false);
+
Ok(())
}
@@ -3224,6 +3074,7 @@ pub(crate) mod tests {
];
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3239,9 +3090,9 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3258,9 +3109,9 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3281,6 +3132,7 @@ pub(crate) mod tests {
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3300,9 +3152,9 @@ pub(crate) mod tests {
"SortExec: expr=[c@2 ASC]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3328,6 +3180,7 @@ pub(crate) mod tests {
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3355,9 +3208,9 @@ pub(crate) mod tests {
// Expect no repartition to happen for local limit
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3377,6 +3230,7 @@ pub(crate) mod tests {
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3395,7 +3249,6 @@ pub(crate) mod tests {
"SortExec: expr=[c@2 ASC]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
@@ -3430,6 +3283,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3462,6 +3316,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3487,6 +3342,7 @@ pub(crate) mod tests {
assert_optimized!(expected, plan.clone(), true, true);
assert_optimized!(expected, plan, false, true);
+
Ok(())
}
@@ -3522,9 +3378,9 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3567,6 +3423,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
assert_optimized!(expected_first_sort_enforcement, plan, false);
+
Ok(())
}
@@ -3594,9 +3451,9 @@ pub(crate) mod tests {
"ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]",
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3623,9 +3480,9 @@ pub(crate) mod tests {
"ProjectionExec: expr=[a@0 as a]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, plan.clone(), true);
assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -3658,6 +3515,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
assert_optimized!(expected_first_sort_enforcement, plan, false);
+
Ok(())
}
@@ -3704,6 +3562,7 @@ pub(crate) mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
assert_optimized!(expected_first_sort_enforcement, plan, false);
+
Ok(())
}
@@ -3725,9 +3584,9 @@ pub(crate) mod tests {
"AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]",
"CsvExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true, false, 2, true, 10);
assert_optimized!(expected_csv, plan_csv, true, false, 2, true, 10);
+
Ok(())
}
@@ -3838,7 +3697,6 @@ pub(crate) mod tests {
)),
vec![("a".to_string(), "a".to_string())],
);
-
assert_optimized!(expected, plan, true, false, 2, true, 10);
}
Ok(())
@@ -3865,7 +3723,6 @@ pub(crate) mod tests {
// Plan already has two partitions
"CsvExec: file_groups={2 groups: [[x:0..100], [y:0..100]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true, false, 2, true, 10);
assert_optimized!(expected_csv, plan_csv, true, false, 2, true, 10);
Ok(())
@@ -3892,9 +3749,9 @@ pub(crate) mod tests {
// Multiple source files splitted across partitions
"CsvExec: file_groups={4 groups: [[x:0..50], [x:50..100], [y:0..50], [y:50..100]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true, false, 4, true, 10);
assert_optimized!(expected_csv, plan_csv, true, false, 4, true, 10);
+
Ok(())
}
@@ -3924,9 +3781,9 @@ pub(crate) mod tests {
// Doesn't parallelize for SortExec without preserve_partitioning
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -3969,9 +3826,9 @@ pub(crate) mod tests {
// SortExec doesn't benefit from input partitioning
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4019,9 +3876,9 @@ pub(crate) mod tests {
"LocalLimitExec: fetch=100",
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4048,9 +3905,9 @@ pub(crate) mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4078,9 +3935,9 @@ pub(crate) mod tests {
let expected_csv = &[
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4112,9 +3969,9 @@ pub(crate) mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4140,9 +3997,9 @@ pub(crate) mod tests {
"SortRequiredExec: [c@2 ASC]",
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4182,8 +4039,8 @@ pub(crate) mod tests {
"ProjectionExec: expr=[a@0 as a2, c@2 as c2]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
assert_optimized!(expected_parquet, plan_parquet, true);
+
Ok(())
}
@@ -4221,8 +4078,8 @@ pub(crate) mod tests {
"ProjectionExec: expr=[a@0 as a2, c@2 as c2]",
"CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], has_header=false",
];
-
assert_optimized!(expected_csv, plan_csv, true);
+
Ok(())
}
@@ -4245,7 +4102,6 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, physical_plan.clone(), true);
assert_optimized!(expected, physical_plan, false);
@@ -4268,7 +4124,6 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c@2 ASC",
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
// last flag sets config.optimizer.PREFER_EXISTING_SORT
assert_optimized!(expected, physical_plan.clone(), true, true);
assert_optimized!(expected, physical_plan, false, true);
@@ -4326,7 +4181,6 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
assert_optimized!(expected, physical_plan.clone(), true);
assert_optimized!(expected, physical_plan, false);
@@ -4386,7 +4240,6 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
assert_optimized!(expected, physical_plan.clone(), true);
assert_optimized!(expected, physical_plan, false);
@@ -4424,9 +4277,8 @@ pub(crate) mod tests {
config.execution.target_partitions = 10;
config.optimizer.enable_round_robin_repartition = true;
config.optimizer.prefer_existing_sort = false;
- let distribution_plan =
- EnforceDistribution::new().optimize(physical_plan, &config)?;
- assert_plan_txt!(expected, distribution_plan);
+ let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?;
+ assert_plan_txt!(expected, dist_plan);
Ok(())
}
@@ -4462,9 +4314,8 @@ pub(crate) mod tests {
config.execution.target_partitions = 10;
config.optimizer.enable_round_robin_repartition = true;
config.optimizer.prefer_existing_sort = false;
- let distribution_plan =
- EnforceDistribution::new().optimize(physical_plan, &config)?;
- assert_plan_txt!(expected, distribution_plan);
+ let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?;
+ assert_plan_txt!(expected, dist_plan);
Ok(())
}
@@ -4485,7 +4336,6 @@ pub(crate) mod tests {
"AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
// Make sure target partition number is 1. In this case hash repartition is unnecessary
assert_optimized!(expected, physical_plan.clone(), true, false, 1, false, 1024);
assert_optimized!(expected, physical_plan, false, false, 1, false, 1024);
@@ -4516,7 +4366,6 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2",
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC]",
];
-
// Make sure target partition number is larger than 2 (e.g partition number at the source).
assert_optimized!(expected, physical_plan.clone(), true, false, 4, false, 1024);
assert_optimized!(expected, physical_plan, false, false, 4, false, 1024);
@@ -4536,7 +4385,6 @@ pub(crate) mod tests {
let expected =
&["ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]"];
-
assert_optimized!(expected, physical_plan.clone(), true);
assert_optimized!(expected, physical_plan, false);
@@ -4564,7 +4412,6 @@ pub(crate) mod tests {
"RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
-
assert_optimized!(expected, physical_plan.clone(), true);
assert_optimized!(expected, physical_plan, false);
diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs
index f609ddea66..3aa9cdad18 100644
--- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs
+++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs
@@ -34,29 +34,30 @@
//! in the physical plan. The first sort is unnecessary since its result is overwritten
//! by another [`SortExec`]. Therefore, this rule removes it from the physical plan.
-use std::borrow::Cow;
use std::sync::Arc;
+use super::utils::add_sort_above;
use crate::config::ConfigOptions;
use crate::error::Result;
use crate::physical_optimizer::replace_with_order_preserving_variants::{
replace_with_order_preserving_variants, OrderPreservationContext,
};
-use crate::physical_optimizer::sort_pushdown::{pushdown_sorts, SortPushDown};
+use crate::physical_optimizer::sort_pushdown::{
+ assign_initial_requirements, pushdown_sorts, SortPushDown,
+};
use crate::physical_optimizer::utils::{
- add_sort_above, is_coalesce_partitions, is_limit, is_repartition, is_sort,
- is_sort_preserving_merge, is_union, is_window,
+ is_coalesce_partitions, is_limit, is_repartition, is_sort, is_sort_preserving_merge,
+ is_union, is_window,
};
use crate::physical_optimizer::PhysicalOptimizerRule;
use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use crate::physical_plan::tree_node::PlanContext;
use crate::physical_plan::windows::{
get_best_fitting_window, BoundedWindowAggExec, WindowAggExec,
};
-use crate::physical_plan::{
- with_new_children_if_necessary, Distribution, ExecutionPlan, InputOrderMode,
-};
+use crate::physical_plan::{Distribution, ExecutionPlan, InputOrderMode};
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::{plan_err, DataFusionError};
@@ -78,177 +79,70 @@ impl EnforceSorting {
}
/// This object is used within the [`EnforceSorting`] rule to track the closest
-/// [`SortExec`] descendant(s) for every child of a plan.
-#[derive(Debug, Clone)]
-struct PlanWithCorrespondingSort {
- plan: Arc<dyn ExecutionPlan>,
- // For every child, track `ExecutionPlan`s starting from the child until
- // the `SortExec`(s). If the child has no connection to any sort, it simply
- // stores false.
- sort_connection: bool,
- children_nodes: Vec<Self>,
-}
-
-impl PlanWithCorrespondingSort {
- fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- sort_connection: false,
- children_nodes: children.into_iter().map(Self::new).collect(),
+/// [`SortExec`] descendant(s) for every child of a plan. The data attribute
+/// stores whether the plan is a `SortExec` or is connected to a `SortExec`
+/// via its children.
+type PlanWithCorrespondingSort = PlanContext<bool>;
+
+fn update_sort_ctx_children(
+ mut node: PlanWithCorrespondingSort,
+ data: bool,
+) -> Result<PlanWithCorrespondingSort> {
+ for child_node in node.children.iter_mut() {
+ let plan = &child_node.plan;
+ child_node.data = if is_sort(plan) {
+ // Initiate connection:
+ true
+ } else if is_limit(plan) {
+ // There is no sort linkage for this path, it starts at a limit.
+ false
+ } else {
+ let is_spm = is_sort_preserving_merge(plan);
+ let required_orderings = plan.required_input_ordering();
+ let flags = plan.maintains_input_order();
+ // Add parent node to the tree if there is at least one child with
+ // a sort connection:
+ izip!(flags, required_orderings).any(|(maintains, required_ordering)| {
+ let propagates_ordering =
+ (maintains && required_ordering.is_none()) || is_spm;
+ let connected_to_sort =
+ child_node.children.iter().any(|child| child.data);
+ propagates_ordering && connected_to_sort
+ })
}
}
- fn update_children(
- parent_plan: Arc<dyn ExecutionPlan>,
- mut children_nodes: Vec<Self>,
- ) -> Result<Self> {
- for node in children_nodes.iter_mut() {
- let plan = &node.plan;
- // Leaves of `sort_onwards` are `SortExec` operators, which impose
- // an ordering. This tree collects all the intermediate executors
- // that maintain this ordering. If we just saw a order imposing
- // operator, we reset the tree and start accumulating.
- node.sort_connection = if is_sort(plan) {
- // Initiate connection
- true
- } else if is_limit(plan) {
- // There is no sort linkage for this path, it starts at a limit.
- false
- } else {
- let is_spm = is_sort_preserving_merge(plan);
- let required_orderings = plan.required_input_ordering();
- let flags = plan.maintains_input_order();
- // Add parent node to the tree if there is at least one
- // child with a sort connection:
- izip!(flags, required_orderings).any(|(maintains, required_ordering)| {
- let propagates_ordering =
- (maintains && required_ordering.is_none()) || is_spm;
- let connected_to_sort =
- node.children_nodes.iter().any(|item| item.sort_connection);
- propagates_ordering && connected_to_sort
- })
- }
- }
-
- let children_plans = children_nodes
- .iter()
- .map(|item| item.plan.clone())
- .collect::<Vec<_>>();
- let plan = with_new_children_if_necessary(parent_plan, children_plans)?.into();
-
- Ok(Self {
- plan,
- sort_connection: false,
- children_nodes,
- })
- }
-}
-
-impl TreeNode for PlanWithCorrespondingSort {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children_nodes.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children_nodes.is_empty() {
- self.children_nodes = self
- .children_nodes
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- }
- Ok(self)
- }
+ node.data = data;
+ node.update_plan_from_children()
}
/// This object is used within the [`EnforceSorting`] rule to track the closest
-/// [`CoalescePartitionsExec`] descendant(s) for every child of a plan.
-#[derive(Debug, Clone)]
-struct PlanWithCorrespondingCoalescePartitions {
- plan: Arc<dyn ExecutionPlan>,
- // Stores whether the plan is a `CoalescePartitionsExec` or it is connected to
- // a `CoalescePartitionsExec` via its children.
- coalesce_connection: bool,
- children_nodes: Vec<Self>,
-}
-
-impl PlanWithCorrespondingCoalescePartitions {
- /// Creates an empty tree with empty connections.
- fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- coalesce_connection: false,
- children_nodes: children.into_iter().map(Self::new).collect(),
- }
- }
-
- fn update_children(mut self) -> Result<Self> {
- self.coalesce_connection = if self.plan.children().is_empty() {
- // Plan has no children, it cannot be a `CoalescePartitionsExec`.
- false
- } else if is_coalesce_partitions(&self.plan) {
- // Initiate a connection
- true
- } else {
- self.children_nodes
- .iter()
- .enumerate()
- .map(|(idx, node)| {
- // Only consider operators that don't require a
- // single partition, and connected to any coalesce
- node.coalesce_connection
- && !matches!(
- self.plan.required_input_distribution()[idx],
- Distribution::SinglePartition
- )
- // If all children are None. There is nothing to track, set connection false.
- })
- .any(|c| c)
- };
-
- let children_plans = self
- .children_nodes
- .iter()
- .map(|item| item.plan.clone())
- .collect();
- self.plan = with_new_children_if_necessary(self.plan, children_plans)?.into();
- Ok(self)
- }
-}
-
-impl TreeNode for PlanWithCorrespondingCoalescePartitions {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children_nodes.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children_nodes.is_empty() {
- self.children_nodes = self
- .children_nodes
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- }
- Ok(self)
- }
+/// [`CoalescePartitionsExec`] descendant(s) for every child of a plan. The data
+/// attribute stores whether the plan is a `CoalescePartitionsExec` or is
+/// connected to a `CoalescePartitionsExec` via its children.
+type PlanWithCorrespondingCoalescePartitions = PlanContext<bool>;
+
+fn update_coalesce_ctx_children(
+ coalesce_context: &mut PlanWithCorrespondingCoalescePartitions,
+) {
+ let children = &coalesce_context.children;
+ coalesce_context.data = if children.is_empty() {
+ // Plan has no children, it cannot be a `CoalescePartitionsExec`.
+ false
+ } else if is_coalesce_partitions(&coalesce_context.plan) {
+ // Initiate a connection:
+ true
+ } else {
+ children.iter().enumerate().any(|(idx, node)| {
+ // Only consider operators that don't require a single partition,
+ // and connected to some `CoalescePartitionsExec`:
+ node.data
+ && !matches!(
+ coalesce_context.plan.required_input_distribution()[idx],
+ Distribution::SinglePartition
+ )
+ })
+ };
}
/// The boolean flag `repartition_sorts` defined in the config indicates
@@ -261,13 +155,13 @@ impl PhysicalOptimizerRule for EnforceSorting {
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
- let plan_requirements = PlanWithCorrespondingSort::new(plan);
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
// Execute a bottom-up traversal to enforce sorting requirements,
// remove unnecessary sorts, and optimize sort-sensitive operators:
let adjusted = plan_requirements.transform_up(&ensure_sorting)?;
let new_plan = if config.optimizer.repartition_sorts {
let plan_with_coalesce_partitions =
- PlanWithCorrespondingCoalescePartitions::new(adjusted.plan);
+ PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
let parallel =
plan_with_coalesce_partitions.transform_up(¶llelize_sorts)?;
parallel.plan
@@ -275,7 +169,7 @@ impl PhysicalOptimizerRule for EnforceSorting {
adjusted.plan
};
- let plan_with_pipeline_fixer = OrderPreservationContext::new(new_plan);
+ let plan_with_pipeline_fixer = OrderPreservationContext::new_default(new_plan);
let updated_plan =
plan_with_pipeline_fixer.transform_up(&|plan_with_pipeline_fixer| {
replace_with_order_preserving_variants(
@@ -288,8 +182,8 @@ impl PhysicalOptimizerRule for EnforceSorting {
// Execute a top-down traversal to exploit sort push-down opportunities
// missed by the bottom-up traversal:
- let mut sort_pushdown = SortPushDown::new(updated_plan.plan);
- sort_pushdown.assign_initial_requirements();
+ let mut sort_pushdown = SortPushDown::new_default(updated_plan.plan);
+ assign_initial_requirements(&mut sort_pushdown);
let adjusted = sort_pushdown.transform_down(&pushdown_sorts)?;
Ok(adjusted.plan)
}
@@ -318,414 +212,347 @@ impl PhysicalOptimizerRule for EnforceSorting {
/// by following connections from [`CoalescePartitionsExec`]s to [`SortExec`]s.
/// By performing sorting in parallel, we can increase performance in some scenarios.
fn parallelize_sorts(
- requirements: PlanWithCorrespondingCoalescePartitions,
+ mut requirements: PlanWithCorrespondingCoalescePartitions,
) -> Result<Transformed<PlanWithCorrespondingCoalescePartitions>> {
- let PlanWithCorrespondingCoalescePartitions {
- mut plan,
- coalesce_connection,
- mut children_nodes,
- } = requirements.update_children()?;
-
- if plan.children().is_empty() || !children_nodes[0].coalesce_connection {
- // We only take an action when the plan is either a SortExec, a
- // SortPreservingMergeExec or a CoalescePartitionsExec, and they
- // all have a single child. Therefore, if the first child is `None`,
- // we can return immediately.
- return Ok(Transformed::No(PlanWithCorrespondingCoalescePartitions {
- plan,
- coalesce_connection,
- children_nodes,
- }));
- } else if (is_sort(&plan) || is_sort_preserving_merge(&plan))
- && plan.output_partitioning().partition_count() <= 1
+ update_coalesce_ctx_children(&mut requirements);
+
+ if requirements.children.is_empty() || !requirements.children[0].data {
+ // We only take an action when the plan is either a `SortExec`, a
+ // `SortPreservingMergeExec` or a `CoalescePartitionsExec`, and they
+ // all have a single child. Therefore, if the first child has no
+ // connection, we can return immediately.
+ Ok(Transformed::No(requirements))
+ } else if (is_sort(&requirements.plan)
+ || is_sort_preserving_merge(&requirements.plan))
+ && requirements.plan.output_partitioning().partition_count() <= 1
{
- // If there is a connection between a CoalescePartitionsExec and a
- // global sort that satisfy the requirements (i.e. intermediate
- // executors don't require single partition), then we can replace
- // the CoalescePartitionsExec + Sort cascade with a SortExec +
- // SortPreservingMergeExec cascade to parallelize sorting.
- let (sort_exprs, fetch) = get_sort_exprs(&plan)?;
+ // Take the initial sort expressions and requirements
+ let (sort_exprs, fetch) = get_sort_exprs(&requirements.plan)?;
let sort_reqs = PhysicalSortRequirement::from_sort_exprs(sort_exprs);
let sort_exprs = sort_exprs.to_vec();
- update_child_to_remove_coalesce(&mut plan, &mut children_nodes[0])?;
- add_sort_above(&mut plan, &sort_reqs, fetch);
- let spm = SortPreservingMergeExec::new(sort_exprs, plan).with_fetch(fetch);
-
- return Ok(Transformed::Yes(
- PlanWithCorrespondingCoalescePartitions::new(Arc::new(spm)),
- ));
- } else if is_coalesce_partitions(&plan) {
- // There is an unnecessary `CoalescePartitionsExec` in the plan.
- update_child_to_remove_coalesce(&mut plan, &mut children_nodes[0])?;
- let new_plan = Arc::new(CoalescePartitionsExec::new(plan)) as _;
- return Ok(Transformed::Yes(
- PlanWithCorrespondingCoalescePartitions::new(new_plan),
- ));
- }
+ // If there is a connection between a `CoalescePartitionsExec` and a
+ // global sort that satisfy the requirements (i.e. intermediate
+ // executors don't require single partition), then we can replace
+ // the `CoalescePartitionsExec` + `SortExec` cascade with a `SortExec`
+ // + `SortPreservingMergeExec` cascade to parallelize sorting.
+ requirements = remove_corresponding_coalesce_in_sub_plan(requirements)?;
+ // We also need to remove the self node since `remove_corresponding_coalesce_in_sub_plan`
+ // deals with the children and their children and so on.
+ requirements = requirements.children.swap_remove(0);
+
+ if !requirements
+ .plan
+ .equivalence_properties()
+ .ordering_satisfy_requirement(&sort_reqs)
+ {
+ requirements = add_sort_above(requirements, sort_reqs, fetch);
+ }
- Ok(Transformed::Yes(PlanWithCorrespondingCoalescePartitions {
- plan,
- coalesce_connection,
- children_nodes,
- }))
+ let spm = SortPreservingMergeExec::new(sort_exprs, requirements.plan.clone());
+ Ok(Transformed::Yes(
+ PlanWithCorrespondingCoalescePartitions::new(
+ Arc::new(spm.with_fetch(fetch)),
+ false,
+ vec![requirements],
+ ),
+ ))
+ } else if is_coalesce_partitions(&requirements.plan) {
+ // There is an unnecessary `CoalescePartitionsExec` in the plan.
+ // This will handle the recursive `CoalescePartitionsExec` plans.
+ requirements = remove_corresponding_coalesce_in_sub_plan(requirements)?;
+ // For the removal of self node which is also a `CoalescePartitionsExec`.
+ requirements = requirements.children.swap_remove(0);
+
+ Ok(Transformed::Yes(
+ PlanWithCorrespondingCoalescePartitions::new(
+ Arc::new(CoalescePartitionsExec::new(requirements.plan.clone())),
+ false,
+ vec![requirements],
+ ),
+ ))
+ } else {
+ Ok(Transformed::Yes(requirements))
+ }
}
/// This function enforces sorting requirements and makes optimizations without
/// violating these requirements whenever possible.
fn ensure_sorting(
- requirements: PlanWithCorrespondingSort,
+ mut requirements: PlanWithCorrespondingSort,
) -> Result<Transformed<PlanWithCorrespondingSort>> {
- let requirements = PlanWithCorrespondingSort::update_children(
- requirements.plan,
- requirements.children_nodes,
- )?;
+ requirements = update_sort_ctx_children(requirements, false)?;
// Perform naive analysis at the beginning -- remove already-satisfied sorts:
- if requirements.plan.children().is_empty() {
+ if requirements.children.is_empty() {
return Ok(Transformed::No(requirements));
}
- if let Some(result) = analyze_immediate_sort_removal(&requirements) {
- return Ok(Transformed::Yes(result));
- }
-
- let plan = requirements.plan;
- let mut children_nodes = requirements.children_nodes;
+ let maybe_requirements = analyze_immediate_sort_removal(requirements);
+ let Transformed::No(mut requirements) = maybe_requirements else {
+ return Ok(maybe_requirements);
+ };
- for (idx, (child_node, required_ordering)) in
- izip!(children_nodes.iter_mut(), plan.required_input_ordering()).enumerate()
+ let plan = &requirements.plan;
+ let mut updated_children = vec![];
+ for (idx, (required_ordering, mut child)) in plan
+ .required_input_ordering()
+ .into_iter()
+ .zip(requirements.children.into_iter())
+ .enumerate()
{
- let mut child_plan = child_node.plan.clone();
- let physical_ordering = child_plan.output_ordering();
- match (required_ordering, physical_ordering) {
- (Some(required_ordering), Some(_)) => {
- if !child_plan
- .equivalence_properties()
- .ordering_satisfy_requirement(&required_ordering)
- {
- // Make sure we preserve the ordering requirements:
- update_child_to_remove_unnecessary_sort(idx, child_node, &plan)?;
- add_sort_above(&mut child_plan, &required_ordering, None);
- if is_sort(&child_plan) {
- *child_node = PlanWithCorrespondingSort::update_children(
- child_plan,
- vec![child_node.clone()],
- )?;
- child_node.sort_connection = true;
- }
+ let physical_ordering = child.plan.output_ordering();
+
+ if let Some(required) = required_ordering {
+ let eq_properties = child.plan.equivalence_properties();
+ if !eq_properties.ordering_satisfy_requirement(&required) {
+ // Make sure we preserve the ordering requirements:
+ if physical_ordering.is_some() {
+ child = update_child_to_remove_unnecessary_sort(idx, child, plan)?;
}
+ child = add_sort_above(child, required, None);
+ child = update_sort_ctx_children(child, true)?;
}
- (Some(required), None) => {
- // Ordering requirement is not met, we should add a `SortExec` to the plan.
- add_sort_above(&mut child_plan, &required, None);
- *child_node = PlanWithCorrespondingSort::update_children(
- child_plan,
- vec![child_node.clone()],
- )?;
- child_node.sort_connection = true;
- }
- (None, Some(_)) => {
- // We have a `SortExec` whose effect may be neutralized by
- // another order-imposing operator. Remove this sort.
- if !plan.maintains_input_order()[idx] || is_union(&plan) {
- update_child_to_remove_unnecessary_sort(idx, child_node, &plan)?;
- }
- }
- (None, None) => {
- update_child_to_remove_unnecessary_sort(idx, child_node, &plan)?;
- }
+ } else if physical_ordering.is_none()
+ || !plan.maintains_input_order()[idx]
+ || is_union(plan)
+ {
+ // We have a `SortExec` whose effect may be neutralized by another
+ // order-imposing operator, remove this sort:
+ child = update_child_to_remove_unnecessary_sort(idx, child, plan)?;
}
+ updated_children.push(child);
}
+ requirements.children = updated_children;
// For window expressions, we can remove some sorts when we can
// calculate the result in reverse:
- if is_window(&plan) && children_nodes[0].sort_connection {
- if let Some(result) = analyze_window_sort_removal(&mut children_nodes[0], &plan)?
- {
- return Ok(Transformed::Yes(result));
- }
- } else if is_sort_preserving_merge(&plan)
- && children_nodes[0]
- .plan
- .output_partitioning()
- .partition_count()
- <= 1
+ let child_node = &requirements.children[0];
+ if is_window(plan) && child_node.data {
+ return adjust_window_sort_removal(requirements).map(Transformed::Yes);
+ } else if is_sort_preserving_merge(plan)
+ && child_node.plan.output_partitioning().partition_count() <= 1
{
- // This SortPreservingMergeExec is unnecessary, input already has a
+ // This `SortPreservingMergeExec` is unnecessary, input already has a
// single partition.
- let child_node = children_nodes.swap_remove(0);
+ let child_node = requirements.children.swap_remove(0);
return Ok(Transformed::Yes(child_node));
}
- Ok(Transformed::Yes(
- PlanWithCorrespondingSort::update_children(plan, children_nodes)?,
- ))
+
+ update_sort_ctx_children(requirements, false).map(Transformed::Yes)
}
/// Analyzes a given [`SortExec`] (`plan`) to determine whether its input
/// already has a finer ordering than it enforces.
fn analyze_immediate_sort_removal(
- node: &PlanWithCorrespondingSort,
-) -> Option<PlanWithCorrespondingSort> {
- let PlanWithCorrespondingSort {
- plan,
- children_nodes,
- ..
- } = node;
- if let Some(sort_exec) = plan.as_any().downcast_ref::<SortExec>() {
- let sort_input = sort_exec.input().clone();
+ mut node: PlanWithCorrespondingSort,
+) -> Transformed<PlanWithCorrespondingSort> {
+ if let Some(sort_exec) = node.plan.as_any().downcast_ref::<SortExec>() {
+ let sort_input = sort_exec.input();
// If this sort is unnecessary, we should remove it:
if sort_input
.equivalence_properties()
.ordering_satisfy(sort_exec.output_ordering().unwrap_or(&[]))
{
- // Since we know that a `SortExec` has exactly one child,
- // we can use the zero index safely:
- return Some(
- if !sort_exec.preserve_partitioning()
- && sort_input.output_partitioning().partition_count() > 1
- {
- // Replace the sort with a sort-preserving merge:
- let new_plan: Arc<dyn ExecutionPlan> =
- Arc::new(SortPreservingMergeExec::new(
- sort_exec.expr().to_vec(),
- sort_input,
- ));
- PlanWithCorrespondingSort {
- plan: new_plan,
- // SortPreservingMergeExec has single child.
- sort_connection: false,
- children_nodes: children_nodes
- .iter()
- .cloned()
- .map(|mut node| {
- node.sort_connection = false;
- node
- })
- .collect(),
- }
- } else {
- // Remove the sort:
- PlanWithCorrespondingSort {
- plan: sort_input,
- sort_connection: false,
- children_nodes: children_nodes[0]
- .children_nodes
- .iter()
- .cloned()
- .map(|mut node| {
- node.sort_connection = false;
- node
- })
- .collect(),
- }
- },
- );
+ node.plan = if !sort_exec.preserve_partitioning()
+ && sort_input.output_partitioning().partition_count() > 1
+ {
+ // Replace the sort with a sort-preserving merge:
+ let expr = sort_exec.expr().to_vec();
+ Arc::new(SortPreservingMergeExec::new(expr, sort_input.clone())) as _
+ } else {
+ // Remove the sort:
+ node.children = node.children.swap_remove(0).children;
+ sort_input.clone()
+ };
+ for child in node.children.iter_mut() {
+ child.data = false;
+ }
+ node.data = false;
+ return Transformed::Yes(node);
}
}
- None
+ Transformed::No(node)
}
-/// Analyzes a [`WindowAggExec`] or a [`BoundedWindowAggExec`] to determine
+/// Adjusts a [`WindowAggExec`] or a [`BoundedWindowAggExec`] to determine
/// whether it may allow removing a sort.
-fn analyze_window_sort_removal(
- sort_tree: &mut PlanWithCorrespondingSort,
- window_exec: &Arc<dyn ExecutionPlan>,
-) -> Result<Option<PlanWithCorrespondingSort>> {
- let requires_single_partition = matches!(
- window_exec.required_input_distribution()[0],
- Distribution::SinglePartition
- );
- remove_corresponding_sort_from_sub_plan(sort_tree, requires_single_partition)?;
- let mut window_child = sort_tree.plan.clone();
+fn adjust_window_sort_removal(
+ mut window_tree: PlanWithCorrespondingSort,
+) -> Result<PlanWithCorrespondingSort> {
+ // Window operators have a single child we need to adjust:
+ let child_node = remove_corresponding_sort_from_sub_plan(
+ window_tree.children.swap_remove(0),
+ matches!(
+ window_tree.plan.required_input_distribution()[0],
+ Distribution::SinglePartition
+ ),
+ )?;
+ window_tree.children.push(child_node);
+
+ let plan = window_tree.plan.as_any();
+ let child_plan = &window_tree.children[0].plan;
let (window_expr, new_window) =
- if let Some(exec) = window_exec.as_any().downcast_ref::<BoundedWindowAggExec>() {
- (
- exec.window_expr(),
- get_best_fitting_window(
- exec.window_expr(),
- &window_child,
- &exec.partition_keys,
- )?,
- )
- } else if let Some(exec) = window_exec.as_any().downcast_ref::<WindowAggExec>() {
- (
- exec.window_expr(),
- get_best_fitting_window(
- exec.window_expr(),
- &window_child,
- &exec.partition_keys,
- )?,
- )
+ if let Some(exec) = plan.downcast_ref::<WindowAggExec>() {
+ let window_expr = exec.window_expr();
+ let new_window =
+ get_best_fitting_window(window_expr, child_plan, &exec.partition_keys)?;
+ (window_expr, new_window)
+ } else if let Some(exec) = plan.downcast_ref::<BoundedWindowAggExec>() {
+ let window_expr = exec.window_expr();
+ let new_window =
+ get_best_fitting_window(window_expr, child_plan, &exec.partition_keys)?;
+ (window_expr, new_window)
} else {
- return plan_err!(
- "Expects to receive either WindowAggExec of BoundedWindowAggExec"
- );
+ return plan_err!("Expected WindowAggExec or BoundedWindowAggExec");
};
- let partitionby_exprs = window_expr[0].partition_by();
- if let Some(new_window) = new_window {
+ window_tree.plan = if let Some(new_window) = new_window {
// We were able to change the window to accommodate the input, use it:
- Ok(Some(PlanWithCorrespondingSort::new(new_window)))
+ new_window
} else {
// We were unable to change the window to accommodate the input, so we
// will insert a sort.
- let reqs = window_exec
+ let reqs = window_tree
+ .plan
.required_input_ordering()
.swap_remove(0)
.unwrap_or_default();
+
// Satisfy the ordering requirement so that the window can run:
- add_sort_above(&mut window_child, &reqs, None);
+ let mut child_node = window_tree.children.swap_remove(0);
+ child_node = add_sort_above(child_node, reqs, None);
+ let child_plan = child_node.plan.clone();
+ window_tree.children.push(child_node);
- let uses_bounded_memory = window_expr.iter().all(|e| e.uses_bounded_memory());
- let new_window = if uses_bounded_memory {
+ if window_expr.iter().all(|e| e.uses_bounded_memory()) {
Arc::new(BoundedWindowAggExec::try_new(
window_expr.to_vec(),
- window_child,
- partitionby_exprs.to_vec(),
+ child_plan,
+ window_expr[0].partition_by().to_vec(),
InputOrderMode::Sorted,
)?) as _
} else {
Arc::new(WindowAggExec::try_new(
window_expr.to_vec(),
- window_child,
- partitionby_exprs.to_vec(),
+ child_plan,
+ window_expr[0].partition_by().to_vec(),
)?) as _
- };
- Ok(Some(PlanWithCorrespondingSort::new(new_window)))
- }
-}
+ }
+ };
-/// Updates child to remove the unnecessary [`CoalescePartitionsExec`] below it.
-fn update_child_to_remove_coalesce(
- child: &mut Arc<dyn ExecutionPlan>,
- coalesce_onwards: &mut PlanWithCorrespondingCoalescePartitions,
-) -> Result<()> {
- if coalesce_onwards.coalesce_connection {
- *child = remove_corresponding_coalesce_in_sub_plan(coalesce_onwards, child)?;
- }
- Ok(())
+ window_tree.data = false;
+ Ok(window_tree)
}
-/// Removes the [`CoalescePartitionsExec`] from the plan in `coalesce_onwards`.
+/// Removes the [`CoalescePartitionsExec`] from the plan in `node`.
fn remove_corresponding_coalesce_in_sub_plan(
- coalesce_onwards: &mut PlanWithCorrespondingCoalescePartitions,
- parent: &Arc<dyn ExecutionPlan>,
-) -> Result<Arc<dyn ExecutionPlan>> {
- if is_coalesce_partitions(&coalesce_onwards.plan) {
+ mut requirements: PlanWithCorrespondingCoalescePartitions,
+) -> Result<PlanWithCorrespondingCoalescePartitions> {
+ let plan = &requirements.plan;
+ let children = &mut requirements.children;
+ if is_coalesce_partitions(&children[0].plan) {
// We can safely use the 0th index since we have a `CoalescePartitionsExec`.
- let mut new_plan = coalesce_onwards.plan.children()[0].clone();
- while new_plan.output_partitioning() == parent.output_partitioning()
- && is_repartition(&new_plan)
- && is_repartition(parent)
+ let mut new_child_node = children[0].children.swap_remove(0);
+ while new_child_node.plan.output_partitioning() == plan.output_partitioning()
+ && is_repartition(&new_child_node.plan)
+ && is_repartition(plan)
{
- new_plan = new_plan.children().swap_remove(0)
+ new_child_node = new_child_node.children.swap_remove(0)
}
- Ok(new_plan)
+ children[0] = new_child_node;
} else {
- let plan = coalesce_onwards.plan.clone();
- let mut children = plan.children();
- for (idx, node) in coalesce_onwards.children_nodes.iter_mut().enumerate() {
- if node.coalesce_connection {
- children[idx] = remove_corresponding_coalesce_in_sub_plan(node, &plan)?;
- }
- }
- plan.with_new_children(children)
+ requirements.children = requirements
+ .children
+ .into_iter()
+ .map(|node| {
+ if node.data {
+ remove_corresponding_coalesce_in_sub_plan(node)
+ } else {
+ Ok(node)
+ }
+ })
+ .collect::<Result<_>>()?;
}
+
+ requirements.update_plan_from_children()
}
/// Updates child to remove the unnecessary sort below it.
fn update_child_to_remove_unnecessary_sort(
child_idx: usize,
- sort_onwards: &mut PlanWithCorrespondingSort,
+ mut node: PlanWithCorrespondingSort,
parent: &Arc<dyn ExecutionPlan>,
-) -> Result<()> {
- if sort_onwards.sort_connection {
+) -> Result<PlanWithCorrespondingSort> {
+ if node.data {
let requires_single_partition = matches!(
parent.required_input_distribution()[child_idx],
Distribution::SinglePartition
);
- remove_corresponding_sort_from_sub_plan(sort_onwards, requires_single_partition)?;
+ node = remove_corresponding_sort_from_sub_plan(node, requires_single_partition)?;
}
- sort_onwards.sort_connection = false;
- Ok(())
+ node.data = false;
+ Ok(node)
}
-/// Removes the sort from the plan in `sort_onwards`.
+/// Removes the sort from the plan in `node`.
fn remove_corresponding_sort_from_sub_plan(
- sort_onwards: &mut PlanWithCorrespondingSort,
+ mut node: PlanWithCorrespondingSort,
requires_single_partition: bool,
-) -> Result<()> {
+) -> Result<PlanWithCorrespondingSort> {
// A `SortExec` is always at the bottom of the tree.
- if is_sort(&sort_onwards.plan) {
- *sort_onwards = sort_onwards.children_nodes.swap_remove(0);
+ if is_sort(&node.plan) {
+ node = node.children.swap_remove(0);
} else {
- let PlanWithCorrespondingSort {
- plan,
- sort_connection: _,
- children_nodes,
- } = sort_onwards;
let mut any_connection = false;
- for (child_idx, child_node) in children_nodes.iter_mut().enumerate() {
- if child_node.sort_connection {
- any_connection = true;
- let requires_single_partition = matches!(
- plan.required_input_distribution()[child_idx],
- Distribution::SinglePartition
- );
- remove_corresponding_sort_from_sub_plan(
- child_node,
- requires_single_partition,
- )?;
- }
- }
- if any_connection || children_nodes.is_empty() {
- *sort_onwards = PlanWithCorrespondingSort::update_children(
- plan.clone(),
- children_nodes.clone(),
- )?;
+ let required_dist = node.plan.required_input_distribution();
+ node.children = node
+ .children
+ .into_iter()
+ .enumerate()
+ .map(|(idx, child)| {
+ if child.data {
+ any_connection = true;
+ remove_corresponding_sort_from_sub_plan(
+ child,
+ matches!(required_dist[idx], Distribution::SinglePartition),
+ )
+ } else {
+ Ok(child)
+ }
+ })
+ .collect::<Result<_>>()?;
+ if any_connection || node.children.is_empty() {
+ node = update_sort_ctx_children(node, false)?;
}
- let PlanWithCorrespondingSort {
- plan,
- children_nodes,
- ..
- } = sort_onwards;
+
// Replace with variants that do not preserve order.
- if is_sort_preserving_merge(plan) {
- children_nodes.swap_remove(0);
- *plan = plan.children().swap_remove(0);
- } else if let Some(repartition) = plan.as_any().downcast_ref::<RepartitionExec>()
+ if is_sort_preserving_merge(&node.plan) {
+ node.children = node.children.swap_remove(0).children;
+ node.plan = node.plan.children().swap_remove(0);
+ } else if let Some(repartition) =
+ node.plan.as_any().downcast_ref::<RepartitionExec>()
{
- *plan = Arc::new(RepartitionExec::try_new(
- children_nodes[0].plan.clone(),
+ node.plan = Arc::new(RepartitionExec::try_new(
+ node.children[0].plan.clone(),
repartition.output_partitioning(),
)?) as _;
}
};
// Deleting a merging sort may invalidate distribution requirements.
// Ensure that we stay compliant with such requirements:
- if requires_single_partition
- && sort_onwards.plan.output_partitioning().partition_count() > 1
+ if requires_single_partition && node.plan.output_partitioning().partition_count() > 1
{
- // If there is existing ordering, to preserve ordering use SortPreservingMergeExec
- // instead of CoalescePartitionsExec.
- if let Some(ordering) = sort_onwards.plan.output_ordering() {
- let plan = Arc::new(SortPreservingMergeExec::new(
- ordering.to_vec(),
- sort_onwards.plan.clone(),
- )) as _;
- *sort_onwards = PlanWithCorrespondingSort::update_children(
- plan,
- vec![sort_onwards.clone()],
- )?;
+ // If there is existing ordering, to preserve ordering use
+ // `SortPreservingMergeExec` instead of a `CoalescePartitionsExec`.
+ let plan = node.plan.clone();
+ let plan = if let Some(ordering) = plan.output_ordering() {
+ Arc::new(SortPreservingMergeExec::new(ordering.to_vec(), plan)) as _
} else {
- let plan =
- Arc::new(CoalescePartitionsExec::new(sort_onwards.plan.clone())) as _;
- *sort_onwards = PlanWithCorrespondingSort::update_children(
- plan,
- vec![sort_onwards.clone()],
- )?;
- }
+ Arc::new(CoalescePartitionsExec::new(plan)) as _
+ };
+ node = PlanWithCorrespondingSort::new(plan, false, vec![node]);
+ node = update_sort_ctx_children(node, false)?;
}
- Ok(())
+ Ok(node)
}
/// Converts an [ExecutionPlan] trait object to a [PhysicalSortExpr] slice when possible.
@@ -734,13 +561,9 @@ fn get_sort_exprs(
) -> Result<(&[PhysicalSortExpr], Option<usize>)> {
if let Some(sort_exec) = sort_any.as_any().downcast_ref::<SortExec>() {
Ok((sort_exec.expr(), sort_exec.fetch()))
- } else if let Some(sort_preserving_merge_exec) =
- sort_any.as_any().downcast_ref::<SortPreservingMergeExec>()
+ } else if let Some(spm) = sort_any.as_any().downcast_ref::<SortPreservingMergeExec>()
{
- Ok((
- sort_preserving_merge_exec.expr(),
- sort_preserving_merge_exec.fetch(),
- ))
+ Ok((spm.expr(), spm.fetch()))
} else {
plan_err!("Given ExecutionPlan is not a SortExec or a SortPreservingMergeExec")
}
@@ -753,7 +576,7 @@ mod tests {
use super::*;
use crate::physical_optimizer::enforce_distribution::EnforceDistribution;
use crate::physical_optimizer::test_utils::{
- aggregate_exec, bounded_window_exec, coalesce_batches_exec,
+ aggregate_exec, bounded_window_exec, check_integrity, coalesce_batches_exec,
coalesce_partitions_exec, filter_exec, global_limit_exec, hash_join_exec,
limit_exec, local_limit_exec, memory_exec, parquet_exec, parquet_exec_sorted,
repartition_exec, sort_exec, sort_expr, sort_expr_options, sort_merge_join_exec,
@@ -776,7 +599,6 @@ mod tests {
let nullable_column = Field::new("nullable_col", DataType::Int32, true);
let non_nullable_column = Field::new("non_nullable_col", DataType::Int32, false);
let schema = Arc::new(Schema::new(vec![nullable_column, non_nullable_column]));
-
Ok(schema)
}
@@ -812,6 +634,50 @@ mod tests {
let session_ctx = SessionContext::new_with_config(config);
let state = session_ctx.state();
+ // This file has 4 rules that use tree node, apply these rules as in the
+ // EnforSorting::optimize implementation
+ // After these operations tree nodes should be in a consistent state.
+ // This code block makes sure that these rules doesn't violate tree node integrity.
+ {
+ let plan_requirements = PlanWithCorrespondingSort::new_default($PLAN.clone());
+ let adjusted = plan_requirements
+ .transform_up(&ensure_sorting)
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+
+ let new_plan = if state.config_options().optimizer.repartition_sorts {
+ let plan_with_coalesce_partitions =
+ PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ let parallel = plan_with_coalesce_partitions
+ .transform_up(¶llelize_sorts)
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+ parallel.plan
+ } else {
+ adjusted.plan
+ };
+
+ let plan_with_pipeline_fixer = OrderPreservationContext::new_default(new_plan);
+ let updated_plan = plan_with_pipeline_fixer
+ .transform_up(&|plan_with_pipeline_fixer| {
+ replace_with_order_preserving_variants(
+ plan_with_pipeline_fixer,
+ false,
+ true,
+ state.config_options(),
+ )
+ })
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+
+ let mut sort_pushdown = SortPushDown::new_default(updated_plan.plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ sort_pushdown
+ .transform_down(&pushdown_sorts)
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+ }
+
let physical_plan = $PLAN;
let formatted = displayable(physical_plan.as_ref()).indent(true).to_string();
let actual: Vec<&str> = formatted.trim().lines().collect();
@@ -858,6 +724,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -917,6 +784,7 @@ mod tests {
" SortExec: expr=[non_nullable_col@1 DESC]",
" MemoryExec: partitions=1, partition_sizes=[0]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -938,6 +806,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -964,6 +833,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1003,6 +873,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1047,6 +918,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1100,6 +972,7 @@ mod tests {
" RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
" MemoryExec: partitions=1, partition_sizes=[0]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1127,6 +1000,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1156,6 +1030,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1198,6 +1073,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC,non_nullable_col@1 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1221,6 +1097,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1248,6 +1125,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1275,6 +1153,7 @@ mod tests {
// should not add a sort at the output of the union, input plan should not be changed
let expected_optimized = expected_input.clone();
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1306,6 +1185,7 @@ mod tests {
// should not add a sort at the output of the union, input plan should not be changed
let expected_optimized = expected_input.clone();
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1342,6 +1222,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC,non_nullable_col@1 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1383,6 +1264,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1424,6 +1306,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC,non_nullable_col@1 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1470,6 +1353,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1518,6 +1402,7 @@ mod tests {
" RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1552,6 +1437,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_output, physical_plan, true);
+
Ok(())
}
@@ -1600,6 +1486,7 @@ mod tests {
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1649,6 +1536,7 @@ mod tests {
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1685,6 +1573,7 @@ mod tests {
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1736,6 +1625,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC,non_nullable_col@1 DESC NULLS LAST]",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -1979,6 +1869,7 @@ mod tests {
" SortExec: expr=[nullable_col@0 ASC,non_nullable_col@1 ASC]",
" MemoryExec: partitions=1, partition_sizes=[0]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -2013,6 +1904,7 @@ mod tests {
" RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
" ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col]"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -2105,6 +1997,7 @@ mod tests {
" MemoryExec: partitions=1, partition_sizes=[0]",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -2132,6 +2025,7 @@ mod tests {
" RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
" CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], has_header=false"];
assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
Ok(())
}
@@ -2221,6 +2115,7 @@ mod tests {
physical_plan,
true
);
+
Ok(())
}
@@ -2242,6 +2137,7 @@ mod tests {
" RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
" CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], has_header=false",];
assert_optimized!(expected_input, expected_optimized, physical_plan, false);
+
Ok(())
}
@@ -2270,6 +2166,7 @@ mod tests {
" RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
" CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], has_header=false",];
assert_optimized!(expected_input, expected_optimized, physical_plan, false);
+
Ok(())
}
@@ -2302,6 +2199,7 @@ mod tests {
" CsvExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], has_header=false",
];
assert_optimized!(expected_input, expected_optimized, physical_plan, false);
+
Ok(())
}
}
diff --git a/datafusion/core/src/physical_optimizer/join_selection.rs b/datafusion/core/src/physical_optimizer/join_selection.rs
index f9b9fdf85c..083cd5ecab 100644
--- a/datafusion/core/src/physical_optimizer/join_selection.rs
+++ b/datafusion/core/src/physical_optimizer/join_selection.rs
@@ -27,7 +27,9 @@ use std::sync::Arc;
use crate::config::ConfigOptions;
use crate::error::Result;
-use crate::physical_optimizer::pipeline_checker::PipelineStatePropagator;
+use crate::physical_optimizer::pipeline_checker::{
+ children_unbounded, PipelineStatePropagator,
+};
use crate::physical_optimizer::PhysicalOptimizerRule;
use crate::physical_plan::joins::utils::{ColumnIndex, JoinFilter};
use crate::physical_plan::joins::{
@@ -229,7 +231,7 @@ impl PhysicalOptimizerRule for JoinSelection {
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
- let pipeline = PipelineStatePropagator::new(plan);
+ let pipeline = PipelineStatePropagator::new_default(plan);
// First, we make pipeline-fixing modifications to joins so as to accommodate
// unbounded inputs. Each pipeline-fixing subrule, which is a function
// of type `PipelineFixerSubrule`, takes a single [`PipelineStatePropagator`]
@@ -352,7 +354,7 @@ fn try_collect_left(
)?))),
(false, true) => {
if supports_swap(*hash_join.join_type()) {
- Ok(Some(swap_hash_join(hash_join, PartitionMode::CollectLeft)?))
+ swap_hash_join(hash_join, PartitionMode::CollectLeft).map(Some)
} else {
Ok(None)
}
@@ -443,10 +445,8 @@ fn statistical_join_selection_subrule(
}
/// Pipeline-fixing join selection subrule.
-pub type PipelineFixerSubrule = dyn Fn(
- PipelineStatePropagator,
- &ConfigOptions,
-) -> Option<Result<PipelineStatePropagator>>;
+pub type PipelineFixerSubrule =
+ dyn Fn(PipelineStatePropagator, &ConfigOptions) -> Result<PipelineStatePropagator>;
/// Converts a hash join to a symmetric hash join in the case of infinite inputs on both sides.
///
@@ -466,16 +466,16 @@ pub type PipelineFixerSubrule = dyn Fn(
fn hash_join_convert_symmetric_subrule(
mut input: PipelineStatePropagator,
config_options: &ConfigOptions,
-) -> Option<Result<PipelineStatePropagator>> {
+) -> Result<PipelineStatePropagator> {
// Check if the current plan node is a HashJoinExec.
if let Some(hash_join) = input.plan.as_any().downcast_ref::<HashJoinExec>() {
// Determine if left and right children are unbounded.
- let ub_flags = input.children_unbounded();
+ let ub_flags = children_unbounded(&input);
let (left_unbounded, right_unbounded) = (ub_flags[0], ub_flags[1]);
// Update the unbounded flag of the input.
- input.unbounded = left_unbounded || right_unbounded;
+ input.data = left_unbounded || right_unbounded;
// Process only if both left and right sides are unbounded.
- let result = if left_unbounded && right_unbounded {
+ if left_unbounded && right_unbounded {
// Determine the partition mode based on configuration.
let mode = if config_options.optimizer.repartition_joins {
StreamJoinPartitionMode::Partitioned
@@ -519,7 +519,7 @@ fn hash_join_convert_symmetric_subrule(
let name = schema.field(*index).name();
let col = Arc::new(Column::new(name, *index)) as _;
// Check if the column is ordered.
- equivalence.get_expr_ordering(col).state
+ equivalence.get_expr_ordering(col).data
!= SortProperties::Unordered
},
)
@@ -539,7 +539,7 @@ fn hash_join_convert_symmetric_subrule(
let left_order = determine_order(JoinSide::Left);
let right_order = determine_order(JoinSide::Right);
- SymmetricHashJoinExec::try_new(
+ return SymmetricHashJoinExec::try_new(
hash_join.left().clone(),
hash_join.right().clone(),
hash_join.on().to_vec(),
@@ -553,14 +553,10 @@ fn hash_join_convert_symmetric_subrule(
.map(|exec| {
input.plan = Arc::new(exec) as _;
input
- })
- } else {
- Ok(input)
- };
- Some(result)
- } else {
- None
+ });
+ }
}
+ Ok(input)
}
/// This subrule will swap build/probe sides of a hash join depending on whether
@@ -607,12 +603,12 @@ fn hash_join_convert_symmetric_subrule(
fn hash_join_swap_subrule(
mut input: PipelineStatePropagator,
_config_options: &ConfigOptions,
-) -> Option<Result<PipelineStatePropagator>> {
+) -> Result<PipelineStatePropagator> {
if let Some(hash_join) = input.plan.as_any().downcast_ref::<HashJoinExec>() {
- let ub_flags = input.children_unbounded();
+ let ub_flags = children_unbounded(&input);
let (left_unbounded, right_unbounded) = (ub_flags[0], ub_flags[1]);
- input.unbounded = left_unbounded || right_unbounded;
- let result = if left_unbounded
+ input.data = left_unbounded || right_unbounded;
+ if left_unbounded
&& !right_unbounded
&& matches!(
*hash_join.join_type(),
@@ -620,18 +616,12 @@ fn hash_join_swap_subrule(
| JoinType::Left
| JoinType::LeftSemi
| JoinType::LeftAnti
- ) {
- swap_join_according_to_unboundedness(hash_join).map(|plan| {
- input.plan = plan;
- input
- })
- } else {
- Ok(input)
- };
- Some(result)
- } else {
- None
+ )
+ {
+ input.plan = swap_join_according_to_unboundedness(hash_join)?;
+ }
}
+ Ok(input)
}
/// This function swaps sides of a hash join to make it runnable even if one of
@@ -669,13 +659,11 @@ fn apply_subrules(
config_options: &ConfigOptions,
) -> Result<Transformed<PipelineStatePropagator>> {
for subrule in subrules {
- if let Some(value) = subrule(input.clone(), config_options).transpose()? {
- input = value;
- }
+ input = subrule(input, config_options)?;
}
- let is_unbounded = input
+ input.data = input
.plan
- .unbounded_output(&input.children_unbounded())
+ .unbounded_output(&children_unbounded(&input))
// Treat the case where an operator can not run on unbounded data as
// if it can and it outputs unbounded data. Do not raise an error yet.
// Such operators may be fixed, adjusted or replaced later on during
@@ -683,7 +671,6 @@ fn apply_subrules(
// etc. If this doesn't happen, the final `PipelineChecker` rule will
// catch this and raise an error anyway.
.unwrap_or(true);
- input.unbounded = is_unbounded;
Ok(Transformed::Yes(input))
}
@@ -693,6 +680,7 @@ mod tests_statistical {
use super::*;
use crate::{
+ physical_optimizer::test_utils::check_integrity,
physical_plan::{
displayable, joins::PartitionMode, ColumnStatistics, Statistics,
},
@@ -840,26 +828,51 @@ mod tests_statistical {
(big, medium, small)
}
+ pub(crate) fn crosscheck_plans(plan: Arc<dyn ExecutionPlan>) -> Result<()> {
+ let pipeline = PipelineStatePropagator::new_default(plan);
+ let subrules: Vec<Box<PipelineFixerSubrule>> = vec![
+ Box::new(hash_join_convert_symmetric_subrule),
+ Box::new(hash_join_swap_subrule),
+ ];
+ let state = pipeline
+ .transform_up(&|p| apply_subrules(p, &subrules, &ConfigOptions::new()))
+ .and_then(check_integrity)?;
+ // TODO: End state payloads will be checked here.
+ let config = ConfigOptions::new().optimizer;
+ let collect_left_threshold = config.hash_join_single_partition_threshold;
+ let collect_threshold_num_rows = config.hash_join_single_partition_threshold_rows;
+ let _ = state.plan.transform_up(&|plan| {
+ statistical_join_selection_subrule(
+ plan,
+ collect_left_threshold,
+ collect_threshold_num_rows,
+ )
+ })?;
+ Ok(())
+ }
+
#[tokio::test]
async fn test_join_with_swap() {
let (big, small) = create_big_and_small();
- let join = HashJoinExec::try_new(
- Arc::clone(&big),
- Arc::clone(&small),
- vec![(
- Column::new_with_schema("big_col", &big.schema()).unwrap(),
- Column::new_with_schema("small_col", &small.schema()).unwrap(),
- )],
- None,
- &JoinType::Left,
- PartitionMode::CollectLeft,
- false,
- )
- .unwrap();
+ let join = Arc::new(
+ HashJoinExec::try_new(
+ Arc::clone(&big),
+ Arc::clone(&small),
+ vec![(
+ Column::new_with_schema("big_col", &big.schema()).unwrap(),
+ Column::new_with_schema("small_col", &small.schema()).unwrap(),
+ )],
+ None,
+ &JoinType::Left,
+ PartitionMode::CollectLeft,
+ false,
+ )
+ .unwrap(),
+ );
let optimized_join = JoinSelection::new()
- .optimize(Arc::new(join), &ConfigOptions::new())
+ .optimize(join.clone(), &ConfigOptions::new())
.unwrap();
let swapping_projection = optimized_join
@@ -889,28 +902,31 @@ mod tests_statistical {
swapped_join.right().statistics().unwrap().total_byte_size,
Precision::Inexact(2097152)
);
+ crosscheck_plans(join.clone()).unwrap();
}
#[tokio::test]
async fn test_left_join_with_swap() {
let (big, small) = create_big_and_small();
// Left out join should alway swap when the mode is PartitionMode::CollectLeft, even left side is small and right side is large
- let join = HashJoinExec::try_new(
- Arc::clone(&small),
- Arc::clone(&big),
- vec![(
- Column::new_with_schema("small_col", &small.schema()).unwrap(),
- Column::new_with_schema("big_col", &big.schema()).unwrap(),
- )],
- None,
- &JoinType::Left,
- PartitionMode::CollectLeft,
- false,
- )
- .unwrap();
+ let join = Arc::new(
+ HashJoinExec::try_new(
+ Arc::clone(&small),
+ Arc::clone(&big),
+ vec![(
+ Column::new_with_schema("small_col", &small.schema()).unwrap(),
+ Column::new_with_schema("big_col", &big.schema()).unwrap(),
+ )],
+ None,
+ &JoinType::Left,
+ PartitionMode::CollectLeft,
+ false,
+ )
+ .unwrap(),
+ );
let optimized_join = JoinSelection::new()
- .optimize(Arc::new(join), &ConfigOptions::new())
+ .optimize(join.clone(), &ConfigOptions::new())
.unwrap();
let swapping_projection = optimized_join
@@ -940,6 +956,7 @@ mod tests_statistical {
swapped_join.right().statistics().unwrap().total_byte_size,
Precision::Inexact(8192)
);
+ crosscheck_plans(join.clone()).unwrap();
}
#[tokio::test]
@@ -948,24 +965,26 @@ mod tests_statistical {
for join_type in join_types {
let (big, small) = create_big_and_small();
- let join = HashJoinExec::try_new(
- Arc::clone(&big),
- Arc::clone(&small),
- vec![(
- Column::new_with_schema("big_col", &big.schema()).unwrap(),
- Column::new_with_schema("small_col", &small.schema()).unwrap(),
- )],
- None,
- &join_type,
- PartitionMode::Partitioned,
- false,
- )
- .unwrap();
+ let join = Arc::new(
+ HashJoinExec::try_new(
+ Arc::clone(&big),
+ Arc::clone(&small),
+ vec![(
+ Column::new_with_schema("big_col", &big.schema()).unwrap(),
+ Column::new_with_schema("small_col", &small.schema()).unwrap(),
+ )],
+ None,
+ &join_type,
+ PartitionMode::Partitioned,
+ false,
+ )
+ .unwrap(),
+ );
let original_schema = join.schema();
let optimized_join = JoinSelection::new()
- .optimize(Arc::new(join), &ConfigOptions::new())
+ .optimize(join.clone(), &ConfigOptions::new())
.unwrap();
let swapped_join = optimized_join
@@ -976,7 +995,6 @@ mod tests_statistical {
);
assert_eq!(swapped_join.schema().fields().len(), 1);
-
assert_eq!(
swapped_join.left().statistics().unwrap().total_byte_size,
Precision::Inexact(8192)
@@ -985,8 +1003,8 @@ mod tests_statistical {
swapped_join.right().statistics().unwrap().total_byte_size,
Precision::Inexact(2097152)
);
-
assert_eq!(original_schema, swapped_join.schema());
+ crosscheck_plans(join).unwrap();
}
}
@@ -996,18 +1014,20 @@ mod tests_statistical {
let expected_lines =
$EXPECTED_LINES.iter().map(|s| *s).collect::<Vec<&str>>();
+ let plan = Arc::new($PLAN);
let optimized = JoinSelection::new()
- .optimize(Arc::new($PLAN), &ConfigOptions::new())
+ .optimize(plan.clone(), &ConfigOptions::new())
.unwrap();
- let plan = displayable(optimized.as_ref()).indent(true).to_string();
- let actual_lines = plan.split("\n").collect::<Vec<&str>>();
+ let plan_string = displayable(optimized.as_ref()).indent(true).to_string();
+ let actual_lines = plan_string.split("\n").collect::<Vec<&str>>();
assert_eq!(
&expected_lines, &actual_lines,
"\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
expected_lines, actual_lines
);
+ crosscheck_plans(plan).unwrap();
};
}
@@ -1069,22 +1089,24 @@ mod tests_statistical {
#[tokio::test]
async fn test_join_no_swap() {
let (big, small) = create_big_and_small();
- let join = HashJoinExec::try_new(
- Arc::clone(&small),
- Arc::clone(&big),
- vec![(
- Column::new_with_schema("small_col", &small.schema()).unwrap(),
- Column::new_with_schema("big_col", &big.schema()).unwrap(),
- )],
- None,
- &JoinType::Inner,
- PartitionMode::CollectLeft,
- false,
- )
- .unwrap();
+ let join = Arc::new(
+ HashJoinExec::try_new(
+ Arc::clone(&small),
+ Arc::clone(&big),
+ vec![(
+ Column::new_with_schema("small_col", &small.schema()).unwrap(),
+ Column::new_with_schema("big_col", &big.schema()).unwrap(),
+ )],
+ None,
+ &JoinType::Inner,
+ PartitionMode::CollectLeft,
+ false,
+ )
+ .unwrap(),
+ );
let optimized_join = JoinSelection::new()
- .optimize(Arc::new(join), &ConfigOptions::new())
+ .optimize(join.clone(), &ConfigOptions::new())
.unwrap();
let swapped_join = optimized_join
@@ -1100,6 +1122,7 @@ mod tests_statistical {
swapped_join.right().statistics().unwrap().total_byte_size,
Precision::Inexact(2097152)
);
+ crosscheck_plans(join).unwrap();
}
#[tokio::test]
@@ -1270,19 +1293,21 @@ mod tests_statistical {
is_swapped: bool,
expected_mode: PartitionMode,
) {
- let join = HashJoinExec::try_new(
- left,
- right,
- on,
- None,
- &JoinType::Inner,
- PartitionMode::Auto,
- false,
- )
- .unwrap();
+ let join = Arc::new(
+ HashJoinExec::try_new(
+ left,
+ right,
+ on,
+ None,
+ &JoinType::Inner,
+ PartitionMode::Auto,
+ false,
+ )
+ .unwrap(),
+ );
let optimized_join = JoinSelection::new()
- .optimize(Arc::new(join), &ConfigOptions::new())
+ .optimize(join.clone(), &ConfigOptions::new())
.unwrap();
if !is_swapped {
@@ -1306,6 +1331,7 @@ mod tests_statistical {
assert_eq!(*swapped_join.partition_mode(), expected_mode);
}
+ crosscheck_plans(join).unwrap();
}
}
@@ -1350,6 +1376,8 @@ mod util_tests {
#[cfg(test)]
mod hash_join_tests {
+ use self::tests_statistical::crosscheck_plans;
+
use super::*;
use crate::physical_optimizer::join_selection::swap_join_type;
use crate::physical_optimizer::test_utils::SourceType;
@@ -1716,7 +1744,7 @@ mod hash_join_tests {
2,
)) as Arc<dyn ExecutionPlan>;
- let join = HashJoinExec::try_new(
+ let join = Arc::new(HashJoinExec::try_new(
Arc::clone(&left_exec),
Arc::clone(&right_exec),
vec![(
@@ -1727,29 +1755,19 @@ mod hash_join_tests {
&t.initial_join_type,
t.initial_mode,
false,
- )?;
+ )?);
+ let left_child = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
+ let right_child = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
let children = vec![
- PipelineStatePropagator {
- plan: Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
- unbounded: left_unbounded,
- children: vec![],
- },
- PipelineStatePropagator {
- plan: Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
- unbounded: right_unbounded,
- children: vec![],
- },
+ PipelineStatePropagator::new(left_child, left_unbounded, vec![]),
+ PipelineStatePropagator::new(right_child, right_unbounded, vec![]),
];
- let initial_hash_join_state = PipelineStatePropagator {
- plan: Arc::new(join),
- unbounded: false,
- children,
- };
+ let initial_hash_join_state =
+ PipelineStatePropagator::new(join.clone(), false, children);
let optimized_hash_join =
- hash_join_swap_subrule(initial_hash_join_state, &ConfigOptions::new())
- .unwrap()?;
+ hash_join_swap_subrule(initial_hash_join_state, &ConfigOptions::new())?;
let optimized_join_plan = optimized_hash_join.plan;
// If swap did happen
@@ -1805,6 +1823,7 @@ mod hash_join_tests {
)
);
};
+ crosscheck_plans(plan).unwrap();
Ok(())
}
}
diff --git a/datafusion/core/src/physical_optimizer/pipeline_checker.rs b/datafusion/core/src/physical_optimizer/pipeline_checker.rs
index e281d0e7c2..bb0665c10b 100644
--- a/datafusion/core/src/physical_optimizer/pipeline_checker.rs
+++ b/datafusion/core/src/physical_optimizer/pipeline_checker.rs
@@ -19,19 +19,19 @@
//! infinite sources, if there are any. It will reject non-runnable query plans
//! that use pipeline-breaking operators on infinite input(s).
-use std::borrow::Cow;
use std::sync::Arc;
use crate::config::ConfigOptions;
use crate::error::Result;
use crate::physical_optimizer::PhysicalOptimizerRule;
-use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
+use crate::physical_plan::ExecutionPlan;
use datafusion_common::config::OptimizerOptions;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::{plan_err, DataFusionError};
use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported};
use datafusion_physical_plan::joins::SymmetricHashJoinExec;
+use datafusion_physical_plan::tree_node::PlanContext;
/// The PipelineChecker rule rejects non-runnable query plans that use
/// pipeline-breaking operators on infinite input(s).
@@ -51,7 +51,7 @@ impl PhysicalOptimizerRule for PipelineChecker {
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
- let pipeline = PipelineStatePropagator::new(plan);
+ let pipeline = PipelineStatePropagator::new_default(plan);
let state = pipeline
.transform_up(&|p| check_finiteness_requirements(p, &config.optimizer))?;
Ok(state.plan)
@@ -66,54 +66,12 @@ impl PhysicalOptimizerRule for PipelineChecker {
}
}
-/// [PipelineStatePropagator] propagates the [ExecutionPlan] pipelining information.
-#[derive(Clone, Debug)]
-pub struct PipelineStatePropagator {
- pub(crate) plan: Arc<dyn ExecutionPlan>,
- pub(crate) unbounded: bool,
- pub(crate) children: Vec<Self>,
-}
-
-impl PipelineStatePropagator {
- /// Constructs a new, default pipelining state.
- pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- unbounded: false,
- children: children.into_iter().map(Self::new).collect(),
- }
- }
+/// This object propagates the [`ExecutionPlan`] pipelining information.
+pub type PipelineStatePropagator = PlanContext<bool>;
- /// Returns the children unboundedness information.
- pub fn children_unbounded(&self) -> Vec<bool> {
- self.children.iter().map(|c| c.unbounded).collect()
- }
-}
-
-impl TreeNode for PipelineStatePropagator {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children.is_empty() {
- self.children = self
- .children
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- }
- Ok(self)
- }
+/// Collects unboundedness flags of all the children of the plan in `pipeline`.
+pub fn children_unbounded(pipeline: &PipelineStatePropagator) -> Vec<bool> {
+ pipeline.children.iter().map(|c| c.data).collect()
}
/// This function propagates finiteness information and rejects any plan with
@@ -126,16 +84,15 @@ pub fn check_finiteness_requirements(
if !(optimizer_options.allow_symmetric_joins_without_pruning
|| (exec.check_if_order_information_available()? && is_prunable(exec)))
{
- const MSG: &str = "Join operation cannot operate on a non-prunable stream without enabling \
- the 'allow_symmetric_joins_without_pruning' configuration flag";
- return plan_err!("{}", MSG);
+ return plan_err!("Join operation cannot operate on a non-prunable stream without enabling \
+ the 'allow_symmetric_joins_without_pruning' configuration flag");
}
}
input
.plan
- .unbounded_output(&input.children_unbounded())
+ .unbounded_output(&children_unbounded(&input))
.map(|value| {
- input.unbounded = value;
+ input.data = value;
Transformed::Yes(input)
})
}
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 e49b358608..4656b5b270 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
@@ -19,121 +19,78 @@
//! order-preserving variants when it is helpful; either in terms of
//! performance or to accommodate unbounded streams by fixing the pipeline.
-use std::borrow::Cow;
use std::sync::Arc;
-use super::utils::is_repartition;
+use super::utils::{is_repartition, is_sort_preserving_merge};
use crate::error::Result;
use crate::physical_optimizer::utils::{is_coalesce_partitions, is_sort};
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
-use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
use datafusion_common::config::ConfigOptions;
-use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::tree_node::Transformed;
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::tree_node::PlanContext;
use datafusion_physical_plan::unbounded_output;
+use itertools::izip;
+
/// For a given `plan`, this object carries the information one needs from its
/// descendants to decide whether it is beneficial to replace order-losing (but
/// somewhat faster) variants of certain operators with their order-preserving
/// (but somewhat slower) cousins.
-#[derive(Debug, Clone)]
-pub(crate) struct OrderPreservationContext {
- pub(crate) plan: Arc<dyn ExecutionPlan>,
- ordering_connection: bool,
- children_nodes: Vec<Self>,
-}
+pub type OrderPreservationContext = PlanContext<bool>;
+
+/// Updates order-preservation data for all children of the given node.
+pub fn update_children(opc: &mut OrderPreservationContext) {
+ for PlanContext {
+ plan,
+ children,
+ data,
+ } in opc.children.iter_mut()
+ {
+ let maintains_input_order = plan.maintains_input_order();
+ let inspect_child = |idx| {
+ maintains_input_order[idx]
+ || is_coalesce_partitions(plan)
+ || is_repartition(plan)
+ };
-impl OrderPreservationContext {
- /// Creates an empty context tree. Each node has `false` connections.
- pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- ordering_connection: false,
- children_nodes: children.into_iter().map(Self::new).collect(),
+ // We cut the path towards nodes that do not maintain ordering.
+ for (idx, c) in children.iter_mut().enumerate() {
+ c.data &= inspect_child(idx);
}
- }
-
- /// Creates a new order-preservation context from those of children nodes.
- pub fn update_children(mut self) -> Result<Self> {
- for node in self.children_nodes.iter_mut() {
- let plan = node.plan.clone();
- let children = plan.children();
- let maintains_input_order = plan.maintains_input_order();
- let inspect_child = |idx| {
- maintains_input_order[idx]
- || is_coalesce_partitions(&plan)
- || is_repartition(&plan)
- };
-
- // We cut the path towards nodes that do not maintain ordering.
- for (idx, c) in node.children_nodes.iter_mut().enumerate() {
- c.ordering_connection &= inspect_child(idx);
- }
-
- node.ordering_connection = if children.is_empty() {
- false
- } else if !node.children_nodes[0].ordering_connection
- && ((is_repartition(&plan) && !maintains_input_order[0])
- || (is_coalesce_partitions(&plan)
- && children[0].output_ordering().is_some()))
- {
- // We either have a RepartitionExec or a CoalescePartitionsExec
- // and they lose their input ordering, so initiate connection:
- true
- } else {
- // Maintain connection if there is a child with a connection,
- // and operator can possibly maintain that connection (either
- // in its current form or when we replace it with the corresponding
- // order preserving operator).
- node.children_nodes
- .iter()
- .enumerate()
- .any(|(idx, c)| c.ordering_connection && inspect_child(idx))
- }
- }
-
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- self.ordering_connection = false;
- Ok(self)
- }
-}
-
-impl TreeNode for OrderPreservationContext {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children_nodes.iter().map(Cow::Borrowed).collect()
- }
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children_nodes.is_empty() {
- self.children_nodes = self
- .children_nodes
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
+ let plan_children = plan.children();
+ *data = if plan_children.is_empty() {
+ false
+ } else if !children[0].data
+ && ((is_repartition(plan) && !maintains_input_order[0])
+ || (is_coalesce_partitions(plan)
+ && plan_children[0].output_ordering().is_some()))
+ {
+ // We either have a RepartitionExec or a CoalescePartitionsExec
+ // and they lose their input ordering, so initiate connection:
+ true
+ } else {
+ // Maintain connection if there is a child with a connection,
+ // and operator can possibly maintain that connection (either
+ // in its current form or when we replace it with the corresponding
+ // order preserving operator).
+ children
+ .iter()
+ .enumerate()
+ .any(|(idx, c)| c.data && inspect_child(idx))
}
- Ok(self)
}
+ opc.data = false;
}
/// Calculates the updated plan by replacing operators that lose ordering
/// inside `sort_input` with their order-preserving variants. This will
/// generate an alternative plan, which will be accepted or rejected later on
/// depending on whether it helps us remove a `SortExec`.
-fn get_updated_plan(
+fn plan_with_order_preserving_variants(
mut sort_input: OrderPreservationContext,
// Flag indicating that it is desirable to replace `RepartitionExec`s with
// `SortPreservingRepartitionExec`s:
@@ -142,74 +99,115 @@ fn get_updated_plan(
// with `SortPreservingMergeExec`s:
is_spm_better: bool,
) -> Result<OrderPreservationContext> {
- let updated_children = sort_input
- .children_nodes
- .clone()
+ sort_input.children = sort_input
+ .children
.into_iter()
- .map(|item| {
- // Update children and their descendants in the given tree if the connection is open:
- if item.ordering_connection {
- get_updated_plan(item, is_spr_better, is_spm_better)
+ .map(|node| {
+ // Update descendants in the given tree if there is a connection:
+ if node.data {
+ plan_with_order_preserving_variants(node, is_spr_better, is_spm_better)
} else {
- Ok(item)
+ Ok(node)
}
})
- .collect::<Result<Vec<_>>>()?;
+ .collect::<Result<_>>()?;
+ sort_input.data = false;
- sort_input.plan = sort_input
- .plan
- .with_new_children(updated_children.iter().map(|c| c.plan.clone()).collect())?;
- sort_input.ordering_connection = false;
- sort_input.children_nodes = updated_children;
-
- // When a `RepartitionExec` doesn't preserve ordering, replace it with
- // a sort-preserving variant if appropriate:
if is_repartition(&sort_input.plan)
&& !sort_input.plan.maintains_input_order()[0]
&& is_spr_better
{
- let child = sort_input.plan.children().swap_remove(0);
- let repartition =
- RepartitionExec::try_new(child, sort_input.plan.output_partitioning())?
- .with_preserve_order();
- sort_input.plan = Arc::new(repartition) as _;
- sort_input.children_nodes[0].ordering_connection = true;
+ // When a `RepartitionExec` doesn't preserve ordering, replace it with
+ // a sort-preserving variant if appropriate:
+ let child = sort_input.children[0].plan.clone();
+ let partitioning = sort_input.plan.output_partitioning();
+ sort_input.plan = Arc::new(
+ RepartitionExec::try_new(child, partitioning)?.with_preserve_order(),
+ ) as _;
+ sort_input.children[0].data = true;
+ return Ok(sort_input);
} else if is_coalesce_partitions(&sort_input.plan) && is_spm_better {
- // When the input of a `CoalescePartitionsExec` has an ordering, replace it
- // with a `SortPreservingMergeExec` if appropriate:
- if let Some(ordering) = sort_input.children_nodes[0]
- .plan
- .output_ordering()
- .map(|o| o.to_vec())
+ let child = &sort_input.children[0].plan;
+ if let Some(ordering) = child.output_ordering().map(Vec::from) {
+ // When the input of a `CoalescePartitionsExec` has an ordering,
+ // replace it with a `SortPreservingMergeExec` if appropriate:
+ let spm = SortPreservingMergeExec::new(ordering, child.clone());
+ sort_input.plan = Arc::new(spm) as _;
+ sort_input.children[0].data = true;
+ return Ok(sort_input);
+ }
+ }
+
+ sort_input.update_plan_from_children()
+}
+
+/// Calculates the updated plan by replacing operators that preserve ordering
+/// inside `sort_input` with their order-breaking variants. This will restore
+/// the original plan modified by [`plan_with_order_preserving_variants`].
+fn plan_with_order_breaking_variants(
+ mut sort_input: OrderPreservationContext,
+) -> Result<OrderPreservationContext> {
+ let plan = &sort_input.plan;
+ sort_input.children = izip!(
+ sort_input.children,
+ plan.maintains_input_order(),
+ plan.required_input_ordering()
+ )
+ .map(|(node, maintains, required_ordering)| {
+ // Replace with non-order preserving variants as long as ordering is
+ // not required by intermediate operators:
+ if maintains
+ && (is_sort_preserving_merge(plan)
+ || !required_ordering.map_or(false, |required_ordering| {
+ node.plan
+ .equivalence_properties()
+ .ordering_satisfy_requirement(&required_ordering)
+ }))
{
- // Now we can mutate `new_node.children_nodes` safely
- let child = sort_input.children_nodes.clone().swap_remove(0);
- sort_input.plan =
- Arc::new(SortPreservingMergeExec::new(ordering, child.plan)) as _;
- sort_input.children_nodes[0].ordering_connection = true;
+ plan_with_order_breaking_variants(node)
+ } else {
+ Ok(node)
}
+ })
+ .collect::<Result<_>>()?;
+ sort_input.data = false;
+
+ if is_repartition(plan) && plan.maintains_input_order()[0] {
+ // When a `RepartitionExec` preserves ordering, replace it with a
+ // non-sort-preserving variant:
+ let child = sort_input.children[0].plan.clone();
+ let partitioning = plan.output_partitioning();
+ sort_input.plan = Arc::new(RepartitionExec::try_new(child, partitioning)?) as _;
+ } else if is_sort_preserving_merge(plan) {
+ // Replace `SortPreservingMergeExec` with a `CoalescePartitionsExec`:
+ let child = sort_input.children[0].plan.clone();
+ let coalesce = CoalescePartitionsExec::new(child);
+ sort_input.plan = Arc::new(coalesce) as _;
+ } else {
+ return sort_input.update_plan_from_children();
}
+ sort_input.children[0].data = false;
Ok(sort_input)
}
/// The `replace_with_order_preserving_variants` optimizer sub-rule tries to
/// remove `SortExec`s from the physical plan by replacing operators that do
/// not preserve ordering with their order-preserving variants; i.e. by replacing
-/// `RepartitionExec`s with `SortPreservingRepartitionExec`s or by replacing
+/// ordinary `RepartitionExec`s with their sort-preserving variants or by replacing
/// `CoalescePartitionsExec`s with `SortPreservingMergeExec`s.
///
/// If this replacement is helpful for removing a `SortExec`, it updates the plan.
/// Otherwise, it leaves the plan unchanged.
///
-/// Note: this optimizer sub-rule will only produce `SortPreservingRepartitionExec`s
-/// if the query is bounded or if the config option `bounded_order_preserving_variants`
-/// is set to `true`.
+/// NOTE: This optimizer sub-rule will only produce sort-preserving `RepartitionExec`s
+/// if the query is bounded or if the config option `prefer_existing_sort` is
+/// set to `true`.
///
/// The algorithm flow is simply like this:
/// 1. Visit nodes of the physical plan bottom-up and look for `SortExec` nodes.
-/// 1_1. During the traversal, keep track of operators that maintain ordering
-/// (or can maintain ordering when replaced by an order-preserving variant) until
+/// During the traversal, keep track of operators that maintain ordering (or
+/// can maintain ordering when replaced by an order-preserving variant) until
/// a `SortExec` is found.
/// 2. When a `SortExec` is found, update the child of the `SortExec` by replacing
/// operators that do not preserve ordering in the tree with their order
@@ -218,57 +216,56 @@ fn get_updated_plan(
/// its input ordering with the output ordering it imposes. We do this because
/// replacing operators that lose ordering with their order-preserving variants
/// enables us to preserve the previously lost ordering at the input of `SortExec`.
-/// 4. If the `SortExec` in question turns out to be unnecessary, remove it and use
-/// updated plan. Otherwise, use the original plan.
-/// 5. Continue the bottom-up traversal until another `SortExec` is seen, or the traversal
-/// is complete.
+/// 4. If the `SortExec` in question turns out to be unnecessary, remove it and
+/// use updated plan. Otherwise, use the original plan.
+/// 5. Continue the bottom-up traversal until another `SortExec` is seen, or the
+/// traversal is complete.
pub(crate) fn replace_with_order_preserving_variants(
- requirements: OrderPreservationContext,
- // A flag indicating that replacing `RepartitionExec`s with
- // `SortPreservingRepartitionExec`s is desirable when it helps
- // to remove a `SortExec` from the plan. If this flag is `false`,
- // this replacement should only be made to fix the pipeline (streaming).
+ mut requirements: OrderPreservationContext,
+ // A flag indicating that replacing `RepartitionExec`s with sort-preserving
+ // variants is desirable when it helps to remove a `SortExec` from the plan.
+ // If this flag is `false`, this replacement should only be made to fix the
+ // pipeline (streaming).
is_spr_better: bool,
// A flag indicating that replacing `CoalescePartitionsExec`s with
- // `SortPreservingMergeExec`s is desirable when it helps to remove
- // a `SortExec` from the plan. If this flag is `false`, this replacement
+ // `SortPreservingMergeExec`s is desirable when it helps to remove a
+ // `SortExec` from the plan. If this flag is `false`, this replacement
// should only be made to fix the pipeline (streaming).
is_spm_better: bool,
config: &ConfigOptions,
) -> Result<Transformed<OrderPreservationContext>> {
- let mut requirements = requirements.update_children()?;
- if !(is_sort(&requirements.plan)
- && requirements.children_nodes[0].ordering_connection)
- {
+ update_children(&mut requirements);
+ if !(is_sort(&requirements.plan) && requirements.children[0].data) {
return Ok(Transformed::No(requirements));
}
- // For unbounded cases, replace with the order-preserving variant in
- // any case, as doing so helps fix the pipeline.
- // Also do the replacement if opted-in via config options.
+ // For unbounded cases, we replace with the order-preserving variant in any
+ // case, as doing so helps fix the pipeline. Also replace if config allows.
let use_order_preserving_variant =
config.optimizer.prefer_existing_sort || unbounded_output(&requirements.plan);
- let mut updated_sort_input = get_updated_plan(
- requirements.children_nodes.clone().swap_remove(0),
+ // Create an alternate plan with order-preserving variants:
+ let mut alternate_plan = plan_with_order_preserving_variants(
+ requirements.children.swap_remove(0),
is_spr_better || use_order_preserving_variant,
is_spm_better || use_order_preserving_variant,
)?;
- // If this sort is unnecessary, we should remove it and update the plan:
- if updated_sort_input
+ // If the alternate plan makes this sort unnecessary, accept the alternate:
+ if alternate_plan
.plan
.equivalence_properties()
.ordering_satisfy(requirements.plan.output_ordering().unwrap_or(&[]))
{
- for child in updated_sort_input.children_nodes.iter_mut() {
- child.ordering_connection = false;
+ for child in alternate_plan.children.iter_mut() {
+ child.data = false;
}
- Ok(Transformed::Yes(updated_sort_input))
+ Ok(Transformed::Yes(alternate_plan))
} else {
- for child in requirements.children_nodes.iter_mut() {
- child.ordering_connection = false;
- }
+ // The alternate plan does not help, use faster order-breaking variants:
+ alternate_plan = plan_with_order_breaking_variants(alternate_plan)?;
+ alternate_plan.data = false;
+ requirements.children = vec![alternate_plan];
Ok(Transformed::Yes(requirements))
}
}
@@ -280,6 +277,7 @@ mod tests {
use crate::datasource::file_format::file_compression_type::FileCompressionType;
use crate::datasource::listing::PartitionedFile;
use crate::datasource::physical_plan::{CsvExec, FileScanConfig};
+ use crate::physical_optimizer::test_utils::check_integrity;
use crate::physical_plan::coalesce_batches::CoalesceBatchesExec;
use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use crate::physical_plan::filter::FilterExec;
@@ -287,7 +285,9 @@ mod tests {
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
- use crate::physical_plan::{displayable, get_plan_string, Partitioning};
+ use crate::physical_plan::{
+ displayable, get_plan_string, ExecutionPlan, Partitioning,
+ };
use crate::prelude::SessionConfig;
use crate::test::TestStreamPartition;
@@ -394,8 +394,8 @@ mod tests {
// Run the rule top-down
let config = SessionConfig::new().with_prefer_existing_sort($PREFER_EXISTING_SORT);
- let plan_with_pipeline_fixer = OrderPreservationContext::new(physical_plan);
- let parallel = plan_with_pipeline_fixer.transform_up(&|plan_with_pipeline_fixer| replace_with_order_preserving_variants(plan_with_pipeline_fixer, false, false, config.options()))?;
+ let plan_with_pipeline_fixer = OrderPreservationContext::new_default(physical_plan);
+ let parallel = plan_with_pipeline_fixer.transform_up(&|plan_with_pipeline_fixer| replace_with_order_preserving_variants(plan_with_pipeline_fixer, false, false, config.options())).and_then(check_integrity)?;
let optimized_physical_plan = parallel.plan;
// Get string representation of the plan
diff --git a/datafusion/core/src/physical_optimizer/sort_pushdown.rs b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
index f0a8c8cfd3..3413486c6b 100644
--- a/datafusion/core/src/physical_optimizer/sort_pushdown.rs
+++ b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
@@ -15,11 +15,11 @@
// specific language governing permissions and limitations
// under the License.
-use std::borrow::Cow;
use std::sync::Arc;
+use super::utils::add_sort_above;
use crate::physical_optimizer::utils::{
- add_sort_above, is_limit, is_sort_preserving_merge, is_union, is_window,
+ is_limit, is_sort_preserving_merge, is_union, is_window,
};
use crate::physical_plan::filter::FilterExec;
use crate::physical_plan::joins::utils::calculate_join_output_ordering;
@@ -27,9 +27,10 @@ use crate::physical_plan::joins::{HashJoinExec, SortMergeJoinExec};
use crate::physical_plan::projection::ProjectionExec;
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort::SortExec;
-use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
+use crate::physical_plan::tree_node::PlanContext;
+use crate::physical_plan::ExecutionPlan;
-use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::tree_node::Transformed;
use datafusion_common::{plan_err, DataFusionError, JoinSide, Result};
use datafusion_expr::JoinType;
use datafusion_physical_expr::expressions::Column;
@@ -39,60 +40,17 @@ use datafusion_physical_expr::{
/// This is a "data class" we use within the [`EnforceSorting`] rule to push
/// down [`SortExec`] in the plan. In some cases, we can reduce the total
-/// computational cost by pushing down `SortExec`s through some executors.
+/// computational cost by pushing down `SortExec`s through some executors. The
+/// object carries the parent required ordering as its data.
///
/// [`EnforceSorting`]: crate::physical_optimizer::enforce_sorting::EnforceSorting
-#[derive(Debug, Clone)]
-pub(crate) struct SortPushDown {
- /// Current plan
- pub plan: Arc<dyn ExecutionPlan>,
- /// Parent required sort ordering
- required_ordering: Option<Vec<PhysicalSortRequirement>>,
- children_nodes: Vec<Self>,
-}
+pub type SortPushDown = PlanContext<Option<Vec<PhysicalSortRequirement>>>;
-impl SortPushDown {
- /// Creates an empty tree with empty `required_ordering`'s.
- pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
- let children = plan.children();
- Self {
- plan,
- required_ordering: None,
- children_nodes: children.into_iter().map(Self::new).collect(),
- }
- }
-
- /// Assigns the ordering requirement of the root node to the its children.
- pub fn assign_initial_requirements(&mut self) {
- let reqs = self.plan.required_input_ordering();
- for (child, requirement) in self.children_nodes.iter_mut().zip(reqs) {
- child.required_ordering = requirement;
- }
- }
-}
-
-impl TreeNode for SortPushDown {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children_nodes.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children_nodes.is_empty() {
- self.children_nodes = self
- .children_nodes
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- self.plan = with_new_children_if_necessary(
- self.plan,
- self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
- )?
- .into();
- }
- Ok(self)
+/// Assigns the ordering requirement of the root node to the its children.
+pub fn assign_initial_requirements(node: &mut SortPushDown) {
+ let reqs = node.plan.required_input_ordering();
+ for (child, requirement) in node.children.iter_mut().zip(reqs) {
+ child.data = requirement;
}
}
@@ -100,71 +58,61 @@ pub(crate) fn pushdown_sorts(
mut requirements: SortPushDown,
) -> Result<Transformed<SortPushDown>> {
let plan = &requirements.plan;
- let parent_required = requirements.required_ordering.as_deref().unwrap_or(&[]);
+ let parent_reqs = requirements.data.as_deref().unwrap_or(&[]);
+ let satisfy_parent = plan
+ .equivalence_properties()
+ .ordering_satisfy_requirement(parent_reqs);
if let Some(sort_exec) = plan.as_any().downcast_ref::<SortExec>() {
- if !plan
- .equivalence_properties()
- .ordering_satisfy_requirement(parent_required)
- {
- // If the current plan is a SortExec, modify it to satisfy parent requirements:
- let mut new_plan = sort_exec.input().clone();
- add_sort_above(&mut new_plan, parent_required, sort_exec.fetch());
- requirements.plan = new_plan;
- };
-
- let required_ordering = requirements
- .plan
+ let required_ordering = plan
.output_ordering()
.map(PhysicalSortRequirement::from_sort_exprs)
.unwrap_or_default();
- // Since new_plan is a SortExec, we can safely get the 0th index.
- let mut child = requirements.children_nodes.swap_remove(0);
+
+ if !satisfy_parent {
+ // Make sure this `SortExec` satisfies parent requirements:
+ let fetch = sort_exec.fetch();
+ let sort_reqs = requirements.data.unwrap_or_default();
+ requirements = requirements.children.swap_remove(0);
+ requirements = add_sort_above(requirements, sort_reqs, fetch);
+ };
+
+ // We can safely get the 0th index as we are dealing with a `SortExec`.
+ let mut child = requirements.children.swap_remove(0);
if let Some(adjusted) =
pushdown_requirement_to_children(&child.plan, &required_ordering)?
{
- for (c, o) in child.children_nodes.iter_mut().zip(adjusted) {
- c.required_ordering = o;
+ for (grand_child, order) in child.children.iter_mut().zip(adjusted) {
+ grand_child.data = order;
}
// Can push down requirements
- child.required_ordering = None;
- Ok(Transformed::Yes(child))
+ child.data = None;
+ return Ok(Transformed::Yes(child));
} else {
// Can not push down requirements
- let mut empty_node = SortPushDown::new(requirements.plan);
- empty_node.assign_initial_requirements();
- Ok(Transformed::Yes(empty_node))
+ requirements.children = vec![child];
+ assign_initial_requirements(&mut requirements);
}
- } else {
- // Executors other than SortExec
- if plan
- .equivalence_properties()
- .ordering_satisfy_requirement(parent_required)
- {
- // Satisfies parent requirements, immediately return.
- let reqs = requirements.plan.required_input_ordering();
- for (child, order) in requirements.children_nodes.iter_mut().zip(reqs) {
- child.required_ordering = order;
- }
- return Ok(Transformed::Yes(requirements));
+ } else if satisfy_parent {
+ // For non-sort operators, immediately return if parent requirements are met:
+ let reqs = plan.required_input_ordering();
+ for (child, order) in requirements.children.iter_mut().zip(reqs) {
+ child.data = order;
}
- // 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)? {
- for (c, o) in requirements.children_nodes.iter_mut().zip(adjusted) {
- c.required_ordering = o;
- }
- requirements.required_ordering = None;
- Ok(Transformed::Yes(requirements))
- } else {
- // Can not push down requirements, add new SortExec:
- let mut new_plan = requirements.plan;
- add_sort_above(&mut new_plan, parent_required, None);
- let mut new_empty = SortPushDown::new(new_plan);
- new_empty.assign_initial_requirements();
- // Can not push down requirements
- Ok(Transformed::Yes(new_empty))
+ } else if let Some(adjusted) = pushdown_requirement_to_children(plan, parent_reqs)? {
+ // Can not satisfy the parent requirements, check whether we can push
+ // requirements down:
+ for (child, order) in requirements.children.iter_mut().zip(adjusted) {
+ child.data = order;
}
+ requirements.data = None;
+ } else {
+ // Can not push down requirements, add new `SortExec`:
+ let sort_reqs = requirements.data.clone().unwrap_or_default();
+ requirements = add_sort_above(requirements, sort_reqs, None);
+ assign_initial_requirements(&mut requirements);
}
+ Ok(Transformed::Yes(requirements))
}
fn pushdown_requirement_to_children(
@@ -178,11 +126,7 @@ fn pushdown_requirement_to_children(
let child_plan = plan.children().swap_remove(0);
match determine_children_requirement(parent_required, request_child, child_plan) {
RequirementsCompatibility::Satisfy => {
- let req = if request_child.is_empty() {
- None
- } else {
- Some(request_child.to_vec())
- };
+ let req = (!request_child.is_empty()).then(|| request_child.to_vec());
Ok(Some(vec![req]))
}
RequirementsCompatibility::Compatible(adjusted) => Ok(Some(vec![adjusted])),
@@ -191,20 +135,14 @@ fn pushdown_requirement_to_children(
} else if is_union(plan) {
// UnionExec does not have real sort requirements for its input. Here we change the adjusted_request_ordering to UnionExec's output ordering and
// propagate the sort requirements down to correct the unnecessary descendant SortExec under the UnionExec
- let req = if parent_required.is_empty() {
- None
- } else {
- Some(parent_required.to_vec())
- };
+ let req = (!parent_required.is_empty()).then(|| parent_required.to_vec());
Ok(Some(vec![req; plan.children().len()]))
} else if let Some(smj) = plan.as_any().downcast_ref::<SortMergeJoinExec>() {
// If the current plan is SortMergeJoinExec
let left_columns_len = smj.left().schema().fields().len();
let parent_required_expr =
PhysicalSortRequirement::to_sort_exprs(parent_required.iter().cloned());
- let expr_source_side =
- expr_source_sides(&parent_required_expr, smj.join_type(), left_columns_len);
- match expr_source_side {
+ match expr_source_side(&parent_required_expr, smj.join_type(), left_columns_len) {
Some(JoinSide::Left) => try_pushdown_requirements_to_join(
smj,
parent_required,
@@ -256,11 +194,7 @@ fn pushdown_requirement_to_children(
} else {
// Can push-down through SortPreservingMergeExec, because parent requirement is finer
// than SortPreservingMergeExec output ordering.
- let req = if parent_required.is_empty() {
- None
- } else {
- Some(parent_required.to_vec())
- };
+ let req = (!parent_required.is_empty()).then(|| parent_required.to_vec());
Ok(Some(vec![req]))
}
} else {
@@ -268,11 +202,8 @@ fn pushdown_requirement_to_children(
maintains_input_order
.into_iter()
.map(|flag| {
- if flag && !parent_required.is_empty() {
- Some(parent_required.to_vec())
- } else {
- None
- }
+ (flag && !parent_required.is_empty())
+ .then(|| parent_required.to_vec())
})
.collect(),
))
@@ -350,7 +281,7 @@ fn try_pushdown_requirements_to_join(
}))
}
-fn expr_source_sides(
+fn expr_source_side(
required_exprs: &[PhysicalSortExpr],
join_type: JoinType,
left_columns_len: usize,
@@ -402,18 +333,14 @@ fn shift_right_required(
parent_required: LexRequirementRef,
left_columns_len: usize,
) -> Result<Vec<PhysicalSortRequirement>> {
- let new_right_required: Vec<PhysicalSortRequirement> = parent_required
+ let new_right_required = parent_required
.iter()
.filter_map(|r| {
let col = r.expr.as_any().downcast_ref::<Column>()?;
-
- if col.index() < left_columns_len {
- return None;
- }
-
- let new_col =
- Arc::new(Column::new(col.name(), col.index() - left_columns_len));
- Some(r.clone().with_expr(new_col))
+ col.index().checked_sub(left_columns_len).map(|offset| {
+ r.clone()
+ .with_expr(Arc::new(Column::new(col.name(), offset)))
+ })
})
.collect::<Vec<_>>();
if new_right_required.len() == parent_required.len() {
diff --git a/datafusion/core/src/physical_optimizer/test_utils.rs b/datafusion/core/src/physical_optimizer/test_utils.rs
index 68ac9598a3..5de6cff0b4 100644
--- a/datafusion/core/src/physical_optimizer/test_utils.rs
+++ b/datafusion/core/src/physical_optimizer/test_utils.rs
@@ -21,6 +21,7 @@ use std::sync::Arc;
use crate::datasource::listing::PartitionedFile;
use crate::datasource::physical_plan::{FileScanConfig, ParquetExec};
+use crate::datasource::stream::{StreamConfig, StreamTable};
use crate::error::Result;
use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
use crate::physical_plan::coalesce_batches::CoalesceBatchesExec;
@@ -39,13 +40,15 @@ use crate::physical_plan::{ExecutionPlan, InputOrderMode, Partitioning};
use crate::prelude::{CsvReadOptions, SessionContext};
use arrow_schema::{Schema, SchemaRef, SortOptions};
+use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::{JoinType, Statistics};
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_expr::{AggregateFunction, WindowFrame, WindowFunctionDefinition};
use datafusion_physical_expr::expressions::col;
use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
+use datafusion_physical_plan::displayable;
+use datafusion_physical_plan::tree_node::PlanContext;
-use crate::datasource::stream::{StreamConfig, StreamTable};
use async_trait::async_trait;
async fn register_current_csv(
@@ -361,3 +364,25 @@ pub fn sort_exec(
let sort_exprs = sort_exprs.into_iter().collect();
Arc::new(SortExec::new(sort_exprs, input))
}
+
+/// A [`PlanContext`] object is susceptible to being left in an inconsistent state after
+/// untested mutable operations. It is crucial that there be no discrepancies between a plan
+/// associated with the root node and the plan generated after traversing all nodes
+/// within the [`PlanContext`] tree. In addition to verifying the plans resulting from optimizer
+/// rules, it is essential to ensure that the overall tree structure corresponds with the plans
+/// contained within the node contexts.
+/// TODO: Once [`ExecutionPlan`] implements [`PartialEq`], string comparisons should be
+/// replaced with direct plan equality checks.
+pub fn check_integrity<T: Clone>(context: PlanContext<T>) -> Result<PlanContext<T>> {
+ context.transform_up(&|node| {
+ let children_plans = node.plan.children();
+ assert_eq!(node.children.len(), children_plans.len());
+ for (child_plan, child_node) in children_plans.iter().zip(node.children.iter()) {
+ assert_eq!(
+ displayable(child_plan.as_ref()).one_line().to_string(),
+ displayable(child_node.plan.as_ref()).one_line().to_string()
+ );
+ }
+ Ok(Transformed::No(node))
+ })
+}
diff --git a/datafusion/core/src/physical_optimizer/utils.rs b/datafusion/core/src/physical_optimizer/utils.rs
index f8063e9694..4f4b17345e 100644
--- a/datafusion/core/src/physical_optimizer/utils.rs
+++ b/datafusion/core/src/physical_optimizer/utils.rs
@@ -27,30 +27,23 @@ use crate::physical_plan::union::UnionExec;
use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
use crate::physical_plan::ExecutionPlan;
-use datafusion_physical_expr::{LexRequirementRef, PhysicalSortRequirement};
+use datafusion_physical_expr::{LexRequirement, PhysicalSortRequirement};
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
+use datafusion_physical_plan::tree_node::PlanContext;
/// This utility function adds a `SortExec` above an operator according to the
/// given ordering requirements while preserving the original partitioning.
-pub fn add_sort_above(
- node: &mut Arc<dyn ExecutionPlan>,
- sort_requirement: LexRequirementRef,
+pub fn add_sort_above<T: Clone + Default>(
+ node: PlanContext<T>,
+ sort_requirements: LexRequirement,
fetch: Option<usize>,
-) {
- // If the ordering requirement is already satisfied, do not add a sort.
- if !node
- .equivalence_properties()
- .ordering_satisfy_requirement(sort_requirement)
- {
- 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 {
- new_sort.with_preserve_partitioning(true)
- } else {
- new_sort
- }) as _
+) -> PlanContext<T> {
+ let sort_expr = PhysicalSortRequirement::to_sort_exprs(sort_requirements);
+ let mut new_sort = SortExec::new(sort_expr, node.plan.clone()).with_fetch(fetch);
+ if node.plan.output_partitioning().partition_count() > 1 {
+ new_sort = new_sort.with_preserve_partitioning(true);
}
+ PlanContext::new(Arc::new(new_sort), T::default(), vec![node])
}
/// Checks whether the given operator is a limit;
diff --git a/datafusion/expr/src/tree_node/expr.rs b/datafusion/expr/src/tree_node/expr.rs
index 56388be58b..05464c96d0 100644
--- a/datafusion/expr/src/tree_node/expr.rs
+++ b/datafusion/expr/src/tree_node/expr.rs
@@ -23,15 +23,17 @@ use crate::expr::{
ScalarFunction, ScalarFunctionDefinition, Sort, TryCast, WindowFunction,
};
use crate::{Expr, GetFieldAccess};
-use std::borrow::Cow;
-use datafusion_common::tree_node::TreeNode;
+use datafusion_common::tree_node::{TreeNode, VisitRecursion};
use datafusion_common::{internal_err, DataFusionError, Result};
impl TreeNode for Expr {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- match self {
- Expr::Alias(Alias { expr, .. })
+ fn apply_children<F: FnMut(&Self) -> Result<VisitRecursion>>(
+ &self,
+ op: &mut F,
+ ) -> Result<VisitRecursion> {
+ let children = match self {
+ Expr::Alias(Alias{expr, .. })
| Expr::Not(expr)
| Expr::IsNotNull(expr)
| Expr::IsTrue(expr)
@@ -45,79 +47,67 @@ impl TreeNode for Expr {
| Expr::Cast(Cast { expr, .. })
| Expr::TryCast(TryCast { expr, .. })
| Expr::Sort(Sort { expr, .. })
- | Expr::InSubquery(InSubquery { expr, .. }) => vec![Cow::Borrowed(expr)],
+ | Expr::InSubquery(InSubquery{ expr, .. }) => vec![expr.as_ref()],
Expr::GetIndexedField(GetIndexedField { expr, field }) => {
- let expr = Cow::Borrowed(expr.as_ref());
+ let expr = expr.as_ref();
match field {
- GetFieldAccess::ListIndex { key } => {
- vec![Cow::Borrowed(key.as_ref()), expr]
- }
- GetFieldAccess::ListRange { start, stop } => {
- vec![Cow::Borrowed(start), Cow::Borrowed(stop), expr]
- }
- GetFieldAccess::NamedStructField { name: _name } => {
- vec![expr]
+ GetFieldAccess::ListIndex {key} => vec![key.as_ref(), expr],
+ GetFieldAccess::ListRange {start, stop} => {
+ vec![start.as_ref(), stop.as_ref(), expr]
}
+ GetFieldAccess::NamedStructField { .. } => vec![expr],
}
}
Expr::GroupingSet(GroupingSet::Rollup(exprs))
- | Expr::GroupingSet(GroupingSet::Cube(exprs)) => exprs.iter().map(Cow::Borrowed).collect(),
- Expr::ScalarFunction(ScalarFunction { args, .. }) => args.iter().map(Cow::Borrowed).collect(),
+ | Expr::GroupingSet(GroupingSet::Cube(exprs)) => exprs.iter().collect(),
+ Expr::ScalarFunction (ScalarFunction{ args, .. } ) => {
+ args.iter().collect()
+ }
Expr::GroupingSet(GroupingSet::GroupingSets(lists_of_exprs)) => {
- lists_of_exprs.iter().flatten().map(Cow::Borrowed).collect()
+ lists_of_exprs.iter().flatten().collect()
}
Expr::Column(_)
// Treat OuterReferenceColumn as a leaf expression
| Expr::OuterReferenceColumn(_, _)
| Expr::ScalarVariable(_, _)
| Expr::Literal(_)
- | Expr::Exists { .. }
+ | Expr::Exists {..}
| Expr::ScalarSubquery(_)
- | Expr::Wildcard { .. }
- | Expr::Placeholder(_) => vec![],
+ | Expr::Wildcard {..}
+ | Expr::Placeholder (_) => vec![],
Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
- vec![Cow::Borrowed(left), Cow::Borrowed(right)]
+ vec![left.as_ref(), right.as_ref()]
}
Expr::Like(Like { expr, pattern, .. })
| Expr::SimilarTo(Like { expr, pattern, .. }) => {
- vec![Cow::Borrowed(expr), Cow::Borrowed(pattern)]
+ vec![expr.as_ref(), pattern.as_ref()]
}
Expr::Between(Between {
expr, low, high, ..
- }) => vec![
- Cow::Borrowed(expr),
- Cow::Borrowed(low),
- Cow::Borrowed(high),
- ],
+ }) => vec![expr.as_ref(), low.as_ref(), high.as_ref()],
Expr::Case(case) => {
let mut expr_vec = vec![];
if let Some(expr) = case.expr.as_ref() {
- expr_vec.push(Cow::Borrowed(expr.as_ref()));
+ expr_vec.push(expr.as_ref());
};
for (when, then) in case.when_then_expr.iter() {
- expr_vec.push(Cow::Borrowed(when));
- expr_vec.push(Cow::Borrowed(then));
+ expr_vec.push(when.as_ref());
+ expr_vec.push(then.as_ref());
}
if let Some(else_expr) = case.else_expr.as_ref() {
- expr_vec.push(Cow::Borrowed(else_expr));
+ expr_vec.push(else_expr.as_ref());
}
expr_vec
}
- Expr::AggregateFunction(AggregateFunction {
- args,
- filter,
- order_by,
- ..
- }) => {
- let mut expr_vec: Vec<_> = args.iter().map(Cow::Borrowed).collect();
-
+ Expr::AggregateFunction(AggregateFunction { args, filter, order_by, .. })
+ => {
+ let mut expr_vec = args.iter().collect::<Vec<_>>();
if let Some(f) = filter {
- expr_vec.push(Cow::Borrowed(f));
+ expr_vec.push(f.as_ref());
}
- if let Some(o) = order_by {
- expr_vec.extend(o.iter().map(Cow::Borrowed).collect::<Vec<_>>());
+ if let Some(order_by) = order_by {
+ expr_vec.extend(order_by);
}
-
expr_vec
}
Expr::WindowFunction(WindowFunction {
@@ -126,34 +116,47 @@ impl TreeNode for Expr {
order_by,
..
}) => {
- let mut expr_vec: Vec<_> = args.iter().map(Cow::Borrowed).collect();
- expr_vec.extend(partition_by.iter().map(Cow::Borrowed).collect::<Vec<_>>());
- expr_vec.extend(order_by.iter().map(Cow::Borrowed).collect::<Vec<_>>());
+ let mut expr_vec = args.iter().collect::<Vec<_>>();
+ expr_vec.extend(partition_by);
+ expr_vec.extend(order_by);
expr_vec
}
Expr::InList(InList { expr, list, .. }) => {
- let mut expr_vec = vec![Cow::Borrowed(expr.as_ref())];
- expr_vec.extend(list.iter().map(Cow::Borrowed).collect::<Vec<_>>());
+ let mut expr_vec = vec![expr.as_ref()];
+ expr_vec.extend(list);
expr_vec
}
+ };
+
+ for child in children {
+ match op(child)? {
+ VisitRecursion::Continue => {}
+ VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
+ VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
+ }
}
- }
- fn map_children<F>(self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- let mut transform = transform;
+ Ok(VisitRecursion::Continue)
+ }
+ fn map_children<F: FnMut(Self) -> Result<Self>>(
+ self,
+ mut transform: F,
+ ) -> Result<Self> {
Ok(match self {
+ Expr::Column(_)
+ | Expr::Wildcard { .. }
+ | Expr::Placeholder(Placeholder { .. })
+ | Expr::OuterReferenceColumn(_, _)
+ | Expr::Exists { .. }
+ | Expr::ScalarSubquery(_)
+ | Expr::ScalarVariable(_, _)
+ | Expr::Literal(_) => self,
Expr::Alias(Alias {
expr,
relation,
name,
}) => Expr::Alias(Alias::new(transform(*expr)?, relation, name)),
- Expr::Column(_) => self,
- Expr::OuterReferenceColumn(_, _) => self,
- Expr::Exists { .. } => self,
Expr::InSubquery(InSubquery {
expr,
subquery,
@@ -163,9 +166,6 @@ impl TreeNode for Expr {
subquery,
negated,
)),
- Expr::ScalarSubquery(_) => self,
- Expr::ScalarVariable(ty, names) => Expr::ScalarVariable(ty, names),
- Expr::Literal(value) => Expr::Literal(value),
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
Expr::BinaryExpr(BinaryExpr::new(
transform_boxed(left, &mut transform)?,
@@ -244,7 +244,6 @@ impl TreeNode for Expr {
))
})
.collect::<Result<Vec<_>>>()?;
-
let else_expr = transform_option_box(case.else_expr, &mut transform)?;
Expr::Case(Case::new(expr, when_then_expr, else_expr))
@@ -273,9 +272,7 @@ impl TreeNode for Expr {
ScalarFunction::new_udf(fun, transform_vec(args, &mut transform)?),
),
ScalarFunctionDefinition::Name(_) => {
- return internal_err!(
- "Function `Expr` with name should be resolved."
- );
+ return internal_err!("Function `Expr` with name should be resolved.")
}
},
Expr::WindowFunction(WindowFunction {
@@ -308,11 +305,9 @@ impl TreeNode for Expr {
))
}
AggregateFunctionDefinition::UDF(fun) => {
- let order_by = if let Some(order_by) = order_by {
- Some(transform_vec(order_by, &mut transform)?)
- } else {
- None
- };
+ let order_by = order_by
+ .map(|order_by| transform_vec(order_by, &mut transform))
+ .transpose()?;
Expr::AggregateFunction(AggregateFunction::new_udf(
fun,
transform_vec(args, &mut transform)?,
@@ -322,9 +317,7 @@ impl TreeNode for Expr {
))
}
AggregateFunctionDefinition::Name(_) => {
- return internal_err!(
- "Function `Expr` with name should be resolved."
- );
+ return internal_err!("Function `Expr` with name should be resolved.")
}
},
Expr::GroupingSet(grouping_set) => match grouping_set {
@@ -337,13 +330,12 @@ impl TreeNode for Expr {
GroupingSet::GroupingSets(lists_of_exprs) => {
Expr::GroupingSet(GroupingSet::GroupingSets(
lists_of_exprs
- .iter()
- .map(|exprs| transform_vec(exprs.clone(), &mut transform))
+ .into_iter()
+ .map(|exprs| transform_vec(exprs, &mut transform))
.collect::<Result<Vec<_>>>()?,
))
}
},
-
Expr::InList(InList {
expr,
list,
@@ -353,62 +345,47 @@ impl TreeNode for Expr {
transform_vec(list, &mut transform)?,
negated,
)),
- Expr::Wildcard { qualifier } => Expr::Wildcard { qualifier },
Expr::GetIndexedField(GetIndexedField { expr, field }) => {
Expr::GetIndexedField(GetIndexedField::new(
transform_boxed(expr, &mut transform)?,
field,
))
}
- Expr::Placeholder(Placeholder { id, data_type }) => {
- Expr::Placeholder(Placeholder { id, data_type })
- }
})
}
}
-fn transform_boxed<F>(boxed_expr: Box<Expr>, transform: &mut F) -> Result<Box<Expr>>
-where
- F: FnMut(Expr) -> Result<Expr>,
-{
- // TODO:
- // It might be possible to avoid an allocation (the Box::new) below by reusing the box.
- let expr: Expr = *boxed_expr;
- let rewritten_expr = transform(expr)?;
- Ok(Box::new(rewritten_expr))
+fn transform_boxed<F: FnMut(Expr) -> Result<Expr>>(
+ boxed_expr: Box<Expr>,
+ transform: &mut F,
+) -> Result<Box<Expr>> {
+ // TODO: It might be possible to avoid an allocation (the Box::new) below by reusing the box.
+ transform(*boxed_expr).map(Box::new)
}
-fn transform_option_box<F>(
+fn transform_option_box<F: FnMut(Expr) -> Result<Expr>>(
option_box: Option<Box<Expr>>,
transform: &mut F,
-) -> Result<Option<Box<Expr>>>
-where
- F: FnMut(Expr) -> Result<Expr>,
-{
+) -> Result<Option<Box<Expr>>> {
option_box
.map(|expr| transform_boxed(expr, transform))
.transpose()
}
/// &mut transform a Option<`Vec` of `Expr`s>
-fn transform_option_vec<F>(
+fn transform_option_vec<F: FnMut(Expr) -> Result<Expr>>(
option_box: Option<Vec<Expr>>,
transform: &mut F,
-) -> Result<Option<Vec<Expr>>>
-where
- F: FnMut(Expr) -> Result<Expr>,
-{
- Ok(if let Some(exprs) = option_box {
- Some(transform_vec(exprs, transform)?)
- } else {
- None
- })
+) -> Result<Option<Vec<Expr>>> {
+ option_box
+ .map(|exprs| transform_vec(exprs, transform))
+ .transpose()
}
/// &mut transform a `Vec` of `Expr`s
-fn transform_vec<F>(v: Vec<Expr>, transform: &mut F) -> Result<Vec<Expr>>
-where
- F: FnMut(Expr) -> Result<Expr>,
-{
+fn transform_vec<F: FnMut(Expr) -> Result<Expr>>(
+ v: Vec<Expr>,
+ transform: &mut F,
+) -> Result<Vec<Expr>> {
v.into_iter().map(transform).collect()
}
diff --git a/datafusion/expr/src/tree_node/plan.rs b/datafusion/expr/src/tree_node/plan.rs
index 208a8b57d7..589bb917a9 100644
--- a/datafusion/expr/src/tree_node/plan.rs
+++ b/datafusion/expr/src/tree_node/plan.rs
@@ -18,33 +18,19 @@
//! Tree node implementation for logical plan
use crate::LogicalPlan;
-use datafusion_common::tree_node::{TreeNodeVisitor, VisitRecursion};
-use datafusion_common::{tree_node::TreeNode, Result};
-use std::borrow::Cow;
-impl TreeNode for LogicalPlan {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.inputs().into_iter().map(Cow::Borrowed).collect()
- }
+use datafusion_common::tree_node::{TreeNode, TreeNodeVisitor, VisitRecursion};
+use datafusion_common::{handle_tree_recursion, Result};
- fn apply<F>(&self, op: &mut F) -> Result<VisitRecursion>
- where
- F: FnMut(&Self) -> Result<VisitRecursion>,
- {
- // Note,
- //
+impl TreeNode for LogicalPlan {
+ fn apply<F: FnMut(&Self) -> Result<VisitRecursion>>(
+ &self,
+ op: &mut F,
+ ) -> Result<VisitRecursion> {
// Compared to the default implementation, we need to invoke
// [`Self::apply_subqueries`] before visiting its children
- match op(self)? {
- VisitRecursion::Continue => {}
- // If the recursion should skip, do not apply to its children. And let the recursion continue
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- // If the recursion should stop, do not apply to its children
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
- };
-
+ handle_tree_recursion!(op(self)?);
self.apply_subqueries(op)?;
-
self.apply_children(&mut |node| node.apply(op))
}
@@ -74,26 +60,20 @@ impl TreeNode for LogicalPlan {
) -> Result<VisitRecursion> {
// Compared to the default implementation, we need to invoke
// [`Self::visit_subqueries`] before visiting its children
-
- match visitor.pre_visit(self)? {
- VisitRecursion::Continue => {}
- // If the recursion should skip, do not apply to its children. And let the recursion continue
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- // If the recursion should stop, do not apply to its children
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
- };
-
+ handle_tree_recursion!(visitor.pre_visit(self)?);
self.visit_subqueries(visitor)?;
+ handle_tree_recursion!(self.apply_children(&mut |node| node.visit(visitor))?);
+ visitor.post_visit(self)
+ }
- match self.apply_children(&mut |node| node.visit(visitor))? {
- VisitRecursion::Continue => {}
- // If the recursion should skip, do not apply to its children. And let the recursion continue
- VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
- // If the recursion should stop, do not apply to its children
- VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
+ fn apply_children<F: FnMut(&Self) -> Result<VisitRecursion>>(
+ &self,
+ op: &mut F,
+ ) -> Result<VisitRecursion> {
+ for child in self.inputs() {
+ handle_tree_recursion!(op(child)?)
}
-
- visitor.post_visit(self)
+ Ok(VisitRecursion::Continue)
}
fn map_children<F>(self, transform: F) -> Result<Self>
diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs
index 31c1cf6119..94650c0a62 100644
--- a/datafusion/physical-expr/src/equivalence/properties.rs
+++ b/datafusion/physical-expr/src/equivalence/properties.rs
@@ -15,28 +15,27 @@
// specific language governing permissions and limitations
// under the License.
-use crate::expressions::Column;
-use arrow_schema::SchemaRef;
-use datafusion_common::{JoinSide, JoinType};
-use indexmap::IndexSet;
-use itertools::Itertools;
use std::collections::{HashMap, HashSet};
use std::hash::{Hash, Hasher};
use std::sync::Arc;
+use super::ordering::collapse_lex_ordering;
use crate::equivalence::{
collapse_lex_req, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping,
};
-
-use crate::expressions::Literal;
+use crate::expressions::{Column, Literal};
use crate::sort_properties::{ExprOrdering, SortProperties};
use crate::{
physical_exprs_contains, LexOrdering, LexOrderingRef, LexRequirement,
LexRequirementRef, PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement,
};
+
+use arrow_schema::SchemaRef;
use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::{JoinSide, JoinType};
-use super::ordering::collapse_lex_ordering;
+use indexmap::IndexSet;
+use itertools::Itertools;
/// A `EquivalenceProperties` object stores useful information related to a schema.
/// Currently, it keeps track of:
@@ -314,8 +313,8 @@ impl EquivalenceProperties {
/// Returns `true` if the specified ordering is satisfied, `false` otherwise.
fn ordering_satisfy_single(&self, req: &PhysicalSortRequirement) -> bool {
let expr_ordering = self.get_expr_ordering(req.expr.clone());
- let ExprOrdering { expr, state, .. } = expr_ordering;
- match state {
+ let ExprOrdering { expr, data, .. } = expr_ordering;
+ match data {
SortProperties::Ordered(options) => {
let sort_expr = PhysicalSortExpr { expr, options };
sort_expr.satisfy(req, self.schema())
@@ -708,9 +707,9 @@ impl EquivalenceProperties {
let ordered_exprs = search_indices
.iter()
.flat_map(|&idx| {
- let ExprOrdering { expr, state, .. } =
+ let ExprOrdering { expr, data, .. } =
eq_properties.get_expr_ordering(exprs[idx].clone());
- if let SortProperties::Ordered(options) = state {
+ if let SortProperties::Ordered(options) = data {
Some((PhysicalSortExpr { expr, options }, idx))
} else {
None
@@ -776,7 +775,7 @@ impl EquivalenceProperties {
/// Returns an `ExprOrdering` object containing the ordering information for
/// the given expression.
pub fn get_expr_ordering(&self, expr: Arc<dyn PhysicalExpr>) -> ExprOrdering {
- ExprOrdering::new(expr.clone())
+ ExprOrdering::new_default(expr.clone())
.transform_up(&|expr| Ok(update_ordering(expr, self)))
// Guaranteed to always return `Ok`.
.unwrap()
@@ -802,18 +801,19 @@ fn update_ordering(
// We have a Column, which is one of the two possible leaf node types:
let normalized_expr = eq_properties.eq_group.normalize_expr(node.expr.clone());
if eq_properties.is_expr_constant(&normalized_expr) {
- node.state = SortProperties::Singleton;
+ node.data = SortProperties::Singleton;
} else if let Some(options) = eq_properties
.normalized_oeq_class()
.get_options(&normalized_expr)
{
- node.state = SortProperties::Ordered(options);
+ node.data = SortProperties::Ordered(options);
} else if !node.expr.children().is_empty() {
// We have an intermediate (non-leaf) node, account for its children:
- node.state = node.expr.get_ordering(&node.children_state());
+ let children_orderings = node.children.iter().map(|c| c.data).collect_vec();
+ node.data = node.expr.get_ordering(&children_orderings);
} else if node.expr.as_any().is::<Literal>() {
// We have a Literal, which is the other possible leaf node type:
- node.state = node.expr.get_ordering(&[]);
+ node.data = node.expr.get_ordering(&[]);
} else {
return Transformed::No(node);
}
@@ -1683,9 +1683,9 @@ mod tests {
let expr_ordering = eq_properties.get_expr_ordering(expr.clone());
let err_msg = format!(
"expr:{:?}, expected: {:?}, actual: {:?}, leading_orderings: {leading_orderings:?}",
- expr, expected, expr_ordering.state
+ expr, expected, expr_ordering.data
);
- assert_eq!(expr_ordering.state, expected, "{}", err_msg);
+ assert_eq!(expr_ordering.data, expected, "{}", err_msg);
}
Ok(())
diff --git a/datafusion/physical-expr/src/intervals/cp_solver.rs b/datafusion/physical-expr/src/intervals/cp_solver.rs
index 5064ad8d5c..b2403dadf0 100644
--- a/datafusion/physical-expr/src/intervals/cp_solver.rs
+++ b/datafusion/physical-expr/src/intervals/cp_solver.rs
@@ -180,7 +180,7 @@ impl ExprIntervalGraphNode {
/// object. Literals are created with definite, singleton intervals while
/// any other expression starts with an indefinite interval ([-∞, ∞]).
pub fn make_node(node: &ExprTreeNode<NodeIndex>, schema: &Schema) -> Result<Self> {
- let expr = node.expression().clone();
+ let expr = node.expr.clone();
if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
let value = literal.value();
Interval::try_new(value.clone(), value.clone())
diff --git a/datafusion/physical-expr/src/sort_properties.rs b/datafusion/physical-expr/src/sort_properties.rs
index 0205f85dce..4df29ced2f 100644
--- a/datafusion/physical-expr/src/sort_properties.rs
+++ b/datafusion/physical-expr/src/sort_properties.rs
@@ -15,16 +15,13 @@
// specific language governing permissions and limitations
// under the License.
-use std::borrow::Cow;
-use std::{ops::Neg, sync::Arc};
+use std::ops::Neg;
-use arrow_schema::SortOptions;
+use crate::tree_node::ExprContext;
-use crate::PhysicalExpr;
-use datafusion_common::tree_node::TreeNode;
-use datafusion_common::Result;
+use arrow_schema::SortOptions;
-/// To propagate [`SortOptions`] across the [`PhysicalExpr`], it is insufficient
+/// To propagate [`SortOptions`] across the `PhysicalExpr`, it is insufficient
/// to simply use `Option<SortOptions>`: There must be a differentiation between
/// unordered columns and literal values, since literals may not break the ordering
/// when they are used as a child of some binary expression when the other child has
@@ -139,56 +136,13 @@ impl Neg for SortProperties {
}
/// The `ExprOrdering` struct is designed to aid in the determination of ordering (represented
-/// by [`SortProperties`]) for a given [`PhysicalExpr`]. When analyzing the orderings
-/// of a [`PhysicalExpr`], the process begins by assigning the ordering of its leaf nodes.
+/// by [`SortProperties`]) for a given `PhysicalExpr`. When analyzing the orderings
+/// of a `PhysicalExpr`, the process begins by assigning the ordering of its leaf nodes.
/// By propagating these leaf node orderings upwards in the expression tree, the overall
-/// ordering of the entire [`PhysicalExpr`] can be derived.
+/// ordering of the entire `PhysicalExpr` can be derived.
///
-/// This struct holds the necessary state information for each expression in the [`PhysicalExpr`].
-/// It encapsulates the orderings (`state`) associated with the expression (`expr`), and
-/// orderings of the children expressions (`children_states`). The [`ExprOrdering`] of a parent
+/// This struct holds the necessary state information for each expression in the `PhysicalExpr`.
+/// It encapsulates the orderings (`data`) associated with the expression (`expr`), and
+/// orderings of the children expressions (`children`). The [`ExprOrdering`] of a parent
/// expression is determined based on the [`ExprOrdering`] states of its children expressions.
-#[derive(Debug, Clone)]
-pub struct ExprOrdering {
- pub expr: Arc<dyn PhysicalExpr>,
- pub state: SortProperties,
- pub children: Vec<Self>,
-}
-
-impl ExprOrdering {
- /// Creates a new [`ExprOrdering`] with [`SortProperties::Unordered`] states
- /// for `expr` and its children.
- pub fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
- let children = expr.children();
- Self {
- expr,
- state: Default::default(),
- children: children.into_iter().map(Self::new).collect(),
- }
- }
-
- /// Get a reference to each child state.
- pub fn children_state(&self) -> Vec<SortProperties> {
- self.children.iter().map(|c| c.state).collect()
- }
-}
-
-impl TreeNode for ExprOrdering {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children.iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- if !self.children.is_empty() {
- self.children = self
- .children
- .into_iter()
- .map(transform)
- .collect::<Result<_>>()?;
- }
- Ok(self)
- }
-}
+pub type ExprOrdering = ExprContext<SortProperties>;
diff --git a/datafusion/physical-expr/src/tree_node.rs b/datafusion/physical-expr/src/tree_node.rs
index 742846cf56..42dc6673af 100644
--- a/datafusion/physical-expr/src/tree_node.rs
+++ b/datafusion/physical-expr/src/tree_node.rs
@@ -16,12 +16,15 @@
// under the License.
//! This module provides common traits for visiting or rewriting tree nodes easily.
-use crate::physical_expr::with_new_children_if_necessary;
-use crate::PhysicalExpr;
-use datafusion_common::tree_node::DynTreeNode;
-use datafusion_common::Result;
+
+use std::fmt::{self, Display, Formatter};
use std::sync::Arc;
+use crate::physical_expr::{with_new_children_if_necessary, PhysicalExpr};
+
+use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode};
+use datafusion_common::Result;
+
impl DynTreeNode for dyn PhysicalExpr {
fn arc_children(&self) -> Vec<Arc<Self>> {
self.children()
@@ -35,3 +38,63 @@ impl DynTreeNode for dyn PhysicalExpr {
with_new_children_if_necessary(arc_self, new_children)
}
}
+
+/// A node object encapsulating a [`PhysicalExpr`] node with a payload. Since there are
+/// two ways to access child plans—directly from the plan and through child nodes—it's
+/// recommended to perform mutable operations via [`Self::update_expr_from_children`].
+#[derive(Debug)]
+pub struct ExprContext<T: Sized> {
+ /// The physical expression associated with this context.
+ pub expr: Arc<dyn PhysicalExpr>,
+ /// Custom data payload of the node.
+ pub data: T,
+ /// Child contexts of this node.
+ pub children: Vec<Self>,
+}
+
+impl<T> ExprContext<T> {
+ pub fn new(expr: Arc<dyn PhysicalExpr>, data: T, children: Vec<Self>) -> Self {
+ Self {
+ expr,
+ data,
+ children,
+ }
+ }
+
+ pub fn update_expr_from_children(mut self) -> Result<Self> {
+ let children_expr = self.children.iter().map(|c| c.expr.clone()).collect();
+ self.expr = with_new_children_if_necessary(self.expr, children_expr)?;
+ Ok(self)
+ }
+}
+
+impl<T: Default> ExprContext<T> {
+ pub fn new_default(plan: Arc<dyn PhysicalExpr>) -> Self {
+ let children = plan.children().into_iter().map(Self::new_default).collect();
+ Self::new(plan, Default::default(), children)
+ }
+}
+
+impl<T: Display> Display for ExprContext<T> {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ write!(f, "expr: {:?}", self.expr)?;
+ write!(f, "data:{}", self.data)?;
+ write!(f, "")
+ }
+}
+
+impl<T> ConcreteTreeNode for ExprContext<T> {
+ fn children(&self) -> Vec<&Self> {
+ self.children.iter().collect()
+ }
+
+ fn take_children(mut self) -> (Self, Vec<Self>) {
+ let children = std::mem::take(&mut self.children);
+ (self, children)
+ }
+
+ fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> {
+ self.children = children;
+ self.update_expr_from_children()
+ }
+}
diff --git a/datafusion/physical-expr/src/utils/mod.rs b/datafusion/physical-expr/src/utils/mod.rs
index 64a62dc782..e14ff26921 100644
--- a/datafusion/physical-expr/src/utils/mod.rs
+++ b/datafusion/physical-expr/src/utils/mod.rs
@@ -18,11 +18,12 @@
mod guarantee;
pub use guarantee::{Guarantee, LiteralGuarantee};
-use std::borrow::{Borrow, Cow};
+use std::borrow::Borrow;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use crate::expressions::{BinaryExpr, Column};
+use crate::tree_node::ExprContext;
use crate::{PhysicalExpr, PhysicalSortExpr};
use arrow::array::{make_array, Array, ArrayRef, BooleanArray, MutableArrayData};
@@ -127,49 +128,7 @@ pub fn get_indices_of_exprs_strict<T: Borrow<Arc<dyn PhysicalExpr>>>(
.collect()
}
-#[derive(Clone, Debug)]
-pub struct ExprTreeNode<T> {
- expr: Arc<dyn PhysicalExpr>,
- data: Option<T>,
- child_nodes: Vec<ExprTreeNode<T>>,
-}
-
-impl<T> ExprTreeNode<T> {
- pub fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
- let children = expr.children();
- ExprTreeNode {
- expr,
- data: None,
- child_nodes: children.into_iter().map(Self::new).collect_vec(),
- }
- }
-
- pub fn expression(&self) -> &Arc<dyn PhysicalExpr> {
- &self.expr
- }
-
- pub fn children(&self) -> &[ExprTreeNode<T>] {
- &self.child_nodes
- }
-}
-
-impl<T: Clone> TreeNode for ExprTreeNode<T> {
- fn children_nodes(&self) -> Vec<Cow<Self>> {
- self.children().iter().map(Cow::Borrowed).collect()
- }
-
- fn map_children<F>(mut self, transform: F) -> Result<Self>
- where
- F: FnMut(Self) -> Result<Self>,
- {
- self.child_nodes = self
- .child_nodes
- .into_iter()
- .map(transform)
- .collect::<Result<Vec<_>>>()?;
- Ok(self)
- }
-}
+pub type ExprTreeNode<T> = ExprContext<Option<T>>;
/// This struct facilitates the [TreeNodeRewriter] mechanism to convert a
/// [PhysicalExpr] tree into a DAEG (i.e. an expression DAG) by collecting
@@ -207,7 +166,7 @@ impl<'a, T, F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>> TreeNodeRewriter
// of visited expressions and return the newly created node index.
None => {
let node_idx = self.graph.add_node((self.constructor)(&node)?);
- for expr_node in node.child_nodes.iter() {
+ for expr_node in node.children.iter() {
self.graph.add_edge(node_idx, expr_node.data.unwrap(), 0);
}
self.visited_plans.push((expr.clone(), node_idx));
@@ -230,7 +189,7 @@ where
F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>,
{
// Create a new expression tree node from the input expression.
- let init = ExprTreeNode::new(expr);
+ let init = ExprTreeNode::new_default(expr);
// Create a new `PhysicalExprDAEGBuilder` instance.
let mut builder = PhysicalExprDAEGBuilder {
graph: StableGraph::<T, usize>::new(),
@@ -388,7 +347,7 @@ mod tests {
}
fn make_dummy_node(node: &ExprTreeNode<NodeIndex>) -> Result<PhysicalExprDummyNode> {
- let expr = node.expression().clone();
+ let expr = node.expr.clone();
let dummy_property = if expr.as_any().is::<BinaryExpr>() {
"Binary"
} else if expr.as_any().is::<Column>() {
diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs
index bce906a00c..b8a5f95c53 100644
--- a/datafusion/physical-plan/src/tree_node.rs
+++ b/datafusion/physical-plan/src/tree_node.rs
@@ -17,11 +17,14 @@
//! This module provides common traits for visiting or rewriting tree nodes easily.
-use crate::{with_new_children_if_necessary, ExecutionPlan};
-use datafusion_common::tree_node::{DynTreeNode, Transformed};
-use datafusion_common::Result;
+use std::fmt::{self, Display, Formatter};
use std::sync::Arc;
+use crate::{displayable, with_new_children_if_necessary, ExecutionPlan};
+
+use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode, Transformed};
+use datafusion_common::Result;
+
impl DynTreeNode for dyn ExecutionPlan {
fn arc_children(&self) -> Vec<Arc<Self>> {
self.children()
@@ -35,3 +38,64 @@ impl DynTreeNode for dyn ExecutionPlan {
with_new_children_if_necessary(arc_self, new_children).map(Transformed::into)
}
}
+
+/// A node object beneficial for writing optimizer rules, encapsulating an [`ExecutionPlan`] node with a payload.
+/// Since there are two ways to access child plans—directly from the plan and through child nodes—it's recommended
+/// to perform mutable operations via [`Self::update_plan_from_children`].
+#[derive(Debug)]
+pub struct PlanContext<T: Sized> {
+ /// The execution plan associated with this context.
+ pub plan: Arc<dyn ExecutionPlan>,
+ /// Custom data payload of the node.
+ pub data: T,
+ /// Child contexts of this node.
+ pub children: Vec<Self>,
+}
+
+impl<T> PlanContext<T> {
+ pub fn new(plan: Arc<dyn ExecutionPlan>, data: T, children: Vec<Self>) -> Self {
+ Self {
+ plan,
+ data,
+ children,
+ }
+ }
+
+ pub fn update_plan_from_children(mut self) -> Result<Self> {
+ let children_plans = self.children.iter().map(|c| c.plan.clone()).collect();
+ self.plan = with_new_children_if_necessary(self.plan, children_plans)?.into();
+ Ok(self)
+ }
+}
+
+impl<T: Default> PlanContext<T> {
+ pub fn new_default(plan: Arc<dyn ExecutionPlan>) -> Self {
+ let children = plan.children().into_iter().map(Self::new_default).collect();
+ Self::new(plan, Default::default(), children)
+ }
+}
+
+impl<T: Display> Display for PlanContext<T> {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ let node_string = displayable(self.plan.as_ref()).one_line();
+ write!(f, "Node plan: {}", node_string)?;
+ write!(f, "Node data: {}", self.data)?;
+ write!(f, "")
+ }
+}
+
+impl<T> ConcreteTreeNode for PlanContext<T> {
+ fn children(&self) -> Vec<&Self> {
+ self.children.iter().collect()
+ }
+
+ fn take_children(mut self) -> (Self, Vec<Self>) {
+ let children = std::mem::take(&mut self.children);
+ (self, children)
+ }
+
+ fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> {
+ self.children = children;
+ self.update_plan_from_children()
+ }
+}