You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/09 18:07:33 UTC
(datafusion) 01/01: perf: optimize aggregate functions over literal arguments (#25849)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25849-5218379fac30106965295f39a692da15b556b6c9
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1c49b7f6d38a6c38ff0763855814db2ef5d4a352
Author: Xiaozhe Yu <66...@qq.com>
AuthorDate: Fri Oct 9 18:07:09 2026 +0000
perf: optimize aggregate functions over literal arguments (#25849)
## Which issue does this PR close?
- Closes #25847.
## Rationale for this change
Unlike scalar functions, aggregate functions cannot be passed scalar
values directly. When an aggregate function is passed a literal argument
(e.g., `count(1)`), for every input batch we expand the literal value
into an 8k-element array before invoking the aggregate. Since we started
with a constant, this repeated scalar -> array conversion is wasted
work, and the cost can be non-trivial: for example, scalar -> array
conversion takes up about 20% of total query execution time for a simple
`count(*)` query on a 5M row in-memory table.
## What changes are included in this PR?
- Add `AggregateArgument` to identify literal aggregate arguments when
aggregate execution is created.
- Evaluate literal arguments once and reuse their expanded arrays across
batches with zero-copy slices.
- Retain expanded arrays up to 1 MiB in ungrouped and hash-table
aggregation paths.
## What is the testing strategy for this PR?
Unit tests cover literal reuse, larger batches, non-literal arguments,
selection handling, and the 1 MiB limit.
A local `release-nonlto` benchmark used one partition, the default batch
size, and a 5,000,000-row in-memory table with a non-null `Int64` column
`v`:
```sql
SELECT COUNT(*) FROM t WHERE v < 900;
```
Cached constant-argument `COUNT` was about 24% faster than uncached
evaluation.
The `aggregate_filter` Criterion benchmark now includes literal
`COUNT(1)` cases. Local comparisons against main show 2–5% improvements
for grouped aggregation.
## Are there any user-facing changes?
No.
---
.../physical-plan/benches/aggregate_filter.rs | 44 ++--
.../src/aggregates/aggregate_argument.rs | 234 +++++++++++++++++++++
.../src/aggregates/aggregate_hash_table/common.rs | 50 ++---
.../aggregate_hash_table/common_ordered.rs | 4 +-
.../aggregate_hash_table/partial_table.rs | 4 +-
.../src/aggregates/aggregate_stream.rs | 43 ++--
datafusion/physical-plan/src/aggregates/mod.rs | 1 +
7 files changed, 321 insertions(+), 59 deletions(-)
diff --git a/datafusion/physical-plan/benches/aggregate_filter.rs b/datafusion/physical-plan/benches/aggregate_filter.rs
index ae97cb7388..417d153abc 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<usize>] = &[
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<AggregateFunctionExpr> {
let value = col("value", schema).unwrap();
- let argument: Arc<dyn PhysicalExpr> = match argument_kind {
- ArgumentKind::Column => value,
- ArgumentKind::Multiply => Arc::new(BinaryExpr::new(
- Arc::clone(&value),
- Operator::Multiply,
- value,
- )),
- };
+ let (function, argument): (Arc<AggregateUDF>, Arc<dyn PhysicalExpr>) =
+ 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 0000000000..56ca9e2837
--- /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<dyn PhysicalExpr>,
+ /// Set when `expr` is a [`Literal`].
+ literal: Option<LiteralArray>,
+}
+
+impl AggregateArgument {
+ pub(in crate::aggregates) fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
+ let literal = expr
+ .downcast_ref::<Literal>()
+ .map(|literal| LiteralArray::new(literal.value().clone()));
+ Self { expr, literal }
+ }
+
+ pub(in crate::aggregates) fn expr(&self) -> &Arc<dyn PhysicalExpr> {
+ &self.expr
+ }
+
+ /// Evaluates the argument against `batch`, returning one value per row.
+ pub(in crate::aggregates) fn evaluate(
+ &mut self,
+ batch: &RecordBatch,
+ ) -> Result<ArrayRef> {
+ 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<ArrayRef> {
+ 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<ArrayRef>,
+}
+
+impl LiteralArray {
+ fn new(value: ScalarValue) -> Self {
+ Self { value, array: None }
+ }
+
+ fn array_of_size(&mut self, num_rows: usize) -> Result<ArrayRef> {
+ 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::<Int32Type>().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::<Int32Type>(),
+ &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::<Int32Type>(),
+ &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::<i32>().values().as_ptr(),
+ second.as_string::<i32>().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::<Int32Type>(),
+ &Int32Array::from(vec![0, 1, 2])
+ );
+ let second = argument.evaluate(&batch(2))?;
+ assert_eq!(
+ second.as_primitive::<Int32Type>(),
+ &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 020196e7da..76e0b0dbe3 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<AggrMode> AggregateHashTable<AggrMode> {
/// See comments in [`EvaluatedAggregateBatch`]
pub(super) fn evaluate_batch(
- &self,
+ &mut self,
batch: &RecordBatch,
) -> Result<EvaluatedAggregateBatch> {
- 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<AggrMode> AggregateHashTable<AggrMode> {
let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| {
state
.accumulators
- .iter()
+ .iter_mut()
.enumerate()
.map(|(idx, acc)| {
self.aggregate_argument_metrics
@@ -522,8 +523,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<Arc<dyn PhysicalExpr>>,
+ /// Example: `CORR(x, y)` stores two arguments here, while `SUM(x)` stores one.
+ arguments: Vec<AggregateArgument>,
/// Optional `FILTER` expression for this accumulator.
///
@@ -716,7 +717,7 @@ impl HashAggregateAccumulator {
) -> Self {
Self {
aggregate_expr,
- arguments,
+ arguments: arguments.into_iter().map(AggregateArgument::new).collect(),
filter,
accumulator,
submetrics,
@@ -730,7 +731,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),
@@ -746,7 +750,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<CompactedAccumulatorArgs> {
let selection = self.evaluate_filter(batch)?;
@@ -767,13 +771,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))
}
})
@@ -791,21 +795,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<RowAlignedAccumulatorArgs> {
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::<Result<_>>()?;
@@ -895,8 +895,8 @@ impl HashAggregateAccumulator {
) -> Result<Vec<ArrayRef>> {
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()
@@ -993,7 +993,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 86c052009f..c25b7485b8 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<AggrMode> OrderedAggregateTable<AggrMode> {
/// 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<EvaluatedAggregateBatch> {
let grouping_set_args =
@@ -250,7 +250,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
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 2e5f92ed2d..144c8cd2fe 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
@@ -130,7 +130,7 @@ impl AggregateHashTable<PartialSkipMarker> {
&mut self,
batch: &RecordBatch,
) -> Result<RecordBatch> {
- 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))?;
@@ -143,7 +143,7 @@ impl AggregateHashTable<PartialSkipMarker> {
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 5811b98d74..fb24c8c313 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<Vec<Arc<dyn PhysicalExpr>>>,
+ aggregate_arguments: Vec<Vec<AggregateArgument>>,
filter_expressions: Arc<[Option<Arc<dyn PhysicalExpr>>]>,
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<dyn PhysicalExpr> = 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<Arc<dyn PhysicalExpr>>],
+ arguments: &mut [Vec<AggregateArgument>],
filters: &[Option<Arc<dyn PhysicalExpr>>],
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::<Result<Vec<_>>>()
})?;
// 1.4
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 6b56bd308b..d4fcccbc9f 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org