You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ak...@apache.org on 2023/09/10 10:00:45 UTC
[arrow-datafusion] branch main updated: Refactor the EnforceDistribution Rule (#7488)
This is an automated email from the ASF dual-hosted git repository.
akurmustafa 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 60e00b3a65 Refactor the EnforceDistribution Rule (#7488)
60e00b3a65 is described below
commit 60e00b3a6542b0a88ea392bfddf12a5295470225
Author: Mustafa Akur <10...@users.noreply.github.com>
AuthorDate: Sun Sep 10 13:00:39 2023 +0300
Refactor the EnforceDistribution Rule (#7488)
* Initial version of the plan
* Minor changes
* Minor changes
* Minor changes
* Fix 3rd case
* Case 4
* Add exec tree implementation
* Minor changes
* Change benefits API
* Add union handling
* add new test cases
* Add new test
* Add new test
* Add new tests
* All slt tests pass
* Add csv and parquet repartition
* All tests pass
* minor changes
* minor changes
* Add new tests
* Add stat check
* all tests pass
* run enfordist rule test
* match existing tests, in the new rule
* Minor changes
* Minor changes
* Minor changes
* Make sure that enforce sorting rule generates valid plan
* Minor changes
* Replace old rule with new one
* Remove unnecessary code
* all tests pass
* All tests pass, resolve linter errors.
* Update comments
* Update comments, simplifications
* Update documentation
* Update tests, add ordering invalidated check
* Update tests
* Fix failing tpch plan
* Fix failing test
* Fix indeterministic test
* Apply suggestions from code review
Co-authored-by: Metehan Yıldırım <10...@users.noreply.github.com>
* Minor changes
* Add ordering satify check, add immediate roundrobin removal
* Update comment, add should preserve check
* Update failing .slt test
* Review changes
* Update docstring
* Add new test
* Add parallelization by preserving ordering during repartition
* add bounded flag check
* minor changes, rename, update comment
* Add output ordering check, Use config, move sort adding inside if,
* Address reviews
* Address reviews
---------
Co-authored-by: Metehan Yıldırım <10...@users.noreply.github.com>
Co-authored-by: Mehmet Ozan Kabak <oz...@gmail.com>
---
...dist_enforcement.rs => enforce_distribution.rs} | 1302 +++++++++++++++++---
.../{sort_enforcement.rs => enforce_sorting.rs} | 96 +-
datafusion/core/src/physical_optimizer/mod.rs | 4 +-
.../core/src/physical_optimizer/optimizer.rs | 19 +-
.../core/src/physical_optimizer/repartition.rs | 681 +++++-----
.../replace_with_order_preserving_variants.rs | 2 +-
.../core/src/physical_optimizer/sort_pushdown.rs | 19 +-
.../core/src/physical_optimizer/test_utils.rs | 16 +-
datafusion/core/src/physical_optimizer/utils.rs | 10 +-
.../core/src/physical_plan/coalesce_batches.rs | 4 +
.../core/src/physical_plan/repartition/mod.rs | 9 +-
.../physical_plan/sorts/sort_preserving_merge.rs | 4 +
datafusion/core/tests/sql/joins.rs | 28 +-
datafusion/sqllogictest/test_files/decimal.slt | 7 +-
datafusion/sqllogictest/test_files/explain.slt | 1 -
datafusion/sqllogictest/test_files/groupby.slt | 28 +-
.../sqllogictest/test_files/intersection.slt | 1 -
datafusion/sqllogictest/test_files/joins.slt | 63 +-
.../test_files/tpch/drop_tables.slt.part | 2 -
.../sqllogictest/test_files/tpch/q15.slt.part | 28 +-
datafusion/sqllogictest/test_files/union.slt | 15 +-
datafusion/sqllogictest/test_files/window.slt | 1 +
22 files changed, 1656 insertions(+), 684 deletions(-)
diff --git a/datafusion/core/src/physical_optimizer/dist_enforcement.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
similarity index 60%
rename from datafusion/core/src/physical_optimizer/dist_enforcement.rs
rename to datafusion/core/src/physical_optimizer/enforce_distribution.rs
index e8a2786dc1..5be53cf81b 100644
--- a/datafusion/core/src/physical_optimizer/dist_enforcement.rs
+++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs
@@ -16,10 +16,20 @@
// under the License.
//! EnforceDistribution optimizer rule inspects the physical plan with respect
-//! to distribution requirements and adds [RepartitionExec]s to satisfy them
-//! when necessary.
+//! to distribution requirements and adds [`RepartitionExec`]s to satisfy them
+//! when necessary. If increasing parallelism is beneficial (and also desirable
+//! according to the configuration), this rule increases partition counts in
+//! the physical plan.
+
+use std::fmt;
+use std::fmt::Formatter;
+use std::sync::Arc;
+
use crate::config::ConfigOptions;
-use crate::error::Result;
+use crate::datasource::physical_plan::{CsvExec, ParquetExec};
+use crate::error::{DataFusionError, Result};
+use crate::physical_optimizer::enforce_sorting::{unbounded_output, ExecTree};
+use crate::physical_optimizer::utils::{add_sort_above, get_plan_string};
use crate::physical_optimizer::PhysicalOptimizerRule;
use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
@@ -29,29 +39,144 @@ use crate::physical_plan::joins::{
use crate::physical_plan::projection::ProjectionExec;
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort::SortOptions;
+use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use crate::physical_plan::union::{can_interleave, InterleaveExec, UnionExec};
use crate::physical_plan::windows::WindowAggExec;
use crate::physical_plan::Partitioning;
use crate::physical_plan::{with_new_children_if_necessary, Distribution, ExecutionPlan};
+
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
use datafusion_expr::logical_plan::JoinType;
use datafusion_physical_expr::equivalence::EquivalenceProperties;
-use datafusion_physical_expr::expressions::Column;
-use datafusion_physical_expr::expressions::NoOp;
-use datafusion_physical_expr::utils::map_columns_before_projection;
+use datafusion_physical_expr::expressions::{Column, NoOp};
+use datafusion_physical_expr::utils::{
+ map_columns_before_projection, ordering_satisfy_requirement_concrete,
+};
use datafusion_physical_expr::{
expr_list_eq_strict_order, normalize_expr_with_equivalence_properties, PhysicalExpr,
+ PhysicalSortRequirement,
};
-use std::sync::Arc;
-/// The EnforceDistribution rule ensures that distribution requirements are met
-/// in the strictest way. It might add additional [RepartitionExec] to the plan tree
-/// and give a non-optimal plan, but it can avoid the possible data skew in joins.
+use datafusion_common::internal_err;
+use itertools::izip;
+
+/// The `EnforceDistribution` rule ensures that distribution requirements are
+/// met. In doing so, this rule will increase the parallelism in the plan by
+/// introducing repartitioning operators to the physical plan.
+///
+/// For example, given an input such as:
+///
///
-/// For example for a HashJoin with keys(a, b, c), the required Distribution(a, b, c) can be satisfied by
-/// several alternative partitioning ways: [(a, b, c), (a, b), (a, c), (b, c), (a), (b), (c), ( )].
+/// ```text
+/// ┌─────────────────────────────────┐
+/// │ │
+/// │ ExecutionPlan │
+/// │ │
+/// └─────────────────────────────────┘
+/// ▲ ▲
+/// │ │
+/// ┌─────┘ └─────┐
+/// │ │
+/// │ │
+/// │ │
+/// ┌───────────┐ ┌───────────┐
+/// │ │ │ │
+/// │ batch A1 │ │ batch B1 │
+/// │ │ │ │
+/// ├───────────┤ ├───────────┤
+/// │ │ │ │
+/// │ batch A2 │ │ batch B2 │
+/// │ │ │ │
+/// ├───────────┤ ├───────────┤
+/// │ │ │ │
+/// │ batch A3 │ │ batch B3 │
+/// │ │ │ │
+/// └───────────┘ └───────────┘
///
-/// This rule only chooses the exactly match and satisfies the Distribution(a, b, c) by a HashPartition(a, b, c).
+/// Input Input
+/// A B
+/// ```
+///
+/// This rule will attempt to add a `RepartitionExec` to increase parallelism
+/// (to 3, in this case) and create the following arrangement:
+///
+/// ```text
+/// ┌─────────────────────────────────┐
+/// │ │
+/// │ ExecutionPlan │
+/// │ │
+/// └─────────────────────────────────┘
+/// ▲ ▲ ▲ Input now has 3
+/// │ │ │ partitions
+/// ┌───────┘ │ └───────┐
+/// │ │ │
+/// │ │ │
+/// ┌───────────┐ ┌───────────┐ ┌───────────┐
+/// │ │ │ │ │ │
+/// │ batch A1 │ │ batch A3 │ │ batch B3 │
+/// │ │ │ │ │ │
+/// ├───────────┤ ├───────────┤ ├───────────┤
+/// │ │ │ │ │ │
+/// │ batch B2 │ │ batch B1 │ │ batch A2 │
+/// │ │ │ │ │ │
+/// └───────────┘ └───────────┘ └───────────┘
+/// ▲ ▲ ▲
+/// │ │ │
+/// └─────────┐ │ ┌──────────┘
+/// │ │ │
+/// │ │ │
+/// ┌─────────────────────────────────┐ batches are
+/// │ RepartitionExec(3) │ repartitioned
+/// │ RoundRobin │
+/// │ │
+/// └─────────────────────────────────┘
+/// ▲ ▲
+/// │ │
+/// ┌─────┘ └─────┐
+/// │ │
+/// │ │
+/// │ │
+/// ┌───────────┐ ┌───────────┐
+/// │ │ │ │
+/// │ batch A1 │ │ batch B1 │
+/// │ │ │ │
+/// ├───────────┤ ├───────────┤
+/// │ │ │ │
+/// │ batch A2 │ │ batch B2 │
+/// │ │ │ │
+/// ├───────────┤ ├───────────┤
+/// │ │ │ │
+/// │ batch A3 │ │ batch B3 │
+/// │ │ │ │
+/// └───────────┘ └───────────┘
+///
+///
+/// Input Input
+/// A B
+/// ```
+///
+/// The `EnforceDistribution` rule
+/// - is idempotent; i.e. it can be applied multiple times, each time producing
+/// the same result.
+/// - always produces a valid plan in terms of distribution requirements. Its
+/// input plan can be valid or invalid with respect to distribution requirements,
+/// but the output plan will always be valid.
+/// - produces a valid plan in terms of ordering requirements, *if* its input is
+/// a valid plan in terms of ordering requirements. If the input plan is invalid,
+/// this rule does not attempt to fix it as doing so is the responsibility of the
+/// `EnforceSorting` rule.
+///
+/// Note that distribution requirements are met in the strictest way. This may
+/// result in more than strictly necessary [`RepartitionExec`]s in the plan, but
+/// meeting the requirements in the strictest way may help avoid possible data
+/// skew in joins.
+///
+/// For example for a hash join with keys (a, b, c), the required Distribution(a, b, c)
+/// can be satisfied by several alternative partitioning ways: (a, b, c), (a, b),
+/// (a, c), (b, c), (a), (b), (c) and ( ).
+///
+/// This rule only chooses the exact match and satisfies the Distribution(a, b, c)
+/// by a HashPartition(a, b, c).
#[derive(Default)]
pub struct EnforceDistribution {}
@@ -68,26 +193,30 @@ impl PhysicalOptimizerRule for EnforceDistribution {
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
- let target_partitions = config.execution.target_partitions;
let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
- let new_plan = if top_down_join_key_reordering {
+
+ 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 adjusted =
plan_requirements.transform_down(&adjust_input_keys_ordering)?;
adjusted.plan
} else {
- plan
+ // Run a bottom-up process
+ plan.transform_up(&|plan| {
+ Ok(Transformed::Yes(reorder_join_keys_to_inputs(plan)?))
+ })?
};
+
+ let distribution_context = DistributionContext::new(adjusted);
// Distribution enforcement needs to be applied bottom-up.
- new_plan.transform_up(&|plan| {
- let adjusted = if !top_down_join_key_reordering {
- reorder_join_keys_to_inputs(plan)?
- } else {
- plan
- };
- ensure_distribution(adjusted, target_partitions)
- })
+ let distribution_context =
+ distribution_context.transform_up(&|distribution_context| {
+ ensure_distribution(distribution_context, config)
+ })?;
+
+ // If output ordering is not necessary, removes it
+ update_plan_to_remove_unnecessary_final_order(distribution_context)
}
fn name(&self) -> &str {
@@ -529,7 +658,7 @@ fn shift_right_required(
/// The Bottom-Up approach will be useful in future if we plan to support storage partition-wised Joins.
/// In that case, the datasources/tables might be pre-partitioned and we can't adjust the key ordering of the datasources
/// and then can't apply the Top-Down reordering process.
-fn reorder_join_keys_to_inputs(
+pub(crate) fn reorder_join_keys_to_inputs(
plan: Arc<dyn crate::physical_plan::ExecutionPlan>,
) -> Result<Arc<dyn crate::physical_plan::ExecutionPlan>> {
let plan_any = plan.as_any();
@@ -799,78 +928,670 @@ fn new_join_conditions(
new_join_on
}
+/// Updates `dist_onward` such that, to keep track of
+/// `input` in the `exec_tree`.
+///
+/// # Arguments
+///
+/// * `input`: Current execution plan
+/// * `dist_onward`: It keeps track of executors starting from a distribution
+/// changing operator (e.g Repartition, SortPreservingMergeExec, etc.)
+/// until child of `input` (`input` should have single child).
+/// * `input_idx`: index of the `input`, for its parent.
+///
+fn update_distribution_onward(
+ input: Arc<dyn ExecutionPlan>,
+ dist_onward: &mut Option<ExecTree>,
+ input_idx: usize,
+) {
+ // Update the onward tree if there is an active branch
+ if let Some(exec_tree) = dist_onward {
+ // When we add a new operator to change distribution
+ // we add RepartitionExec, SortPreservingMergeExec, CoalescePartitionsExec
+ // in this case, we need to update exec tree idx such that exec tree is now child of these
+ // operators (change the 0, since all of the operators have single child).
+ exec_tree.idx = 0;
+ *exec_tree = ExecTree::new(input, input_idx, vec![exec_tree.clone()]);
+ } else {
+ *dist_onward = Some(ExecTree::new(input, input_idx, vec![]));
+ }
+}
+
+/// Adds RoundRobin repartition operator to the plan increase parallelism.
+///
+/// # Arguments
+///
+/// * `input`: Current execution plan
+/// * `n_target`: desired target partition number, if partition number of the
+/// current executor is less than this value. Partition number will be increased.
+/// * `dist_onward`: It keeps track of executors starting from a distribution
+/// changing operator (e.g Repartition, SortPreservingMergeExec, etc.)
+/// until `input` plan.
+/// * `input_idx`: index of the `input`, for its parent.
+///
+/// # Returns
+///
+/// A [Result] object that contains new execution plan, where desired partition number
+/// is achieved by adding RoundRobin Repartition.
+fn add_roundrobin_on_top(
+ input: Arc<dyn ExecutionPlan>,
+ n_target: usize,
+ dist_onward: &mut Option<ExecTree>,
+ input_idx: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ // Adding repartition is helpful
+ if input.output_partitioning().partition_count() < n_target {
+ // When there is an existing ordering, we preserve ordering
+ // during repartition. This will be un-done in the future
+ // If any of the following conditions is true
+ // - Preserving ordering is not helpful in terms of satisfying ordering requirements
+ // - Usage of order preserving variants is not desirable
+ // (determined by flag `config.optimizer.bounded_order_preserving_variants`)
+ let should_preserve_ordering = input.output_ordering().is_some();
+
+ let new_plan = Arc::new(
+ RepartitionExec::try_new(input, Partitioning::RoundRobinBatch(n_target))?
+ .with_preserve_order(should_preserve_ordering),
+ ) as Arc<dyn ExecutionPlan>;
+ if let Some(exec_tree) = dist_onward {
+ return internal_err!(
+ "ExecTree should have been empty, but got:{:?}",
+ exec_tree
+ );
+ }
+
+ // update distribution onward with new operator
+ update_distribution_onward(new_plan.clone(), dist_onward, input_idx);
+ Ok(new_plan)
+ } else {
+ // Partition is not helpful, we already have desired number of partitions.
+ Ok(input)
+ }
+}
+
+/// Adds a hash repartition operator:
+/// - to increase parallelism, and/or
+/// - to satisfy requirements of the subsequent operators.
+/// Repartition(Hash) is added on top of operator `input`.
+///
+/// # Arguments
+///
+/// * `input`: Current execution plan
+/// * `hash_exprs`: Stores Physical Exprs that are used during hashing.
+/// * `n_target`: desired target partition number, if partition number of the
+/// current executor is less than this value. Partition number will be increased.
+/// * `dist_onward`: It keeps track of executors starting from a distribution
+/// changing operator (e.g Repartition, SortPreservingMergeExec, etc.)
+/// until `input` plan.
+/// * `input_idx`: index of the `input`, for its parent.
+///
+/// # Returns
+///
+/// A [Result] object that contains new execution plan, where desired distribution is
+/// satisfied by adding Hash Repartition.
+fn add_hash_on_top(
+ input: Arc<dyn ExecutionPlan>,
+ hash_exprs: Vec<Arc<dyn PhysicalExpr>>,
+ // Repartition(Hash) will have `n_target` partitions at the output.
+ n_target: usize,
+ // Stores executors starting from Repartition(RoundRobin) until
+ // current executor. When Repartition(Hash) is added, `dist_onward`
+ // is updated such that it stores connection from Repartition(RoundRobin)
+ // until Repartition(Hash).
+ dist_onward: &mut Option<ExecTree>,
+ input_idx: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ // When there is an existing ordering, we preserve ordering
+ // during repartition. This will be un-done in the future
+ // If any of the following conditions is true
+ // - Preserving ordering is not helpful in terms of satisfying ordering requirements
+ // - Usage of order preserving variants is not desirable
+ // (determined by flag `config.optimizer.bounded_order_preserving_variants`)
+ let should_preserve_ordering = input.output_ordering().is_some();
+ // Since hashing benefits from partitioning, add a round-robin repartition
+ // before it:
+ let mut new_plan = add_roundrobin_on_top(input, n_target, dist_onward, 0)?;
+ new_plan = Arc::new(
+ RepartitionExec::try_new(new_plan, Partitioning::Hash(hash_exprs, n_target))?
+ .with_preserve_order(should_preserve_ordering),
+ ) as _;
+
+ // update distribution onward with new operator
+ update_distribution_onward(new_plan.clone(), dist_onward, input_idx);
+ Ok(new_plan)
+}
+
+/// Adds a `SortPreservingMergeExec` operator on top of input executor:
+/// - to satisfy single distribution requirement.
+///
+/// # Arguments
+///
+/// * `input`: Current execution plan
+/// * `dist_onward`: It keeps track of executors starting from a distribution
+/// changing operator (e.g Repartition, SortPreservingMergeExec, etc.)
+/// until `input` plan.
+/// * `input_idx`: index of the `input`, for its parent.
+///
+/// # Returns
+///
+/// New execution plan, where desired single
+/// distribution is satisfied by adding `SortPreservingMergeExec`.
+fn add_spm_on_top(
+ input: Arc<dyn ExecutionPlan>,
+ dist_onward: &mut Option<ExecTree>,
+ input_idx: usize,
+) -> Arc<dyn ExecutionPlan> {
+ // Add SortPreservingMerge only when partition count is larger than 1.
+ if input.output_partitioning().partition_count() > 1 {
+ // When there is an existing ordering, we preserve ordering
+ // during decreasıng partıtıons. This will be un-done in the future
+ // If any of the following conditions is true
+ // - Preserving ordering is not helpful in terms of satisfying ordering requirements
+ // - Usage of order preserving variants is not desirable
+ // (determined by flag `config.optimizer.bounded_order_preserving_variants`)
+ let should_preserve_ordering = input.output_ordering().is_some();
+ let new_plan: Arc<dyn ExecutionPlan> = if should_preserve_ordering {
+ let existing_ordering = input.output_ordering().unwrap_or(&[]);
+ Arc::new(SortPreservingMergeExec::new(
+ existing_ordering.to_vec(),
+ input,
+ )) as _
+ } else {
+ Arc::new(CoalescePartitionsExec::new(input)) as _
+ };
+
+ // update repartition onward with new operator
+ update_distribution_onward(new_plan.clone(), dist_onward, input_idx);
+ new_plan
+ } else {
+ input
+ }
+}
+
+/// Updates the physical plan inside `distribution_context` if having a
+/// `RepartitionExec(RoundRobin)` is not helpful.
+///
+/// Assume that following plan is given:
+/// ```text
+/// "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10",
+/// " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+/// " ParquetExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC]",
+/// ```
+///
+/// `RepartitionExec` at the top is unnecessary. Since it doesn't help with increasing parallelism.
+/// This function removes top repartition, and returns following plan.
+///
+/// ```text
+/// "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+/// " ParquetExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC]",
+/// ```
+fn remove_unnecessary_repartition(
+ distribution_context: DistributionContext,
+) -> Result<DistributionContext> {
+ let DistributionContext {
+ mut plan,
+ mut distribution_onwards,
+ } = distribution_context;
+
+ // Remove any redundant RoundRobin at the start:
+ if let Some(repartition) = plan.as_any().downcast_ref::<RepartitionExec>() {
+ if let Partitioning::RoundRobinBatch(n_out) = repartition.partitioning() {
+ // Repartition is useless:
+ if *n_out <= repartition.input().output_partitioning().partition_count() {
+ let mut new_distribution_onwards =
+ vec![None; repartition.input().children().len()];
+ if let Some(exec_tree) = &distribution_onwards[0] {
+ for child in &exec_tree.children {
+ new_distribution_onwards[child.idx] = Some(child.clone());
+ }
+ }
+ plan = repartition.input().clone();
+ distribution_onwards = new_distribution_onwards;
+ }
+ }
+ }
+
+ // Create a plan with the updated children:
+ Ok(DistributionContext {
+ plan,
+ distribution_onwards,
+ })
+}
+
+/// Changes each child of the `dist_context.plan` such that they no longer
+/// use order preserving variants, if no ordering is required at the output
+/// of the physical plan (there is no global ordering requirement by the query).
+fn update_plan_to_remove_unnecessary_final_order(
+ dist_context: DistributionContext,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let DistributionContext {
+ plan,
+ distribution_onwards,
+ } = dist_context;
+ let new_children = izip!(plan.children(), distribution_onwards)
+ .map(|(mut child, mut dist_onward)| {
+ replace_order_preserving_variants(&mut child, &mut dist_onward)?;
+ Ok(child)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ if !new_children.is_empty() {
+ plan.with_new_children(new_children)
+ } else {
+ Ok(plan)
+ }
+}
+
+/// Updates the physical plan `input` by using `dist_onward` replace order preserving operator variants
+/// with their corresponding operators that do not preserve order. It is a wrapper for `replace_order_preserving_variants_helper`
+fn replace_order_preserving_variants(
+ input: &mut Arc<dyn ExecutionPlan>,
+ dist_onward: &mut Option<ExecTree>,
+) -> Result<()> {
+ if let Some(dist_onward) = dist_onward {
+ *input = replace_order_preserving_variants_helper(dist_onward)?;
+ }
+ *dist_onward = None;
+ Ok(())
+}
+
+/// Updates the physical plan inside `ExecTree` if preserving ordering while changing partitioning
+/// is not helpful or desirable.
+///
+/// Assume that following plan is given:
+/// ```text
+/// "SortPreservingMergeExec: \[a@0 ASC]"
+/// " SortPreservingRepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10",
+/// " SortPreservingRepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+/// " ParquetExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC]",
+/// ```
+///
+/// This function converts plan above (inside `ExecTree`) to the following:
+///
+/// ```text
+/// "CoalescePartitionsExec"
+/// " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10",
+/// " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+/// " ParquetExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC]",
+/// ```
+fn replace_order_preserving_variants_helper(
+ exec_tree: &ExecTree,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let mut updated_children = exec_tree.plan.children();
+ for child in &exec_tree.children {
+ updated_children[child.idx] = replace_order_preserving_variants_helper(child)?;
+ }
+ if let Some(spm) = exec_tree
+ .plan
+ .as_any()
+ .downcast_ref::<SortPreservingMergeExec>()
+ {
+ return Ok(Arc::new(CoalescePartitionsExec::new(spm.input().clone())));
+ }
+ if let Some(repartition) = exec_tree.plan.as_any().downcast_ref::<RepartitionExec>() {
+ if repartition.preserve_order() {
+ return Ok(Arc::new(
+ RepartitionExec::try_new(
+ repartition.input().clone(),
+ repartition.partitioning().clone(),
+ )?
+ .with_preserve_order(false),
+ ));
+ }
+ }
+ exec_tree.plan.clone().with_new_children(updated_children)
+}
+
/// This function checks whether we need to add additional data exchange
/// operators to satisfy distribution requirements. Since this function
/// takes care of such requirements, we should avoid manually adding data
/// exchange operators in other places.
fn ensure_distribution(
- plan: Arc<dyn ExecutionPlan>,
- target_partitions: usize,
-) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
- if plan.children().is_empty() {
- return Ok(Transformed::No(plan));
+ dist_context: DistributionContext,
+ config: &ConfigOptions,
+) -> Result<Transformed<DistributionContext>> {
+ let target_partitions = config.execution.target_partitions;
+ // When `false`, round robin repartition will not be added to increase parallelism
+ let enable_round_robin = config.optimizer.enable_round_robin_repartition;
+ let repartition_file_scans = config.optimizer.repartition_file_scans;
+ let repartition_file_min_size = config.optimizer.repartition_file_min_size;
+ let is_unbounded = unbounded_output(&dist_context.plan);
+ // Use order preserving variants either of the conditions true
+ // - it is desired according to config
+ // - when plan is unbounded
+ let order_preserving_variants_desirable =
+ is_unbounded || config.optimizer.bounded_order_preserving_variants;
+
+ if dist_context.plan.children().is_empty() {
+ return Ok(Transformed::No(dist_context));
}
-
- // special case for UnionExec: 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
- if let Some(union_exec) = plan.as_any().downcast_ref::<UnionExec>() {
- if can_interleave(union_exec.inputs()) {
- let plan = InterleaveExec::try_new(union_exec.inputs().clone())?;
- return Ok(Transformed::Yes(Arc::new(plan)));
- }
+ // Don't need to apply when the returned row count is not greater than 1:
+ let stats = dist_context.plan.statistics();
+ let mut repartition_beneficial_stat = true;
+ if stats.is_exact {
+ repartition_beneficial_stat =
+ stats.num_rows.map(|num_rows| num_rows > 1).unwrap_or(true);
}
- let required_input_distributions = plan.required_input_distribution();
- let children: Vec<Arc<dyn ExecutionPlan>> = plan.children();
- assert_eq!(children.len(), required_input_distributions.len());
+ // Remove unnecessary repartition from the physical plan if any
+ let DistributionContext {
+ plan,
+ mut distribution_onwards,
+ } = remove_unnecessary_repartition(dist_context)?;
+
+ let n_children = plan.children().len();
+ // This loop iterates over all the children to:
+ // - Increase parallelism for every child if it is beneficial.
+ // - Satisfy the distribution requirements of every child, if it is not
+ // already satisfied.
+ // We store the updated children in `new_children`.
+ let new_children = izip!(
+ plan.children().into_iter(),
+ plan.required_input_distribution().iter(),
+ plan.required_input_ordering().iter(),
+ distribution_onwards.iter_mut(),
+ plan.benefits_from_input_partitioning(),
+ plan.maintains_input_order(),
+ 0..n_children
+ )
+ .map(
+ |(
+ mut child,
+ requirement,
+ required_input_ordering,
+ dist_onward,
+ would_benefit,
+ maintains,
+ child_idx,
+ )| {
+ if enable_round_robin
+ // Operator benefits from partitioning (e.g. filter):
+ && (would_benefit && repartition_beneficial_stat)
+ // Unless partitioning doesn't increase the partition count, it is not beneficial:
+ && child.output_partitioning().partition_count() < target_partitions
+ {
+ // When `repartition_file_scans` is set, leverage source operators
+ // (`ParquetExec`, `CsvExec` etc.) to increase parallelism at the source.
+ if repartition_file_scans {
+ if let Some(parquet_exec) =
+ child.as_any().downcast_ref::<ParquetExec>()
+ {
+ child = Arc::new(parquet_exec.get_repartitioned(
+ target_partitions,
+ repartition_file_min_size,
+ ));
+ } else if let Some(csv_exec) =
+ child.as_any().downcast_ref::<CsvExec>()
+ {
+ if let Some(csv_exec) = csv_exec.get_repartitioned(
+ target_partitions,
+ repartition_file_min_size,
+ ) {
+ child = Arc::new(csv_exec);
+ }
+ }
+ }
+ // Increase parallelism by adding round-robin repartitioning
+ // on top of the operator. Note that we only do this if the
+ // partition count is not already equal to the desired partition
+ // count.
+ child = add_roundrobin_on_top(
+ child,
+ target_partitions,
+ dist_onward,
+ child_idx,
+ )?;
+ }
- // Add RepartitionExec to guarantee output partitioning
- let new_children: Result<Vec<Arc<dyn ExecutionPlan>>> = children
- .into_iter()
- .zip(required_input_distributions)
- .map(|(child, required)| {
- if child
+ if !child
.output_partitioning()
- .satisfy(required.clone(), || child.equivalence_properties())
+ .satisfy(requirement.clone(), || child.equivalence_properties())
{
- Ok(child)
- } else {
- let new_child: Result<Arc<dyn ExecutionPlan>> = match required {
- Distribution::SinglePartition
- if child.output_partitioning().partition_count() > 1 =>
- {
- Ok(Arc::new(CoalescePartitionsExec::new(child.clone())))
+ // Satisfy the distribution requirement if it is unmet.
+ match requirement {
+ Distribution::SinglePartition => {
+ child = add_spm_on_top(child, dist_onward, child_idx);
}
- _ => {
- let partition = required.create_partitioning(target_partitions);
- Ok(Arc::new(RepartitionExec::try_new(child, partition)?))
+ Distribution::HashPartitioned(exprs) => {
+ child = add_hash_on_top(
+ child,
+ exprs.to_vec(),
+ target_partitions,
+ dist_onward,
+ child_idx,
+ )?;
}
+ Distribution::UnspecifiedDistribution => {}
};
- new_child
}
+ // There is an ordering requirement of the operator:
+ if let Some(required_input_ordering) = required_input_ordering {
+ let existing_ordering = child.output_ordering().unwrap_or(&[]);
+ // Either:
+ // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or
+ // - using order preserving variant is not desirable.
+ if !ordering_satisfy_requirement_concrete(
+ existing_ordering,
+ required_input_ordering,
+ || child.equivalence_properties(),
+ || child.ordering_equivalence_properties(),
+ ) || !order_preserving_variants_desirable
+ {
+ replace_order_preserving_variants(&mut child, dist_onward)?;
+ let sort_expr = PhysicalSortRequirement::to_sort_exprs(
+ required_input_ordering.clone(),
+ );
+ // Make sure to satisfy ordering requirement
+ add_sort_above(&mut child, sort_expr, None)?;
+ }
+ // Stop tracking distribution changing operators
+ *dist_onward = None;
+ } else {
+ // no ordering requirement
+ match requirement {
+ // Operator requires specific distribution.
+ Distribution::SinglePartition | Distribution::HashPartitioned(_) => {
+ // Since there is no ordering requirement, preserving ordering is pointless
+ replace_order_preserving_variants(&mut child, dist_onward)?;
+ }
+ Distribution::UnspecifiedDistribution => {
+ // Since ordering is lost, trying to preserve ordering is pointless
+ if !maintains {
+ replace_order_preserving_variants(&mut child, dist_onward)?;
+ }
+ }
+ }
+ }
+ Ok(child)
+ },
+ )
+ .collect::<Result<Vec<_>>>()?;
+
+ let new_distribution_context = DistributionContext {
+ plan: if plan.as_any().is::<UnionExec>() && can_interleave(&new_children) {
+ // 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(new_children)?)
+ } else {
+ plan.clone().with_new_children(new_children)?
+ },
+ distribution_onwards,
+ };
+ Ok(Transformed::Yes(new_distribution_context))
+}
+
+/// A struct to keep track of distribution changing executors
+/// (`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 {
+ pub(crate) plan: Arc<dyn ExecutionPlan>,
+ /// Keep track of associations for each child of the plan. If `None`,
+ /// there is no distribution changing operator in its descendants.
+ distribution_onwards: Vec<Option<ExecTree>>,
+}
+
+impl DistributionContext {
+ /// Creates an empty context.
+ fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+ let length = plan.children().len();
+ DistributionContext {
+ plan,
+ distribution_onwards: vec![None; length],
+ }
+ }
+
+ /// Constructs a new context from children contexts.
+ fn new_from_children_nodes(
+ children_nodes: Vec<DistributionContext>,
+ parent_plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Self> {
+ let children_plans = children_nodes
+ .iter()
+ .map(|item| item.plan.clone())
+ .collect();
+ let distribution_onwards = children_nodes
+ .into_iter()
+ .enumerate()
+ .map(|(idx, context)| {
+ let DistributionContext {
+ plan,
+ // The `distribution_onwards` tree keeps track of operators
+ // that change distribution, or preserves the existing
+ // distribution (starting from an operator that change distribution).
+ distribution_onwards,
+ } = context;
+ if plan.children().is_empty() {
+ // Plan has no children, there is nothing to propagate.
+ None
+ } else if distribution_onwards[0].is_none() {
+ if let Some(repartition) =
+ plan.as_any().downcast_ref::<RepartitionExec>()
+ {
+ match repartition.partitioning() {
+ Partitioning::RoundRobinBatch(_)
+ | Partitioning::Hash(_, _) => {
+ // Start tracking operators starting from this repartition (either roundrobin or hash):
+ return Some(ExecTree::new(plan, idx, vec![]));
+ }
+ _ => {}
+ }
+ } else if plan.as_any().is::<SortPreservingMergeExec>()
+ || plan.as_any().is::<CoalescePartitionsExec>()
+ {
+ // Start tracking operators starting from this sort preserving merge:
+ return Some(ExecTree::new(plan, idx, vec![]));
+ }
+ None
+ } else {
+ // Propagate children distribution tracking to the above
+ let new_distribution_onwards = izip!(
+ plan.required_input_distribution().iter(),
+ distribution_onwards.into_iter()
+ )
+ .flat_map(|(required_dist, distribution_onwards)| {
+ if let Some(distribution_onwards) = distribution_onwards {
+ // Operator can safely propagate the distribution above.
+ // This is similar to maintaining order in the EnforceSorting rule.
+ if let Distribution::UnspecifiedDistribution = required_dist {
+ return Some(distribution_onwards);
+ }
+ }
+ None
+ })
+ .collect::<Vec<_>>();
+ // Either:
+ // - None of the children has a connection to an operator that modifies distribution, or
+ // - The current operator requires distribution at its input so doesn't propagate it above.
+ if new_distribution_onwards.is_empty() {
+ None
+ } else {
+ Some(ExecTree::new(plan, idx, new_distribution_onwards))
+ }
+ }
+ })
+ .collect();
+ Ok(DistributionContext {
+ plan: with_new_children_if_necessary(parent_plan, children_plans)?.into(),
+ distribution_onwards,
})
- .collect();
- with_new_children_if_necessary(plan, new_children?)
+ }
+
+ /// Computes distribution tracking contexts for every child of the plan.
+ fn children(&self) -> Vec<DistributionContext> {
+ self.plan
+ .children()
+ .into_iter()
+ .map(|child| DistributionContext::new(child))
+ .collect()
+ }
+}
+
+impl TreeNode for DistributionContext {
+ fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
+ where
+ F: FnMut(&Self) -> Result<VisitRecursion>,
+ {
+ for child in self.children() {
+ match op(&child)? {
+ VisitRecursion::Continue => {}
+ VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
+ VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
+ }
+ }
+ Ok(VisitRecursion::Continue)
+ }
+
+ fn map_children<F>(self, transform: F) -> Result<Self>
+ where
+ F: FnMut(Self) -> Result<Self>,
+ {
+ let children = self.children();
+ if children.is_empty() {
+ Ok(self)
+ } else {
+ let children_nodes = children
+ .into_iter()
+ .map(transform)
+ .collect::<Result<Vec<_>>>()?;
+ DistributionContext::new_from_children_nodes(children_nodes, self.plan)
+ }
+ }
+}
+
+/// 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)?;
+ for (idx, child) in self.distribution_onwards.iter().enumerate() {
+ if let Some(child) = child {
+ write!(f, "idx:{:?}, exec_tree:{}", idx, child)?;
+ }
+ }
+ write!(f, "")
+ }
}
#[derive(Debug, Clone)]
@@ -881,7 +1602,7 @@ struct JoinKeyPairs {
#[derive(Debug, Clone)]
struct PlanWithKeyRequirements {
- plan: Arc<dyn ExecutionPlan>,
+ pub plan: Arc<dyn ExecutionPlan>,
/// Parent required key ordering
required_key_ordering: Vec<Arc<dyn PhysicalExpr>>,
/// The request key ordering to children
@@ -961,34 +1682,35 @@ impl TreeNode for PlanWithKeyRequirements {
#[cfg(test)]
mod tests {
- use crate::physical_plan::coalesce_batches::CoalesceBatchesExec;
- use crate::physical_plan::filter::FilterExec;
- use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
- use arrow::compute::SortOptions;
- use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
- use datafusion_expr::logical_plan::JoinType;
- use datafusion_expr::Operator;
- use datafusion_physical_expr::{
- expressions, expressions::binary, expressions::lit, expressions::Column,
- PhysicalExpr, PhysicalSortExpr,
- };
use std::ops::Deref;
use super::*;
use crate::datasource::listing::PartitionedFile;
use crate::datasource::object_store::ObjectStoreUrl;
use crate::datasource::physical_plan::{FileScanConfig, ParquetExec};
- use crate::physical_optimizer::sort_enforcement::EnforceSorting;
+ use crate::physical_optimizer::enforce_sorting::EnforceSorting;
use crate::physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
+ use crate::physical_plan::coalesce_batches::CoalesceBatchesExec;
use crate::physical_plan::expressions::col;
+ use crate::physical_plan::filter::FilterExec;
use crate::physical_plan::joins::{
utils::JoinOn, HashJoinExec, PartitionMode, SortMergeJoinExec,
};
use crate::physical_plan::projection::ProjectionExec;
+ use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use crate::physical_plan::{displayable, Statistics};
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+ use datafusion_expr::logical_plan::JoinType;
+ use datafusion_expr::Operator;
+ use datafusion_physical_expr::{
+ expressions, expressions::binary, expressions::lit, expressions::Column,
+ PhysicalExpr, PhysicalSortExpr,
+ };
+
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::Int64, true),
@@ -1154,24 +1876,61 @@ mod tests {
.collect()
}
+ fn ensure_distribution_helper(
+ plan: Arc<dyn ExecutionPlan>,
+ target_partitions: usize,
+ bounded_order_preserving_variants: bool,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let distribution_context = DistributionContext::new(plan);
+ let mut config = ConfigOptions::new();
+ config.execution.target_partitions = target_partitions;
+ config.optimizer.enable_round_robin_repartition = false;
+ config.optimizer.repartition_file_scans = false;
+ config.optimizer.repartition_file_min_size = 1024;
+ config.optimizer.bounded_order_preserving_variants =
+ bounded_order_preserving_variants;
+ ensure_distribution(distribution_context, &config).map(|item| item.into().plan)
+ }
+
/// Runs the repartition optimizer and asserts the plan against the expected
macro_rules! assert_optimized {
- ($EXPECTED_LINES: expr, $PLAN: expr) => {
+ ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr) => {
let expected_lines: Vec<&str> = $EXPECTED_LINES.iter().map(|s| *s).collect();
let mut config = ConfigOptions::new();
config.execution.target_partitions = 10;
+ config.optimizer.bounded_order_preserving_variants = true;
- // run optimizer
- let optimizer = EnforceDistribution {};
- let optimized = optimizer.optimize($PLAN, &config)?;
// NOTE: These tests verify the joint `EnforceDistribution` + `EnforceSorting` cascade
// because they were written prior to the separation of `BasicEnforcement` into
- // `EnforceSorting` and `EnfoceDistribution`.
+ // `EnforceSorting` and `EnforceDistribution`.
// TODO: Orthogonalize the tests here just to verify `EnforceDistribution` and create
// new tests for the cascade.
- let optimizer = EnforceSorting::new();
- let optimized = optimizer.optimize(optimized, &config)?;
+ let optimized = if $FIRST_ENFORCE_DIST {
+ // Run enforce distribution rule first:
+ let optimizer = EnforceDistribution::new();
+ let optimized = optimizer.optimize($PLAN.clone(), &config)?;
+ // The rule should be idempotent.
+ // Re-running this rule shouldn't introduce unnecessary operators.
+ let optimizer = EnforceDistribution::new();
+ let optimized = optimizer.optimize(optimized, &config)?;
+ // Run the enforce sorting rule:
+ let optimizer = EnforceSorting::new();
+ let optimized = optimizer.optimize(optimized, &config)?;
+ optimized
+ } else {
+ // Run the enforce sorting rule first:
+ let optimizer = EnforceSorting::new();
+ let optimized = optimizer.optimize($PLAN.clone(), &config)?;
+ // Run enforce distribution rule:
+ let optimizer = EnforceDistribution::new();
+ let optimized = optimizer.optimize(optimized, &config)?;
+ // The rule should be idempotent.
+ // Re-running this rule shouldn't introduce unnecessary operators.
+ let optimizer = EnforceDistribution::new();
+ let optimized = optimizer.optimize(optimized, &config)?;
+ optimized
+ };
// Now format correctly
let plan = displayable(optimized.as_ref()).indent(true).to_string();
@@ -1260,12 +2019,15 @@ mod tests {
JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti => vec![
top_join_plan.as_str(),
join_plan.as_str(),
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
// Should include 4 RepartitionExecs
@@ -1273,16 +2035,20 @@ mod tests {
top_join_plan.as_str(),
"RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
join_plan.as_str(),
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
};
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+ assert_optimized!(expected, top_join, false);
}
JoinType::RightSemi | JoinType::RightAnti => {}
}
@@ -1316,12 +2082,15 @@ mod tests {
vec![
top_join_plan.as_str(),
join_plan.as_str(),
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
// Should include 4 RepartitionExecs
@@ -1330,16 +2099,20 @@ mod tests {
top_join_plan.as_str(),
"RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10",
join_plan.as_str(),
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
};
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+ assert_optimized!(expected, top_join, false);
}
JoinType::LeftSemi | JoinType::LeftAnti => {}
}
@@ -1385,14 +2158,18 @@ mod tests {
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a1@0, c@2)]",
"ProjectionExec: expr=[a@0 as a1, a@0 as a2]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+ assert_optimized!(expected, top_join, false);
// Join on (a2 == c)
let top_join_on = vec![(
@@ -1407,15 +2184,19 @@ mod tests {
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a2@1, c@2)]",
"ProjectionExec: expr=[a@0 as a1, a@0 as a2]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+ assert_optimized!(expected, top_join, false);
Ok(())
}
@@ -1458,15 +2239,19 @@ mod tests {
"ProjectionExec: expr=[c1@0 as a]",
"ProjectionExec: expr=[c@2 as c1]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+ assert_optimized!(expected, top_join, false);
Ok(())
}
@@ -1494,15 +2279,18 @@ mod tests {
let expected = &[
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a1@0, a2@0)]",
"AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]",
- "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]",
- "RepartitionExec: partitioning=Hash([a2@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a2@0], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[a@0 as a2], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, join);
+ assert_optimized!(expected, join.clone(), true);
+ assert_optimized!(expected, join, false);
Ok(())
}
@@ -1543,15 +2331,18 @@ mod tests {
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b1@1, b@0), (a1@0, a@1)]",
"ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]",
"AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]",
- "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]",
- "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, join);
+ assert_optimized!(expected, join.clone(), true);
+ assert_optimized!(expected, join, false);
Ok(())
}
@@ -1650,19 +2441,24 @@ mod tests {
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(B@2, b1@6), (C@3, c@2), (AA@1, a1@5)]",
"ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (c@2, c1@2), (a@0, a1@0)]",
- "RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (c@2, c1@2), (a@0, a1@0)]",
- "RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, filter_top_join);
+ assert_optimized!(expected, filter_top_join.clone(), true);
+ assert_optimized!(expected, filter_top_join, false);
Ok(())
}
@@ -1691,11 +2487,12 @@ mod tests {
Column::new_with_schema("c1", &right.schema()).unwrap(),
),
];
- let bottom_left_join = ensure_distribution(
+
+ let bottom_left_join = ensure_distribution_helper(
hash_join_exec(left.clone(), right.clone(), &join_on, &JoinType::Inner),
10,
- )?
- .into();
+ true,
+ )?;
// Projection(a as A, a as AA, b as B, c as C)
let alias_pairs: Vec<(String, String)> = vec![
@@ -1722,11 +2519,11 @@ mod tests {
Column::new_with_schema("a1", &right.schema()).unwrap(),
),
];
- let bottom_right_join = ensure_distribution(
+ let bottom_right_join = ensure_distribution_helper(
hash_join_exec(left, right.clone(), &join_on, &JoinType::Inner),
10,
- )?
- .into();
+ true,
+ )?;
// Join on (B == b1 and C == c and AA = a1)
let top_join_on = vec![
@@ -1772,15 +2569,19 @@ mod tests {
top_join_plan.as_str(),
"ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1, b1@1), (c@2, c1@2)]",
- "RepartitionExec: partitioning=Hash([a@0, b@1, c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0, b@1, c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([a1@0, b1@1, c1@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a1@0, b1@1, c1@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c@2, c1@2), (b@1, b1@1), (a@0, a1@0)]",
- "RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
@@ -1812,11 +2613,11 @@ mod tests {
Column::new_with_schema("b1", &right.schema()).unwrap(),
),
];
- let bottom_left_join = ensure_distribution(
+ let bottom_left_join = ensure_distribution_helper(
hash_join_exec(left.clone(), right.clone(), &join_on, &JoinType::Inner),
10,
- )?
- .into();
+ true,
+ )?;
// Projection(a as A, a as AA, b as B, c as C)
let alias_pairs: Vec<(String, String)> = vec![
@@ -1843,11 +2644,11 @@ mod tests {
Column::new_with_schema("a1", &right.schema()).unwrap(),
),
];
- let bottom_right_join = ensure_distribution(
+ let bottom_right_join = ensure_distribution_helper(
hash_join_exec(left, right.clone(), &join_on, &JoinType::Inner),
10,
- )?
- .into();
+ true,
+ )?;
// Join on (B == b1 and C == c and AA = a1)
let top_join_on = vec![
@@ -1893,15 +2694,19 @@ mod tests {
top_join_plan.as_str(),
"ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1, b1@1)]",
- "RepartitionExec: partitioning=Hash([a@0, b@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0, b@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([a1@0, b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a1@0, b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c@2, c1@2), (b@1, b1@1), (a@0, a1@0)]",
- "RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
- "RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
@@ -1961,41 +2766,92 @@ mod tests {
format!("SortMergeJoin: join_type={join_type}, on=[(a@0, c@2)]");
let expected = match join_type {
- // Should include 3 RepartitionExecs 3 SortExecs
+ // Should include 6 RepartitionExecs 3 SortExecs
JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti =>
vec![
top_join_plan.as_str(),
join_plan.as_str(),
"SortExec: expr=[a@0 ASC]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[b1@1 ASC]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[c@2 ASC]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
- // Should include 4 RepartitionExecs
+ // Should include 7 RepartitionExecs
_ => vec![
top_join_plan.as_str(),
"SortExec: expr=[a@0 ASC]",
"RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
join_plan.as_str(),
"SortExec: expr=[a@0 ASC]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[b1@1 ASC]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[c@2 ASC]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
};
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+
+ let expected_first_sort_enforcement = match join_type {
+ // Should include 3 RepartitionExecs 3 SortExecs
+ JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti =>
+ vec![
+ top_join_plan.as_str(),
+ join_plan.as_str(),
+ "SortPreservingRepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[a@0 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[b1@1 ASC]",
+ "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[c@2 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ ],
+ // Should include 8 RepartitionExecs (4 of them preserves ordering)
+ _ => vec![
+ top_join_plan.as_str(),
+ "SortPreservingRepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[a@0 ASC]",
+ "CoalescePartitionsExec",
+ join_plan.as_str(),
+ "SortPreservingRepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[a@0 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[b1@1 ASC]",
+ "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[c@2 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ ],
+ };
+ assert_optimized!(expected_first_sort_enforcement, top_join, false);
match join_type {
JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {
@@ -2020,14 +2876,17 @@ mod tests {
top_join_plan.as_str(),
join_plan.as_str(),
"SortExec: expr=[a@0 ASC]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[b1@1 ASC]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[c@2 ASC]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
// Should include 4 RepartitionExecs and 4 SortExecs
@@ -2037,18 +2896,65 @@ mod tests {
"RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10",
join_plan.as_str(),
"SortExec: expr=[a@0 ASC]",
- "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[b1@1 ASC]",
- "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"SortExec: expr=[c@2 ASC]",
- "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
],
};
- assert_optimized!(expected, top_join);
+ assert_optimized!(expected, top_join.clone(), true);
+
+ let expected_first_sort_enforcement = match join_type {
+ // Should include 6 RepartitionExecs (3 of them preserves order) and 3 SortExecs
+ JoinType::Inner | JoinType::Right => vec![
+ top_join_plan.as_str(),
+ join_plan.as_str(),
+ "SortPreservingRepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[a@0 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[b1@1 ASC]",
+ "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[c@2 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ ],
+ // Should include 8 RepartitionExecs (4 of them preserves order) and 4 SortExecs
+ _ => vec![
+ top_join_plan.as_str(),
+ "SortPreservingRepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[b1@6 ASC]",
+ "CoalescePartitionsExec",
+ join_plan.as_str(),
+ "SortPreservingRepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[a@0 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[b1@1 ASC]",
+ "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[c@2 ASC]",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ ],
+ };
+ assert_optimized!(expected_first_sort_enforcement, top_join, false);
}
_ => {}
}
@@ -2110,17 +3016,45 @@ mod tests {
"ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]",
"ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]",
"AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]",
- "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10",
+ "AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortExec: expr=[b2@1 ASC,a2@0 ASC]",
+ "ProjectionExec: expr=[a@1 as a2, b@0 as b2]",
+ "AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]",
+ "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10",
+ "AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ ];
+ assert_optimized!(expected, join.clone(), true);
+
+ let expected_first_sort_enforcement = &[
+ "SortMergeJoin: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)]",
+ "SortPreservingRepartitionExec: partitioning=Hash([b3@1, a3@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "SortExec: expr=[b3@1 ASC,a3@0 ASC]",
+ "CoalescePartitionsExec",
+ "ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]",
+ "ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]",
+ "AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]",
+ "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+ "SortPreservingRepartitionExec: partitioning=Hash([b2@1, a2@0], 10), input_partitions=10",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"SortExec: expr=[b2@1 ASC,a2@0 ASC]",
+ "CoalescePartitionsExec",
"ProjectionExec: expr=[a@1 as a2, b@0 as b2]",
"AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]",
- "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, join);
+ assert_optimized!(expected_first_sort_enforcement, join, false);
Ok(())
}
@@ -2149,7 +3083,8 @@ mod tests {
"CoalesceBatchesExec: target_batch_size=4096",
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]",
];
- assert_optimized!(expected, exec);
+ assert_optimized!(expected, exec, true);
+ assert_optimized!(expected, exec, false);
Ok(())
}
@@ -2173,21 +3108,24 @@ mod tests {
let plan =
aggregate_exec_with_alias(plan, vec![("a1".to_string(), "a2".to_string())]);
- // Only two RepartitionExecs added, no final RepartionExec required
+ // Only two RepartitionExecs added, no final RepartitionExec required
let expected = &[
"AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]",
"AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]",
"InterleaveExec",
"AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]",
- "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
"AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]",
- "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=1",
+ "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10",
"AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
}
diff --git a/datafusion/core/src/physical_optimizer/sort_enforcement.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs
similarity index 96%
rename from datafusion/core/src/physical_optimizer/sort_enforcement.rs
rename to datafusion/core/src/physical_optimizer/enforce_sorting.rs
index 261bec7232..a976105d36 100644
--- a/datafusion/core/src/physical_optimizer/sort_enforcement.rs
+++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs
@@ -34,6 +34,10 @@
//! 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::fmt;
+use std::fmt::Formatter;
+use std::sync::Arc;
+
use crate::config::ConfigOptions;
use crate::error::Result;
use crate::physical_optimizer::replace_with_order_preserving_variants::{
@@ -41,8 +45,8 @@ use crate::physical_optimizer::replace_with_order_preserving_variants::{
};
use crate::physical_optimizer::sort_pushdown::{pushdown_sorts, SortPushDown};
use crate::physical_optimizer::utils::{
- add_sort_above, find_indices, is_coalesce_partitions, is_limit, is_repartition,
- is_sort, is_sort_preserving_merge, is_sorted, is_union, is_window,
+ add_sort_above, find_indices, get_plan_string, is_coalesce_partitions, is_limit,
+ is_repartition, is_sort, is_sort_preserving_merge, is_sorted, is_union, is_window,
merge_and_order_indices, set_difference,
};
use crate::physical_optimizer::PhysicalOptimizerRule;
@@ -53,6 +57,7 @@ use crate::physical_plan::windows::{
BoundedWindowAggExec, PartitionSearchMode, WindowAggExec,
};
use crate::physical_plan::{with_new_children_if_necessary, Distribution, ExecutionPlan};
+
use arrow::datatypes::SchemaRef;
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
use datafusion_common::utils::{get_at_indices, longest_consecutive_prefix};
@@ -62,8 +67,8 @@ use datafusion_physical_expr::utils::{
ordering_satisfy_requirement_concrete,
};
use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement};
+
use itertools::{concat, izip, Itertools};
-use std::sync::Arc;
/// This rule inspects [`SortExec`]'s in the given physical plan and removes the
/// ones it can prove unnecessary.
@@ -89,6 +94,18 @@ pub(crate) struct ExecTree {
pub children: Vec<ExecTree>,
}
+impl fmt::Display for ExecTree {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ let plan_string = get_plan_string(&self.plan);
+ write!(f, "\nidx: {:?}", self.idx)?;
+ write!(f, "\nplan: {:?}", plan_string)?;
+ for child in self.children.iter() {
+ write!(f, "\nexec_tree:{}", child)?;
+ }
+ writeln!(f)
+ }
+}
+
impl ExecTree {
/// Create new Exec tree
pub fn new(
@@ -981,13 +998,12 @@ pub(crate) fn unbounded_output(plan: &Arc<dyn ExecutionPlan>) -> bool {
#[cfg(test)]
mod tests {
use super::*;
- use crate::physical_optimizer::dist_enforcement::EnforceDistribution;
use crate::physical_optimizer::test_utils::{
aggregate_exec, bounded_window_exec, coalesce_batches_exec,
- coalesce_partitions_exec, filter_exec, get_plan_string, 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, sort_preserving_merge_exec, union_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,
+ sort_preserving_merge_exec, union_exec,
};
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::windows::PartitionSearchMode::{
@@ -996,6 +1012,7 @@ mod tests {
use crate::physical_plan::{displayable, Partitioning};
use crate::prelude::{SessionConfig, SessionContext};
use crate::test::csv_exec_sorted;
+
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion_common::Result;
@@ -1003,6 +1020,8 @@ mod tests {
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::expressions::{col, NotExpr};
use datafusion_physical_expr::PhysicalSortExpr;
+
+ use crate::physical_optimizer::enforce_distribution::EnforceDistribution;
use std::sync::Arc;
fn create_test_schema() -> Result<SchemaRef> {
@@ -2624,6 +2643,18 @@ mod tests {
let orig_plan =
Arc::new(SortExec::new(sort_exprs, repartition)) as Arc<dyn ExecutionPlan>;
+ let actual = get_plan_string(&orig_plan);
+ println!("{:?}", actual);
+ let expected_input = vec![
+ "SortExec: expr=[nullable_col@0 ASC]",
+ " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ " BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow }], mode=[Sorted]",
+ " MemoryExec: partitions=1, partition_sizes=[0]",
+ ];
+ assert_eq!(
+ expected_input, actual,
+ "\n**Original Plan Mismatch\n\nexpected:\n\n{expected_input:#?}\nactual:\n\n{actual:#?}\n\n"
+ );
let mut plan = orig_plan.clone();
let rules = vec![
@@ -2764,4 +2795,53 @@ mod tests {
assert_optimized!(expected_input, expected_optimized, physical_plan, false);
Ok(())
}
+
+ #[tokio::test]
+ async fn test_do_not_pushdown_through_spm() -> Result<()> {
+ let schema = create_test_schema3()?;
+ let sort_exprs = vec![sort_expr("a", &schema), sort_expr("b", &schema)];
+ let source = csv_exec_sorted(&schema, sort_exprs.clone(), false);
+ let repartition_rr = repartition_exec(source);
+ let spm = sort_preserving_merge_exec(sort_exprs, repartition_rr);
+ let physical_plan = sort_exec(vec![sort_expr("b", &schema)], spm);
+
+ let expected_input = ["SortExec: expr=[b@1 ASC]",
+ " SortPreservingMergeExec: [a@0 ASC,b@1 ASC]",
+ " 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",];
+ let expected_optimized = ["SortExec: expr=[b@1 ASC]",
+ " SortPreservingMergeExec: [a@0 ASC,b@1 ASC]",
+ " 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(())
+ }
+
+ #[tokio::test]
+ async fn test_pushdown_through_spm() -> Result<()> {
+ let schema = create_test_schema3()?;
+ let sort_exprs = vec![sort_expr("a", &schema), sort_expr("b", &schema)];
+ let source = csv_exec_sorted(&schema, sort_exprs.clone(), false);
+ let repartition_rr = repartition_exec(source);
+ let spm = sort_preserving_merge_exec(sort_exprs, repartition_rr);
+ let physical_plan = sort_exec(
+ vec![
+ sort_expr("a", &schema),
+ sort_expr("b", &schema),
+ sort_expr("c", &schema),
+ ],
+ spm,
+ );
+
+ let expected_input = ["SortExec: expr=[a@0 ASC,b@1 ASC,c@2 ASC]",
+ " SortPreservingMergeExec: [a@0 ASC,b@1 ASC]",
+ " 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",];
+ let expected_optimized = ["SortPreservingMergeExec: [a@0 ASC,b@1 ASC]",
+ " SortExec: expr=[a@0 ASC,b@1 ASC,c@2 ASC]",
+ " 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(())
+ }
}
diff --git a/datafusion/core/src/physical_optimizer/mod.rs b/datafusion/core/src/physical_optimizer/mod.rs
index f74d4ea0c9..1a966eb009 100644
--- a/datafusion/core/src/physical_optimizer/mod.rs
+++ b/datafusion/core/src/physical_optimizer/mod.rs
@@ -24,14 +24,14 @@
pub mod aggregate_statistics;
pub mod coalesce_batches;
pub mod combine_partial_final_agg;
-pub mod dist_enforcement;
+pub mod enforce_distribution;
+pub mod enforce_sorting;
pub mod join_selection;
pub mod optimizer;
pub mod pipeline_checker;
pub mod pruning;
pub mod repartition;
pub mod replace_with_order_preserving_variants;
-pub mod sort_enforcement;
mod sort_pushdown;
mod utils;
diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs
index 3f6698c6cf..cd78555f52 100644
--- a/datafusion/core/src/physical_optimizer/optimizer.rs
+++ b/datafusion/core/src/physical_optimizer/optimizer.rs
@@ -23,11 +23,10 @@ use crate::config::ConfigOptions;
use crate::physical_optimizer::aggregate_statistics::AggregateStatistics;
use crate::physical_optimizer::coalesce_batches::CoalesceBatches;
use crate::physical_optimizer::combine_partial_final_agg::CombinePartialFinalAggregate;
-use crate::physical_optimizer::dist_enforcement::EnforceDistribution;
+use crate::physical_optimizer::enforce_distribution::EnforceDistribution;
+use crate::physical_optimizer::enforce_sorting::EnforceSorting;
use crate::physical_optimizer::join_selection::JoinSelection;
use crate::physical_optimizer::pipeline_checker::PipelineChecker;
-use crate::physical_optimizer::repartition::Repartition;
-use crate::physical_optimizer::sort_enforcement::EnforceSorting;
use crate::{error::Result, physical_plan::ExecutionPlan};
/// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which
@@ -75,16 +74,10 @@ impl PhysicalOptimizer {
// repartitioning and local sorting steps to meet distribution and ordering requirements.
// Therefore, it should run before EnforceDistribution and EnforceSorting.
Arc::new(JoinSelection::new()),
- // In order to increase the parallelism, the Repartition rule will change the
- // output partitioning of some operators in the plan tree, which will influence
- // other rules. Therefore, it should run as soon as possible. It is optional because:
- // - It's not used for the distributed engine, Ballista.
- // - It's conflicted with some parts of the EnforceDistribution, since it will
- // introduce additional repartitioning while EnforceDistribution aims to
- // reduce unnecessary repartitioning.
- Arc::new(Repartition::new()),
- // The EnforceDistribution rule is for adding essential repartition to satisfy the required
- // distribution. Please make sure that the whole plan tree is determined before this rule.
+ // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
+ // requirements. Please make sure that the whole plan tree is determined before this rule.
+ // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
+ // least one of the operators in the plan benefits from increased parallelism.
Arc::new(EnforceDistribution::new()),
// The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
Arc::new(CombinePartialFinalAggregate::new()),
diff --git a/datafusion/core/src/physical_optimizer/repartition.rs b/datafusion/core/src/physical_optimizer/repartition.rs
index 97cbc7ca5b..93bde05546 100644
--- a/datafusion/core/src/physical_optimizer/repartition.rs
+++ b/datafusion/core/src/physical_optimizer/repartition.rs
@@ -15,311 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-//! Repartition optimizer that introduces repartition nodes to increase the level of parallelism available
-use datafusion_common::tree_node::Transformed;
-use std::sync::Arc;
-
-use super::optimizer::PhysicalOptimizerRule;
-use crate::config::ConfigOptions;
-use crate::datasource::physical_plan::{CsvExec, ParquetExec};
-use crate::error::Result;
-use crate::physical_plan::Partitioning::*;
-use crate::physical_plan::{
- repartition::RepartitionExec, with_new_children_if_necessary, ExecutionPlan,
-};
-
-/// Optimizer that introduces repartition to introduce more
-/// parallelism in the plan
-///
-/// For example, given an input such as:
-///
-///
-/// ```text
-/// ┌─────────────────────────────────┐
-/// │ │
-/// │ ExecutionPlan │
-/// │ │
-/// └─────────────────────────────────┘
-/// ▲ ▲
-/// │ │
-/// ┌─────┘ └─────┐
-/// │ │
-/// │ │
-/// │ │
-/// ┌───────────┐ ┌───────────┐
-/// │ │ │ │
-/// │ batch A1 │ │ batch B1 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A2 │ │ batch B2 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A3 │ │ batch B3 │
-/// │ │ │ │
-/// └───────────┘ └───────────┘
-///
-/// Input Input
-/// A B
-/// ```
-///
-/// This optimizer will attempt to add a `RepartitionExec` to increase
-/// the parallelism (to 3 in this case)
-///
-/// ```text
-/// ┌─────────────────────────────────┐
-/// │ │
-/// │ ExecutionPlan │
-/// │ │
-/// └─────────────────────────────────┘
-/// ▲ ▲ ▲ Input now has 3
-/// │ │ │ partitions
-/// ┌───────┘ │ └───────┐
-/// │ │ │
-/// │ │ │
-/// ┌───────────┐ ┌───────────┐ ┌───────────┐
-/// │ │ │ │ │ │
-/// │ batch A1 │ │ batch A3 │ │ batch B3 │
-/// │ │ │ │ │ │
-/// ├───────────┤ ├───────────┤ ├───────────┤
-/// │ │ │ │ │ │
-/// │ batch B2 │ │ batch B1 │ │ batch A2 │
-/// │ │ │ │ │ │
-/// └───────────┘ └───────────┘ └───────────┘
-/// ▲ ▲ ▲
-/// │ │ │
-/// └─────────┐ │ ┌──────────┘
-/// │ │ │
-/// │ │ │
-/// ┌─────────────────────────────────┐ batches are
-/// │ RepartitionExec(3) │ repartitioned
-/// │ RoundRobin │
-/// │ │
-/// └─────────────────────────────────┘
-/// ▲ ▲
-/// │ │
-/// ┌─────┘ └─────┐
-/// │ │
-/// │ │
-/// │ │
-/// ┌───────────┐ ┌───────────┐
-/// │ │ │ │
-/// │ batch A1 │ │ batch B1 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A2 │ │ batch B2 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A3 │ │ batch B3 │
-/// │ │ │ │
-/// └───────────┘ └───────────┘
-///
-///
-/// Input Input
-/// A B
-/// ```
-#[derive(Default)]
-pub struct Repartition {}
-
-impl Repartition {
- #[allow(missing_docs)]
- pub fn new() -> Self {
- Self {}
- }
-}
-
-/// Recursively attempts to increase the overall parallelism of the
-/// plan, while respecting ordering, by adding a `RepartitionExec` at
-/// the output of `plan` if it would help parallelism and not destroy
-/// any possibly useful ordering.
-///
-/// It does so using a depth first scan of the tree, and repartitions
-/// any plan that:
-///
-/// 1. Has fewer partitions than `target_partitions`
-///
-/// 2. Has a direct parent that `benefits_from_input_partitioning`
-///
-/// 3. Does not destroy any existing sort order if the parent is
-/// relying on it.
-///
-/// if `can_reorder` is false, it means the parent node of `plan` is
-/// trying to take advantage of the output sort order of plan, so it
-/// should not be repartitioned if doing so would destroy the output
-/// sort order.
-///
-/// (Parent) - If can_reorder is false, means this parent node is
-/// trying to use the sort ouder order this plan. If true
-/// means parent doesn't care about sort order
-///
-/// (plan) - We are deciding to add a partition above here
-///
-/// (children) - Recursively visit all children first
-///
-/// If 'would_benefit` is true, the upstream operator would benefit
-/// from additional partitions and thus repatitioning is considered.
-///
-/// if `is_root` is true, no repartition is added.
-fn optimize_partitions(
- target_partitions: usize,
- plan: Arc<dyn ExecutionPlan>,
- is_root: bool,
- can_reorder: bool,
- would_benefit: bool,
- repartition_file_scans: bool,
- repartition_file_min_size: usize,
-) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
- // Recurse into children bottom-up (attempt to repartition as
- // early as possible)
- let new_plan = if plan.children().is_empty() {
- // leaf node - don't replace children
- Transformed::No(plan)
- } else {
- let children = plan
- .children()
- .iter()
- .enumerate()
- .map(|(idx, child)| {
- // Does plan itself (not its parent) require its input to
- // be sorted in some way?
- let required_input_ordering =
- plan_has_required_input_ordering(plan.as_ref());
-
- // We can reorder a child if:
- // - It has no ordering to preserve, or
- // - Its parent has no required input ordering and does not
- // maintain input ordering.
- // Check if this condition holds:
- let can_reorder_child = child.output_ordering().is_none()
- || (!required_input_ordering
- && (can_reorder || !plan.maintains_input_order()[idx]));
-
- optimize_partitions(
- target_partitions,
- child.clone(),
- false, // child is not root
- can_reorder_child,
- plan.benefits_from_input_partitioning()[idx],
- repartition_file_scans,
- repartition_file_min_size,
- )
- .map(Transformed::into)
- })
- .collect::<Result<_>>()?;
- with_new_children_if_necessary(plan, children)?
- };
-
- let (new_plan, transformed) = new_plan.into_pair();
-
- // decide if we should bother trying to repartition the output of this plan
- let mut could_repartition = match new_plan.output_partitioning() {
- // Apply when underlying node has less than `self.target_partitions` amount of concurrency
- RoundRobinBatch(x) => x < target_partitions,
- UnknownPartitioning(x) => x < target_partitions,
- // we don't want to introduce partitioning after hash partitioning
- // as the plan will likely depend on this
- Hash(_, _) => false,
- };
-
- // Don't need to apply when the returned row count is not greater than 1
- let stats = new_plan.statistics();
- if stats.is_exact {
- could_repartition = could_repartition
- && stats.num_rows.map(|num_rows| num_rows > 1).unwrap_or(true);
- }
-
- // don't repartition root of the plan
- if is_root {
- could_repartition = false;
- }
-
- let repartition_allowed = would_benefit && could_repartition && can_reorder;
-
- // If repartition is not allowed - return plan as it is
- if !repartition_allowed {
- return Ok(if transformed {
- Transformed::Yes(new_plan)
- } else {
- Transformed::No(new_plan)
- });
- }
-
- // For ParquetExec return internally repartitioned version of the plan in case `repartition_file_scans` is set
- if let Some(parquet_exec) = new_plan.as_any().downcast_ref::<ParquetExec>() {
- if repartition_file_scans {
- return Ok(Transformed::Yes(Arc::new(
- parquet_exec
- .get_repartitioned(target_partitions, repartition_file_min_size),
- )));
- }
- }
-
- if let Some(csv_exec) = new_plan.as_any().downcast_ref::<CsvExec>() {
- if repartition_file_scans {
- let repartitioned_exec_option =
- csv_exec.get_repartitioned(target_partitions, repartition_file_min_size);
- if let Some(repartitioned_exec) = repartitioned_exec_option {
- return Ok(Transformed::Yes(Arc::new(repartitioned_exec)));
- }
- }
- }
-
- // Otherwise - return plan wrapped up in RepartitionExec
- Ok(Transformed::Yes(Arc::new(RepartitionExec::try_new(
- new_plan,
- RoundRobinBatch(target_partitions),
- )?)))
-}
-
-/// Returns true if `plan` requires any of inputs to be sorted in some
-/// way for correctness. If this is true, its output should not be
-/// repartitioned if it would destroy the required order.
-fn plan_has_required_input_ordering(plan: &dyn ExecutionPlan) -> bool {
- // NB: checking `is_empty()` is not the right check!
- plan.required_input_ordering().iter().any(Option::is_some)
-}
-
-impl PhysicalOptimizerRule for Repartition {
- fn optimize(
- &self,
- plan: Arc<dyn ExecutionPlan>,
- config: &ConfigOptions,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- let target_partitions = config.execution.target_partitions;
- let enabled = config.optimizer.enable_round_robin_repartition;
- let repartition_file_scans = config.optimizer.repartition_file_scans;
- let repartition_file_min_size = config.optimizer.repartition_file_min_size;
- // Don't run optimizer if target_partitions == 1
- if !enabled || target_partitions == 1 {
- Ok(plan)
- } else {
- let is_root = true;
- let can_reorder = plan.output_ordering().is_none();
- let would_benefit = false;
- optimize_partitions(
- target_partitions,
- plan.clone(),
- is_root,
- can_reorder,
- would_benefit,
- repartition_file_scans,
- repartition_file_min_size,
- )
- .map(Transformed::into)
- }
- }
-
- fn name(&self) -> &str {
- "repartition"
- }
-
- fn schema_check(&self) -> bool {
- true
- }
-}
+//! Tests to check whether parallelism can be increased as desired
#[cfg(test)]
#[ctor::ctor]
@@ -329,15 +25,15 @@ fn init() {
#[cfg(test)]
mod tests {
- use arrow::compute::SortOptions;
- use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+ use std::sync::Arc;
- use super::*;
use crate::datasource::listing::PartitionedFile;
use crate::datasource::object_store::ObjectStoreUrl;
- use crate::datasource::physical_plan::{FileScanConfig, ParquetExec};
- use crate::physical_optimizer::dist_enforcement::EnforceDistribution;
- use crate::physical_optimizer::sort_enforcement::EnforceSorting;
+ use crate::datasource::physical_plan::{CsvExec, FileScanConfig, ParquetExec};
+ use crate::error::Result;
+ use crate::physical_optimizer::enforce_distribution::EnforceDistribution;
+ use crate::physical_optimizer::enforce_sorting::EnforceSorting;
+ use crate::physical_optimizer::PhysicalOptimizerRule;
use crate::physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
@@ -348,7 +44,13 @@ mod tests {
use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use crate::physical_plan::union::UnionExec;
+ use crate::physical_plan::ExecutionPlan;
use crate::physical_plan::{displayable, DisplayAs, DisplayFormatType, Statistics};
+
+ use crate::physical_optimizer::test_utils::repartition_exec;
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+ use datafusion_common::config::ConfigOptions;
use datafusion_common::FileCompressionType;
use datafusion_physical_expr::PhysicalSortRequirement;
@@ -356,6 +58,13 @@ mod tests {
Arc::new(Schema::new(vec![Field::new("c1", DataType::Boolean, true)]))
}
+ fn schema2() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("c1", DataType::Boolean, true),
+ Field::new("c2", DataType::Boolean, true),
+ ]))
+ }
+
/// Generate FileScanConfig for file scan executors like 'ParquetExec'
fn scan_config(sorted: bool, single_file: bool) -> FileScanConfig {
let sort_exprs = vec![PhysicalSortExpr {
@@ -385,6 +94,35 @@ mod tests {
}
}
+ /// Generate FileScanConfig for file scan executors like 'ParquetExec'
+ fn scan_config2(sorted: bool, single_file: bool) -> FileScanConfig {
+ let sort_exprs = vec![PhysicalSortExpr {
+ expr: col("c1", &schema2()).unwrap(),
+ options: SortOptions::default(),
+ }];
+
+ let file_groups = if single_file {
+ vec![vec![PartitionedFile::new("x".to_string(), 100)]]
+ } else {
+ vec![
+ vec![PartitionedFile::new("x".to_string(), 100)],
+ vec![PartitionedFile::new("y".to_string(), 200)],
+ ]
+ };
+
+ FileScanConfig {
+ object_store_url: ObjectStoreUrl::parse("test:///").unwrap(),
+ file_schema: schema2(),
+ file_groups,
+ statistics: Statistics::default(),
+ projection: None,
+ limit: None,
+ table_partition_cols: vec![],
+ output_ordering: if sorted { vec![sort_exprs] } else { vec![] },
+ infinite_source: false,
+ }
+ }
+
/// Create a non sorted parquet exec
fn parquet_exec() -> Arc<ParquetExec> {
Arc::new(ParquetExec::new(scan_config(false, true), None, None))
@@ -441,6 +179,11 @@ mod tests {
Arc::new(ParquetExec::new(scan_config(true, false), None, None))
}
+ // Created a sorted parquet exec with multiple files
+ fn parquet_exec_multiple_sorted2() -> Arc<ParquetExec> {
+ Arc::new(ParquetExec::new(scan_config2(true, false), None, None))
+ }
+
fn sort_preserving_merge_exec(
input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
@@ -452,6 +195,17 @@ mod tests {
Arc::new(SortPreservingMergeExec::new(expr, input))
}
+ fn sort_preserving_merge_exec2(
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let expr = vec![PhysicalSortExpr {
+ expr: col("c2", &schema2()).unwrap(),
+ options: arrow::compute::SortOptions::default(),
+ }];
+
+ Arc::new(SortPreservingMergeExec::new(expr, input))
+ }
+
fn filter_exec(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
Arc::new(FilterExec::try_new(col("c1", &schema()).unwrap(), input).unwrap())
}
@@ -534,31 +288,50 @@ mod tests {
/// Runs the repartition optimizer and asserts the plan against the expected
macro_rules! assert_optimized {
- ($EXPECTED_LINES: expr, $PLAN: expr) => {
- assert_optimized!($EXPECTED_LINES, $PLAN, 10, false, 1024);
+ ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr) => {
+ assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, 10, false, 1024, false);
+ };
+
+ ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr, $TARGET_PARTITIONS: expr, $REPARTITION_FILE_SCANS: expr, $REPARTITION_FILE_MIN_SIZE: expr) => {
+ assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, $TARGET_PARTITIONS, $REPARTITION_FILE_SCANS, $REPARTITION_FILE_MIN_SIZE, false);
};
- ($EXPECTED_LINES: expr, $PLAN: expr, $TARGET_PARTITIONS: expr, $REPARTITION_FILE_SCANS: expr, $REPARTITION_FILE_MIN_SIZE: expr) => {
+ ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr, $TARGET_PARTITIONS: expr, $REPARTITION_FILE_SCANS: expr, $REPARTITION_FILE_MIN_SIZE: expr, $BOUNDED_ORDER_PRESERVING_VARIANTS: expr) => {
let expected_lines: Vec<&str> = $EXPECTED_LINES.iter().map(|s| *s).collect();
let mut config = ConfigOptions::new();
config.execution.target_partitions = $TARGET_PARTITIONS;
config.optimizer.repartition_file_scans = $REPARTITION_FILE_SCANS;
config.optimizer.repartition_file_min_size = $REPARTITION_FILE_MIN_SIZE;
-
- // run optimizer
- let optimizers: Vec<Arc<dyn PhysicalOptimizerRule + Sync + Send>> = vec![
- Arc::new(Repartition::new()),
- // EnforceDistribution is an essential rule to be applied.
- // Otherwise, the correctness of the generated optimized plan cannot be guaranteed
- Arc::new(EnforceDistribution::new()),
- // EnforceSorting is an essential rule to be applied.
- // Otherwise, the correctness of the generated optimized plan cannot be guaranteed
- Arc::new(EnforceSorting::new()),
- ];
- let optimized = optimizers.into_iter().fold($PLAN, |plan, optimizer| {
- optimizer.optimize(plan, &config).unwrap()
- });
+ config.optimizer.bounded_order_preserving_variants = $BOUNDED_ORDER_PRESERVING_VARIANTS;
+
+ let optimized = if $FIRST_ENFORCE_DIST{
+ // Run the optimizer to first apply distribution enforcement
+ // and then sort enforcement.
+ let optimizers: Vec<Arc<dyn PhysicalOptimizerRule + Sync + Send>> = vec![
+ // Verify that distribution enforcement is idempotent by running it twice.
+ Arc::new(EnforceDistribution::new()),
+ Arc::new(EnforceDistribution::new()),
+ Arc::new(EnforceSorting::new()),
+ ];
+ let optimized = optimizers.into_iter().fold($PLAN, |plan, optimizer| {
+ optimizer.optimize(plan, &config).unwrap()
+ });
+ optimized
+ } else {
+ // Run the optimizer to first apply sort enforcement and then
+ // distribution enforcement.
+ let optimizers: Vec<Arc<dyn PhysicalOptimizerRule + Sync + Send>> = vec![
+ Arc::new(EnforceSorting::new()),
+ // Verify that distribution enforcement is idempotent by running it twice.
+ Arc::new(EnforceDistribution::new()),
+ Arc::new(EnforceDistribution::new()),
+ ];
+ let optimized = optimizers.into_iter().fold($PLAN, |plan, optimizer| {
+ optimizer.optimize(plan, &config).unwrap()
+ });
+ optimized
+ };
// Now format correctly
let plan = displayable(optimized.as_ref()).indent(true).to_string();
@@ -584,7 +357,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -601,7 +375,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -619,7 +394,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -637,7 +413,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -653,25 +430,27 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
#[test]
fn repartition_sorted_limit_with_filter() -> Result<()> {
- let plan = limit_exec(filter_exec(sort_exec(parquet_exec(), false)));
+ let plan = sort_required_exec(filter_exec(sort_exec(parquet_exec(), false)));
let expected = &[
- "GlobalLimitExec: skip=0, fetch=100",
- "LocalLimitExec: fetch=100",
+ "SortRequiredExec",
"FilterExec: c1@0",
- // data is sorted so can't repartition here even though
- // filter would benefit from parallelism, the answers might be wrong
+ // We can use repartition here, ordering requirement by SortRequiredExec
+ // is still satisfied.
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"SortExec: expr=[c1@0 ASC]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -696,7 +475,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -723,7 +503,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -743,7 +524,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -754,13 +536,13 @@ mod tests {
// need repartiton and resort as the data was not sorted correctly
let expected = &[
- "SortPreservingMergeExec: [c1@0 ASC]",
"SortExec: expr=[c1@0 ASC]",
- "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
+
Ok(())
}
@@ -775,7 +557,8 @@ mod tests {
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[c1], output_ordering=[c1@0 ASC]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -793,7 +576,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1], output_ordering=[c1@0 ASC]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -810,7 +594,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1], output_ordering=[c1@0 ASC]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -843,7 +628,8 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -854,14 +640,13 @@ mod tests {
// needs to repartition / sort as the data was not sorted correctly
let expected = &[
- "SortPreservingMergeExec: [c1@0 ASC]",
"SortExec: expr=[c1@0 ASC]",
- "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
"ProjectionExec: expr=[c1@0 as c1]",
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -878,7 +663,8 @@ mod tests {
"ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[c1], output_ordering=[c1@0 ASC]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
@@ -893,14 +679,14 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+ assert_optimized!(expected, plan, false);
Ok(())
}
#[test]
fn repartition_transitively_past_sort_with_filter() -> Result<()> {
- let plan =
- sort_preserving_merge_exec(sort_exec(filter_exec(parquet_exec()), true));
+ let plan = sort_exec(filter_exec(parquet_exec()), false);
let expected = &[
"SortPreservingMergeExec: [c1@0 ASC]",
@@ -911,16 +697,23 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+
+ let expected_first_sort_enforcement = &[
+ "SortExec: expr=[c1@0 ASC]",
+ "CoalescePartitionsExec",
+ "FilterExec: c1@0",
+ // Expect repartition on the input of the filter (as it can benefit from additional parallelism)
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
+ ];
+ assert_optimized!(expected_first_sort_enforcement, plan, false);
Ok(())
}
#[test]
fn repartition_transitively_past_sort_with_projection_and_filter() -> Result<()> {
- let plan = sort_preserving_merge_exec(sort_exec(
- projection_exec(filter_exec(parquet_exec())),
- true,
- ));
+ let plan = sort_exec(projection_exec(filter_exec(parquet_exec())), false);
let expected = &[
"SortPreservingMergeExec: [c1@0 ASC]",
@@ -933,7 +726,17 @@ mod tests {
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
- assert_optimized!(expected, plan);
+ assert_optimized!(expected, plan.clone(), true);
+
+ let expected_first_sort_enforcement = &[
+ "SortExec: expr=[c1@0 ASC]",
+ "CoalescePartitionsExec",
+ "ProjectionExec: expr=[c1@0 as c1]",
+ "FilterExec: c1@0",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
+ ];
+ assert_optimized!(expected_first_sort_enforcement, plan, false);
Ok(())
}
@@ -955,8 +758,8 @@ mod tests {
"CsvExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1003,7 +806,7 @@ mod tests {
compression_type,
)));
- assert_optimized!(expected, plan, 2, true, 10);
+ assert_optimized!(expected, plan, true, 2, true, 10);
}
Ok(())
}
@@ -1028,8 +831,8 @@ mod tests {
"CsvExec: file_groups={2 groups: [[x], [y]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1053,8 +856,8 @@ mod tests {
"CsvExec: file_groups={4 groups: [[x:0..75], [x:75..100, y:0..50], [y:50..125], [y:125..200]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 4, true, 10);
- assert_optimized!(expected_csv, plan_csv, 4, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 4, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 4, true, 10);
Ok(())
}
@@ -1080,8 +883,8 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1092,27 +895,31 @@ mod tests {
let expected_parquet = &[
"GlobalLimitExec: skip=0, fetch=100",
+ "CoalescePartitionsExec",
"LocalLimitExec: fetch=100",
"FilterExec: c1@0",
- // data is sorted so can't repartition here even though
- // filter would benefit from parallelism, the answers might be wrong
+ // even though data is sorted, we can use repartition here. Since
+ // ordering is not used in subsequent stages anyway.
+ "RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1",
"SortExec: expr=[c1@0 ASC]",
// SortExec doesn't benefit from input partitioning
"ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
];
let expected_csv = &[
"GlobalLimitExec: skip=0, fetch=100",
+ "CoalescePartitionsExec",
"LocalLimitExec: fetch=100",
"FilterExec: c1@0",
- // data is sorted so can't repartition here even though
- // filter would benefit from parallelism, the answers might be wrong
+ // even though data is sorted, we can use repartition here. Since
+ // ordering is not used in subsequent stages anyway.
+ "RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1",
"SortExec: expr=[c1@0 ASC]",
// SortExec doesn't benefit from input partitioning
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1154,8 +961,8 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1183,8 +990,8 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1202,8 +1009,8 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], output_ordering=[c1@0 ASC], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1229,8 +1036,8 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], output_ordering=[c1@0 ASC], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1251,8 +1058,8 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], output_ordering=[c1@0 ASC], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
Ok(())
}
@@ -1273,8 +1080,100 @@ mod tests {
"CsvExec: file_groups={1 group: [[x]]}, projection=[c1], output_ordering=[c1@0 ASC], has_header=false",
];
- assert_optimized!(expected_parquet, plan_parquet, 2, true, 10);
- assert_optimized!(expected_csv, plan_csv, 2, true, 10);
+ assert_optimized!(expected_parquet, plan_parquet, true, 2, true, 10);
+ assert_optimized!(expected_csv, plan_csv, true, 2, true, 10);
+ Ok(())
+ }
+
+ #[test]
+ fn remove_redundant_roundrobins() -> Result<()> {
+ let input = parquet_exec();
+ let repartition = repartition_exec(repartition_exec(input));
+ let physical_plan = repartition_exec(filter_exec(repartition));
+
+ let expected = &[
+ "FilterExec: c1@0",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1",
+ "ParquetExec: file_groups={1 group: [[x]]}, projection=[c1]",
+ ];
+
+ assert_optimized!(expected, physical_plan.clone(), true, 2, true, 10);
+ assert_optimized!(expected, physical_plan, false, 2, true, 10);
+
+ Ok(())
+ }
+
+ #[test]
+ fn preserve_ordering_through_repartition() -> Result<()> {
+ let input = parquet_exec_multiple_sorted();
+ let physical_plan = sort_preserving_merge_exec(filter_exec(input));
+
+ let expected = &[
+ "SortPreservingMergeExec: [c1@0 ASC]",
+ "FilterExec: c1@0",
+ "SortPreservingRepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+ "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[c1], output_ordering=[c1@0 ASC]",
+ ];
+
+ // last flag sets config.optimizer.bounded_order_preserving_variants
+ assert_optimized!(expected, physical_plan.clone(), true, 10, false, 10, true);
+ assert_optimized!(expected, physical_plan, false, 10, false, 10, true);
+
+ Ok(())
+ }
+
+ #[test]
+ fn do_not_preserve_ordering_through_repartition() -> Result<()> {
+ let input = parquet_exec_multiple_sorted();
+ let physical_plan = sort_preserving_merge_exec(filter_exec(input));
+
+ let expected = &[
+ "SortPreservingMergeExec: [c1@0 ASC]",
+ "SortExec: expr=[c1@0 ASC]",
+ "FilterExec: c1@0",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+ "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[c1], output_ordering=[c1@0 ASC]",
+ ];
+
+ assert_optimized!(expected, physical_plan.clone(), true, 10, false, 10);
+ assert_optimized!(expected, physical_plan, false, 10, false, 10);
+
+ Ok(())
+ }
+
+ #[test]
+ fn do_not_preserve_ordering_through_repartition2() -> Result<()> {
+ let input = parquet_exec_multiple_sorted2();
+ let physical_plan = sort_preserving_merge_exec2(filter_exec(input));
+
+ let expected = &[
+ "SortPreservingMergeExec: [c2@1 ASC]",
+ "SortExec: expr=[c2@1 ASC]",
+ "FilterExec: c1@0",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+ "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[c1, c2], output_ordering=[c1@0 ASC]",
+ ];
+
+ assert_optimized!(expected, physical_plan.clone(), true, 10, false, 10);
+ assert_optimized!(expected, physical_plan, false, 10, false, 10);
+
+ Ok(())
+ }
+
+ #[test]
+ fn do_not_preserve_ordering_through_repartition3() -> Result<()> {
+ let input = parquet_exec_multiple_sorted();
+ let physical_plan = filter_exec(input);
+
+ let expected = &[
+ "FilterExec: c1@0",
+ "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
+ "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[c1], output_ordering=[c1@0 ASC]",
+ ];
+
+ assert_optimized!(expected, physical_plan.clone(), true, 10, false, 10);
+ assert_optimized!(expected, physical_plan, false, 10, false, 10);
+
Ok(())
}
@@ -1314,6 +1213,10 @@ mod tests {
self.input.output_partitioning()
}
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ vec![false]
+ }
+
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
self.input.output_ordering()
}
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 231fbf4074..59ded45f7b 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
@@ -20,7 +20,7 @@
//! performance or to accommodate unbounded streams by fixing the pipeline.
use crate::error::Result;
-use crate::physical_optimizer::sort_enforcement::{unbounded_output, ExecTree};
+use crate::physical_optimizer::enforce_sorting::{unbounded_output, ExecTree};
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;
diff --git a/datafusion/core/src/physical_optimizer/sort_pushdown.rs b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
index aaa2b6de2b..81c2e76b74 100644
--- a/datafusion/core/src/physical_optimizer/sort_pushdown.rs
+++ b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
@@ -17,7 +17,9 @@
use std::sync::Arc;
-use crate::physical_optimizer::utils::{add_sort_above, is_limit, is_union, is_window};
+use crate::physical_optimizer::utils::{
+ add_sort_above, 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, JoinSide};
use crate::physical_plan::joins::{HashJoinExec, SortMergeJoinExec};
@@ -31,7 +33,7 @@ use datafusion_common::{plan_err, DataFusionError, Result};
use datafusion_expr::JoinType;
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::utils::{
- ordering_satisfy_requirement, requirements_compatible,
+ ordering_satisfy, ordering_satisfy_requirement, requirements_compatible,
};
use datafusion_physical_expr::{PhysicalSortExpr, PhysicalSortRequirement};
@@ -41,7 +43,7 @@ use itertools::izip;
/// down [`SortExec`] in the plan. In some cases, we can reduce the total
/// computational cost by pushing down `SortExec`s through some executors.
///
-/// [`EnforceSorting`]: crate::physical_optimizer::sort_enforcement::EnforceSorting
+/// [`EnforceSorting`]: crate::physical_optimizer::enforce_sorting::EnforceSorting
#[derive(Debug, Clone)]
pub(crate) struct SortPushDown {
/// Current plan
@@ -260,6 +262,17 @@ fn pushdown_requirement_to_children(
|| plan.as_any().is::<ProjectionExec>()
|| is_limit(plan)
|| plan.as_any().is::<HashJoinExec>()
+ // Do not push-down through SortPreservingMergeExec when
+ // ordering requirement invalidates requirement of sort preserving merge exec.
+ || (is_sort_preserving_merge(plan) && !ordering_satisfy(
+ parent_required
+ .map(|req| PhysicalSortRequirement::to_sort_exprs(req.to_vec()))
+ .as_deref(),
+ plan.output_ordering(),
+ || plan.equivalence_properties(),
+ || plan.ordering_equivalence_properties(),
+ )
+ )
{
// If the current plan is a leaf node or can not maintain any of the input ordering, can not pushed down requirements.
// For RepartitionExec, we always choose to not push down the sort requirements even the RepartitionExec(input_partition=1) could maintain input ordering.
diff --git a/datafusion/core/src/physical_optimizer/test_utils.rs b/datafusion/core/src/physical_optimizer/test_utils.rs
index 5c8d8b1e64..e021cda2c8 100644
--- a/datafusion/core/src/physical_optimizer/test_utils.rs
+++ b/datafusion/core/src/physical_optimizer/test_utils.rs
@@ -17,6 +17,8 @@
//! Collection of testing utility functions that are leveraged by the query optimizer rules
+use std::sync::Arc;
+
use crate::datasource::listing::PartitionedFile;
use crate::datasource::physical_plan::{FileScanConfig, ParquetExec};
use crate::error::Result;
@@ -33,16 +35,17 @@ use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use crate::physical_plan::union::UnionExec;
use crate::physical_plan::windows::create_window_expr;
-use crate::physical_plan::{displayable, ExecutionPlan, Partitioning};
+use crate::physical_plan::{ExecutionPlan, Partitioning};
use crate::prelude::{CsvReadOptions, SessionContext};
+
use arrow_schema::{Schema, SchemaRef, SortOptions};
-use async_trait::async_trait;
use datafusion_common::{JoinType, Statistics};
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_expr::{AggregateFunction, WindowFrame, WindowFunction};
use datafusion_physical_expr::expressions::col;
use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
-use std::sync::Arc;
+
+use async_trait::async_trait;
async fn register_current_csv(
ctx: &SessionContext,
@@ -251,13 +254,6 @@ pub fn filter_exec(
Arc::new(FilterExec::try_new(predicate, input).unwrap())
}
-// Util function to get string representation of a physical plan
-pub fn get_plan_string(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
- let formatted = displayable(plan.as_ref()).indent(true).to_string();
- let actual: Vec<&str> = formatted.trim().lines().collect();
- actual.iter().map(|elem| elem.to_string()).collect()
-}
-
pub fn sort_preserving_merge_exec(
sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
input: Arc<dyn ExecutionPlan>,
diff --git a/datafusion/core/src/physical_optimizer/utils.rs b/datafusion/core/src/physical_optimizer/utils.rs
index 82cfe76b68..5d150ae60f 100644
--- a/datafusion/core/src/physical_optimizer/utils.rs
+++ b/datafusion/core/src/physical_optimizer/utils.rs
@@ -29,7 +29,8 @@ use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use crate::physical_plan::union::UnionExec;
use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
-use crate::physical_plan::ExecutionPlan;
+use crate::physical_plan::{displayable, ExecutionPlan};
+
use datafusion_common::DataFusionError;
use datafusion_physical_expr::utils::ordering_satisfy;
use datafusion_physical_expr::PhysicalSortExpr;
@@ -154,6 +155,13 @@ pub fn is_repartition(plan: &Arc<dyn ExecutionPlan>) -> bool {
plan.as_any().is::<RepartitionExec>()
}
+/// Utility function yielding a string representation of the given [`ExecutionPlan`].
+pub fn get_plan_string(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
+ let formatted = displayable(plan.as_ref()).indent(true).to_string();
+ let actual: Vec<&str> = formatted.trim().lines().collect();
+ actual.iter().map(|elem| elem.to_string()).collect()
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/datafusion/core/src/physical_plan/coalesce_batches.rs b/datafusion/core/src/physical_plan/coalesce_batches.rs
index 71ad39658f..cc32d4163b 100644
--- a/datafusion/core/src/physical_plan/coalesce_batches.rs
+++ b/datafusion/core/src/physical_plan/coalesce_batches.rs
@@ -132,6 +132,10 @@ impl ExecutionPlan for CoalesceBatchesExec {
vec![true]
}
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ vec![false]
+ }
+
fn equivalence_properties(&self) -> EquivalenceProperties {
self.input.equivalence_properties()
}
diff --git a/datafusion/core/src/physical_plan/repartition/mod.rs b/datafusion/core/src/physical_plan/repartition/mod.rs
index b23b495674..a69f33aa19 100644
--- a/datafusion/core/src/physical_plan/repartition/mod.rs
+++ b/datafusion/core/src/physical_plan/repartition/mod.rs
@@ -637,7 +637,14 @@ impl RepartitionExec {
/// Set Order preserving flag
pub fn with_preserve_order(mut self, preserve_order: bool) -> Self {
- self.preserve_order = preserve_order;
+ // Set "preserve order" mode only if the input partition count is larger than 1
+ // Because in these cases naive `RepartitionExec` cannot maintain ordering. Using
+ // `SortPreservingRepartitionExec` is necessity. However, when input partition number
+ // is 1, `RepartitionExec` can maintain ordering. In this case, we don't need to use
+ // `SortPreservingRepartitionExec` variant to maintain ordering.
+ if self.input.output_partitioning().partition_count() > 1 {
+ self.preserve_order = preserve_order
+ }
self
}
diff --git a/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs b/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs
index 5b12f1cdd9..c9f693a8a2 100644
--- a/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs
+++ b/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs
@@ -158,6 +158,10 @@ impl ExecutionPlan for SortPreservingMergeExec {
vec![Distribution::UnspecifiedDistribution]
}
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ vec![false]
+ }
+
fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> {
vec![Some(PhysicalSortRequirement::from_sort_exprs(&self.expr))]
}
diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs
index 1aaa01bf4f..d08f09d3b6 100644
--- a/datafusion/core/tests/sql/joins.rs
+++ b/datafusion/core/tests/sql/joins.rs
@@ -127,17 +127,19 @@ async fn join_change_in_planner() -> Result<()> {
[
"SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(a2@1, a2@1)], filter=CAST(a1@0 AS Int64) > CAST(a1@1 AS Int64) + 3 AND CAST(a1@0 AS Int64) < CAST(a1@1 AS Int64) + 10",
" CoalesceBatchesExec: target_batch_size=8192",
- " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=1",
- // " CsvExec: file_groups={1 group: [[tempdir/left.csv]]}, projection=[a1, a2], has_header=false",
+ " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=8",
+ " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
+ // " CsvExec: file_groups={1 group: [[tempdir/left.csv]]}, projection=[a1, a2], has_header=false",
" CoalesceBatchesExec: target_batch_size=8192",
- " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=1",
- // " CsvExec: file_groups={1 group: [[tempdir/right.csv]]}, projection=[a1, a2], has_header=false"
+ " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=8",
+ " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
+ // " CsvExec: file_groups={1 group: [[tempdir/right.csv]]}, projection=[a1, a2], has_header=false"
]
};
let mut actual: Vec<&str> = formatted.trim().lines().collect();
// Remove CSV lines
- actual.remove(3);
- actual.remove(5);
+ actual.remove(4);
+ actual.remove(7);
assert_eq!(
expected,
@@ -180,17 +182,19 @@ async fn join_change_in_planner_without_sort() -> Result<()> {
[
"SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(a2@1, a2@1)], filter=CAST(a1@0 AS Int64) > CAST(a1@1 AS Int64) + 3 AND CAST(a1@0 AS Int64) < CAST(a1@1 AS Int64) + 10",
" CoalesceBatchesExec: target_batch_size=8192",
- " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=1",
- // " CsvExec: file_groups={1 group: [[tempdir/left.csv]]}, projection=[a1, a2], has_header=false",
+ " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=8",
+ " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
+ // " CsvExec: file_groups={1 group: [[tempdir/left.csv]]}, projection=[a1, a2], has_header=false",
" CoalesceBatchesExec: target_batch_size=8192",
- " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=1",
- // " CsvExec: file_groups={1 group: [[tempdir/right.csv]]}, projection=[a1, a2], has_header=false"
+ " RepartitionExec: partitioning=Hash([a2@1], 8), input_partitions=8",
+ " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
+ // " CsvExec: file_groups={1 group: [[tempdir/right.csv]]}, projection=[a1, a2], has_header=false"
]
};
let mut actual: Vec<&str> = formatted.trim().lines().collect();
// Remove CSV lines
- actual.remove(3);
- actual.remove(5);
+ actual.remove(4);
+ actual.remove(7);
assert_eq!(
expected,
diff --git a/datafusion/sqllogictest/test_files/decimal.slt b/datafusion/sqllogictest/test_files/decimal.slt
index 76743e444e..3405401f7c 100644
--- a/datafusion/sqllogictest/test_files/decimal.slt
+++ b/datafusion/sqllogictest/test_files/decimal.slt
@@ -618,7 +618,12 @@ select a / b from foo;
statement ok
create table t as values (arrow_cast(123, 'Decimal256(5,2)'));
-query error DataFusion error: This feature is not implemented: AvgAccumulator for \(Decimal256\(5, 2\) \-\-> Decimal256\(9, 6\)\)
+# make sure query below runs in single partition
+# otherwise error message may not be deterministic
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+query error DataFusion error: This feature is not implemented: AvgAccumulator for \(Decimal256\(5, 2\) --> Decimal256\(9, 6\)\)
select AVG(column1) from t;
statement ok
diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt
index d44d19737a..b6b1ca9acd 100644
--- a/datafusion/sqllogictest/test_files/explain.slt
+++ b/datafusion/sqllogictest/test_files/explain.slt
@@ -244,7 +244,6 @@ logical_plan TableScan: simple_explain_test projection=[a, b, c]
initial_physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true
physical_plan after aggregate_statistics SAME TEXT AS ABOVE
physical_plan after join_selection SAME TEXT AS ABOVE
-physical_plan after repartition SAME TEXT AS ABOVE
physical_plan after EnforceDistribution SAME TEXT AS ABOVE
physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE
physical_plan after EnforceSorting SAME TEXT AS ABOVE
diff --git a/datafusion/sqllogictest/test_files/groupby.slt b/datafusion/sqllogictest/test_files/groupby.slt
index 8edf2b9ded..0c53a58ccd 100644
--- a/datafusion/sqllogictest/test_files/groupby.slt
+++ b/datafusion/sqllogictest/test_files/groupby.slt
@@ -1957,9 +1957,9 @@ Sort: l.col0 ASC NULLS LAST
----------TableScan: tab0 projection=[col0, col1]
physical_plan
SortPreservingMergeExec: [col0@0 ASC NULLS LAST]
---SortExec: expr=[col0@0 ASC NULLS LAST]
-----ProjectionExec: expr=[col0@0 as col0, LAST_VALUE(r.col1) ORDER BY [r.col0 ASC NULLS LAST]@3 as last_col1]
-------AggregateExec: mode=FinalPartitioned, gby=[col0@0 as col0, col1@1 as col1, col2@2 as col2], aggr=[LAST_VALUE(r.col1)]
+--ProjectionExec: expr=[col0@0 as col0, LAST_VALUE(r.col1) ORDER BY [r.col0 ASC NULLS LAST]@3 as last_col1]
+----AggregateExec: mode=FinalPartitioned, gby=[col0@0 as col0, col1@1 as col1, col2@2 as col2], aggr=[LAST_VALUE(r.col1)], ordering_mode=PartiallyOrdered
+------SortExec: expr=[col0@0 ASC NULLS LAST]
--------CoalesceBatchesExec: target_batch_size=8192
----------RepartitionExec: partitioning=Hash([col0@0, col1@1, col2@2], 4), input_partitions=4
------------AggregateExec: mode=Partial, gby=[col0@0 as col0, col1@1 as col1, col2@2 as col2], aggr=[LAST_VALUE(r.col1)], ordering_mode=PartiallyOrdered
@@ -2127,8 +2127,9 @@ Projection: annotated_data_infinite2.a, annotated_data_infinite2.b, FIRST_VALUE(
----TableScan: annotated_data_infinite2 projection=[a, b, c]
physical_plan
ProjectionExec: expr=[a@0 as a, b@1 as b, FIRST_VALUE(annotated_data_infinite2.c) ORDER BY [annotated_data_infinite2.a DESC NULLS FIRST]@2 as first_c]
---AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[LAST_VALUE(annotated_data_infinite2.c)], ordering_mode=FullyOrdered
-----CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], infinite_source=true, output_ordering=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST], has_header=true
+--AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[LAST_VALUE(annotated_data_infinite2.c)], ordering_mode=FullyOrdered
+----AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[LAST_VALUE(annotated_data_infinite2.c)], ordering_mode=FullyOrdered
+------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], infinite_source=true, output_ordering=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST], has_header=true
query III
SELECT a, b, FIRST_VALUE(c ORDER BY a DESC) as first_c
@@ -2153,8 +2154,9 @@ Projection: annotated_data_infinite2.a, annotated_data_infinite2.b, LAST_VALUE(a
----TableScan: annotated_data_infinite2 projection=[a, b, c]
physical_plan
ProjectionExec: expr=[a@0 as a, b@1 as b, LAST_VALUE(annotated_data_infinite2.c) ORDER BY [annotated_data_infinite2.a DESC NULLS FIRST]@2 as last_c]
---AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[FIRST_VALUE(annotated_data_infinite2.c)], ordering_mode=FullyOrdered
-----CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], infinite_source=true, output_ordering=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST], has_header=true
+--AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[FIRST_VALUE(annotated_data_infinite2.c)], ordering_mode=FullyOrdered
+----AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[FIRST_VALUE(annotated_data_infinite2.c)], ordering_mode=FullyOrdered
+------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], infinite_source=true, output_ordering=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST], has_header=true
query III
SELECT a, b, LAST_VALUE(c ORDER BY a DESC) as last_c
@@ -3284,12 +3286,14 @@ SortPreservingMergeExec: [sn@2 ASC NULLS LAST]
----ProjectionExec: expr=[zip_code@1 as zip_code, country@2 as country, sn@0 as sn, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount]
------AggregateExec: mode=FinalPartitioned, gby=[sn@0 as sn, zip_code@1 as zip_code, country@2 as country, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount], aggr=[]
--------CoalesceBatchesExec: target_batch_size=8192
-----------RepartitionExec: partitioning=Hash([sn@0, zip_code@1, country@2, ts@3, currency@4, amount@5, sum_amount@6], 8), input_partitions=1
+----------RepartitionExec: partitioning=Hash([sn@0, zip_code@1, country@2, ts@3, currency@4, amount@5, sum_amount@6], 8), input_partitions=8
------------AggregateExec: mode=Partial, gby=[sn@2 as sn, zip_code@0 as zip_code, country@1 as country, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount], aggr=[]
---------------ProjectionExec: expr=[zip_code@0 as zip_code, country@1 as country, sn@2 as sn, ts@3 as ts, currency@4 as currency, amount@5 as amount, SUM(l.amount) ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING@6 as sum_amount]
-----------------BoundedWindowAggExec: wdw=[SUM(l.amount) ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING: Ok(Field { name: "SUM(l.amount) ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(1)), end_bound: Following(UInt64(1)) }], mode=[Sorted]
-------------------CoalescePartitionsExec
---------------------MemoryExec: partitions=8, partition_sizes=[1, 0, 0, 0, 0, 0, 0, 0]
+--------------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
+----------------ProjectionExec: expr=[zip_code@0 as zip_code, country@1 as country, sn@2 as sn, ts@3 as ts, currency@4 as currency, amount@5 as amount, SUM(l.amount) ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING@6 as sum_amount]
+------------------BoundedWindowAggExec: wdw=[SUM(l.amount) ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING: Ok(Field { name: "SUM(l.amount) ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(1)), end_bound: Following(UInt64(1)) }], mode=[Sorted]
+--------------------CoalescePartitionsExec
+----------------------MemoryExec: partitions=8, partition_sizes=[1, 0, 0, 0, 0, 0, 0, 0]
+
query ITIPTRR
SELECT *
diff --git a/datafusion/sqllogictest/test_files/intersection.slt b/datafusion/sqllogictest/test_files/intersection.slt
index 31121a333d..301878cc98 100644
--- a/datafusion/sqllogictest/test_files/intersection.slt
+++ b/datafusion/sqllogictest/test_files/intersection.slt
@@ -30,7 +30,6 @@ SELECT * FROM (SELECT null AS id1, 1 AS id2) t1
----
NULL 1
-
query IR
SELECT int_col, double_col FROM alltypes_plain where int_col > 0 INTERSECT ALL SELECT int_col, double_col FROM alltypes_plain LIMIT 4
----
diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt
index cd93e51951..01c0131fdb 100644
--- a/datafusion/sqllogictest/test_files/joins.slt
+++ b/datafusion/sqllogictest/test_files/joins.slt
@@ -3125,14 +3125,18 @@ Sort: l_table.rn1 ASC NULLS LAST
physical_plan
SortPreservingMergeExec: [rn1@5 ASC NULLS LAST]
--SortMergeJoin: join_type=Inner, on=[(a@1, a@1)]
-----CoalesceBatchesExec: target_batch_size=4096
-------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=1
---------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
-----------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
+----SortExec: expr=[rn1@5 ASC NULLS LAST]
+------CoalesceBatchesExec: target_batch_size=4096
+--------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=2
+----------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+------------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
+--------------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
+----------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
+----SortExec: expr=[a@1 ASC]
+------CoalesceBatchesExec: target_batch_size=4096
+--------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=2
+----------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
-----CoalesceBatchesExec: target_batch_size=4096
-------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=1
---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
# sort merge join should propagate ordering equivalence of the right side
# for right join. Hence final requirement rn1 ASC is already satisfied at
@@ -3157,14 +3161,18 @@ Sort: r_table.rn1 ASC NULLS LAST
physical_plan
SortPreservingMergeExec: [rn1@10 ASC NULLS LAST]
--SortMergeJoin: join_type=Right, on=[(a@1, a@1)]
-----CoalesceBatchesExec: target_batch_size=4096
-------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=1
---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
-----CoalesceBatchesExec: target_batch_size=4096
-------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=1
---------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
-----------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
+----SortExec: expr=[a@1 ASC]
+------CoalesceBatchesExec: target_batch_size=4096
+--------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=2
+----------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
+----SortExec: expr=[rn1@5 ASC NULLS LAST]
+------CoalesceBatchesExec: target_batch_size=4096
+--------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=2
+----------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+------------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
+--------------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
+----------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
# SortMergeJoin should add ordering equivalences of
# right table as lexicographical append to the global ordering
@@ -3192,17 +3200,22 @@ Sort: l_table.a ASC NULLS FIRST, l_table.b ASC NULLS LAST, l_table.c ASC NULLS L
----------TableScan: annotated_data projection=[a0, a, b, c, d]
physical_plan
SortPreservingMergeExec: [a@1 ASC,b@2 ASC NULLS LAST,c@3 ASC NULLS LAST,rn1@11 ASC NULLS LAST]
---SortMergeJoin: join_type=Inner, on=[(a@1, a@1)]
-----CoalesceBatchesExec: target_batch_size=4096
-------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=1
---------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
-----------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
-------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
-----CoalesceBatchesExec: target_batch_size=4096
-------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=1
---------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
-----------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
-------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
+--SortExec: expr=[a@1 ASC,b@2 ASC NULLS LAST,c@3 ASC NULLS LAST,rn1@11 ASC NULLS LAST]
+----SortMergeJoin: join_type=Inner, on=[(a@1, a@1)]
+------SortExec: expr=[a@1 ASC]
+--------CoalesceBatchesExec: target_batch_size=4096
+----------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=2
+------------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+--------------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
+----------------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
+------------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
+------SortExec: expr=[a@1 ASC]
+--------CoalesceBatchesExec: target_batch_size=4096
+----------RepartitionExec: partitioning=Hash([a@1], 2), input_partitions=2
+------------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+--------------ProjectionExec: expr=[a0@0 as a0, a@1 as a, b@2 as b, c@3 as c, d@4 as d, ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@5 as rn1]
+----------------BoundedWindowAggExec: wdw=[ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "ROW_NUMBER() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)) }], mode=[Sorted]
+------------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
statement ok
set datafusion.optimizer.prefer_hash_join = true;
diff --git a/datafusion/sqllogictest/test_files/tpch/drop_tables.slt.part b/datafusion/sqllogictest/test_files/tpch/drop_tables.slt.part
index 6e6acee5b5..35faf3719d 100644
--- a/datafusion/sqllogictest/test_files/tpch/drop_tables.slt.part
+++ b/datafusion/sqllogictest/test_files/tpch/drop_tables.slt.part
@@ -1,5 +1,3 @@
-
-
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
diff --git a/datafusion/sqllogictest/test_files/tpch/q15.slt.part b/datafusion/sqllogictest/test_files/tpch/q15.slt.part
index 9ac8596f26..4515b8ae1f 100644
--- a/datafusion/sqllogictest/test_files/tpch/q15.slt.part
+++ b/datafusion/sqllogictest/test_files/tpch/q15.slt.part
@@ -95,20 +95,20 @@ SortPreservingMergeExec: [s_suppkey@0 ASC NULLS LAST]
----------------------------------FilterExec: l_shipdate@3 >= 9496 AND l_shipdate@3 < 9587
------------------------------------CsvExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_discou [...]
----------CoalesceBatchesExec: target_batch_size=8192
-------------RepartitionExec: partitioning=Hash([MAX(revenue0.total_revenue)@0], 4), input_partitions=1
---------------AggregateExec: mode=Final, gby=[], aggr=[MAX(revenue0.total_revenue)]
-----------------CoalescePartitionsExec
-------------------AggregateExec: mode=Partial, gby=[], aggr=[MAX(revenue0.total_revenue)]
---------------------ProjectionExec: expr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 as total_revenue]
-----------------------AggregateExec: mode=FinalPartitioned, gby=[l_suppkey@0 as l_suppkey], aggr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]
-------------------------CoalesceBatchesExec: target_batch_size=8192
---------------------------RepartitionExec: partitioning=Hash([l_suppkey@0], 4), input_partitions=4
-----------------------------AggregateExec: mode=Partial, gby=[l_suppkey@0 as l_suppkey], aggr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]
-------------------------------ProjectionExec: expr=[l_suppkey@0 as l_suppkey, l_extendedprice@1 as l_extendedprice, l_discount@2 as l_discount]
---------------------------------CoalesceBatchesExec: target_batch_size=8192
-----------------------------------FilterExec: l_shipdate@3 >= 9496 AND l_shipdate@3 < 9587
-------------------------------------CsvExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_discou [...]
-
+------------RepartitionExec: partitioning=Hash([MAX(revenue0.total_revenue)@0], 4), input_partitions=4
+--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+----------------AggregateExec: mode=Final, gby=[], aggr=[MAX(revenue0.total_revenue)]
+------------------CoalescePartitionsExec
+--------------------AggregateExec: mode=Partial, gby=[], aggr=[MAX(revenue0.total_revenue)]
+----------------------ProjectionExec: expr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 as total_revenue]
+------------------------AggregateExec: mode=FinalPartitioned, gby=[l_suppkey@0 as l_suppkey], aggr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]
+--------------------------CoalesceBatchesExec: target_batch_size=8192
+----------------------------RepartitionExec: partitioning=Hash([l_suppkey@0], 4), input_partitions=4
+------------------------------AggregateExec: mode=Partial, gby=[l_suppkey@0 as l_suppkey], aggr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]
+--------------------------------ProjectionExec: expr=[l_suppkey@0 as l_suppkey, l_extendedprice@1 as l_extendedprice, l_discount@2 as l_discount]
+----------------------------------CoalesceBatchesExec: target_batch_size=8192
+------------------------------------FilterExec: l_shipdate@3 >= 9496 AND l_shipdate@3 < 9587
+--------------------------------------CsvExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_disc [...]
query ITTTR
with revenue0 (supplier_no, total_revenue) as (
diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt
index c8aaa36fbc..05eaa10dab 100644
--- a/datafusion/sqllogictest/test_files/union.slt
+++ b/datafusion/sqllogictest/test_files/union.slt
@@ -466,8 +466,8 @@ UnionExec
----AggregateExec: mode=FinalPartitioned, gby=[Int64(1)@0 as Int64(1)], aggr=[]
------CoalesceBatchesExec: target_batch_size=8192
--------RepartitionExec: partitioning=Hash([Int64(1)@0], 4), input_partitions=4
-----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-------------AggregateExec: mode=Partial, gby=[1 as Int64(1)], aggr=[]
+----------AggregateExec: mode=Partial, gby=[1 as Int64(1)], aggr=[]
+------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
--------------EmptyExec: produce_one_row=true
--ProjectionExec: expr=[2 as a]
----EmptyExec: produce_one_row=true
@@ -498,11 +498,14 @@ UnionExec
----AggregateExec: mode=FinalPartitioned, gby=[n@0 as n], aggr=[COUNT(*)]
------CoalesceBatchesExec: target_batch_size=8192
--------RepartitionExec: partitioning=Hash([n@0], 4), input_partitions=4
-----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-------------AggregateExec: mode=Partial, gby=[n@0 as n], aggr=[COUNT(*)]
+----------AggregateExec: mode=Partial, gby=[n@0 as n], aggr=[COUNT(*)]
+------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
--------------ProjectionExec: expr=[5 as n]
----------------EmptyExec: produce_one_row=true
--ProjectionExec: expr=[x@0 as count, y@1 as n]
----ProjectionExec: expr=[1 as x, MAX(Int64(10))@0 as y]
-------AggregateExec: mode=Single, gby=[], aggr=[MAX(Int64(10))]
---------EmptyExec: produce_one_row=true
+------AggregateExec: mode=Final, gby=[], aggr=[MAX(Int64(10))]
+--------CoalescePartitionsExec
+----------AggregateExec: mode=Partial, gby=[], aggr=[MAX(Int64(10))]
+------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+--------------EmptyExec: produce_one_row=true
diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt
index 207bf21101..f8f8f30ade 100644
--- a/datafusion/sqllogictest/test_files/window.slt
+++ b/datafusion/sqllogictest/test_files/window.slt
@@ -417,6 +417,7 @@ sum(amount) - lag(sum(amount), 1) over (order by idx) as difference from (
select * from (values ('a', 1, 100), ('a', 2, 150)) as t (col1, idx, amount)
) a
group by col1, idx
+ORDER BY idx
----
a 1 1 100 NULL NULL
a 2 1 150 100 50