From e5066624df768ea01f1d6155bad09c5ee55fb21a Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 09:25:27 +0800 Subject: [PATCH 01/33] move out the ordering ruel Signed-off-by: jayzhan211 --- datafusion/core/src/physical_planner.rs | 286 +++++++++++++++++- .../physical-plan/src/aggregates/mod.rs | 48 ++- datafusion/physical-plan/src/windows/mod.rs | 2 +- .../sqllogictest/test_files/aggregate.slt | 13 +- 4 files changed, 329 insertions(+), 20 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index c25523c5ae33e..60b79b336f6e2 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -90,13 +90,18 @@ use datafusion_expr::{ DescribeTable, DmlStatement, RecursiveQuery, ScalarFunctionDefinition, StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, }; -use datafusion_physical_expr::expressions::Literal; +use datafusion_physical_expr::aggregate::is_order_sensitive; +use datafusion_physical_expr::expressions::{FirstValue, LastValue, Literal}; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_sql::utils::window_expr_common_partition_keys; use async_trait::async_trait; use datafusion_common::config::FormatOptions; -use datafusion_physical_expr::LexOrdering; +use datafusion_physical_expr::{ + physical_exprs_contains, reverse_order_bys, EquivalenceProperties, LexOrdering, + LexRequirement, PhysicalSortRequirement, +}; use futures::future::BoxFuture; use futures::{FutureExt, StreamExt, TryStreamExt}; use itertools::{multiunzip, Itertools}; @@ -450,6 +455,182 @@ pub struct DefaultPhysicalPlanner { extension_planners: Vec>, } +fn get_aggregate_expr_req( + aggr_expr: &Arc, + group_by: &PhysicalGroupBy, + agg_mode: &AggregateMode, +) -> LexOrdering { + // If the aggregation function is not order sensitive, or the aggregation + // is performing a "second stage" calculation, or all aggregate function + // requirements are inside the GROUP BY expression, then ignore the ordering + // requirement. + if !is_order_sensitive(aggr_expr) || !agg_mode.is_first_stage() { + return vec![]; + } + + let mut req = aggr_expr.order_bys().unwrap_or_default().to_vec(); + + // In non-first stage modes, we accumulate data (using `merge_batch`) from + // different partitions (i.e. merge partial results). During this merge, we + // consider the ordering of each partial result. Hence, we do not need to + // use the ordering requirement in such modes as long as partial results are + // generated with the correct ordering. + if group_by.is_single() { + // Remove all orderings that occur in the group by. These requirements + // will definitely be satisfied -- Each group by expression will have + // distinct values per group, hence all requirements are satisfied. + let physical_exprs = group_by.input_exprs(); + req.retain(|sort_expr| { + !physical_exprs_contains(&physical_exprs, &sort_expr.expr) + }); + } + req +} + +fn finer_ordering( + existing_req: &LexOrdering, + aggr_expr: &Arc, + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Option { + let aggr_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); + eq_properties.get_finer_ordering(existing_req, &aggr_req) +} + +/// Concatenates the given slices. +fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { + [lhs, rhs].concat() +} +fn get_aggregate_exprs_requirement( + prefix_requirement: &[PhysicalSortRequirement], + aggr_exprs: &mut [Arc], + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Result { + let mut requirement = vec![]; + for aggr_expr in aggr_exprs.iter_mut() { + let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); + let reverse_aggr_req = reverse_order_bys(aggr_req); + let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); + let reverse_aggr_req = + PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); + + use datafusion_functions_aggregate::first_last::FirstValuePhysicalExpr; + use datafusion_functions_aggregate::first_last::LastValuePhysicalExpr; + + if let Some(first_value) = + aggr_expr.as_any().downcast_ref::() + { + let mut first_value = first_value.clone(); + if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &aggr_req, + )) { + first_value = first_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(first_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to LAST_VALUE enables more efficient execution + // given the existing ordering: + let mut last_value = first_value.convert_to_last(); + last_value = last_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(last_value) as _; + } else { + // Requirement is not satisfied with existing ordering. + first_value = first_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(first_value) as _; + } + continue; + } + if let Some(last_value) = + aggr_expr.as_any().downcast_ref::() + { + let mut last_value = last_value.clone(); + if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &aggr_req, + )) { + last_value = last_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(last_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to FIRST_VALUE enables more efficient execution + // given the existing ordering: + let mut first_value = last_value.convert_to_first(); + first_value = first_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(first_value) as _; + } else { + // Requirement is not satisfied with existing ordering. + last_value = last_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(last_value) as _; + } + continue; + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Requirement is satisfied by existing ordering + requirement = finer_ordering; + continue; + } + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Reverse requirement is satisfied by exiting ordering. + // Hence reverse the aggregator + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + // There is a requirement that both satisfies existing requirement and current + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + continue; + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + // There is a requirement that both satisfies existing requirement and reverse + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + // Neither the existing requirement and current aggregate requirement satisfy the other, this means + // requirements are conflicting. Currently, we do not support + // conflicting requirements. + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ); + } + Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) +} + #[async_trait] impl PhysicalPlanner for DefaultPhysicalPlanner { /// Create a physical plan from a logical plan @@ -464,6 +645,107 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { let plan = self .create_initial_plan(logical_plan, session_state) .await?; + + println!("try my rule"); + // if let Some(AggregateExec { group_by, mut aggr_expr, input, input_schema, .. }) = plan.as_any().downcast_ref::() { + + let plan = if let Some(aggr_exec) = + plan.as_any().downcast_ref::() + { + let input = if let Some(aggr_exec) = + aggr_exec.input().as_any().downcast_ref::() + { + println!("AggregateExec"); + let input = aggr_exec.input().clone(); + let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); + let group_by = aggr_exec.group_by(); + let mode = aggr_exec.mode().clone(); + + use datafusion_physical_expr::equivalence::collapse_lex_req; + use datafusion_physical_plan::windows::get_ordered_partition_by_indices; + + let input_eq_properties = input.equivalence_properties(); + let groupby_exprs = group_by.input_exprs(); + let indices = + get_ordered_partition_by_indices(&groupby_exprs, &input); + let mut new_requirement = indices + .iter() + .map(|&idx| PhysicalSortRequirement { + expr: groupby_exprs[idx].clone(), + options: None, + }) + .collect::>(); + + print!("my new_requirement: {:?}", new_requirement); + println!("my aggr_expr: {:?}", aggr_expr); + println!("my group_by: {:?}", group_by); + println!("my input_eq_properties: {:?}", input_eq_properties); + println!("my mode: {:?}", mode); + let req = get_aggregate_exprs_requirement( + &new_requirement, + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + println!("my req: {:?}", req); + new_requirement.extend(req); + new_requirement = collapse_lex_req(new_requirement); + let required_input_ordering = + (!new_requirement.is_empty()).then_some(new_requirement); + + println!( + "my required_input_ordering: {:?}", + required_input_ordering + ); + let filter_expr = aggr_exec.filter_expr().to_vec(); + + // let metrics = aggr_exec.metrics(); + + let p = AggregateExec { + mode, + group_by: group_by.clone(), + aggr_expr, + filter_expr, + input, + schema: aggr_exec.schema().clone(), + input_schema: aggr_exec.input_schema().clone(), + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering, + limit: None, + input_order_mode: aggr_exec.input_order_mode().clone(), + cache: aggr_exec.cache().clone(), + }; + + Arc::new(p) as Arc + } else { + aggr_exec.input().clone() + }; + + // TODO: modify the input of aggr_exec + + // aggr_exec as Arc + + Arc::new(AggregateExec { + mode: aggr_exec.mode().clone(), + group_by: aggr_exec.group_by().clone(), + aggr_expr: aggr_exec.aggr_expr().to_vec(), + filter_expr: aggr_exec.filter_expr().to_vec(), + input, + schema: aggr_exec.schema().clone(), + input_schema: aggr_exec.input_schema().clone(), + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering: aggr_exec.required_input_ordering()[0].clone(), + limit: aggr_exec.limit().clone(), + input_order_mode: aggr_exec.input_order_mode().clone(), + cache: aggr_exec.cache().clone(), + }) as Arc + } else { + plan + }; + + println!("start optimize"); + self.optimize_internal(plan, session_state, |_, _| {}) } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 98c44e23c6c77..63ba87ee44346 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -241,19 +241,19 @@ impl From for SendableRecordBatchStream { #[derive(Debug)] pub struct AggregateExec { /// Aggregation mode (full, partial) - mode: AggregateMode, + pub mode: AggregateMode, /// Group by expressions - group_by: PhysicalGroupBy, + pub group_by: PhysicalGroupBy, /// Aggregate expressions - aggr_expr: Vec>, + pub aggr_expr: Vec>, /// FILTER (WHERE clause) expression for each aggregate expression - filter_expr: Vec>>, + pub filter_expr: Vec>>, /// Set if the output of this aggregation is truncated by a upstream sort/limit clause - limit: Option, + pub limit: Option, /// Input plan, could be a partial aggregate or the input to the aggregate pub input: Arc, /// Schema after the aggregate is applied - schema: SchemaRef, + pub schema: SchemaRef, /// Input schema before any aggregation is applied. For partial aggregate this will be the /// same as input.schema() but for the final aggregate it will be the same as the input /// to the partial aggregate, i.e., partial and final aggregates have same `input_schema`. @@ -261,14 +261,28 @@ pub struct AggregateExec { /// expressions from protobuf for final aggregate. pub input_schema: SchemaRef, /// Execution metrics - metrics: ExecutionPlanMetricsSet, - required_input_ordering: Option, + pub metrics: ExecutionPlanMetricsSet, + pub required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns - input_order_mode: InputOrderMode, - cache: PlanProperties, + pub input_order_mode: InputOrderMode, + pub cache: PlanProperties, } impl AggregateExec { + pub fn with_required_input_ordering( + self, + required_input_ordering: Option, + ) -> Self { + Self { + required_input_ordering, + ..self + } + } + + pub fn cache(&self) -> &PlanProperties { + &self.cache + } + /// Create a new hash aggregate execution plan pub fn try_new( mode: AggregateMode, @@ -336,6 +350,11 @@ impl AggregateExec { }) .collect::>(); + println!("correct new_requirement: {:?}", new_requirement); + println!("correct aggr_expr: {:?}", aggr_expr); + println!("correct group_by: {:?}", group_by); + println!("correct input_eq_properties: {:?}", input_eq_properties); + println!("correct mode: {:?}", mode); let req = get_aggregate_exprs_requirement( &new_requirement, &mut aggr_expr, @@ -362,6 +381,15 @@ impl AggregateExec { let required_input_ordering = (!new_requirement.is_empty()).then_some(new_requirement); + if required_input_ordering.is_some() { + println!( + "correct required_input_ordering: {:?}", + required_input_ordering + ); + } + + let required_input_ordering = None; + let cache = Self::compute_properties( &input, schema.clone(), diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index c5c845614c7b6..e01ee06a12b8f 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -373,7 +373,7 @@ pub(crate) fn calc_requirements< /// For instance, if input is ordered by a, b, c and PARTITION BY b, a is used, /// this vector will be [1, 0]. It means that when we iterate b, a columns with the order [1, 0] /// resulting vector (a, b) is a preset of the existing ordering (a, b, c). -pub(crate) fn get_ordered_partition_by_indices( +pub fn get_ordered_partition_by_indices( partition_by_exprs: &[Arc], input: &Arc, ) -> Vec { diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 4929ab485d6d7..5827ad77ad76f 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -123,7 +123,7 @@ LOCATION '../core/tests/data/aggregate_agg_multi_order.csv'; query ? select array_agg(c1 order by c2 desc, c3) from agg_order; ---- -[5, 6, 7, 8, 9, 1, 2, 3, 4, 10] +[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] query TT explain select array_agg(c1 order by c2 desc, c3) from agg_order; @@ -135,9 +135,8 @@ physical_plan AggregateExec: mode=Final, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] --CoalescePartitionsExec ----AggregateExec: mode=Partial, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] -------SortExec: expr=[c2@1 DESC,c3@2 ASC NULLS LAST] ---------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -----------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true +------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true # test array_agg_order with list data type statement ok @@ -152,8 +151,8 @@ CREATE TABLE array_agg_order_list_table AS VALUES query T? rowsort select column1, array_agg(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [[7, 8, 9], [4, 5, 6]] -w [[3, 2, 5], [9, 5, 2], [1, 2, 3]] +b [[4, 5, 6], [7, 8, 9]] +w [[1, 2, 3], [9, 5, 2], [3, 2, 5]] query T?? rowsort select column1, first_value(column3 order by column2, column4 desc), last_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; @@ -164,7 +163,7 @@ w [3, 2, 5] [1, 2, 3] query T? rowsort select column1, nth_value(column3, 2 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [4, 5, 6] +b [7, 8, 9] w [9, 5, 2] statement ok From 281288687d27d76d20bcbfab40bb37f8c800eb2e Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 10:20:22 +0800 Subject: [PATCH 02/33] introduce rule Signed-off-by: jayzhan211 --- datafusion/core/src/physical_optimizer/mod.rs | 1 + .../core/src/physical_optimizer/optimizer.rs | 2 + .../physical_optimizer/simplify_ordering.rs | 391 ++++++++++++++++++ datafusion/core/src/physical_planner.rs | 287 +------------ 4 files changed, 397 insertions(+), 284 deletions(-) create mode 100644 datafusion/core/src/physical_optimizer/simplify_ordering.rs diff --git a/datafusion/core/src/physical_optimizer/mod.rs b/datafusion/core/src/physical_optimizer/mod.rs index e990fead610d1..06f8974719634 100644 --- a/datafusion/core/src/physical_optimizer/mod.rs +++ b/datafusion/core/src/physical_optimizer/mod.rs @@ -34,6 +34,7 @@ pub mod pipeline_checker; mod projection_pushdown; pub mod pruning; pub mod replace_with_order_preserving_variants; +mod simplify_ordering; mod sort_pushdown; pub mod topk_aggregation; mod utils; diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 48da68cb2e376..931e20fff953e 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -20,6 +20,7 @@ use std::sync::Arc; use super::projection_pushdown::ProjectionPushdown; +use super::simplify_ordering::SimplifyOrdering; use crate::config::ConfigOptions; use crate::physical_optimizer::aggregate_statistics::AggregateStatistics; use crate::physical_optimizer::coalesce_batches::CoalesceBatches; @@ -126,6 +127,7 @@ impl PhysicalOptimizer { // are not present, the load of executors such as join or union will be // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), + Arc::new(SimplifyOrdering::new()), ]; Self::with_rules(rules) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs new file mode 100644 index 0000000000000..fd0dd746abd71 --- /dev/null +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -0,0 +1,391 @@ +// 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. + +use std::sync::Arc; + +use datafusion_common::{ + config::ConfigOptions, + not_impl_err, + tree_node::{Transformed, TransformedResult, TreeNode}, +}; +use datafusion_physical_expr::{ + aggregate::is_order_sensitive, + expressions::{FirstValue, LastValue}, + physical_exprs_contains, reverse_order_bys, AggregateExpr, EquivalenceProperties, + LexOrdering, LexRequirement, PhysicalSortRequirement, +}; +use datafusion_physical_plan::{ + aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, + metrics::ExecutionPlanMetricsSet, + ExecutionPlan, ExecutionPlanProperties, +}; + +use crate::error::Result; + +use super::PhysicalOptimizerRule; + +#[derive(Default)] +pub struct SimplifyOrdering {} + +impl SimplifyOrdering { + pub fn new() -> Self { + Self::default() + } +} + +impl PhysicalOptimizerRule for SimplifyOrdering { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + plan.transform_down(&get_common_requirement_of_aggregate_input) + .data() + } + + fn name(&self) -> &str { + "SimpleOrdering" + } + + fn schema_check(&self) -> bool { + true + } +} + +fn get_common_requirement_of_aggregate_input( + plan: Arc, +) -> Result>> { + let mut is_transformed = false; + + let plan = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + let input = if let Some(aggr_exec) = + aggr_exec.input().as_any().downcast_ref::() + { + // println!("AggregateExec"); + let input = aggr_exec.input().clone(); + let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); + let group_by = aggr_exec.group_by(); + let mode = aggr_exec.mode().clone(); + + use datafusion_physical_expr::equivalence::collapse_lex_req; + use datafusion_physical_plan::windows::get_ordered_partition_by_indices; + + let input_eq_properties = input.equivalence_properties(); + let groupby_exprs = group_by.input_exprs(); + let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); + let mut new_requirement = indices + .iter() + .map(|&idx| PhysicalSortRequirement { + expr: groupby_exprs[idx].clone(), + options: None, + }) + .collect::>(); + + // print!("my new_requirement: {:?}", new_requirement); + // println!("my aggr_expr: {:?}", aggr_expr); + // println!("my group_by: {:?}", group_by); + // println!("my input_eq_properties: {:?}", input_eq_properties); + // println!("my mode: {:?}", mode); + let req = get_aggregate_exprs_requirement( + &new_requirement, + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + // println!("my req: {:?}", req); + new_requirement.extend(req); + new_requirement = collapse_lex_req(new_requirement); + let required_input_ordering = + (!new_requirement.is_empty()).then_some(new_requirement); + + // println!( + // "my required_input_ordering: {:?}", + // required_input_ordering + // ); + let filter_expr = aggr_exec.filter_expr().to_vec(); + + // let metrics = aggr_exec.metrics(); + + let p = AggregateExec { + mode, + group_by: group_by.clone(), + aggr_expr, + filter_expr, + input, + schema: aggr_exec.schema().clone(), + input_schema: aggr_exec.input_schema().clone(), + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering, + limit: None, + input_order_mode: aggr_exec.input_order_mode().clone(), + cache: aggr_exec.cache().clone(), + }; + + is_transformed = true; + + Arc::new(p) as Arc + } else { + aggr_exec.input().clone() + }; + + // TODO: modify the input of aggr_exec + + // aggr_exec as Arc + + Arc::new(AggregateExec { + mode: aggr_exec.mode().clone(), + group_by: aggr_exec.group_by().clone(), + aggr_expr: aggr_exec.aggr_expr().to_vec(), + filter_expr: aggr_exec.filter_expr().to_vec(), + input, + schema: aggr_exec.schema().clone(), + input_schema: aggr_exec.input_schema().clone(), + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering: aggr_exec.required_input_ordering()[0].clone(), + limit: aggr_exec.limit().clone(), + input_order_mode: aggr_exec.input_order_mode().clone(), + cache: aggr_exec.cache().clone(), + }) as Arc + } else { + plan + }; + + if is_transformed { + Ok(Transformed::yes(plan)) + } else { + Ok(Transformed::no(plan)) + } +} + +/// Determines the lexical ordering requirement for an aggregate expression. +/// +/// # Parameters +/// +/// - `aggr_expr`: A reference to an `Arc` representing the +/// aggregate expression. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// +/// # Returns +/// +/// A `LexOrdering` instance indicating the lexical ordering requirement for +/// the aggregate expression. +fn get_aggregate_expr_req( + aggr_expr: &Arc, + group_by: &PhysicalGroupBy, + agg_mode: &AggregateMode, +) -> LexOrdering { + // If the aggregation function is not order sensitive, or the aggregation + // is performing a "second stage" calculation, or all aggregate function + // requirements are inside the GROUP BY expression, then ignore the ordering + // requirement. + if !is_order_sensitive(aggr_expr) || !agg_mode.is_first_stage() { + return vec![]; + } + + let mut req = aggr_expr.order_bys().unwrap_or_default().to_vec(); + + // In non-first stage modes, we accumulate data (using `merge_batch`) from + // different partitions (i.e. merge partial results). During this merge, we + // consider the ordering of each partial result. Hence, we do not need to + // use the ordering requirement in such modes as long as partial results are + // generated with the correct ordering. + if group_by.is_single() { + // Remove all orderings that occur in the group by. These requirements + // will definitely be satisfied -- Each group by expression will have + // distinct values per group, hence all requirements are satisfied. + let physical_exprs = group_by.input_exprs(); + req.retain(|sort_expr| { + !physical_exprs_contains(&physical_exprs, &sort_expr.expr) + }); + } + req +} + +/// Computes the finer ordering for between given existing ordering requirement +/// of aggregate expression. +/// +/// # Parameters +/// +/// * `existing_req` - The existing lexical ordering that needs refinement. +/// * `aggr_expr` - A reference to an aggregate expression trait object. +/// * `group_by` - Information about the physical grouping (e.g group by expression). +/// * `eq_properties` - Equivalence properties relevant to the computation. +/// * `agg_mode` - The mode of aggregation (e.g., Partial, Final, etc.). +/// +/// # Returns +/// +/// An `Option` representing the computed finer lexical ordering, +/// or `None` if there is no finer ordering; e.g. the existing requirement and +/// the aggregator requirement is incompatible. +fn finer_ordering( + existing_req: &LexOrdering, + aggr_expr: &Arc, + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Option { + let aggr_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); + eq_properties.get_finer_ordering(existing_req, &aggr_req) +} + +/// Concatenates the given slices. +fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { + [lhs, rhs].concat() +} + +/// Get the common requirement that satisfies all the aggregate expressions. +/// +/// # Parameters +/// +/// - `aggr_exprs`: A slice of `Arc` containing all the +/// aggregate expressions. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `eq_properties`: A reference to an `EquivalenceProperties` instance +/// representing equivalence properties for ordering. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// +/// # Returns +/// +/// A `LexRequirement` instance, which is the requirement that satisfies all the +/// aggregate requirements. Returns an error in case of conflicting requirements. +fn get_aggregate_exprs_requirement( + prefix_requirement: &[PhysicalSortRequirement], + aggr_exprs: &mut [Arc], + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Result { + let mut requirement = vec![]; + for aggr_expr in aggr_exprs.iter_mut() { + let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); + let reverse_aggr_req = reverse_order_bys(aggr_req); + let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); + let reverse_aggr_req = + PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); + + if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { + let mut first_value = first_value.clone(); + if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &aggr_req, + )) { + first_value = first_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(first_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to LAST_VALUE enables more efficient execution + // given the existing ordering: + let mut last_value = first_value.convert_to_last(); + last_value = last_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(last_value) as _; + } else { + // Requirement is not satisfied with existing ordering. + first_value = first_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(first_value) as _; + } + continue; + } + if let Some(last_value) = aggr_expr.as_any().downcast_ref::() { + let mut last_value = last_value.clone(); + if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &aggr_req, + )) { + last_value = last_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(last_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to FIRST_VALUE enables more efficient execution + // given the existing ordering: + let mut first_value = last_value.convert_to_first(); + first_value = first_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(first_value) as _; + } else { + // Requirement is not satisfied with existing ordering. + last_value = last_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(last_value) as _; + } + continue; + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Requirement is satisfied by existing ordering + requirement = finer_ordering; + continue; + } + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Reverse requirement is satisfied by exiting ordering. + // Hence reverse the aggregator + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + // There is a requirement that both satisfies existing requirement and current + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + continue; + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + // There is a requirement that both satisfies existing requirement and reverse + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + // Neither the existing requirement and current aggregate requirement satisfy the other, this means + // requirements are conflicting. Currently, we do not support + // conflicting requirements. + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ); + } + Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) +} diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 60b79b336f6e2..798dfa2d5bfbc 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -74,6 +74,7 @@ use arrow::compute::SortOptions; use arrow::datatypes::{Schema, SchemaRef}; use arrow_array::builder::StringBuilder; use arrow_array::RecordBatch; +use async_trait::async_trait; use datafusion_common::display::ToStringifiedPlan; use datafusion_common::{ exec_err, internal_err, not_impl_err, plan_err, DFSchema, FileType, ScalarValue, @@ -90,18 +91,12 @@ use datafusion_expr::{ DescribeTable, DmlStatement, RecursiveQuery, ScalarFunctionDefinition, StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, }; -use datafusion_physical_expr::aggregate::is_order_sensitive; -use datafusion_physical_expr::expressions::{FirstValue, LastValue, Literal}; -use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_expr::expressions::Literal; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_sql::utils::window_expr_common_partition_keys; -use async_trait::async_trait; use datafusion_common::config::FormatOptions; -use datafusion_physical_expr::{ - physical_exprs_contains, reverse_order_bys, EquivalenceProperties, LexOrdering, - LexRequirement, PhysicalSortRequirement, -}; +use datafusion_physical_expr::LexOrdering; use futures::future::BoxFuture; use futures::{FutureExt, StreamExt, TryStreamExt}; use itertools::{multiunzip, Itertools}; @@ -455,182 +450,6 @@ pub struct DefaultPhysicalPlanner { extension_planners: Vec>, } -fn get_aggregate_expr_req( - aggr_expr: &Arc, - group_by: &PhysicalGroupBy, - agg_mode: &AggregateMode, -) -> LexOrdering { - // If the aggregation function is not order sensitive, or the aggregation - // is performing a "second stage" calculation, or all aggregate function - // requirements are inside the GROUP BY expression, then ignore the ordering - // requirement. - if !is_order_sensitive(aggr_expr) || !agg_mode.is_first_stage() { - return vec![]; - } - - let mut req = aggr_expr.order_bys().unwrap_or_default().to_vec(); - - // In non-first stage modes, we accumulate data (using `merge_batch`) from - // different partitions (i.e. merge partial results). During this merge, we - // consider the ordering of each partial result. Hence, we do not need to - // use the ordering requirement in such modes as long as partial results are - // generated with the correct ordering. - if group_by.is_single() { - // Remove all orderings that occur in the group by. These requirements - // will definitely be satisfied -- Each group by expression will have - // distinct values per group, hence all requirements are satisfied. - let physical_exprs = group_by.input_exprs(); - req.retain(|sort_expr| { - !physical_exprs_contains(&physical_exprs, &sort_expr.expr) - }); - } - req -} - -fn finer_ordering( - existing_req: &LexOrdering, - aggr_expr: &Arc, - group_by: &PhysicalGroupBy, - eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Option { - let aggr_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); - eq_properties.get_finer_ordering(existing_req, &aggr_req) -} - -/// Concatenates the given slices. -fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { - [lhs, rhs].concat() -} -fn get_aggregate_exprs_requirement( - prefix_requirement: &[PhysicalSortRequirement], - aggr_exprs: &mut [Arc], - group_by: &PhysicalGroupBy, - eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Result { - let mut requirement = vec![]; - for aggr_expr in aggr_exprs.iter_mut() { - let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); - let reverse_aggr_req = reverse_order_bys(aggr_req); - let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); - let reverse_aggr_req = - PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); - - use datafusion_functions_aggregate::first_last::FirstValuePhysicalExpr; - use datafusion_functions_aggregate::first_last::LastValuePhysicalExpr; - - if let Some(first_value) = - aggr_expr.as_any().downcast_ref::() - { - let mut first_value = first_value.clone(); - if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &aggr_req, - )) { - first_value = first_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(first_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to LAST_VALUE enables more efficient execution - // given the existing ordering: - let mut last_value = first_value.convert_to_last(); - last_value = last_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(last_value) as _; - } else { - // Requirement is not satisfied with existing ordering. - first_value = first_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(first_value) as _; - } - continue; - } - if let Some(last_value) = - aggr_expr.as_any().downcast_ref::() - { - let mut last_value = last_value.clone(); - if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &aggr_req, - )) { - last_value = last_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(last_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to FIRST_VALUE enables more efficient execution - // given the existing ordering: - let mut first_value = last_value.convert_to_first(); - first_value = first_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(first_value) as _; - } else { - // Requirement is not satisfied with existing ordering. - last_value = last_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(last_value) as _; - } - continue; - } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Requirement is satisfied by existing ordering - requirement = finer_ordering; - continue; - } - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Reverse requirement is satisfied by exiting ordering. - // Hence reverse the aggregator - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } - } - } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - // There is a requirement that both satisfies existing requirement and current - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - continue; - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - // There is a requirement that both satisfies existing requirement and reverse - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } - } - // Neither the existing requirement and current aggregate requirement satisfy the other, this means - // requirements are conflicting. Currently, we do not support - // conflicting requirements. - return not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" - ); - } - Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) -} - #[async_trait] impl PhysicalPlanner for DefaultPhysicalPlanner { /// Create a physical plan from a logical plan @@ -646,106 +465,6 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { .create_initial_plan(logical_plan, session_state) .await?; - println!("try my rule"); - // if let Some(AggregateExec { group_by, mut aggr_expr, input, input_schema, .. }) = plan.as_any().downcast_ref::() { - - let plan = if let Some(aggr_exec) = - plan.as_any().downcast_ref::() - { - let input = if let Some(aggr_exec) = - aggr_exec.input().as_any().downcast_ref::() - { - println!("AggregateExec"); - let input = aggr_exec.input().clone(); - let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); - let group_by = aggr_exec.group_by(); - let mode = aggr_exec.mode().clone(); - - use datafusion_physical_expr::equivalence::collapse_lex_req; - use datafusion_physical_plan::windows::get_ordered_partition_by_indices; - - let input_eq_properties = input.equivalence_properties(); - let groupby_exprs = group_by.input_exprs(); - let indices = - get_ordered_partition_by_indices(&groupby_exprs, &input); - let mut new_requirement = indices - .iter() - .map(|&idx| PhysicalSortRequirement { - expr: groupby_exprs[idx].clone(), - options: None, - }) - .collect::>(); - - print!("my new_requirement: {:?}", new_requirement); - println!("my aggr_expr: {:?}", aggr_expr); - println!("my group_by: {:?}", group_by); - println!("my input_eq_properties: {:?}", input_eq_properties); - println!("my mode: {:?}", mode); - let req = get_aggregate_exprs_requirement( - &new_requirement, - &mut aggr_expr, - &group_by, - input_eq_properties, - &mode, - )?; - println!("my req: {:?}", req); - new_requirement.extend(req); - new_requirement = collapse_lex_req(new_requirement); - let required_input_ordering = - (!new_requirement.is_empty()).then_some(new_requirement); - - println!( - "my required_input_ordering: {:?}", - required_input_ordering - ); - let filter_expr = aggr_exec.filter_expr().to_vec(); - - // let metrics = aggr_exec.metrics(); - - let p = AggregateExec { - mode, - group_by: group_by.clone(), - aggr_expr, - filter_expr, - input, - schema: aggr_exec.schema().clone(), - input_schema: aggr_exec.input_schema().clone(), - metrics: ExecutionPlanMetricsSet::new(), - required_input_ordering, - limit: None, - input_order_mode: aggr_exec.input_order_mode().clone(), - cache: aggr_exec.cache().clone(), - }; - - Arc::new(p) as Arc - } else { - aggr_exec.input().clone() - }; - - // TODO: modify the input of aggr_exec - - // aggr_exec as Arc - - Arc::new(AggregateExec { - mode: aggr_exec.mode().clone(), - group_by: aggr_exec.group_by().clone(), - aggr_expr: aggr_exec.aggr_expr().to_vec(), - filter_expr: aggr_exec.filter_expr().to_vec(), - input, - schema: aggr_exec.schema().clone(), - input_schema: aggr_exec.input_schema().clone(), - metrics: ExecutionPlanMetricsSet::new(), - required_input_ordering: aggr_exec.required_input_ordering()[0].clone(), - limit: aggr_exec.limit().clone(), - input_order_mode: aggr_exec.input_order_mode().clone(), - cache: aggr_exec.cache().clone(), - }) as Arc - } else { - plan - }; - - println!("start optimize"); - self.optimize_internal(plan, session_state, |_, _| {}) } } From d20dde38f85fcfdee198f530b3d01e41541892c2 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 10:56:46 +0800 Subject: [PATCH 03/33] revert test result Signed-off-by: jayzhan211 --- datafusion/sqllogictest/test_files/aggregate.slt | 13 +++++++------ datafusion/sqllogictest/test_files/explain.slt | 3 +++ datafusion/sqllogictest/test_files/join.slt | 2 +- datafusion/sqllogictest/test_files/predicates.slt | 2 +- datafusion/sqllogictest/test_files/timestamps.slt | 2 +- 5 files changed, 13 insertions(+), 9 deletions(-) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 5827ad77ad76f..4929ab485d6d7 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -123,7 +123,7 @@ LOCATION '../core/tests/data/aggregate_agg_multi_order.csv'; query ? select array_agg(c1 order by c2 desc, c3) from agg_order; ---- -[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] +[5, 6, 7, 8, 9, 1, 2, 3, 4, 10] query TT explain select array_agg(c1 order by c2 desc, c3) from agg_order; @@ -135,8 +135,9 @@ physical_plan AggregateExec: mode=Final, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] --CoalescePartitionsExec ----AggregateExec: mode=Partial, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] -------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 ---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true +------SortExec: expr=[c2@1 DESC,c3@2 ASC NULLS LAST] +--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +----------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true # test array_agg_order with list data type statement ok @@ -151,8 +152,8 @@ CREATE TABLE array_agg_order_list_table AS VALUES query T? rowsort select column1, array_agg(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [[4, 5, 6], [7, 8, 9]] -w [[1, 2, 3], [9, 5, 2], [3, 2, 5]] +b [[7, 8, 9], [4, 5, 6]] +w [[3, 2, 5], [9, 5, 2], [1, 2, 3]] query T?? rowsort select column1, first_value(column3 order by column2, column4 desc), last_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; @@ -163,7 +164,7 @@ w [3, 2, 5] [1, 2, 3] query T? rowsort select column1, nth_value(column3, 2 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [7, 8, 9] +b [4, 5, 6] w [9, 5, 2] statement ok diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index b7ad36dace162..57ee8c311f6c6 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -254,6 +254,7 @@ physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPAC physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] @@ -312,6 +313,7 @@ GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Co physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -348,6 +350,7 @@ GlobalLimitExec: skip=0, fetch=10 physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10 --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 diff --git a/datafusion/sqllogictest/test_files/join.slt b/datafusion/sqllogictest/test_files/join.slt index da9b4168e7e09..135ab80754253 100644 --- a/datafusion/sqllogictest/test_files/join.slt +++ b/datafusion/sqllogictest/test_files/join.slt @@ -587,7 +587,7 @@ FROM t1 ---- 11 11 11 -# subsequent inner join +# subsequent inner join query III rowsort SELECT t1.t1_id, t2.t2_id, t3.t3_id FROM t1 diff --git a/datafusion/sqllogictest/test_files/predicates.slt b/datafusion/sqllogictest/test_files/predicates.slt index 33c9ff7c3eedf..4c9254beef6b2 100644 --- a/datafusion/sqllogictest/test_files/predicates.slt +++ b/datafusion/sqllogictest/test_files/predicates.slt @@ -781,4 +781,4 @@ logical_plan EmptyRelation physical_plan EmptyExec statement ok -drop table t; \ No newline at end of file +drop table t; diff --git a/datafusion/sqllogictest/test_files/timestamps.slt b/datafusion/sqllogictest/test_files/timestamps.slt index f0e04b522a785..491b9b810687a 100644 --- a/datafusion/sqllogictest/test_files/timestamps.slt +++ b/datafusion/sqllogictest/test_files/timestamps.slt @@ -2794,4 +2794,4 @@ SELECT '2000-12-01 04:04:12' AT TIME ZONE 'America/New York'; # abbreviated timezone is not supported statement error -SELECT '2023-03-12 02:00:00' AT TIME ZONE 'EDT'; \ No newline at end of file +SELECT '2023-03-12 02:00:00' AT TIME ZONE 'EDT'; From c671033424490ce69b07d1027d97d2d3d37e11fa Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 14:53:01 +0800 Subject: [PATCH 04/33] pass mulit order test Signed-off-by: jayzhan211 --- .../core/src/physical_optimizer/optimizer.rs | 7 +- .../physical_optimizer/output_requirements.rs | 9 + .../physical_optimizer/simplify_ordering.rs | 302 ++++++++----- datafusion/core/src/physical_planner.rs | 4 +- .../physical-plan/src/aggregates/mod.rs | 400 ++---------------- .../physical-plan/src/coalesce_partitions.rs | 8 + .../sqllogictest/test_files/aggregate.slt | 26 +- 7 files changed, 290 insertions(+), 466 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 931e20fff953e..b777933e120e0 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -76,10 +76,12 @@ impl PhysicalOptimizer { /// Create a new optimizer using the recommended list of rules pub fn new() -> Self { let rules: Vec> = vec![ + Arc::new(SimplifyOrdering::new()), // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), Arc::new(AggregateStatistics::new()), + Arc::new(SimplifyOrdering::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the // EnforceDistribution and EnforceSorting rules as they decide whether to add additional @@ -90,13 +92,16 @@ impl PhysicalOptimizer { // as that rule may inject other operations in between the different AggregateExecs. // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), + Arc::new(SimplifyOrdering::new()), // 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()), + Arc::new(SimplifyOrdering::new()), // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule Arc::new(CombinePartialFinalAggregate::new()), + Arc::new(SimplifyOrdering::new()), // The EnforceSorting rule is for adding essential local sorting to satisfy the required // ordering. Please make sure that the whole plan tree is determined before this rule. // Note that one should always run this rule after running the EnforceDistribution rule @@ -120,6 +125,7 @@ impl PhysicalOptimizer { // into an `order by max(x) limit y`. In this case it will copy the limit value down // to the aggregation, allowing it to use only y number of accumulators. Arc::new(TopKAggregation::new()), + Arc::new(SimplifyOrdering::new()), // The ProjectionPushdown rule tries to push projections towards // the sources in the execution plan. As a result of this process, // a projection can disappear if it reaches the source providers, and @@ -127,7 +133,6 @@ impl PhysicalOptimizer { // are not present, the load of executors such as join or union will be // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), - Arc::new(SimplifyOrdering::new()), ]; Self::with_rules(rules) diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 829d523c990ca..386340bbeda69 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -94,6 +94,15 @@ pub(crate) struct OutputRequirementExec { } impl OutputRequirementExec { + pub fn clone_with_input(&self, input: Arc) -> Self { + Self { + input, + order_requirement: self.order_requirement.clone(), + dist_requirement: self.dist_requirement.clone(), + cache: self.cache.clone(), + } + } + pub(crate) fn new( input: Arc, requirements: Option, diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index fd0dd746abd71..bb02df69b92be 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -17,10 +17,11 @@ use std::sync::Arc; +use arrow::compute::kernels::aggregate; use datafusion_common::{ config::ConfigOptions, not_impl_err, - tree_node::{Transformed, TransformedResult, TreeNode}, + tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, }; use datafusion_physical_expr::{ aggregate::is_order_sensitive, @@ -30,13 +31,16 @@ use datafusion_physical_expr::{ }; use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, - metrics::ExecutionPlanMetricsSet, + coalesce_partitions::CoalescePartitionsExec, ExecutionPlan, ExecutionPlanProperties, }; +use datafusion_physical_expr::equivalence::collapse_lex_req; +use datafusion_physical_plan::windows::get_ordered_partition_by_indices; + use crate::error::Result; -use super::PhysicalOptimizerRule; +use super::{output_requirements::OutputRequirementExec, PhysicalOptimizerRule}; #[derive(Default)] pub struct SimplifyOrdering {} @@ -53,8 +57,12 @@ impl PhysicalOptimizerRule for SimplifyOrdering { plan: Arc, _config: &ConfigOptions, ) -> Result> { - plan.transform_down(&get_common_requirement_of_aggregate_input) - .data() + let res = plan + .transform_down(&get_common_requirement_of_aggregate_input) + .data(); + + // println!("res: {:?}", res); + res } fn name(&self) -> &str { @@ -69,104 +77,100 @@ impl PhysicalOptimizerRule for SimplifyOrdering { fn get_common_requirement_of_aggregate_input( plan: Arc, ) -> Result>> { - let mut is_transformed = false; + let children = plan.children(); + + let new_c: Option>> = if children.is_empty() { + None + } else { + assert_eq!(children.len(), 1, "AggregateExec should have one child"); + let c = children[0].clone(); + + // for c in children { + // let new_c = get_common_requirement_of_aggregate_input(c.clone())?; + // if new_c.transformed { + // is_transformed = true; + // } + // } - let plan = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let input = if let Some(aggr_exec) = - aggr_exec.input().as_any().downcast_ref::() + let new_c = get_common_requirement_of_aggregate_input(c)?; + Some(new_c) + }; + + let plan = optimize_internal(plan)?; + // println!("t: {} plan: {:?}", plan.transformed, plan); + + if let Some(c) = new_c { + if !c.transformed { + return Ok(plan); + } + + let plan = plan.data; + + // TODO: support more types of ExecutionPlan + if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + let p = aggr_exec.clone_with_input(c.data); + return Ok(Transformed::yes(Arc::new(p) as Arc)); + } else if let Some(coalesce_exec) = + plan.as_any().downcast_ref::() { - // println!("AggregateExec"); - let input = aggr_exec.input().clone(); - let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); - let group_by = aggr_exec.group_by(); - let mode = aggr_exec.mode().clone(); - - use datafusion_physical_expr::equivalence::collapse_lex_req; - use datafusion_physical_plan::windows::get_ordered_partition_by_indices; - - let input_eq_properties = input.equivalence_properties(); - let groupby_exprs = group_by.input_exprs(); - let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); - let mut new_requirement = indices - .iter() - .map(|&idx| PhysicalSortRequirement { - expr: groupby_exprs[idx].clone(), - options: None, - }) - .collect::>(); - - // print!("my new_requirement: {:?}", new_requirement); - // println!("my aggr_expr: {:?}", aggr_expr); - // println!("my group_by: {:?}", group_by); - // println!("my input_eq_properties: {:?}", input_eq_properties); - // println!("my mode: {:?}", mode); - let req = get_aggregate_exprs_requirement( - &new_requirement, - &mut aggr_expr, - &group_by, - input_eq_properties, - &mode, - )?; - // println!("my req: {:?}", req); - new_requirement.extend(req); - new_requirement = collapse_lex_req(new_requirement); - let required_input_ordering = - (!new_requirement.is_empty()).then_some(new_requirement); - - // println!( - // "my required_input_ordering: {:?}", - // required_input_ordering - // ); - let filter_expr = aggr_exec.filter_expr().to_vec(); - - // let metrics = aggr_exec.metrics(); - - let p = AggregateExec { - mode, - group_by: group_by.clone(), - aggr_expr, - filter_expr, - input, - schema: aggr_exec.schema().clone(), - input_schema: aggr_exec.input_schema().clone(), - metrics: ExecutionPlanMetricsSet::new(), - required_input_ordering, - limit: None, - input_order_mode: aggr_exec.input_order_mode().clone(), - cache: aggr_exec.cache().clone(), - }; - - is_transformed = true; - - Arc::new(p) as Arc + let p = coalesce_exec.clone_with_input(c.data); + return Ok(Transformed::yes(Arc::new(p) as Arc)); + } else if let Some(out_req_exec) = + plan.as_any().downcast_ref::() + { + let p = out_req_exec.clone_with_input(c.data); + return Ok(Transformed::yes(Arc::new(p) as Arc)); } else { - aggr_exec.input().clone() - }; + return not_impl_err!("Unsupported ExecutionPlan type: {}", plan.name()); + } + } - // TODO: modify the input of aggr_exec - - // aggr_exec as Arc - - Arc::new(AggregateExec { - mode: aggr_exec.mode().clone(), - group_by: aggr_exec.group_by().clone(), - aggr_expr: aggr_exec.aggr_expr().to_vec(), - filter_expr: aggr_exec.filter_expr().to_vec(), - input, - schema: aggr_exec.schema().clone(), - input_schema: aggr_exec.input_schema().clone(), - metrics: ExecutionPlanMetricsSet::new(), - required_input_ordering: aggr_exec.required_input_ordering()[0].clone(), - limit: aggr_exec.limit().clone(), - input_order_mode: aggr_exec.input_order_mode().clone(), - cache: aggr_exec.cache().clone(), - }) as Arc - } else { - plan - }; + return Ok(plan); +} + +fn optimize_internal( + plan: Arc, +) -> Result>> { + if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + if aggr_exec.mode() != &AggregateMode::Partial { + return Ok(Transformed::no(plan)); + } + + let input = aggr_exec.input().clone(); + let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); + let group_by = aggr_exec.group_by(); + let mode = aggr_exec.mode(); + + let input_eq_properties = input.equivalence_properties(); + let groupby_exprs = group_by.input_exprs(); + // If existing ordering satisfies a prefix of the GROUP BY expressions, + // prefix requirements with this section. In this case, aggregation will + // work more efficiently. + let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); + let mut new_requirement = indices + .iter() + .map(|&idx| PhysicalSortRequirement { + expr: groupby_exprs[idx].clone(), + options: None, + }) + .collect::>(); + + let req = get_aggregate_exprs_requirement( + &new_requirement, + &mut aggr_expr, + &group_by, + input_eq_properties, + mode, + )?; + new_requirement.extend(req); + new_requirement = collapse_lex_req(new_requirement); + let required_input_ordering = + (!new_requirement.is_empty()).then_some(new_requirement); + + let p = aggr_exec.clone_with_required_input_ordering(required_input_ordering); - if is_transformed { - Ok(Transformed::yes(plan)) + let res = Arc::new(p) as Arc; + Ok(Transformed::yes(res)) } else { Ok(Transformed::no(plan)) } @@ -389,3 +393,105 @@ fn get_aggregate_exprs_requirement( } Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) } + + +mod tests { + use super::*; + use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions}; + use datafusion_physical_expr::{expressions::{col, OrderSensitiveArrayAgg}, PhysicalSortExpr}; + + fn create_test_schema() -> Result { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, true); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, true); + let e = Field::new("e", DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); + + Ok(schema) + } + + #[tokio::test] + async fn test_get_finest_requirements() -> Result<()> { + let test_schema = create_test_schema()?; + // Assume column a and b are aliases + // Assume also that a ASC and c DESC describe the same global ordering for the table. (Since they are ordering equivalent). + let options1 = SortOptions { + descending: false, + nulls_first: false, + }; + let col_a = &col("a", &test_schema)?; + let col_b = &col("b", &test_schema)?; + let col_c = &col("c", &test_schema)?; + let mut eq_properties = EquivalenceProperties::new(test_schema); + // Columns a and b are equal. + eq_properties.add_equal_conditions(col_a, col_b); + // Aggregate requirements are + // [None], [a ASC], [a ASC, b ASC, c ASC], [a ASC, b ASC] respectively + let order_by_exprs = vec![ + None, + Some(vec![PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }]), + Some(vec![ + PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_b.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_c.clone(), + options: options1, + }, + ]), + Some(vec![ + PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_b.clone(), + options: options1, + }, + ]), + ]; + let common_requirement = vec![ + PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_c.clone(), + options: options1, + }, + ]; + let mut aggr_exprs = order_by_exprs + .into_iter() + .map(|order_by_expr| { + Arc::new(OrderSensitiveArrayAgg::new( + col_a.clone(), + "array_agg", + DataType::Int32, + false, + vec![], + order_by_expr.unwrap_or_default(), + )) as _ + }) + .collect::>(); + let group_by = PhysicalGroupBy::new_single(vec![]); + let res = get_aggregate_exprs_requirement( + &[], + &mut aggr_exprs, + &group_by, + &eq_properties, + &AggregateMode::Partial, + )?; + let res = PhysicalSortRequirement::to_sort_exprs(res); + assert_eq!(res, common_requirement); + Ok(()) + } +} \ No newline at end of file diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 798dfa2d5bfbc..0159b460f3109 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -465,7 +465,9 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { .create_initial_plan(logical_plan, session_state) .await?; - self.optimize_internal(plan, session_state, |_, _| {}) + let res = self.optimize_internal(plan, session_state, |_, _| {}); + println!("optimized done"); + res } } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 63ba87ee44346..4f0cc4495ecb9 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -36,15 +36,13 @@ use arrow::array::ArrayRef; use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::stats::Precision; -use datafusion_common::{internal_err, not_impl_err, Result}; +use datafusion_common::{internal_err, Result}; use datafusion_execution::TaskContext; use datafusion_expr::Accumulator; use datafusion_physical_expr::{ - aggregate::is_order_sensitive, - equivalence::{collapse_lex_req, ProjectionMapping}, - expressions::{Column, FirstValue, LastValue, Max, Min, UnKnownColumn}, - physical_exprs_contains, reverse_order_bys, AggregateExpr, EquivalenceProperties, - LexOrdering, LexRequirement, PhysicalExpr, PhysicalSortRequirement, + equivalence::ProjectionMapping, + expressions::{Column, Max, Min, UnKnownColumn}, + AggregateExpr, LexRequirement, PhysicalExpr, }; use itertools::Itertools; @@ -269,13 +267,42 @@ pub struct AggregateExec { } impl AggregateExec { - pub fn with_required_input_ordering( - self, + pub fn clone_with_input(&self, input: Arc) -> Self { + Self { + input, + // clone the rest of the fields + mode: self.mode, + group_by: self.group_by.clone(), + aggr_expr: self.aggr_expr.clone(), + filter_expr: self.filter_expr.clone(), + limit: self.limit, + schema: self.schema.clone(), + input_schema: self.input_schema.clone(), + metrics: self.metrics.clone(), + input_order_mode: self.input_order_mode.clone(), + cache: self.cache.clone(), + required_input_ordering: self.required_input_ordering.clone(), + } + } + + pub fn clone_with_required_input_ordering( + &self, required_input_ordering: Option, ) -> Self { Self { required_input_ordering, - ..self + // clone the rest of the fields + mode: self.mode, + group_by: self.group_by.clone(), + aggr_expr: self.aggr_expr.clone(), + filter_expr: self.filter_expr.clone(), + limit: self.limit, + input: self.input.clone(), + schema: self.schema.clone(), + input_schema: self.input_schema.clone(), + metrics: self.metrics.clone(), + input_order_mode: self.input_order_mode.clone(), + cache: self.cache.clone(), } } @@ -324,7 +351,7 @@ impl AggregateExec { fn try_new_with_schema( mode: AggregateMode, group_by: PhysicalGroupBy, - mut aggr_expr: Vec>, + aggr_expr: Vec>, filter_expr: Vec>>, input: Arc, input_schema: SchemaRef, @@ -335,35 +362,13 @@ impl AggregateExec { return internal_err!("Inconsistent aggregate expr: {:?} and filter expr: {:?} for AggregateExec, their size should match", aggr_expr, filter_expr); } - let input_eq_properties = input.equivalence_properties(); // Get GROUP BY expressions: let groupby_exprs = group_by.input_exprs(); // If existing ordering satisfies a prefix of the GROUP BY expressions, // prefix requirements with this section. In this case, aggregation will // work more efficiently. let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); - let mut new_requirement = indices - .iter() - .map(|&idx| PhysicalSortRequirement { - expr: groupby_exprs[idx].clone(), - options: None, - }) - .collect::>(); - - println!("correct new_requirement: {:?}", new_requirement); - println!("correct aggr_expr: {:?}", aggr_expr); - println!("correct group_by: {:?}", group_by); - println!("correct input_eq_properties: {:?}", input_eq_properties); - println!("correct mode: {:?}", mode); - let req = get_aggregate_exprs_requirement( - &new_requirement, - &mut aggr_expr, - &group_by, - input_eq_properties, - &mode, - )?; - new_requirement.extend(req); - new_requirement = collapse_lex_req(new_requirement); + let copied_aggr_expr = aggr_expr.clone(); let input_order_mode = if indices.len() == groupby_exprs.len() && !indices.is_empty() { @@ -378,16 +383,6 @@ impl AggregateExec { let projection_mapping = ProjectionMapping::try_new(&group_by.expr, &input.schema())?; - let required_input_ordering = - (!new_requirement.is_empty()).then_some(new_requirement); - - if required_input_ordering.is_some() { - println!( - "correct required_input_ordering: {:?}", - required_input_ordering - ); - } - let required_input_ordering = None; let cache = Self::compute_properties( @@ -400,7 +395,7 @@ impl AggregateExec { Ok(AggregateExec { mode, group_by, - aggr_expr, + aggr_expr: copied_aggr_expr, filter_expr, input, schema, @@ -823,224 +818,6 @@ fn group_schema(schema: &Schema, group_count: usize) -> SchemaRef { Arc::new(Schema::new(group_fields)) } -/// Determines the lexical ordering requirement for an aggregate expression. -/// -/// # Parameters -/// -/// - `aggr_expr`: A reference to an `Arc` representing the -/// aggregate expression. -/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the -/// physical GROUP BY expression. -/// - `agg_mode`: A reference to an `AggregateMode` instance representing the -/// mode of aggregation. -/// -/// # Returns -/// -/// A `LexOrdering` instance indicating the lexical ordering requirement for -/// the aggregate expression. -fn get_aggregate_expr_req( - aggr_expr: &Arc, - group_by: &PhysicalGroupBy, - agg_mode: &AggregateMode, -) -> LexOrdering { - // If the aggregation function is not order sensitive, or the aggregation - // is performing a "second stage" calculation, or all aggregate function - // requirements are inside the GROUP BY expression, then ignore the ordering - // requirement. - if !is_order_sensitive(aggr_expr) || !agg_mode.is_first_stage() { - return vec![]; - } - - let mut req = aggr_expr.order_bys().unwrap_or_default().to_vec(); - - // In non-first stage modes, we accumulate data (using `merge_batch`) from - // different partitions (i.e. merge partial results). During this merge, we - // consider the ordering of each partial result. Hence, we do not need to - // use the ordering requirement in such modes as long as partial results are - // generated with the correct ordering. - if group_by.is_single() { - // Remove all orderings that occur in the group by. These requirements - // will definitely be satisfied -- Each group by expression will have - // distinct values per group, hence all requirements are satisfied. - let physical_exprs = group_by.input_exprs(); - req.retain(|sort_expr| { - !physical_exprs_contains(&physical_exprs, &sort_expr.expr) - }); - } - req -} - -/// Computes the finer ordering for between given existing ordering requirement -/// of aggregate expression. -/// -/// # Parameters -/// -/// * `existing_req` - The existing lexical ordering that needs refinement. -/// * `aggr_expr` - A reference to an aggregate expression trait object. -/// * `group_by` - Information about the physical grouping (e.g group by expression). -/// * `eq_properties` - Equivalence properties relevant to the computation. -/// * `agg_mode` - The mode of aggregation (e.g., Partial, Final, etc.). -/// -/// # Returns -/// -/// An `Option` representing the computed finer lexical ordering, -/// or `None` if there is no finer ordering; e.g. the existing requirement and -/// the aggregator requirement is incompatible. -fn finer_ordering( - existing_req: &LexOrdering, - aggr_expr: &Arc, - group_by: &PhysicalGroupBy, - eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Option { - let aggr_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); - eq_properties.get_finer_ordering(existing_req, &aggr_req) -} - -/// Concatenates the given slices. -fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { - [lhs, rhs].concat() -} - -/// Get the common requirement that satisfies all the aggregate expressions. -/// -/// # Parameters -/// -/// - `aggr_exprs`: A slice of `Arc` containing all the -/// aggregate expressions. -/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the -/// physical GROUP BY expression. -/// - `eq_properties`: A reference to an `EquivalenceProperties` instance -/// representing equivalence properties for ordering. -/// - `agg_mode`: A reference to an `AggregateMode` instance representing the -/// mode of aggregation. -/// -/// # Returns -/// -/// A `LexRequirement` instance, which is the requirement that satisfies all the -/// aggregate requirements. Returns an error in case of conflicting requirements. -fn get_aggregate_exprs_requirement( - prefix_requirement: &[PhysicalSortRequirement], - aggr_exprs: &mut [Arc], - group_by: &PhysicalGroupBy, - eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Result { - let mut requirement = vec![]; - for aggr_expr in aggr_exprs.iter_mut() { - let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); - let reverse_aggr_req = reverse_order_bys(aggr_req); - let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); - let reverse_aggr_req = - PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); - - if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { - let mut first_value = first_value.clone(); - if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &aggr_req, - )) { - first_value = first_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(first_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to LAST_VALUE enables more efficient execution - // given the existing ordering: - let mut last_value = first_value.convert_to_last(); - last_value = last_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(last_value) as _; - } else { - // Requirement is not satisfied with existing ordering. - first_value = first_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(first_value) as _; - } - continue; - } - if let Some(last_value) = aggr_expr.as_any().downcast_ref::() { - let mut last_value = last_value.clone(); - if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &aggr_req, - )) { - last_value = last_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(last_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to FIRST_VALUE enables more efficient execution - // given the existing ordering: - let mut first_value = last_value.convert_to_first(); - first_value = first_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(first_value) as _; - } else { - // Requirement is not satisfied with existing ordering. - last_value = last_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(last_value) as _; - } - continue; - } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Requirement is satisfied by existing ordering - requirement = finer_ordering; - continue; - } - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Reverse requirement is satisfied by exiting ordering. - // Hence reverse the aggregator - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } - } - } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - // There is a requirement that both satisfies existing requirement and current - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - continue; - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - // There is a requirement that both satisfies existing requirement and reverse - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } - } - // Neither the existing requirement and current aggregate requirement satisfy the other, this means - // requirements are conflicting. Currently, we do not support - // conflicting requirements. - return not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" - ); - } - Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) -} - /// returns physical expressions for arguments to evaluate against a batch /// The expressions are different depending on `mode`: /// * Partial: AggregateExpr::expressions @@ -1263,27 +1040,15 @@ mod tests { use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_physical_expr::expressions::{ - lit, ApproxDistinct, Count, LastValue, Median, OrderSensitiveArrayAgg, + lit, ApproxDistinct, Count, FirstValue, LastValue, Median, OrderSensitiveArrayAgg }; use datafusion_physical_expr::{ reverse_order_bys, AggregateExpr, EquivalenceProperties, PhysicalExpr, - PhysicalSortExpr, + PhysicalSortExpr, PhysicalSortRequirement, }; use futures::{FutureExt, Stream}; - // Generate a schema which consists of 5 columns (a, b, c, d, e) - fn create_test_schema() -> Result { - let a = Field::new("a", DataType::Int32, true); - let b = Field::new("b", DataType::Int32, true); - let c = Field::new("c", DataType::Int32, true); - let d = Field::new("d", DataType::Int32, true); - let e = Field::new("e", DataType::Int32, true); - let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); - - Ok(schema) - } - /// some mock data to aggregates fn some_data() -> (Arc, Vec) { // define a schema. @@ -2127,89 +1892,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_get_finest_requirements() -> Result<()> { - let test_schema = create_test_schema()?; - // Assume column a and b are aliases - // Assume also that a ASC and c DESC describe the same global ordering for the table. (Since they are ordering equivalent). - let options1 = SortOptions { - descending: false, - nulls_first: false, - }; - let col_a = &col("a", &test_schema)?; - let col_b = &col("b", &test_schema)?; - let col_c = &col("c", &test_schema)?; - let mut eq_properties = EquivalenceProperties::new(test_schema); - // Columns a and b are equal. - eq_properties.add_equal_conditions(col_a, col_b); - // Aggregate requirements are - // [None], [a ASC], [a ASC, b ASC, c ASC], [a ASC, b ASC] respectively - let order_by_exprs = vec![ - None, - Some(vec![PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }]), - Some(vec![ - PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_b.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_c.clone(), - options: options1, - }, - ]), - Some(vec![ - PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_b.clone(), - options: options1, - }, - ]), - ]; - let common_requirement = vec![ - PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_c.clone(), - options: options1, - }, - ]; - let mut aggr_exprs = order_by_exprs - .into_iter() - .map(|order_by_expr| { - Arc::new(OrderSensitiveArrayAgg::new( - col_a.clone(), - "array_agg", - DataType::Int32, - false, - vec![], - order_by_expr.unwrap_or_default(), - )) as _ - }) - .collect::>(); - let group_by = PhysicalGroupBy::new_single(vec![]); - let res = get_aggregate_exprs_requirement( - &[], - &mut aggr_exprs, - &group_by, - &eq_properties, - &AggregateMode::Partial, - )?; - let res = PhysicalSortRequirement::to_sort_exprs(res); - assert_eq!(res, common_requirement); - Ok(()) - } #[test] fn test_agg_exec_same_schema() -> Result<()> { diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 1c725ce31f146..b95d17b627b96 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -72,6 +72,14 @@ impl CoalescePartitionsExec { input.execution_mode(), // Execution Mode ) } + + pub fn clone_with_input(&self, input: Arc) -> Self { + CoalescePartitionsExec { + input, + metrics: self.metrics.clone(), + cache: self.cache.clone(), + } + } } impl DisplayAs for CoalescePartitionsExec { diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 4929ab485d6d7..135cded2e94fe 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -123,7 +123,7 @@ LOCATION '../core/tests/data/aggregate_agg_multi_order.csv'; query ? select array_agg(c1 order by c2 desc, c3) from agg_order; ---- -[5, 6, 7, 8, 9, 1, 2, 3, 4, 10] +[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] query TT explain select array_agg(c1 order by c2 desc, c3) from agg_order; @@ -135,9 +135,21 @@ physical_plan AggregateExec: mode=Final, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] --CoalescePartitionsExec ----AggregateExec: mode=Partial, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] -------SortExec: expr=[c2@1 DESC,c3@2 ASC NULLS LAST] ---------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -----------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true +------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true + +query TT +explain select array_agg(c1 order by c2 desc, c3) from agg_order; +---- +logical_plan +Aggregate: groupBy=[[]], aggr=[[ARRAY_AGG(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]]] +--TableScan: agg_order projection=[c1, c2, c3] +physical_plan +AggregateExec: mode=Final, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] +--CoalescePartitionsExec +----AggregateExec: mode=Partial, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] +------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true # test array_agg_order with list data type statement ok @@ -152,8 +164,8 @@ CREATE TABLE array_agg_order_list_table AS VALUES query T? rowsort select column1, array_agg(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [[7, 8, 9], [4, 5, 6]] -w [[3, 2, 5], [9, 5, 2], [1, 2, 3]] +b [[4, 5, 6], [7, 8, 9]] +w [[1, 2, 3], [9, 5, 2], [3, 2, 5]] query T?? rowsort select column1, first_value(column3 order by column2, column4 desc), last_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; @@ -164,7 +176,7 @@ w [3, 2, 5] [1, 2, 3] query T? rowsort select column1, nth_value(column3, 2 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [4, 5, 6] +b [7, 8, 9] w [9, 5, 2] statement ok From 3f43dc19d1de68e323e216e50c18ce77a6c26eea Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 14:59:57 +0800 Subject: [PATCH 05/33] cleanup Signed-off-by: jayzhan211 --- datafusion/core/src/physical_optimizer/simplify_ordering.rs | 2 +- datafusion/core/src/physical_planner.rs | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index bb02df69b92be..67c9641687877 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -82,7 +82,7 @@ fn get_common_requirement_of_aggregate_input( let new_c: Option>> = if children.is_empty() { None } else { - assert_eq!(children.len(), 1, "AggregateExec should have one child"); + assert_eq!(children.len(), 1, "children: {:?}", children); let c = children[0].clone(); // for c in children { diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 0159b460f3109..798dfa2d5bfbc 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -465,9 +465,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { .create_initial_plan(logical_plan, session_state) .await?; - let res = self.optimize_internal(plan, session_state, |_, _| {}); - println!("optimized done"); - res + self.optimize_internal(plan, session_state, |_, _| {}) } } } From c8b46a49d4a0cc7a5e74ec89733e20e5f7234473 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 15:40:45 +0800 Subject: [PATCH 06/33] with new childes Signed-off-by: jayzhan211 --- .../core/src/physical_optimizer/optimizer.rs | 8 ++ .../physical_optimizer/simplify_ordering.rs | 85 +++++++++---------- .../physical-plan/src/aggregates/mod.rs | 4 +- 3 files changed, 53 insertions(+), 44 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index b777933e120e0..e7fe76d429691 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -76,11 +76,13 @@ impl PhysicalOptimizer { /// Create a new optimizer using the recommended list of rules pub fn new() -> Self { let rules: Vec> = vec![ + // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), Arc::new(AggregateStatistics::new()), + // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the @@ -92,15 +94,18 @@ impl PhysicalOptimizer { // as that rule may inject other operations in between the different AggregateExecs. // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), + // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), // 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()), + // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule Arc::new(CombinePartialFinalAggregate::new()), + // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), // The EnforceSorting rule is for adding essential local sorting to satisfy the required // ordering. Please make sure that the whole plan tree is determined before this rule. @@ -115,6 +120,8 @@ impl PhysicalOptimizer { // Remove the ancillary output requirement operator since we are done with the planning // phase. Arc::new(OutputRequirements::new_remove_mode()), + // Appears after AggregateExec is created + Arc::new(SimplifyOrdering::new()), // The PipelineChecker rule will reject non-runnable query plans that use // pipeline-breaking operators on infinite input(s). The rule generates a // diagnostic error message when this happens. It makes no changes to the @@ -125,6 +132,7 @@ impl PhysicalOptimizer { // into an `order by max(x) limit y`. In this case it will copy the limit value down // to the aggregation, allowing it to use only y number of accumulators. Arc::new(TopKAggregation::new()), + // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), // The ProjectionPushdown rule tries to push projections towards // the sources in the execution plan. As a result of this process, diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index 67c9641687877..5f7faa92b789a 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -79,53 +79,52 @@ fn get_common_requirement_of_aggregate_input( ) -> Result>> { let children = plan.children(); - let new_c: Option>> = if children.is_empty() { - None + let mut is_transformed = false; + let mut new_children: Vec> = vec![]; + for c in children.iter() { + let res = get_common_requirement_of_aggregate_input(c.clone())?; + if res.transformed { + is_transformed = true; + } + new_children.push(res.data); + } + + let plan = if is_transformed { + plan.with_new_children(new_children)? } else { - assert_eq!(children.len(), 1, "children: {:?}", children); - let c = children[0].clone(); - - // for c in children { - // let new_c = get_common_requirement_of_aggregate_input(c.clone())?; - // if new_c.transformed { - // is_transformed = true; - // } - // } - - let new_c = get_common_requirement_of_aggregate_input(c)?; - Some(new_c) + plan }; let plan = optimize_internal(plan)?; - // println!("t: {} plan: {:?}", plan.transformed, plan); - - if let Some(c) = new_c { - if !c.transformed { - return Ok(plan); - } - - let plan = plan.data; - - // TODO: support more types of ExecutionPlan - if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let p = aggr_exec.clone_with_input(c.data); - return Ok(Transformed::yes(Arc::new(p) as Arc)); - } else if let Some(coalesce_exec) = - plan.as_any().downcast_ref::() - { - let p = coalesce_exec.clone_with_input(c.data); - return Ok(Transformed::yes(Arc::new(p) as Arc)); - } else if let Some(out_req_exec) = - plan.as_any().downcast_ref::() - { - let p = out_req_exec.clone_with_input(c.data); - return Ok(Transformed::yes(Arc::new(p) as Arc)); - } else { - return not_impl_err!("Unsupported ExecutionPlan type: {}", plan.name()); - } - } - - return Ok(plan); + Ok(plan) + + // if let Some(c) = new_c { + // if !c.transformed { + // return Ok(plan); + // } + + // let plan = plan.data; + + // // TODO: support more types of ExecutionPlan + // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.clone_with_input(c.data); + // return Ok(Transformed::yes(Arc::new(p) as Arc)); + // } else if let Some(coalesce_exec) = + // plan.as_any().downcast_ref::() + // { + // let p = coalesce_exec.clone_with_input(c.data); + // return Ok(Transformed::yes(Arc::new(p) as Arc)); + // } else if let Some(out_req_exec) = + // plan.as_any().downcast_ref::() + // { + // let p = out_req_exec.clone_with_input(c.data); + // return Ok(Transformed::yes(Arc::new(p) as Arc)); + // } else { + // return not_impl_err!("Unsupported ExecutionPlan type: {}", plan.name()); + // } + // } + + // return Ok(plan); } fn optimize_internal( diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 4f0cc4495ecb9..30fd28b538e82 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -698,6 +698,9 @@ impl ExecutionPlan for AggregateExec { self: Arc, children: Vec>, ) -> Result> { + // let aggr_exec = self.clone_with_input(children[0].clone()); + // Ok(Arc::new(aggr_exec)) + let mut me = AggregateExec::try_new_with_schema( self.mode, self.group_by.clone(), @@ -706,7 +709,6 @@ impl ExecutionPlan for AggregateExec { children[0].clone(), self.input_schema.clone(), self.schema.clone(), - //self.original_schema.clone(), )?; me.limit = self.limit; Ok(Arc::new(me)) From fcb297b7c9943d0d85b5996ec0eec1efc1e2c524 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 16:09:34 +0800 Subject: [PATCH 07/33] revert slt Signed-off-by: jayzhan211 --- .../sqllogictest/test_files/aggregate.slt | 28 ++++++------------- 1 file changed, 8 insertions(+), 20 deletions(-) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 135cded2e94fe..b3a58dd582e7d 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -123,7 +123,7 @@ LOCATION '../core/tests/data/aggregate_agg_multi_order.csv'; query ? select array_agg(c1 order by c2 desc, c3) from agg_order; ---- -[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] +[5, 6, 7, 8, 9, 1, 2, 3, 4, 10] query TT explain select array_agg(c1 order by c2 desc, c3) from agg_order; @@ -135,21 +135,9 @@ physical_plan AggregateExec: mode=Final, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] --CoalescePartitionsExec ----AggregateExec: mode=Partial, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] -------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 ---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true - -query TT -explain select array_agg(c1 order by c2 desc, c3) from agg_order; ----- -logical_plan -Aggregate: groupBy=[[]], aggr=[[ARRAY_AGG(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]]] ---TableScan: agg_order projection=[c1, c2, c3] -physical_plan -AggregateExec: mode=Final, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] ---CoalescePartitionsExec -----AggregateExec: mode=Partial, gby=[], aggr=[ARRAY_AGG(agg_order.c1)] -------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 ---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true +------SortExec: expr=[c2@1 DESC,c3@2 ASC NULLS LAST] +--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +----------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], has_header=true # test array_agg_order with list data type statement ok @@ -164,8 +152,8 @@ CREATE TABLE array_agg_order_list_table AS VALUES query T? rowsort select column1, array_agg(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [[4, 5, 6], [7, 8, 9]] -w [[1, 2, 3], [9, 5, 2], [3, 2, 5]] +b [[7, 8, 9], [4, 5, 6]] +w [[3, 2, 5], [9, 5, 2], [1, 2, 3]] query T?? rowsort select column1, first_value(column3 order by column2, column4 desc), last_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; @@ -176,7 +164,7 @@ w [3, 2, 5] [1, 2, 3] query T? rowsort select column1, nth_value(column3, 2 order by column2, column4 desc) from array_agg_order_list_table group by column1; ---- -b [7, 8, 9] +b [4, 5, 6] w [9, 5, 2] statement ok @@ -3439,4 +3427,4 @@ SELECT LAST_VALUE(column1 ORDER BY column2 DESC) IGNORE NULLS FROM t; 3 statement ok -DROP TABLE t; +DROP TABLE t; \ No newline at end of file From ab10168bf050b549bea9a0b3823a3f8b759161ec Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 16:24:44 +0800 Subject: [PATCH 08/33] revert back Signed-off-by: jayzhan211 --- .../physical_optimizer/simplify_ordering.rs | 11 +- .../physical-plan/src/aggregates/mod.rs | 303 +++++++++++++++++- 2 files changed, 302 insertions(+), 12 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index 5f7faa92b789a..e6ddcdc257208 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -57,12 +57,13 @@ impl PhysicalOptimizerRule for SimplifyOrdering { plan: Arc, _config: &ConfigOptions, ) -> Result> { - let res = plan - .transform_down(&get_common_requirement_of_aggregate_input) - .data(); + Ok(plan) + // let res = plan + // .transform_down(&get_common_requirement_of_aggregate_input) + // .data(); - // println!("res: {:?}", res); - res + // // println!("res: {:?}", res); + // res } fn name(&self) -> &str { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 30fd28b538e82..c9ade05711c70 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -36,9 +36,13 @@ use arrow::array::ArrayRef; use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::stats::Precision; -use datafusion_common::{internal_err, Result}; +use datafusion_common::{internal_err, not_impl_err, Result}; use datafusion_execution::TaskContext; use datafusion_expr::Accumulator; +use datafusion_physical_expr::aggregate::is_order_sensitive; +use datafusion_physical_expr::equivalence::collapse_lex_req; +use datafusion_physical_expr::expressions::{FirstValue, LastValue}; +use datafusion_physical_expr::{physical_exprs_contains, reverse_order_bys, EquivalenceProperties, LexOrdering, PhysicalSortRequirement}; use datafusion_physical_expr::{ equivalence::ProjectionMapping, expressions::{Column, Max, Min, UnKnownColumn}, @@ -328,7 +332,7 @@ impl AggregateExec { )?; let schema = Arc::new(schema); - AggregateExec::try_new_with_schema( + let res = AggregateExec::try_new_with_schema( mode, group_by, aggr_expr, @@ -336,7 +340,50 @@ impl AggregateExec { input, input_schema, schema, - ) + )?; + + res.rewrite_ordering() + } + + fn rewrite_ordering(mut self) -> Result { + let input = self.input(); + let group_by = self.group_by(); + let mut aggr_expr = self.aggr_expr().to_vec(); + let mode = self.mode(); + + let input_eq_properties = input.equivalence_properties(); + // Get GROUP BY expressions: + let groupby_exprs = group_by.input_exprs(); + // If existing ordering satisfies a prefix of the GROUP BY expressions, + // prefix requirements with this section. In this case, aggregation will + // work more efficiently. + let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); + + let mut new_requirement = indices + .iter() + .map(|&idx| PhysicalSortRequirement { + expr: groupby_exprs[idx].clone(), + options: None, + }) + .collect::>(); + + let req = get_aggregate_exprs_requirement( + &new_requirement, + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + new_requirement.extend(req); + new_requirement = collapse_lex_req(new_requirement); + + let required_input_ordering = + (!new_requirement.is_empty()).then_some(new_requirement); + + self.aggr_expr = aggr_expr; + self.required_input_ordering = required_input_ordering; + + Ok(self) } /// Create a new hash aggregate execution plan with the given schema. @@ -351,7 +398,7 @@ impl AggregateExec { fn try_new_with_schema( mode: AggregateMode, group_by: PhysicalGroupBy, - aggr_expr: Vec>, + mut aggr_expr: Vec>, filter_expr: Vec>>, input: Arc, input_schema: SchemaRef, @@ -362,13 +409,30 @@ impl AggregateExec { return internal_err!("Inconsistent aggregate expr: {:?} and filter expr: {:?} for AggregateExec, their size should match", aggr_expr, filter_expr); } + let input_eq_properties = input.equivalence_properties(); // Get GROUP BY expressions: let groupby_exprs = group_by.input_exprs(); // If existing ordering satisfies a prefix of the GROUP BY expressions, // prefix requirements with this section. In this case, aggregation will // work more efficiently. let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); - let copied_aggr_expr = aggr_expr.clone(); + let mut new_requirement = indices + .iter() + .map(|&idx| PhysicalSortRequirement { + expr: groupby_exprs[idx].clone(), + options: None, + }) + .collect::>(); + + let req = get_aggregate_exprs_requirement( + &new_requirement, + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + new_requirement.extend(req); + new_requirement = collapse_lex_req(new_requirement); let input_order_mode = if indices.len() == groupby_exprs.len() && !indices.is_empty() { @@ -383,7 +447,8 @@ impl AggregateExec { let projection_mapping = ProjectionMapping::try_new(&group_by.expr, &input.schema())?; - let required_input_ordering = None; + let required_input_ordering = + (!new_requirement.is_empty()).then_some(new_requirement); let cache = Self::compute_properties( &input, @@ -392,10 +457,12 @@ impl AggregateExec { &mode, &input_order_mode, ); + + let required_input_ordering = None; Ok(AggregateExec { mode, group_by, - aggr_expr: copied_aggr_expr, + aggr_expr, filter_expr, input, schema, @@ -408,6 +475,7 @@ impl AggregateExec { }) } + /// Aggregation mode (full, partial) pub fn mode(&self) -> &AggregateMode { &self.mode @@ -711,6 +779,9 @@ impl ExecutionPlan for AggregateExec { self.schema.clone(), )?; me.limit = self.limit; + + me = me.rewrite_ordering()?; + Ok(Arc::new(me)) } @@ -820,6 +891,224 @@ fn group_schema(schema: &Schema, group_count: usize) -> SchemaRef { Arc::new(Schema::new(group_fields)) } +/// Determines the lexical ordering requirement for an aggregate expression. +/// +/// # Parameters +/// +/// - `aggr_expr`: A reference to an `Arc` representing the +/// aggregate expression. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// +/// # Returns +/// +/// A `LexOrdering` instance indicating the lexical ordering requirement for +/// the aggregate expression. +fn get_aggregate_expr_req( + aggr_expr: &Arc, + group_by: &PhysicalGroupBy, + agg_mode: &AggregateMode, +) -> LexOrdering { + // If the aggregation function is not order sensitive, or the aggregation + // is performing a "second stage" calculation, or all aggregate function + // requirements are inside the GROUP BY expression, then ignore the ordering + // requirement. + if !is_order_sensitive(aggr_expr) || !agg_mode.is_first_stage() { + return vec![]; + } + + let mut req = aggr_expr.order_bys().unwrap_or_default().to_vec(); + + // In non-first stage modes, we accumulate data (using `merge_batch`) from + // different partitions (i.e. merge partial results). During this merge, we + // consider the ordering of each partial result. Hence, we do not need to + // use the ordering requirement in such modes as long as partial results are + // generated with the correct ordering. + if group_by.is_single() { + // Remove all orderings that occur in the group by. These requirements + // will definitely be satisfied -- Each group by expression will have + // distinct values per group, hence all requirements are satisfied. + let physical_exprs = group_by.input_exprs(); + req.retain(|sort_expr| { + !physical_exprs_contains(&physical_exprs, &sort_expr.expr) + }); + } + req +} + +/// Computes the finer ordering for between given existing ordering requirement +/// of aggregate expression. +/// +/// # Parameters +/// +/// * `existing_req` - The existing lexical ordering that needs refinement. +/// * `aggr_expr` - A reference to an aggregate expression trait object. +/// * `group_by` - Information about the physical grouping (e.g group by expression). +/// * `eq_properties` - Equivalence properties relevant to the computation. +/// * `agg_mode` - The mode of aggregation (e.g., Partial, Final, etc.). +/// +/// # Returns +/// +/// An `Option` representing the computed finer lexical ordering, +/// or `None` if there is no finer ordering; e.g. the existing requirement and +/// the aggregator requirement is incompatible. +fn finer_ordering( + existing_req: &LexOrdering, + aggr_expr: &Arc, + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Option { + let aggr_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); + eq_properties.get_finer_ordering(existing_req, &aggr_req) +} + +/// Concatenates the given slices. +fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { + [lhs, rhs].concat() +} + +/// Get the common requirement that satisfies all the aggregate expressions. +/// +/// # Parameters +/// +/// - `aggr_exprs`: A slice of `Arc` containing all the +/// aggregate expressions. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `eq_properties`: A reference to an `EquivalenceProperties` instance +/// representing equivalence properties for ordering. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// +/// # Returns +/// +/// A `LexRequirement` instance, which is the requirement that satisfies all the +/// aggregate requirements. Returns an error in case of conflicting requirements. +fn get_aggregate_exprs_requirement( + prefix_requirement: &[PhysicalSortRequirement], + aggr_exprs: &mut [Arc], + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Result { + let mut requirement = vec![]; + for aggr_expr in aggr_exprs.iter_mut() { + let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); + let reverse_aggr_req = reverse_order_bys(aggr_req); + let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); + let reverse_aggr_req = + PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); + + if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { + let mut first_value = first_value.clone(); + if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &aggr_req, + )) { + first_value = first_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(first_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to LAST_VALUE enables more efficient execution + // given the existing ordering: + let mut last_value = first_value.convert_to_last(); + last_value = last_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(last_value) as _; + } else { + // Requirement is not satisfied with existing ordering. + first_value = first_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(first_value) as _; + } + continue; + } + if let Some(last_value) = aggr_expr.as_any().downcast_ref::() { + let mut last_value = last_value.clone(); + if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &aggr_req, + )) { + last_value = last_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(last_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to FIRST_VALUE enables more efficient execution + // given the existing ordering: + let mut first_value = last_value.convert_to_first(); + first_value = first_value.with_requirement_satisfied(true); + *aggr_expr = Arc::new(first_value) as _; + } else { + // Requirement is not satisfied with existing ordering. + last_value = last_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(last_value) as _; + } + continue; + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Requirement is satisfied by existing ordering + requirement = finer_ordering; + continue; + } + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Reverse requirement is satisfied by exiting ordering. + // Hence reverse the aggregator + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + // There is a requirement that both satisfies existing requirement and current + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + continue; + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + // There is a requirement that both satisfies existing requirement and reverse + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + // Neither the existing requirement and current aggregate requirement satisfy the other, this means + // requirements are conflicting. Currently, we do not support + // conflicting requirements. + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ); + } + Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) +} + /// returns physical expressions for arguments to evaluate against a batch /// The expressions are different depending on `mode`: /// * Partial: AggregateExpr::expressions From fffb54a88403b58e8f33875fbf840c1f25857742 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 17:07:47 +0800 Subject: [PATCH 09/33] rm rewrite in new child Signed-off-by: jayzhan211 --- .../enforce_distribution.rs | 9 ++++++- .../physical_optimizer/output_requirements.rs | 13 +++++++++- .../physical-plan/src/aggregates/mod.rs | 25 +++++++++++++++---- datafusion/physical-plan/src/lib.rs | 12 ++++++++- .../physical-plan/src/recursive_query.rs | 12 ++++++++- 5 files changed, 62 insertions(+), 9 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index 145f08af76dd1..81e7c2ba6c85a 100644 --- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs @@ -1217,7 +1217,14 @@ fn ensure_distribution( // Data Arc::new(InterleaveExec::try_new(children_plans)?) } else { - plan.with_new_children(children_plans)? + let plan = plan.with_new_children(children_plans)?; + + if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + let p = aggr_exec.rewrite_ordering()?; + Arc::new(p) + } else { + plan + } }; Ok(Transformed::yes(DistributionContext::new( diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 386340bbeda69..0f7e86b387478 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -32,6 +32,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Result, Statistics}; use datafusion_physical_expr::{Distribution, LexRequirement, PhysicalSortRequirement}; +use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ExecutionPlanProperties, PlanProperties}; @@ -283,7 +284,17 @@ fn require_top_ordering_helper( // be responsible for (i.e. the originator of) the global ordering. let (new_child, is_changed) = require_top_ordering_helper(children.swap_remove(0))?; - Ok((plan.with_new_children(vec![new_child])?, is_changed)) + + let plan = plan.with_new_children(vec![new_child])?; + + if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + let p = aggr_exec.rewrite_ordering()?; + Ok((Arc::new(p) as _, is_changed)) + } else { + Ok((plan, is_changed)) + } + + // Ok((plan.with_new_children(vec![new_child])?, is_changed)) } else { // Stop searching, there is no global ordering desired for the query. Ok((plan, false)) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index c9ade05711c70..1f0919e3ba739 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -345,7 +345,7 @@ impl AggregateExec { res.rewrite_ordering() } - fn rewrite_ordering(mut self) -> Result { + pub fn rewrite_ordering(&self) -> Result { let input = self.input(); let group_by = self.group_by(); let mut aggr_expr = self.aggr_expr().to_vec(); @@ -380,10 +380,25 @@ impl AggregateExec { let required_input_ordering = (!new_requirement.is_empty()).then_some(new_requirement); - self.aggr_expr = aggr_expr; - self.required_input_ordering = required_input_ordering; + Ok(Self { + mode: self.mode().clone(), + group_by: self.group_by().clone(), + aggr_expr, + filter_expr: self.filter_expr().to_vec(), + input: self.input().clone(), + schema: self.schema.clone(), + input_schema: self.input_schema.clone(), + metrics: self.metrics.clone(), + required_input_ordering, + limit: self.limit, + input_order_mode: self.input_order_mode.clone(), + cache: self.cache.clone(), + }) + + // self.aggr_expr = aggr_expr; + // self.required_input_ordering = required_input_ordering; - Ok(self) + // Ok(()) } /// Create a new hash aggregate execution plan with the given schema. @@ -780,7 +795,7 @@ impl ExecutionPlan for AggregateExec { )?; me.limit = self.limit; - me = me.rewrite_ordering()?; + // let me = me.rewrite_ordering()?; Ok(Arc::new(me)) } diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index e1c8489655bf5..4407e8aa2cd69 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -27,6 +27,7 @@ use crate::metrics::MetricsSet; use crate::repartition::RepartitionExec; use crate::sorts::sort_preserving_merge::SortPreservingMergeExec; +use aggregates::AggregateExec; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; @@ -673,7 +674,16 @@ pub fn with_new_children_if_necessary( .zip(old_children.iter()) .any(|(c1, c2)| !Arc::data_ptr_eq(c1, c2)) { - plan.with_new_children(children) + let plan = plan.with_new_children(children)?; + + if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + let p = aggr_exec.rewrite_ordering()?; + Ok(Arc::new(p)) + } else { + Ok(plan) + } + + } else { Ok(plan) } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index ba7d1a54548a4..f8b8c6b5e1c92 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -26,6 +26,7 @@ use super::{ work_table::{ReservedBatches, WorkTable, WorkTableExec}, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; +use crate::aggregates::AggregateExec; use crate::{DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan}; use arrow::datatypes::SchemaRef; @@ -359,7 +360,16 @@ fn reset_plan_states(plan: Arc) -> Result() { + let p = aggr_exec.rewrite_ordering()?; + Ok(Transformed::yes(Arc::new(p))) + } else { + Ok(Transformed::yes(new_plan)) + } + + // Ok(Transformed::yes(new_plan)) } }) .data() From fec3eaf4012a31cdf06c849939afa76b2d23b42f Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 18:16:13 +0800 Subject: [PATCH 10/33] backup Signed-off-by: jayzhan211 --- .../enforce_distribution.rs | 13 ++++--- .../core/src/physical_optimizer/optimizer.rs | 2 + .../physical_optimizer/output_requirements.rs | 13 ++++--- .../replace_with_order_preserving_variants.rs | 23 +++++++++++ .../physical_optimizer/simplify_ordering.rs | 38 ++++++++++++++----- datafusion/physical-plan/src/lib.rs | 13 ++++--- datafusion/physical-plan/src/tree_node.rs | 13 ++++++- .../sqllogictest/test_files/aggregate.slt | 2 +- .../sqllogictest/test_files/explain.slt | 32 ++++++++++++++-- .../sqllogictest/test_files/group_by.slt | 3 ++ 10 files changed, 120 insertions(+), 32 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index 81e7c2ba6c85a..bb1704a56c2bd 100644 --- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs @@ -1219,12 +1219,13 @@ fn ensure_distribution( } else { let plan = plan.with_new_children(children_plans)?; - if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let p = aggr_exec.rewrite_ordering()?; - Arc::new(p) - } else { - plan - } + // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.rewrite_ordering()?; + // Arc::new(p) + // } else { + // plan + // } + plan }; Ok(Transformed::yes(DistributionContext::new( diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index e7fe76d429691..92393940cc63d 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -81,6 +81,7 @@ impl PhysicalOptimizer { // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), + Arc::new(SimplifyOrdering::new()), Arc::new(AggregateStatistics::new()), // Appears after AggregateExec is created Arc::new(SimplifyOrdering::new()), @@ -112,6 +113,7 @@ impl PhysicalOptimizer { // Note that one should always run this rule after running the EnforceDistribution rule // as the latter may break local sorting requirements. Arc::new(EnforceSorting::new()), + Arc::new(SimplifyOrdering::new()), // TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future. Arc::new(ProjectionPushdown::new()), // The CoalesceBatches rule will not influence the distribution and ordering of the diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 0f7e86b387478..62588629f0a67 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -287,12 +287,13 @@ fn require_top_ordering_helper( let plan = plan.with_new_children(vec![new_child])?; - if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let p = aggr_exec.rewrite_ordering()?; - Ok((Arc::new(p) as _, is_changed)) - } else { - Ok((plan, is_changed)) - } + // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.rewrite_ordering()?; + // Ok((Arc::new(p) as _, is_changed)) + // } else { + // } + + Ok((plan, is_changed)) // Ok((plan.with_new_children(vec![new_child])?, is_changed)) } else { 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 ad19215fbf67e..7479dec4efbc3 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 @@ -29,6 +29,7 @@ use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::Transformed; +use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::ExecutionPlanProperties; @@ -138,6 +139,18 @@ fn plan_with_order_preserving_variants( } } + // let mut plan_context = sort_input.update_plan_from_children()?; + // let plan = plan_context.plan.clone(); + // let p = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.rewrite_ordering()?; + // Arc::new(p) + // } else { + // plan + // }; + // plan_context.plan = p; + + // Ok(plan_context) + sort_input.update_plan_from_children() } @@ -184,6 +197,16 @@ fn plan_with_order_breaking_variants( let coalesce = CoalescePartitionsExec::new(child); sort_input.plan = Arc::new(coalesce) as _; } else { + // let mut pc = sort_input.update_plan_from_children()?; + // let plan = pc.plan.clone(); + // let p = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.rewrite_ordering()?; + // Arc::new(p) + // } else { + // plan + // }; + // pc.plan = p; + // return Ok(pc); return sort_input.update_plan_from_children(); } diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index e6ddcdc257208..6d76d4ff4d02a 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -57,13 +57,13 @@ impl PhysicalOptimizerRule for SimplifyOrdering { plan: Arc, _config: &ConfigOptions, ) -> Result> { - Ok(plan) - // let res = plan - // .transform_down(&get_common_requirement_of_aggregate_input) - // .data(); + // Ok(plan) + let res = plan + .transform_down(&get_common_requirement_of_aggregate_input) + .data(); - // // println!("res: {:?}", res); - // res + // println!("res: {:?}", res); + res } fn name(&self) -> &str { @@ -90,13 +90,33 @@ fn get_common_requirement_of_aggregate_input( new_children.push(res.data); } + let mode1 = plan.properties().execution_mode(); + // println!("mode 1: {:?}", mode); + let plan = if is_transformed { + + // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.clone_with_input(new_children[0].clone()); + // Arc::new(p) as Arc + // } else { + // plan.with_new_children(new_children)? + // } plan.with_new_children(new_children)? + } else { plan }; + let mode2 = plan.properties().execution_mode(); + if mode1 != mode2 { + println!("mode1: {:?}, mode2: {:?}", mode1, mode2); + } + let plan = optimize_internal(plan)?; + let mode3 = plan.data.properties().execution_mode(); + if mode1 != mode3 { + println!("mode1: {:?}, mode3: {:?}", mode1, mode3); + } Ok(plan) // if let Some(c) = new_c { @@ -132,9 +152,9 @@ fn optimize_internal( plan: Arc, ) -> Result>> { if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - if aggr_exec.mode() != &AggregateMode::Partial { - return Ok(Transformed::no(plan)); - } + // if aggr_exec.mode() != &AggregateMode::Partial { + // return Ok(Transformed::no(plan)); + // } let input = aggr_exec.input().clone(); let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 4407e8aa2cd69..3f25855230b0d 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -676,12 +676,13 @@ pub fn with_new_children_if_necessary( { let plan = plan.with_new_children(children)?; - if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let p = aggr_exec.rewrite_ordering()?; - Ok(Arc::new(p)) - } else { - Ok(plan) - } + // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.rewrite_ordering()?; + // Ok(Arc::new(p)) + // } else { + // Ok(plan) + // } + Ok(plan) } else { diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index 52a52f81bdaf0..435aa625710b8 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -20,6 +20,7 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; +use crate::aggregates::AggregateExec; use crate::{displayable, with_new_children_if_necessary, ExecutionPlan}; use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode}; @@ -63,7 +64,17 @@ impl PlanContext { pub fn update_plan_from_children(mut self) -> Result { let children_plans = self.children.iter().map(|c| c.plan.clone()).collect(); - self.plan = with_new_children_if_necessary(self.plan, children_plans)?; + let plan = with_new_children_if_necessary(self.plan, children_plans)?; + + let plan = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + let p = aggr_exec.rewrite_ordering()?; + Arc::new(p) + } else { + plan + }; + + self.plan = plan; + Ok(self) } } diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index b3a58dd582e7d..4929ab485d6d7 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3427,4 +3427,4 @@ SELECT LAST_VALUE(column1 ORDER BY column2 DESC) IGNORE NULLS FROM t; 3 statement ok -DROP TABLE t; \ No newline at end of file +DROP TABLE t; diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 57ee8c311f6c6..848ba40963592 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -239,22 +239,30 @@ logical_plan after optimize_projections SAME TEXT AS ABOVE 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 initial_physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] +physical_plan after SimpleOrdering CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan after OutputRequirements OutputRequirementExec --CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceSorting SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] @@ -295,25 +303,33 @@ EXPLAIN VERBOSE SELECT * FROM alltypes_plain limit 10; initial_physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after OutputRequirements OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] ----ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceSorting SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -332,25 +348,35 @@ GlobalLimitExec: skip=0, fetch=10 initial_physical_plan_with_stats GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after SimpleOrdering +GlobalLimitExec: skip=0, fetch=10 +--ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 physical_plan after OutputRequirements OutputRequirementExec --GlobalLimitExec: skip=0, fetch=10 ----ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceSorting SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements GlobalLimitExec: skip=0, fetch=10 --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10 --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 869462b4722a8..748a02aa0afe6 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -4249,6 +4249,9 @@ SELECT date_bin('15 minutes', ts) as time_chunks 2018-12-13T12:00:00 2018-11-13T17:00:00 + + + # Since extract is not a monotonic function, below query should not run. # when source is unbounded. query error From 5b68b482b20778ef0046c7772549faf1288d6194 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Wed, 10 Apr 2024 22:07:56 +0800 Subject: [PATCH 11/33] only move conversion to optimizer Signed-off-by: jayzhan211 --- .../enforce_distribution.rs | 15 ++- .../core/src/physical_optimizer/optimizer.rs | 21 ++-- .../physical_optimizer/output_requirements.rs | 12 +-- .../replace_with_order_preserving_variants.rs | 22 ---- .../physical_optimizer/simplify_ordering.rs | 6 -- .../physical-plan/src/aggregates/mod.rs | 100 ++++-------------- datafusion/physical-plan/src/lib.rs | 12 +-- .../physical-plan/src/recursive_query.rs | 11 +- datafusion/physical-plan/src/tree_node.rs | 9 +- .../sqllogictest/test_files/explain.slt | 29 ----- .../sqllogictest/test_files/group_by.slt | 15 +-- 11 files changed, 50 insertions(+), 202 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index bb1704a56c2bd..fe43890bc8501 100644 --- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs @@ -1080,6 +1080,11 @@ fn ensure_distribution( } }; + // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { + // let p = aggr_exec.rewrite_ordering()?; + // plan = Arc::new(p); + // } + // 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 @@ -1217,15 +1222,7 @@ fn ensure_distribution( // Data Arc::new(InterleaveExec::try_new(children_plans)?) } else { - let plan = plan.with_new_children(children_plans)?; - - // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.rewrite_ordering()?; - // Arc::new(p) - // } else { - // plan - // } - plan + plan.with_new_children(children_plans)? }; Ok(Transformed::yes(DistributionContext::new( diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 92393940cc63d..28d7e9b17fed0 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -77,14 +77,14 @@ impl PhysicalOptimizer { pub fn new() -> Self { let rules: Vec> = vec![ // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), Arc::new(AggregateStatistics::new()), // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the // EnforceDistribution and EnforceSorting rules as they decide whether to add additional @@ -96,24 +96,24 @@ impl PhysicalOptimizer { // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // 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()), // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule Arc::new(CombinePartialFinalAggregate::new()), // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // The EnforceSorting rule is for adding essential local sorting to satisfy the required // ordering. Please make sure that the whole plan tree is determined before this rule. // Note that one should always run this rule after running the EnforceDistribution rule // as the latter may break local sorting requirements. Arc::new(EnforceSorting::new()), - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future. Arc::new(ProjectionPushdown::new()), // The CoalesceBatches rule will not influence the distribution and ordering of the @@ -123,7 +123,7 @@ impl PhysicalOptimizer { // phase. Arc::new(OutputRequirements::new_remove_mode()), // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // The PipelineChecker rule will reject non-runnable query plans that use // pipeline-breaking operators on infinite input(s). The rule generates a // diagnostic error message when this happens. It makes no changes to the @@ -134,8 +134,7 @@ impl PhysicalOptimizer { // into an `order by max(x) limit y`. In this case it will copy the limit value down // to the aggregation, allowing it to use only y number of accumulators. Arc::new(TopKAggregation::new()), - // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), + // Arc::new(SimplifyOrdering::new()), // The ProjectionPushdown rule tries to push projections towards // the sources in the execution plan. As a result of this process, // a projection can disappear if it reaches the source providers, and @@ -143,6 +142,8 @@ impl PhysicalOptimizer { // are not present, the load of executors such as join or union will be // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), + // Appears after AggregateExec is created + // Arc::new(SimplifyOrdering::new()), ]; Self::with_rules(rules) diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 62588629f0a67..0d2a888ee2d29 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -285,17 +285,7 @@ fn require_top_ordering_helper( let (new_child, is_changed) = require_top_ordering_helper(children.swap_remove(0))?; - let plan = plan.with_new_children(vec![new_child])?; - - // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.rewrite_ordering()?; - // Ok((Arc::new(p) as _, is_changed)) - // } else { - // } - - Ok((plan, is_changed)) - - // Ok((plan.with_new_children(vec![new_child])?, is_changed)) + Ok((plan.with_new_children(vec![new_child])?, is_changed)) } else { // Stop searching, there is no global ordering desired for the query. Ok((plan, false)) 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 7479dec4efbc3..c486d13caa178 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 @@ -139,18 +139,6 @@ fn plan_with_order_preserving_variants( } } - // let mut plan_context = sort_input.update_plan_from_children()?; - // let plan = plan_context.plan.clone(); - // let p = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.rewrite_ordering()?; - // Arc::new(p) - // } else { - // plan - // }; - // plan_context.plan = p; - - // Ok(plan_context) - sort_input.update_plan_from_children() } @@ -197,16 +185,6 @@ fn plan_with_order_breaking_variants( let coalesce = CoalescePartitionsExec::new(child); sort_input.plan = Arc::new(coalesce) as _; } else { - // let mut pc = sort_input.update_plan_from_children()?; - // let plan = pc.plan.clone(); - // let p = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.rewrite_ordering()?; - // Arc::new(p) - // } else { - // plan - // }; - // pc.plan = p; - // return Ok(pc); return sort_input.update_plan_from_children(); } diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index 6d76d4ff4d02a..8375e825c6334 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -95,12 +95,6 @@ fn get_common_requirement_of_aggregate_input( let plan = if is_transformed { - // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.clone_with_input(new_children[0].clone()); - // Arc::new(p) as Arc - // } else { - // plan.with_new_children(new_children)? - // } plan.with_new_children(new_children)? } else { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 1f0919e3ba739..fcc1400c21f2f 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -332,7 +332,7 @@ impl AggregateExec { )?; let schema = Arc::new(schema); - let res = AggregateExec::try_new_with_schema( + AggregateExec::try_new_with_schema( mode, group_by, aggr_expr, @@ -340,65 +340,7 @@ impl AggregateExec { input, input_schema, schema, - )?; - - res.rewrite_ordering() - } - - pub fn rewrite_ordering(&self) -> Result { - let input = self.input(); - let group_by = self.group_by(); - let mut aggr_expr = self.aggr_expr().to_vec(); - let mode = self.mode(); - - let input_eq_properties = input.equivalence_properties(); - // Get GROUP BY expressions: - let groupby_exprs = group_by.input_exprs(); - // If existing ordering satisfies a prefix of the GROUP BY expressions, - // prefix requirements with this section. In this case, aggregation will - // work more efficiently. - let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); - - let mut new_requirement = indices - .iter() - .map(|&idx| PhysicalSortRequirement { - expr: groupby_exprs[idx].clone(), - options: None, - }) - .collect::>(); - - let req = get_aggregate_exprs_requirement( - &new_requirement, - &mut aggr_expr, - &group_by, - input_eq_properties, - &mode, - )?; - new_requirement.extend(req); - new_requirement = collapse_lex_req(new_requirement); - - let required_input_ordering = - (!new_requirement.is_empty()).then_some(new_requirement); - - Ok(Self { - mode: self.mode().clone(), - group_by: self.group_by().clone(), - aggr_expr, - filter_expr: self.filter_expr().to_vec(), - input: self.input().clone(), - schema: self.schema.clone(), - input_schema: self.input_schema.clone(), - metrics: self.metrics.clone(), - required_input_ordering, - limit: self.limit, - input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), - }) - - // self.aggr_expr = aggr_expr; - // self.required_input_ordering = required_input_ordering; - - // Ok(()) + ) } /// Create a new hash aggregate execution plan with the given schema. @@ -473,7 +415,7 @@ impl AggregateExec { &input_order_mode, ); - let required_input_ordering = None; + // let required_input_ordering = None; Ok(AggregateExec { mode, group_by, @@ -1025,15 +967,15 @@ fn get_aggregate_exprs_requirement( )) { first_value = first_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(first_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to LAST_VALUE enables more efficient execution - // given the existing ordering: - let mut last_value = first_value.convert_to_last(); - last_value = last_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(last_value) as _; + // } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + // prefix_requirement, + // &reverse_aggr_req, + // )) { + // // Converting to LAST_VALUE enables more efficient execution + // // given the existing ordering: + // let mut last_value = first_value.convert_to_last(); + // last_value = last_value.with_requirement_satisfied(true); + // *aggr_expr = Arc::new(last_value) as _; } else { // Requirement is not satisfied with existing ordering. first_value = first_value.with_requirement_satisfied(false); @@ -1049,15 +991,15 @@ fn get_aggregate_exprs_requirement( )) { last_value = last_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(last_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to FIRST_VALUE enables more efficient execution - // given the existing ordering: - let mut first_value = last_value.convert_to_first(); - first_value = first_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(first_value) as _; + // } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + // prefix_requirement, + // &reverse_aggr_req, + // )) { + // // Converting to FIRST_VALUE enables more efficient execution + // // given the existing ordering: + // let mut first_value = last_value.convert_to_first(); + // first_value = first_value.with_requirement_satisfied(true); + // *aggr_expr = Arc::new(first_value) as _; } else { // Requirement is not satisfied with existing ordering. last_value = last_value.with_requirement_satisfied(false); diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 3f25855230b0d..8890777b40eb8 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -674,17 +674,7 @@ pub fn with_new_children_if_necessary( .zip(old_children.iter()) .any(|(c1, c2)| !Arc::data_ptr_eq(c1, c2)) { - let plan = plan.with_new_children(children)?; - - // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.rewrite_ordering()?; - // Ok(Arc::new(p)) - // } else { - // Ok(plan) - // } - Ok(plan) - - + plan.with_new_children(children) } else { Ok(plan) } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index f8b8c6b5e1c92..91abd5cfdf284 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -360,16 +360,7 @@ fn reset_plan_states(plan: Arc) -> Result() { - let p = aggr_exec.rewrite_ordering()?; - Ok(Transformed::yes(Arc::new(p))) - } else { - Ok(Transformed::yes(new_plan)) - } - - // Ok(Transformed::yes(new_plan)) + Ok(Transformed::yes(new_plan)) } }) .data() diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index 435aa625710b8..86cf2d0554742 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -17,6 +17,7 @@ //! This module provides common traits for visiting or rewriting tree nodes easily. +use core::panic; use std::fmt::{self, Display, Formatter}; use std::sync::Arc; @@ -65,14 +66,6 @@ impl PlanContext { pub fn update_plan_from_children(mut self) -> Result { let children_plans = self.children.iter().map(|c| c.plan.clone()).collect(); let plan = with_new_children_if_necessary(self.plan, children_plans)?; - - let plan = if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let p = aggr_exec.rewrite_ordering()?; - Arc::new(p) - } else { - plan - }; - self.plan = plan; Ok(self) diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 848ba40963592..b7ad36dace162 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -239,29 +239,20 @@ logical_plan after optimize_projections SAME TEXT AS ABOVE 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 initial_physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] -physical_plan after SimpleOrdering CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan after OutputRequirements OutputRequirementExec --CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceDistribution SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceSorting SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] @@ -303,32 +294,23 @@ EXPLAIN VERBOSE SELECT * FROM alltypes_plain limit 10; initial_physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after OutputRequirements OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] ----ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceDistribution SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceSorting SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -348,34 +330,23 @@ GlobalLimitExec: skip=0, fetch=10 initial_physical_plan_with_stats GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] -physical_plan after SimpleOrdering -GlobalLimitExec: skip=0, fetch=10 ---ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 physical_plan after OutputRequirements OutputRequirementExec --GlobalLimitExec: skip=0, fetch=10 ----ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceDistribution SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after EnforceSorting SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements GlobalLimitExec: skip=0, fetch=10 --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10 diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 748a02aa0afe6..56b17a976dc36 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2677,7 +2677,7 @@ Projection: sales_global.country, ARRAY_AGG(sales_global.amount) ORDER BY [sales ----TableScan: sales_global projection=[country, amount] physical_plan ProjectionExec: expr=[country@0 as country, ARRAY_AGG(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@1 as amounts, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@2 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@3 as fv2] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), LAST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] ----SortExec: expr=[amount@1 DESC] ------MemoryExec: partitions=1, partition_sizes=[1] @@ -2708,7 +2708,7 @@ Projection: sales_global.country, ARRAY_AGG(sales_global.amount) ORDER BY [sales ----TableScan: sales_global projection=[country, amount] physical_plan ProjectionExec: expr=[country@0 as country, ARRAY_AGG(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@1 as amounts, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@2 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@3 as fv2] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] ----SortExec: expr=[amount@1 ASC NULLS LAST] ------MemoryExec: partitions=1, partition_sizes=[1] @@ -2740,7 +2740,7 @@ Projection: sales_global.country, FIRST_VALUE(sales_global.amount) ORDER BY [sal ----TableScan: sales_global projection=[country, amount] physical_plan ProjectionExec: expr=[country@0 as country, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@1 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@2 as fv2, ARRAY_AGG(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@3 as amounts] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount), ARRAY_AGG(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount), ARRAY_AGG(sales_global.amount)] ----SortExec: expr=[amount@1 ASC NULLS LAST] ------MemoryExec: partitions=1, partition_sizes=[1] @@ -2805,7 +2805,7 @@ Projection: sales_global.country, FIRST_VALUE(sales_global.amount) ORDER BY [sal ------TableScan: sales_global projection=[country, ts, amount] physical_plan ProjectionExec: expr=[country@0 as country, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@1 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as lv1, SUM(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@3 as sum1] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[LAST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount), SUM(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount), SUM(sales_global.amount)] ----MemoryExec: partitions=1, partition_sizes=[1] query TRRR rowsort @@ -3157,7 +3157,7 @@ SortPreservingMergeExec: [country@0 ASC NULLS LAST] ------AggregateExec: mode=FinalPartitioned, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] --------CoalesceBatchesExec: target_batch_size=4 ----------RepartitionExec: partitioning=Hash([country@0], 8), input_partitions=8 -------------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), LAST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] +------------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] --------------SortExec: expr=[amount@1 DESC] ----------------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1 ------------------MemoryExec: partitions=1, partition_sizes=[1] @@ -3800,10 +3800,10 @@ Projection: FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_tab ----TableScan: multiple_ordered_table projection=[a, c, d] physical_plan ProjectionExec: expr=[FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_table.a ASC NULLS LAST]@1 as first_a, LAST_VALUE(multiple_ordered_table.c) ORDER BY [multiple_ordered_table.c DESC NULLS FIRST]@2 as last_c] ---AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] +--AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), LAST_VALUE(multiple_ordered_table.c)] ----CoalesceBatchesExec: target_batch_size=2 ------RepartitionExec: partitioning=Hash([d@0], 8), input_partitions=8 ---------AggregateExec: mode=Partial, gby=[d@2 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] +--------AggregateExec: mode=Partial, gby=[d@2 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), LAST_VALUE(multiple_ordered_table.c)] ----------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1 ------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, c, d], output_orderings=[[a@0 ASC NULLS LAST], [c@1 ASC NULLS LAST]], has_header=true @@ -4252,6 +4252,7 @@ SELECT date_bin('15 minutes', ts) as time_chunks + # Since extract is not a monotonic function, below query should not run. # when source is unbounded. query error From 2ab74c94fa7f2d4dd5a0c11a50cad26ca5b06274 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Thu, 11 Apr 2024 09:34:03 +0800 Subject: [PATCH 12/33] find test that do reverse Signed-off-by: jayzhan211 --- .../core/src/physical_optimizer/optimizer.rs | 2 +- .../physical_optimizer/simplify_ordering.rs | 3 ++ .../physical-plan/src/aggregates/mod.rs | 24 ++++++++------ .../sqllogictest/test_files/explain.slt | 3 ++ datafusion/sqllogictest/test_files/test1.slt | 32 +++++++++++++++++++ 5 files changed, 54 insertions(+), 10 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/test1.slt diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 28d7e9b17fed0..636d9a20e7a03 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -143,7 +143,7 @@ impl PhysicalOptimizer { // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), + Arc::new(SimplifyOrdering::new()), ]; Self::with_rules(rules) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index 8375e825c6334..8e07436215e1e 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -176,6 +176,9 @@ fn optimize_internal( input_eq_properties, mode, )?; + + // println!("req: {:?}", req); + new_requirement.extend(req); new_requirement = collapse_lex_req(new_requirement); let required_input_ordering = diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index fcc1400c21f2f..3500ae4c10c55 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -967,15 +967,21 @@ fn get_aggregate_exprs_requirement( )) { first_value = first_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(first_value) as _; - // } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - // prefix_requirement, - // &reverse_aggr_req, - // )) { - // // Converting to LAST_VALUE enables more efficient execution - // // given the existing ordering: - // let mut last_value = first_value.convert_to_last(); - // last_value = last_value.with_requirement_satisfied(true); - // *aggr_expr = Arc::new(last_value) as _; + } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + prefix_requirement, + &reverse_aggr_req, + )) { + // Converting to LAST_VALUE enables more efficient execution + // given the existing ordering: + // let mut last_value = first_value.convert_to_last(); + // last_value = last_value.with_requirement_satisfied(true); + // *aggr_expr = Arc::new(last_value) as _; + + + println!("can convertt {:?}", aggr_expr); + + first_value = first_value.with_requirement_satisfied(false); + *aggr_expr = Arc::new(first_value) as _; } else { // Requirement is not satisfied with existing ordering. first_value = first_value.with_requirement_satisfied(false); diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index b7ad36dace162..57ee8c311f6c6 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -254,6 +254,7 @@ physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPAC physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] @@ -312,6 +313,7 @@ GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Co physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -348,6 +350,7 @@ GlobalLimitExec: skip=0, fetch=10 physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10 --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 diff --git a/datafusion/sqllogictest/test_files/test1.slt b/datafusion/sqllogictest/test_files/test1.slt new file mode 100644 index 0000000000000..2b3775ff3a36c --- /dev/null +++ b/datafusion/sqllogictest/test_files/test1.slt @@ -0,0 +1,32 @@ +# 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. + +statement ok +CREATE EXTERNAL TABLE agg_order ( +c1 INT NOT NULL, +c2 INT NOT NULL, +c3 INT NOT NULL +) +STORED AS CSV +WITH HEADER ROW +WITH ORDER (c1 ASC, c2 ASC) +WITH ORDER (c3 ASC) +LOCATION '../core/tests/data/aggregate_agg_multi_order.csv'; + +# test array_agg with order by multiple columns +query I +select first_value(c1 order by c3 desc) from agg_order; From 03abf1999dbf06d1fc43be6932f65ed28c11fa4e Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Thu, 11 Apr 2024 20:20:02 +0800 Subject: [PATCH 13/33] add test for first and last Signed-off-by: jayzhan211 --- .../core/tests/data/convert_first_last.csv | 11 +++++++++++ datafusion/sqllogictest/test_files/test1.slt | 19 ++++++++++++++----- 2 files changed, 25 insertions(+), 5 deletions(-) create mode 100644 datafusion/core/tests/data/convert_first_last.csv diff --git a/datafusion/core/tests/data/convert_first_last.csv b/datafusion/core/tests/data/convert_first_last.csv new file mode 100644 index 0000000000000..059b631e57118 --- /dev/null +++ b/datafusion/core/tests/data/convert_first_last.csv @@ -0,0 +1,11 @@ +c1,c2,c3 +1,9,0 +2,8,1 +3,7,2 +4,6,3 +5,5,4 +6,4,5 +7,3,6 +8,2,7 +9,1,8 +10,0,9 \ No newline at end of file diff --git a/datafusion/sqllogictest/test_files/test1.slt b/datafusion/sqllogictest/test_files/test1.slt index 2b3775ff3a36c..5fc11514e6710 100644 --- a/datafusion/sqllogictest/test_files/test1.slt +++ b/datafusion/sqllogictest/test_files/test1.slt @@ -16,17 +16,26 @@ # under the License. statement ok -CREATE EXTERNAL TABLE agg_order ( +CREATE EXTERNAL TABLE convert_first_last_table ( c1 INT NOT NULL, c2 INT NOT NULL, c3 INT NOT NULL ) STORED AS CSV WITH HEADER ROW -WITH ORDER (c1 ASC, c2 ASC) +WITH ORDER (c1 ASC) +WITH ORDER (c2 DESC) WITH ORDER (c3 ASC) -LOCATION '../core/tests/data/aggregate_agg_multi_order.csv'; +LOCATION '../core/tests/data/convert_first_last.csv'; -# test array_agg with order by multiple columns +# test first to last query I -select first_value(c1 order by c3 desc) from agg_order; +select first_value(c1 order by c3 desc) from convert_first_last_table; +---- +10 + +# test last to first +query I +select last_value(c1 order by c2 asc) from convert_first_last_table; +---- +1 From 83c721a5218af1f214add92f3dde606a5791797a Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Thu, 11 Apr 2024 22:24:21 +0800 Subject: [PATCH 14/33] pass all test Signed-off-by: jayzhan211 --- .../core/src/physical_optimizer/optimizer.rs | 2 +- .../physical_optimizer/simplify_ordering.rs | 144 +++++++++++++++--- datafusion/core/src/physical_planner.rs | 2 + .../src/equivalence/properties.rs | 3 + .../physical-plan/src/aggregates/mod.rs | 88 +++++------ .../physical-plan/src/coalesce_partitions.rs | 14 +- datafusion/sqllogictest/test_files/test1.slt | 22 ++- 7 files changed, 196 insertions(+), 79 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 636d9a20e7a03..494fd4d490026 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -76,6 +76,7 @@ impl PhysicalOptimizer { /// Create a new optimizer using the recommended list of rules pub fn new() -> Self { let rules: Vec> = vec![ + Arc::new(SimplifyOrdering::new()), // Appears after AggregateExec is created // Arc::new(SimplifyOrdering::new()), // If there is a output requirement of the query, make sure that @@ -143,7 +144,6 @@ impl PhysicalOptimizer { // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), // Appears after AggregateExec is created - Arc::new(SimplifyOrdering::new()), ]; Self::with_rules(rules) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index 8e07436215e1e..041399102eb22 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -15,9 +15,8 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; - use arrow::compute::kernels::aggregate; +use datafusion_common::Result; use datafusion_common::{ config::ConfigOptions, not_impl_err, @@ -25,6 +24,7 @@ use datafusion_common::{ }; use datafusion_physical_expr::{ aggregate::is_order_sensitive, + equivalence::ProjectionMapping, expressions::{FirstValue, LastValue}, physical_exprs_contains, reverse_order_bys, AggregateExpr, EquivalenceProperties, LexOrdering, LexRequirement, PhysicalSortRequirement, @@ -32,14 +32,13 @@ use datafusion_physical_expr::{ use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, coalesce_partitions::CoalescePartitionsExec, - ExecutionPlan, ExecutionPlanProperties, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, }; +use std::sync::Arc; use datafusion_physical_expr::equivalence::collapse_lex_req; use datafusion_physical_plan::windows::get_ordered_partition_by_indices; -use crate::error::Result; - use super::{output_requirements::OutputRequirementExec, PhysicalOptimizerRule}; #[derive(Default)] @@ -80,12 +79,12 @@ fn get_common_requirement_of_aggregate_input( ) -> Result>> { let children = plan.children(); - let mut is_transformed = false; + let mut is_child_transformed = false; let mut new_children: Vec> = vec![]; for c in children.iter() { let res = get_common_requirement_of_aggregate_input(c.clone())?; if res.transformed { - is_transformed = true; + is_child_transformed = true; } new_children.push(res.data); } @@ -93,10 +92,12 @@ fn get_common_requirement_of_aggregate_input( let mode1 = plan.properties().execution_mode(); // println!("mode 1: {:?}", mode); - let plan = if is_transformed { + println!("plan name: {:?}", plan.name()); + println!("is_child_transformed: {:?}", is_child_transformed); + println!("new_children: {:?}", new_children); + let plan = if is_child_transformed { plan.with_new_children(new_children)? - } else { plan }; @@ -107,11 +108,35 @@ fn get_common_requirement_of_aggregate_input( } let plan = optimize_internal(plan)?; + + if plan.transformed { + println!("transformed plan: {:?}", plan.data.name()); + } else { + println!("not transformed plan: {:?}", plan.data.name()); + } + let mode3 = plan.data.properties().execution_mode(); if mode1 != mode3 { println!("mode1: {:?}, mode3: {:?}", mode1, mode3); } - Ok(plan) + + if !plan.transformed { + let name = plan.data.name(); + println!( + "not transformed plan: {:?} and is_child_transformed: {:?}", + name, is_child_transformed + ); + } + + // If one of the children is transformed, then the plan is considered transformed, then we update + // the children of the plan from bottom to top. + if plan.transformed || is_child_transformed { + Ok(Transformed::yes(plan.data)) + } else { + Ok(Transformed::no(plan.data)) + } + + // Ok(plan) // if let Some(c) = new_c { // if !c.transformed { @@ -142,25 +167,52 @@ fn get_common_requirement_of_aggregate_input( // return Ok(plan); } +fn try_get_updated_aggr_expr_from_child( + aggr_exec: &AggregateExec, +) -> Vec> { + let input = aggr_exec.input(); + if aggr_exec.mode() != &AggregateMode::Partial { + // Some aggregators may be modified during initialization for + // optimization purposes. For example, a FIRST_VALUE may turn + // into a LAST_VALUE with the reverse ordering requirement. + // To reflect such changes to subsequent stages, use the updated + // `AggregateExpr`/`PhysicalSortExpr` objects. + // + // The bottom up transformation is the mirror of LogicalPlan::Aggregate creation in [create_initial_plan] + + if let Some(c_aggr_exec) = input.as_any().downcast_ref::() { + if c_aggr_exec.mode() == &AggregateMode::Partial { + // If the input is an AggregateExec in Partial mode, then the + // input is a CoalescePartitionsExec. In this case, the + // AggregateExec is the second stage of aggregation. The + // requirements of the second stage are the requirements of + // the first stage. + return c_aggr_exec.aggr_expr().to_vec(); + } + } + } + + aggr_exec.aggr_expr().to_vec() +} + fn optimize_internal( plan: Arc, ) -> Result>> { if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // if aggr_exec.mode() != &AggregateMode::Partial { - // return Ok(Transformed::no(plan)); - // } + println!("mode: {:?}", aggr_exec.mode()); - let input = aggr_exec.input().clone(); - let mut aggr_expr = aggr_exec.aggr_expr().to_vec(); + let input = aggr_exec.input(); + let mut aggr_expr = try_get_updated_aggr_expr_from_child(aggr_exec); let group_by = aggr_exec.group_by(); let mode = aggr_exec.mode(); let input_eq_properties = input.equivalence_properties(); + // println!("input_eq_properties: {:?}", input_eq_properties); let groupby_exprs = group_by.input_exprs(); // If existing ordering satisfies a prefix of the GROUP BY expressions, // prefix requirements with this section. In this case, aggregation will // work more efficiently. - let indices = get_ordered_partition_by_indices(&groupby_exprs, &input); + let indices = get_ordered_partition_by_indices(&groupby_exprs, input); let mut new_requirement = indices .iter() .map(|&idx| PhysicalSortRequirement { @@ -169,6 +221,7 @@ fn optimize_internal( }) .collect::>(); + println!("1 aggr_expr: {:?}", aggr_expr); let req = get_aggregate_exprs_requirement( &new_requirement, &mut aggr_expr, @@ -176,6 +229,7 @@ fn optimize_internal( input_eq_properties, mode, )?; + println!("2 aggr_expr: {:?}", aggr_expr); // println!("req: {:?}", req); @@ -184,7 +238,48 @@ fn optimize_internal( let required_input_ordering = (!new_requirement.is_empty()).then_some(new_requirement); - let p = aggr_exec.clone_with_required_input_ordering(required_input_ordering); + println!("required_input_ordering: {:?}", required_input_ordering); + println!("agg_expr: {:?}", aggr_expr); + + let input_order_mode = + if indices.len() == groupby_exprs.len() && !indices.is_empty() { + InputOrderMode::Sorted + } else if !indices.is_empty() { + InputOrderMode::PartiallySorted(indices) + } else { + InputOrderMode::Linear + }; + let projection_mapping = + ProjectionMapping::try_new(&group_by.expr(), &input.schema())?; + + let cache = AggregateExec::compute_properties( + &input, + plan.schema().clone(), + &projection_mapping, + &mode, + &input_order_mode, + ); + + // let p = AggregateExec { + // mode: mode.clone(), + // group_by: group_by.clone(), + // aggr_expr, + // filter_expr: aggr_exec.filter_expr().to_vec(), + // input: input.clone(), + // schema: aggr_exec.schema(), + // input_schema: aggr_exec.input_schema(), + // metrics: aggr_exec.metrics().clone(), + // required_input_ordering, + // limit: None, + // input_order_mode, + // cache, + // }; + let p = aggr_exec.clone_with_required_input_ordering_and_aggr_expr( + required_input_ordering, + aggr_expr, + cache, + input_order_mode, + ); let res = Arc::new(p) as Arc; Ok(Transformed::yes(res)) @@ -306,16 +401,23 @@ fn get_aggregate_exprs_requirement( if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { let mut first_value = first_value.clone(); + + // println!("prefix_requirement: {:?}", prefix_requirement); + // println!("aggr_req: {:?}", aggr_req); + // println!("reverse_aggr_req: {:?}", reverse_aggr_req); + if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, &aggr_req, )) { + println!("1st step"); first_value = first_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(first_value) as _; } else if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, &reverse_aggr_req, )) { + println!("2nd step"); // Converting to LAST_VALUE enables more efficient execution // given the existing ordering: let mut last_value = first_value.convert_to_last(); @@ -411,11 +513,13 @@ fn get_aggregate_exprs_requirement( Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) } - mod tests { use super::*; use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions}; - use datafusion_physical_expr::{expressions::{col, OrderSensitiveArrayAgg}, PhysicalSortExpr}; + use datafusion_physical_expr::{ + expressions::{col, OrderSensitiveArrayAgg}, + PhysicalSortExpr, + }; fn create_test_schema() -> Result { let a = Field::new("a", DataType::Int32, true); @@ -511,4 +615,4 @@ mod tests { assert_eq!(res, common_requirement); Ok(()) } -} \ No newline at end of file +} diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 798dfa2d5bfbc..40a1a60d6865f 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -465,6 +465,8 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { .create_initial_plan(logical_plan, session_state) .await?; + // println!("first plan: {:#?}", plan); + self.optimize_internal(plan, session_state, |_, _| {}) } } diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs index c14c88d6c69bf..1f434d0149a9d 100644 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ b/datafusion/physical-expr/src/equivalence/properties.rs @@ -287,9 +287,11 @@ impl EquivalenceProperties { let mut eq_properties = self.clone(); // First, standardize the given requirement: let normalized_reqs = eq_properties.normalize_sort_requirements(reqs); + // println!("normalized_reqs: {:?}", normalized_reqs); for normalized_req in normalized_reqs { // Check whether given ordering is satisfied if !eq_properties.ordering_satisfy_single(&normalized_req) { + // println!("not single"); return false; } // Treat satisfied keys as constants in subsequent iterations. We @@ -325,6 +327,7 @@ impl EquivalenceProperties { fn ordering_satisfy_single(&self, req: &PhysicalSortRequirement) -> bool { let expr_ordering = self.get_expr_ordering(req.expr.clone()); let ExprOrdering { expr, data, .. } = expr_ordering; + // println!("expr: {:?}, data: {:?}", expr, data); match data { SortProperties::Ordered(options) => { let sort_expr = PhysicalSortExpr { expr, options }; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3500ae4c10c55..d089e89c3cb1d 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -42,12 +42,15 @@ use datafusion_expr::Accumulator; use datafusion_physical_expr::aggregate::is_order_sensitive; use datafusion_physical_expr::equivalence::collapse_lex_req; use datafusion_physical_expr::expressions::{FirstValue, LastValue}; -use datafusion_physical_expr::{physical_exprs_contains, reverse_order_bys, EquivalenceProperties, LexOrdering, PhysicalSortRequirement}; use datafusion_physical_expr::{ equivalence::ProjectionMapping, expressions::{Column, Max, Min, UnKnownColumn}, AggregateExpr, LexRequirement, PhysicalExpr, }; +use datafusion_physical_expr::{ + physical_exprs_contains, reverse_order_bys, EquivalenceProperties, LexOrdering, + PhysicalSortRequirement, +}; use itertools::Itertools; @@ -271,42 +274,45 @@ pub struct AggregateExec { } impl AggregateExec { - pub fn clone_with_input(&self, input: Arc) -> Self { - Self { - input, - // clone the rest of the fields - mode: self.mode, - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), - limit: self.limit, - schema: self.schema.clone(), - input_schema: self.input_schema.clone(), - metrics: self.metrics.clone(), - input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), - required_input_ordering: self.required_input_ordering.clone(), - } - } - - pub fn clone_with_required_input_ordering( + // pub fn clone_with_input(&self, input: Arc) -> Self { + // Self { + // input, + // // clone the rest of the fields + // mode: self.mode, + // group_by: self.group_by.clone(), + // aggr_expr: self.aggr_expr.clone(), + // filter_expr: self.filter_expr.clone(), + // limit: self.limit, + // schema: self.schema.clone(), + // input_schema: self.input_schema.clone(), + // metrics: self.metrics.clone(), + // input_order_mode: self.input_order_mode.clone(), + // cache: self.cache.clone(), + // required_input_ordering: self.required_input_ordering.clone(), + // } + // } + + pub fn clone_with_required_input_ordering_and_aggr_expr( &self, required_input_ordering: Option, + aggr_expr: Vec>, + cache: PlanProperties, + input_order_mode: InputOrderMode, ) -> Self { Self { required_input_ordering, // clone the rest of the fields mode: self.mode, group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), + aggr_expr, filter_expr: self.filter_expr.clone(), limit: self.limit, input: self.input.clone(), schema: self.schema.clone(), input_schema: self.input_schema.clone(), - metrics: self.metrics.clone(), - input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), + metrics: ExecutionPlanMetricsSet::new(), + input_order_mode, + cache, } } @@ -432,7 +438,6 @@ impl AggregateExec { }) } - /// Aggregation mode (full, partial) pub fn mode(&self) -> &AggregateMode { &self.mode @@ -555,7 +560,7 @@ impl AggregateExec { } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties( + pub fn compute_properties( input: &Arc, schema: SchemaRef, projection_mapping: &ProjectionMapping, @@ -960,6 +965,9 @@ fn get_aggregate_exprs_requirement( PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { + println!("prefix_requirement {:?}", prefix_requirement); + println!("reverse_aggr_req {:?}", reverse_aggr_req); + let mut first_value = first_value.clone(); if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, @@ -967,21 +975,15 @@ fn get_aggregate_exprs_requirement( )) { first_value = first_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(first_value) as _; - } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &reverse_aggr_req, - )) { - // Converting to LAST_VALUE enables more efficient execution - // given the existing ordering: - // let mut last_value = first_value.convert_to_last(); - // last_value = last_value.with_requirement_satisfied(true); - // *aggr_expr = Arc::new(last_value) as _; - - - println!("can convertt {:?}", aggr_expr); - - first_value = first_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(first_value) as _; + // } else if eq_properties.ordering_satisfy_requirement(&concat_slices( + // prefix_requirement, + // &reverse_aggr_req, + // )) { + // // Converting to LAST_VALUE enables more efficient execution + // // given the existing ordering: + // let mut last_value = first_value.convert_to_last(); + // last_value = last_value.with_requirement_satisfied(true); + // *aggr_expr = Arc::new(last_value) as _; } else { // Requirement is not satisfied with existing ordering. first_value = first_value.with_requirement_satisfied(false); @@ -1001,6 +1003,7 @@ fn get_aggregate_exprs_requirement( // prefix_requirement, // &reverse_aggr_req, // )) { + // println!("last to first {:?}", aggr_expr); // // Converting to FIRST_VALUE enables more efficient execution // // given the existing ordering: // let mut first_value = last_value.convert_to_first(); @@ -1294,7 +1297,7 @@ mod tests { use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_physical_expr::expressions::{ - lit, ApproxDistinct, Count, FirstValue, LastValue, Median, OrderSensitiveArrayAgg + lit, ApproxDistinct, Count, FirstValue, LastValue, Median, OrderSensitiveArrayAgg, }; use datafusion_physical_expr::{ reverse_order_bys, AggregateExpr, EquivalenceProperties, PhysicalExpr, @@ -2146,7 +2149,6 @@ mod tests { Ok(()) } - #[test] fn test_agg_exec_same_schema() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index b95d17b627b96..67ceb07aeeb6d 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -73,13 +73,13 @@ impl CoalescePartitionsExec { ) } - pub fn clone_with_input(&self, input: Arc) -> Self { - CoalescePartitionsExec { - input, - metrics: self.metrics.clone(), - cache: self.cache.clone(), - } - } + // pub fn clone_with_input(&self, input: Arc) -> Self { + // CoalescePartitionsExec { + // input, + // metrics: self.metrics.clone(), + // cache: self.cache.clone(), + // } + // } } impl DisplayAs for CoalescePartitionsExec { diff --git a/datafusion/sqllogictest/test_files/test1.slt b/datafusion/sqllogictest/test_files/test1.slt index 5fc11514e6710..0c903fdeb955b 100644 --- a/datafusion/sqllogictest/test_files/test1.slt +++ b/datafusion/sqllogictest/test_files/test1.slt @@ -28,14 +28,20 @@ WITH ORDER (c2 DESC) WITH ORDER (c3 ASC) LOCATION '../core/tests/data/convert_first_last.csv'; -# test first to last -query I -select first_value(c1 order by c3 desc) from convert_first_last_table; +# test first to last, the result does not show difference, we need to check the conversion by `explain` +query TT +explain select first_value(c1 order by c3 desc) from convert_first_last_table; ---- -10 +logical_plan +Aggregate: groupBy=[[]], aggr=[[FIRST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c3 DESC NULLS FIRST]]] +--TableScan: convert_first_last_table projection=[c1, c3] +physical_plan +AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +--CoalescePartitionsExec +----AggregateExec: mode=Partial, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c3], output_orderings=[[c1@0 ASC NULLS LAST], [c3@1 ASC NULLS LAST]], has_header=true # test last to first -query I -select last_value(c1 order by c2 asc) from convert_first_last_table; ----- -1 +# query TT +# explain select last_value(c1 order by c2 asc) from convert_first_last_table; \ No newline at end of file From 34b8f01f057be7baf6dbea70b44ba5987fe91892 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Thu, 11 Apr 2024 22:27:15 +0800 Subject: [PATCH 15/33] upd test Signed-off-by: jayzhan211 --- .../core/src/physical_optimizer/simplify_ordering.rs | 2 +- datafusion/sqllogictest/test_files/explain.slt | 8 +++++--- datafusion/sqllogictest/test_files/group_by.slt | 6 +++--- datafusion/sqllogictest/test_files/test1.slt | 2 +- 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index 041399102eb22..d62621c7af534 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -171,7 +171,7 @@ fn try_get_updated_aggr_expr_from_child( aggr_exec: &AggregateExec, ) -> Vec> { let input = aggr_exec.input(); - if aggr_exec.mode() != &AggregateMode::Partial { + if aggr_exec.mode() == &AggregateMode::Final || aggr_exec.mode() == &AggregateMode::FinalPartitioned { // Some aggregators may be modified during initialization for // optimization purposes. For example, a FIRST_VALUE may turn // into a LAST_VALUE with the reverse ordering requirement. diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 57ee8c311f6c6..1e15f401be4dc 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -239,6 +239,7 @@ logical_plan after optimize_projections SAME TEXT AS ABOVE 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 initial_physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] +physical_plan after SimpleOrdering CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan after OutputRequirements OutputRequirementExec --CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true @@ -254,7 +255,6 @@ physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPAC physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] @@ -295,6 +295,7 @@ EXPLAIN VERBOSE SELECT * FROM alltypes_plain limit 10; initial_physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after OutputRequirements OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -313,7 +314,6 @@ GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Co physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -332,6 +332,9 @@ GlobalLimitExec: skip=0, fetch=10 initial_physical_plan_with_stats GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after SimpleOrdering +GlobalLimitExec: skip=0, fetch=10 +--ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 physical_plan after OutputRequirements OutputRequirementExec --GlobalLimitExec: skip=0, fetch=10 @@ -350,7 +353,6 @@ GlobalLimitExec: skip=0, fetch=10 physical_plan after PipelineChecker SAME TEXT AS ABOVE physical_plan after LimitAggregation SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan GlobalLimitExec: skip=0, fetch=10 --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 56b17a976dc36..8945f3b01c283 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2805,7 +2805,7 @@ Projection: sales_global.country, FIRST_VALUE(sales_global.amount) ORDER BY [sal ------TableScan: sales_global projection=[country, ts, amount] physical_plan ProjectionExec: expr=[country@0 as country, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@1 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as lv1, SUM(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@3 as sum1] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount), SUM(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[LAST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount), SUM(sales_global.amount)] ----MemoryExec: partitions=1, partition_sizes=[1] query TRRR rowsort @@ -3800,10 +3800,10 @@ Projection: FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_tab ----TableScan: multiple_ordered_table projection=[a, c, d] physical_plan ProjectionExec: expr=[FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_table.a ASC NULLS LAST]@1 as first_a, LAST_VALUE(multiple_ordered_table.c) ORDER BY [multiple_ordered_table.c DESC NULLS FIRST]@2 as last_c] ---AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), LAST_VALUE(multiple_ordered_table.c)] +--AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] ----CoalesceBatchesExec: target_batch_size=2 ------RepartitionExec: partitioning=Hash([d@0], 8), input_partitions=8 ---------AggregateExec: mode=Partial, gby=[d@2 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), LAST_VALUE(multiple_ordered_table.c)] +--------AggregateExec: mode=Partial, gby=[d@2 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] ----------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1 ------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, c, d], output_orderings=[[a@0 ASC NULLS LAST], [c@1 ASC NULLS LAST]], has_header=true diff --git a/datafusion/sqllogictest/test_files/test1.slt b/datafusion/sqllogictest/test_files/test1.slt index 0c903fdeb955b..4ca847fa01d8e 100644 --- a/datafusion/sqllogictest/test_files/test1.slt +++ b/datafusion/sqllogictest/test_files/test1.slt @@ -44,4 +44,4 @@ AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1) # test last to first # query TT -# explain select last_value(c1 order by c2 asc) from convert_first_last_table; \ No newline at end of file +# explain select last_value(c1 order by c2 asc) from convert_first_last_table; From b17e8f7bb0c839f6139fdee07e5dd6871d503159 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 08:57:15 +0800 Subject: [PATCH 16/33] upd test Signed-off-by: jayzhan211 --- .../enforce_distribution.rs | 5 --- .../core/src/physical_optimizer/optimizer.rs | 1 + .../physical_optimizer/output_requirements.rs | 9 ---- .../replace_with_order_preserving_variants.rs | 1 - .../physical_optimizer/simplify_ordering.rs | 12 +++--- datafusion/core/src/physical_planner.rs | 2 - .../src/equivalence/properties.rs | 2 - .../physical-plan/src/aggregates/mod.rs | 41 +++++++++++-------- .../sqllogictest/test_files/explain.slt | 3 ++ .../sqllogictest/test_files/group_by.slt | 8 ++-- 10 files changed, 37 insertions(+), 47 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_distribution.rs b/datafusion/core/src/physical_optimizer/enforce_distribution.rs index fe43890bc8501..145f08af76dd1 100644 --- a/datafusion/core/src/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/src/physical_optimizer/enforce_distribution.rs @@ -1080,11 +1080,6 @@ fn ensure_distribution( } }; - // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.rewrite_ordering()?; - // plan = Arc::new(p); - // } - // 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 diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 494fd4d490026..e0263ced38909 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -114,6 +114,7 @@ impl PhysicalOptimizer { // Note that one should always run this rule after running the EnforceDistribution rule // as the latter may break local sorting requirements. Arc::new(EnforceSorting::new()), + Arc::new(SimplifyOrdering::new()), // Arc::new(SimplifyOrdering::new()), // TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future. Arc::new(ProjectionPushdown::new()), diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 0d2a888ee2d29..412976ca9e43e 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -95,15 +95,6 @@ pub(crate) struct OutputRequirementExec { } impl OutputRequirementExec { - pub fn clone_with_input(&self, input: Arc) -> Self { - Self { - input, - order_requirement: self.order_requirement.clone(), - dist_requirement: self.dist_requirement.clone(), - cache: self.cache.clone(), - } - } - pub(crate) fn new( input: Arc, requirements: Option, 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 c486d13caa178..ad19215fbf67e 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 @@ -29,7 +29,6 @@ use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::Transformed; -use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::ExecutionPlanProperties; diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index d62621c7af534..f70bae8834099 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -199,7 +199,6 @@ fn optimize_internal( plan: Arc, ) -> Result>> { if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - println!("mode: {:?}", aggr_exec.mode()); let input = aggr_exec.input(); let mut aggr_expr = try_get_updated_aggr_expr_from_child(aggr_exec); @@ -207,7 +206,6 @@ fn optimize_internal( let mode = aggr_exec.mode(); let input_eq_properties = input.equivalence_properties(); - // println!("input_eq_properties: {:?}", input_eq_properties); let groupby_exprs = group_by.input_exprs(); // If existing ordering satisfies a prefix of the GROUP BY expressions, // prefix requirements with this section. In this case, aggregation will @@ -221,7 +219,6 @@ fn optimize_internal( }) .collect::>(); - println!("1 aggr_expr: {:?}", aggr_expr); let req = get_aggregate_exprs_requirement( &new_requirement, &mut aggr_expr, @@ -229,9 +226,6 @@ fn optimize_internal( input_eq_properties, mode, )?; - println!("2 aggr_expr: {:?}", aggr_expr); - - // println!("req: {:?}", req); new_requirement.extend(req); new_requirement = collapse_lex_req(new_requirement); @@ -281,6 +275,12 @@ fn optimize_internal( input_order_mode, ); + // let aggr_exec = aggr_exec.to_owned().to_owned(); + // aggr_exec.with_aggr_expr(aggr_expr); + // aggr_exec.with_required_input_ordering(required_input_ordering); + // aggr_exec.with_cache(cache); + // aggr_exec.with_input_order_mode(input_order_mode); + let res = Arc::new(p) as Arc; Ok(Transformed::yes(res)) } else { diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 40a1a60d6865f..798dfa2d5bfbc 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -465,8 +465,6 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { .create_initial_plan(logical_plan, session_state) .await?; - // println!("first plan: {:#?}", plan); - self.optimize_internal(plan, session_state, |_, _| {}) } } diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs index 1f434d0149a9d..367cc35435771 100644 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ b/datafusion/physical-expr/src/equivalence/properties.rs @@ -287,11 +287,9 @@ impl EquivalenceProperties { let mut eq_properties = self.clone(); // First, standardize the given requirement: let normalized_reqs = eq_properties.normalize_sort_requirements(reqs); - // println!("normalized_reqs: {:?}", normalized_reqs); for normalized_req in normalized_reqs { // Check whether given ordering is satisfied if !eq_properties.ordering_satisfy_single(&normalized_req) { - // println!("not single"); return false; } // Treat satisfied keys as constants in subsequent iterations. We diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index d089e89c3cb1d..be4ba9fe8a164 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -274,24 +274,6 @@ pub struct AggregateExec { } impl AggregateExec { - // pub fn clone_with_input(&self, input: Arc) -> Self { - // Self { - // input, - // // clone the rest of the fields - // mode: self.mode, - // group_by: self.group_by.clone(), - // aggr_expr: self.aggr_expr.clone(), - // filter_expr: self.filter_expr.clone(), - // limit: self.limit, - // schema: self.schema.clone(), - // input_schema: self.input_schema.clone(), - // metrics: self.metrics.clone(), - // input_order_mode: self.input_order_mode.clone(), - // cache: self.cache.clone(), - // required_input_ordering: self.required_input_ordering.clone(), - // } - // } - pub fn clone_with_required_input_ordering_and_aggr_expr( &self, required_input_ordering: Option, @@ -438,6 +420,29 @@ impl AggregateExec { }) } + pub fn with_cache(mut self, cache: PlanProperties) -> Self { + self.cache = cache; + self + } + + pub fn with_input_order_mode(mut self, input_order_mode: InputOrderMode) -> Self { + self.input_order_mode = input_order_mode; + self + } + + pub fn with_required_input_ordering( + mut self, + required_input_ordering: Option, + ) -> Self { + self.required_input_ordering = required_input_ordering; + self + } + + pub fn with_aggr_expr(mut self, aggr_expr: Vec>) -> Self { + self.aggr_expr = aggr_expr; + self + } + /// Aggregation mode (full, partial) pub fn mode(&self) -> &AggregateMode { &self.mode diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 1e15f401be4dc..2c82f2e6a68d6 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -249,6 +249,7 @@ physical_plan after LimitedDistinctAggregation 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 +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true @@ -306,6 +307,7 @@ physical_plan after LimitedDistinctAggregation 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 +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements @@ -345,6 +347,7 @@ physical_plan after LimitedDistinctAggregation 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 +physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 8945f3b01c283..88e0272ef9bca 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2677,7 +2677,7 @@ Projection: sales_global.country, ARRAY_AGG(sales_global.amount) ORDER BY [sales ----TableScan: sales_global projection=[country, amount] physical_plan ProjectionExec: expr=[country@0 as country, ARRAY_AGG(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@1 as amounts, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@2 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@3 as fv2] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), LAST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] ----SortExec: expr=[amount@1 DESC] ------MemoryExec: partitions=1, partition_sizes=[1] @@ -2708,7 +2708,7 @@ Projection: sales_global.country, ARRAY_AGG(sales_global.amount) ORDER BY [sales ----TableScan: sales_global projection=[country, amount] physical_plan ProjectionExec: expr=[country@0 as country, ARRAY_AGG(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@1 as amounts, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@2 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@3 as fv2] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount)] ----SortExec: expr=[amount@1 ASC NULLS LAST] ------MemoryExec: partitions=1, partition_sizes=[1] @@ -2740,7 +2740,7 @@ Projection: sales_global.country, FIRST_VALUE(sales_global.amount) ORDER BY [sal ----TableScan: sales_global projection=[country, amount] physical_plan ProjectionExec: expr=[country@0 as country, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@1 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@2 as fv2, ARRAY_AGG(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@3 as amounts] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount), ARRAY_AGG(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount), ARRAY_AGG(sales_global.amount)] ----SortExec: expr=[amount@1 ASC NULLS LAST] ------MemoryExec: partitions=1, partition_sizes=[1] @@ -3157,7 +3157,7 @@ SortPreservingMergeExec: [country@0 ASC NULLS LAST] ------AggregateExec: mode=FinalPartitioned, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] --------CoalesceBatchesExec: target_batch_size=4 ----------RepartitionExec: partitioning=Hash([country@0], 8), input_partitions=8 -------------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] +------------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[ARRAY_AGG(sales_global.amount), LAST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount)] --------------SortExec: expr=[amount@1 DESC] ----------------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1 ------------------MemoryExec: partitions=1, partition_sizes=[1] From a313af086357e0edc4ae83c0e9fb4039ea6c9ae4 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 09:28:54 +0800 Subject: [PATCH 17/33] cleanup Signed-off-by: jayzhan211 --- .../core/src/physical_optimizer/optimizer.rs | 24 ++--- .../physical_optimizer/output_requirements.rs | 1 - .../physical_optimizer/simplify_ordering.rs | 88 ++----------------- datafusion/core/src/physical_planner.rs | 2 + .../src/equivalence/properties.rs | 1 - .../physical-plan/src/aggregates/mod.rs | 15 ++-- .../sqllogictest/test_files/explain.slt | 8 +- 7 files changed, 24 insertions(+), 115 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index e0263ced38909..134d11e43a0e5 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use super::projection_pushdown::ProjectionPushdown; -use super::simplify_ordering::SimplifyOrdering; +use super::simplify_ordering::ConvertFirstLast; use crate::config::ConfigOptions; use crate::physical_optimizer::aggregate_statistics::AggregateStatistics; use crate::physical_optimizer::coalesce_batches::CoalesceBatches; @@ -76,16 +76,10 @@ impl PhysicalOptimizer { /// Create a new optimizer using the recommended list of rules pub fn new() -> Self { let rules: Vec> = vec![ - Arc::new(SimplifyOrdering::new()), - // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), - // Arc::new(SimplifyOrdering::new()), Arc::new(AggregateStatistics::new()), - // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the // EnforceDistribution and EnforceSorting rules as they decide whether to add additional @@ -96,26 +90,22 @@ impl PhysicalOptimizer { // as that rule may inject other operations in between the different AggregateExecs. // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), - // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), + // Run once before PartialFinalAggregation is rewritten to ensure the rule is applied correctly + Arc::new(ConvertFirstLast::new()), // 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()), - // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule Arc::new(CombinePartialFinalAggregate::new()), - // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), // The EnforceSorting rule is for adding essential local sorting to satisfy the required // ordering. Please make sure that the whole plan tree is determined before this rule. // Note that one should always run this rule after running the EnforceDistribution rule // as the latter may break local sorting requirements. Arc::new(EnforceSorting::new()), - Arc::new(SimplifyOrdering::new()), - // Arc::new(SimplifyOrdering::new()), + // Run once after the local sorting requirement is changed + Arc::new(ConvertFirstLast::new()), // TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future. Arc::new(ProjectionPushdown::new()), // The CoalesceBatches rule will not influence the distribution and ordering of the @@ -124,8 +114,6 @@ impl PhysicalOptimizer { // Remove the ancillary output requirement operator since we are done with the planning // phase. Arc::new(OutputRequirements::new_remove_mode()), - // Appears after AggregateExec is created - // Arc::new(SimplifyOrdering::new()), // The PipelineChecker rule will reject non-runnable query plans that use // pipeline-breaking operators on infinite input(s). The rule generates a // diagnostic error message when this happens. It makes no changes to the @@ -136,7 +124,6 @@ impl PhysicalOptimizer { // into an `order by max(x) limit y`. In this case it will copy the limit value down // to the aggregation, allowing it to use only y number of accumulators. Arc::new(TopKAggregation::new()), - // Arc::new(SimplifyOrdering::new()), // The ProjectionPushdown rule tries to push projections towards // the sources in the execution plan. As a result of this process, // a projection can disappear if it reaches the source providers, and @@ -144,7 +131,6 @@ impl PhysicalOptimizer { // are not present, the load of executors such as join or union will be // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), - // Appears after AggregateExec is created ]; Self::with_rules(rules) diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 412976ca9e43e..9f50309985856 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -32,7 +32,6 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Result, Statistics}; use datafusion_physical_expr::{Distribution, LexRequirement, PhysicalSortRequirement}; -use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ExecutionPlanProperties, PlanProperties}; diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index f70bae8834099..adf4b59b74a20 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -15,23 +15,21 @@ // specific language governing permissions and limitations // under the License. -use arrow::compute::kernels::aggregate; use datafusion_common::Result; use datafusion_common::{ config::ConfigOptions, not_impl_err, - tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, + tree_node::{Transformed, TransformedResult, TreeNode}, }; +use datafusion_physical_expr::expressions::{FirstValue, LastValue}; use datafusion_physical_expr::{ aggregate::is_order_sensitive, equivalence::ProjectionMapping, - expressions::{FirstValue, LastValue}, physical_exprs_contains, reverse_order_bys, AggregateExpr, EquivalenceProperties, LexOrdering, LexRequirement, PhysicalSortRequirement, }; use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, - coalesce_partitions::CoalescePartitionsExec, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, }; use std::sync::Arc; @@ -39,29 +37,28 @@ use std::sync::Arc; use datafusion_physical_expr::equivalence::collapse_lex_req; use datafusion_physical_plan::windows::get_ordered_partition_by_indices; -use super::{output_requirements::OutputRequirementExec, PhysicalOptimizerRule}; +use super::PhysicalOptimizerRule; #[derive(Default)] -pub struct SimplifyOrdering {} +pub struct ConvertFirstLast {} -impl SimplifyOrdering { +impl ConvertFirstLast { pub fn new() -> Self { Self::default() } } -impl PhysicalOptimizerRule for SimplifyOrdering { +impl PhysicalOptimizerRule for ConvertFirstLast { fn optimize( &self, plan: Arc, _config: &ConfigOptions, ) -> Result> { - // Ok(plan) + let res = plan .transform_down(&get_common_requirement_of_aggregate_input) .data(); - // println!("res: {:?}", res); res } @@ -89,45 +86,14 @@ fn get_common_requirement_of_aggregate_input( new_children.push(res.data); } - let mode1 = plan.properties().execution_mode(); - // println!("mode 1: {:?}", mode); - - println!("plan name: {:?}", plan.name()); - println!("is_child_transformed: {:?}", is_child_transformed); - println!("new_children: {:?}", new_children); - let plan = if is_child_transformed { plan.with_new_children(new_children)? } else { plan }; - let mode2 = plan.properties().execution_mode(); - if mode1 != mode2 { - println!("mode1: {:?}, mode2: {:?}", mode1, mode2); - } - let plan = optimize_internal(plan)?; - if plan.transformed { - println!("transformed plan: {:?}", plan.data.name()); - } else { - println!("not transformed plan: {:?}", plan.data.name()); - } - - let mode3 = plan.data.properties().execution_mode(); - if mode1 != mode3 { - println!("mode1: {:?}, mode3: {:?}", mode1, mode3); - } - - if !plan.transformed { - let name = plan.data.name(); - println!( - "not transformed plan: {:?} and is_child_transformed: {:?}", - name, is_child_transformed - ); - } - // If one of the children is transformed, then the plan is considered transformed, then we update // the children of the plan from bottom to top. if plan.transformed || is_child_transformed { @@ -135,36 +101,6 @@ fn get_common_requirement_of_aggregate_input( } else { Ok(Transformed::no(plan.data)) } - - // Ok(plan) - - // if let Some(c) = new_c { - // if !c.transformed { - // return Ok(plan); - // } - - // let plan = plan.data; - - // // TODO: support more types of ExecutionPlan - // if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - // let p = aggr_exec.clone_with_input(c.data); - // return Ok(Transformed::yes(Arc::new(p) as Arc)); - // } else if let Some(coalesce_exec) = - // plan.as_any().downcast_ref::() - // { - // let p = coalesce_exec.clone_with_input(c.data); - // return Ok(Transformed::yes(Arc::new(p) as Arc)); - // } else if let Some(out_req_exec) = - // plan.as_any().downcast_ref::() - // { - // let p = out_req_exec.clone_with_input(c.data); - // return Ok(Transformed::yes(Arc::new(p) as Arc)); - // } else { - // return not_impl_err!("Unsupported ExecutionPlan type: {}", plan.name()); - // } - // } - - // return Ok(plan); } fn try_get_updated_aggr_expr_from_child( @@ -232,9 +168,6 @@ fn optimize_internal( let required_input_ordering = (!new_requirement.is_empty()).then_some(new_requirement); - println!("required_input_ordering: {:?}", required_input_ordering); - println!("agg_expr: {:?}", aggr_expr); - let input_order_mode = if indices.len() == groupby_exprs.len() && !indices.is_empty() { InputOrderMode::Sorted @@ -402,22 +335,17 @@ fn get_aggregate_exprs_requirement( if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { let mut first_value = first_value.clone(); - // println!("prefix_requirement: {:?}", prefix_requirement); - // println!("aggr_req: {:?}", aggr_req); - // println!("reverse_aggr_req: {:?}", reverse_aggr_req); - + println!("reverse_aggr_req: {:?}", reverse_aggr_req); if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, &aggr_req, )) { - println!("1st step"); first_value = first_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(first_value) as _; } else if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, &reverse_aggr_req, )) { - println!("2nd step"); // Converting to LAST_VALUE enables more efficient execution // given the existing ordering: let mut last_value = first_value.convert_to_last(); diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 798dfa2d5bfbc..40a1a60d6865f 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -465,6 +465,8 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { .create_initial_plan(logical_plan, session_state) .await?; + // println!("first plan: {:#?}", plan); + self.optimize_internal(plan, session_state, |_, _| {}) } } diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs index 367cc35435771..c14c88d6c69bf 100644 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ b/datafusion/physical-expr/src/equivalence/properties.rs @@ -325,7 +325,6 @@ impl EquivalenceProperties { fn ordering_satisfy_single(&self, req: &PhysicalSortRequirement) -> bool { let expr_ordering = self.get_expr_ordering(req.expr.clone()); let ExprOrdering { expr, data, .. } = expr_ordering; - // println!("expr: {:?}, data: {:?}", expr, data); match data { SortProperties::Ordered(options) => { let sort_expr = PhysicalSortExpr { expr, options }; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index be4ba9fe8a164..6ce68cb27e6be 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -48,7 +48,7 @@ use datafusion_physical_expr::{ AggregateExpr, LexRequirement, PhysicalExpr, }; use datafusion_physical_expr::{ - physical_exprs_contains, reverse_order_bys, EquivalenceProperties, LexOrdering, + physical_exprs_contains, EquivalenceProperties, LexOrdering, PhysicalSortRequirement, }; @@ -964,15 +964,12 @@ fn get_aggregate_exprs_requirement( let mut requirement = vec![]; for aggr_expr in aggr_exprs.iter_mut() { let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); - let reverse_aggr_req = reverse_order_bys(aggr_req); + // let reverse_aggr_req = reverse_order_bys(aggr_req); let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); - let reverse_aggr_req = - PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); + // let reverse_aggr_req = + // PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { - println!("prefix_requirement {:?}", prefix_requirement); - println!("reverse_aggr_req {:?}", reverse_aggr_req); - let mut first_value = first_value.clone(); if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, @@ -1302,11 +1299,11 @@ mod tests { use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_physical_expr::expressions::{ - lit, ApproxDistinct, Count, FirstValue, LastValue, Median, OrderSensitiveArrayAgg, + lit, ApproxDistinct, Count, Median, }; use datafusion_physical_expr::{ reverse_order_bys, AggregateExpr, EquivalenceProperties, PhysicalExpr, - PhysicalSortExpr, PhysicalSortRequirement, + PhysicalSortExpr, }; use futures::{FutureExt, Stream}; diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 2c82f2e6a68d6..3ec626aed698f 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -239,13 +239,13 @@ logical_plan after optimize_projections SAME TEXT AS ABOVE 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 initial_physical_plan_with_stats CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] -physical_plan after SimpleOrdering CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true physical_plan after OutputRequirements OutputRequirementExec --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 LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after SimpleOrdering 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 @@ -296,7 +296,6 @@ EXPLAIN VERBOSE SELECT * FROM alltypes_plain limit 10; initial_physical_plan GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] -physical_plan after SimpleOrdering SAME TEXT AS ABOVE physical_plan after OutputRequirements OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] @@ -304,6 +303,7 @@ OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after SimpleOrdering 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 @@ -334,9 +334,6 @@ GlobalLimitExec: skip=0, fetch=10 initial_physical_plan_with_stats GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] --ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] -physical_plan after SimpleOrdering -GlobalLimitExec: skip=0, fetch=10 ---ParquetExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10 physical_plan after OutputRequirements OutputRequirementExec --GlobalLimitExec: skip=0, fetch=10 @@ -344,6 +341,7 @@ OutputRequirementExec physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE +physical_plan after SimpleOrdering 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 From a37fb3e3e0db63219a89c6c57926369d3dfab008 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 09:29:43 +0800 Subject: [PATCH 18/33] cleanup Signed-off-by: jayzhan211 --- datafusion/physical-plan/src/coalesce_partitions.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 67ceb07aeeb6d..1c725ce31f146 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -72,14 +72,6 @@ impl CoalescePartitionsExec { input.execution_mode(), // Execution Mode ) } - - // pub fn clone_with_input(&self, input: Arc) -> Self { - // CoalescePartitionsExec { - // input, - // metrics: self.metrics.clone(), - // cache: self.cache.clone(), - // } - // } } impl DisplayAs for CoalescePartitionsExec { From 8de501c857b9a3b782ea9cc5bb77d18ebd643d18 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 19:50:09 +0800 Subject: [PATCH 19/33] cleanup Signed-off-by: jayzhan211 --- .../physical_optimizer/simplify_ordering.rs | 113 ++---------------- .../physical-plan/src/aggregates/mod.rs | 39 +++--- 2 files changed, 27 insertions(+), 125 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/simplify_ordering.rs index adf4b59b74a20..5f67e0b2edc68 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/simplify_ordering.rs @@ -23,11 +23,10 @@ use datafusion_common::{ }; use datafusion_physical_expr::expressions::{FirstValue, LastValue}; use datafusion_physical_expr::{ - aggregate::is_order_sensitive, - equivalence::ProjectionMapping, - physical_exprs_contains, reverse_order_bys, AggregateExpr, EquivalenceProperties, - LexOrdering, LexRequirement, PhysicalSortRequirement, + equivalence::ProjectionMapping, reverse_order_bys, AggregateExpr, + EquivalenceProperties, LexRequirement, PhysicalSortRequirement, }; +use datafusion_physical_plan::aggregates::{concat_slices, finer_ordering}; use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, @@ -54,7 +53,6 @@ impl PhysicalOptimizerRule for ConvertFirstLast { plan: Arc, _config: &ConfigOptions, ) -> Result> { - let res = plan .transform_down(&get_common_requirement_of_aggregate_input) .data(); @@ -107,7 +105,9 @@ fn try_get_updated_aggr_expr_from_child( aggr_exec: &AggregateExec, ) -> Vec> { let input = aggr_exec.input(); - if aggr_exec.mode() == &AggregateMode::Final || aggr_exec.mode() == &AggregateMode::FinalPartitioned { + if aggr_exec.mode() == &AggregateMode::Final + || aggr_exec.mode() == &AggregateMode::FinalPartitioned + { // Some aggregators may be modified during initialization for // optimization purposes. For example, a FIRST_VALUE may turn // into a LAST_VALUE with the reverse ordering requirement. @@ -115,7 +115,6 @@ fn try_get_updated_aggr_expr_from_child( // `AggregateExpr`/`PhysicalSortExpr` objects. // // The bottom up transformation is the mirror of LogicalPlan::Aggregate creation in [create_initial_plan] - if let Some(c_aggr_exec) = input.as_any().downcast_ref::() { if c_aggr_exec.mode() == &AggregateMode::Partial { // If the input is an AggregateExec in Partial mode, then the @@ -135,7 +134,6 @@ fn optimize_internal( plan: Arc, ) -> Result>> { if let Some(aggr_exec) = plan.as_any().downcast_ref::() { - let input = aggr_exec.input(); let mut aggr_expr = try_get_updated_aggr_expr_from_child(aggr_exec); let group_by = aggr_exec.group_by(); @@ -187,20 +185,6 @@ fn optimize_internal( &input_order_mode, ); - // let p = AggregateExec { - // mode: mode.clone(), - // group_by: group_by.clone(), - // aggr_expr, - // filter_expr: aggr_exec.filter_expr().to_vec(), - // input: input.clone(), - // schema: aggr_exec.schema(), - // input_schema: aggr_exec.input_schema(), - // metrics: aggr_exec.metrics().clone(), - // required_input_ordering, - // limit: None, - // input_order_mode, - // cache, - // }; let p = aggr_exec.clone_with_required_input_ordering_and_aggr_expr( required_input_ordering, aggr_expr, @@ -221,85 +205,6 @@ fn optimize_internal( } } -/// Determines the lexical ordering requirement for an aggregate expression. -/// -/// # Parameters -/// -/// - `aggr_expr`: A reference to an `Arc` representing the -/// aggregate expression. -/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the -/// physical GROUP BY expression. -/// - `agg_mode`: A reference to an `AggregateMode` instance representing the -/// mode of aggregation. -/// -/// # Returns -/// -/// A `LexOrdering` instance indicating the lexical ordering requirement for -/// the aggregate expression. -fn get_aggregate_expr_req( - aggr_expr: &Arc, - group_by: &PhysicalGroupBy, - agg_mode: &AggregateMode, -) -> LexOrdering { - // If the aggregation function is not order sensitive, or the aggregation - // is performing a "second stage" calculation, or all aggregate function - // requirements are inside the GROUP BY expression, then ignore the ordering - // requirement. - if !is_order_sensitive(aggr_expr) || !agg_mode.is_first_stage() { - return vec![]; - } - - let mut req = aggr_expr.order_bys().unwrap_or_default().to_vec(); - - // In non-first stage modes, we accumulate data (using `merge_batch`) from - // different partitions (i.e. merge partial results). During this merge, we - // consider the ordering of each partial result. Hence, we do not need to - // use the ordering requirement in such modes as long as partial results are - // generated with the correct ordering. - if group_by.is_single() { - // Remove all orderings that occur in the group by. These requirements - // will definitely be satisfied -- Each group by expression will have - // distinct values per group, hence all requirements are satisfied. - let physical_exprs = group_by.input_exprs(); - req.retain(|sort_expr| { - !physical_exprs_contains(&physical_exprs, &sort_expr.expr) - }); - } - req -} - -/// Computes the finer ordering for between given existing ordering requirement -/// of aggregate expression. -/// -/// # Parameters -/// -/// * `existing_req` - The existing lexical ordering that needs refinement. -/// * `aggr_expr` - A reference to an aggregate expression trait object. -/// * `group_by` - Information about the physical grouping (e.g group by expression). -/// * `eq_properties` - Equivalence properties relevant to the computation. -/// * `agg_mode` - The mode of aggregation (e.g., Partial, Final, etc.). -/// -/// # Returns -/// -/// An `Option` representing the computed finer lexical ordering, -/// or `None` if there is no finer ordering; e.g. the existing requirement and -/// the aggregator requirement is incompatible. -fn finer_ordering( - existing_req: &LexOrdering, - aggr_expr: &Arc, - group_by: &PhysicalGroupBy, - eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Option { - let aggr_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); - eq_properties.get_finer_ordering(existing_req, &aggr_req) -} - -/// Concatenates the given slices. -fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { - [lhs, rhs].concat() -} - /// Get the common requirement that satisfies all the aggregate expressions. /// /// # Parameters @@ -335,7 +240,6 @@ fn get_aggregate_exprs_requirement( if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { let mut first_value = first_value.clone(); - println!("reverse_aggr_req: {:?}", reverse_aggr_req); if eq_properties.ordering_satisfy_requirement(&concat_slices( prefix_requirement, &aggr_req, @@ -443,13 +347,14 @@ fn get_aggregate_exprs_requirement( mod tests { use super::*; + use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions}; use datafusion_physical_expr::{ expressions::{col, OrderSensitiveArrayAgg}, PhysicalSortExpr, }; - fn create_test_schema() -> Result { + fn create_test_schema_v4() -> Result { let a = Field::new("a", DataType::Int32, true); let b = Field::new("b", DataType::Int32, true); let c = Field::new("c", DataType::Int32, true); @@ -462,7 +367,7 @@ mod tests { #[tokio::test] async fn test_get_finest_requirements() -> Result<()> { - let test_schema = create_test_schema()?; + let test_schema = create_test_schema_v4()?; // Assume column a and b are aliases // Assume also that a ASC and c DESC describe the same global ordering for the table. (Since they are ordering equivalent). let options1 = SortOptions { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6ce68cb27e6be..11c2c7d8a5941 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -48,8 +48,7 @@ use datafusion_physical_expr::{ AggregateExpr, LexRequirement, PhysicalExpr, }; use datafusion_physical_expr::{ - physical_exprs_contains, EquivalenceProperties, LexOrdering, - PhysicalSortRequirement, + physical_exprs_contains, EquivalenceProperties, LexOrdering, PhysicalSortRequirement, }; use itertools::Itertools; @@ -246,19 +245,19 @@ impl From for SendableRecordBatchStream { #[derive(Debug)] pub struct AggregateExec { /// Aggregation mode (full, partial) - pub mode: AggregateMode, + mode: AggregateMode, /// Group by expressions - pub group_by: PhysicalGroupBy, + group_by: PhysicalGroupBy, /// Aggregate expressions - pub aggr_expr: Vec>, + aggr_expr: Vec>, /// FILTER (WHERE clause) expression for each aggregate expression - pub filter_expr: Vec>>, + filter_expr: Vec>>, /// Set if the output of this aggregation is truncated by a upstream sort/limit clause - pub limit: Option, + limit: Option, /// Input plan, could be a partial aggregate or the input to the aggregate pub input: Arc, /// Schema after the aggregate is applied - pub schema: SchemaRef, + schema: SchemaRef, /// Input schema before any aggregation is applied. For partial aggregate this will be the /// same as input.schema() but for the final aggregate it will be the same as the input /// to the partial aggregate, i.e., partial and final aggregates have same `input_schema`. @@ -266,11 +265,11 @@ pub struct AggregateExec { /// expressions from protobuf for final aggregate. pub input_schema: SchemaRef, /// Execution metrics - pub metrics: ExecutionPlanMetricsSet, - pub required_input_ordering: Option, + metrics: ExecutionPlanMetricsSet, + required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns - pub input_order_mode: InputOrderMode, - pub cache: PlanProperties, + input_order_mode: InputOrderMode, + cache: PlanProperties, } impl AggregateExec { @@ -282,19 +281,19 @@ impl AggregateExec { input_order_mode: InputOrderMode, ) -> Self { Self { + aggr_expr, required_input_ordering, + metrics: ExecutionPlanMetricsSet::new(), + input_order_mode, + cache, // clone the rest of the fields mode: self.mode, group_by: self.group_by.clone(), - aggr_expr, filter_expr: self.filter_expr.clone(), limit: self.limit, input: self.input.clone(), schema: self.schema.clone(), input_schema: self.input_schema.clone(), - metrics: ExecutionPlanMetricsSet::new(), - input_order_mode, - cache, } } @@ -921,7 +920,7 @@ fn get_aggregate_expr_req( /// An `Option` representing the computed finer lexical ordering, /// or `None` if there is no finer ordering; e.g. the existing requirement and /// the aggregator requirement is incompatible. -fn finer_ordering( +pub fn finer_ordering( existing_req: &LexOrdering, aggr_expr: &Arc, group_by: &PhysicalGroupBy, @@ -933,7 +932,7 @@ fn finer_ordering( } /// Concatenates the given slices. -fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { +pub fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { [lhs, rhs].concat() } @@ -1298,9 +1297,7 @@ mod tests { use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; - use datafusion_physical_expr::expressions::{ - lit, ApproxDistinct, Count, Median, - }; + use datafusion_physical_expr::expressions::{lit, ApproxDistinct, Count, Median}; use datafusion_physical_expr::{ reverse_order_bys, AggregateExpr, EquivalenceProperties, PhysicalExpr, PhysicalSortExpr, From 107f4c82b9e5b86796e59eb7e6b02c2a2ee87035 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 19:52:30 +0800 Subject: [PATCH 20/33] add aggregate test Signed-off-by: jayzhan211 --- .../sqllogictest/test_files/aggregate.slt | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 4929ab485d6d7..69fb4ea38e9fb 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3428,3 +3428,45 @@ SELECT LAST_VALUE(column1 ORDER BY column2 DESC) IGNORE NULLS FROM t; statement ok DROP TABLE t; + +# Test Convert FirstLast optimizer rule +statement ok +CREATE EXTERNAL TABLE convert_first_last_table ( +c1 INT NOT NULL, +c2 INT NOT NULL, +c3 INT NOT NULL +) +STORED AS CSV +WITH HEADER ROW +WITH ORDER (c1 ASC) +WITH ORDER (c2 DESC) +WITH ORDER (c3 ASC) +LOCATION '../core/tests/data/convert_first_last.csv'; + +# test first to last, the result does not show difference, we need to check the conversion by `explain` +query TT +explain select first_value(c1 order by c3 desc) from convert_first_last_table; +---- +logical_plan +Aggregate: groupBy=[[]], aggr=[[FIRST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c3 DESC NULLS FIRST]]] +--TableScan: convert_first_last_table projection=[c1, c3] +physical_plan +AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +--CoalescePartitionsExec +----AggregateExec: mode=Partial, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c3], output_orderings=[[c1@0 ASC NULLS LAST], [c3@1 ASC NULLS LAST]], has_header=true + +# test last to first +query TT +explain select last_value(c1 order by c2 asc) from convert_first_last_table; +---- +logical_plan +Aggregate: groupBy=[[]], aggr=[[LAST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c2 ASC NULLS LAST]]] +--TableScan: convert_first_last_table projection=[c1, c2] +physical_plan +AggregateExec: mode=Final, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] +--CoalescePartitionsExec +----AggregateExec: mode=Partial, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] +------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c2], output_orderings=[[c1@0 ASC NULLS LAST], [c2@1 DESC]], has_header=true From 345d4c60cde8c05aaee218f7a7f3c6eee39c6c82 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 20:06:52 +0800 Subject: [PATCH 21/33] cleanup Signed-off-by: jayzhan211 --- ...simplify_ordering.rs => convert_first_last.rs} | 15 +++++---------- datafusion/core/src/physical_optimizer/mod.rs | 2 +- .../core/src/physical_optimizer/optimizer.rs | 2 +- datafusion/physical-plan/src/aggregates/mod.rs | 4 +++- datafusion/physical-plan/src/recursive_query.rs | 1 - datafusion/physical-plan/src/tree_node.rs | 2 -- 6 files changed, 10 insertions(+), 16 deletions(-) rename datafusion/core/src/physical_optimizer/{simplify_ordering.rs => convert_first_last.rs} (96%) diff --git a/datafusion/core/src/physical_optimizer/simplify_ordering.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs similarity index 96% rename from datafusion/core/src/physical_optimizer/simplify_ordering.rs rename to datafusion/core/src/physical_optimizer/convert_first_last.rs index 5f67e0b2edc68..a7bb01550a4b3 100644 --- a/datafusion/core/src/physical_optimizer/simplify_ordering.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -185,20 +185,14 @@ fn optimize_internal( &input_order_mode, ); - let p = aggr_exec.clone_with_required_input_ordering_and_aggr_expr( + let aggr_exec = aggr_exec.new_with_aggr_expr_and_ordering_info( required_input_ordering, aggr_expr, cache, input_order_mode, ); - // let aggr_exec = aggr_exec.to_owned().to_owned(); - // aggr_exec.with_aggr_expr(aggr_expr); - // aggr_exec.with_required_input_ordering(required_input_ordering); - // aggr_exec.with_cache(cache); - // aggr_exec.with_input_order_mode(input_order_mode); - - let res = Arc::new(p) as Arc; + let res = Arc::new(aggr_exec) as Arc; Ok(Transformed::yes(res)) } else { Ok(Transformed::no(plan)) @@ -345,6 +339,7 @@ fn get_aggregate_exprs_requirement( Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) } +#[cfg(test)] mod tests { use super::*; @@ -354,7 +349,7 @@ mod tests { PhysicalSortExpr, }; - fn create_test_schema_v4() -> Result { + fn create_test_schema() -> Result { let a = Field::new("a", DataType::Int32, true); let b = Field::new("b", DataType::Int32, true); let c = Field::new("c", DataType::Int32, true); @@ -367,7 +362,7 @@ mod tests { #[tokio::test] async fn test_get_finest_requirements() -> Result<()> { - let test_schema = create_test_schema_v4()?; + let test_schema = create_test_schema()?; // Assume column a and b are aliases // Assume also that a ASC and c DESC describe the same global ordering for the table. (Since they are ordering equivalent). let options1 = SortOptions { diff --git a/datafusion/core/src/physical_optimizer/mod.rs b/datafusion/core/src/physical_optimizer/mod.rs index 06f8974719634..c552523d1c1d2 100644 --- a/datafusion/core/src/physical_optimizer/mod.rs +++ b/datafusion/core/src/physical_optimizer/mod.rs @@ -34,7 +34,7 @@ pub mod pipeline_checker; mod projection_pushdown; pub mod pruning; pub mod replace_with_order_preserving_variants; -mod simplify_ordering; +mod convert_first_last; mod sort_pushdown; pub mod topk_aggregation; mod utils; diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 134d11e43a0e5..897a315e64b76 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use super::projection_pushdown::ProjectionPushdown; -use super::simplify_ordering::ConvertFirstLast; +use super::convert_first_last::ConvertFirstLast; use crate::config::ConfigOptions; use crate::physical_optimizer::aggregate_statistics::AggregateStatistics; use crate::physical_optimizer::coalesce_batches::CoalesceBatches; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 11c2c7d8a5941..a99fe45fc8566 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -273,7 +273,9 @@ pub struct AggregateExec { } impl AggregateExec { - pub fn clone_with_required_input_ordering_and_aggr_expr( + /// Function used in `ConvertFirstLast` optimizer rule, + /// where we need parts of the new value, others cloned from the old one + pub fn new_with_aggr_expr_and_ordering_info( &self, required_input_ordering: Option, aggr_expr: Vec>, diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 91abd5cfdf284..ba7d1a54548a4 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -26,7 +26,6 @@ use super::{ work_table::{ReservedBatches, WorkTable, WorkTableExec}, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::aggregates::AggregateExec; use crate::{DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan}; use arrow::datatypes::SchemaRef; diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index 86cf2d0554742..746df9dd8755f 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -17,11 +17,9 @@ //! This module provides common traits for visiting or rewriting tree nodes easily. -use core::panic; use std::fmt::{self, Display, Formatter}; use std::sync::Arc; -use crate::aggregates::AggregateExec; use crate::{displayable, with_new_children_if_necessary, ExecutionPlan}; use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode}; From 030d7b9cbb4dc8bbcb8b0ad684de4c46f7b6391d Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 20:24:29 +0800 Subject: [PATCH 22/33] final draft Signed-off-by: jayzhan211 --- .../physical_optimizer/convert_first_last.rs | 32 ++++++++----- datafusion/core/src/physical_optimizer/mod.rs | 2 +- .../core/src/physical_optimizer/optimizer.rs | 2 +- .../physical_optimizer/output_requirements.rs | 1 - datafusion/core/src/physical_planner.rs | 5 +- .../physical-plan/src/aggregates/mod.rs | 47 +------------------ datafusion/physical-plan/src/lib.rs | 1 - datafusion/sqllogictest/test_files/test1.slt | 47 ------------------- 8 files changed, 25 insertions(+), 112 deletions(-) delete mode 100644 datafusion/sqllogictest/test_files/test1.slt diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index a7bb01550a4b3..d3b97cda844b5 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -38,6 +38,15 @@ use datafusion_physical_plan::windows::get_ordered_partition_by_indices; use super::PhysicalOptimizerRule; +/// The optimizer rule check the ordering requirements of the aggregate expressions. +/// And convert between FIRST_VALUE and LAST_VALUE if possible. +/// For example, If we have an ascending values and we want LastValue from the descending requirement, +/// it is equivalent to FirstValue with the current ascending ordering. +/// +/// The concrete example is that, says we have values c1 with [1, 2, 3], which is an ascending order. +/// If we want LastValue(c1 order by desc), which is the first value of reversed c1 [3, 2, 1], +/// so we can convert the aggregate expression to FirstValue(c1 order by asc), +/// since the current ordering is already satisfied, it saves our time! #[derive(Default)] pub struct ConvertFirstLast {} @@ -53,11 +62,8 @@ impl PhysicalOptimizerRule for ConvertFirstLast { plan: Arc, _config: &ConfigOptions, ) -> Result> { - let res = plan - .transform_down(&get_common_requirement_of_aggregate_input) - .data(); - - res + plan.transform_down(&get_common_requirement_of_aggregate_input) + .data() } fn name(&self) -> &str { @@ -156,7 +162,7 @@ fn optimize_internal( let req = get_aggregate_exprs_requirement( &new_requirement, &mut aggr_expr, - &group_by, + group_by, input_eq_properties, mode, )?; @@ -175,13 +181,13 @@ fn optimize_internal( InputOrderMode::Linear }; let projection_mapping = - ProjectionMapping::try_new(&group_by.expr(), &input.schema())?; + ProjectionMapping::try_new(group_by.expr(), &input.schema())?; let cache = AggregateExec::compute_properties( - &input, + input, plan.schema().clone(), &projection_mapping, - &mode, + mode, &input_order_mode, ); @@ -192,8 +198,9 @@ fn optimize_internal( input_order_mode, ); - let res = Arc::new(aggr_exec) as Arc; - Ok(Transformed::yes(res)) + Ok(Transformed::yes( + Arc::new(aggr_exec) as Arc + )) } else { Ok(Transformed::no(plan)) } @@ -216,6 +223,9 @@ fn optimize_internal( /// /// A `LexRequirement` instance, which is the requirement that satisfies all the /// aggregate requirements. Returns an error in case of conflicting requirements. +/// +/// Similar to the one in datafusion/physical-plan/src/aggregates/mod.rs, but this +/// function care about the possible of optimization of FIRST_VALUE and LAST_VALUE fn get_aggregate_exprs_requirement( prefix_requirement: &[PhysicalSortRequirement], aggr_exprs: &mut [Arc], diff --git a/datafusion/core/src/physical_optimizer/mod.rs b/datafusion/core/src/physical_optimizer/mod.rs index c552523d1c1d2..c80668c6da74a 100644 --- a/datafusion/core/src/physical_optimizer/mod.rs +++ b/datafusion/core/src/physical_optimizer/mod.rs @@ -24,6 +24,7 @@ pub mod aggregate_statistics; pub mod coalesce_batches; pub mod combine_partial_final_agg; +mod convert_first_last; pub mod enforce_distribution; pub mod enforce_sorting; pub mod join_selection; @@ -34,7 +35,6 @@ pub mod pipeline_checker; mod projection_pushdown; pub mod pruning; pub mod replace_with_order_preserving_variants; -mod convert_first_last; mod sort_pushdown; pub mod topk_aggregation; mod utils; diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 897a315e64b76..e8ae03e11c62a 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -19,8 +19,8 @@ use std::sync::Arc; -use super::projection_pushdown::ProjectionPushdown; use super::convert_first_last::ConvertFirstLast; +use super::projection_pushdown::ProjectionPushdown; use crate::config::ConfigOptions; use crate::physical_optimizer::aggregate_statistics::AggregateStatistics; use crate::physical_optimizer::coalesce_batches::CoalesceBatches; diff --git a/datafusion/core/src/physical_optimizer/output_requirements.rs b/datafusion/core/src/physical_optimizer/output_requirements.rs index 9f50309985856..829d523c990ca 100644 --- a/datafusion/core/src/physical_optimizer/output_requirements.rs +++ b/datafusion/core/src/physical_optimizer/output_requirements.rs @@ -274,7 +274,6 @@ fn require_top_ordering_helper( // be responsible for (i.e. the originator of) the global ordering. let (new_child, is_changed) = require_top_ordering_helper(children.swap_remove(0))?; - Ok((plan.with_new_children(vec![new_child])?, is_changed)) } else { // Stop searching, there is no global ordering desired for the query. diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 40a1a60d6865f..c25523c5ae33e 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -74,7 +74,6 @@ use arrow::compute::SortOptions; use arrow::datatypes::{Schema, SchemaRef}; use arrow_array::builder::StringBuilder; use arrow_array::RecordBatch; -use async_trait::async_trait; use datafusion_common::display::ToStringifiedPlan; use datafusion_common::{ exec_err, internal_err, not_impl_err, plan_err, DFSchema, FileType, ScalarValue, @@ -95,6 +94,7 @@ use datafusion_physical_expr::expressions::Literal; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_sql::utils::window_expr_common_partition_keys; +use async_trait::async_trait; use datafusion_common::config::FormatOptions; use datafusion_physical_expr::LexOrdering; use futures::future::BoxFuture; @@ -464,9 +464,6 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { let plan = self .create_initial_plan(logical_plan, session_state) .await?; - - // println!("first plan: {:#?}", plan); - self.optimize_internal(plan, session_state, |_, _| {}) } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index a99fe45fc8566..66488d1d10e35 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -273,7 +273,7 @@ pub struct AggregateExec { } impl AggregateExec { - /// Function used in `ConvertFirstLast` optimizer rule, + /// Function used in `ConvertFirstLast` optimizer rule, /// where we need parts of the new value, others cloned from the old one pub fn new_with_aggr_expr_and_ordering_info( &self, @@ -421,29 +421,6 @@ impl AggregateExec { }) } - pub fn with_cache(mut self, cache: PlanProperties) -> Self { - self.cache = cache; - self - } - - pub fn with_input_order_mode(mut self, input_order_mode: InputOrderMode) -> Self { - self.input_order_mode = input_order_mode; - self - } - - pub fn with_required_input_ordering( - mut self, - required_input_ordering: Option, - ) -> Self { - self.required_input_ordering = required_input_ordering; - self - } - - pub fn with_aggr_expr(mut self, aggr_expr: Vec>) -> Self { - self.aggr_expr = aggr_expr; - self - } - /// Aggregation mode (full, partial) pub fn mode(&self) -> &AggregateMode { &self.mode @@ -965,10 +942,7 @@ fn get_aggregate_exprs_requirement( let mut requirement = vec![]; for aggr_expr in aggr_exprs.iter_mut() { let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); - // let reverse_aggr_req = reverse_order_bys(aggr_req); let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); - // let reverse_aggr_req = - // PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { let mut first_value = first_value.clone(); @@ -978,15 +952,6 @@ fn get_aggregate_exprs_requirement( )) { first_value = first_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(first_value) as _; - // } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - // prefix_requirement, - // &reverse_aggr_req, - // )) { - // // Converting to LAST_VALUE enables more efficient execution - // // given the existing ordering: - // let mut last_value = first_value.convert_to_last(); - // last_value = last_value.with_requirement_satisfied(true); - // *aggr_expr = Arc::new(last_value) as _; } else { // Requirement is not satisfied with existing ordering. first_value = first_value.with_requirement_satisfied(false); @@ -1002,16 +967,6 @@ fn get_aggregate_exprs_requirement( )) { last_value = last_value.with_requirement_satisfied(true); *aggr_expr = Arc::new(last_value) as _; - // } else if eq_properties.ordering_satisfy_requirement(&concat_slices( - // prefix_requirement, - // &reverse_aggr_req, - // )) { - // println!("last to first {:?}", aggr_expr); - // // Converting to FIRST_VALUE enables more efficient execution - // // given the existing ordering: - // let mut first_value = last_value.convert_to_first(); - // first_value = first_value.with_requirement_satisfied(true); - // *aggr_expr = Arc::new(first_value) as _; } else { // Requirement is not satisfied with existing ordering. last_value = last_value.with_requirement_satisfied(false); diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 8890777b40eb8..e1c8489655bf5 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -27,7 +27,6 @@ use crate::metrics::MetricsSet; use crate::repartition::RepartitionExec; use crate::sorts::sort_preserving_merge::SortPreservingMergeExec; -use aggregates::AggregateExec; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; diff --git a/datafusion/sqllogictest/test_files/test1.slt b/datafusion/sqllogictest/test_files/test1.slt deleted file mode 100644 index 4ca847fa01d8e..0000000000000 --- a/datafusion/sqllogictest/test_files/test1.slt +++ /dev/null @@ -1,47 +0,0 @@ -# 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. - -statement ok -CREATE EXTERNAL TABLE convert_first_last_table ( -c1 INT NOT NULL, -c2 INT NOT NULL, -c3 INT NOT NULL -) -STORED AS CSV -WITH HEADER ROW -WITH ORDER (c1 ASC) -WITH ORDER (c2 DESC) -WITH ORDER (c3 ASC) -LOCATION '../core/tests/data/convert_first_last.csv'; - -# test first to last, the result does not show difference, we need to check the conversion by `explain` -query TT -explain select first_value(c1 order by c3 desc) from convert_first_last_table; ----- -logical_plan -Aggregate: groupBy=[[]], aggr=[[FIRST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c3 DESC NULLS FIRST]]] ---TableScan: convert_first_last_table projection=[c1, c3] -physical_plan -AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] ---CoalescePartitionsExec -----AggregateExec: mode=Partial, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] -------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 ---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c3], output_orderings=[[c1@0 ASC NULLS LAST], [c3@1 ASC NULLS LAST]], has_header=true - -# test last to first -# query TT -# explain select last_value(c1 order by c2 asc) from convert_first_last_table; From cf3f33e0e67cc5edf63df40b316f2ef9cba72c61 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 12 Apr 2024 20:27:25 +0800 Subject: [PATCH 23/33] cleanup again Signed-off-by: jayzhan211 --- datafusion/physical-plan/src/aggregates/mod.rs | 6 ------ datafusion/physical-plan/src/tree_node.rs | 3 +-- datafusion/sqllogictest/test_files/group_by.slt | 4 ---- 3 files changed, 1 insertion(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 66488d1d10e35..9a9a29971de90 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -404,7 +404,6 @@ impl AggregateExec { &input_order_mode, ); - // let required_input_ordering = None; Ok(AggregateExec { mode, group_by, @@ -711,9 +710,6 @@ impl ExecutionPlan for AggregateExec { self: Arc, children: Vec>, ) -> Result> { - // let aggr_exec = self.clone_with_input(children[0].clone()); - // Ok(Arc::new(aggr_exec)) - let mut me = AggregateExec::try_new_with_schema( self.mode, self.group_by.clone(), @@ -725,8 +721,6 @@ impl ExecutionPlan for AggregateExec { )?; me.limit = self.limit; - // let me = me.rewrite_ordering()?; - Ok(Arc::new(me)) } diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index 746df9dd8755f..46460cbb66843 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -63,8 +63,7 @@ impl PlanContext { pub fn update_plan_from_children(mut self) -> Result { let children_plans = self.children.iter().map(|c| c.plan.clone()).collect(); - let plan = with_new_children_if_necessary(self.plan, children_plans)?; - self.plan = plan; + self.plan = with_new_children_if_necessary(self.plan, children_plans)?; Ok(self) } diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 88e0272ef9bca..869462b4722a8 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -4249,10 +4249,6 @@ SELECT date_bin('15 minutes', ts) as time_chunks 2018-12-13T12:00:00 2018-11-13T17:00:00 - - - - # Since extract is not a monotonic function, below query should not run. # when source is unbounded. query error From 6d9ee9f9ee4ebedfeb0afb688a87723a2b1cf3ef Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 09:10:07 +0800 Subject: [PATCH 24/33] pull out finer ordering code and reuse Signed-off-by: jayzhan211 --- .../physical_optimizer/convert_first_last.rs | 59 +------ .../physical-plan/src/aggregates/mod.rs | 156 +++++++++--------- 2 files changed, 79 insertions(+), 136 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index d3b97cda844b5..97b95a94dca6c 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -26,7 +26,7 @@ use datafusion_physical_expr::{ equivalence::ProjectionMapping, reverse_order_bys, AggregateExpr, EquivalenceProperties, LexRequirement, PhysicalSortRequirement, }; -use datafusion_physical_plan::aggregates::{concat_slices, finer_ordering}; +use datafusion_physical_plan::aggregates::{concat_slices, finer_ordering, optimize_for_finer_ordering}; use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, @@ -290,61 +290,8 @@ fn get_aggregate_exprs_requirement( } continue; } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Requirement is satisfied by existing ordering - requirement = finer_ordering; - continue; - } - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Reverse requirement is satisfied by exiting ordering. - // Hence reverse the aggregator - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } - } - } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - // There is a requirement that both satisfies existing requirement and current - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - continue; - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - // There is a requirement that both satisfies existing requirement and reverse - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } - } - // Neither the existing requirement and current aggregate requirement satisfy the other, this means - // requirements are conflicting. Currently, we do not support - // conflicting requirements. - return not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" - ); + + optimize_for_finer_ordering(&mut requirement, aggr_expr, group_by, eq_properties, agg_mode)?; } Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 9a9a29971de90..9d848ae9ed1f7 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -938,93 +938,89 @@ fn get_aggregate_exprs_requirement( let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); - if let Some(first_value) = aggr_expr.as_any().downcast_ref::() { - let mut first_value = first_value.clone(); - if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &aggr_req, - )) { - first_value = first_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(first_value) as _; - } else { - // Requirement is not satisfied with existing ordering. - first_value = first_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(first_value) as _; - } - continue; - } - if let Some(last_value) = aggr_expr.as_any().downcast_ref::() { - let mut last_value = last_value.clone(); - if eq_properties.ordering_satisfy_requirement(&concat_slices( - prefix_requirement, - &aggr_req, - )) { - last_value = last_value.with_requirement_satisfied(true); - *aggr_expr = Arc::new(last_value) as _; - } else { - // Requirement is not satisfied with existing ordering. - last_value = last_value.with_requirement_satisfied(false); - *aggr_expr = Arc::new(last_value) as _; - } - continue; + optimize_for_finer_ordering( + &mut requirement, + aggr_expr, + group_by, + eq_properties, + agg_mode, + )?; + } + Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) +} + +/// Optimize the requirement for finer ordering based on various parameters. +/// +/// This function takes in a mutable reference to a `LexOrdering`, a mutable reference to an `AggregateExpr`, +/// a `PhysicalGroupBy`, an `EquivalenceProperties`, and an `AggregateMode`, and attempts to optimize the +/// requirement for finer ordering. It checks if the finer ordering satisfies existing requirements or +/// if a reverse ordering satisfies them, updating the references accordingly. If neither satisfies the +/// requirements, it returns an error indicating conflicting ordering requirements. +pub fn optimize_for_finer_ordering( + requirement: &mut LexOrdering, + aggr_expr: &mut Arc, + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Result<()> { + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Requirement is satisfied by existing ordering + *requirement = finer_ordering; + return Ok(()); } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { if eq_properties.ordering_satisfy(&finer_ordering) { - // Requirement is satisfied by existing ordering - requirement = finer_ordering; - continue; - } - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Reverse requirement is satisfied by exiting ordering. - // Hence reverse the aggregator - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } + // Reverse requirement is satisfied by exiting ordering. + // Hence reverse the aggregator + *requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + return Ok(()); } } - if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - // There is a requirement that both satisfies existing requirement and current + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + // There is a requirement that both satisfies existing requirement and current + // aggregate requirement. Use updated requirement + // return (Some(finer_ordering), None); + *requirement = finer_ordering; + return Ok(()); + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + // There is a requirement that both satisfies existing requirement and reverse // aggregate requirement. Use updated requirement - requirement = finer_ordering; - continue; - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - &requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - // There is a requirement that both satisfies existing requirement and reverse - // aggregate requirement. Use updated requirement - requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - continue; - } + // return (Some(finer_ordering), Some(reverse_aggr_expr)); + *requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + return Ok(()); } - // Neither the existing requirement and current aggregate requirement satisfy the other, this means - // requirements are conflicting. Currently, we do not support - // conflicting requirements. - return not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" - ); } - Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) + + // Neither the existing requirement and current aggregate requirement satisfy the other, this means + // requirements are conflicting. Currently, we do not support + // conflicting requirements. + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ); } /// returns physical expressions for arguments to evaluate against a batch From abd40c607ee50184c66952b0bb364eb20dcab170 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 09:22:31 +0800 Subject: [PATCH 25/33] clippy Signed-off-by: jayzhan211 --- .../physical_optimizer/convert_first_last.rs | 11 ++++++--- .../physical-plan/src/aggregates/mod.rs | 24 ++++++++----------- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index 97b95a94dca6c..28aebe049e447 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -18,7 +18,6 @@ use datafusion_common::Result; use datafusion_common::{ config::ConfigOptions, - not_impl_err, tree_node::{Transformed, TransformedResult, TreeNode}, }; use datafusion_physical_expr::expressions::{FirstValue, LastValue}; @@ -26,7 +25,7 @@ use datafusion_physical_expr::{ equivalence::ProjectionMapping, reverse_order_bys, AggregateExpr, EquivalenceProperties, LexRequirement, PhysicalSortRequirement, }; -use datafusion_physical_plan::aggregates::{concat_slices, finer_ordering, optimize_for_finer_ordering}; +use datafusion_physical_plan::aggregates::{concat_slices, optimize_for_finer_ordering}; use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, @@ -291,7 +290,13 @@ fn get_aggregate_exprs_requirement( continue; } - optimize_for_finer_ordering(&mut requirement, aggr_expr, group_by, eq_properties, agg_mode)?; + optimize_for_finer_ordering( + &mut requirement, + aggr_expr, + group_by, + eq_properties, + agg_mode, + )?; } Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 9d848ae9ed1f7..6fd8a887bad22 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -41,7 +41,6 @@ use datafusion_execution::TaskContext; use datafusion_expr::Accumulator; use datafusion_physical_expr::aggregate::is_order_sensitive; use datafusion_physical_expr::equivalence::collapse_lex_req; -use datafusion_physical_expr::expressions::{FirstValue, LastValue}; use datafusion_physical_expr::{ equivalence::ProjectionMapping, expressions::{Column, Max, Min, UnKnownColumn}, @@ -371,7 +370,6 @@ impl AggregateExec { .collect::>(); let req = get_aggregate_exprs_requirement( - &new_requirement, &mut aggr_expr, &group_by, input_eq_properties, @@ -927,7 +925,6 @@ pub fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { /// A `LexRequirement` instance, which is the requirement that satisfies all the /// aggregate requirements. Returns an error in case of conflicting requirements. fn get_aggregate_exprs_requirement( - prefix_requirement: &[PhysicalSortRequirement], aggr_exprs: &mut [Arc], group_by: &PhysicalGroupBy, eq_properties: &EquivalenceProperties, @@ -935,9 +932,6 @@ fn get_aggregate_exprs_requirement( ) -> Result { let mut requirement = vec![]; for aggr_expr in aggr_exprs.iter_mut() { - let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); - let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); - optimize_for_finer_ordering( &mut requirement, aggr_expr, @@ -955,7 +949,7 @@ fn get_aggregate_exprs_requirement( /// a `PhysicalGroupBy`, an `EquivalenceProperties`, and an `AggregateMode`, and attempts to optimize the /// requirement for finer ordering. It checks if the finer ordering satisfies existing requirements or /// if a reverse ordering satisfies them, updating the references accordingly. If neither satisfies the -/// requirements, it returns an error indicating conflicting ordering requirements. +/// requirements, it returns an error indicating conflicting ordering requirements. pub fn optimize_for_finer_ordering( requirement: &mut LexOrdering, aggr_expr: &mut Arc, @@ -964,7 +958,7 @@ pub fn optimize_for_finer_ordering( agg_mode: &AggregateMode, ) -> Result<()> { if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + finer_ordering(requirement, aggr_expr, group_by, eq_properties, agg_mode) { if eq_properties.ordering_satisfy(&finer_ordering) { // Requirement is satisfied by existing ordering @@ -974,7 +968,7 @@ pub fn optimize_for_finer_ordering( } if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { if let Some(finer_ordering) = finer_ordering( - &requirement, + requirement, &reverse_aggr_expr, group_by, eq_properties, @@ -990,7 +984,7 @@ pub fn optimize_for_finer_ordering( } } if let Some(finer_ordering) = - finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + finer_ordering(requirement, aggr_expr, group_by, eq_properties, agg_mode) { // There is a requirement that both satisfies existing requirement and current // aggregate requirement. Use updated requirement @@ -1000,7 +994,7 @@ pub fn optimize_for_finer_ordering( } if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { if let Some(finer_ordering) = finer_ordering( - &requirement, + requirement, &reverse_aggr_expr, group_by, eq_properties, @@ -1018,9 +1012,9 @@ pub fn optimize_for_finer_ordering( // Neither the existing requirement and current aggregate requirement satisfy the other, this means // requirements are conflicting. Currently, we do not support // conflicting requirements. - return not_impl_err!( + not_impl_err!( "Conflicting ordering requirements in aggregate functions is not supported" - ); + ) } /// returns physical expressions for arguments to evaluate against a batch @@ -1244,7 +1238,9 @@ mod tests { use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; - use datafusion_physical_expr::expressions::{lit, ApproxDistinct, Count, Median}; + use datafusion_physical_expr::expressions::{ + lit, ApproxDistinct, Count, FirstValue, LastValue, Median, + }; use datafusion_physical_expr::{ reverse_order_bys, AggregateExpr, EquivalenceProperties, PhysicalExpr, PhysicalSortExpr, From 3e2c18606236f47572a1534c1e325c8c84884c6e Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 15:16:01 +0800 Subject: [PATCH 26/33] remove finer in optimize rule Signed-off-by: jayzhan211 --- .../physical_optimizer/convert_first_last.rs | 153 +++--------------- .../physical-plan/src/aggregates/mod.rs | 101 +++++++++++- 2 files changed, 116 insertions(+), 138 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index 28aebe049e447..f95fe2eb27b09 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -23,16 +23,15 @@ use datafusion_common::{ use datafusion_physical_expr::expressions::{FirstValue, LastValue}; use datafusion_physical_expr::{ equivalence::ProjectionMapping, reverse_order_bys, AggregateExpr, - EquivalenceProperties, LexRequirement, PhysicalSortRequirement, + EquivalenceProperties, PhysicalSortRequirement, }; -use datafusion_physical_plan::aggregates::{concat_slices, optimize_for_finer_ordering}; +use datafusion_physical_plan::aggregates::concat_slices; use datafusion_physical_plan::{ - aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, + aggregates::{AggregateExec, AggregateMode}, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, }; use std::sync::Arc; -use datafusion_physical_expr::equivalence::collapse_lex_req; use datafusion_physical_plan::windows::get_ordered_partition_by_indices; use super::PhysicalOptimizerRule; @@ -61,7 +60,7 @@ impl PhysicalOptimizerRule for ConvertFirstLast { plan: Arc, _config: &ConfigOptions, ) -> Result> { - plan.transform_down(&get_common_requirement_of_aggregate_input) + plan.transform_up(&get_common_requirement_of_aggregate_input) .data() } @@ -77,24 +76,26 @@ impl PhysicalOptimizerRule for ConvertFirstLast { fn get_common_requirement_of_aggregate_input( plan: Arc, ) -> Result>> { + // Optimize children let children = plan.children(); - let mut is_child_transformed = false; let mut new_children: Vec> = vec![]; for c in children.iter() { - let res = get_common_requirement_of_aggregate_input(c.clone())?; + let res = optimize_internal(c.clone())?; if res.transformed { is_child_transformed = true; } new_children.push(res.data); } + // Update children if transformed let plan = if is_child_transformed { plan.with_new_children(new_children)? } else { plan }; + // Update itself let plan = optimize_internal(plan)?; // If one of the children is transformed, then the plan is considered transformed, then we update @@ -106,6 +107,9 @@ fn get_common_requirement_of_aggregate_input( } } +/// In `create_initial_plan` for LogicalPlan::Aggregate, we have a nested AggregateExec where the first layer +/// is in Partial mode and the second layer is in Final or Finalpartitioned mode. +/// If the first layer of aggregate plan is transformed, we need to update the child of the layer with final mode. fn try_get_updated_aggr_expr_from_child( aggr_exec: &AggregateExec, ) -> Vec> { @@ -150,7 +154,7 @@ fn optimize_internal( // prefix requirements with this section. In this case, aggregation will // work more efficiently. let indices = get_ordered_partition_by_indices(&groupby_exprs, input); - let mut new_requirement = indices + let requirement = indices .iter() .map(|&idx| PhysicalSortRequirement { expr: groupby_exprs[idx].clone(), @@ -158,18 +162,13 @@ fn optimize_internal( }) .collect::>(); - let req = get_aggregate_exprs_requirement( - &new_requirement, + get_aggregate_exprs_requirement( + &requirement, &mut aggr_expr, - group_by, input_eq_properties, - mode, )?; - new_requirement.extend(req); - new_requirement = collapse_lex_req(new_requirement); - let required_input_ordering = - (!new_requirement.is_empty()).then_some(new_requirement); + let required_input_ordering = (!requirement.is_empty()).then_some(requirement); let input_order_mode = if indices.len() == groupby_exprs.len() && !indices.is_empty() { @@ -224,15 +223,12 @@ fn optimize_internal( /// aggregate requirements. Returns an error in case of conflicting requirements. /// /// Similar to the one in datafusion/physical-plan/src/aggregates/mod.rs, but this -/// function care about the possible of optimization of FIRST_VALUE and LAST_VALUE +/// function care only the possible conversion between FIRST_VALUE and LAST_VALUE fn get_aggregate_exprs_requirement( prefix_requirement: &[PhysicalSortRequirement], aggr_exprs: &mut [Arc], - group_by: &PhysicalGroupBy, eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Result { - let mut requirement = vec![]; +) -> Result<()> { for aggr_expr in aggr_exprs.iter_mut() { let aggr_req = aggr_expr.order_bys().unwrap_or(&[]); let reverse_aggr_req = reverse_order_bys(aggr_req); @@ -289,120 +285,7 @@ fn get_aggregate_exprs_requirement( } continue; } - - optimize_for_finer_ordering( - &mut requirement, - aggr_expr, - group_by, - eq_properties, - agg_mode, - )?; - } - Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) -} - -#[cfg(test)] -mod tests { - use super::*; - - use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions}; - use datafusion_physical_expr::{ - expressions::{col, OrderSensitiveArrayAgg}, - PhysicalSortExpr, - }; - - fn create_test_schema() -> Result { - let a = Field::new("a", DataType::Int32, true); - let b = Field::new("b", DataType::Int32, true); - let c = Field::new("c", DataType::Int32, true); - let d = Field::new("d", DataType::Int32, true); - let e = Field::new("e", DataType::Int32, true); - let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); - - Ok(schema) } - #[tokio::test] - async fn test_get_finest_requirements() -> Result<()> { - let test_schema = create_test_schema()?; - // Assume column a and b are aliases - // Assume also that a ASC and c DESC describe the same global ordering for the table. (Since they are ordering equivalent). - let options1 = SortOptions { - descending: false, - nulls_first: false, - }; - let col_a = &col("a", &test_schema)?; - let col_b = &col("b", &test_schema)?; - let col_c = &col("c", &test_schema)?; - let mut eq_properties = EquivalenceProperties::new(test_schema); - // Columns a and b are equal. - eq_properties.add_equal_conditions(col_a, col_b); - // Aggregate requirements are - // [None], [a ASC], [a ASC, b ASC, c ASC], [a ASC, b ASC] respectively - let order_by_exprs = vec![ - None, - Some(vec![PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }]), - Some(vec![ - PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_b.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_c.clone(), - options: options1, - }, - ]), - Some(vec![ - PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_b.clone(), - options: options1, - }, - ]), - ]; - let common_requirement = vec![ - PhysicalSortExpr { - expr: col_a.clone(), - options: options1, - }, - PhysicalSortExpr { - expr: col_c.clone(), - options: options1, - }, - ]; - let mut aggr_exprs = order_by_exprs - .into_iter() - .map(|order_by_expr| { - Arc::new(OrderSensitiveArrayAgg::new( - col_a.clone(), - "array_agg", - DataType::Int32, - false, - vec![], - order_by_expr.unwrap_or_default(), - )) as _ - }) - .collect::>(); - let group_by = PhysicalGroupBy::new_single(vec![]); - let res = get_aggregate_exprs_requirement( - &[], - &mut aggr_exprs, - &group_by, - &eq_properties, - &AggregateMode::Partial, - )?; - let res = PhysicalSortRequirement::to_sort_exprs(res); - assert_eq!(res, common_requirement); - Ok(()) - } + Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6fd8a887bad22..3f484b8580f3b 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -891,7 +891,7 @@ fn get_aggregate_expr_req( /// An `Option` representing the computed finer lexical ordering, /// or `None` if there is no finer ordering; e.g. the existing requirement and /// the aggregator requirement is incompatible. -pub fn finer_ordering( +fn finer_ordering( existing_req: &LexOrdering, aggr_expr: &Arc, group_by: &PhysicalGroupBy, @@ -950,7 +950,7 @@ fn get_aggregate_exprs_requirement( /// requirement for finer ordering. It checks if the finer ordering satisfies existing requirements or /// if a reverse ordering satisfies them, updating the references accordingly. If neither satisfies the /// requirements, it returns an error indicating conflicting ordering requirements. -pub fn optimize_for_finer_ordering( +fn optimize_for_finer_ordering( requirement: &mut LexOrdering, aggr_expr: &mut Arc, group_by: &PhysicalGroupBy, @@ -1239,7 +1239,7 @@ mod tests { use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_physical_expr::expressions::{ - lit, ApproxDistinct, Count, FirstValue, LastValue, Median, + lit, ApproxDistinct, Count, FirstValue, LastValue, Median, OrderSensitiveArrayAgg, }; use datafusion_physical_expr::{ reverse_order_bys, AggregateExpr, EquivalenceProperties, PhysicalExpr, @@ -1248,6 +1248,18 @@ mod tests { use futures::{FutureExt, Stream}; + // Generate a schema which consists of 5 columns (a, b, c, d, e) + fn create_test_schema() -> Result { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, true); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, true); + let e = Field::new("e", DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); + + Ok(schema) + } + /// some mock data to aggregates fn some_data() -> (Arc, Vec) { // define a schema. @@ -2091,6 +2103,89 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_get_finest_requirements() -> Result<()> { + let test_schema = create_test_schema()?; + // Assume column a and b are aliases + // Assume also that a ASC and c DESC describe the same global ordering for the table. (Since they are ordering equivalent). + let options1 = SortOptions { + descending: false, + nulls_first: false, + }; + let col_a = &col("a", &test_schema)?; + let col_b = &col("b", &test_schema)?; + let col_c = &col("c", &test_schema)?; + let mut eq_properties = EquivalenceProperties::new(test_schema); + // Columns a and b are equal. + eq_properties.add_equal_conditions(col_a, col_b); + // Aggregate requirements are + // [None], [a ASC], [a ASC, b ASC, c ASC], [a ASC, b ASC] respectively + let order_by_exprs = vec![ + None, + Some(vec![PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }]), + Some(vec![ + PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_b.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_c.clone(), + options: options1, + }, + ]), + Some(vec![ + PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_b.clone(), + options: options1, + }, + ]), + ]; + let common_requirement = vec![ + PhysicalSortExpr { + expr: col_a.clone(), + options: options1, + }, + PhysicalSortExpr { + expr: col_c.clone(), + options: options1, + }, + ]; + let mut aggr_exprs = order_by_exprs + .into_iter() + .map(|order_by_expr| { + Arc::new(OrderSensitiveArrayAgg::new( + col_a.clone(), + "array_agg", + DataType::Int32, + false, + vec![], + order_by_expr.unwrap_or_default(), + )) as _ + }) + .collect::>(); + let group_by = PhysicalGroupBy::new_single(vec![]); + let res = get_aggregate_exprs_requirement( + &mut aggr_exprs, + &group_by, + &eq_properties, + &AggregateMode::Partial, + )?; + let res = PhysicalSortRequirement::to_sort_exprs(res); + assert_eq!(res, common_requirement); + Ok(()) + } + #[test] fn test_agg_exec_same_schema() -> Result<()> { let schema = Arc::new(Schema::new(vec![ From d7c259035ec8ff0e1edc0366db2e0d59a27d9414 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 15:28:15 +0800 Subject: [PATCH 27/33] add comments and clenaup Signed-off-by: jayzhan211 --- datafusion/core/src/physical_optimizer/convert_first_last.rs | 3 +++ datafusion/physical-plan/src/aggregates/mod.rs | 2 -- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index f95fe2eb27b09..c067bd2f64526 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -110,6 +110,9 @@ fn get_common_requirement_of_aggregate_input( /// In `create_initial_plan` for LogicalPlan::Aggregate, we have a nested AggregateExec where the first layer /// is in Partial mode and the second layer is in Final or Finalpartitioned mode. /// If the first layer of aggregate plan is transformed, we need to update the child of the layer with final mode. +/// Therefore, we check it and get the updated aggregate expressions. +/// +/// If AggregateExec is created from elsewhere, we skip the check and return the original aggregate expressions. fn try_get_updated_aggr_expr_from_child( aggr_exec: &AggregateExec, ) -> Vec> { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3f484b8580f3b..d457119b5b9b5 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -988,7 +988,6 @@ fn optimize_for_finer_ordering( { // There is a requirement that both satisfies existing requirement and current // aggregate requirement. Use updated requirement - // return (Some(finer_ordering), None); *requirement = finer_ordering; return Ok(()); } @@ -1002,7 +1001,6 @@ fn optimize_for_finer_ordering( ) { // There is a requirement that both satisfies existing requirement and reverse // aggregate requirement. Use updated requirement - // return (Some(finer_ordering), Some(reverse_aggr_expr)); *requirement = finer_ordering; *aggr_expr = reverse_aggr_expr; return Ok(()); From fed356a35857a6b17258efee13dab7601f502916 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 15:34:06 +0800 Subject: [PATCH 28/33] rename fun Signed-off-by: jayzhan211 --- datafusion/core/src/physical_optimizer/convert_first_last.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index c067bd2f64526..45e95827c773d 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -165,7 +165,7 @@ fn optimize_internal( }) .collect::>(); - get_aggregate_exprs_requirement( + try_convert_first_last_if_better( &requirement, &mut aggr_expr, input_eq_properties, @@ -227,7 +227,7 @@ fn optimize_internal( /// /// Similar to the one in datafusion/physical-plan/src/aggregates/mod.rs, but this /// function care only the possible conversion between FIRST_VALUE and LAST_VALUE -fn get_aggregate_exprs_requirement( +fn try_convert_first_last_if_better( prefix_requirement: &[PhysicalSortRequirement], aggr_exprs: &mut [Arc], eq_properties: &EquivalenceProperties, From 0d3d4618a90b99fc642094021ea76ecd5f6224a2 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 15:39:54 +0800 Subject: [PATCH 29/33] rename fun Signed-off-by: jayzhan211 --- .../physical-plan/src/aggregates/mod.rs | 132 ++++++++---------- 1 file changed, 55 insertions(+), 77 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index d457119b5b9b5..e24e71e735401 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -369,7 +369,7 @@ impl AggregateExec { }) .collect::>(); - let req = get_aggregate_exprs_requirement( + let req = get_finer_aggregate_exprs_requirement( &mut aggr_expr, &group_by, input_eq_properties, @@ -924,7 +924,7 @@ pub fn concat_slices(lhs: &[T], rhs: &[T]) -> Vec { /// /// A `LexRequirement` instance, which is the requirement that satisfies all the /// aggregate requirements. Returns an error in case of conflicting requirements. -fn get_aggregate_exprs_requirement( +fn get_finer_aggregate_exprs_requirement( aggr_exprs: &mut [Arc], group_by: &PhysicalGroupBy, eq_properties: &EquivalenceProperties, @@ -932,87 +932,65 @@ fn get_aggregate_exprs_requirement( ) -> Result { let mut requirement = vec![]; for aggr_expr in aggr_exprs.iter_mut() { - optimize_for_finer_ordering( - &mut requirement, - aggr_expr, - group_by, - eq_properties, - agg_mode, - )?; - } - Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) -} - -/// Optimize the requirement for finer ordering based on various parameters. -/// -/// This function takes in a mutable reference to a `LexOrdering`, a mutable reference to an `AggregateExpr`, -/// a `PhysicalGroupBy`, an `EquivalenceProperties`, and an `AggregateMode`, and attempts to optimize the -/// requirement for finer ordering. It checks if the finer ordering satisfies existing requirements or -/// if a reverse ordering satisfies them, updating the references accordingly. If neither satisfies the -/// requirements, it returns an error indicating conflicting ordering requirements. -fn optimize_for_finer_ordering( - requirement: &mut LexOrdering, - aggr_expr: &mut Arc, - group_by: &PhysicalGroupBy, - eq_properties: &EquivalenceProperties, - agg_mode: &AggregateMode, -) -> Result<()> { - if let Some(finer_ordering) = - finer_ordering(requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - if eq_properties.ordering_satisfy(&finer_ordering) { - // Requirement is satisfied by existing ordering - *requirement = finer_ordering; - return Ok(()); - } - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { if eq_properties.ordering_satisfy(&finer_ordering) { - // Reverse requirement is satisfied by exiting ordering. - // Hence reverse the aggregator - *requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - return Ok(()); + // Requirement is satisfied by existing ordering + requirement = finer_ordering; + continue; } } - } - if let Some(finer_ordering) = - finer_ordering(requirement, aggr_expr, group_by, eq_properties, agg_mode) - { - // There is a requirement that both satisfies existing requirement and current - // aggregate requirement. Use updated requirement - *requirement = finer_ordering; - return Ok(()); - } - if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if let Some(finer_ordering) = finer_ordering( - requirement, - &reverse_aggr_expr, - group_by, - eq_properties, - agg_mode, - ) { - // There is a requirement that both satisfies existing requirement and reverse + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + if eq_properties.ordering_satisfy(&finer_ordering) { + // Reverse requirement is satisfied by exiting ordering. + // Hence reverse the aggregator + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + } + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + // There is a requirement that both satisfies existing requirement and current // aggregate requirement. Use updated requirement - *requirement = finer_ordering; - *aggr_expr = reverse_aggr_expr; - return Ok(()); + requirement = finer_ordering; + continue; } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + // There is a requirement that both satisfies existing requirement and reverse + // aggregate requirement. Use updated requirement + requirement = finer_ordering; + *aggr_expr = reverse_aggr_expr; + continue; + } + } + + // Neither the existing requirement and current aggregate requirement satisfy the other, this means + // requirements are conflicting. Currently, we do not support + // conflicting requirements. + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ) } - // Neither the existing requirement and current aggregate requirement satisfy the other, this means - // requirements are conflicting. Currently, we do not support - // conflicting requirements. - not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" - ) + Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) } /// returns physical expressions for arguments to evaluate against a batch @@ -2173,7 +2151,7 @@ mod tests { }) .collect::>(); let group_by = PhysicalGroupBy::new_single(vec![]); - let res = get_aggregate_exprs_requirement( + let res = get_finer_aggregate_exprs_requirement( &mut aggr_exprs, &group_by, &eq_properties, From d3daa4e025ecf97671bd3e93ec55d4599fbda84a Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 15:45:07 +0800 Subject: [PATCH 30/33] fmt Signed-off-by: jayzhan211 --- datafusion/physical-plan/src/aggregates/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index e24e71e735401..ba9a6b1be0efc 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -987,7 +987,7 @@ fn get_finer_aggregate_exprs_requirement( // conflicting requirements. return not_impl_err!( "Conflicting ordering requirements in aggregate functions is not supported" - ) + ); } Ok(PhysicalSortRequirement::from_sort_exprs(&requirement)) From dcb3e90f630b971b57c9e19799e94109ba8f18f3 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 13 Apr 2024 22:04:31 +0800 Subject: [PATCH 31/33] avoid unnecessary recursion and rename Signed-off-by: jayzhan211 --- .../physical_optimizer/convert_first_last.rs | 112 ++++++------------ .../core/src/physical_optimizer/optimizer.rs | 6 +- .../sqllogictest/test_files/explain.slt | 12 +- 3 files changed, 48 insertions(+), 82 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/convert_first_last.rs b/datafusion/core/src/physical_optimizer/convert_first_last.rs index 45e95827c773d..4102313d31268 100644 --- a/datafusion/core/src/physical_optimizer/convert_first_last.rs +++ b/datafusion/core/src/physical_optimizer/convert_first_last.rs @@ -46,15 +46,15 @@ use super::PhysicalOptimizerRule; /// so we can convert the aggregate expression to FirstValue(c1 order by asc), /// since the current ordering is already satisfied, it saves our time! #[derive(Default)] -pub struct ConvertFirstLast {} +pub struct OptimizeAggregateOrder {} -impl ConvertFirstLast { +impl OptimizeAggregateOrder { pub fn new() -> Self { Self::default() } } -impl PhysicalOptimizerRule for ConvertFirstLast { +impl PhysicalOptimizerRule for OptimizeAggregateOrder { fn optimize( &self, plan: Arc, @@ -65,7 +65,7 @@ impl PhysicalOptimizerRule for ConvertFirstLast { } fn name(&self) -> &str { - "SimpleOrdering" + "OptimizeAggregateOrder" } fn schema_check(&self) -> bool { @@ -75,75 +75,6 @@ impl PhysicalOptimizerRule for ConvertFirstLast { fn get_common_requirement_of_aggregate_input( plan: Arc, -) -> Result>> { - // Optimize children - let children = plan.children(); - let mut is_child_transformed = false; - let mut new_children: Vec> = vec![]; - for c in children.iter() { - let res = optimize_internal(c.clone())?; - if res.transformed { - is_child_transformed = true; - } - new_children.push(res.data); - } - - // Update children if transformed - let plan = if is_child_transformed { - plan.with_new_children(new_children)? - } else { - plan - }; - - // Update itself - let plan = optimize_internal(plan)?; - - // If one of the children is transformed, then the plan is considered transformed, then we update - // the children of the plan from bottom to top. - if plan.transformed || is_child_transformed { - Ok(Transformed::yes(plan.data)) - } else { - Ok(Transformed::no(plan.data)) - } -} - -/// In `create_initial_plan` for LogicalPlan::Aggregate, we have a nested AggregateExec where the first layer -/// is in Partial mode and the second layer is in Final or Finalpartitioned mode. -/// If the first layer of aggregate plan is transformed, we need to update the child of the layer with final mode. -/// Therefore, we check it and get the updated aggregate expressions. -/// -/// If AggregateExec is created from elsewhere, we skip the check and return the original aggregate expressions. -fn try_get_updated_aggr_expr_from_child( - aggr_exec: &AggregateExec, -) -> Vec> { - let input = aggr_exec.input(); - if aggr_exec.mode() == &AggregateMode::Final - || aggr_exec.mode() == &AggregateMode::FinalPartitioned - { - // Some aggregators may be modified during initialization for - // optimization purposes. For example, a FIRST_VALUE may turn - // into a LAST_VALUE with the reverse ordering requirement. - // To reflect such changes to subsequent stages, use the updated - // `AggregateExpr`/`PhysicalSortExpr` objects. - // - // The bottom up transformation is the mirror of LogicalPlan::Aggregate creation in [create_initial_plan] - if let Some(c_aggr_exec) = input.as_any().downcast_ref::() { - if c_aggr_exec.mode() == &AggregateMode::Partial { - // If the input is an AggregateExec in Partial mode, then the - // input is a CoalescePartitionsExec. In this case, the - // AggregateExec is the second stage of aggregation. The - // requirements of the second stage are the requirements of - // the first stage. - return c_aggr_exec.aggr_expr().to_vec(); - } - } - } - - aggr_exec.aggr_expr().to_vec() -} - -fn optimize_internal( - plan: Arc, ) -> Result>> { if let Some(aggr_exec) = plan.as_any().downcast_ref::() { let input = aggr_exec.input(); @@ -207,6 +138,41 @@ fn optimize_internal( } } +/// In `create_initial_plan` for LogicalPlan::Aggregate, we have a nested AggregateExec where the first layer +/// is in Partial mode and the second layer is in Final or Finalpartitioned mode. +/// If the first layer of aggregate plan is transformed, we need to update the child of the layer with final mode. +/// Therefore, we check it and get the updated aggregate expressions. +/// +/// If AggregateExec is created from elsewhere, we skip the check and return the original aggregate expressions. +fn try_get_updated_aggr_expr_from_child( + aggr_exec: &AggregateExec, +) -> Vec> { + let input = aggr_exec.input(); + if aggr_exec.mode() == &AggregateMode::Final + || aggr_exec.mode() == &AggregateMode::FinalPartitioned + { + // Some aggregators may be modified during initialization for + // optimization purposes. For example, a FIRST_VALUE may turn + // into a LAST_VALUE with the reverse ordering requirement. + // To reflect such changes to subsequent stages, use the updated + // `AggregateExpr`/`PhysicalSortExpr` objects. + // + // The bottom up transformation is the mirror of LogicalPlan::Aggregate creation in [create_initial_plan] + if let Some(c_aggr_exec) = input.as_any().downcast_ref::() { + if c_aggr_exec.mode() == &AggregateMode::Partial { + // If the input is an AggregateExec in Partial mode, then the + // input is a CoalescePartitionsExec. In this case, the + // AggregateExec is the second stage of aggregation. The + // requirements of the second stage are the requirements of + // the first stage. + return c_aggr_exec.aggr_expr().to_vec(); + } + } + } + + aggr_exec.aggr_expr().to_vec() +} + /// Get the common requirement that satisfies all the aggregate expressions. /// /// # Parameters diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index e8ae03e11c62a..70ac12344810e 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -19,7 +19,7 @@ use std::sync::Arc; -use super::convert_first_last::ConvertFirstLast; +use super::convert_first_last::OptimizeAggregateOrder; use super::projection_pushdown::ProjectionPushdown; use crate::config::ConfigOptions; use crate::physical_optimizer::aggregate_statistics::AggregateStatistics; @@ -91,7 +91,7 @@ impl PhysicalOptimizer { // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), // Run once before PartialFinalAggregation is rewritten to ensure the rule is applied correctly - Arc::new(ConvertFirstLast::new()), + Arc::new(OptimizeAggregateOrder::new()), // 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 @@ -105,7 +105,7 @@ impl PhysicalOptimizer { // as the latter may break local sorting requirements. Arc::new(EnforceSorting::new()), // Run once after the local sorting requirement is changed - Arc::new(ConvertFirstLast::new()), + Arc::new(OptimizeAggregateOrder::new()), // TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future. Arc::new(ProjectionPushdown::new()), // The CoalesceBatches rule will not influence the distribution and ordering of the diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 3ec626aed698f..467bbca3c0bfb 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -245,11 +245,11 @@ OutputRequirementExec physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after OptimizeAggregateOrder 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 -physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], has_header=true @@ -303,11 +303,11 @@ OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after OptimizeAggregateOrder 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 -physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements @@ -341,11 +341,11 @@ OutputRequirementExec physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after OptimizeAggregateOrder 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 -physical_plan after SimpleOrdering SAME TEXT AS ABOVE +physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements From 2951fbc938cd5c3bb442e452a8a8d2c7399b1ee3 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Mon, 15 Apr 2024 20:22:06 +0800 Subject: [PATCH 32/33] remove unnecessary rule Signed-off-by: jayzhan211 --- datafusion/core/src/physical_optimizer/optimizer.rs | 2 -- datafusion/sqllogictest/test_files/aggregate.slt | 4 ++-- datafusion/sqllogictest/test_files/explain.slt | 3 --- datafusion/sqllogictest/test_files/group_by.slt | 4 ++-- 4 files changed, 4 insertions(+), 9 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/optimizer.rs b/datafusion/core/src/physical_optimizer/optimizer.rs index 70ac12344810e..08cbf68fa6179 100644 --- a/datafusion/core/src/physical_optimizer/optimizer.rs +++ b/datafusion/core/src/physical_optimizer/optimizer.rs @@ -90,8 +90,6 @@ impl PhysicalOptimizer { // as that rule may inject other operations in between the different AggregateExecs. // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), - // Run once before PartialFinalAggregation is rewritten to ensure the rule is applied correctly - Arc::new(OptimizeAggregateOrder::new()), // 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 diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 69fb4ea38e9fb..3b044d980a91a 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3451,7 +3451,7 @@ logical_plan Aggregate: groupBy=[[]], aggr=[[FIRST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c3 DESC NULLS FIRST]]] --TableScan: convert_first_last_table projection=[c1, c3] physical_plan -AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +AggregateExec: mode=Final, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] --CoalescePartitionsExec ----AggregateExec: mode=Partial, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] ------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 @@ -3465,7 +3465,7 @@ logical_plan Aggregate: groupBy=[[]], aggr=[[LAST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c2 ASC NULLS LAST]]] --TableScan: convert_first_last_table projection=[c1, c2] physical_plan -AggregateExec: mode=Final, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] +AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] --CoalescePartitionsExec ----AggregateExec: mode=Partial, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] ------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 467bbca3c0bfb..0e5195388c839 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -245,7 +245,6 @@ OutputRequirementExec physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after OptimizeAggregateOrder 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 @@ -303,7 +302,6 @@ OutputRequirementExec, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]:),(Col[ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after OptimizeAggregateOrder 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 @@ -341,7 +339,6 @@ OutputRequirementExec physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE -physical_plan after OptimizeAggregateOrder 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/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 869462b4722a8..c15879c756365 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2805,7 +2805,7 @@ Projection: sales_global.country, FIRST_VALUE(sales_global.amount) ORDER BY [sal ------TableScan: sales_global projection=[country, ts, amount] physical_plan ProjectionExec: expr=[country@0 as country, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@1 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as lv1, SUM(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@3 as sum1] ---AggregateExec: mode=Single, gby=[country@0 as country], aggr=[LAST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount), SUM(sales_global.amount)] +--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount), SUM(sales_global.amount)] ----MemoryExec: partitions=1, partition_sizes=[1] query TRRR rowsort @@ -3800,7 +3800,7 @@ Projection: FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_tab ----TableScan: multiple_ordered_table projection=[a, c, d] physical_plan ProjectionExec: expr=[FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_table.a ASC NULLS LAST]@1 as first_a, LAST_VALUE(multiple_ordered_table.c) ORDER BY [multiple_ordered_table.c DESC NULLS FIRST]@2 as last_c] ---AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] +--AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), LAST_VALUE(multiple_ordered_table.c)] ----CoalesceBatchesExec: target_batch_size=2 ------RepartitionExec: partitioning=Hash([d@0], 8), input_partitions=8 --------AggregateExec: mode=Partial, gby=[d@2 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] From fffcd3be1492b12439c3d3f19d8c633849cb8be6 Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Mon, 15 Apr 2024 20:26:39 +0800 Subject: [PATCH 33/33] fix merge Signed-off-by: jayzhan211 --- .../sqllogictest/test_files/aggregate.slt | 28 +++++++++---------- .../sqllogictest/test_files/group_by.slt | 4 +-- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index e11b4ef370dce..3d24fe3888d77 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3475,25 +3475,25 @@ query TT explain select first_value(c1 order by c3 desc) from convert_first_last_table; ---- logical_plan -Aggregate: groupBy=[[]], aggr=[[FIRST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c3 DESC NULLS FIRST]]] ---TableScan: convert_first_last_table projection=[c1, c3] +01)Aggregate: groupBy=[[]], aggr=[[FIRST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c3 DESC NULLS FIRST]]] +02)--TableScan: convert_first_last_table projection=[c1, c3] physical_plan -AggregateExec: mode=Final, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] ---CoalescePartitionsExec -----AggregateExec: mode=Partial, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] -------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 ---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c3], output_orderings=[[c1@0 ASC NULLS LAST], [c3@1 ASC NULLS LAST]], has_header=true +01)AggregateExec: mode=Final, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c3], output_orderings=[[c1@0 ASC NULLS LAST], [c3@1 ASC NULLS LAST]], has_header=true # test last to first query TT explain select last_value(c1 order by c2 asc) from convert_first_last_table; ---- logical_plan -Aggregate: groupBy=[[]], aggr=[[LAST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c2 ASC NULLS LAST]]] ---TableScan: convert_first_last_table projection=[c1, c2] +01)Aggregate: groupBy=[[]], aggr=[[LAST_VALUE(convert_first_last_table.c1) ORDER BY [convert_first_last_table.c2 ASC NULLS LAST]]] +02)--TableScan: convert_first_last_table projection=[c1, c2] physical_plan -AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] ---CoalescePartitionsExec -----AggregateExec: mode=Partial, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] -------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 ---------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c2], output_orderings=[[c1@0 ASC NULLS LAST], [c2@1 DESC]], has_header=true +01)AggregateExec: mode=Final, gby=[], aggr=[LAST_VALUE(convert_first_last_table.c1)] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[], aggr=[FIRST_VALUE(convert_first_last_table.c1)] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/convert_first_last.csv]]}, projection=[c1, c2], output_orderings=[[c1@0 ASC NULLS LAST], [c2@1 DESC]], has_header=true diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 1acdcde9c8eea..5c5bf58dd0496 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2805,7 +2805,7 @@ logical_plan 04)------TableScan: sales_global projection=[country, ts, amount] physical_plan 01)ProjectionExec: expr=[country@0 as country, FIRST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@1 as fv1, LAST_VALUE(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as lv1, SUM(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@3 as sum1] -02)--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[LAST_VALUE(sales_global.amount), FIRST_VALUE(sales_global.amount), SUM(sales_global.amount)] +02)--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[FIRST_VALUE(sales_global.amount), LAST_VALUE(sales_global.amount), SUM(sales_global.amount)] 03)----MemoryExec: partitions=1, partition_sizes=[1] query TRRR rowsort @@ -3800,7 +3800,7 @@ logical_plan 03)----TableScan: multiple_ordered_table projection=[a, c, d] physical_plan 01)ProjectionExec: expr=[FIRST_VALUE(multiple_ordered_table.a) ORDER BY [multiple_ordered_table.a ASC NULLS LAST]@1 as first_a, LAST_VALUE(multiple_ordered_table.c) ORDER BY [multiple_ordered_table.c DESC NULLS FIRST]@2 as last_c] -02)--AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)] +02)--AggregateExec: mode=FinalPartitioned, gby=[d@0 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), LAST_VALUE(multiple_ordered_table.c)] 03)----CoalesceBatchesExec: target_batch_size=2 04)------RepartitionExec: partitioning=Hash([d@0], 8), input_partitions=8 05)--------AggregateExec: mode=Partial, gby=[d@2 as d], aggr=[FIRST_VALUE(multiple_ordered_table.a), FIRST_VALUE(multiple_ordered_table.c)]