diff --git a/datafusion/physical-plan/benches/aggregate_filter.rs b/datafusion/physical-plan/benches/aggregate_filter.rs index ae97cb7388609..417d153abce5c 100644 --- a/datafusion/physical-plan/benches/aggregate_filter.rs +++ b/datafusion/physical-plan/benches/aggregate_filter.rs @@ -29,11 +29,12 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_execution::TaskContext; -use datafusion_expr::Operator; +use datafusion_expr::{AggregateUDF, Operator}; +use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; -use datafusion_physical_expr::expressions::{BinaryExpr, col}; +use datafusion_physical_expr::expressions::{BinaryExpr, col, lit}; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; @@ -61,6 +62,7 @@ const FILTER_CASES: &[Option] = &[ enum ArgumentKind { Column, Multiply, + Literal, } impl ArgumentKind { @@ -68,6 +70,7 @@ impl ArgumentKind { match self { Self::Column => "column", Self::Multiply => "multiply", + Self::Literal => "literal", } } } @@ -151,19 +154,30 @@ fn aggregate_expr( argument_kind: ArgumentKind, ) -> Arc { let value = col("value", schema).unwrap(); - let argument: Arc = match argument_kind { - ArgumentKind::Column => value, - ArgumentKind::Multiply => Arc::new(BinaryExpr::new( - Arc::clone(&value), - Operator::Multiply, - value, - )), - }; + let (function, argument): (Arc, Arc) = + match argument_kind { + ArgumentKind::Column => (sum_udaf(), value), + ArgumentKind::Multiply => ( + sum_udaf(), + Arc::new(BinaryExpr::new( + Arc::clone(&value), + Operator::Multiply, + value, + )), + ), + // `COUNT(*)` is planned as `COUNT(1)` + ArgumentKind::Literal => (count_udaf(), lit(1i64)), + }; + let alias = format!( + "{}_{}_{aggregate_index}", + function.name(), + argument_kind.name() + ); Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![argument]) + AggregateExprBuilder::new(function, vec![argument]) .schema(Arc::clone(schema)) - .alias(format!("sum_{}_{aggregate_index}", argument_kind.name())) + .alias(alias) .build() .unwrap(), ) @@ -266,7 +280,11 @@ fn aggregate_filter_benchmark(c: &mut Criterion) { c, &runtime, AggregateLayout::OneAggregate, - &[ArgumentKind::Column, ArgumentKind::Multiply], + &[ + ArgumentKind::Column, + ArgumentKind::Multiply, + ArgumentKind::Literal, + ], ); benchmark_case( c, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_argument.rs b/datafusion/physical-plan/src/aggregates/aggregate_argument.rs new file mode 100644 index 0000000000000..56ca9e283754b --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_argument.rs @@ -0,0 +1,234 @@ +// 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. + +//! Evaluation of aggregate function arguments against input batches. + +use std::sync::Arc; + +use crate::PhysicalExpr; +use arrow::array::{ArrayRef, BooleanArray}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, ScalarValue}; +use datafusion_physical_expr::expressions::Literal; + +/// Largest literal array, in bytes, that an [`AggregateArgument`] keeps for +/// later batches. +/// +/// The kept array lives as long as the aggregate stream and is not tracked by +/// a memory pool, so larger arrays are rebuilt for every batch instead. +const MAX_KEPT_LITERAL_ARRAY_BYTES: usize = 1024 * 1024; + +/// One argument of an aggregate function, such as `x` in `SUM(x)`. +/// +/// Accumulators take one value per input row, so an argument that evaluates to +/// a single value must be expanded to an array for every batch. A literal +/// argument, such as the `1` in `COUNT(1)` (which is how `COUNT(*)` is planned) +/// or the `','` in `STRING_AGG(x, ',')`, has the same value for every batch of +/// the query. Its array is built on demand and reused for later batches, which +/// receive a zero-copy slice of it. +/// +/// Only literals are treated this way. Other expressions are evaluated against +/// every batch: a [`ColumnarValue::Scalar`] result holds one value for the rows +/// of a single batch, and the next batch may produce a different value. +/// +/// [`ColumnarValue::Scalar`]: datafusion_expr::ColumnarValue::Scalar +pub(in crate::aggregates) struct AggregateArgument { + expr: Arc, + /// Set when `expr` is a [`Literal`]. + literal: Option, +} + +impl AggregateArgument { + pub(in crate::aggregates) fn new(expr: Arc) -> Self { + let literal = expr + .downcast_ref::() + .map(|literal| LiteralArray::new(literal.value().clone())); + Self { expr, literal } + } + + pub(in crate::aggregates) fn expr(&self) -> &Arc { + &self.expr + } + + /// Evaluates the argument against `batch`, returning one value per row. + pub(in crate::aggregates) fn evaluate( + &mut self, + batch: &RecordBatch, + ) -> Result { + match &mut self.literal { + Some(literal) => literal.array_of_size(batch.num_rows()), + None => self + .expr + .evaluate(batch)? + .into_array_of_size(batch.num_rows()), + } + } + + /// Evaluates the argument for the rows of `batch` where `selection` is + /// true, returning one value per row of `batch`. + /// + /// Rows where `selection` is not true may hold any value, because the + /// selection is also passed to [`GroupsAccumulator::convert_to_state`], + /// which ignores them. + /// + /// [`GroupsAccumulator::convert_to_state`]: datafusion_expr::GroupsAccumulator::convert_to_state + pub(in crate::aggregates) fn evaluate_selection( + &mut self, + batch: &RecordBatch, + selection: &BooleanArray, + ) -> Result { + match &mut self.literal { + Some(literal) if selection.has_true() => { + literal.array_of_size(batch.num_rows()) + } + _ => self + .expr + .evaluate_selection(batch, selection)? + .into_array_of_size(batch.num_rows()), + } + } +} + +/// A literal value and the largest reusable array retained so far. +struct LiteralArray { + value: ScalarValue, + /// `None` until the first array is built, or while every array built has + /// been larger than [`MAX_KEPT_LITERAL_ARRAY_BYTES`]. + array: Option, +} + +impl LiteralArray { + fn new(value: ScalarValue) -> Self { + Self { value, array: None } + } + + fn array_of_size(&mut self, num_rows: usize) -> Result { + if let Some(array) = &self.array + && array.len() >= num_rows + { + return Ok(array.slice(0, num_rows)); + } + + let array = self.value.to_array_of_size(num_rows)?; + if array.get_array_memory_size() <= MAX_KEPT_LITERAL_ARRAY_BYTES { + self.array = Some(Arc::clone(&array)); + } + Ok(array) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{AsArray, Int32Array}; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use datafusion_physical_expr::expressions::{col, lit}; + + fn batch(num_rows: usize) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let values = Int32Array::from_iter_values(0..num_rows as i32); + RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap() + } + + fn values_ptr(array: &ArrayRef) -> *const i32 { + array.as_primitive::().values().as_ptr() + } + + #[test] + fn literal_array_is_reused_across_batches() -> Result<()> { + let mut argument = AggregateArgument::new(lit(1i32)); + + let first = argument.evaluate(&batch(4))?; + let smaller = argument.evaluate(&batch(3))?; + assert_eq!( + smaller.as_primitive::(), + &Int32Array::from(vec![1; 3]) + ); + assert_eq!(values_ptr(&smaller), values_ptr(&first)); + + // A larger batch builds a longer array, which serves later batches + let larger = argument.evaluate(&batch(6))?; + assert_eq!( + larger.as_primitive::(), + &Int32Array::from(vec![1; 6]) + ); + let after_larger = argument.evaluate(&batch(5))?; + assert_eq!(values_ptr(&after_larger), values_ptr(&larger)); + + assert!(argument.evaluate(&batch(0))?.is_empty()); + Ok(()) + } + + #[test] + fn large_literal_array_is_not_kept() -> Result<()> { + let mut argument = AggregateArgument::new(lit("x".repeat(1024))); + + let first = argument.evaluate(&batch(2048))?; + let second = argument.evaluate(&batch(2048))?; + assert_eq!(first.as_ref(), second.as_ref()); + assert_ne!( + first.as_string::().values().as_ptr(), + second.as_string::().values().as_ptr() + ); + Ok(()) + } + + #[test] + fn non_literal_is_evaluated_for_every_batch() -> Result<()> { + let schema = batch(0).schema(); + let mut argument = AggregateArgument::new(col("a", &schema)?); + + let first = argument.evaluate(&batch(3))?; + assert_eq!( + first.as_primitive::(), + &Int32Array::from(vec![0, 1, 2]) + ); + let second = argument.evaluate(&batch(2))?; + assert_eq!( + second.as_primitive::(), + &Int32Array::from(vec![0, 1]) + ); + Ok(()) + } + + #[test] + fn literal_selection_matches_evaluate_selection() -> Result<()> { + let batch = batch(3); + let expr = lit(7i32); + let mut argument = AggregateArgument::new(Arc::clone(&expr)); + + for selection in [ + BooleanArray::from(vec![true, true, true]), + BooleanArray::from(vec![false, true, false]), + BooleanArray::from(vec![Some(false), None, Some(true)]), + BooleanArray::from(vec![false, false, false]), + BooleanArray::from(vec![None, Some(false), None]), + ] { + let expected = expr + .evaluate_selection(&batch, &selection)? + .into_array(batch.num_rows())?; + let actual = argument.evaluate_selection(&batch, &selection)?; + assert_eq!( + actual.as_ref(), + expected.as_ref(), + "selection {selection:?}" + ); + } + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index b25caf815eed1..8ad11ddd3f4e3 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -34,6 +34,7 @@ use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use log::debug; use crate::PhysicalExpr; +use crate::aggregates::aggregate_argument::AggregateArgument; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, @@ -203,10 +204,10 @@ impl AggregateHashTable { /// See comments in [`EvaluatedAggregateBatch`] pub(super) fn evaluate_batch( - &self, + &mut self, batch: &RecordBatch, ) -> Result { - let state = self.state.building(); + let state = self.state.building_mut(); // Outer vec: one per grouping set; inner vec: group-by expressions. let grouping_set_args = self .group_by_metrics @@ -216,7 +217,7 @@ impl AggregateHashTable { let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| { state .accumulators - .iter() + .iter_mut() .enumerate() .map(|(idx, acc)| { self.aggregate_argument_metrics @@ -526,8 +527,8 @@ pub(super) struct HashAggregateAccumulator { /// Arguments to pass to this accumulator. /// - /// Example: `CORR(x, y)` stores two expressions here, while `SUM(x)` stores one. - arguments: Vec>, + /// Example: `CORR(x, y)` stores two arguments here, while `SUM(x)` stores one. + arguments: Vec, /// Optional `FILTER` expression for this accumulator. /// @@ -720,7 +721,7 @@ impl HashAggregateAccumulator { ) -> Self { Self { aggregate_expr, - arguments, + arguments: arguments.into_iter().map(AggregateArgument::new).collect(), filter, accumulator, submetrics, @@ -734,7 +735,10 @@ impl HashAggregateAccumulator { create_group_accumulator(&self.aggregate_expr, Arc::clone(&self.submetrics))?; Ok(Self::new( Arc::clone(&self.aggregate_expr), - self.arguments.clone(), + self.arguments + .iter() + .map(|argument| Arc::clone(argument.expr())) + .collect(), self.filter.clone(), accumulator, Arc::clone(&self.submetrics), @@ -750,7 +754,7 @@ impl HashAggregateAccumulator { /// Before updating [`GroupsAccumulator`], the retained selection is used to /// compact the matching group IDs and is not passed through. pub(super) fn evaluate_compacted_args( - &self, + &mut self, batch: &RecordBatch, ) -> Result { let selection = self.evaluate_filter(batch)?; @@ -771,13 +775,13 @@ impl HashAggregateAccumulator { }; let arguments = self .arguments - .iter() - .map(|expr| { + .iter_mut() + .map(|argument| { if let Some(argument_batch) = argument_batch { - expr.evaluate(argument_batch) - .and_then(|value| value.into_array(argument_batch.num_rows())) + argument.evaluate(argument_batch) } else { - let data_type = expr.data_type(batch.schema_ref().as_ref())?; + let data_type = + argument.expr().data_type(batch.schema_ref().as_ref())?; Ok(new_empty_array(&data_type)) } }) @@ -795,21 +799,17 @@ impl HashAggregateAccumulator { /// rows remain as null argument values and the filter is passed to /// [`GroupsAccumulator::convert_to_state`]. pub(super) fn evaluate_row_aligned_args( - &self, + &mut self, batch: &RecordBatch, ) -> Result { let filter = self.evaluate_filter(batch)?; let selection = filter.as_ref(); let arguments = self .arguments - .iter() - .map(|expr| { - selection - .map_or_else( - || expr.evaluate(batch), - |selection| expr.evaluate_selection(batch, selection), - ) - .and_then(|value| value.into_array(batch.num_rows())) + .iter_mut() + .map(|argument| match selection { + Some(selection) => argument.evaluate_selection(batch, selection), + None => argument.evaluate(batch), }) .collect::>()?; @@ -899,8 +899,8 @@ impl HashAggregateAccumulator { ) -> Result> { self.arguments .iter() - .map(|expr| { - let data_type = expr.data_type(input_schema)?; + .map(|argument| { + let data_type = argument.expr().data_type(input_schema)?; Ok(new_null_array(&data_type, num_rows)) }) .collect() @@ -997,7 +997,7 @@ mod tests { let submetrics = aggregate_sub_metrics(&metrics, 0, ["SUM(value)"]) .pop() .expect("one aggregate submetric factory"); - let accumulator = sum_accumulator(&schema, "include", 1, submetrics)?; + let mut accumulator = sum_accumulator(&schema, "include", 1, submetrics)?; let group_by_metrics = GroupByMetrics::new(&metrics, 0); let argument_metrics = AggregateArgumentMetrics::new(&metrics, 0, ["SUM(value)"]); let accumulator_metrics = AggregateAccumulatorMetrics::new( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 86c052009faea..c25b7485b88b1 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -239,7 +239,7 @@ impl OrderedAggregateTable { /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function /// evaluates `k+1`, `v*v`. pub(super) fn evaluate_batch( - &self, + &mut self, batch: &RecordBatch, ) -> Result { let grouping_set_args = @@ -250,7 +250,7 @@ impl OrderedAggregateTable { let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| { self.buffer .accumulators - .iter() + .iter_mut() .enumerate() .map(|(idx, acc)| { self.aggregate_argument_metrics diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 397b766f41697..7ff529c9e096a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -131,7 +131,7 @@ impl AggregateHashTable { &mut self, batch: &RecordBatch, ) -> Result { - let state = self.state.building(); + let state = self.state.building_mut(); let grouping_set_args = self .group_by_metrics .time_group_key_preparation(|| evaluate_group_by(&state.group_by, batch))?; @@ -144,7 +144,7 @@ impl AggregateHashTable { let mut output = grouping_set_args.into_iter().next().unwrap_or_default(); let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - for (idx, acc) in state.accumulators.iter().enumerate() { + for (idx, acc) in state.accumulators.iter_mut().enumerate() { let values = self.group_by_metrics.time_aggregate_arguments(|| { self.aggregate_argument_metrics .time(idx, || acc.evaluate_row_aligned_args(batch)) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 5811b98d74d76..fb24c8c313cf4 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -17,6 +17,7 @@ //! Aggregate without grouping columns +use crate::aggregates::aggregate_argument::AggregateArgument; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, aggregate_sub_metrics, @@ -47,7 +48,6 @@ use std::task::{Context, Poll}; use super::AggregateExec; use crate::filter::batch_filter; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::stream::{Stream, StreamExt}; /// stream struct for aggregation without grouping columns @@ -68,7 +68,7 @@ struct AggregateStreamInner { schema: SchemaRef, mode: AggregateMode, input: SendableRecordBatchStream, - aggregate_expressions: Vec>>, + aggregate_arguments: Vec>, filter_expressions: Arc<[Option>]>, aggregate_argument_metrics: AggregateArgumentMetrics, aggregate_accumulator_metrics: AggregateAccumulatorMetrics, @@ -122,8 +122,8 @@ impl AggregateStreamInner { guard.clone() }; - let agg_exprs = self - .aggregate_expressions + let agg_args = self + .aggregate_arguments .get(acc_info.aggr_index) .ok_or_else(|| { internal_datafusion_err!( @@ -132,12 +132,15 @@ impl AggregateStreamInner { ) })?; // Only aggregates with a single argument are supported. - let column_expr = agg_exprs.first().ok_or_else(|| { - internal_datafusion_err!( - "Aggregate expression at index {} expected a single argument", - acc_info.aggr_index - ) - })?; + let column_expr = agg_args + .first() + .ok_or_else(|| { + internal_datafusion_err!( + "Aggregate expression at index {} expected a single argument", + acc_info.aggr_index + ) + })? + .expr(); let literal = lit(bound); let predicate: Arc = match acc_info.aggr_type { @@ -300,7 +303,10 @@ impl AggregateStream { let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let input = agg.input.execute(partition, Arc::clone(context))?; - let aggregate_expressions = aggregate_expressions(agg.aggr_expr(), &agg.mode, 0)?; + let aggregate_arguments = aggregate_expressions(agg.aggr_expr(), &agg.mode, 0)? + .into_iter() + .map(|exprs| exprs.into_iter().map(AggregateArgument::new).collect()) + .collect(); let filter_expressions = match agg.mode.input_mode() { AggregateInputMode::Raw => agg_filter_expr, AggregateInputMode::Partial => vec![None; agg.aggr_expr().len()].into(), @@ -361,7 +367,7 @@ impl AggregateStream { mode: agg.mode, input, baseline_metrics, - aggregate_expressions, + aggregate_arguments, filter_expressions, aggregate_argument_metrics, aggregate_accumulator_metrics, @@ -386,7 +392,7 @@ impl AggregateStream { &this.mode, &batch, &mut this.accumulators, - &this.aggregate_expressions, + &mut this.aggregate_arguments, &this.filter_expressions, &this.aggregate_argument_metrics, &this.aggregate_accumulator_metrics, @@ -476,7 +482,7 @@ fn aggregate_batch( mode: &AggregateMode, batch: &RecordBatch, accumulators: &mut [AccumulatorItem], - expressions: &[Vec>], + arguments: &mut [Vec], filters: &[Option>], aggregate_argument_metrics: &AggregateArgumentMetrics, aggregate_accumulator_metrics: &AggregateAccumulatorMetrics, @@ -491,17 +497,20 @@ fn aggregate_batch( // 1.1 accumulators .iter_mut() - .zip(expressions) + .zip(arguments) .zip(filters) .enumerate() - .try_for_each(|(index, ((accum, expr), filter))| { + .try_for_each(|(index, ((accum, arguments), filter))| { // 1.2 and 1.3 let values = aggregate_argument_metrics.time(index, || { let batch = match filter { Some(filter) => Cow::Owned(batch_filter(batch, filter)?), None => Cow::Borrowed(batch), }; - evaluate_expressions_to_arrays(expr, batch.as_ref()) + arguments + .iter_mut() + .map(|argument| argument.evaluate(batch.as_ref())) + .collect::>>() })?; // 1.4 diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6552bd95e06a5..b6f48a9967b3a 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -210,6 +210,7 @@ use itertools::Itertools; use topk::hash_table::is_supported_hash_key_type; use topk::heap::is_supported_heap_type; +mod aggregate_argument; mod aggregate_hash_table; mod aggregate_stream; pub mod group_values;