alamb commented on code in PR #7671: URL: https://github.com/apache/arrow-datafusion/pull/7671#discussion_r1340487999
########## datafusion/core/src/physical_optimizer/global_requirements.rs: ########## @@ -0,0 +1,268 @@ +// 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 +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The GlobalOrderRequire optimizer rule either: +//! - Adds an auxiliary `GlobalRequirementExec` operator to keep track of global +//! ordering and distribution requirement across rules, or +//! - Removes the auxiliary `GlobalRequirementExec` operator from the physical plan. +//! Since the `GlobalRequirementExec` operator is only a helper operator, it +//! shouldn't occur in the final plan (i.e. the executed plan). + +use std::sync::Arc; + +use crate::physical_optimizer::PhysicalOptimizerRule; +use crate::physical_plan::sorts::sort::SortExec; +use crate::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; + +use arrow_schema::SchemaRef; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Result, Statistics}; +use datafusion_physical_expr::{ + Distribution, LexOrderingReq, PhysicalSortExpr, PhysicalSortRequirement, +}; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; + +/// This rule either adds or removes [`GlobalRequirements`]s to/from the physical +/// plan according to its `mode` attribute, which is set by the constructors +/// `new_add_mode` and `new_remove_mode`. With this rule, we can keep track of +/// the global requirements (ordering and distribution) across rules. +#[derive(Debug)] +pub struct GlobalRequirements { + mode: RuleMode, +} + +impl GlobalRequirements { + /// Create a new rule which works in `Add` mode; i.e. it simply adds a + /// top-level [`GlobalRequirementExec`] into the physical plan to keep track + /// of global ordering and distribution requirements if there are any. + /// Note that this rule should run at the beginning. + pub fn new_add_mode() -> Self { + Self { + mode: RuleMode::Add, + } + } + + /// Create a new rule which works in `Remove` mode; i.e. it simply removes + /// the top-level [`GlobalRequirementExec`] from the physical plan if there is + /// any. We do this because a `GlobalRequirementExec` is an ancillary, + /// non-executable operator whose sole purpose is to track global + /// requirements during optimization. Therefore, a + /// `GlobalRequirementExec` should not appear in the final plan. + pub fn new_remove_mode() -> Self { + Self { + mode: RuleMode::Remove, + } + } +} + +#[derive(Debug, Ord, PartialOrd, PartialEq, Eq, Hash)] +enum RuleMode { + Add, + Remove, +} + +/// An ancillary, non-executable operator whose sole purpose is to track global +/// requirements during optimization. It imposes +/// - the ordering requirement in its `order_requirement` attribute. +/// - the distribution requirement in its `dist_requirement` attribute. Review Comment: ```suggestion /// - the distribution requirement in its `dist_requirement` attribute. /// /// See [`GlobalRequirements`] for more details ``` ########## datafusion/core/src/physical_optimizer/global_requirements.rs: ########## @@ -0,0 +1,268 @@ +// 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 +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The GlobalOrderRequire optimizer rule either: +//! - Adds an auxiliary `GlobalRequirementExec` operator to keep track of global +//! ordering and distribution requirement across rules, or +//! - Removes the auxiliary `GlobalRequirementExec` operator from the physical plan. +//! Since the `GlobalRequirementExec` operator is only a helper operator, it +//! shouldn't occur in the final plan (i.e. the executed plan). + +use std::sync::Arc; + +use crate::physical_optimizer::PhysicalOptimizerRule; +use crate::physical_plan::sorts::sort::SortExec; +use crate::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; + +use arrow_schema::SchemaRef; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Result, Statistics}; +use datafusion_physical_expr::{ + Distribution, LexOrderingReq, PhysicalSortExpr, PhysicalSortRequirement, +}; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; + +/// This rule either adds or removes [`GlobalRequirements`]s to/from the physical +/// plan according to its `mode` attribute, which is set by the constructors +/// `new_add_mode` and `new_remove_mode`. With this rule, we can keep track of +/// the global requirements (ordering and distribution) across rules. Review Comment: ```suggestion /// This rule either adds or removes [`GlobalRequirements`]s to/from the physical /// plan according to its `mode` attribute, which is set by the constructors /// `new_add_mode` and `new_remove_mode`. With this rule, we can keep track of /// the global requirements (ordering and distribution) across rules. /// /// The primary usecase of this node and rule is to specify and preserve the desired output /// ordering and distribution the entire plan. When sending to a single client, a single partition may /// be desirable, but when sending to a multi-partitioned writer, keeping multiple partitions may be /// better. ``` ########## datafusion/core/src/physical_optimizer/enforce_distribution.rs: ########## @@ -2984,38 +2960,38 @@ mod tests { vec![ top_join_plan.as_str(), join_plan.as_str(), - "SortExec: expr=[a@0 ASC]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + "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]", - "SortExec: expr=[b1@1 ASC]", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + "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]", - "SortExec: expr=[c@2 ASC]", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + "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 7 RepartitionExecs Review Comment: this comment perhaps might be out of date ########## datafusion/core/src/physical_optimizer/global_requirements.rs: ########## @@ -0,0 +1,268 @@ +// 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 +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The GlobalOrderRequire optimizer rule either: +//! - Adds an auxiliary `GlobalRequirementExec` operator to keep track of global +//! ordering and distribution requirement across rules, or +//! - Removes the auxiliary `GlobalRequirementExec` operator from the physical plan. +//! Since the `GlobalRequirementExec` operator is only a helper operator, it +//! shouldn't occur in the final plan (i.e. the executed plan). + +use std::sync::Arc; + +use crate::physical_optimizer::PhysicalOptimizerRule; +use crate::physical_plan::sorts::sort::SortExec; +use crate::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; + +use arrow_schema::SchemaRef; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Result, Statistics}; +use datafusion_physical_expr::{ + Distribution, LexOrderingReq, PhysicalSortExpr, PhysicalSortRequirement, +}; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; + +/// This rule either adds or removes [`GlobalRequirements`]s to/from the physical +/// plan according to its `mode` attribute, which is set by the constructors +/// `new_add_mode` and `new_remove_mode`. With this rule, we can keep track of +/// the global requirements (ordering and distribution) across rules. +#[derive(Debug)] +pub struct GlobalRequirements { + mode: RuleMode, +} + +impl GlobalRequirements { + /// Create a new rule which works in `Add` mode; i.e. it simply adds a + /// top-level [`GlobalRequirementExec`] into the physical plan to keep track + /// of global ordering and distribution requirements if there are any. + /// Note that this rule should run at the beginning. + pub fn new_add_mode() -> Self { + Self { + mode: RuleMode::Add, + } + } + + /// Create a new rule which works in `Remove` mode; i.e. it simply removes + /// the top-level [`GlobalRequirementExec`] from the physical plan if there is + /// any. We do this because a `GlobalRequirementExec` is an ancillary, + /// non-executable operator whose sole purpose is to track global + /// requirements during optimization. Therefore, a + /// `GlobalRequirementExec` should not appear in the final plan. + pub fn new_remove_mode() -> Self { + Self { + mode: RuleMode::Remove, + } + } +} + +#[derive(Debug, Ord, PartialOrd, PartialEq, Eq, Hash)] +enum RuleMode { + Add, + Remove, +} + +/// An ancillary, non-executable operator whose sole purpose is to track global +/// requirements during optimization. It imposes +/// - the ordering requirement in its `order_requirement` attribute. +/// - the distribution requirement in its `dist_requirement` attribute. +#[derive(Debug)] +struct GlobalRequirementExec { + input: Arc<dyn ExecutionPlan>, + order_requirement: Option<LexOrderingReq>, + dist_requirement: Distribution, +} + +impl GlobalRequirementExec { + fn new( + input: Arc<dyn ExecutionPlan>, + requirements: Option<LexOrderingReq>, + dist_requirement: Distribution, + ) -> Self { + Self { + input, + order_requirement: requirements, + dist_requirement, + } + } + + fn input(&self) -> Arc<dyn ExecutionPlan> { + self.input.clone() + } +} + +impl DisplayAs for GlobalRequirementExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "GlobalRequirementExec") + } +} + +impl ExecutionPlan for GlobalRequirementExec { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn output_partitioning(&self) -> crate::physical_plan::Partitioning { + self.input.output_partitioning() + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + vec![self.dist_requirement.clone()] + } + + fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { + self.input.output_ordering() + } + + fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> { + vec![self.input.clone()] + } + + fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> { + vec![self.order_requirement.clone()] + } + + fn unbounded_output(&self, children: &[bool]) -> Result<bool> { + // Has a single child + Ok(children[0]) + } + + fn with_new_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(Self::new( + children.remove(0), // has a single child + self.order_requirement.clone(), + self.dist_requirement.clone(), + ))) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<crate::execution::context::TaskContext>, + ) -> Result<crate::physical_plan::SendableRecordBatchStream> { + unreachable!(); + } + + fn statistics(&self) -> Statistics { + self.input.statistics() + } +} + +impl PhysicalOptimizerRule for GlobalRequirements { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + _config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + match self.mode { + RuleMode::Add => require_top_ordering(plan), + RuleMode::Remove => plan.transform_up(&|plan| { + if let Some(sort_req) = + plan.as_any().downcast_ref::<GlobalRequirementExec>() + { + Ok(Transformed::Yes(sort_req.input())) + } else { + Ok(Transformed::No(plan)) + } + }), + } + } + + fn name(&self) -> &str { + "GlobalRequirements" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// This functions adds ancillary `GlobalRequirementExec` to the the physical plan, so that +/// global requirements are not lost during optimization. +fn require_top_ordering(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> { + let (new_plan, is_changed) = require_top_ordering_helper(plan)?; + if is_changed { + Ok(new_plan) + } else { + // Add `GlobalRequirementExec` to the top, with no specified ordering and distribution requirement. + Ok(Arc::new(GlobalRequirementExec::new( + new_plan, + // there is no ordering requirement + None, + Distribution::UnspecifiedDistribution, + )) as _) + } +} + +/// Helper function that adds an ancillary `GlobalRequirementExec` to the given plan. +/// First entry in the tuple is resulting plan, second entry indicates whether any +/// `GlobalRequirementExec` is added to the plan. +fn require_top_ordering_helper( Review Comment: I see this is basically trying to reverse engineer the top level requirements based on the pattern of the plan. Given `GlobalRequirementExec` can specify plan output requirements, I wonder if it would make more sense to directly create `GlobalRequirementExec` as pat of the initial physical plan, for example in https://github.com/apache/arrow-datafusion/blob/7b12666ec87ea741c3f5b56ddf1647f6d794f9e3/datafusion/core/src/physical_planner.rs#L536 That would also offer a nice way to control how inputs to write plans might look 🤔 ########## datafusion/core/src/physical_optimizer/global_requirements.rs: ########## @@ -0,0 +1,268 @@ +// 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 +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The GlobalOrderRequire optimizer rule either: +//! - Adds an auxiliary `GlobalRequirementExec` operator to keep track of global +//! ordering and distribution requirement across rules, or +//! - Removes the auxiliary `GlobalRequirementExec` operator from the physical plan. +//! Since the `GlobalRequirementExec` operator is only a helper operator, it +//! shouldn't occur in the final plan (i.e. the executed plan). + +use std::sync::Arc; + +use crate::physical_optimizer::PhysicalOptimizerRule; +use crate::physical_plan::sorts::sort::SortExec; +use crate::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; + +use arrow_schema::SchemaRef; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Result, Statistics}; +use datafusion_physical_expr::{ + Distribution, LexOrderingReq, PhysicalSortExpr, PhysicalSortRequirement, +}; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; + +/// This rule either adds or removes [`GlobalRequirements`]s to/from the physical +/// plan according to its `mode` attribute, which is set by the constructors +/// `new_add_mode` and `new_remove_mode`. With this rule, we can keep track of +/// the global requirements (ordering and distribution) across rules. +#[derive(Debug)] +pub struct GlobalRequirements { + mode: RuleMode, +} + +impl GlobalRequirements { + /// Create a new rule which works in `Add` mode; i.e. it simply adds a + /// top-level [`GlobalRequirementExec`] into the physical plan to keep track + /// of global ordering and distribution requirements if there are any. + /// Note that this rule should run at the beginning. + pub fn new_add_mode() -> Self { + Self { + mode: RuleMode::Add, + } + } + + /// Create a new rule which works in `Remove` mode; i.e. it simply removes + /// the top-level [`GlobalRequirementExec`] from the physical plan if there is + /// any. We do this because a `GlobalRequirementExec` is an ancillary, + /// non-executable operator whose sole purpose is to track global + /// requirements during optimization. Therefore, a + /// `GlobalRequirementExec` should not appear in the final plan. + pub fn new_remove_mode() -> Self { + Self { + mode: RuleMode::Remove, + } + } +} + +#[derive(Debug, Ord, PartialOrd, PartialEq, Eq, Hash)] +enum RuleMode { + Add, + Remove, +} + +/// An ancillary, non-executable operator whose sole purpose is to track global +/// requirements during optimization. It imposes +/// - the ordering requirement in its `order_requirement` attribute. +/// - the distribution requirement in its `dist_requirement` attribute. +#[derive(Debug)] +struct GlobalRequirementExec { + input: Arc<dyn ExecutionPlan>, + order_requirement: Option<LexOrderingReq>, + dist_requirement: Distribution, +} + +impl GlobalRequirementExec { + fn new( + input: Arc<dyn ExecutionPlan>, + requirements: Option<LexOrderingReq>, + dist_requirement: Distribution, + ) -> Self { + Self { + input, + order_requirement: requirements, + dist_requirement, + } + } + + fn input(&self) -> Arc<dyn ExecutionPlan> { + self.input.clone() + } +} + +impl DisplayAs for GlobalRequirementExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "GlobalRequirementExec") + } +} + +impl ExecutionPlan for GlobalRequirementExec { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn output_partitioning(&self) -> crate::physical_plan::Partitioning { + self.input.output_partitioning() + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + vec![self.dist_requirement.clone()] + } + + fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { + self.input.output_ordering() + } + + fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> { + vec![self.input.clone()] + } + + fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> { + vec![self.order_requirement.clone()] + } + + fn unbounded_output(&self, children: &[bool]) -> Result<bool> { + // Has a single child + Ok(children[0]) + } + + fn with_new_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(Self::new( + children.remove(0), // has a single child + self.order_requirement.clone(), + self.dist_requirement.clone(), + ))) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<crate::execution::context::TaskContext>, + ) -> Result<crate::physical_plan::SendableRecordBatchStream> { + unreachable!(); + } + + fn statistics(&self) -> Statistics { + self.input.statistics() + } +} + +impl PhysicalOptimizerRule for GlobalRequirements { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + _config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + match self.mode { + RuleMode::Add => require_top_ordering(plan), + RuleMode::Remove => plan.transform_up(&|plan| { + if let Some(sort_req) = + plan.as_any().downcast_ref::<GlobalRequirementExec>() + { + Ok(Transformed::Yes(sort_req.input())) + } else { + Ok(Transformed::No(plan)) + } + }), + } + } + + fn name(&self) -> &str { + "GlobalRequirements" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// This functions adds ancillary `GlobalRequirementExec` to the the physical plan, so that +/// global requirements are not lost during optimization. +fn require_top_ordering(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> { + let (new_plan, is_changed) = require_top_ordering_helper(plan)?; + if is_changed { + Ok(new_plan) + } else { + // Add `GlobalRequirementExec` to the top, with no specified ordering and distribution requirement. + Ok(Arc::new(GlobalRequirementExec::new( + new_plan, + // there is no ordering requirement + None, + Distribution::UnspecifiedDistribution, + )) as _) + } +} + +/// Helper function that adds an ancillary `GlobalRequirementExec` to the given plan. +/// First entry in the tuple is resulting plan, second entry indicates whether any +/// `GlobalRequirementExec` is added to the plan. +fn require_top_ordering_helper( Review Comment: I see this is basically trying to reverse engineer the top level requirements based on the pattern of the plan. Given `GlobalRequirementExec` can specify plan output requirements, I wonder if it would make more sense to directly create `GlobalRequirementExec` as pat of the initial physical plan, for example in https://github.com/apache/arrow-datafusion/blob/7b12666ec87ea741c3f5b56ddf1647f6d794f9e3/datafusion/core/src/physical_planner.rs#L536 That would also offer a nice way to control how inputs to write plans might look 🤔 ########## datafusion/core/src/physical_optimizer/enforce_distribution.rs: ########## @@ -3302,6 +3278,12 @@ mod tests { "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]", ]; assert_optimized!(expected, exec, true); + let expected = &[ Review Comment: This seems like a regression -- the comments in the code say "the optimizer should not add a new sort exec" but then in this code the optimizer has added the new sort exec. Maybe we need to update the test to add a `GlobalRequirements` node so that the sorts are not added? the same comments apply to the other tests in this file -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
