feat: optimize count star materialization

datafusion/physical-expr-common/src/utils.rs

@@ -36,7 +36,8 @@ use arrow::datatypes::{
 };
 use arrow::record_batch::RecordBatch;
 use arrow::{downcast_dictionary_array, downcast_primitive_array};
-use datafusion_common::Result;
+use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err};
+use datafusion_expr_common::columnar_value::ColumnarValue;
 use datafusion_expr_common::sort_properties::ExprProperties;
 
 /// Represents a [`PhysicalExpr`] node with associated properties (order and
@@ -425,6 +426,68 @@ pub fn evaluate_expressions_to_arrays_with_metrics<'a>(
         .collect::<Result<Vec<ArrayRef>>>()
 }
 
+/// Largest array, in bytes, that a [`ScalarArrayCache`] keeps between calls.
+const MAX_CACHED_SCALAR_ARRAY_BYTES: usize = 1024 * 1024;
+
+/// Reuses the array expanded from a scalar expression result across batches.
+#[derive(Debug, Default)]
+pub struct ScalarArrayCache {
+    cached: Option<(ScalarValue, ArrayRef)>,
+}
+
+impl ScalarArrayCache {
+    /// Like [`ColumnarValue::into_array_of_size`], but reuses the array built
+    /// for an earlier, equal scalar.
+    pub fn into_array_of_size(
+        &mut self,
+        value: ColumnarValue,
+        num_rows: usize,
+    ) -> Result<ArrayRef> {
+        let scalar = match value {
+            ColumnarValue::Scalar(scalar) => scalar,
+            array @ ColumnarValue::Array(_) => {
+                return array.into_array_of_size(num_rows);
+            }
+        };
+
+        if let Some((cached_scalar, array)) = &self.cached
+            && array.len() >= num_rows
+            && *cached_scalar == scalar
+        {
+            return Ok(array.slice(0, num_rows));
+        }
+
+        let array = scalar.to_array_of_size(num_rows)?;
+        if array.get_array_memory_size() <= MAX_CACHED_SCALAR_ARRAY_BYTES {
+            self.cached = Some((scalar, Arc::clone(&array)));
+        }
+        Ok(array)
+    }
+}
+
+/// Like [`evaluate_expressions_to_arrays`], but expands scalar results through
+/// one cache per expression.
+pub fn evaluate_expressions_to_arrays_with_cache(
+    exprs: &[Arc<dyn PhysicalExpr>],
+    caches: &mut [ScalarArrayCache],
+    batch: &RecordBatch,
+) -> Result<Vec<ArrayRef>> {
+    assert_eq_or_internal_err!(
+        exprs.len(),
+        caches.len(),
+        "expected one scalar array cache per expression"
+    );
+    let num_rows = batch.num_rows();
+    exprs
+        .iter()
+        .zip(caches)
+        .map(|(expr, cache)| {
+            expr.evaluate(batch)
+                .and_then(|value| cache.into_array_of_size(value, num_rows))
+        })
+        .collect()
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -656,4 +719,56 @@ mod tests {
         assert_eq!(scattered.value(4), 50);
         Ok(())
     }
+
+    #[test]
+    fn scalar_array_cache_reuses_equal_scalars() -> Result<()> {
+        let mut cache = ScalarArrayCache::default();
+        let one = || ColumnarValue::Scalar(ScalarValue::Int32(Some(1)));
+
+        let first = cache.into_array_of_size(one(), 4)?;
+        let second = cache.into_array_of_size(one(), 3)?;
+        assert_eq!(as_int32_array(&second)?, &Int32Array::from(vec![1, 1, 1]));
+        assert_eq!(
+            as_int32_array(&second)?.values().as_ptr(),
+            as_int32_array(&first)?.values().as_ptr()
+        );
+
+        let larger = cache.into_array_of_size(one(), 6)?;
+        assert_eq!(as_int32_array(&larger)?, &Int32Array::from(vec![1; 6]));
+        let smaller = cache.into_array_of_size(one(), 5)?;
+        assert_eq!(
+            as_int32_array(&smaller)?.values().as_ptr(),
+            as_int32_array(&larger)?.values().as_ptr()
+        );
+
+        let two = ColumnarValue::Scalar(ScalarValue::Int32(Some(2)));
+        let two = cache.into_array_of_size(two, 2)?;
+        assert_eq!(as_int32_array(&two)?, &Int32Array::from(vec![2, 2]));
+
+        let array: ArrayRef = Arc::new(Int32Array::from(vec![7, 8]));
+        let result =
+            cache.into_array_of_size(ColumnarValue::Array(Arc::clone(&array)), 2)?;
+        assert!(Arc::ptr_eq(&result, &array));
+        assert!(
+            cache
+                .into_array_of_size(ColumnarValue::Array(array), 3)
+                .is_err()
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn scalar_array_cache_skips_large_arrays() -> Result<()> {
+        let mut cache = ScalarArrayCache::default();
+        let value = || ColumnarValue::Scalar(ScalarValue::from("x".repeat(1024)));
+
+        let first = cache.into_array_of_size(value(), 2048)?;
+        let second = cache.into_array_of_size(value(), 2048)?;
+        assert_eq!(first.as_ref(), second.as_ref());
+        assert_ne!(
+            as_string_array(&first).values().as_ptr(),
+            as_string_array(&second).values().as_ptr()
+        );
+        Ok(())
+    }
 }

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -31,6 +31,7 @@ use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use datafusion_expr::{AggregateMetrics, EmitTo, GroupsAccumulator};
 use datafusion_physical_expr::GroupsAccumulatorAdapter;
 use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
+use datafusion_physical_expr_common::utils::ScalarArrayCache;
 use log::debug;
 
 use crate::PhysicalExpr;
@@ -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
@@ -529,6 +530,9 @@ pub(super) struct HashAggregateAccumulator {
     /// Example: `CORR(x, y)` stores two expressions here, while `SUM(x)` stores one.
     arguments: Vec<Arc<dyn PhysicalExpr>>,
 
+    /// One cache per aggregate argument.
+    argument_caches: Vec<ScalarArrayCache>,
+
     /// Optional `FILTER` expression for this accumulator.
     ///
     /// Example: `SUM(x) FILTER (WHERE x > 10)` stores the `x > 10` predicate.
@@ -718,9 +722,14 @@ impl HashAggregateAccumulator {
         accumulator: Box<dyn GroupsAccumulator>,
         submetrics: Arc<dyn AggregateMetrics>,
     ) -> Self {
+        let argument_caches = arguments
+            .iter()
+            .map(|_| ScalarArrayCache::default())
+            .collect();
         Self {
             aggregate_expr,
             arguments,
+            argument_caches,
             filter,
             accumulator,
             submetrics,
@@ -750,7 +759,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)?;
@@ -772,10 +781,12 @@ impl HashAggregateAccumulator {
         let arguments = self
             .arguments
             .iter()
-            .map(|expr| {
+            .zip(&mut self.argument_caches)
+            .map(|(expr, cache)| {
                 if let Some(argument_batch) = argument_batch {
-                    expr.evaluate(argument_batch)
-                        .and_then(|value| value.into_array(argument_batch.num_rows()))
+                    expr.evaluate(argument_batch).and_then(|value| {
+                        cache.into_array_of_size(value, argument_batch.num_rows())
+                    })
                 } else {
                     let data_type = expr.data_type(batch.schema_ref().as_ref())?;
                     Ok(new_empty_array(&data_type))
@@ -795,21 +806,22 @@ 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| {
+            .zip(&mut self.argument_caches)
+            .map(|(expr, cache)| {
                 selection
                     .map_or_else(
                         || expr.evaluate(batch),
                         |selection| expr.evaluate_selection(batch, selection),
                     )
-                    .and_then(|value| value.into_array(batch.num_rows()))
+                    .and_then(|value| cache.into_array_of_size(value, batch.num_rows()))
             })
             .collect::<Result<_>>()?;
 
@@ -997,7 +1009,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(

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

datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs

@@ -131,7 +131,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))?;
@@ -144,7 +144,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))

datafusion/physical-plan/src/aggregates/aggregate_stream.rs

@@ -47,7 +47,9 @@ 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 datafusion_physical_expr_common::utils::{
+    ScalarArrayCache, evaluate_expressions_to_arrays_with_cache,
+};
 use futures::stream::{Stream, StreamExt};
 
 /// stream struct for aggregation without grouping columns
@@ -69,6 +71,7 @@ struct AggregateStreamInner {
     mode: AggregateMode,
     input: SendableRecordBatchStream,
     aggregate_expressions: Vec<Vec<Arc<dyn PhysicalExpr>>>,
+    aggregate_argument_caches: Vec<Vec<ScalarArrayCache>>,
     filter_expressions: Arc<[Option<Arc<dyn PhysicalExpr>>]>,
     aggregate_argument_metrics: AggregateArgumentMetrics,
     aggregate_accumulator_metrics: AggregateAccumulatorMetrics,
@@ -301,6 +304,10 @@ impl AggregateStream {
         let input = agg.input.execute(partition, Arc::clone(context))?;
 
         let aggregate_expressions = aggregate_expressions(agg.aggr_expr(), &agg.mode, 0)?;
+        let aggregate_argument_caches = aggregate_expressions
+            .iter()
+            .map(|exprs| exprs.iter().map(|_| ScalarArrayCache::default()).collect())
+            .collect();
         let filter_expressions = match agg.mode.input_mode() {
             AggregateInputMode::Raw => agg_filter_expr,
             AggregateInputMode::Partial => vec![None; agg.aggr_expr().len()].into(),
@@ -362,6 +369,7 @@ impl AggregateStream {
             input,
             baseline_metrics,
             aggregate_expressions,
+            aggregate_argument_caches,
             filter_expressions,
             aggregate_argument_metrics,
             aggregate_accumulator_metrics,
@@ -387,6 +395,7 @@ impl AggregateStream {
                                 &batch,
                                 &mut this.accumulators,
                                 &this.aggregate_expressions,
+                                &mut this.aggregate_argument_caches,
                                 &this.filter_expressions,
                                 &this.aggregate_argument_metrics,
                                 &this.aggregate_accumulator_metrics,
@@ -472,11 +481,13 @@ impl RecordBatchStream for AggregateStream {
 /// If successful, this returns the additional number of bytes that were allocated during this process.
 ///
 /// TODO: Make this a member function
+#[expect(clippy::too_many_arguments)]
 fn aggregate_batch(
     mode: &AggregateMode,
     batch: &RecordBatch,
     accumulators: &mut [AccumulatorItem],
     expressions: &[Vec<Arc<dyn PhysicalExpr>>],
+    argument_caches: &mut [Vec<ScalarArrayCache>],
     filters: &[Option<Arc<dyn PhysicalExpr>>],
     aggregate_argument_metrics: &AggregateArgumentMetrics,
     aggregate_accumulator_metrics: &AggregateAccumulatorMetrics,
@@ -492,16 +503,17 @@ fn aggregate_batch(
     accumulators
         .iter_mut()
         .zip(expressions)
+        .zip(argument_caches)
         .zip(filters)
         .enumerate()
-        .try_for_each(|(index, ((accum, expr), filter))| {
+        .try_for_each(|(index, (((accum, expr), caches), 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())
+                evaluate_expressions_to_arrays_with_cache(expr, caches, batch.as_ref())
             })?;
 
             // 1.4