mustafasrepo commented on code in PR #7671: URL: https://github.com/apache/arrow-datafusion/pull/7671#discussion_r1341213456
########## 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 thought about this design. I think there are a couple of reason why current design is better - Currently `GlobalRequirementExec` is not necessarily the top executor in the physical plan. It is the executor above that defines output ordering. Such as plan below ```sql ProjectionExec OutputRequirementExec SortExec (specifies output ordering) ``` The reason for this strategy is that putting `OutputRequirementExec` at the top has couple of problems. Consider alternative strategy where plan above turns the following ```sql OutputRequirementExec ProjectionExec SortExec (specifies output ordering) ``` In this new version `OutputRequirementExec` may not have the necessary column to satisfy ordering desired at the output. Ordering should be satisfied before projection(These column may get lost during projection). Consider another plan ```sql WindowExec(OVER() sum()) Source (have an ordering) ``` If we were to require output ordering for it above plan would turn to ``` OutputRequirementExec WindowExec(OVER() sum()) Source (have an ordering) ``` However, output ordering in the plan above is accidental. It is not a requirement by the query. And planner should be free to mess this output ordering, if if is helpful. In short output ordering requirement only when query contains `ORDER BY` clause at the end. In this light, another strategy might be during plan creation we can insert `OutputRequirementExec` when we see `LogicalPlan::Sort`. However, in following kind of queries ```sql SELECT a, b FROM ( SELECT a, b FROM table1 ORDER BY a ) ORDER BY b ``` This would introduce unnecessary requirement for the subquery. Hence to require only appropriate absolutely necessary output ordering, we need to traverse physical plan from top to bottom, as long as ordering is maintained. Then put the `OutputRequirementExec` on top of it (this is the current approach). For this we need to have a top down pass on the physical plan. This can be done in `create_initial_plan` stage also. After initial physical plan is created. However, current implementation generally insert corresponding operator(s) for the `logicalPlan` node. Hence it is not easy to integrate top-down traversal code here. However, I think we can move OutputRequirementExec insertion code on top of here. https://github.com/apache/arrow-datafusion/blob/7b12666ec87ea741c3f5b56ddf1647f6d794f9e3/datafusion/core/src/physical_planner.rs#L450 After initial plan is created. What do you think, about this place? -- 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]
