From a99de23b9afd5b5c1ea386d87cc46d57a42ea3c1 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 10:51:56 +0300 Subject: [PATCH 01/31] Initial commit --- .../src/physical_optimizer/enforce_sorting.rs | 82 +++++ datafusion/expr/src/aggregate_function.rs | 9 +- .../expr/src/type_coercion/aggregates.rs | 1 + .../physical-expr/src/aggregate/build_in.rs | 7 + .../physical-expr/src/aggregate/first_last.rs | 293 ++++++++++++++++++ .../physical-expr/src/expressions/mod.rs | 2 +- 6 files changed, 391 insertions(+), 3 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index f609ddea66cff..a56bc1ed29d4d 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2305,3 +2305,85 @@ mod tests { Ok(()) } } + +#[cfg(test)] +mod tmp_tests { + use crate::assert_batches_eq; + use crate::physical_plan::{collect, displayable, ExecutionPlan}; + use crate::prelude::SessionContext; + use arrow::util::pretty::print_batches; + use datafusion_common::Result; + use datafusion_execution::config::SessionConfig; + use datafusion_physical_plan::get_plan_string; + use std::sync::Arc; + + fn print_plan(plan: &Arc) -> Result<()> { + let formatted = displayable(plan.as_ref()).indent(true).to_string(); + let actual: Vec<&str> = formatted.trim().lines().collect(); + println!("{:#?}", actual); + Ok(()) + } + + const MULTIPLE_ORDERED_TABLE: &str = "CREATE EXTERNAL TABLE multiple_ordered_table ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER + ) + STORED AS CSV + WITH HEADER ROW + WITH ORDER (a ASC, b ASC) + WITH ORDER (c ASC) + LOCATION '../core/tests/data/window_2.csv'"; + + #[tokio::test] + async fn test_query() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(1); + let ctx = SessionContext::new_with_config(config); + + ctx.sql(MULTIPLE_ORDERED_TABLE).await?; + + // let sql = "SELECT a, FIRST_VALUE(c) + // FROM multiple_ordered_table + // GROUP BY a"; + + let sql = "SELECT a, NTH_VALUE(c, 2) + FROM multiple_ordered_table + GROUP BY a"; + + let msg = format!("Creating logical plan for '{sql}'"); + let dataframe = ctx.sql(sql).await.expect(&msg); + let physical_plan = dataframe.create_physical_plan().await?; + print_plan(&physical_plan)?; + let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; + print_batches(&batches)?; + + let expected = vec![ + "ProjectionExec: expr=[NTH_VALUE(annotated_data_finite2.d,Int64(2)) ORDER BY [annotated_data_finite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as nth_value1]", + " GlobalLimitExec: skip=0, fetch=5", + " BoundedWindowAggExec: wdw=[NTH_VALUE(annotated_data_finite2.d,Int64(2)) ORDER BY [annotated_data_finite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Ok(Field { name: \"NTH_VALUE(annotated_data_finite2.d,Int64(2)) ORDER BY [annotated_data_finite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW\", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(Int32(NULL)), end_bound: CurrentRow }], mode=[Sorted]", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, d], output_ordering=[a@0 ASC NULLS LAST], has_header=true", + ]; + // Get string representation of the plan + let actual = get_plan_string(&physical_plan); + assert_eq!( + expected, actual, + "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" + ); + + let expected = [ + "+------------+", + "| nth_value1 |", + "+------------+", + "| 2 |", + "| 2 |", + "| 2 |", + "| 2 |", + "| 2 |", + "+------------+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) + } +} diff --git a/datafusion/expr/src/aggregate_function.rs b/datafusion/expr/src/aggregate_function.rs index cea72c3cb5e6b..df8a9017ac864 100644 --- a/datafusion/expr/src/aggregate_function.rs +++ b/datafusion/expr/src/aggregate_function.rs @@ -17,7 +17,7 @@ //! Aggregate function module contains all built-in aggregate functions definitions -use crate::utils; +use crate::{BuiltInWindowFunction, utils}; use crate::{type_coercion::aggregates::*, Signature, TypeSignature, Volatility}; use arrow::datatypes::{DataType, Field}; use datafusion_common::{plan_datafusion_err, plan_err, DataFusionError, Result}; @@ -50,6 +50,8 @@ pub enum AggregateFunction { FirstValue, /// last_value LastValue, + /// nth_value + NthValue, /// Variance (Sample) Variance, /// Variance (Population) @@ -118,6 +120,7 @@ impl AggregateFunction { ArrayAgg => "ARRAY_AGG", FirstValue => "FIRST_VALUE", LastValue => "LAST_VALUE", + NthValue => "NTH_VALUE", Variance => "VAR", VariancePop => "VAR_POP", Stddev => "STDDEV", @@ -174,6 +177,7 @@ impl FromStr for AggregateFunction { "array_agg" => AggregateFunction::ArrayAgg, "first_value" => AggregateFunction::FirstValue, "last_value" => AggregateFunction::LastValue, + "nth_value" => AggregateFunction::NthValue, "string_agg" => AggregateFunction::StringAgg, // statistical "corr" => AggregateFunction::Correlation, @@ -300,7 +304,7 @@ impl AggregateFunction { Ok(coerced_data_types[0].clone()) } AggregateFunction::Grouping => Ok(DataType::Int32), - AggregateFunction::FirstValue | AggregateFunction::LastValue => { + AggregateFunction::FirstValue | AggregateFunction::LastValue | AggregateFunction::NthValue => { Ok(coerced_data_types[0].clone()) } AggregateFunction::StringAgg => Ok(DataType::LargeUtf8), @@ -371,6 +375,7 @@ impl AggregateFunction { | AggregateFunction::LastValue => { Signature::uniform(1, NUMERICS.to_vec(), Volatility::Immutable) } + AggregateFunction::NthValue => Signature::any(2, Volatility::Immutable), AggregateFunction::Covariance | AggregateFunction::CovariancePop | AggregateFunction::Correlation diff --git a/datafusion/expr/src/type_coercion/aggregates.rs b/datafusion/expr/src/type_coercion/aggregates.rs index 7128b575978a3..917cd02554506 100644 --- a/datafusion/expr/src/type_coercion/aggregates.rs +++ b/datafusion/expr/src/type_coercion/aggregates.rs @@ -297,6 +297,7 @@ pub fn coerce_types( AggregateFunction::Median | AggregateFunction::FirstValue | AggregateFunction::LastValue => Ok(input_types.to_vec()), + AggregateFunction::NthValue => Ok(input_types.to_vec()), AggregateFunction::Grouping => Ok(vec![input_types[0].clone()]), AggregateFunction::StringAgg => { if !is_string_agg_supported_arg_type(&input_types[0]) { diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs index c40f0db194055..11a60758fc076 100644 --- a/datafusion/physical-expr/src/aggregate/build_in.rs +++ b/datafusion/physical-expr/src/aggregate/build_in.rs @@ -369,6 +369,13 @@ pub fn create_aggregate_expr( ordering_req.to_vec(), ordering_types, )), + (AggregateFunction::NthValue, _) => Arc::new(expressions::NthValueAgg::new( + input_phy_exprs[0].clone(), + name, + input_phy_types[0].clone(), + ordering_req.to_vec(), + ordering_types, + )), (AggregateFunction::StringAgg, false) => { if !ordering_req.is_empty() { return not_impl_err!( diff --git a/datafusion/physical-expr/src/aggregate/first_last.rs b/datafusion/physical-expr/src/aggregate/first_last.rs index 4afa8d0dd5eca..9ec26d8d1117a 100644 --- a/datafusion/physical-expr/src/aggregate/first_last.rs +++ b/datafusion/physical-expr/src/aggregate/first_last.rs @@ -667,6 +667,299 @@ impl Accumulator for LastValueAccumulator { } } +/// NTH_VALUE aggregate expression +#[derive(Debug, Clone)] +pub struct NthValueAgg { + name: String, + input_data_type: DataType, + order_by_data_types: Vec, + expr: Arc, + ordering_req: LexOrdering, + requirement_satisfied: bool, +} + +impl NthValueAgg { + /// Creates a new LAST_VALUE aggregation function. + pub fn new( + expr: Arc, + name: impl Into, + input_data_type: DataType, + ordering_req: LexOrdering, + order_by_data_types: Vec, + ) -> Self { + let requirement_satisfied = ordering_req.is_empty(); + Self { + name: name.into(), + input_data_type, + order_by_data_types, + expr, + ordering_req, + requirement_satisfied, + } + } + + /// Returns the name of the aggregate expression. + pub fn name(&self) -> &str { + &self.name + } + + /// Returns the input data type of the aggregate expression. + pub fn input_data_type(&self) -> &DataType { + &self.input_data_type + } + + /// Returns the data types of the order-by columns. + pub fn order_by_data_types(&self) -> &Vec { + &self.order_by_data_types + } + + /// Returns the expression associated with the aggregate function. + pub fn expr(&self) -> &Arc { + &self.expr + } + + /// Returns the lexical ordering requirements of the aggregate expression. + pub fn ordering_req(&self) -> &LexOrdering { + &self.ordering_req + } + + pub fn with_requirement_satisfied(mut self, requirement_satisfied: bool) -> Self { + self.requirement_satisfied = requirement_satisfied; + self + } +} + +impl AggregateExpr for NthValueAgg { + /// Return a reference to Any that can be used for downcasting + fn as_any(&self) -> &dyn Any { + self + } + + fn field(&self) -> Result { + Ok(Field::new(&self.name, self.input_data_type.clone(), true)) + } + + fn create_accumulator(&self) -> Result> { + crate::aggregate::first_last::LastValueAccumulator::try_new( + &self.input_data_type, + &self.order_by_data_types, + self.ordering_req.clone(), + ) + .map(|acc| { + Box::new(acc.with_requirement_satisfied(self.requirement_satisfied)) as _ + }) + } + + fn state_fields(&self) -> Result> { + let mut fields = vec![Field::new( + format_state_name(&self.name, "last_value"), + self.input_data_type.clone(), + true, + )]; + fields.extend(ordering_fields( + &self.ordering_req, + &self.order_by_data_types, + )); + fields.push(Field::new( + format_state_name(&self.name, "is_set"), + DataType::Boolean, + true, + )); + Ok(fields) + } + + fn expressions(&self) -> Vec> { + vec![self.expr.clone()] + } + + fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { + (!self.ordering_req.is_empty()).then_some(&self.ordering_req) + } + + fn name(&self) -> &str { + &self.name + } + + fn create_sliding_accumulator(&self) -> Result> { + crate::aggregate::first_last::LastValueAccumulator::try_new( + &self.input_data_type, + &self.order_by_data_types, + self.ordering_req.clone(), + ) + .map(|acc| { + Box::new(acc.with_requirement_satisfied(self.requirement_satisfied)) as _ + }) + } +} + +impl PartialEq for NthValueAgg { + fn eq(&self, other: &dyn Any) -> bool { + down_cast_any_ref(other) + .downcast_ref::() + .map(|x| { + self.name == x.name + && self.input_data_type == x.input_data_type + && self.order_by_data_types == x.order_by_data_types + && self.expr.eq(&x.expr) + }) + .unwrap_or(false) + } +} + +#[derive(Debug)] +struct NthValueAccumulator { + last: ScalarValue, + // The `is_set` flag keeps track of whether the last value is finalized. + // This information is used to discriminate genuine NULLs and NULLS that + // occur due to empty partitions. + is_set: bool, + orderings: Vec, + // Stores the applicable ordering requirement. + ordering_req: LexOrdering, + // Stores whether incoming data already satisfies the ordering requirement. + requirement_satisfied: bool, +} + +impl NthValueAccumulator { + /// Creates a new `NthValueAccumulator` for the given `data_type`. + pub fn try_new( + data_type: &DataType, + ordering_dtypes: &[DataType], + ordering_req: LexOrdering, + ) -> Result { + let orderings = ordering_dtypes + .iter() + .map(ScalarValue::try_from) + .collect::>>()?; + let requirement_satisfied = ordering_req.is_empty(); + ScalarValue::try_from(data_type).map(|last| Self { + last, + is_set: false, + orderings, + ordering_req, + requirement_satisfied, + }) + } + + // Updates state with the values in the given row. + fn update_with_new_row(&mut self, row: &[ScalarValue]) { + self.last = row[0].clone(); + self.orderings = row[1..].to_vec(); + self.is_set = true; + } + + fn get_last_idx(&self, values: &[ArrayRef]) -> Result> { + let [value, ordering_values @ ..] = values else { + return internal_err!("Empty row in LAST_VALUE"); + }; + if self.requirement_satisfied { + // Get last entry according to the order of data: + return Ok((!value.is_empty()).then_some(value.len() - 1)); + } + let sort_columns = ordering_values + .iter() + .zip(self.ordering_req.iter()) + .map(|(values, req)| { + // Take the reverse ordering requirement. This enables us to + // use "fetch = 1" to get the last value. + SortColumn { + values: values.clone(), + options: Some(!req.options), + } + }) + .collect::>(); + let indices = lexsort_to_indices(&sort_columns, Some(1))?; + Ok((!indices.is_empty()).then_some(indices.value(0) as _)) + } + + fn with_requirement_satisfied(mut self, requirement_satisfied: bool) -> Self { + self.requirement_satisfied = requirement_satisfied; + self + } +} + +impl Accumulator for NthValueAccumulator { + fn state(&self) -> Result> { + let mut result = vec![self.last.clone()]; + result.extend(self.orderings.clone()); + result.push(ScalarValue::Boolean(Some(self.is_set))); + Ok(result) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if !self.is_set || self.requirement_satisfied { + if let Some(last_idx) = self.get_last_idx(values)? { + let row = get_row_at_idx(values, last_idx)?; + self.update_with_new_row(&row); + } + } else if let Some(last_idx) = self.get_last_idx(values)? { + let row = get_row_at_idx(values, last_idx)?; + let orderings = &row[1..]; + // Update when there is a more recent entry + if compare_rows( + &self.orderings, + orderings, + &get_sort_options(&self.ordering_req), + )? + .is_lt() + { + self.update_with_new_row(&row); + } + } + + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + // LAST_VALUE(last1, last2, last3, ...) + // last index contains is_set flag. + let is_set_idx = states.len() - 1; + let flags = states[is_set_idx].as_boolean(); + let filtered_states = filter_states_according_to_is_set(states, flags)?; + // 1..is_set_idx range corresponds to ordering section + let sort_cols = + convert_to_sort_cols(&filtered_states[1..is_set_idx], &self.ordering_req); + + let ordered_states = if sort_cols.is_empty() { + // When no ordering is given, use existing state as is: + filtered_states + } else { + let indices = lexsort_to_indices(&sort_cols, None)?; + get_arrayref_at_indices(&filtered_states, &indices)? + }; + + if !ordered_states[0].is_empty() { + let last_idx = ordered_states[0].len() - 1; + let last_row = get_row_at_idx(&ordered_states, last_idx)?; + // When collecting orderings, we exclude the is_set flag from the state. + let last_ordering = &last_row[1..is_set_idx]; + let sort_options = get_sort_options(&self.ordering_req); + // Either there is no existing value, or there is a newer (latest) + // version in the new data: + if !self.is_set + || compare_rows(&self.orderings, last_ordering, &sort_options)?.is_lt() + { + // Update with last value in the state. Note that we should exclude the + // is_set flag from the state. Otherwise, we will end up with a state + // containing two is_set flags. + self.update_with_new_row(&last_row[0..is_set_idx]); + } + } + Ok(()) + } + + fn evaluate(&self) -> Result { + Ok(self.last.clone()) + } + + fn size(&self) -> usize { + std::mem::size_of_val(self) - std::mem::size_of_val(&self.last) + + self.last.size() + + ScalarValue::size_of_vec(&self.orderings) + - std::mem::size_of_val(&self.orderings) + } +} + /// Filters states according to the `is_set` flag at the last column and returns /// the resulting states. fn filter_states_according_to_is_set( diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index b6d0ad5b91043..2ae4a08026e9e 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -55,7 +55,7 @@ pub use crate::aggregate::correlation::Correlation; pub use crate::aggregate::count::Count; pub use crate::aggregate::count_distinct::DistinctCount; pub use crate::aggregate::covariance::{Covariance, CovariancePop}; -pub use crate::aggregate::first_last::{FirstValue, LastValue}; +pub use crate::aggregate::first_last::{FirstValue, LastValue, NthValueAgg}; pub use crate::aggregate::grouping::Grouping; pub use crate::aggregate::median::Median; pub use crate::aggregate::min_max::{Max, Min}; From da9adeeb81b764bfb113f49426d071ae91f91d27 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 10:52:48 +0300 Subject: [PATCH 02/31] minor changes --- datafusion/core/src/physical_optimizer/enforce_sorting.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index a56bc1ed29d4d..6c5ca3492b316 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2360,10 +2360,8 @@ mod tmp_tests { print_batches(&batches)?; let expected = vec![ - "ProjectionExec: expr=[NTH_VALUE(annotated_data_finite2.d,Int64(2)) ORDER BY [annotated_data_finite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as nth_value1]", - " GlobalLimitExec: skip=0, fetch=5", - " BoundedWindowAggExec: wdw=[NTH_VALUE(annotated_data_finite2.d,Int64(2)) ORDER BY [annotated_data_finite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Ok(Field { name: \"NTH_VALUE(annotated_data_finite2.d,Int64(2)) ORDER BY [annotated_data_finite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW\", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(Int32(NULL)), end_bound: CurrentRow }], mode=[Sorted]", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, d], output_ordering=[a@0 ASC NULLS LAST], has_header=true", + "AggregateExec: mode=Single, gby=[a@0 as a], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, c], output_orderings=[[a@0 ASC NULLS LAST], [c@1 ASC NULLS LAST]], has_header=true", ]; // Get string representation of the plan let actual = get_plan_string(&physical_plan); From 4c033ca9a157da15922fa14d239ee9b183f0a6b4 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 10:59:58 +0300 Subject: [PATCH 03/31] Parse index argument --- .../physical-expr/src/aggregate/build_in.rs | 24 ++++++++++++------- .../physical-expr/src/aggregate/first_last.rs | 3 +++ 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs index 11a60758fc076..9343a705d8521 100644 --- a/datafusion/physical-expr/src/aggregate/build_in.rs +++ b/datafusion/physical-expr/src/aggregate/build_in.rs @@ -29,9 +29,10 @@ use crate::aggregate::regr::RegrType; use crate::{expressions, AggregateExpr, PhysicalExpr, PhysicalSortExpr}; use arrow::datatypes::Schema; -use datafusion_common::{not_impl_err, DataFusionError, Result}; +use datafusion_common::{not_impl_err, DataFusionError, Result, plan_datafusion_err}; pub use datafusion_expr::AggregateFunction; use std::sync::Arc; +use crate::expressions::Literal; /// Create a physical aggregation expression. /// This function errors when `input_phy_exprs`' can't be coerced to a valid argument type of the aggregation function. @@ -369,13 +370,20 @@ pub fn create_aggregate_expr( ordering_req.to_vec(), ordering_types, )), - (AggregateFunction::NthValue, _) => Arc::new(expressions::NthValueAgg::new( - input_phy_exprs[0].clone(), - name, - input_phy_types[0].clone(), - ordering_req.to_vec(), - ordering_types, - )), + (AggregateFunction::NthValue, _) => { + let n = input_phy_exprs[1].as_any().downcast_ref::().unwrap().value(); + let n: i64 = n + .clone() + .try_into()?; + Arc::new(expressions::NthValueAgg::new( + input_phy_exprs[0].clone(), + n, + name, + input_phy_types[0].clone(), + ordering_req.to_vec(), + ordering_types, + )) + }, (AggregateFunction::StringAgg, false) => { if !ordering_req.is_empty() { return not_impl_err!( diff --git a/datafusion/physical-expr/src/aggregate/first_last.rs b/datafusion/physical-expr/src/aggregate/first_last.rs index 9ec26d8d1117a..5f54a8acb0050 100644 --- a/datafusion/physical-expr/src/aggregate/first_last.rs +++ b/datafusion/physical-expr/src/aggregate/first_last.rs @@ -674,6 +674,7 @@ pub struct NthValueAgg { input_data_type: DataType, order_by_data_types: Vec, expr: Arc, + n: i64, ordering_req: LexOrdering, requirement_satisfied: bool, } @@ -682,6 +683,7 @@ impl NthValueAgg { /// Creates a new LAST_VALUE aggregation function. pub fn new( expr: Arc, + n: i64, name: impl Into, input_data_type: DataType, ordering_req: LexOrdering, @@ -693,6 +695,7 @@ impl NthValueAgg { input_data_type, order_by_data_types, expr, + n, ordering_req, requirement_satisfied, } From 97019c0060bc710f11ee7b87a9edc3ca2690efc5 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 11:21:00 +0300 Subject: [PATCH 04/31] Move nth_value to array_agg --- .../src/aggregate/array_agg_ordered.rs | 287 +++++++++++++++++ .../physical-expr/src/aggregate/build_in.rs | 7 +- .../physical-expr/src/aggregate/first_last.rs | 296 ------------------ .../physical-expr/src/expressions/mod.rs | 4 +- 4 files changed, 294 insertions(+), 300 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index eb5ae8b0b0c3f..d6fb75dfd2216 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -323,6 +323,293 @@ impl OrderSensitiveArrayAggAccumulator { } } +/// Expression for a ARRAY_AGG(ORDER BY) aggregation. +/// When aggregation works in multiple partitions +/// aggregations are split into multiple partitions, +/// then their results are merged. This aggregator +/// is a version of ARRAY_AGG that can support producing +/// intermediate aggregation (with necessary side information) +/// and that can merge aggregations from multiple partitions. +#[derive(Debug)] +pub struct NthValueAgg { + /// Column name + name: String, + /// The DataType for the input expression + input_data_type: DataType, + /// The input expression + expr: Arc, + n: i64, + /// If the input expression can have NULLs + nullable: bool, + /// Ordering data types + order_by_data_types: Vec, + /// Ordering requirement + ordering_req: LexOrdering, +} + +impl NthValueAgg { + /// Create a new `OrderSensitiveArrayAgg` aggregate function + pub fn new( + expr: Arc, + n: i64, + name: impl Into, + input_data_type: DataType, + nullable: bool, + order_by_data_types: Vec, + ordering_req: LexOrdering, + ) -> Self { + Self { + name: name.into(), + input_data_type, + expr, + n, + nullable, + order_by_data_types, + ordering_req, + } + } +} + +impl AggregateExpr for NthValueAgg { + fn as_any(&self) -> &dyn Any { + self + } + + fn field(&self) -> Result { + Ok(Field::new_list( + &self.name, + // This should be the same as return type of AggregateFunction::ArrayAgg + Field::new("item", self.input_data_type.clone(), true), + self.nullable, + )) + } + + fn create_accumulator(&self) -> Result> { + Ok(Box::new(crate::aggregate::array_agg_ordered::NthValueAccumulator::try_new( + &self.input_data_type, + &self.order_by_data_types, + self.ordering_req.clone(), + )?)) + } + + fn state_fields(&self) -> Result> { + let mut fields = vec![Field::new_list( + format_state_name(&self.name, "array_agg"), + Field::new("item", self.input_data_type.clone(), true), + self.nullable, // This should be the same as field() + )]; + let orderings = ordering_fields(&self.ordering_req, &self.order_by_data_types); + fields.push(Field::new_list( + format_state_name(&self.name, "array_agg_orderings"), + Field::new("item", DataType::Struct(Fields::from(orderings)), true), + self.nullable, + )); + Ok(fields) + } + + fn expressions(&self) -> Vec> { + vec![self.expr.clone()] + } + + fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { + if self.ordering_req.is_empty() { + None + } else { + Some(&self.ordering_req) + } + } + + fn name(&self) -> &str { + &self.name + } +} + +impl PartialEq for NthValueAgg { + fn eq(&self, other: &dyn Any) -> bool { + down_cast_any_ref(other) + .downcast_ref::() + .map(|x| { + self.name == x.name + && self.input_data_type == x.input_data_type + && self.order_by_data_types == x.order_by_data_types + && self.expr.eq(&x.expr) + }) + .unwrap_or(false) + } +} + +#[derive(Debug)] +pub(crate) struct NthValueAccumulator { + // `values` stores entries in the ARRAY_AGG result. + values: Vec, + // `ordering_values` stores values of ordering requirement expression + // corresponding to each value in the ARRAY_AGG. + // For each `ScalarValue` inside `values`, there will be a corresponding + // `Vec` inside `ordering_values` which stores it ordering. + // This information is used during merging results of the different partitions. + // For detailed information how merging is done see [`merge_ordered_arrays`] + ordering_values: Vec>, + // `datatypes` stores, datatype of expression inside ARRAY_AGG and ordering requirement expressions. + datatypes: Vec, + // Stores ordering requirement of the Accumulator + ordering_req: LexOrdering, +} + +impl NthValueAccumulator { + /// Create a new order-sensitive ARRAY_AGG accumulator based on the given + /// item data type. + pub fn try_new( + datatype: &DataType, + ordering_dtypes: &[DataType], + ordering_req: LexOrdering, + ) -> Result { + let mut datatypes = vec![datatype.clone()]; + datatypes.extend(ordering_dtypes.iter().cloned()); + Ok(Self { + values: vec![], + ordering_values: vec![], + datatypes, + ordering_req, + }) + } +} + +impl Accumulator for NthValueAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.is_empty() { + return Ok(()); + } + + let n_row = values[0].len(); + for index in 0..n_row { + let row = get_row_at_idx(values, index)?; + self.values.push(row[0].clone()); + self.ordering_values.push(row[1..].to_vec()); + } + + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + if states.is_empty() { + return Ok(()); + } + // First entry in the state is the aggregation result. + let array_agg_values = &states[0]; + // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. + // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores + // values received from its ordering requirement expression. (This information is necessary for during merging). + let agg_orderings = &states[1]; + + if let Some(agg_orderings) = agg_orderings.as_list_opt::() { + // Stores ARRAY_AGG results coming from each partition + let mut partition_values = vec![]; + // Stores ordering requirement expression results coming from each partition + let mut partition_ordering_values = vec![]; + + // Existing values should be merged also. + partition_values.push(self.values.clone()); + partition_ordering_values.push(self.ordering_values.clone()); + + let array_agg_res = + ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; + + for v in array_agg_res.into_iter() { + partition_values.push(v); + } + + let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; + + for partition_ordering_rows in orderings.into_iter() { + // Extract value from struct to ordering_rows for each group/partition + let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { + if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { + Ok(ordering_columns_per_row) + } else { + exec_err!( + "Expects to receive ScalarValue::Struct(Some(..), _) but got:{:?}", + ordering_row.data_type() + ) + } + }).collect::>>()?; + + partition_ordering_values.push(ordering_value); + } + + let sort_options = self + .ordering_req + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(); + let (new_values, new_orderings) = merge_ordered_arrays( + &partition_values, + &partition_ordering_values, + &sort_options, + )?; + self.values = new_values; + self.ordering_values = new_orderings; + } else { + return exec_err!("Expects to receive a list array"); + } + Ok(()) + } + + fn state(&self) -> Result> { + let mut result = vec![self.evaluate()?]; + result.push(self.evaluate_orderings()?); + Ok(result) + } + + fn evaluate(&self) -> Result { + let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); + Ok(ScalarValue::List(arr)) + } + + fn size(&self) -> usize { + let mut total = std::mem::size_of_val(self) + + ScalarValue::size_of_vec(&self.values) + - std::mem::size_of_val(&self.values); + + // Add size of the `self.ordering_values` + total += + std::mem::size_of::>() * self.ordering_values.capacity(); + for row in &self.ordering_values { + total += ScalarValue::size_of_vec(row) - std::mem::size_of_val(row); + } + + // Add size of the `self.datatypes` + total += std::mem::size_of::() * self.datatypes.capacity(); + for dtype in &self.datatypes { + total += dtype.size() - std::mem::size_of_val(dtype); + } + + // Add size of the `self.ordering_req` + total += std::mem::size_of::() * self.ordering_req.capacity(); + // TODO: Calculate size of each `PhysicalSortExpr` more accurately. + total + } +} + +impl NthValueAccumulator { + fn evaluate_orderings(&self) -> Result { + let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]); + let struct_field = Fields::from(fields.clone()); + + let orderings: Vec = self + .ordering_values + .iter() + .map(|ordering| { + ScalarValue::Struct(Some(ordering.clone()), struct_field.clone()) + }) + .collect(); + let struct_type = DataType::Struct(Fields::from(fields)); + + // Wrap in List, so we have the same data structure ListArray(StructArray..) for group by cases + let arr = ScalarValue::new_list(&orderings, &struct_type); + Ok(ScalarValue::List(arr)) + } +} + /// This is a wrapper struct to be able to correctly merge ARRAY_AGG /// data from multiple partitions using `BinaryHeap`. /// When used inside `BinaryHeap` this struct returns smallest `CustomElement`, diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs index 9343a705d8521..554d0dde9f11b 100644 --- a/datafusion/physical-expr/src/aggregate/build_in.rs +++ b/datafusion/physical-expr/src/aggregate/build_in.rs @@ -371,17 +371,20 @@ pub fn create_aggregate_expr( ordering_types, )), (AggregateFunction::NthValue, _) => { + let expr = &input_phy_exprs[0]; let n = input_phy_exprs[1].as_any().downcast_ref::().unwrap().value(); let n: i64 = n .clone() .try_into()?; + let nullable = expr.nullable(input_schema)?; Arc::new(expressions::NthValueAgg::new( - input_phy_exprs[0].clone(), + expr.clone(), n, name, input_phy_types[0].clone(), - ordering_req.to_vec(), + nullable, ordering_types, + ordering_req.to_vec(), )) }, (AggregateFunction::StringAgg, false) => { diff --git a/datafusion/physical-expr/src/aggregate/first_last.rs b/datafusion/physical-expr/src/aggregate/first_last.rs index 5f54a8acb0050..4afa8d0dd5eca 100644 --- a/datafusion/physical-expr/src/aggregate/first_last.rs +++ b/datafusion/physical-expr/src/aggregate/first_last.rs @@ -667,302 +667,6 @@ impl Accumulator for LastValueAccumulator { } } -/// NTH_VALUE aggregate expression -#[derive(Debug, Clone)] -pub struct NthValueAgg { - name: String, - input_data_type: DataType, - order_by_data_types: Vec, - expr: Arc, - n: i64, - ordering_req: LexOrdering, - requirement_satisfied: bool, -} - -impl NthValueAgg { - /// Creates a new LAST_VALUE aggregation function. - pub fn new( - expr: Arc, - n: i64, - name: impl Into, - input_data_type: DataType, - ordering_req: LexOrdering, - order_by_data_types: Vec, - ) -> Self { - let requirement_satisfied = ordering_req.is_empty(); - Self { - name: name.into(), - input_data_type, - order_by_data_types, - expr, - n, - ordering_req, - requirement_satisfied, - } - } - - /// Returns the name of the aggregate expression. - pub fn name(&self) -> &str { - &self.name - } - - /// Returns the input data type of the aggregate expression. - pub fn input_data_type(&self) -> &DataType { - &self.input_data_type - } - - /// Returns the data types of the order-by columns. - pub fn order_by_data_types(&self) -> &Vec { - &self.order_by_data_types - } - - /// Returns the expression associated with the aggregate function. - pub fn expr(&self) -> &Arc { - &self.expr - } - - /// Returns the lexical ordering requirements of the aggregate expression. - pub fn ordering_req(&self) -> &LexOrdering { - &self.ordering_req - } - - pub fn with_requirement_satisfied(mut self, requirement_satisfied: bool) -> Self { - self.requirement_satisfied = requirement_satisfied; - self - } -} - -impl AggregateExpr for NthValueAgg { - /// Return a reference to Any that can be used for downcasting - fn as_any(&self) -> &dyn Any { - self - } - - fn field(&self) -> Result { - Ok(Field::new(&self.name, self.input_data_type.clone(), true)) - } - - fn create_accumulator(&self) -> Result> { - crate::aggregate::first_last::LastValueAccumulator::try_new( - &self.input_data_type, - &self.order_by_data_types, - self.ordering_req.clone(), - ) - .map(|acc| { - Box::new(acc.with_requirement_satisfied(self.requirement_satisfied)) as _ - }) - } - - fn state_fields(&self) -> Result> { - let mut fields = vec![Field::new( - format_state_name(&self.name, "last_value"), - self.input_data_type.clone(), - true, - )]; - fields.extend(ordering_fields( - &self.ordering_req, - &self.order_by_data_types, - )); - fields.push(Field::new( - format_state_name(&self.name, "is_set"), - DataType::Boolean, - true, - )); - Ok(fields) - } - - fn expressions(&self) -> Vec> { - vec![self.expr.clone()] - } - - fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { - (!self.ordering_req.is_empty()).then_some(&self.ordering_req) - } - - fn name(&self) -> &str { - &self.name - } - - fn create_sliding_accumulator(&self) -> Result> { - crate::aggregate::first_last::LastValueAccumulator::try_new( - &self.input_data_type, - &self.order_by_data_types, - self.ordering_req.clone(), - ) - .map(|acc| { - Box::new(acc.with_requirement_satisfied(self.requirement_satisfied)) as _ - }) - } -} - -impl PartialEq for NthValueAgg { - fn eq(&self, other: &dyn Any) -> bool { - down_cast_any_ref(other) - .downcast_ref::() - .map(|x| { - self.name == x.name - && self.input_data_type == x.input_data_type - && self.order_by_data_types == x.order_by_data_types - && self.expr.eq(&x.expr) - }) - .unwrap_or(false) - } -} - -#[derive(Debug)] -struct NthValueAccumulator { - last: ScalarValue, - // The `is_set` flag keeps track of whether the last value is finalized. - // This information is used to discriminate genuine NULLs and NULLS that - // occur due to empty partitions. - is_set: bool, - orderings: Vec, - // Stores the applicable ordering requirement. - ordering_req: LexOrdering, - // Stores whether incoming data already satisfies the ordering requirement. - requirement_satisfied: bool, -} - -impl NthValueAccumulator { - /// Creates a new `NthValueAccumulator` for the given `data_type`. - pub fn try_new( - data_type: &DataType, - ordering_dtypes: &[DataType], - ordering_req: LexOrdering, - ) -> Result { - let orderings = ordering_dtypes - .iter() - .map(ScalarValue::try_from) - .collect::>>()?; - let requirement_satisfied = ordering_req.is_empty(); - ScalarValue::try_from(data_type).map(|last| Self { - last, - is_set: false, - orderings, - ordering_req, - requirement_satisfied, - }) - } - - // Updates state with the values in the given row. - fn update_with_new_row(&mut self, row: &[ScalarValue]) { - self.last = row[0].clone(); - self.orderings = row[1..].to_vec(); - self.is_set = true; - } - - fn get_last_idx(&self, values: &[ArrayRef]) -> Result> { - let [value, ordering_values @ ..] = values else { - return internal_err!("Empty row in LAST_VALUE"); - }; - if self.requirement_satisfied { - // Get last entry according to the order of data: - return Ok((!value.is_empty()).then_some(value.len() - 1)); - } - let sort_columns = ordering_values - .iter() - .zip(self.ordering_req.iter()) - .map(|(values, req)| { - // Take the reverse ordering requirement. This enables us to - // use "fetch = 1" to get the last value. - SortColumn { - values: values.clone(), - options: Some(!req.options), - } - }) - .collect::>(); - let indices = lexsort_to_indices(&sort_columns, Some(1))?; - Ok((!indices.is_empty()).then_some(indices.value(0) as _)) - } - - fn with_requirement_satisfied(mut self, requirement_satisfied: bool) -> Self { - self.requirement_satisfied = requirement_satisfied; - self - } -} - -impl Accumulator for NthValueAccumulator { - fn state(&self) -> Result> { - let mut result = vec![self.last.clone()]; - result.extend(self.orderings.clone()); - result.push(ScalarValue::Boolean(Some(self.is_set))); - Ok(result) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - if !self.is_set || self.requirement_satisfied { - if let Some(last_idx) = self.get_last_idx(values)? { - let row = get_row_at_idx(values, last_idx)?; - self.update_with_new_row(&row); - } - } else if let Some(last_idx) = self.get_last_idx(values)? { - let row = get_row_at_idx(values, last_idx)?; - let orderings = &row[1..]; - // Update when there is a more recent entry - if compare_rows( - &self.orderings, - orderings, - &get_sort_options(&self.ordering_req), - )? - .is_lt() - { - self.update_with_new_row(&row); - } - } - - Ok(()) - } - - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - // LAST_VALUE(last1, last2, last3, ...) - // last index contains is_set flag. - let is_set_idx = states.len() - 1; - let flags = states[is_set_idx].as_boolean(); - let filtered_states = filter_states_according_to_is_set(states, flags)?; - // 1..is_set_idx range corresponds to ordering section - let sort_cols = - convert_to_sort_cols(&filtered_states[1..is_set_idx], &self.ordering_req); - - let ordered_states = if sort_cols.is_empty() { - // When no ordering is given, use existing state as is: - filtered_states - } else { - let indices = lexsort_to_indices(&sort_cols, None)?; - get_arrayref_at_indices(&filtered_states, &indices)? - }; - - if !ordered_states[0].is_empty() { - let last_idx = ordered_states[0].len() - 1; - let last_row = get_row_at_idx(&ordered_states, last_idx)?; - // When collecting orderings, we exclude the is_set flag from the state. - let last_ordering = &last_row[1..is_set_idx]; - let sort_options = get_sort_options(&self.ordering_req); - // Either there is no existing value, or there is a newer (latest) - // version in the new data: - if !self.is_set - || compare_rows(&self.orderings, last_ordering, &sort_options)?.is_lt() - { - // Update with last value in the state. Note that we should exclude the - // is_set flag from the state. Otherwise, we will end up with a state - // containing two is_set flags. - self.update_with_new_row(&last_row[0..is_set_idx]); - } - } - Ok(()) - } - - fn evaluate(&self) -> Result { - Ok(self.last.clone()) - } - - fn size(&self) -> usize { - std::mem::size_of_val(self) - std::mem::size_of_val(&self.last) - + self.last.size() - + ScalarValue::size_of_vec(&self.orderings) - - std::mem::size_of_val(&self.orderings) - } -} - /// Filters states according to the `is_set` flag at the last column and returns /// the resulting states. fn filter_states_according_to_is_set( diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 2ae4a08026e9e..a436321b3836d 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -46,7 +46,7 @@ pub use crate::aggregate::approx_percentile_cont::ApproxPercentileCont; pub use crate::aggregate::approx_percentile_cont_with_weight::ApproxPercentileContWithWeight; pub use crate::aggregate::array_agg::ArrayAgg; pub use crate::aggregate::array_agg_distinct::DistinctArrayAgg; -pub use crate::aggregate::array_agg_ordered::OrderSensitiveArrayAgg; +pub use crate::aggregate::array_agg_ordered::{OrderSensitiveArrayAgg, NthValueAgg}; pub use crate::aggregate::average::{Avg, AvgAccumulator}; pub use crate::aggregate::bit_and_or_xor::{BitAnd, BitOr, BitXor, DistinctBitXor}; pub use crate::aggregate::bool_and_or::{BoolAnd, BoolOr}; @@ -55,7 +55,7 @@ pub use crate::aggregate::correlation::Correlation; pub use crate::aggregate::count::Count; pub use crate::aggregate::count_distinct::DistinctCount; pub use crate::aggregate::covariance::{Covariance, CovariancePop}; -pub use crate::aggregate::first_last::{FirstValue, LastValue, NthValueAgg}; +pub use crate::aggregate::first_last::{FirstValue, LastValue}; pub use crate::aggregate::grouping::Grouping; pub use crate::aggregate::median::Median; pub use crate::aggregate::min_max::{Max, Min}; From b5672fad5610efaeea0649e6ce287fae88ec1d1d Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 11:36:40 +0300 Subject: [PATCH 05/31] Initial implementation (with redundant data) --- .../src/physical_optimizer/enforce_sorting.rs | 25 ++++++----- datafusion/expr/src/aggregate_function.rs | 8 ++-- .../src/aggregate/array_agg_ordered.rs | 43 +++++++++++++------ .../physical-expr/src/aggregate/build_in.rs | 16 ++++--- .../physical-expr/src/expressions/mod.rs | 2 +- 5 files changed, 56 insertions(+), 38 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index 6c5ca3492b316..6b4bc487d17a4 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2348,9 +2348,9 @@ mod tmp_tests { // FROM multiple_ordered_table // GROUP BY a"; - let sql = "SELECT a, NTH_VALUE(c, 2) + let sql = "SELECT a, b, NTH_VALUE(c, 2), ARRAY_AGG(c) FROM multiple_ordered_table - GROUP BY a"; + GROUP BY a, b"; let msg = format!("Creating logical plan for '{sql}'"); let dataframe = ctx.sql(sql).await.expect(&msg); @@ -2360,8 +2360,8 @@ mod tmp_tests { print_batches(&batches)?; let expected = vec![ - "AggregateExec: mode=Single, gby=[a@0 as a], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, c], output_orderings=[[a@0 ASC NULLS LAST], [c@1 ASC NULLS LAST]], has_header=true", + "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2)), ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", ]; // Get string representation of the plan let actual = get_plan_string(&physical_plan); @@ -2371,15 +2371,14 @@ mod tmp_tests { ); let expected = [ - "+------------+", - "| nth_value1 |", - "+------------+", - "| 2 |", - "| 2 |", - "| 2 |", - "| 2 |", - "| 2 |", - "+------------+", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) | ARRAY_AGG(multiple_ordered_table.c) |", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "| 0 | 0 | 1 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] |", + "| 0 | 1 | 26 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] |", + "| 1 | 2 | 51 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] |", + "| 1 | 3 | 76 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] |", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", ]; assert_batches_eq!(expected, &batches); Ok(()) diff --git a/datafusion/expr/src/aggregate_function.rs b/datafusion/expr/src/aggregate_function.rs index df8a9017ac864..c5e48a92cf9ec 100644 --- a/datafusion/expr/src/aggregate_function.rs +++ b/datafusion/expr/src/aggregate_function.rs @@ -17,8 +17,8 @@ //! Aggregate function module contains all built-in aggregate functions definitions -use crate::{BuiltInWindowFunction, utils}; use crate::{type_coercion::aggregates::*, Signature, TypeSignature, Volatility}; +use crate::{utils, BuiltInWindowFunction}; use arrow::datatypes::{DataType, Field}; use datafusion_common::{plan_datafusion_err, plan_err, DataFusionError, Result}; use std::sync::Arc; @@ -304,9 +304,9 @@ impl AggregateFunction { Ok(coerced_data_types[0].clone()) } AggregateFunction::Grouping => Ok(DataType::Int32), - AggregateFunction::FirstValue | AggregateFunction::LastValue | AggregateFunction::NthValue => { - Ok(coerced_data_types[0].clone()) - } + AggregateFunction::FirstValue + | AggregateFunction::LastValue + | AggregateFunction::NthValue => Ok(coerced_data_types[0].clone()), AggregateFunction::StringAgg => Ok(DataType::LargeUtf8), } } diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index d6fb75dfd2216..b5b899cea9daa 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -376,20 +376,18 @@ impl AggregateExpr for NthValueAgg { } fn field(&self) -> Result { - Ok(Field::new_list( - &self.name, - // This should be the same as return type of AggregateFunction::ArrayAgg - Field::new("item", self.input_data_type.clone(), true), - self.nullable, - )) + Ok(Field::new(&self.name, self.input_data_type.clone(), true)) } fn create_accumulator(&self) -> Result> { - Ok(Box::new(crate::aggregate::array_agg_ordered::NthValueAccumulator::try_new( - &self.input_data_type, - &self.order_by_data_types, - self.ordering_req.clone(), - )?)) + Ok(Box::new( + crate::aggregate::array_agg_ordered::NthValueAccumulator::try_new( + self.n, + &self.input_data_type, + &self.order_by_data_types, + self.ordering_req.clone(), + )?, + )) } fn state_fields(&self) -> Result> { @@ -440,6 +438,7 @@ impl PartialEq for NthValueAgg { #[derive(Debug)] pub(crate) struct NthValueAccumulator { + n: i64, // `values` stores entries in the ARRAY_AGG result. values: Vec, // `ordering_values` stores values of ordering requirement expression @@ -459,6 +458,7 @@ impl NthValueAccumulator { /// Create a new order-sensitive ARRAY_AGG accumulator based on the given /// item data type. pub fn try_new( + n: i64, datatype: &DataType, ordering_dtypes: &[DataType], ordering_req: LexOrdering, @@ -466,6 +466,7 @@ impl NthValueAccumulator { let mut datatypes = vec![datatype.clone()]; datatypes.extend(ordering_dtypes.iter().cloned()); Ok(Self { + n, values: vec![], ordering_values: vec![], datatypes, @@ -561,8 +562,24 @@ impl Accumulator for NthValueAccumulator { } fn evaluate(&self) -> Result { - let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); - Ok(ScalarValue::List(arr)) + if self.n > 0 { + let forward_idx = self.n as usize - 1; + if let Some(nth_value) = self.values.get(forward_idx) { + Ok(nth_value.clone()) + } else { + ScalarValue::try_from(self.datatypes[0].clone()) + } + } else if self.n < 0 { + let backward_idx = (-self.n) as usize - 1; + if let Some(nth_value) = self.values.get(backward_idx) { + Ok(nth_value.clone()) + } else { + ScalarValue::try_from(self.datatypes[0].clone()) + } + } else { + // n cannot be 0 + unreachable!(); + } } fn size(&self) -> usize { diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs index 554d0dde9f11b..39f729f376277 100644 --- a/datafusion/physical-expr/src/aggregate/build_in.rs +++ b/datafusion/physical-expr/src/aggregate/build_in.rs @@ -27,12 +27,12 @@ //! * Return type: a function `(arg_types) -> return_type`. E.g. for min, ([f32]) -> f32, ([f64]) -> f64. use crate::aggregate::regr::RegrType; +use crate::expressions::Literal; use crate::{expressions, AggregateExpr, PhysicalExpr, PhysicalSortExpr}; use arrow::datatypes::Schema; -use datafusion_common::{not_impl_err, DataFusionError, Result, plan_datafusion_err}; +use datafusion_common::{not_impl_err, plan_datafusion_err, DataFusionError, Result}; pub use datafusion_expr::AggregateFunction; use std::sync::Arc; -use crate::expressions::Literal; /// Create a physical aggregation expression. /// This function errors when `input_phy_exprs`' can't be coerced to a valid argument type of the aggregation function. @@ -372,10 +372,12 @@ pub fn create_aggregate_expr( )), (AggregateFunction::NthValue, _) => { let expr = &input_phy_exprs[0]; - let n = input_phy_exprs[1].as_any().downcast_ref::().unwrap().value(); - let n: i64 = n - .clone() - .try_into()?; + let n = input_phy_exprs[1] + .as_any() + .downcast_ref::() + .unwrap() + .value(); + let n: i64 = n.clone().try_into()?; let nullable = expr.nullable(input_schema)?; Arc::new(expressions::NthValueAgg::new( expr.clone(), @@ -386,7 +388,7 @@ pub fn create_aggregate_expr( ordering_types, ordering_req.to_vec(), )) - }, + } (AggregateFunction::StringAgg, false) => { if !ordering_req.is_empty() { return not_impl_err!( diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index a436321b3836d..d7bb33d5457ce 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -46,7 +46,7 @@ pub use crate::aggregate::approx_percentile_cont::ApproxPercentileCont; pub use crate::aggregate::approx_percentile_cont_with_weight::ApproxPercentileContWithWeight; pub use crate::aggregate::array_agg::ArrayAgg; pub use crate::aggregate::array_agg_distinct::DistinctArrayAgg; -pub use crate::aggregate::array_agg_ordered::{OrderSensitiveArrayAgg, NthValueAgg}; +pub use crate::aggregate::array_agg_ordered::{NthValueAgg, OrderSensitiveArrayAgg}; pub use crate::aggregate::average::{Avg, AvgAccumulator}; pub use crate::aggregate::bit_and_or_xor::{BitAnd, BitOr, BitXor, DistinctBitXor}; pub use crate::aggregate::bool_and_or::{BoolAnd, BoolOr}; From 73ec593cb5ad7c8b16b675d06e21d517ef13968e Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 13:09:33 +0300 Subject: [PATCH 06/31] Add new test --- .../src/physical_optimizer/enforce_sorting.rs | 105 ++++++++++++++++-- datafusion/physical-expr/src/aggregate/mod.rs | 3 +- .../physical-plan/src/aggregates/mod.rs | 15 +++ 3 files changed, 113 insertions(+), 10 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index 6b4bc487d17a4..d2d2db9b066d4 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2337,23 +2337,110 @@ mod tmp_tests { WITH ORDER (c ASC) LOCATION '../core/tests/data/window_2.csv'"; + const QUERY1: &str = "SELECT a, b, NTH_VALUE(c, 2), ARRAY_AGG(c) + FROM multiple_ordered_table + GROUP BY a, b"; + + const QUERY2: &str = + "SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), ARRAY_AGG(c ORDER BY c ASC) + FROM multiple_ordered_table + GROUP BY a, b"; + + const QUERY3: &str = + "SELECT a, b, NTH_VALUE(c, 2 ORDER BY c DESC), ARRAY_AGG(c ORDER BY c ASC) + FROM multiple_ordered_table + GROUP BY a, b"; + #[tokio::test] - async fn test_query() -> Result<()> { + async fn test_query1() -> Result<()> { let config = SessionConfig::new().with_target_partitions(1); let ctx = SessionContext::new_with_config(config); + let query = QUERY1; ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - // let sql = "SELECT a, FIRST_VALUE(c) - // FROM multiple_ordered_table - // GROUP BY a"; + let msg = format!("Creating logical plan for '{query}'"); + let dataframe = ctx.sql(query).await.expect(&msg); + let physical_plan = dataframe.create_physical_plan().await?; + print_plan(&physical_plan)?; + let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; + print_batches(&batches)?; - let sql = "SELECT a, b, NTH_VALUE(c, 2), ARRAY_AGG(c) - FROM multiple_ordered_table - GROUP BY a, b"; + let expected = vec![ + "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2)), ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", + ]; + // Get string representation of the plan + let actual = get_plan_string(&physical_plan); + assert_eq!( + expected, actual, + "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" + ); + + let expected = [ + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) | ARRAY_AGG(multiple_ordered_table.c) |", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "| 0 | 0 | 1 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] |", + "| 0 | 1 | 26 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] |", + "| 1 | 2 | 51 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] |", + "| 1 | 3 | 76 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] |", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) + } + + #[tokio::test] + async fn test_query2() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(1); + let ctx = SessionContext::new_with_config(config); + let query = QUERY2; + + ctx.sql(MULTIPLE_ORDERED_TABLE).await?; + + let msg = format!("Creating logical plan for '{query}'"); + let dataframe = ctx.sql(query).await.expect(&msg); + let physical_plan = dataframe.create_physical_plan().await?; + print_plan(&physical_plan)?; + let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; + print_batches(&batches)?; + + let expected = vec![ + "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2)), ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", + ]; + // Get string representation of the plan + let actual = get_plan_string(&physical_plan); + assert_eq!( + expected, actual, + "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" + ); + + let expected = [ + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) | ARRAY_AGG(multiple_ordered_table.c) |", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "| 0 | 0 | 1 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] |", + "| 0 | 1 | 26 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] |", + "| 1 | 2 | 51 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] |", + "| 1 | 3 | 76 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] |", + "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) + } + + #[tokio::test] + async fn test_query3() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(1); + let ctx = SessionContext::new_with_config(config); + let query = QUERY3; + + ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - let msg = format!("Creating logical plan for '{sql}'"); - let dataframe = ctx.sql(sql).await.expect(&msg); + let msg = format!("Creating logical plan for '{query}'"); + let dataframe = ctx.sql(query).await.expect(&msg); let physical_plan = dataframe.create_physical_plan().await?; print_plan(&physical_plan)?; let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; diff --git a/datafusion/physical-expr/src/aggregate/mod.rs b/datafusion/physical-expr/src/aggregate/mod.rs index 5bd1fca385b11..1e265c65e754a 100644 --- a/datafusion/physical-expr/src/aggregate/mod.rs +++ b/datafusion/physical-expr/src/aggregate/mod.rs @@ -20,7 +20,7 @@ use std::fmt::Debug; use std::sync::Arc; use self::groups_accumulator::GroupsAccumulator; -use crate::expressions::OrderSensitiveArrayAgg; +use crate::expressions::{NthValueAgg, OrderSensitiveArrayAgg}; use crate::{PhysicalExpr, PhysicalSortExpr}; use arrow::datatypes::Field; @@ -140,4 +140,5 @@ pub trait AggregateExpr: Send + Sync + Debug + PartialEq { /// However, an `ARRAY_AGG` with `ORDER BY` depends on the input ordering. pub fn is_order_sensitive(aggr_expr: &Arc) -> bool { aggr_expr.as_any().is::() + || aggr_expr.as_any().is::() } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0b94dd01cfd40..5eb323cc0fa3e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -952,6 +952,21 @@ fn get_aggregate_exprs_requirement( finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) { requirement = finer_ordering; + } else 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, + ) { + *aggr_expr = reverse_aggr_expr; + requirement = finer_ordering; + } else { + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ); + } } else { // If neither of the requirements satisfy the other, this means // requirements are conflicting. Currently, we do not support From b04fdefbf733e689ebff24264936df38e766e7ae Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 13:20:54 +0300 Subject: [PATCH 07/31] Add reverse support --- .../src/physical_optimizer/enforce_sorting.rs | 20 ++++++++-------- .../src/aggregate/array_agg_ordered.rs | 23 +++++++++++++++---- .../physical-plan/src/aggregates/mod.rs | 3 ++- 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index d2d2db9b066d4..c40671859a185 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2347,7 +2347,7 @@ mod tmp_tests { GROUP BY a, b"; const QUERY3: &str = - "SELECT a, b, NTH_VALUE(c, 2 ORDER BY c DESC), ARRAY_AGG(c ORDER BY c ASC) + "SELECT a, b, ARRAY_AGG(c ORDER BY c ASC), NTH_VALUE(c, 2 ORDER BY c DESC) FROM multiple_ordered_table GROUP BY a, b"; @@ -2447,7 +2447,7 @@ mod tmp_tests { print_batches(&batches)?; let expected = vec![ - "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2)), ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted", + "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[ARRAY_AGG(multiple_ordered_table.c), NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", ]; // Get string representation of the plan @@ -2458,14 +2458,14 @@ mod tmp_tests { ); let expected = [ - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) | ARRAY_AGG(multiple_ordered_table.c) |", - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - "| 0 | 0 | 1 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] |", - "| 0 | 1 | 26 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] |", - "| 1 | 2 | 51 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] |", - "| 1 | 3 | 76 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] |", - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", + "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", + "| a | b | ARRAY_AGG(multiple_ordered_table.c) | NTH_VALUE(multiple_ordered_table.c,Int64(2)) |", + "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", + "| 0 | 0 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] | 23 |", + "| 0 | 1 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] | 48 |", + "| 1 | 2 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] | 73 |", + "| 1 | 3 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] | 98 |", + "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", ]; assert_batches_eq!(expected, &batches); Ok(()) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index b5b899cea9daa..388df723db513 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -26,7 +26,9 @@ use std::sync::Arc; use crate::aggregate::utils::{down_cast_any_ref, ordering_fields}; use crate::expressions::format_state_name; -use crate::{AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr}; +use crate::{ + reverse_order_bys, AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr, +}; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field}; @@ -420,6 +422,18 @@ impl AggregateExpr for NthValueAgg { fn name(&self) -> &str { &self.name } + + fn reverse_expr(&self) -> Option> { + Some(Arc::new(Self { + name: self.name.to_string(), + input_data_type: self.input_data_type.clone(), + expr: self.expr.clone(), + n: -self.n, + nullable: self.nullable, + order_by_data_types: self.order_by_data_types.clone(), + ordering_req: reverse_order_bys(&self.ordering_req), + }) as _) + } } impl PartialEq for NthValueAgg { @@ -570,9 +584,10 @@ impl Accumulator for NthValueAccumulator { ScalarValue::try_from(self.datatypes[0].clone()) } } else if self.n < 0 { - let backward_idx = (-self.n) as usize - 1; - if let Some(nth_value) = self.values.get(backward_idx) { - Ok(nth_value.clone()) + let backward_idx = (-self.n) as usize; + if self.values.len() >= backward_idx { + let idx = self.values.len() - backward_idx; + Ok(self.values[idx].clone()) } else { ScalarValue::try_from(self.datatypes[0].clone()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 5eb323cc0fa3e..9ea1cc0c63772 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -953,6 +953,7 @@ fn get_aggregate_exprs_requirement( { requirement = finer_ordering; } else if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + println!("reverse_aggr_expr:{:?}", reverse_aggr_expr); if let Some(finer_ordering) = finer_ordering( &requirement, &reverse_aggr_expr, @@ -964,7 +965,7 @@ fn get_aggregate_exprs_requirement( requirement = finer_ordering; } else { return not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" + "Conflicting ordering requirements in aggregate functions is not supported" ); } } else { From 151f3d321d16c25d9daabf07d00f1c2678bcb512 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 13:31:16 +0300 Subject: [PATCH 08/31] Add new slt tests --- .../sqllogictest/test_files/groupby.slt | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/datafusion/sqllogictest/test_files/groupby.slt b/datafusion/sqllogictest/test_files/groupby.slt index b09ff79e88d50..7fa316eb96955 100644 --- a/datafusion/sqllogictest/test_files/groupby.slt +++ b/datafusion/sqllogictest/test_files/groupby.slt @@ -4284,3 +4284,57 @@ LIMIT 5 1 FRA 3 2022-01-02T12:00:00 EUR 200 1 TUR 2 2022-01-01T11:30:00 TRY 75 1 TUR 4 2022-01-03T10:00:00 TRY 100 + +statement ok +set datafusion.execution.target_partitions = 1; + +query III? +SELECT a, b, NTH_VALUE(c, 2), ARRAY_AGG(c) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 1 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] +0 1 26 [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] +1 2 51 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] +1 3 76 [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] + +query III? +SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), ARRAY_AGG(c ORDER BY c ASC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 1 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] +0 1 26 [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] +1 2 51 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] +1 3 76 [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] + +query II?I +SELECT a, b, ARRAY_AGG(c ORDER BY c ASC), NTH_VALUE(c, 2 ORDER BY c DESC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] 23 +0 1 [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] 48 +1 2 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] 73 +1 3 [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] 98 + +query IIIIII +SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), NTH_VALUE(c, 3 ORDER BY c ASC), NTH_VALUE(c, 2 ORDER BY c DESC), NTH_VALUE(c, 3 ORDER BY c DESC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 1 2 23 22 +0 1 26 27 48 47 +1 2 51 52 73 72 +1 3 76 77 98 97 + +# nth value cannot work with conflicting requirements +statement error DataFusion error: This feature is not implemented: Conflicting ordering requirements in aggregate functions is not supported +SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), NTH_VALUE(c, 3 ORDER BY d ASC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; From 80f5c41ccbe36cd47de60b694fc78a6ed060674a Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 15:22:35 +0300 Subject: [PATCH 09/31] Add multi partition support --- .../src/physical_optimizer/enforce_sorting.rs | 94 +++++++++++++++++++ .../src/aggregate/array_agg_ordered.rs | 48 +++++++--- .../physical-expr/src/aggregate/utils.rs | 4 +- .../physical-plan/src/aggregates/mod.rs | 1 - .../sqllogictest/test_files/groupby.slt | 56 +++++++++++ 5 files changed, 186 insertions(+), 17 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index c40671859a185..240632fc0016b 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2351,6 +2351,15 @@ mod tmp_tests { FROM multiple_ordered_table GROUP BY a, b"; + const QUERY4: &str = "SELECT a, b, NTH_VALUE(c, 2) + FROM multiple_ordered_table + GROUP BY a, b"; + + // const QUERY4: &str = + // "SELECT a, b, ARRAY_AGG(c) + // FROM multiple_ordered_table + // GROUP BY a, b"; + #[tokio::test] async fn test_query1() -> Result<()> { let config = SessionConfig::new().with_target_partitions(1); @@ -2470,4 +2479,89 @@ mod tmp_tests { assert_batches_eq!(expected, &batches); Ok(()) } + + #[tokio::test] + async fn test_query1_multi_partition() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(2); + let ctx = SessionContext::new_with_config(config); + let query = QUERY1; + + ctx.sql(MULTIPLE_ORDERED_TABLE).await?; + + let msg = format!("Creating logical plan for '{query}'"); + let dataframe = ctx.sql(query).await.expect(&msg); + let physical_plan = dataframe.create_physical_plan().await?; + print_plan(&physical_plan)?; + let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; + print_batches(&batches)?; + + let expected = vec![ + "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[ARRAY_AGG(multiple_ordered_table.c), NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", + ]; + // Get string representation of the plan + let actual = get_plan_string(&physical_plan); + assert_eq!( + expected, actual, + "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" + ); + + let expected = [ + "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", + "| a | b | ARRAY_AGG(multiple_ordered_table.c) | NTH_VALUE(multiple_ordered_table.c,Int64(2)) |", + "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", + "| 0 | 0 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] | 23 |", + "| 0 | 1 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] | 48 |", + "| 1 | 2 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] | 73 |", + "| 1 | 3 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] | 98 |", + "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) + } + + #[tokio::test] + async fn test_query4_multi_partition() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(2); + let ctx = SessionContext::new_with_config(config); + let query = QUERY4; + + ctx.sql(MULTIPLE_ORDERED_TABLE).await?; + + let msg = format!("Creating logical plan for '{query}'"); + let dataframe = ctx.sql(query).await.expect(&msg); + let physical_plan = dataframe.create_physical_plan().await?; + print_plan(&physical_plan)?; + let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; + print_batches(&batches)?; + + let expected = vec![ + "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", + " SortExec: expr=[a@0 ASC NULLS LAST,b@1 ASC NULLS LAST]", + " CoalesceBatchesExec: target_batch_size=8192", + " RepartitionExec: partitioning=Hash([a@0, b@1], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", + " RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1", + " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", + ]; + // Get string representation of the plan + let actual = get_plan_string(&physical_plan); + assert_eq!( + expected, actual, + "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" + ); + + let expected = [ + "+---+---+----------------------------------------------+", + "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) |", + "+---+---+----------------------------------------------+", + "| 0 | 0 | 1 |", + "| 0 | 1 | 26 |", + "| 1 | 2 | 51 |", + "| 1 | 3 | 76 |", + "+---+---+----------------------------------------------+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) + } } diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 388df723db513..794822777b293 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -394,16 +394,19 @@ impl AggregateExpr for NthValueAgg { fn state_fields(&self) -> Result> { let mut fields = vec![Field::new_list( - format_state_name(&self.name, "array_agg"), + format_state_name(&self.name, "nth_value"), Field::new("item", self.input_data_type.clone(), true), self.nullable, // This should be the same as field() )]; - let orderings = ordering_fields(&self.ordering_req, &self.order_by_data_types); - fields.push(Field::new_list( - format_state_name(&self.name, "array_agg_orderings"), - Field::new("item", DataType::Struct(Fields::from(orderings)), true), - self.nullable, - )); + if !self.ordering_req.is_empty() { + let orderings = + ordering_fields(&self.ordering_req, &self.order_by_data_types); + fields.push(Field::new_list( + format_state_name(&self.name, "nth_value_orderings"), + Field::new("item", DataType::Struct(Fields::from(orderings)), true), + self.nullable, + )); + } Ok(fields) } @@ -511,12 +514,22 @@ impl Accumulator for NthValueAccumulator { } // First entry in the state is the aggregation result. let array_agg_values = &states[0]; - // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. - // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores - // values received from its ordering requirement expression. (This information is necessary for during merging). - let agg_orderings = &states[1]; + if self.ordering_req.is_empty() { + let array_agg_res = + ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; + // Stores ARRAY_AGG results coming from each partition + let mut partition_values = vec![]; + // Existing values should be merged also. + partition_values.extend(self.values.clone()); + for v in array_agg_res.into_iter() { + partition_values.extend(v); + } + self.values = partition_values; + } else if let Some(agg_orderings) = states[1].as_list_opt::() { + // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. + // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores + // values received from its ordering requirement expression. (This information is necessary for during merging). - if let Some(agg_orderings) = agg_orderings.as_list_opt::() { // Stores ARRAY_AGG results coming from each partition let mut partition_values = vec![]; // Stores ordering requirement expression results coming from each partition @@ -570,8 +583,10 @@ impl Accumulator for NthValueAccumulator { } fn state(&self) -> Result> { - let mut result = vec![self.evaluate()?]; - result.push(self.evaluate_orderings()?); + let mut result = vec![self.evaluate_values()?]; + if !self.ordering_req.is_empty() { + result.push(self.evaluate_orderings()?); + } Ok(result) } @@ -640,6 +655,11 @@ impl NthValueAccumulator { let arr = ScalarValue::new_list(&orderings, &struct_type); Ok(ScalarValue::List(arr)) } + + fn evaluate_values(&self) -> Result { + let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); + Ok(ScalarValue::List(arr)) + } } /// This is a wrapper struct to be able to correctly merge ARRAY_AGG diff --git a/datafusion/physical-expr/src/aggregate/utils.rs b/datafusion/physical-expr/src/aggregate/utils.rs index d73c46a0f6870..6dd586bfb8cea 100644 --- a/datafusion/physical-expr/src/aggregate/utils.rs +++ b/datafusion/physical-expr/src/aggregate/utils.rs @@ -196,9 +196,9 @@ pub(crate) fn ordering_fields( ordering_req .iter() .zip(data_types.iter()) - .map(|(expr, dtype)| { + .map(|(sort_expr, dtype)| { Field::new( - expr.to_string().as_str(), + sort_expr.expr.to_string().as_str(), dtype.clone(), // Multi partitions may be empty hence field should be nullable. true, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 9ea1cc0c63772..22b252368a62c 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -953,7 +953,6 @@ fn get_aggregate_exprs_requirement( { requirement = finer_ordering; } else if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - println!("reverse_aggr_expr:{:?}", reverse_aggr_expr); if let Some(finer_ordering) = finer_ordering( &requirement, &reverse_aggr_expr, diff --git a/datafusion/sqllogictest/test_files/groupby.slt b/datafusion/sqllogictest/test_files/groupby.slt index 7fa316eb96955..e4f388afa2b0f 100644 --- a/datafusion/sqllogictest/test_files/groupby.slt +++ b/datafusion/sqllogictest/test_files/groupby.slt @@ -4338,3 +4338,59 @@ SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), NTH_VALUE(c, 3 ORDER BY d ASC) FROM multiple_ordered_table GROUP BY a, b ORDER BY a, b; + + +statement ok +set datafusion.execution.target_partitions = 8; + +# NTH_VALUE(c, 2) and ARRAY_AGG(c)[2] should produce same result +query III +SELECT a, b, NTH_VALUE(c, 2) - ARRAY_AGG(c)[2] +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 0 +0 1 0 +1 2 0 +1 3 0 + +query III? +SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), ARRAY_AGG(c ORDER BY c ASC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 1 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] +0 1 26 [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] +1 2 51 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] +1 3 76 [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] + +query II?I +SELECT a, b, ARRAY_AGG(c ORDER BY c ASC), NTH_VALUE(c, 2 ORDER BY c DESC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] 23 +0 1 [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] 48 +1 2 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] 73 +1 3 [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] 98 + +query IIIIII +SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), NTH_VALUE(c, 3 ORDER BY c ASC), NTH_VALUE(c, 2 ORDER BY c DESC), NTH_VALUE(c, 3 ORDER BY c DESC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 1 2 23 22 +0 1 26 27 48 47 +1 2 51 52 73 72 +1 3 76 77 98 97 + +# nth value cannot work with conflicting requirements +statement error DataFusion error: This feature is not implemented: Conflicting ordering requirements in aggregate functions is not supported +SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), NTH_VALUE(c, 3 ORDER BY d ASC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; From ec1847ff4b38190d323cb5df9ff456fa9b39fe2b Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 15:23:07 +0300 Subject: [PATCH 10/31] Minor changes --- datafusion/expr/src/aggregate_function.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/expr/src/aggregate_function.rs b/datafusion/expr/src/aggregate_function.rs index c5e48a92cf9ec..fff481e2b3938 100644 --- a/datafusion/expr/src/aggregate_function.rs +++ b/datafusion/expr/src/aggregate_function.rs @@ -17,8 +17,8 @@ //! Aggregate function module contains all built-in aggregate functions definitions +use crate::utils; use crate::{type_coercion::aggregates::*, Signature, TypeSignature, Volatility}; -use crate::{utils, BuiltInWindowFunction}; use arrow::datatypes::{DataType, Field}; use datafusion_common::{plan_datafusion_err, plan_err, DataFusionError, Result}; use std::sync::Arc; From 2b9ab957d5ce9d964851493704ccd158d93ee187 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 15:28:01 +0300 Subject: [PATCH 11/31] Minor changes --- .../src/aggregate/array_agg_ordered.rs | 38 ++++++++++--------- .../physical-expr/src/aggregate/build_in.rs | 2 +- 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 794822777b293..29c49852f1c41 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -36,7 +36,7 @@ use arrow_array::cast::AsArray; use arrow_array::Array; use arrow_schema::{Fields, SortOptions}; use datafusion_common::utils::{compare_rows, get_row_at_idx}; -use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue}; +use datafusion_common::{exec_err, internal_err, DataFusionError, Result, ScalarValue}; use datafusion_expr::Accumulator; use itertools::izip; @@ -591,24 +591,28 @@ impl Accumulator for NthValueAccumulator { } fn evaluate(&self) -> Result { - if self.n > 0 { - let forward_idx = self.n as usize - 1; - if let Some(nth_value) = self.values.get(forward_idx) { - Ok(nth_value.clone()) - } else { - ScalarValue::try_from(self.datatypes[0].clone()) + match self.n.cmp(&0) { + Ordering::Greater => { + let forward_idx = self.n as usize - 1; + if let Some(nth_value) = self.values.get(forward_idx) { + Ok(nth_value.clone()) + } else { + ScalarValue::try_from(self.datatypes[0].clone()) + } } - } else if self.n < 0 { - let backward_idx = (-self.n) as usize; - if self.values.len() >= backward_idx { - let idx = self.values.len() - backward_idx; - Ok(self.values[idx].clone()) - } else { - ScalarValue::try_from(self.datatypes[0].clone()) + Ordering::Less => { + let backward_idx = (-self.n) as usize; + if self.values.len() >= backward_idx { + let idx = self.values.len() - backward_idx; + Ok(self.values[idx].clone()) + } else { + ScalarValue::try_from(self.datatypes[0].clone()) + } + } + Ordering::Equal => { + // n cannot be 0 + internal_err!("Nth value indices are 1 based. 0 is invalid index") } - } else { - // n cannot be 0 - unreachable!(); } } diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs index 39f729f376277..91560bc960b96 100644 --- a/datafusion/physical-expr/src/aggregate/build_in.rs +++ b/datafusion/physical-expr/src/aggregate/build_in.rs @@ -30,7 +30,7 @@ use crate::aggregate::regr::RegrType; use crate::expressions::Literal; use crate::{expressions, AggregateExpr, PhysicalExpr, PhysicalSortExpr}; use arrow::datatypes::Schema; -use datafusion_common::{not_impl_err, plan_datafusion_err, DataFusionError, Result}; +use datafusion_common::{not_impl_err, DataFusionError, Result}; pub use datafusion_expr::AggregateFunction; use std::sync::Arc; From 863d4bb73707056bb8b136a9358108f230d80e8d Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 15:40:26 +0300 Subject: [PATCH 12/31] Add new aggregator to the proto --- datafusion/proto/proto/datafusion.proto | 1 + datafusion/proto/src/generated/pbjson.rs | 3 +++ datafusion/proto/src/generated/prost.rs | 3 +++ datafusion/proto/src/logical_plan/from_proto.rs | 1 + datafusion/proto/src/logical_plan/to_proto.rs | 4 ++++ 5 files changed, 12 insertions(+) diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index d5f8397aa30cf..1e3feda27d961 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -713,6 +713,7 @@ enum AggregateFunction { REGR_SYY = 33; REGR_SXY = 34; STRING_AGG = 35; + NTH_VALUE_AGG = 36; } message AggregateExprNode { diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index 12e834d75adff..128ea9159f120 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -457,6 +457,7 @@ impl serde::Serialize for AggregateFunction { Self::RegrSyy => "REGR_SYY", Self::RegrSxy => "REGR_SXY", Self::StringAgg => "STRING_AGG", + Self::NthValueAgg => "NTH_VALUE_AGG", }; serializer.serialize_str(variant) } @@ -504,6 +505,7 @@ impl<'de> serde::Deserialize<'de> for AggregateFunction { "REGR_SYY", "REGR_SXY", "STRING_AGG", + "NTH_VALUE_AGG", ]; struct GeneratedVisitor; @@ -580,6 +582,7 @@ impl<'de> serde::Deserialize<'de> for AggregateFunction { "REGR_SYY" => Ok(AggregateFunction::RegrSyy), "REGR_SXY" => Ok(AggregateFunction::RegrSxy), "STRING_AGG" => Ok(AggregateFunction::StringAgg), + "NTH_VALUE_AGG" => Ok(AggregateFunction::NthValueAgg), _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 4ee0b70325ca7..bc6747f5424d0 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -3072,6 +3072,7 @@ pub enum AggregateFunction { RegrSyy = 33, RegrSxy = 34, StringAgg = 35, + NthValueAgg = 36, } impl AggregateFunction { /// String value of the enum field names used in the ProtoBuf definition. @@ -3118,6 +3119,7 @@ impl AggregateFunction { AggregateFunction::RegrSyy => "REGR_SYY", AggregateFunction::RegrSxy => "REGR_SXY", AggregateFunction::StringAgg => "STRING_AGG", + AggregateFunction::NthValueAgg => "NTH_VALUE_AGG", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -3161,6 +3163,7 @@ impl AggregateFunction { "REGR_SYY" => Some(Self::RegrSyy), "REGR_SXY" => Some(Self::RegrSxy), "STRING_AGG" => Some(Self::StringAgg), + "NTH_VALUE_AGG" => Some(Self::NthValueAgg), _ => None, } } diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 3f48be0c4d1f3..905c1e84c82a4 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -616,6 +616,7 @@ impl From for AggregateFunction { protobuf::AggregateFunction::Median => Self::Median, protobuf::AggregateFunction::FirstValueAgg => Self::FirstValue, protobuf::AggregateFunction::LastValueAgg => Self::LastValue, + protobuf::AggregateFunction::NthValueAgg => Self::NthValue, protobuf::AggregateFunction::StringAgg => Self::StringAgg, } } diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index b7d1ef225251d..1064f7be5b39c 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -409,6 +409,7 @@ impl From<&AggregateFunction> for protobuf::AggregateFunction { AggregateFunction::Median => Self::Median, AggregateFunction::FirstValue => Self::FirstValueAgg, AggregateFunction::LastValue => Self::LastValueAgg, + AggregateFunction::NthValue => Self::NthValueAgg, AggregateFunction::StringAgg => Self::StringAgg, } } @@ -728,6 +729,9 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode { AggregateFunction::LastValue => { protobuf::AggregateFunction::LastValueAgg } + AggregateFunction::NthValue => { + protobuf::AggregateFunction::NthValueAgg + } AggregateFunction::StringAgg => { protobuf::AggregateFunction::StringAgg } From e22f456c15ae31d6c4b7d33d976d4af3b2b0aafc Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Wed, 10 Jan 2024 15:44:43 +0300 Subject: [PATCH 13/31] Remove redundant tests --- .../src/physical_optimizer/enforce_sorting.rs | 260 ------------------ 1 file changed, 260 deletions(-) diff --git a/datafusion/core/src/physical_optimizer/enforce_sorting.rs b/datafusion/core/src/physical_optimizer/enforce_sorting.rs index 240632fc0016b..f609ddea66cff 100644 --- a/datafusion/core/src/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/src/physical_optimizer/enforce_sorting.rs @@ -2305,263 +2305,3 @@ mod tests { Ok(()) } } - -#[cfg(test)] -mod tmp_tests { - use crate::assert_batches_eq; - use crate::physical_plan::{collect, displayable, ExecutionPlan}; - use crate::prelude::SessionContext; - use arrow::util::pretty::print_batches; - use datafusion_common::Result; - use datafusion_execution::config::SessionConfig; - use datafusion_physical_plan::get_plan_string; - use std::sync::Arc; - - fn print_plan(plan: &Arc) -> Result<()> { - let formatted = displayable(plan.as_ref()).indent(true).to_string(); - let actual: Vec<&str> = formatted.trim().lines().collect(); - println!("{:#?}", actual); - Ok(()) - } - - const MULTIPLE_ORDERED_TABLE: &str = "CREATE EXTERNAL TABLE multiple_ordered_table ( - a0 INTEGER, - a INTEGER, - b INTEGER, - c INTEGER, - d INTEGER - ) - STORED AS CSV - WITH HEADER ROW - WITH ORDER (a ASC, b ASC) - WITH ORDER (c ASC) - LOCATION '../core/tests/data/window_2.csv'"; - - const QUERY1: &str = "SELECT a, b, NTH_VALUE(c, 2), ARRAY_AGG(c) - FROM multiple_ordered_table - GROUP BY a, b"; - - const QUERY2: &str = - "SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), ARRAY_AGG(c ORDER BY c ASC) - FROM multiple_ordered_table - GROUP BY a, b"; - - const QUERY3: &str = - "SELECT a, b, ARRAY_AGG(c ORDER BY c ASC), NTH_VALUE(c, 2 ORDER BY c DESC) - FROM multiple_ordered_table - GROUP BY a, b"; - - const QUERY4: &str = "SELECT a, b, NTH_VALUE(c, 2) - FROM multiple_ordered_table - GROUP BY a, b"; - - // const QUERY4: &str = - // "SELECT a, b, ARRAY_AGG(c) - // FROM multiple_ordered_table - // GROUP BY a, b"; - - #[tokio::test] - async fn test_query1() -> Result<()> { - let config = SessionConfig::new().with_target_partitions(1); - let ctx = SessionContext::new_with_config(config); - let query = QUERY1; - - ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - - let msg = format!("Creating logical plan for '{query}'"); - let dataframe = ctx.sql(query).await.expect(&msg); - let physical_plan = dataframe.create_physical_plan().await?; - print_plan(&physical_plan)?; - let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; - print_batches(&batches)?; - - let expected = vec![ - "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2)), ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", - ]; - // Get string representation of the plan - let actual = get_plan_string(&physical_plan); - assert_eq!( - expected, actual, - "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" - ); - - let expected = [ - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) | ARRAY_AGG(multiple_ordered_table.c) |", - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - "| 0 | 0 | 1 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] |", - "| 0 | 1 | 26 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] |", - "| 1 | 2 | 51 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] |", - "| 1 | 3 | 76 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] |", - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - ]; - assert_batches_eq!(expected, &batches); - Ok(()) - } - - #[tokio::test] - async fn test_query2() -> Result<()> { - let config = SessionConfig::new().with_target_partitions(1); - let ctx = SessionContext::new_with_config(config); - let query = QUERY2; - - ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - - let msg = format!("Creating logical plan for '{query}'"); - let dataframe = ctx.sql(query).await.expect(&msg); - let physical_plan = dataframe.create_physical_plan().await?; - print_plan(&physical_plan)?; - let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; - print_batches(&batches)?; - - let expected = vec![ - "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2)), ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", - ]; - // Get string representation of the plan - let actual = get_plan_string(&physical_plan); - assert_eq!( - expected, actual, - "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" - ); - - let expected = [ - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) | ARRAY_AGG(multiple_ordered_table.c) |", - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - "| 0 | 0 | 1 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] |", - "| 0 | 1 | 26 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] |", - "| 1 | 2 | 51 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] |", - "| 1 | 3 | 76 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] |", - "+---+---+----------------------------------------------+------------------------------------------------------------------------------------------------------+", - ]; - assert_batches_eq!(expected, &batches); - Ok(()) - } - - #[tokio::test] - async fn test_query3() -> Result<()> { - let config = SessionConfig::new().with_target_partitions(1); - let ctx = SessionContext::new_with_config(config); - let query = QUERY3; - - ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - - let msg = format!("Creating logical plan for '{query}'"); - let dataframe = ctx.sql(query).await.expect(&msg); - let physical_plan = dataframe.create_physical_plan().await?; - print_plan(&physical_plan)?; - let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; - print_batches(&batches)?; - - let expected = vec![ - "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[ARRAY_AGG(multiple_ordered_table.c), NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", - ]; - // Get string representation of the plan - let actual = get_plan_string(&physical_plan); - assert_eq!( - expected, actual, - "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" - ); - - let expected = [ - "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", - "| a | b | ARRAY_AGG(multiple_ordered_table.c) | NTH_VALUE(multiple_ordered_table.c,Int64(2)) |", - "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", - "| 0 | 0 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] | 23 |", - "| 0 | 1 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] | 48 |", - "| 1 | 2 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] | 73 |", - "| 1 | 3 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] | 98 |", - "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", - ]; - assert_batches_eq!(expected, &batches); - Ok(()) - } - - #[tokio::test] - async fn test_query1_multi_partition() -> Result<()> { - let config = SessionConfig::new().with_target_partitions(2); - let ctx = SessionContext::new_with_config(config); - let query = QUERY1; - - ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - - let msg = format!("Creating logical plan for '{query}'"); - let dataframe = ctx.sql(query).await.expect(&msg); - let physical_plan = dataframe.create_physical_plan().await?; - print_plan(&physical_plan)?; - let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; - print_batches(&batches)?; - - let expected = vec![ - "AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[ARRAY_AGG(multiple_ordered_table.c), NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", - ]; - // Get string representation of the plan - let actual = get_plan_string(&physical_plan); - assert_eq!( - expected, actual, - "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" - ); - - let expected = [ - "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", - "| a | b | ARRAY_AGG(multiple_ordered_table.c) | NTH_VALUE(multiple_ordered_table.c,Int64(2)) |", - "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", - "| 0 | 0 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] | 23 |", - "| 0 | 1 | [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] | 48 |", - "| 1 | 2 | [50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74] | 73 |", - "| 1 | 3 | [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99] | 98 |", - "+---+---+------------------------------------------------------------------------------------------------------+----------------------------------------------+", - ]; - assert_batches_eq!(expected, &batches); - Ok(()) - } - - #[tokio::test] - async fn test_query4_multi_partition() -> Result<()> { - let config = SessionConfig::new().with_target_partitions(2); - let ctx = SessionContext::new_with_config(config); - let query = QUERY4; - - ctx.sql(MULTIPLE_ORDERED_TABLE).await?; - - let msg = format!("Creating logical plan for '{query}'"); - let dataframe = ctx.sql(query).await.expect(&msg); - let physical_plan = dataframe.create_physical_plan().await?; - print_plan(&physical_plan)?; - let batches = collect(physical_plan.clone(), ctx.task_ctx()).await?; - print_batches(&batches)?; - - let expected = vec![ - "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", - " SortExec: expr=[a@0 ASC NULLS LAST,b@1 ASC NULLS LAST]", - " CoalesceBatchesExec: target_batch_size=8192", - " RepartitionExec: partitioning=Hash([a@0, b@1], 2), input_partitions=2", - " AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[NTH_VALUE(multiple_ordered_table.c,Int64(2))], ordering_mode=Sorted", - " RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1", - " CsvExec: file_groups={1 group: [[Users/akurmustafa/projects/synnada/arrow-datafusion-synnada/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true", - ]; - // Get string representation of the plan - let actual = get_plan_string(&physical_plan); - assert_eq!( - expected, actual, - "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" - ); - - let expected = [ - "+---+---+----------------------------------------------+", - "| a | b | NTH_VALUE(multiple_ordered_table.c,Int64(2)) |", - "+---+---+----------------------------------------------+", - "| 0 | 0 | 1 |", - "| 0 | 1 | 26 |", - "| 1 | 2 | 51 |", - "| 1 | 3 | 76 |", - "+---+---+----------------------------------------------+", - ]; - assert_batches_eq!(expected, &batches); - Ok(()) - } -} From 22ff28e873d2b3e5cf2eccb56cd6d758542e30d3 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Thu, 11 Jan 2024 10:18:05 +0300 Subject: [PATCH 14/31] Keep n entries in the state for nth value --- .../src/aggregate/array_agg_ordered.rs | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 29c49852f1c41..feb1da5ff1dd1 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -498,11 +498,28 @@ impl Accumulator for NthValueAccumulator { return Ok(()); } - let n_row = values[0].len(); - for index in 0..n_row { - let row = get_row_at_idx(values, index)?; - self.values.push(row[0].clone()); - self.ordering_values.push(row[1..].to_vec()); + if self.n > 0 { + let n_required_from_start = self.n as usize; + let n_remaining = n_required_from_start.saturating_sub(self.values.len()); + let n_remaining = std::cmp::min(n_remaining, values[0].len()); + for index in 0..n_remaining { + let row = get_row_at_idx(values, index)?; + self.values.push(row[0].clone()); + self.ordering_values.push(row[1..].to_vec()); + } + } else { + let n_row = values[0].len(); + let n_required_from_end = (-self.n) as usize; + for index in 0..n_row { + let row = get_row_at_idx(values, index)?; + self.values.push(row[0].clone()); + self.ordering_values.push(row[1..].to_vec()); + } + let start_offset = self.values.len().saturating_sub(n_required_from_end); + if start_offset > 0 { + self.values = self.values[start_offset..].to_vec(); + self.ordering_values = self.ordering_values[start_offset..].to_vec(); + } } Ok(()) From d542b367a9f0315de1fba6707e0c79d2396955a0 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Thu, 11 Jan 2024 10:36:57 +0300 Subject: [PATCH 15/31] Change implementation --- .../src/aggregate/array_agg_ordered.rs | 53 ++++++++++--------- 1 file changed, 27 insertions(+), 26 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index feb1da5ff1dd1..fe42625e25812 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -480,6 +480,10 @@ impl NthValueAccumulator { ordering_dtypes: &[DataType], ordering_req: LexOrdering, ) -> Result { + if n == 0 { + // n cannot be 0 + return internal_err!("Nth value indices are 1 based. 0 is invalid index"); + } let mut datatypes = vec![datatype.clone()]; datatypes.extend(ordering_dtypes.iter().cloned()); Ok(Self { @@ -498,9 +502,9 @@ impl Accumulator for NthValueAccumulator { return Ok(()); } + let n_required = self.n.unsigned_abs() as usize; if self.n > 0 { - let n_required_from_start = self.n as usize; - let n_remaining = n_required_from_start.saturating_sub(self.values.len()); + let n_remaining = n_required.saturating_sub(self.values.len()); let n_remaining = std::cmp::min(n_remaining, values[0].len()); for index in 0..n_remaining { let row = get_row_at_idx(values, index)?; @@ -509,13 +513,12 @@ impl Accumulator for NthValueAccumulator { } } else { let n_row = values[0].len(); - let n_required_from_end = (-self.n) as usize; for index in 0..n_row { let row = get_row_at_idx(values, index)?; self.values.push(row[0].clone()); self.ordering_values.push(row[1..].to_vec()); } - let start_offset = self.values.len().saturating_sub(n_required_from_end); + let start_offset = self.values.len().saturating_sub(n_required); if start_offset > 0 { self.values = self.values[start_offset..].to_vec(); self.ordering_values = self.ordering_values[start_offset..].to_vec(); @@ -531,6 +534,7 @@ impl Accumulator for NthValueAccumulator { } // First entry in the state is the aggregation result. let array_agg_values = &states[0]; + let n_required = self.n.unsigned_abs() as usize; if self.ordering_req.is_empty() { let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; @@ -540,6 +544,10 @@ impl Accumulator for NthValueAccumulator { partition_values.extend(self.values.clone()); for v in array_agg_res.into_iter() { partition_values.extend(v); + if partition_values.len() > n_required { + // There is enough data collected can stop merging + break; + } } self.values = partition_values; } else if let Some(agg_orderings) = states[1].as_list_opt::() { @@ -608,28 +616,21 @@ impl Accumulator for NthValueAccumulator { } fn evaluate(&self) -> Result { - match self.n.cmp(&0) { - Ordering::Greater => { - let forward_idx = self.n as usize - 1; - if let Some(nth_value) = self.values.get(forward_idx) { - Ok(nth_value.clone()) - } else { - ScalarValue::try_from(self.datatypes[0].clone()) - } - } - Ordering::Less => { - let backward_idx = (-self.n) as usize; - if self.values.len() >= backward_idx { - let idx = self.values.len() - backward_idx; - Ok(self.values[idx].clone()) - } else { - ScalarValue::try_from(self.datatypes[0].clone()) - } - } - Ordering::Equal => { - // n cannot be 0 - internal_err!("Nth value indices are 1 based. 0 is invalid index") - } + let n_required = self.n.unsigned_abs() as usize; + let from_start = self.n > 0; + let nth_value_idx = if from_start { + // index is from start + let forward_idx = n_required - 1; + (forward_idx < self.values.len()).then_some(forward_idx) + } else { + // index is from end + self.values.len().checked_sub(n_required) + }; + if let Some(idx) = nth_value_idx { + let nth_value = self.values[idx].clone(); + Ok(nth_value.clone()) + } else { + ScalarValue::try_from(self.datatypes[0].clone()) } } From 5cedec57068d4f5249c6748bfa67b166d9243eda Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Thu, 11 Jan 2024 10:48:02 +0300 Subject: [PATCH 16/31] Move nth value to its own file --- .../src/aggregate/array_agg_ordered.rs | 367 +--------------- datafusion/physical-expr/src/aggregate/mod.rs | 1 + .../physical-expr/src/aggregate/nth_value.rs | 390 ++++++++++++++++++ .../physical-expr/src/expressions/mod.rs | 3 +- 4 files changed, 396 insertions(+), 365 deletions(-) create mode 100644 datafusion/physical-expr/src/aggregate/nth_value.rs diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index fe42625e25812..16f1feb12def5 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -26,9 +26,7 @@ use std::sync::Arc; use crate::aggregate::utils::{down_cast_any_ref, ordering_fields}; use crate::expressions::format_state_name; -use crate::{ - reverse_order_bys, AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr, -}; +use crate::{AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr}; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field}; @@ -36,7 +34,7 @@ use arrow_array::cast::AsArray; use arrow_array::Array; use arrow_schema::{Fields, SortOptions}; use datafusion_common::utils::{compare_rows, get_row_at_idx}; -use datafusion_common::{exec_err, internal_err, DataFusionError, Result, ScalarValue}; +use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue}; use datafusion_expr::Accumulator; use itertools::izip; @@ -325,365 +323,6 @@ impl OrderSensitiveArrayAggAccumulator { } } -/// Expression for a ARRAY_AGG(ORDER BY) aggregation. -/// When aggregation works in multiple partitions -/// aggregations are split into multiple partitions, -/// then their results are merged. This aggregator -/// is a version of ARRAY_AGG that can support producing -/// intermediate aggregation (with necessary side information) -/// and that can merge aggregations from multiple partitions. -#[derive(Debug)] -pub struct NthValueAgg { - /// Column name - name: String, - /// The DataType for the input expression - input_data_type: DataType, - /// The input expression - expr: Arc, - n: i64, - /// If the input expression can have NULLs - nullable: bool, - /// Ordering data types - order_by_data_types: Vec, - /// Ordering requirement - ordering_req: LexOrdering, -} - -impl NthValueAgg { - /// Create a new `OrderSensitiveArrayAgg` aggregate function - pub fn new( - expr: Arc, - n: i64, - name: impl Into, - input_data_type: DataType, - nullable: bool, - order_by_data_types: Vec, - ordering_req: LexOrdering, - ) -> Self { - Self { - name: name.into(), - input_data_type, - expr, - n, - nullable, - order_by_data_types, - ordering_req, - } - } -} - -impl AggregateExpr for NthValueAgg { - fn as_any(&self) -> &dyn Any { - self - } - - fn field(&self) -> Result { - Ok(Field::new(&self.name, self.input_data_type.clone(), true)) - } - - fn create_accumulator(&self) -> Result> { - Ok(Box::new( - crate::aggregate::array_agg_ordered::NthValueAccumulator::try_new( - self.n, - &self.input_data_type, - &self.order_by_data_types, - self.ordering_req.clone(), - )?, - )) - } - - fn state_fields(&self) -> Result> { - let mut fields = vec![Field::new_list( - format_state_name(&self.name, "nth_value"), - Field::new("item", self.input_data_type.clone(), true), - self.nullable, // This should be the same as field() - )]; - if !self.ordering_req.is_empty() { - let orderings = - ordering_fields(&self.ordering_req, &self.order_by_data_types); - fields.push(Field::new_list( - format_state_name(&self.name, "nth_value_orderings"), - Field::new("item", DataType::Struct(Fields::from(orderings)), true), - self.nullable, - )); - } - Ok(fields) - } - - fn expressions(&self) -> Vec> { - vec![self.expr.clone()] - } - - fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { - if self.ordering_req.is_empty() { - None - } else { - Some(&self.ordering_req) - } - } - - fn name(&self) -> &str { - &self.name - } - - fn reverse_expr(&self) -> Option> { - Some(Arc::new(Self { - name: self.name.to_string(), - input_data_type: self.input_data_type.clone(), - expr: self.expr.clone(), - n: -self.n, - nullable: self.nullable, - order_by_data_types: self.order_by_data_types.clone(), - ordering_req: reverse_order_bys(&self.ordering_req), - }) as _) - } -} - -impl PartialEq for NthValueAgg { - fn eq(&self, other: &dyn Any) -> bool { - down_cast_any_ref(other) - .downcast_ref::() - .map(|x| { - self.name == x.name - && self.input_data_type == x.input_data_type - && self.order_by_data_types == x.order_by_data_types - && self.expr.eq(&x.expr) - }) - .unwrap_or(false) - } -} - -#[derive(Debug)] -pub(crate) struct NthValueAccumulator { - n: i64, - // `values` stores entries in the ARRAY_AGG result. - values: Vec, - // `ordering_values` stores values of ordering requirement expression - // corresponding to each value in the ARRAY_AGG. - // For each `ScalarValue` inside `values`, there will be a corresponding - // `Vec` inside `ordering_values` which stores it ordering. - // This information is used during merging results of the different partitions. - // For detailed information how merging is done see [`merge_ordered_arrays`] - ordering_values: Vec>, - // `datatypes` stores, datatype of expression inside ARRAY_AGG and ordering requirement expressions. - datatypes: Vec, - // Stores ordering requirement of the Accumulator - ordering_req: LexOrdering, -} - -impl NthValueAccumulator { - /// Create a new order-sensitive ARRAY_AGG accumulator based on the given - /// item data type. - pub fn try_new( - n: i64, - datatype: &DataType, - ordering_dtypes: &[DataType], - ordering_req: LexOrdering, - ) -> Result { - if n == 0 { - // n cannot be 0 - return internal_err!("Nth value indices are 1 based. 0 is invalid index"); - } - let mut datatypes = vec![datatype.clone()]; - datatypes.extend(ordering_dtypes.iter().cloned()); - Ok(Self { - n, - values: vec![], - ordering_values: vec![], - datatypes, - ordering_req, - }) - } -} - -impl Accumulator for NthValueAccumulator { - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - if values.is_empty() { - return Ok(()); - } - - let n_required = self.n.unsigned_abs() as usize; - if self.n > 0 { - let n_remaining = n_required.saturating_sub(self.values.len()); - let n_remaining = std::cmp::min(n_remaining, values[0].len()); - for index in 0..n_remaining { - let row = get_row_at_idx(values, index)?; - self.values.push(row[0].clone()); - self.ordering_values.push(row[1..].to_vec()); - } - } else { - let n_row = values[0].len(); - for index in 0..n_row { - let row = get_row_at_idx(values, index)?; - self.values.push(row[0].clone()); - self.ordering_values.push(row[1..].to_vec()); - } - let start_offset = self.values.len().saturating_sub(n_required); - if start_offset > 0 { - self.values = self.values[start_offset..].to_vec(); - self.ordering_values = self.ordering_values[start_offset..].to_vec(); - } - } - - Ok(()) - } - - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - if states.is_empty() { - return Ok(()); - } - // First entry in the state is the aggregation result. - let array_agg_values = &states[0]; - let n_required = self.n.unsigned_abs() as usize; - if self.ordering_req.is_empty() { - let array_agg_res = - ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - // Stores ARRAY_AGG results coming from each partition - let mut partition_values = vec![]; - // Existing values should be merged also. - partition_values.extend(self.values.clone()); - for v in array_agg_res.into_iter() { - partition_values.extend(v); - if partition_values.len() > n_required { - // There is enough data collected can stop merging - break; - } - } - self.values = partition_values; - } else if let Some(agg_orderings) = states[1].as_list_opt::() { - // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. - // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores - // values received from its ordering requirement expression. (This information is necessary for during merging). - - // Stores ARRAY_AGG results coming from each partition - let mut partition_values = vec![]; - // Stores ordering requirement expression results coming from each partition - let mut partition_ordering_values = vec![]; - - // Existing values should be merged also. - partition_values.push(self.values.clone()); - partition_ordering_values.push(self.ordering_values.clone()); - - let array_agg_res = - ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - - for v in array_agg_res.into_iter() { - partition_values.push(v); - } - - let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; - - for partition_ordering_rows in orderings.into_iter() { - // Extract value from struct to ordering_rows for each group/partition - let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { - if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { - Ok(ordering_columns_per_row) - } else { - exec_err!( - "Expects to receive ScalarValue::Struct(Some(..), _) but got:{:?}", - ordering_row.data_type() - ) - } - }).collect::>>()?; - - partition_ordering_values.push(ordering_value); - } - - let sort_options = self - .ordering_req - .iter() - .map(|sort_expr| sort_expr.options) - .collect::>(); - let (new_values, new_orderings) = merge_ordered_arrays( - &partition_values, - &partition_ordering_values, - &sort_options, - )?; - self.values = new_values; - self.ordering_values = new_orderings; - } else { - return exec_err!("Expects to receive a list array"); - } - Ok(()) - } - - fn state(&self) -> Result> { - let mut result = vec![self.evaluate_values()?]; - if !self.ordering_req.is_empty() { - result.push(self.evaluate_orderings()?); - } - Ok(result) - } - - fn evaluate(&self) -> Result { - let n_required = self.n.unsigned_abs() as usize; - let from_start = self.n > 0; - let nth_value_idx = if from_start { - // index is from start - let forward_idx = n_required - 1; - (forward_idx < self.values.len()).then_some(forward_idx) - } else { - // index is from end - self.values.len().checked_sub(n_required) - }; - if let Some(idx) = nth_value_idx { - let nth_value = self.values[idx].clone(); - Ok(nth_value.clone()) - } else { - ScalarValue::try_from(self.datatypes[0].clone()) - } - } - - fn size(&self) -> usize { - let mut total = std::mem::size_of_val(self) - + ScalarValue::size_of_vec(&self.values) - - std::mem::size_of_val(&self.values); - - // Add size of the `self.ordering_values` - total += - std::mem::size_of::>() * self.ordering_values.capacity(); - for row in &self.ordering_values { - total += ScalarValue::size_of_vec(row) - std::mem::size_of_val(row); - } - - // Add size of the `self.datatypes` - total += std::mem::size_of::() * self.datatypes.capacity(); - for dtype in &self.datatypes { - total += dtype.size() - std::mem::size_of_val(dtype); - } - - // Add size of the `self.ordering_req` - total += std::mem::size_of::() * self.ordering_req.capacity(); - // TODO: Calculate size of each `PhysicalSortExpr` more accurately. - total - } -} - -impl NthValueAccumulator { - fn evaluate_orderings(&self) -> Result { - let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]); - let struct_field = Fields::from(fields.clone()); - - let orderings: Vec = self - .ordering_values - .iter() - .map(|ordering| { - ScalarValue::Struct(Some(ordering.clone()), struct_field.clone()) - }) - .collect(); - let struct_type = DataType::Struct(Fields::from(fields)); - - // Wrap in List, so we have the same data structure ListArray(StructArray..) for group by cases - let arr = ScalarValue::new_list(&orderings, &struct_type); - Ok(ScalarValue::List(arr)) - } - - fn evaluate_values(&self) -> Result { - let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); - Ok(ScalarValue::List(arr)) - } -} - /// This is a wrapper struct to be able to correctly merge ARRAY_AGG /// data from multiple partitions using `BinaryHeap`. /// When used inside `BinaryHeap` this struct returns smallest `CustomElement`, @@ -772,7 +411,7 @@ impl<'a> PartialOrd for CustomElement<'a> { /// For each ScalarValue in the `values` we have a corresponding `Vec` (like timestamp of it) /// for the example above `sort_options` will have size two, that defines ordering requirement of the merge. /// Inner `Vec`s of the `ordering_values` will be compared according `sort_options` (Their sizes should match) -fn merge_ordered_arrays( +pub(crate) fn merge_ordered_arrays( // We will merge values into single `Vec`. values: &[Vec], // `values` will be merged according to `ordering_values`. diff --git a/datafusion/physical-expr/src/aggregate/mod.rs b/datafusion/physical-expr/src/aggregate/mod.rs index 1e265c65e754a..270a8e6f7705a 100644 --- a/datafusion/physical-expr/src/aggregate/mod.rs +++ b/datafusion/physical-expr/src/aggregate/mod.rs @@ -47,6 +47,7 @@ pub(crate) mod covariance; pub(crate) mod first_last; pub(crate) mod grouping; pub(crate) mod median; +pub(crate) mod nth_value; pub(crate) mod string_agg; #[macro_use] pub(crate) mod min_max; diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs new file mode 100644 index 0000000000000..e62d3fa4af5c1 --- /dev/null +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -0,0 +1,390 @@ +// 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::any::Any; +use std::sync::Arc; + +use crate::aggregate::array_agg_ordered::merge_ordered_arrays; +use crate::aggregate::utils::{down_cast_any_ref, ordering_fields}; +use crate::expressions::format_state_name; +use crate::{ + reverse_order_bys, AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr, +}; +use arrow_array::cast::AsArray; +use arrow_array::ArrayRef; +use arrow_schema::{DataType, Field, Fields}; +use datafusion_common::utils::get_row_at_idx; +use datafusion_common::DataFusionError; +use datafusion_common::{exec_err, internal_err, Result, ScalarValue}; +use datafusion_expr::Accumulator; + +/// Expression for a ARRAY_AGG(ORDER BY) aggregation. +/// When aggregation works in multiple partitions +/// aggregations are split into multiple partitions, +/// then their results are merged. This aggregator +/// is a version of ARRAY_AGG that can support producing +/// intermediate aggregation (with necessary side information) +/// and that can merge aggregations from multiple partitions. +#[derive(Debug)] +pub struct NthValueAgg { + /// Column name + name: String, + /// The DataType for the input expression + input_data_type: DataType, + /// The input expression + expr: Arc, + n: i64, + /// If the input expression can have NULLs + nullable: bool, + /// Ordering data types + order_by_data_types: Vec, + /// Ordering requirement + ordering_req: LexOrdering, +} + +impl NthValueAgg { + /// Create a new `OrderSensitiveArrayAgg` aggregate function + pub fn new( + expr: Arc, + n: i64, + name: impl Into, + input_data_type: DataType, + nullable: bool, + order_by_data_types: Vec, + ordering_req: LexOrdering, + ) -> Self { + Self { + name: name.into(), + input_data_type, + expr, + n, + nullable, + order_by_data_types, + ordering_req, + } + } +} + +impl AggregateExpr for NthValueAgg { + fn as_any(&self) -> &dyn Any { + self + } + + fn field(&self) -> Result { + Ok(Field::new(&self.name, self.input_data_type.clone(), true)) + } + + fn create_accumulator(&self) -> Result> { + Ok(Box::new(NthValueAccumulator::try_new( + self.n, + &self.input_data_type, + &self.order_by_data_types, + self.ordering_req.clone(), + )?)) + } + + fn state_fields(&self) -> Result> { + let mut fields = vec![Field::new_list( + format_state_name(&self.name, "nth_value"), + Field::new("item", self.input_data_type.clone(), true), + self.nullable, // This should be the same as field() + )]; + if !self.ordering_req.is_empty() { + let orderings = + ordering_fields(&self.ordering_req, &self.order_by_data_types); + fields.push(Field::new_list( + format_state_name(&self.name, "nth_value_orderings"), + Field::new("item", DataType::Struct(Fields::from(orderings)), true), + self.nullable, + )); + } + Ok(fields) + } + + fn expressions(&self) -> Vec> { + vec![self.expr.clone()] + } + + fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { + if self.ordering_req.is_empty() { + None + } else { + Some(&self.ordering_req) + } + } + + fn name(&self) -> &str { + &self.name + } + + fn reverse_expr(&self) -> Option> { + Some(Arc::new(Self { + name: self.name.to_string(), + input_data_type: self.input_data_type.clone(), + expr: self.expr.clone(), + n: -self.n, + nullable: self.nullable, + order_by_data_types: self.order_by_data_types.clone(), + ordering_req: reverse_order_bys(&self.ordering_req), + }) as _) + } +} + +impl PartialEq for NthValueAgg { + fn eq(&self, other: &dyn Any) -> bool { + down_cast_any_ref(other) + .downcast_ref::() + .map(|x| { + self.name == x.name + && self.input_data_type == x.input_data_type + && self.order_by_data_types == x.order_by_data_types + && self.expr.eq(&x.expr) + }) + .unwrap_or(false) + } +} + +#[derive(Debug)] +pub(crate) struct NthValueAccumulator { + n: i64, + // `values` stores entries in the ARRAY_AGG result. + values: Vec, + // `ordering_values` stores values of ordering requirement expression + // corresponding to each value in the ARRAY_AGG. + // For each `ScalarValue` inside `values`, there will be a corresponding + // `Vec` inside `ordering_values` which stores it ordering. + // This information is used during merging results of the different partitions. + // For detailed information how merging is done see [`merge_ordered_arrays`] + ordering_values: Vec>, + // `datatypes` stores, datatype of expression inside ARRAY_AGG and ordering requirement expressions. + datatypes: Vec, + // Stores ordering requirement of the Accumulator + ordering_req: LexOrdering, +} + +impl NthValueAccumulator { + /// Create a new order-sensitive ARRAY_AGG accumulator based on the given + /// item data type. + pub fn try_new( + n: i64, + datatype: &DataType, + ordering_dtypes: &[DataType], + ordering_req: LexOrdering, + ) -> Result { + if n == 0 { + // n cannot be 0 + return internal_err!("Nth value indices are 1 based. 0 is invalid index"); + } + let mut datatypes = vec![datatype.clone()]; + datatypes.extend(ordering_dtypes.iter().cloned()); + Ok(Self { + n, + values: vec![], + ordering_values: vec![], + datatypes, + ordering_req, + }) + } +} + +impl Accumulator for NthValueAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.is_empty() { + return Ok(()); + } + + let n_required = self.n.unsigned_abs() as usize; + if self.n > 0 { + let n_remaining = n_required.saturating_sub(self.values.len()); + let n_remaining = std::cmp::min(n_remaining, values[0].len()); + for index in 0..n_remaining { + let row = get_row_at_idx(values, index)?; + self.values.push(row[0].clone()); + self.ordering_values.push(row[1..].to_vec()); + } + } else { + let n_row = values[0].len(); + for index in 0..n_row { + let row = get_row_at_idx(values, index)?; + self.values.push(row[0].clone()); + self.ordering_values.push(row[1..].to_vec()); + } + let start_offset = self.values.len().saturating_sub(n_required); + if start_offset > 0 { + self.values = self.values[start_offset..].to_vec(); + self.ordering_values = self.ordering_values[start_offset..].to_vec(); + } + } + + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + if states.is_empty() { + return Ok(()); + } + // First entry in the state is the aggregation result. + let array_agg_values = &states[0]; + let n_required = self.n.unsigned_abs() as usize; + if self.ordering_req.is_empty() { + let array_agg_res = + ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; + // Stores ARRAY_AGG results coming from each partition + let mut partition_values = vec![]; + // Existing values should be merged also. + partition_values.extend(self.values.clone()); + for v in array_agg_res.into_iter() { + partition_values.extend(v); + if partition_values.len() > n_required { + // There is enough data collected can stop merging + break; + } + } + self.values = partition_values; + } else if let Some(agg_orderings) = states[1].as_list_opt::() { + // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. + // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores + // values received from its ordering requirement expression. (This information is necessary for during merging). + + // Stores ARRAY_AGG results coming from each partition + let mut partition_values = vec![]; + // Stores ordering requirement expression results coming from each partition + let mut partition_ordering_values = vec![]; + + // Existing values should be merged also. + partition_values.push(self.values.clone()); + partition_ordering_values.push(self.ordering_values.clone()); + + let array_agg_res = + ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; + + for v in array_agg_res.into_iter() { + partition_values.push(v); + } + + let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; + + for partition_ordering_rows in orderings.into_iter() { + // Extract value from struct to ordering_rows for each group/partition + let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { + if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { + Ok(ordering_columns_per_row) + } else { + exec_err!( + "Expects to receive ScalarValue::Struct(Some(..), _) but got:{:?}", + ordering_row.data_type() + ) + } + }).collect::>>()?; + + partition_ordering_values.push(ordering_value); + } + + let sort_options = self + .ordering_req + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(); + let (new_values, new_orderings) = merge_ordered_arrays( + &partition_values, + &partition_ordering_values, + &sort_options, + )?; + self.values = new_values; + self.ordering_values = new_orderings; + } else { + return exec_err!("Expects to receive a list array"); + } + Ok(()) + } + + fn state(&self) -> Result> { + let mut result = vec![self.evaluate_values()?]; + if !self.ordering_req.is_empty() { + result.push(self.evaluate_orderings()?); + } + Ok(result) + } + + fn evaluate(&self) -> Result { + let n_required = self.n.unsigned_abs() as usize; + let from_start = self.n > 0; + let nth_value_idx = if from_start { + // index is from start + let forward_idx = n_required - 1; + (forward_idx < self.values.len()).then_some(forward_idx) + } else { + // index is from end + self.values.len().checked_sub(n_required) + }; + if let Some(idx) = nth_value_idx { + let nth_value = self.values[idx].clone(); + Ok(nth_value.clone()) + } else { + ScalarValue::try_from(self.datatypes[0].clone()) + } + } + + fn size(&self) -> usize { + let mut total = std::mem::size_of_val(self) + + ScalarValue::size_of_vec(&self.values) + - std::mem::size_of_val(&self.values); + + // Add size of the `self.ordering_values` + total += + std::mem::size_of::>() * self.ordering_values.capacity(); + for row in &self.ordering_values { + total += ScalarValue::size_of_vec(row) - std::mem::size_of_val(row); + } + + // Add size of the `self.datatypes` + total += std::mem::size_of::() * self.datatypes.capacity(); + for dtype in &self.datatypes { + total += dtype.size() - std::mem::size_of_val(dtype); + } + + // Add size of the `self.ordering_req` + total += std::mem::size_of::() * self.ordering_req.capacity(); + // TODO: Calculate size of each `PhysicalSortExpr` more accurately. + total + } +} + +impl NthValueAccumulator { + fn evaluate_orderings(&self) -> Result { + let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]); + let struct_field = Fields::from(fields.clone()); + + let orderings: Vec = self + .ordering_values + .iter() + .map(|ordering| { + ScalarValue::Struct(Some(ordering.clone()), struct_field.clone()) + }) + .collect(); + let struct_type = DataType::Struct(Fields::from(fields)); + + // Wrap in List, so we have the same data structure ListArray(StructArray..) for group by cases + let arr = ScalarValue::new_list(&orderings, &struct_type); + Ok(ScalarValue::List(arr)) + } + + fn evaluate_values(&self) -> Result { + let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); + Ok(ScalarValue::List(arr)) + } +} diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index d7bb33d5457ce..586b6ecd4e54f 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -46,7 +46,7 @@ pub use crate::aggregate::approx_percentile_cont::ApproxPercentileCont; pub use crate::aggregate::approx_percentile_cont_with_weight::ApproxPercentileContWithWeight; pub use crate::aggregate::array_agg::ArrayAgg; pub use crate::aggregate::array_agg_distinct::DistinctArrayAgg; -pub use crate::aggregate::array_agg_ordered::{NthValueAgg, OrderSensitiveArrayAgg}; +pub use crate::aggregate::array_agg_ordered::OrderSensitiveArrayAgg; pub use crate::aggregate::average::{Avg, AvgAccumulator}; pub use crate::aggregate::bit_and_or_xor::{BitAnd, BitOr, BitXor, DistinctBitXor}; pub use crate::aggregate::bool_and_or::{BoolAnd, BoolOr}; @@ -60,6 +60,7 @@ pub use crate::aggregate::grouping::Grouping; pub use crate::aggregate::median::Median; pub use crate::aggregate::min_max::{Max, Min}; pub use crate::aggregate::min_max::{MaxAccumulator, MinAccumulator}; +pub use crate::aggregate::nth_value::NthValueAgg; pub use crate::aggregate::regr::{Regr, RegrType}; pub use crate::aggregate::stats::StatsType; pub use crate::aggregate::stddev::{Stddev, StddevPop}; From 24a0225b94d8a62d8bfed5642ee3f8e4b3bc9467 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Thu, 11 Jan 2024 10:55:10 +0300 Subject: [PATCH 17/31] Minor changes --- .../physical-expr/src/aggregate/nth_value.rs | 40 +++++++++++++------ 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index e62d3fa4af5c1..c7f9dd40ffad9 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -136,9 +136,11 @@ impl AggregateExpr for NthValueAgg { name: self.name.to_string(), input_data_type: self.input_data_type.clone(), expr: self.expr.clone(), + // index should be from the opposite side n: -self.n, nullable: self.nullable, order_by_data_types: self.order_by_data_types.clone(), + // reverse requirement ordering_req: reverse_order_bys(&self.ordering_req), }) as _) } @@ -208,21 +210,14 @@ impl Accumulator for NthValueAccumulator { } let n_required = self.n.unsigned_abs() as usize; - if self.n > 0 { + let from_start = self.n > 0; + if from_start { + // direction is from start let n_remaining = n_required.saturating_sub(self.values.len()); - let n_remaining = std::cmp::min(n_remaining, values[0].len()); - for index in 0..n_remaining { - let row = get_row_at_idx(values, index)?; - self.values.push(row[0].clone()); - self.ordering_values.push(row[1..].to_vec()); - } + self.append_new_data(values, Some(n_remaining))?; } else { - let n_row = values[0].len(); - for index in 0..n_row { - let row = get_row_at_idx(values, index)?; - self.values.push(row[0].clone()); - self.ordering_values.push(row[1..].to_vec()); - } + // direction is from end + self.append_new_data(values, None)?; let start_offset = self.values.len().saturating_sub(n_required); if start_offset > 0 { self.values = self.values[start_offset..].to_vec(); @@ -387,4 +382,23 @@ impl NthValueAccumulator { let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); Ok(ScalarValue::List(arr)) } + + fn append_new_data( + &mut self, + values: &[ArrayRef], + fetch: Option, + ) -> Result<()> { + let n_row = values[0].len(); + let n_to_add = if let Some(fetch) = fetch { + std::cmp::min(fetch, n_row) + } else { + n_row + }; + for index in 0..n_to_add { + let row = get_row_at_idx(values, index)?; + self.values.push(row[0].clone()); + self.ordering_values.push(row[1..].to_vec()); + } + Ok(()) + } } From 3ee2b55e50282e8978b1e39797c3257a8cb153cc Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Thu, 11 Jan 2024 10:58:35 +0300 Subject: [PATCH 18/31] minor changes --- datafusion/sqllogictest/test_files/groupby.slt | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/datafusion/sqllogictest/test_files/groupby.slt b/datafusion/sqllogictest/test_files/groupby.slt index e4f388afa2b0f..3af17a7d08504 100644 --- a/datafusion/sqllogictest/test_files/groupby.slt +++ b/datafusion/sqllogictest/test_files/groupby.slt @@ -4332,14 +4332,8 @@ ORDER BY a, b; 1 2 51 52 73 72 1 3 76 77 98 97 -# nth value cannot work with conflicting requirements -statement error DataFusion error: This feature is not implemented: Conflicting ordering requirements in aggregate functions is not supported -SELECT a, b, NTH_VALUE(c, 2 ORDER BY c ASC), NTH_VALUE(c, 3 ORDER BY d ASC) -FROM multiple_ordered_table -GROUP BY a, b -ORDER BY a, b; - +# increase partition to 8 statement ok set datafusion.execution.target_partitions = 8; From a136973a9daa954012ae149b8c8c00b17d74f04b Mon Sep 17 00:00:00 2001 From: Mehmet Ozan Kabak Date: Fri, 12 Jan 2024 15:27:20 +0300 Subject: [PATCH 19/31] Review --- datafusion/expr/src/aggregate_function.rs | 32 ++++++---- .../expr/src/type_coercion/aggregates.rs | 12 ++-- .../src/aggregate/array_agg_ordered.rs | 7 +- .../physical-expr/src/aggregate/build_in.rs | 29 +++++---- .../physical-expr/src/aggregate/nth_value.rs | 64 ++++++++----------- .../physical-plan/src/aggregates/mod.rs | 1 - .../proto/src/logical_plan/from_proto.rs | 26 ++++---- datafusion/proto/src/logical_plan/to_proto.rs | 1 + 8 files changed, 85 insertions(+), 87 deletions(-) diff --git a/datafusion/expr/src/aggregate_function.rs b/datafusion/expr/src/aggregate_function.rs index fff481e2b3938..c56107104f096 100644 --- a/datafusion/expr/src/aggregate_function.rs +++ b/datafusion/expr/src/aggregate_function.rs @@ -17,12 +17,15 @@ //! Aggregate function module contains all built-in aggregate functions definitions +use std::sync::Arc; +use std::{fmt, str::FromStr}; + use crate::utils; use crate::{type_coercion::aggregates::*, Signature, TypeSignature, Volatility}; + use arrow::datatypes::{DataType, Field}; use datafusion_common::{plan_datafusion_err, plan_err, DataFusionError, Result}; -use std::sync::Arc; -use std::{fmt, str::FromStr}; + use strum_macros::EnumIter; /// Enum of all built-in aggregate functions @@ -30,27 +33,27 @@ use strum_macros::EnumIter; // https://arrow.apache.org/datafusion/contributor-guide/index.html#how-to-add-a-new-aggregate-function #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash, EnumIter)] pub enum AggregateFunction { - /// count + /// Count Count, - /// sum + /// Sum Sum, - /// min + /// Minimum Min, - /// max + /// Maximum Max, - /// avg + /// Average Avg, - /// median + /// Median Median, - /// Approximate aggregate function + /// Approximate distinct function ApproxDistinct, - /// array_agg + /// Aggregation into an array ArrayAgg, - /// first_value + /// First value in a group according to some ordering FirstValue, - /// last_value + /// Last value in a group according to some ordering LastValue, - /// nth_value + /// N'th value in a group according to some ordering NthValue, /// Variance (Sample) Variance, @@ -102,7 +105,7 @@ pub enum AggregateFunction { BoolAnd, /// Bool Or BoolOr, - /// string_agg + /// String aggregation StringAgg, } @@ -428,6 +431,7 @@ impl AggregateFunction { #[cfg(test)] mod tests { use super::*; + use strum::IntoEnumIterator; #[test] diff --git a/datafusion/expr/src/type_coercion/aggregates.rs b/datafusion/expr/src/type_coercion/aggregates.rs index 917cd02554506..d78c5de0a8930 100644 --- a/datafusion/expr/src/type_coercion/aggregates.rs +++ b/datafusion/expr/src/type_coercion/aggregates.rs @@ -15,17 +15,16 @@ // specific language governing permissions and limitations // under the License. +use std::ops::Deref; + +use super::functions::can_coerce_from; +use crate::{AggregateFunction, Signature, TypeSignature}; + use arrow::datatypes::{ DataType, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE, DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, }; - use datafusion_common::{internal_err, plan_err, DataFusionError, Result}; -use std::ops::Deref; - -use crate::{AggregateFunction, Signature, TypeSignature}; - -use super::functions::can_coerce_from; pub static STRINGS: &[DataType] = &[DataType::Utf8, DataType::LargeUtf8]; @@ -595,6 +594,7 @@ pub fn is_string_agg_supported_arg_type(arg_type: &DataType) -> bool { #[cfg(test)] mod tests { use super::*; + use arrow::datatypes::DataType; #[test] diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 16f1feb12def5..78264678990ea 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -28,10 +28,9 @@ use crate::aggregate::utils::{down_cast_any_ref, ordering_fields}; use crate::expressions::format_state_name; use crate::{AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr}; -use arrow::array::ArrayRef; +use arrow::array::{Array, ArrayRef}; use arrow::datatypes::{DataType, Field}; use arrow_array::cast::AsArray; -use arrow_array::Array; use arrow_schema::{Fields, SortOptions}; use datafusion_common::utils::{compare_rows, get_row_at_idx}; use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue}; @@ -500,12 +499,14 @@ pub(crate) fn merge_ordered_arrays( #[cfg(test)] mod tests { + use std::sync::Arc; + use crate::aggregate::array_agg_ordered::merge_ordered_arrays; + use arrow_array::{Array, ArrayRef, Int64Array}; use arrow_schema::SortOptions; use datafusion_common::utils::get_row_at_idx; use datafusion_common::{Result, ScalarValue}; - use std::sync::Arc; #[test] fn test_merge_asc() -> Result<()> { diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs index 91560bc960b96..1a3d21fc40bc5 100644 --- a/datafusion/physical-expr/src/aggregate/build_in.rs +++ b/datafusion/physical-expr/src/aggregate/build_in.rs @@ -26,13 +26,15 @@ //! * Signature: see `Signature` //! * Return type: a function `(arg_types) -> return_type`. E.g. for min, ([f32]) -> f32, ([f64]) -> f64. +use std::sync::Arc; + use crate::aggregate::regr::RegrType; -use crate::expressions::Literal; -use crate::{expressions, AggregateExpr, PhysicalExpr, PhysicalSortExpr}; +use crate::expressions::{self, Literal}; +use crate::{AggregateExpr, PhysicalExpr, PhysicalSortExpr}; + use arrow::datatypes::Schema; -use datafusion_common::{not_impl_err, DataFusionError, Result}; -pub use datafusion_expr::AggregateFunction; -use std::sync::Arc; +use datafusion_common::{internal_err, not_impl_err, DataFusionError, Result}; +use datafusion_expr::AggregateFunction; /// Create a physical aggregation expression. /// This function errors when `input_phy_exprs`' can't be coerced to a valid argument type of the aggregation function. @@ -372,16 +374,19 @@ pub fn create_aggregate_expr( )), (AggregateFunction::NthValue, _) => { let expr = &input_phy_exprs[0]; - let n = input_phy_exprs[1] + let Some(n) = input_phy_exprs[1] .as_any() .downcast_ref::() - .unwrap() - .value(); - let n: i64 = n.clone().try_into()?; + .map(|literal| literal.value()) + else { + return internal_err!( + "Second argument of NTH_VALUE needs to be a literal" + ); + }; let nullable = expr.nullable(input_schema)?; Arc::new(expressions::NthValueAgg::new( expr.clone(), - n, + n.clone().try_into()?, name, input_phy_types[0].clone(), nullable, @@ -416,9 +421,9 @@ mod tests { BitAnd, BitOr, BitXor, BoolAnd, BoolOr, Correlation, Count, Covariance, DistinctArrayAgg, DistinctCount, Max, Min, Stddev, Sum, Variance, }; + use arrow::datatypes::{DataType, Field}; - use datafusion_common::plan_err; - use datafusion_common::ScalarValue; + use datafusion_common::{plan_err, ScalarValue}; use datafusion_expr::type_coercion::aggregates::NUMERICS; use datafusion_expr::{type_coercion, Signature}; diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index c7f9dd40ffad9..f46ec1b601894 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -24,31 +24,28 @@ use crate::expressions::format_state_name; use crate::{ reverse_order_bys, AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr, }; + use arrow_array::cast::AsArray; use arrow_array::ArrayRef; use arrow_schema::{DataType, Field, Fields}; use datafusion_common::utils::get_row_at_idx; -use datafusion_common::DataFusionError; -use datafusion_common::{exec_err, internal_err, Result, ScalarValue}; +use datafusion_common::{exec_err, internal_err, DataFusionError, Result, ScalarValue}; use datafusion_expr::Accumulator; -/// Expression for a ARRAY_AGG(ORDER BY) aggregation. -/// When aggregation works in multiple partitions -/// aggregations are split into multiple partitions, -/// then their results are merged. This aggregator -/// is a version of ARRAY_AGG that can support producing -/// intermediate aggregation (with necessary side information) -/// and that can merge aggregations from multiple partitions. +/// Expression for a `NTH_VALUE(... ORDER BY ..., ...)` aggregation. In a multi +/// partition setting, partial aggregations are computed for every partition, +/// and then their results are merged. #[derive(Debug)] pub struct NthValueAgg { /// Column name name: String, - /// The DataType for the input expression + /// The `DataType` for the input expression input_data_type: DataType, /// The input expression expr: Arc, + /// The `N` value. n: i64, - /// If the input expression can have NULLs + /// If the input expression can have `NULL`s nullable: bool, /// Ordering data types order_by_data_types: Vec, @@ -57,7 +54,7 @@ pub struct NthValueAgg { } impl NthValueAgg { - /// Create a new `OrderSensitiveArrayAgg` aggregate function + /// Create a new `NthValueAgg` aggregate function pub fn new( expr: Arc, n: i64, @@ -120,11 +117,7 @@ impl AggregateExpr for NthValueAgg { } fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { - if self.ordering_req.is_empty() { - None - } else { - Some(&self.ordering_req) - } + (!self.ordering_req.is_empty()).then_some(&self.ordering_req) } fn name(&self) -> &str { @@ -163,23 +156,23 @@ impl PartialEq for NthValueAgg { #[derive(Debug)] pub(crate) struct NthValueAccumulator { n: i64, - // `values` stores entries in the ARRAY_AGG result. + // `values` stores entries in the NTH_VALUE result. values: Vec, // `ordering_values` stores values of ordering requirement expression - // corresponding to each value in the ARRAY_AGG. + // corresponding to each value in the NTH_VALUE. // For each `ScalarValue` inside `values`, there will be a corresponding // `Vec` inside `ordering_values` which stores it ordering. // This information is used during merging results of the different partitions. // For detailed information how merging is done see [`merge_ordered_arrays`] ordering_values: Vec>, - // `datatypes` stores, datatype of expression inside ARRAY_AGG and ordering requirement expressions. + // Stores datatype of the expression inside NTH_VALUE and ordering requirement expressions. datatypes: Vec, // Stores ordering requirement of the Accumulator ordering_req: LexOrdering, } impl NthValueAccumulator { - /// Create a new order-sensitive ARRAY_AGG accumulator based on the given + /// Create a new order-sensitive NTH_VALUE accumulator based on the given /// item data type. pub fn try_new( n: i64, @@ -238,7 +231,7 @@ impl Accumulator for NthValueAccumulator { if self.ordering_req.is_empty() { let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - // Stores ARRAY_AGG results coming from each partition + // Stores NTH_VALUE results coming from each partition let mut partition_values = vec![]; // Existing values should be merged also. partition_values.extend(self.values.clone()); @@ -251,11 +244,11 @@ impl Accumulator for NthValueAccumulator { } self.values = partition_values; } else if let Some(agg_orderings) = states[1].as_list_opt::() { - // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. - // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores + // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside NTH_VALUE list. + // For each `StructArray` inside NTH_VALUE list, we will receive an `Array` that stores // values received from its ordering requirement expression. (This information is necessary for during merging). - // Stores ARRAY_AGG results coming from each partition + // Stores NTH_VALUE results coming from each partition let mut partition_values = vec![]; // Stores ordering requirement expression results coming from each partition let mut partition_ordering_values = vec![]; @@ -308,9 +301,9 @@ impl Accumulator for NthValueAccumulator { } fn state(&self) -> Result> { - let mut result = vec![self.evaluate_values()?]; + let mut result = vec![self.evaluate_values()]; if !self.ordering_req.is_empty() { - result.push(self.evaluate_orderings()?); + result.push(self.evaluate_orderings()); } Ok(result) } @@ -327,8 +320,7 @@ impl Accumulator for NthValueAccumulator { self.values.len().checked_sub(n_required) }; if let Some(idx) = nth_value_idx { - let nth_value = self.values[idx].clone(); - Ok(nth_value.clone()) + Ok(self.values[idx].clone()) } else { ScalarValue::try_from(self.datatypes[0].clone()) } @@ -360,27 +352,25 @@ impl Accumulator for NthValueAccumulator { } impl NthValueAccumulator { - fn evaluate_orderings(&self) -> Result { + fn evaluate_orderings(&self) -> ScalarValue { let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]); let struct_field = Fields::from(fields.clone()); - let orderings: Vec = self + let orderings = self .ordering_values .iter() .map(|ordering| { ScalarValue::Struct(Some(ordering.clone()), struct_field.clone()) }) - .collect(); + .collect::>(); let struct_type = DataType::Struct(Fields::from(fields)); // Wrap in List, so we have the same data structure ListArray(StructArray..) for group by cases - let arr = ScalarValue::new_list(&orderings, &struct_type); - Ok(ScalarValue::List(arr)) + ScalarValue::List(ScalarValue::new_list(&orderings, &struct_type)) } - fn evaluate_values(&self) -> Result { - let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); - Ok(ScalarValue::List(arr)) + fn evaluate_values(&self) -> ScalarValue { + ScalarValue::List(ScalarValue::new_list(&self.values, &self.datatypes[0])) } fn append_new_data( diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 22b252368a62c..110c5aa7d2c84 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -25,7 +25,6 @@ use crate::aggregates::{ no_grouping::AggregateStream, row_hash::GroupedHashAggregateStream, topk_stream::GroupedTopKAggregateStream, }; - use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::windows::get_ordered_partition_by_indices; use crate::{ diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 905c1e84c82a4..c6cfaeea8c140 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use crate::protobuf::{ self, plan_type::PlanTypeEnum::{ @@ -26,6 +28,7 @@ use crate::protobuf::{ AnalyzedLogicalPlanType, CubeNode, GroupingSetNode, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, PlaceholderNode, RollupNode, }; + use arrow::{ array::AsArray, buffer::Buffer, @@ -44,15 +47,15 @@ use datafusion_common::{ use datafusion_expr::window_frame::{check_window_frame, regularize_window_order_by}; use datafusion_expr::{ abs, acos, acosh, array, array_append, array_concat, array_dims, array_distinct, - array_element, array_except, array_has, array_has_all, array_has_any, - array_intersect, array_length, array_ndims, array_position, array_positions, - array_prepend, array_remove, array_remove_all, array_remove_n, array_repeat, - array_replace, array_replace_all, array_replace_n, array_slice, array_sort, - array_to_string, arrow_typeof, ascii, asin, asinh, atan, atan2, atanh, bit_length, - btrim, cardinality, cbrt, ceil, character_length, chr, coalesce, concat_expr, - concat_ws_expr, cos, cosh, cot, current_date, current_time, date_bin, date_part, - date_trunc, decode, degrees, digest, encode, exp, - expr::{self, InList, Sort, WindowFunction}, + array_element, array_empty, array_except, array_has, array_has_all, array_has_any, + array_intersect, array_length, array_ndims, array_pop_back, array_pop_front, + array_position, array_positions, array_prepend, array_remove, array_remove_all, + array_remove_n, array_repeat, array_replace, array_replace_all, array_replace_n, + array_slice, array_sort, array_to_string, arrow_typeof, ascii, asin, asinh, atan, + atan2, atanh, bit_length, btrim, cardinality, cbrt, ceil, character_length, chr, + coalesce, concat_expr, concat_ws_expr, cos, cosh, cot, current_date, current_time, + date_bin, date_part, date_trunc, decode, degrees, digest, encode, exp, + expr::{self, Alias, InList, Placeholder, Sort, WindowFunction}, factorial, find_in_set, flatten, floor, from_unixtime, gcd, gen_range, isnan, iszero, lcm, left, levenshtein, ln, log, log10, log2, logical_plan::{PlanType, StringifiedPlan}, @@ -68,11 +71,6 @@ use datafusion_expr::{ JoinConstraint, JoinType, Like, Operator, TryCast, WindowFrame, WindowFrameBound, WindowFrameUnits, }; -use datafusion_expr::{ - array_empty, array_pop_back, array_pop_front, - expr::{Alias, Placeholder}, -}; -use std::sync::Arc; #[derive(Debug)] pub enum Error { diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 1064f7be5b39c..9f9f4614f9832 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -31,6 +31,7 @@ use crate::protobuf::{ AnalyzedLogicalPlanType, CubeNode, EmptyMessage, GroupingSetNode, LogicalExprList, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, PlaceholderNode, RollupNode, }; + use arrow::{ array::ArrayRef, datatypes::{ From 6a23dd259bfb244f3d06e34c8ce60f1b330a2258 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Fri, 12 Jan 2024 16:33:09 +0300 Subject: [PATCH 20/31] Update comments --- datafusion/physical-expr/src/aggregate/nth_value.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index f46ec1b601894..b25efa9fc7d14 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +//! Defines NTH_VALUE aggregate expression which may specify ordering requirement +//! that can evaluated at runtime during query execution + use std::any::Any; use std::sync::Arc; @@ -161,7 +164,7 @@ pub(crate) struct NthValueAccumulator { // `ordering_values` stores values of ordering requirement expression // corresponding to each value in the NTH_VALUE. // For each `ScalarValue` inside `values`, there will be a corresponding - // `Vec` inside `ordering_values` which stores it ordering. + // `Vec` inside `ordering_values` which stores its ordering. // This information is used during merging results of the different partitions. // For detailed information how merging is done see [`merge_ordered_arrays`] ordering_values: Vec>, @@ -197,6 +200,8 @@ impl NthValueAccumulator { } impl Accumulator for NthValueAccumulator { + /// Updates its state with the `values`. Assumes data in the `values` satisfies the required + /// ordering for the accumulator (across consecutive batches, not just batch-wise). fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { if values.is_empty() { return Ok(()); @@ -373,6 +378,8 @@ impl NthValueAccumulator { ScalarValue::List(ScalarValue::new_list(&self.values, &self.datatypes[0])) } + /// Updates state, with the `values`. Fetch contains missing number of entries for state to be complete + /// None represents all of the new `values` need to be added to the state. fn append_new_data( &mut self, values: &[ArrayRef], From fa709977b4dc59ffab13846de82a0b884b8d9ee1 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Fri, 12 Jan 2024 17:11:55 +0300 Subject: [PATCH 21/31] Use drain method to remove from the beginning. --- datafusion/physical-expr/src/aggregate/nth_value.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index b25efa9fc7d14..8cd72c4737bca 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -218,8 +218,8 @@ impl Accumulator for NthValueAccumulator { self.append_new_data(values, None)?; let start_offset = self.values.len().saturating_sub(n_required); if start_offset > 0 { - self.values = self.values[start_offset..].to_vec(); - self.ordering_values = self.ordering_values[start_offset..].to_vec(); + self.values.drain(0..start_offset); + self.ordering_values.drain(0..start_offset); } } From 2275582e5e5a21e8e6d0312fdbf0745a0e26ca3c Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 15 Jan 2024 17:58:00 +0300 Subject: [PATCH 22/31] Add reverse support, convert buffer to vecdeque --- datafusion/common/src/scalar.rs | 14 +++- .../src/aggregate/array_agg_ordered.rs | 78 +++++++++++++------ .../physical-expr/src/aggregate/nth_value.rs | 46 ++++++----- .../physical-plan/src/aggregates/mod.rs | 62 +++++++++++---- .../sqllogictest/test_files/groupby.slt | 36 +++++++++ 5 files changed, 178 insertions(+), 58 deletions(-) diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs index 8820ca9942fca..2ea7f1d6000ba 100644 --- a/datafusion/common/src/scalar.rs +++ b/datafusion/common/src/scalar.rs @@ -19,7 +19,7 @@ use std::borrow::Borrow; use std::cmp::Ordering; -use std::collections::HashSet; +use std::collections::{HashSet, VecDeque}; use std::convert::{Infallible, TryInto}; use std::hash::Hash; use std::str::FromStr; @@ -2627,6 +2627,18 @@ impl ScalarValue { .sum::() } + /// Estimates [size](Self::size) of [`VecDeque`] in bytes. + /// + /// Includes the size of the [`VecDeque`] container itself. + pub fn size_of_vec_deque(vec_deque: &VecDeque) -> usize { + std::mem::size_of_val(vec_deque) + + (std::mem::size_of::() * vec_deque.capacity()) + + vec_deque + .iter() + .map(|sv| sv.size() - std::mem::size_of_val(sv)) + .sum::() + } + /// Estimates [size](Self::size) of [`HashSet`] in bytes. /// /// Includes the size of the [`HashSet`] container itself. diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 78264678990ea..d9656c268942a 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -26,7 +26,9 @@ use std::sync::Arc; use crate::aggregate::utils::{down_cast_any_ref, ordering_fields}; use crate::expressions::format_state_name; -use crate::{AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr}; +use crate::{ + reverse_order_bys, AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr, +}; use arrow::array::{Array, ArrayRef}; use arrow::datatypes::{DataType, Field}; @@ -59,6 +61,8 @@ pub struct OrderSensitiveArrayAgg { order_by_data_types: Vec, /// Ordering requirement ordering_req: LexOrdering, + /// direction + from_start: bool, } impl OrderSensitiveArrayAgg { @@ -78,6 +82,7 @@ impl OrderSensitiveArrayAgg { nullable, order_by_data_types, ordering_req, + from_start: true, } } } @@ -101,6 +106,7 @@ impl AggregateExpr for OrderSensitiveArrayAgg { &self.input_data_type, &self.order_by_data_types, self.ordering_req.clone(), + self.from_start, )?)) } @@ -134,6 +140,19 @@ impl AggregateExpr for OrderSensitiveArrayAgg { fn name(&self) -> &str { &self.name } + + fn reverse_expr(&self) -> Option> { + Some(Arc::new(Self { + name: self.name.to_string(), + input_data_type: self.input_data_type.clone(), + expr: self.expr.clone(), + nullable: self.nullable, + order_by_data_types: self.order_by_data_types.clone(), + // reverse requirement + ordering_req: reverse_order_bys(&self.ordering_req), + from_start: !self.from_start, + }) as _) + } } impl PartialEq for OrderSensitiveArrayAgg { @@ -165,6 +184,8 @@ pub(crate) struct OrderSensitiveArrayAggAccumulator { datatypes: Vec, // Stores ordering requirement of the Accumulator ordering_req: LexOrdering, + // Whether array direction is from start or not. + from_start: bool, } impl OrderSensitiveArrayAggAccumulator { @@ -174,6 +195,7 @@ impl OrderSensitiveArrayAggAccumulator { datatype: &DataType, ordering_dtypes: &[DataType], ordering_req: LexOrdering, + from_start: bool, ) -> Result { let mut datatypes = vec![datatype.clone()]; datatypes.extend(ordering_dtypes.iter().cloned()); @@ -182,6 +204,7 @@ impl OrderSensitiveArrayAggAccumulator { ordering_values: vec![], datatypes, ordering_req, + from_start, }) } } @@ -215,37 +238,39 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { if let Some(agg_orderings) = agg_orderings.as_list_opt::() { // Stores ARRAY_AGG results coming from each partition - let mut partition_values = vec![]; + let mut partition_values: Vec<&[ScalarValue]> = vec![]; // Stores ordering requirement expression results coming from each partition - let mut partition_ordering_values = vec![]; + let mut partition_ordering_values: Vec<&[Vec]> = vec![]; // Existing values should be merged also. - partition_values.push(self.values.clone()); - partition_ordering_values.push(self.ordering_values.clone()); + partition_values.push(&self.values); + partition_ordering_values.push(&self.ordering_values); let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - for v in array_agg_res.into_iter() { + for v in array_agg_res.iter() { partition_values.push(v); } let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; - for partition_ordering_rows in orderings.into_iter() { + let ordering_values = orderings.into_iter().map(|partition_ordering_rows| { // Extract value from struct to ordering_rows for each group/partition let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { - if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { - Ok(ordering_columns_per_row) - } else { - exec_err!( + if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { + Ok(ordering_columns_per_row) + } else { + exec_err!( "Expects to receive ScalarValue::Struct(Some(..), _) but got:{:?}", ordering_row.data_type() ) - } - }).collect::>>()?; - - partition_ordering_values.push(ordering_value); + } + }).collect::>>()?; + Ok(ordering_value) + }).collect::>>()?; + for ordering_values in ordering_values.iter() { + partition_ordering_values.push(ordering_values); } let sort_options = self @@ -273,8 +298,15 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { } fn evaluate(&self) -> Result { - let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); - Ok(ScalarValue::List(arr)) + if self.from_start { + let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); + Ok(ScalarValue::List(arr)) + } else { + let mut reverse_values = self.values.clone(); + reverse_values.reverse(); + let arr = ScalarValue::new_list(&reverse_values, &self.datatypes[0]); + Ok(ScalarValue::List(arr)) + } } fn size(&self) -> usize { @@ -412,11 +444,11 @@ impl<'a> PartialOrd for CustomElement<'a> { /// Inner `Vec`s of the `ordering_values` will be compared according `sort_options` (Their sizes should match) pub(crate) fn merge_ordered_arrays( // We will merge values into single `Vec`. - values: &[Vec], + values: &[&[ScalarValue]], // `values` will be merged according to `ordering_values`. // Inner `Vec` can be thought as ordering information for the // each `ScalarValue` in the values`. - ordering_values: &[Vec>], + ordering_values: &[&[Vec]], // Defines according to which ordering comparisons should be done. sort_options: &[SortOptions], ) -> Result<(Vec, Vec>)> { @@ -555,8 +587,8 @@ mod tests { ]; let (merged_vals, merged_ts) = merge_ordered_arrays( - &[lhs_vals, rhs_vals], - &[lhs_orderings, rhs_orderings], + &[&lhs_vals, &rhs_vals], + &[&lhs_orderings, &rhs_orderings], &sort_options, )?; let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?; @@ -622,8 +654,8 @@ mod tests { Arc::new(Int64Array::from(vec![4, 4, 3, 3, 2, 2, 1, 1, 0, 0])) as ArrayRef, ]; let (merged_vals, merged_ts) = merge_ordered_arrays( - &[lhs_vals, rhs_vals], - &[lhs_orderings, rhs_orderings], + &[&lhs_vals, &rhs_vals], + &[&lhs_orderings, &rhs_orderings], &sort_options, )?; let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?; diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index 8cd72c4737bca..2b8ac2d516c55 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -19,6 +19,7 @@ //! that can evaluated at runtime during query execution use std::any::Any; +use std::collections::VecDeque; use std::sync::Arc; use crate::aggregate::array_agg_ordered::merge_ordered_arrays; @@ -160,14 +161,14 @@ impl PartialEq for NthValueAgg { pub(crate) struct NthValueAccumulator { n: i64, // `values` stores entries in the NTH_VALUE result. - values: Vec, + values: VecDeque, // `ordering_values` stores values of ordering requirement expression // corresponding to each value in the NTH_VALUE. // For each `ScalarValue` inside `values`, there will be a corresponding // `Vec` inside `ordering_values` which stores its ordering. // This information is used during merging results of the different partitions. // For detailed information how merging is done see [`merge_ordered_arrays`] - ordering_values: Vec>, + ordering_values: VecDeque>, // Stores datatype of the expression inside NTH_VALUE and ordering requirement expressions. datatypes: Vec, // Stores ordering requirement of the Accumulator @@ -191,8 +192,8 @@ impl NthValueAccumulator { datatypes.extend(ordering_dtypes.iter().cloned()); Ok(Self { n, - values: vec![], - ordering_values: vec![], + values: VecDeque::new(), + ordering_values: VecDeque::new(), datatypes, ordering_req, }) @@ -237,7 +238,7 @@ impl Accumulator for NthValueAccumulator { let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; // Stores NTH_VALUE results coming from each partition - let mut partition_values = vec![]; + let mut partition_values = VecDeque::new(); // Existing values should be merged also. partition_values.extend(self.values.clone()); for v in array_agg_res.into_iter() { @@ -254,24 +255,27 @@ impl Accumulator for NthValueAccumulator { // values received from its ordering requirement expression. (This information is necessary for during merging). // Stores NTH_VALUE results coming from each partition - let mut partition_values = vec![]; + let mut partition_values: Vec<&[ScalarValue]> = vec![]; // Stores ordering requirement expression results coming from each partition - let mut partition_ordering_values = vec![]; + let mut partition_ordering_values: Vec<&[Vec]> = vec![]; // Existing values should be merged also. - partition_values.push(self.values.clone()); - partition_ordering_values.push(self.ordering_values.clone()); + let partition_value = self.values.make_contiguous(); + partition_values.push(partition_value); + + let partition_ordering_value = self.ordering_values.make_contiguous(); + partition_ordering_values.push(partition_ordering_value); let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - for v in array_agg_res.into_iter() { + for v in array_agg_res.iter() { partition_values.push(v); } let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; - for partition_ordering_rows in orderings.into_iter() { + let ordering_values = orderings.into_iter().map(|partition_ordering_rows| { // Extract value from struct to ordering_rows for each group/partition let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { @@ -283,8 +287,10 @@ impl Accumulator for NthValueAccumulator { ) } }).collect::>>()?; - - partition_ordering_values.push(ordering_value); + Ok(ordering_value) + }).collect::>>()?; + for ordering_values in ordering_values.iter() { + partition_ordering_values.push(ordering_values); } let sort_options = self @@ -297,8 +303,8 @@ impl Accumulator for NthValueAccumulator { &partition_ordering_values, &sort_options, )?; - self.values = new_values; - self.ordering_values = new_orderings; + self.values = new_values.into(); + self.ordering_values = new_orderings.into(); } else { return exec_err!("Expects to receive a list array"); } @@ -333,7 +339,7 @@ impl Accumulator for NthValueAccumulator { fn size(&self) -> usize { let mut total = std::mem::size_of_val(self) - + ScalarValue::size_of_vec(&self.values) + + ScalarValue::size_of_vec_deque(&self.values) - std::mem::size_of_val(&self.values); // Add size of the `self.ordering_values` @@ -375,7 +381,9 @@ impl NthValueAccumulator { } fn evaluate_values(&self) -> ScalarValue { - ScalarValue::List(ScalarValue::new_list(&self.values, &self.datatypes[0])) + let mut values_cloned = self.values.clone(); + let values_slice = values_cloned.make_contiguous(); + ScalarValue::List(ScalarValue::new_list(values_slice, &self.datatypes[0])) } /// Updates state, with the `values`. Fetch contains missing number of entries for state to be complete @@ -393,8 +401,8 @@ impl NthValueAccumulator { }; for index in 0..n_to_add { let row = get_row_at_idx(values, index)?; - self.values.push(row[0].clone()); - self.ordering_values.push(row[1..].to_vec()); + self.values.push_back(row[0].clone()); + self.ordering_values.push_back(row[1..].to_vec()); } Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 110c5aa7d2c84..0b5042e8570aa 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -903,6 +903,10 @@ fn get_aggregate_exprs_requirement( let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req); let reverse_aggr_req = PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); + + // Get absolutely necessary requirement for the aggregator. + let aggr_hard_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); + let reverse_aggr_hard_req = reverse_order_bys(&aggr_hard_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( @@ -925,7 +929,9 @@ fn get_aggregate_exprs_requirement( first_value = first_value.with_requirement_satisfied(false); *aggr_expr = Arc::new(first_value) as _; } - } else if let Some(last_value) = aggr_expr.as_any().downcast_ref::() { + 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, @@ -947,11 +953,41 @@ fn get_aggregate_exprs_requirement( last_value = last_value.with_requirement_satisfied(false); *aggr_expr = Arc::new(last_value) as _; } - } else if let Some(finer_ordering) = + continue; + } + if aggr_hard_req.is_empty() { + continue; + } + if eq_properties.ordering_satisfy(&aggr_hard_req) { + if let Some(finer_ordering) = + finer_ordering(&requirement, aggr_expr, group_by, eq_properties, agg_mode) + { + requirement = finer_ordering; + continue; + } + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + if eq_properties.ordering_satisfy(&reverse_aggr_hard_req) { + if let Some(finer_ordering) = finer_ordering( + &requirement, + &reverse_aggr_expr, + group_by, + eq_properties, + agg_mode, + ) { + 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) { requirement = finer_ordering; - } else if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + continue; + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { if let Some(finer_ordering) = finer_ordering( &requirement, &reverse_aggr_expr, @@ -959,21 +995,17 @@ fn get_aggregate_exprs_requirement( eq_properties, agg_mode, ) { - *aggr_expr = reverse_aggr_expr; requirement = finer_ordering; - } else { - return not_impl_err!( - "Conflicting ordering requirements in aggregate functions is not supported" - ); + *aggr_expr = reverse_aggr_expr; + continue; } - } else { - // If neither of the requirements 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" - ); } + // If neither of the requirements 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/sqllogictest/test_files/groupby.slt b/datafusion/sqllogictest/test_files/groupby.slt index 3af17a7d08504..1676a90f0a8c3 100644 --- a/datafusion/sqllogictest/test_files/groupby.slt +++ b/datafusion/sqllogictest/test_files/groupby.slt @@ -4332,6 +4332,42 @@ ORDER BY a, b; 1 2 51 52 73 72 1 3 76 77 98 97 +# we should be able to reverse array agg requirement, if it helps to remove a SortExec from plan. +query TT +EXPLAIN SELECT a, b, ARRAY_AGG(c ORDER BY c DESC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +logical_plan +Sort: multiple_ordered_table.a ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +--Aggregate: groupBy=[[multiple_ordered_table.a, multiple_ordered_table.b]], aggr=[[ARRAY_AGG(multiple_ordered_table.c) ORDER BY [multiple_ordered_table.c DESC NULLS FIRST]]] +----TableScan: multiple_ordered_table projection=[a, b, c] +physical_plan +AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[ARRAY_AGG(multiple_ordered_table.c)], ordering_mode=Sorted +--CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], has_header=true + +query II? +SELECT a, b, ARRAY_AGG(c ORDER BY c DESC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 [24, 23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0] +0 1 [49, 48, 47, 46, 45, 44, 43, 42, 41, 40, 39, 38, 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25] +1 2 [74, 73, 72, 71, 70, 69, 68, 67, 66, 65, 64, 63, 62, 61, 60, 59, 58, 57, 56, 55, 54, 53, 52, 51, 50] +1 3 [99, 98, 97, 96, 95, 94, 93, 92, 91, 90, 89, 88, 87, 86, 85, 84, 83, 82, 81, 80, 79, 78, 77, 76, 75] + +query II?II +SELECT a, b, ARRAY_AGG(d ORDER BY d DESC), NTH_VALUE(d, 1 ORDER BY d DESC), NTH_VALUE(d, 1 ORDER BY d ASC) +FROM multiple_ordered_table +GROUP BY a, b +ORDER BY a, b; +---- +0 0 [4, 4, 4, 4, 3, 3, 3, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] 4 0 +0 1 [4, 4, 4, 3, 3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0] 4 0 +1 2 [4, 4, 4, 4, 4, 4, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0] 4 0 +1 3 [4, 4, 4, 4, 4, 4, 3, 3, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] 4 0 # increase partition to 8 statement ok From 7138d7f88b96a9583cd33c1f7f265d1d53208a95 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 15 Jan 2024 18:07:10 +0300 Subject: [PATCH 23/31] Minor changes --- .../physical-plan/src/aggregates/mod.rs | 39 ++++++++++--------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0b5042e8570aa..5feeafc225c63 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -904,9 +904,6 @@ fn get_aggregate_exprs_requirement( let reverse_aggr_req = PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req); - // Get absolutely necessary requirement for the aggregator. - let aggr_hard_req = get_aggregate_expr_req(aggr_expr, group_by, agg_mode); - let reverse_aggr_hard_req = reverse_order_bys(&aggr_hard_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( @@ -955,26 +952,26 @@ fn get_aggregate_exprs_requirement( } continue; } - if aggr_hard_req.is_empty() { - continue; - } - if eq_properties.ordering_satisfy(&aggr_hard_req) { - if let Some(finer_ordering) = - finer_ordering(&requirement, 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) { + // Requirement is satisfied by existing ordering requirement = finer_ordering; continue; } } if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { - if eq_properties.ordering_satisfy(&reverse_aggr_hard_req) { - 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, + &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; @@ -984,6 +981,8 @@ fn get_aggregate_exprs_requirement( 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; } @@ -995,12 +994,14 @@ fn get_aggregate_exprs_requirement( 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; } } - // If neither of the requirements satisfy the other, this means + // 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!( From ebf5c2c6291a620c59e1b22595a7165f16da7819 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 15 Jan 2024 18:16:00 +0300 Subject: [PATCH 24/31] Minor changes --- .../proto/src/logical_plan/from_proto.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index b124bc52b32c9..ccb49f0e5b0f6 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -44,17 +44,19 @@ use datafusion_common::{ Constraints, DFField, DFSchema, DFSchemaRef, DataFusionError, OwnedTableReference, Result, ScalarValue, }; +use datafusion_expr::expr::{Alias, Placeholder}; use datafusion_expr::window_frame::{check_window_frame, regularize_window_order_by}; use datafusion_expr::{ abs, acos, acosh, array, array_append, array_concat, array_dims, array_distinct, - array_element, array_except, array_has, array_has_all, array_has_any, - array_intersect, array_length, array_ndims, array_position, array_positions, - array_prepend, array_remove, array_remove_all, array_remove_n, array_repeat, - array_replace, array_replace_all, array_replace_n, array_resize, array_slice, - array_sort, array_to_string, array_union, arrow_typeof, ascii, asin, asinh, atan, - atan2, atanh, bit_length, btrim, cardinality, cbrt, ceil, character_length, chr, - coalesce, concat_expr, concat_ws_expr, cos, cosh, cot, current_date, current_time, - date_bin, date_part, date_trunc, decode, degrees, digest, encode, exp, + array_element, array_empty, array_except, array_has, array_has_all, array_has_any, + array_intersect, array_length, array_ndims, array_pop_back, array_pop_front, + array_position, array_positions, array_prepend, array_remove, array_remove_all, + array_remove_n, array_repeat, array_replace, array_replace_all, array_replace_n, + array_resize, array_slice, array_sort, array_to_string, array_union, arrow_typeof, + ascii, asin, asinh, atan, atan2, atanh, bit_length, btrim, cardinality, cbrt, ceil, + character_length, chr, coalesce, concat_expr, concat_ws_expr, cos, cosh, cot, + current_date, current_time, date_bin, date_part, date_trunc, decode, degrees, digest, + encode, exp, expr::{self, InList, Sort, WindowFunction}, factorial, find_in_set, flatten, floor, from_unixtime, gcd, gen_range, isnan, iszero, lcm, left, levenshtein, ln, log, log10, log2, From 4cab3dba80905a4910c7a3b1b71cadf14049d805 Mon Sep 17 00:00:00 2001 From: Mehmet Ozan Kabak Date: Mon, 15 Jan 2024 23:58:22 +0300 Subject: [PATCH 25/31] Review Part 2 --- datafusion/common/src/scalar.rs | 32 +- .../src/aggregate/array_agg_ordered.rs | 274 +++++++++--------- .../physical-expr/src/aggregate/nth_value.rs | 17 +- 3 files changed, 155 insertions(+), 168 deletions(-) diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs index 2ea7f1d6000ba..2e067899e1830 100644 --- a/datafusion/common/src/scalar.rs +++ b/datafusion/common/src/scalar.rs @@ -20,10 +20,12 @@ use std::borrow::Borrow; use std::cmp::Ordering; use std::collections::{HashSet, VecDeque}; -use std::convert::{Infallible, TryInto}; +use std::convert::{Infallible, TryFrom, TryInto}; +use std::fmt; use std::hash::Hash; +use std::iter::repeat; use std::str::FromStr; -use std::{convert::TryFrom, fmt, iter::repeat, sync::Arc}; +use std::sync::Arc; use crate::arrow_datafusion_err; use crate::cast::{ @@ -33,23 +35,22 @@ use crate::cast::{ use crate::error::{DataFusionError, Result, _internal_err, _not_impl_err}; use crate::hash_utils::create_hashes; use crate::utils::{array_into_large_list_array, array_into_list_array}; + use arrow::compute::kernels::numeric::*; -use arrow::datatypes::{i256, Fields, SchemaBuilder}; use arrow::util::display::{ArrayFormatter, FormatOptions}; use arrow::{ array::*, compute::kernels::cast::{cast_with_options, CastOptions}, datatypes::{ - ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Float32Type, Int16Type, - Int32Type, Int64Type, Int8Type, IntervalDayTimeType, IntervalMonthDayNanoType, - IntervalUnit, IntervalYearMonthType, TimeUnit, TimestampMicrosecondType, + i256, ArrowDictionaryKeyType, ArrowNativeType, ArrowTimestampType, DataType, + Field, Fields, Float32Type, Int16Type, Int32Type, Int64Type, Int8Type, + IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, + IntervalYearMonthType, SchemaBuilder, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type, DECIMAL128_MAX_PRECISION, }, }; use arrow_array::cast::as_list_array; -use arrow_array::types::ArrowTimestampType; -use arrow_array::{ArrowNativeTypeOp, Scalar}; /// A dynamically typed, nullable single value, (the single-valued counter-part /// to arrow's [`Array`]) @@ -3164,22 +3165,19 @@ impl ScalarType for TimestampNanosecondType { #[cfg(test)] mod tests { - use super::*; - use std::cmp::Ordering; use std::sync::Arc; - use chrono::NaiveDate; - use rand::Rng; + use super::*; + use crate::cast::{as_string_array, as_uint32_array, as_uint64_array}; use arrow::buffer::OffsetBuffer; - use arrow::compute::kernels; - use arrow::compute::{concat, is_null}; - use arrow::datatypes::ArrowPrimitiveType; + use arrow::compute::{concat, is_null, kernels}; + use arrow::datatypes::{ArrowNumericType, ArrowPrimitiveType}; use arrow::util::pretty::pretty_format_columns; - use arrow_array::ArrowNumericType; - use crate::cast::{as_string_array, as_uint32_array, as_uint64_array}; + use chrono::NaiveDate; + use rand::Rng; #[test] fn test_to_array_of_size_for_list() { diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index d9656c268942a..088c83bfc3b2b 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -40,29 +40,25 @@ use datafusion_expr::Accumulator; use itertools::izip; -/// Expression for a ARRAY_AGG(ORDER BY) aggregation. -/// When aggregation works in multiple partitions -/// aggregations are split into multiple partitions, -/// then their results are merged. This aggregator -/// is a version of ARRAY_AGG that can support producing -/// intermediate aggregation (with necessary side information) -/// and that can merge aggregations from multiple partitions. +/// Expression for a `ARRAY_AGG(... ORDER BY ..., ...)` aggregation. In a multi +/// partition setting, partial aggregations are computed for every partition, +/// and then their results are merged. #[derive(Debug)] pub struct OrderSensitiveArrayAgg { /// Column name name: String, - /// The DataType for the input expression + /// The `DataType` for the input expression input_data_type: DataType, /// The input expression expr: Arc, - /// If the input expression can have NULLs + /// If the input expression can have `NULL`s nullable: bool, /// Ordering data types order_by_data_types: Vec, /// Ordering requirement ordering_req: LexOrdering, - /// direction - from_start: bool, + /// Whether the aggregation is running in reverse + reverse: bool, } impl OrderSensitiveArrayAgg { @@ -82,7 +78,7 @@ impl OrderSensitiveArrayAgg { nullable, order_by_data_types, ordering_req, - from_start: true, + reverse: false, } } } @@ -102,12 +98,13 @@ impl AggregateExpr for OrderSensitiveArrayAgg { } fn create_accumulator(&self) -> Result> { - Ok(Box::new(OrderSensitiveArrayAggAccumulator::try_new( + OrderSensitiveArrayAggAccumulator::try_new( &self.input_data_type, &self.order_by_data_types, self.ordering_req.clone(), - self.from_start, - )?)) + self.reverse, + ) + .map(|acc| Box::new(acc) as _) } fn state_fields(&self) -> Result> { @@ -130,11 +127,7 @@ impl AggregateExpr for OrderSensitiveArrayAgg { } fn order_bys(&self) -> Option<&[PhysicalSortExpr]> { - if self.ordering_req.is_empty() { - None - } else { - Some(&self.ordering_req) - } + (!self.ordering_req.is_empty()).then_some(&self.ordering_req) } fn name(&self) -> &str { @@ -148,10 +141,10 @@ impl AggregateExpr for OrderSensitiveArrayAgg { expr: self.expr.clone(), nullable: self.nullable, order_by_data_types: self.order_by_data_types.clone(), - // reverse requirement + // Reverse requirement: ordering_req: reverse_order_bys(&self.ordering_req), - from_start: !self.from_start, - }) as _) + reverse: !self.reverse, + })) } } @@ -171,21 +164,20 @@ impl PartialEq for OrderSensitiveArrayAgg { #[derive(Debug)] pub(crate) struct OrderSensitiveArrayAggAccumulator { - // `values` stores entries in the ARRAY_AGG result. + /// Stores entries in the `ARRAY_AGG` result. values: Vec, - // `ordering_values` stores values of ordering requirement expression - // corresponding to each value in the ARRAY_AGG. - // For each `ScalarValue` inside `values`, there will be a corresponding - // `Vec` inside `ordering_values` which stores it ordering. - // This information is used during merging results of the different partitions. - // For detailed information how merging is done see [`merge_ordered_arrays`] + /// Stores values of ordering requirement expressions corresponding to each + /// entry in `values`. This information is used when merging results from + /// different partitions. For detailed information how merging is done, see + /// [`merge_ordered_arrays`]. ordering_values: Vec>, - // `datatypes` stores, datatype of expression inside ARRAY_AGG and ordering requirement expressions. + /// Stores datatypes of expressions inside values and ordering requirement + /// expressions. datatypes: Vec, - // Stores ordering requirement of the Accumulator + /// Stores the ordering requirement of the `Accumulator`. ordering_req: LexOrdering, - // Whether array direction is from start or not. - from_start: bool, + /// Whether the aggregation is running in reverse. + reverse: bool, } impl OrderSensitiveArrayAggAccumulator { @@ -195,7 +187,7 @@ impl OrderSensitiveArrayAggAccumulator { datatype: &DataType, ordering_dtypes: &[DataType], ordering_req: LexOrdering, - from_start: bool, + reverse: bool, ) -> Result { let mut datatypes = vec![datatype.clone()]; datatypes.extend(ordering_dtypes.iter().cloned()); @@ -204,7 +196,7 @@ impl OrderSensitiveArrayAggAccumulator { ordering_values: vec![], datatypes, ordering_req, - from_start, + reverse, }) } } @@ -229,65 +221,63 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { if states.is_empty() { return Ok(()); } - // First entry in the state is the aggregation result. - let array_agg_values = &states[0]; - // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside ARRAY_AGG list. - // For each `StructArray` inside ARRAY_AGG list, we will receive an `Array` that stores - // values received from its ordering requirement expression. (This information is necessary for during merging). - let agg_orderings = &states[1]; - - if let Some(agg_orderings) = agg_orderings.as_list_opt::() { - // Stores ARRAY_AGG results coming from each partition - let mut partition_values: Vec<&[ScalarValue]> = vec![]; - // Stores ordering requirement expression results coming from each partition - let mut partition_ordering_values: Vec<&[Vec]> = vec![]; - - // Existing values should be merged also. - partition_values.push(&self.values); - partition_ordering_values.push(&self.ordering_values); - - let array_agg_res = - ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - - for v in array_agg_res.iter() { - partition_values.push(v); - } + // First entry in the state is the aggregation result. Second entry + // stores values received for ordering requirement columns for each + // aggregation value inside `ARRAY_AGG` list. For each `StructArray` + // inside `ARRAY_AGG` list, we will receive an `Array` that stores values + // received from its ordering requirement expression. (This information + // is necessary for during merging). + let [array_agg_values, agg_orderings, ..] = &states else { + return exec_err!("State should have two elements"); + }; + let Some(agg_orderings) = agg_orderings.as_list_opt::() else { + return exec_err!("Expects to receive a list array"); + }; - let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; - - let ordering_values = orderings.into_iter().map(|partition_ordering_rows| { - // Extract value from struct to ordering_rows for each group/partition - let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { - if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { - Ok(ordering_columns_per_row) - } else { - exec_err!( - "Expects to receive ScalarValue::Struct(Some(..), _) but got:{:?}", - ordering_row.data_type() - ) - } - }).collect::>>()?; - Ok(ordering_value) - }).collect::>>()?; - for ordering_values in ordering_values.iter() { - partition_ordering_values.push(ordering_values); - } + // Stores ARRAY_AGG results coming from each partition + let mut partition_values = vec![]; + // Stores ordering requirement expression results coming from each partition + let mut partition_ordering_values = vec![]; - let sort_options = self - .ordering_req - .iter() - .map(|sort_expr| sort_expr.options) - .collect::>(); - let (new_values, new_orderings) = merge_ordered_arrays( - &partition_values, - &partition_ordering_values, - &sort_options, - )?; - self.values = new_values; - self.ordering_values = new_orderings; - } else { - return exec_err!("Expects to receive a list array"); + // Existing values should be merged also. + partition_values.push(self.values.as_slice()); + partition_ordering_values.push(self.ordering_values.as_slice()); + + let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; + + for v in array_agg_res.iter() { + partition_values.push(v); } + + let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; + + let ordering_values = orderings.into_iter().map(|partition_ordering_rows| { + // Extract value from struct to ordering_rows for each group/partition + partition_ordering_rows.into_iter().map(|ordering_row| { + if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { + Ok(ordering_columns_per_row) + } else { + exec_err!( + "Expects to receive ScalarValue::Struct(Some(..), _) but got: {:?}", + ordering_row.data_type() + ) + } + }).collect::>>() + }).collect::>>()?; + for ordering_values in ordering_values.iter() { + partition_ordering_values.push(ordering_values); + } + + let sort_options = self + .ordering_req + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(); + (self.values, self.ordering_values) = merge_ordered_arrays( + &partition_values, + &partition_ordering_values, + &sort_options, + )?; Ok(()) } @@ -298,15 +288,14 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { } fn evaluate(&self) -> Result { - if self.from_start { - let arr = ScalarValue::new_list(&self.values, &self.datatypes[0]); - Ok(ScalarValue::List(arr)) - } else { + let array = if self.reverse { let mut reverse_values = self.values.clone(); reverse_values.reverse(); - let arr = ScalarValue::new_list(&reverse_values, &self.datatypes[0]); - Ok(ScalarValue::List(arr)) - } + ScalarValue::new_list(&reverse_values, &self.datatypes[0]) + } else { + ScalarValue::new_list(&self.values, &self.datatypes[0]) + }; + Ok(ScalarValue::List(array)) } fn size(&self) -> usize { @@ -337,7 +326,7 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { impl OrderSensitiveArrayAggAccumulator { fn evaluate_orderings(&self) -> Result { let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]); - let struct_field = Fields::from(fields.clone()); + let struct_field = Fields::from(fields); let orderings: Vec = self .ordering_values @@ -346,7 +335,7 @@ impl OrderSensitiveArrayAggAccumulator { ScalarValue::Struct(Some(ordering.clone()), struct_field.clone()) }) .collect(); - let struct_type = DataType::Struct(Fields::from(fields)); + let struct_type = DataType::Struct(struct_field); // Wrap in List, so we have the same data structure ListArray(StructArray..) for group by cases let arr = ScalarValue::new_list(&orderings, &struct_type); @@ -354,20 +343,19 @@ impl OrderSensitiveArrayAggAccumulator { } } -/// This is a wrapper struct to be able to correctly merge ARRAY_AGG -/// data from multiple partitions using `BinaryHeap`. -/// When used inside `BinaryHeap` this struct returns smallest `CustomElement`, -/// where smallest is determined by `ordering` values (`Vec`) -/// according to `sort_options` +/// This is a wrapper struct to be able to correctly merge `ARRAY_AGG` data from +/// multiple partitions using `BinaryHeap`. When used inside `BinaryHeap`, this +/// struct returns smallest `CustomElement`, where smallest is determined by +/// `ordering` values (`Vec`) according to `sort_options`. #[derive(Debug, PartialEq, Eq)] struct CustomElement<'a> { - // Stores from which partition entry is received + /// Stores the partition this entry came from branch_idx: usize, - // values to be merged + /// Values to merge value: ScalarValue, - // according to `ordering` values, comparisons will be done. + // Comparison "key" ordering: Vec, - // `sort_options` defines, desired ordering by the user + /// Options defining the ordering semantics sort_options: &'a [SortOptions], } @@ -453,13 +441,13 @@ pub(crate) fn merge_ordered_arrays( sort_options: &[SortOptions], ) -> Result<(Vec, Vec>)> { // Keep track the most recent data of each branch, in binary heap data structure. - let mut heap: BinaryHeap = BinaryHeap::new(); + let mut heap = BinaryHeap::::new(); - if !(values.len() == ordering_values.len() - && values + if values.len() != ordering_values.len() + || values .iter() - .zip(ordering_values.iter()) - .all(|(vals, ordering_vals)| vals.len() == ordering_vals.len())) + .zip(ordering_values) + .any(|(vals, ordering_vals)| vals.len() != ordering_vals.len()) { return exec_err!( "Expects values arguments and/or ordering_values arguments to have same size" @@ -476,8 +464,8 @@ pub(crate) fn merge_ordered_arrays( let mut merged_orderings = vec![]; // Continue iterating the loop until consuming data of all branches. loop { - let min_elem = if let Some(min_elem) = heap.pop() { - min_elem + let minimum = if let Some(minimum) = heap.pop() { + minimum } else { // Heap is empty, fill it with the next entries from each branch. for (idx, end_idx, ordering, branch_index) in izip!( @@ -486,43 +474,45 @@ pub(crate) fn merge_ordered_arrays( ordering_values.iter(), 0..n_branch ) { - // We consumed this branch, skip it - if idx == end_idx { - continue; + // If we consumed this branch, skip it. + if idx != end_idx { + // Push the next element to the heap: + heap.push(CustomElement::new( + branch_index, + values[branch_index][*idx].clone(), + ordering[*idx].to_vec(), + sort_options, + )); } - - // Push the next element to the heap. - let elem = CustomElement::new( - branch_index, - values[branch_index][*idx].clone(), - ordering[*idx].to_vec(), - sort_options, - ); - heap.push(elem); } - // Now we have filled the heap, get the largest entry (this will be the next element in merge) - if let Some(min_elem) = heap.pop() { - min_elem + // Now we have filled the heap, get the largest entry (this will be + // the next element in merge). + if let Some(minimum) = heap.pop() { + minimum } else { - // Heap is empty, this means that all indices are same with end_indices. e.g - // We have consumed all of the branches. Merging is completed - // Exit from the loop + // Heap is empty, this means that all indices are same with + // `end_indices`. We have consumed all of the branches, merge + // is completed, exit from the loop: break; } }; - let branch_idx = min_elem.branch_idx; + let branch_idx = minimum.branch_idx; // Increment the index of merged branch, indices[branch_idx] += 1; let row_idx = indices[branch_idx]; - merged_values.push(min_elem.value.clone()); - merged_orderings.push(min_elem.ordering.clone()); + merged_values.push(minimum.value); + merged_orderings.push(minimum.ordering); if row_idx < end_indices[branch_idx] { - // Push next entry in the most recently consumed branch to the heap - // If there is an available entry + // If there is an available entry, push next entry in the most + // recently consumed branch to the heap. let value = values[branch_idx][row_idx].clone(); let ordering_row = ordering_values[branch_idx][row_idx].to_vec(); - let elem = CustomElement::new(branch_idx, value, ordering_row, sort_options); - heap.push(elem); + heap.push(CustomElement::new( + branch_idx, + value, + ordering_row, + sort_options, + )); } } diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index 2b8ac2d516c55..d81e12e00cba9 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -160,18 +160,17 @@ impl PartialEq for NthValueAgg { #[derive(Debug)] pub(crate) struct NthValueAccumulator { n: i64, - // `values` stores entries in the NTH_VALUE result. + /// Stores entries in the `NTH_VALUE` result. values: VecDeque, - // `ordering_values` stores values of ordering requirement expression - // corresponding to each value in the NTH_VALUE. - // For each `ScalarValue` inside `values`, there will be a corresponding - // `Vec` inside `ordering_values` which stores its ordering. - // This information is used during merging results of the different partitions. - // For detailed information how merging is done see [`merge_ordered_arrays`] + /// Stores values of ordering requirement expressions corresponding to each + /// entry in `values`. This information is used when merging results from + /// different partitions. For detailed information how merging is done, see + /// [`merge_ordered_arrays`]. ordering_values: VecDeque>, - // Stores datatype of the expression inside NTH_VALUE and ordering requirement expressions. + /// Stores datatypes of expressions inside values and ordering requirement + /// expressions. datatypes: Vec, - // Stores ordering requirement of the Accumulator + /// Stores the ordering requirement of the `Accumulator`. ordering_req: LexOrdering, } From 3ce91e22195939d0ec5c0982b43d13a4a12e19eb Mon Sep 17 00:00:00 2001 From: Mehmet Ozan Kabak Date: Tue, 16 Jan 2024 00:20:53 +0300 Subject: [PATCH 26/31] Review Part 3 --- .../physical-expr/src/aggregate/nth_value.rs | 24 +++++++------------ .../physical-expr/src/expressions/mod.rs | 10 ++++---- 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index d81e12e00cba9..6066fb6cfa4bd 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -236,18 +236,13 @@ impl Accumulator for NthValueAccumulator { if self.ordering_req.is_empty() { let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - // Stores NTH_VALUE results coming from each partition - let mut partition_values = VecDeque::new(); - // Existing values should be merged also. - partition_values.extend(self.values.clone()); for v in array_agg_res.into_iter() { - partition_values.extend(v); - if partition_values.len() > n_required { + self.values.extend(v); + if self.values.len() > n_required { // There is enough data collected can stop merging break; } } - self.values = partition_values; } else if let Some(agg_orderings) = states[1].as_list_opt::() { // 2nd entry stores values received for ordering requirement columns, for each aggregation value inside NTH_VALUE list. // For each `StructArray` inside NTH_VALUE list, we will receive an `Array` that stores @@ -276,17 +271,16 @@ impl Accumulator for NthValueAccumulator { let ordering_values = orderings.into_iter().map(|partition_ordering_rows| { // Extract value from struct to ordering_rows for each group/partition - let ordering_value = partition_ordering_rows.into_iter().map(|ordering_row| { + partition_ordering_rows.into_iter().map(|ordering_row| { if let ScalarValue::Struct(Some(ordering_columns_per_row), _) = ordering_row { Ok(ordering_columns_per_row) } else { exec_err!( - "Expects to receive ScalarValue::Struct(Some(..), _) but got:{:?}", - ordering_row.data_type() - ) + "Expects to receive ScalarValue::Struct(Some(..), _) but got: {:?}", + ordering_row.data_type() + ) } - }).collect::>>()?; - Ok(ordering_value) + }).collect::>>() }).collect::>>()?; for ordering_values in ordering_values.iter() { partition_ordering_values.push(ordering_values); @@ -364,7 +358,7 @@ impl Accumulator for NthValueAccumulator { impl NthValueAccumulator { fn evaluate_orderings(&self) -> ScalarValue { let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]); - let struct_field = Fields::from(fields.clone()); + let struct_field = Fields::from(fields); let orderings = self .ordering_values @@ -373,7 +367,7 @@ impl NthValueAccumulator { ScalarValue::Struct(Some(ordering.clone()), struct_field.clone()) }) .collect::>(); - let struct_type = DataType::Struct(Fields::from(fields)); + let struct_type = DataType::Struct(struct_field); // Wrap in List, so we have the same data structure ListArray(StructArray..) for group by cases ScalarValue::List(ScalarValue::new_list(&orderings, &struct_type)) diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 586b6ecd4e54f..bbfba4ad83107 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -68,7 +68,6 @@ pub use crate::aggregate::string_agg::StringAgg; pub use crate::aggregate::sum::Sum; pub use crate::aggregate::sum_distinct::DistinctSum; pub use crate::aggregate::variance::{Variance, VariancePop}; - pub use crate::window::cume_dist::cume_dist; pub use crate::window::cume_dist::CumeDist; pub use crate::window::lead_lag::WindowShift; @@ -78,6 +77,7 @@ pub use crate::window::ntile::Ntile; pub use crate::window::rank::{dense_rank, percent_rank, rank}; pub use crate::window::rank::{Rank, RankType}; pub use crate::window::row_number::RowNumber; +pub use crate::PhysicalSortExpr; pub use binary::{binary, BinaryExpr}; pub use case::{case, CaseExpr}; @@ -99,20 +99,20 @@ pub use try_cast::{try_cast, TryCastExpr}; pub fn format_state_name(name: &str, state_name: &str) -> String { format!("{name}[{state_name}]") } -pub use crate::PhysicalSortExpr; #[cfg(test)] pub(crate) mod tests { + use std::sync::Arc; + use crate::expressions::{col, create_aggregate_expr, try_cast}; use crate::{AggregateExpr, EmitTo}; + use arrow::record_batch::RecordBatch; use arrow_array::ArrayRef; use arrow_schema::{Field, Schema}; - use datafusion_common::Result; - use datafusion_common::ScalarValue; + use datafusion_common::{Result, ScalarValue}; use datafusion_expr::type_coercion::aggregates::coerce_types; use datafusion_expr::AggregateFunction; - use std::sync::Arc; /// macro to perform an aggregation using [`datafusion_expr::Accumulator`] and verify the /// result. From d1250bf9d6a31d5933c9564859881e20a29dee80 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 16 Jan 2024 11:41:24 +0300 Subject: [PATCH 27/31] Add new_list from iter --- datafusion/common/src/scalar.rs | 34 +++++++++++++++++++ .../src/aggregate/array_agg_ordered.rs | 7 ++-- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs index 2e067899e1830..a07fc9abec793 100644 --- a/datafusion/common/src/scalar.rs +++ b/datafusion/common/src/scalar.rs @@ -1730,6 +1730,40 @@ impl ScalarValue { Arc::new(array_into_list_array(values)) } + /// Converts `IntoIterator` where each element has type corresponding to + /// `data_type`, to a [`ListArray`]. + /// + /// Example + /// ``` + /// use datafusion_common::ScalarValue; + /// use arrow::array::{ListArray, Int32Array}; + /// use arrow::datatypes::{DataType, Int32Type}; + /// use datafusion_common::cast::as_list_array; + /// + /// let scalars = vec![ + /// ScalarValue::Int32(Some(1)), + /// ScalarValue::Int32(None), + /// ScalarValue::Int32(Some(2)) + /// ]; + /// + /// let result = ScalarValue::new_list_from_iter(scalars.iter(), &DataType::Int32); + /// + /// let expected = ListArray::from_iter_primitive::( + /// vec![ + /// Some(vec![Some(1), None, Some(2)]) + /// ]); + /// + /// assert_eq!(*result, expected); + /// ``` + pub fn new_list_from_iter(values: impl IntoIterator + ExactSizeIterator, data_type: &DataType) -> Arc { + let values = if values.len() == 0 { + new_empty_array(data_type) + } else { + Self::iter_to_array(values).unwrap() + }; + Arc::new(array_into_list_array(values)) + } + /// Converts `Vec` where each element has type corresponding to /// `data_type`, to a [`LargeListArray`]. /// diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 088c83bfc3b2b..77d19259b5a0b 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -288,12 +288,11 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { } fn evaluate(&self) -> Result { + let values = self.values.clone(); let array = if self.reverse { - let mut reverse_values = self.values.clone(); - reverse_values.reverse(); - ScalarValue::new_list(&reverse_values, &self.datatypes[0]) + ScalarValue::new_list_from_iter(values.into_iter().rev(), &self.datatypes[0]) } else { - ScalarValue::new_list(&self.values, &self.datatypes[0]) + ScalarValue::new_list_from_iter(values.into_iter(), &self.datatypes[0]) }; Ok(ScalarValue::List(array)) } From 76841077794c93c7b18d0ca071df0f1ab6478c81 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 16 Jan 2024 12:34:36 +0300 Subject: [PATCH 28/31] Convert API to receive vecdeque --- datafusion/common/src/scalar.rs | 5 ++- .../src/aggregate/array_agg_ordered.rs | 43 ++++++++++--------- .../physical-expr/src/aggregate/nth_value.rs | 18 ++++---- 3 files changed, 34 insertions(+), 32 deletions(-) diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs index a07fc9abec793..89d3d8ded9c92 100644 --- a/datafusion/common/src/scalar.rs +++ b/datafusion/common/src/scalar.rs @@ -1755,7 +1755,10 @@ impl ScalarValue { /// /// assert_eq!(*result, expected); /// ``` - pub fn new_list_from_iter(values: impl IntoIterator + ExactSizeIterator, data_type: &DataType) -> Arc { + pub fn new_list_from_iter( + values: impl IntoIterator + ExactSizeIterator, + data_type: &DataType, + ) -> Arc { let values = if values.len() == 0 { new_empty_array(data_type) } else { diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 77d19259b5a0b..0e27f992c113f 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -20,7 +20,7 @@ use std::any::Any; use std::cmp::Ordering; -use std::collections::BinaryHeap; +use std::collections::{BinaryHeap, VecDeque}; use std::fmt::Debug; use std::sync::Arc; @@ -240,13 +240,13 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { let mut partition_ordering_values = vec![]; // Existing values should be merged also. - partition_values.push(self.values.as_slice()); - partition_ordering_values.push(self.ordering_values.as_slice()); + partition_values.push(self.values.clone().into()); + partition_ordering_values.push(self.ordering_values.clone().into()); let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - for v in array_agg_res.iter() { - partition_values.push(v); + for v in array_agg_res.into_iter() { + partition_values.push(v.into()); } let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; @@ -262,9 +262,9 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { ordering_row.data_type() ) } - }).collect::>>() + }).collect::>>() }).collect::>>()?; - for ordering_values in ordering_values.iter() { + for ordering_values in ordering_values.into_iter() { partition_ordering_values.push(ordering_values); } @@ -431,11 +431,11 @@ impl<'a> PartialOrd for CustomElement<'a> { /// Inner `Vec`s of the `ordering_values` will be compared according `sort_options` (Their sizes should match) pub(crate) fn merge_ordered_arrays( // We will merge values into single `Vec`. - values: &[&[ScalarValue]], + values: &[VecDeque], // `values` will be merged according to `ordering_values`. // Inner `Vec` can be thought as ordering information for the // each `ScalarValue` in the values`. - ordering_values: &[&[Vec]], + ordering_values: &[VecDeque>], // Defines according to which ordering comparisons should be done. sort_options: &[SortOptions], ) -> Result<(Vec, Vec>)> { @@ -520,6 +520,7 @@ pub(crate) fn merge_ordered_arrays( #[cfg(test)] mod tests { + use std::collections::VecDeque; use std::sync::Arc; use crate::aggregate::array_agg_ordered::merge_ordered_arrays; @@ -538,7 +539,7 @@ mod tests { let n_row = lhs_arrays[0].len(); let lhs_orderings = (0..n_row) .map(|idx| get_row_at_idx(&lhs_arrays, idx)) - .collect::>>()?; + .collect::>>()?; let rhs_arrays: Vec = vec![ Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2])), @@ -547,7 +548,7 @@ mod tests { let n_row = rhs_arrays[0].len(); let rhs_orderings = (0..n_row) .map(|idx| get_row_at_idx(&rhs_arrays, idx)) - .collect::>>()?; + .collect::>>()?; let sort_options = vec![ SortOptions { descending: false, @@ -562,12 +563,12 @@ mod tests { let lhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4])) as ArrayRef; let lhs_vals = (0..lhs_vals_arr.len()) .map(|idx| ScalarValue::try_from_array(&lhs_vals_arr, idx)) - .collect::>>()?; + .collect::>>()?; let rhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4])) as ArrayRef; let rhs_vals = (0..rhs_vals_arr.len()) .map(|idx| ScalarValue::try_from_array(&rhs_vals_arr, idx)) - .collect::>>()?; + .collect::>>()?; let expected = Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2, 2, 3, 3, 4, 4])) as ArrayRef; let expected_ts = vec![ @@ -576,8 +577,8 @@ mod tests { ]; let (merged_vals, merged_ts) = merge_ordered_arrays( - &[&lhs_vals, &rhs_vals], - &[&lhs_orderings, &rhs_orderings], + &[lhs_vals, rhs_vals], + &[lhs_orderings, rhs_orderings], &sort_options, )?; let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?; @@ -605,7 +606,7 @@ mod tests { let n_row = lhs_arrays[0].len(); let lhs_orderings = (0..n_row) .map(|idx| get_row_at_idx(&lhs_arrays, idx)) - .collect::>>()?; + .collect::>>()?; let rhs_arrays: Vec = vec![ Arc::new(Int64Array::from(vec![2, 1, 1, 0, 0])), @@ -614,7 +615,7 @@ mod tests { let n_row = rhs_arrays[0].len(); let rhs_orderings = (0..n_row) .map(|idx| get_row_at_idx(&rhs_arrays, idx)) - .collect::>>()?; + .collect::>>()?; let sort_options = vec![ SortOptions { descending: true, @@ -630,12 +631,12 @@ mod tests { let lhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 1, 2])) as ArrayRef; let lhs_vals = (0..lhs_vals_arr.len()) .map(|idx| ScalarValue::try_from_array(&lhs_vals_arr, idx)) - .collect::>>()?; + .collect::>>()?; let rhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 1, 2])) as ArrayRef; let rhs_vals = (0..rhs_vals_arr.len()) .map(|idx| ScalarValue::try_from_array(&rhs_vals_arr, idx)) - .collect::>>()?; + .collect::>>()?; let expected = Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2, 2, 1, 1, 2, 2])) as ArrayRef; let expected_ts = vec![ @@ -643,8 +644,8 @@ mod tests { Arc::new(Int64Array::from(vec![4, 4, 3, 3, 2, 2, 1, 1, 0, 0])) as ArrayRef, ]; let (merged_vals, merged_ts) = merge_ordered_arrays( - &[&lhs_vals, &rhs_vals], - &[&lhs_orderings, &rhs_orderings], + &[lhs_vals, rhs_vals], + &[lhs_orderings, rhs_orderings], &sort_options, )?; let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?; diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index 6066fb6cfa4bd..bf5113cb3f8c7 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -249,22 +249,20 @@ impl Accumulator for NthValueAccumulator { // values received from its ordering requirement expression. (This information is necessary for during merging). // Stores NTH_VALUE results coming from each partition - let mut partition_values: Vec<&[ScalarValue]> = vec![]; + let mut partition_values: Vec> = vec![]; // Stores ordering requirement expression results coming from each partition - let mut partition_ordering_values: Vec<&[Vec]> = vec![]; + let mut partition_ordering_values: Vec>> = vec![]; // Existing values should be merged also. - let partition_value = self.values.make_contiguous(); - partition_values.push(partition_value); + partition_values.push(self.values.clone()); - let partition_ordering_value = self.ordering_values.make_contiguous(); - partition_ordering_values.push(partition_ordering_value); + partition_ordering_values.push(self.ordering_values.clone()); let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?; - for v in array_agg_res.iter() { - partition_values.push(v); + for v in array_agg_res.into_iter() { + partition_values.push(v.into()); } let orderings = ScalarValue::convert_array_to_scalar_vec(agg_orderings)?; @@ -282,8 +280,8 @@ impl Accumulator for NthValueAccumulator { } }).collect::>>() }).collect::>>()?; - for ordering_values in ordering_values.iter() { - partition_ordering_values.push(ordering_values); + for ordering_values in ordering_values.into_iter() { + partition_ordering_values.push(ordering_values.into()); } let sort_options = self From 137c21e7b9268d670dea45c3e089fbd491d336cd Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 16 Jan 2024 12:38:27 +0300 Subject: [PATCH 29/31] Receive mutable argument --- .../src/aggregate/array_agg_ordered.rs | 18 +++++++++--------- .../physical-expr/src/aggregate/nth_value.rs | 4 ++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index 0e27f992c113f..bd87465983991 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -274,8 +274,8 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { .map(|sort_expr| sort_expr.options) .collect::>(); (self.values, self.ordering_values) = merge_ordered_arrays( - &partition_values, - &partition_ordering_values, + &mut partition_values, + &mut partition_ordering_values, &sort_options, )?; Ok(()) @@ -431,11 +431,11 @@ impl<'a> PartialOrd for CustomElement<'a> { /// Inner `Vec`s of the `ordering_values` will be compared according `sort_options` (Their sizes should match) pub(crate) fn merge_ordered_arrays( // We will merge values into single `Vec`. - values: &[VecDeque], + values: &mut [VecDeque], // `values` will be merged according to `ordering_values`. // Inner `Vec` can be thought as ordering information for the // each `ScalarValue` in the values`. - ordering_values: &[VecDeque>], + ordering_values: &mut [VecDeque>], // Defines according to which ordering comparisons should be done. sort_options: &[SortOptions], ) -> Result<(Vec, Vec>)> { @@ -445,7 +445,7 @@ pub(crate) fn merge_ordered_arrays( if values.len() != ordering_values.len() || values .iter() - .zip(ordering_values) + .zip(ordering_values.iter()) .any(|(vals, ordering_vals)| vals.len() != ordering_vals.len()) { return exec_err!( @@ -577,8 +577,8 @@ mod tests { ]; let (merged_vals, merged_ts) = merge_ordered_arrays( - &[lhs_vals, rhs_vals], - &[lhs_orderings, rhs_orderings], + &mut [lhs_vals, rhs_vals], + &mut [lhs_orderings, rhs_orderings], &sort_options, )?; let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?; @@ -644,8 +644,8 @@ mod tests { Arc::new(Int64Array::from(vec![4, 4, 3, 3, 2, 2, 1, 1, 0, 0])) as ArrayRef, ]; let (merged_vals, merged_ts) = merge_ordered_arrays( - &[lhs_vals, rhs_vals], - &[lhs_orderings, rhs_orderings], + &mut [lhs_vals, rhs_vals], + &mut [lhs_orderings, rhs_orderings], &sort_options, )?; let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?; diff --git a/datafusion/physical-expr/src/aggregate/nth_value.rs b/datafusion/physical-expr/src/aggregate/nth_value.rs index bf5113cb3f8c7..5a1ca90b7f5e6 100644 --- a/datafusion/physical-expr/src/aggregate/nth_value.rs +++ b/datafusion/physical-expr/src/aggregate/nth_value.rs @@ -290,8 +290,8 @@ impl Accumulator for NthValueAccumulator { .map(|sort_expr| sort_expr.options) .collect::>(); let (new_values, new_orderings) = merge_ordered_arrays( - &partition_values, - &partition_ordering_values, + &mut partition_values, + &mut partition_ordering_values, &sort_options, )?; self.values = new_values.into(); From 864ee9fb6bab3055f535285e8469fd8c34798f39 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 16 Jan 2024 12:56:12 +0300 Subject: [PATCH 30/31] Refactor merge implementation --- .../src/aggregate/array_agg_ordered.rs | 57 +++++++++---------- 1 file changed, 26 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs index bd87465983991..34f8d20628dc1 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_ordered.rs @@ -38,8 +38,6 @@ use datafusion_common::utils::{compare_rows, get_row_at_idx}; use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue}; use datafusion_expr::Accumulator; -use itertools::izip; - /// Expression for a `ARRAY_AGG(... ORDER BY ..., ...)` aggregation. In a multi /// partition setting, partial aggregations are computed for every partition, /// and then their results are merged. @@ -453,12 +451,6 @@ pub(crate) fn merge_ordered_arrays( ); } let n_branch = values.len(); - // For each branch we keep track of indices of next will be merged entry - let mut indices = vec![0_usize; n_branch]; - // Keep track of sizes of each branch. - let end_indices = (0..n_branch) - .map(|idx| values[idx].len()) - .collect::>(); let mut merged_values = vec![]; let mut merged_orderings = vec![]; // Continue iterating the loop until consuming data of all branches. @@ -467,23 +459,21 @@ pub(crate) fn merge_ordered_arrays( minimum } else { // Heap is empty, fill it with the next entries from each branch. - for (idx, end_idx, ordering, branch_index) in izip!( - indices.iter(), - end_indices.iter(), - ordering_values.iter(), - 0..n_branch - ) { - // If we consumed this branch, skip it. - if idx != end_idx { + for branch_idx in 0..n_branch { + if let Some(orderings) = ordering_values[branch_idx].pop_front() { + // Their size should be same, we can safely .unwrap here. + let value = values[branch_idx].pop_front().unwrap(); // Push the next element to the heap: heap.push(CustomElement::new( - branch_index, - values[branch_index][*idx].clone(), - ordering[*idx].to_vec(), + branch_idx, + value, + orderings, sort_options, )); } + // If None, we consumed this branch, skip it. } + // Now we have filled the heap, get the largest entry (this will be // the next element in merge). if let Some(minimum) = heap.pop() { @@ -495,21 +485,26 @@ pub(crate) fn merge_ordered_arrays( break; } }; - let branch_idx = minimum.branch_idx; - // Increment the index of merged branch, - indices[branch_idx] += 1; - let row_idx = indices[branch_idx]; - merged_values.push(minimum.value); - merged_orderings.push(minimum.ordering); - if row_idx < end_indices[branch_idx] { - // If there is an available entry, push next entry in the most - // recently consumed branch to the heap. - let value = values[branch_idx][row_idx].clone(); - let ordering_row = ordering_values[branch_idx][row_idx].to_vec(); + let CustomElement { + branch_idx, + value, + ordering, + .. + } = minimum; + // Add minimum value in the heap to the result + merged_values.push(value); + merged_orderings.push(ordering); + + // If there is an available entry, push next entry in the most + // recently consumed branch to the heap. + if let Some(orderings) = ordering_values[branch_idx].pop_front() { + // Their size should be same, we can safely .unwrap here. + let value = values[branch_idx].pop_front().unwrap(); + // Push the next element to the heap: heap.push(CustomElement::new( branch_idx, value, - ordering_row, + orderings, sort_options, )); } From 1fad8f9ef3bfd54980dec9b83ed73099101e1170 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Tue, 16 Jan 2024 13:09:34 +0300 Subject: [PATCH 31/31] Fix doctest --- datafusion/common/src/scalar.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs index 89d3d8ded9c92..78d51b7ed6526 100644 --- a/datafusion/common/src/scalar.rs +++ b/datafusion/common/src/scalar.rs @@ -1746,7 +1746,7 @@ impl ScalarValue { /// ScalarValue::Int32(Some(2)) /// ]; /// - /// let result = ScalarValue::new_list_from_iter(scalars.iter(), &DataType::Int32); + /// let result = ScalarValue::new_list_from_iter(scalars.into_iter(), &DataType::Int32); /// /// let expected = ListArray::from_iter_primitive::( /// vec![