This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 3e3a92de29 feat: Calculate non-distinct `sum` from column statistics
when available (#23863)
3e3a92de29 is described below
commit 3e3a92de29ed3d454e72c7bade6328508b6098c6
Author: Adam Gutglick <[email protected]>
AuthorDate: Wed Aug 5 14:20:27 2026 +0100
feat: Calculate non-distinct `sum` from column statistics when available
(#23863)
## Which issue does this PR close?
- Closes #23858.
## Rationale for this change
Seems like a valuable optimization for sources that report the `sum`
statistic.
## What changes are included in this PR?
Implements `Sum::value_from_stats`, supporting direct column expressions
and `cast(col)` expressions introduced by type coercion.
Casts direct-column statistics to the aggregate return type so decimal
sums use the widened `SUM` type.
## Are these changes tested?
Unit coverage for statistics-backed `SUM` values and decimal accumulator
widening, plus core integration coverage for integer and decimal sums
from statistics.
## Are there any user-facing changes?
None
Signed-off-by: Adam Gutglick <[email protected]>
---
.../physical_optimizer/aggregate_statistics.rs | 221 ++++++++++++++++++++-
datafusion/functions-aggregate/src/sum.rs | 129 +++++++++++-
2 files changed, 347 insertions(+), 3 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
index 0fa60ae20d..2d22b60856 100644
--- a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
+++ b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
@@ -29,16 +29,17 @@ use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::physical_plan::ParquetSource;
use datafusion::datasource::source::DataSourceExec;
use datafusion::prelude::{SessionConfig, SessionContext};
-use datafusion_common::assert_batches_eq;
use datafusion_common::cast::as_int64_array;
use datafusion_common::config::ConfigOptions;
use datafusion_common::stats::Precision;
use datafusion_common::{ColumnStatistics, Result, Statistics};
+use datafusion_common::{ScalarValue, assert_batches_eq};
use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
use datafusion_execution::TaskContext;
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_expr::Operator;
use datafusion_functions_aggregate::count::count_udaf;
+use datafusion_functions_aggregate::sum::sum_udaf;
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
use datafusion_physical_expr::expressions::{self, cast};
use datafusion_physical_optimizer::PhysicalOptimizerRule;
@@ -637,3 +638,221 @@ async fn topk_distinct_preserves_nulls() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_sum_from_statistics() -> Result<()> {
+ enum SumArg {
+ ColumnA,
+ ColumnB,
+ CastColumnA(DataType),
+ Binary,
+ }
+
+ struct TestCase {
+ name: &'static str,
+ data_type: DataType,
+ sum_value_a: Precision<ScalarValue>,
+ sum_value_b: Precision<ScalarValue>,
+ sum_arg: SumArg,
+ is_distinct: bool,
+ expected_value: Option<ScalarValue>,
+ }
+
+ for case in [
+ TestCase {
+ name: "exact statistics",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))),
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::ColumnA,
+ is_distinct: false,
+ expected_value: Some(ScalarValue::Int64(Some(10))),
+ },
+ TestCase {
+ name: "second column statistics",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))),
+ sum_value_b: Precision::Exact(ScalarValue::Int64(Some(42))),
+ sum_arg: SumArg::ColumnB,
+ is_distinct: false,
+ expected_value: Some(ScalarValue::Int64(Some(42))),
+ },
+ TestCase {
+ name: "casted int32 column statistics",
+ data_type: DataType::Int32,
+ sum_value_a: Precision::Exact(ScalarValue::Int32(Some(10))),
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::CastColumnA(DataType::Int64),
+ is_distinct: false,
+ expected_value: Some(ScalarValue::Int64(Some(10))),
+ },
+ TestCase {
+ name: "decimal statistics uses aggregate return type",
+ data_type: DataType::Decimal128(5, 2),
+ sum_value_a: Precision::Exact(ScalarValue::Decimal128(Some(12345),
5, 2)),
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::ColumnA,
+ is_distinct: false,
+ expected_value: Some(ScalarValue::Decimal128(Some(12345), 15, 2)),
+ },
+ TestCase {
+ name: "inexact statistics",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Inexact(ScalarValue::Int64(Some(10))),
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::ColumnA,
+ is_distinct: false,
+ expected_value: None,
+ },
+ TestCase {
+ name: "absent statistics",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Absent,
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::ColumnA,
+ is_distinct: false,
+ expected_value: None,
+ },
+ TestCase {
+ name: "null statistics",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Exact(ScalarValue::Int64(None)),
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::ColumnA,
+ is_distinct: false,
+ expected_value: None,
+ },
+ TestCase {
+ name: "binary expr",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))),
+ sum_value_b: Precision::Exact(ScalarValue::Int64(Some(42))),
+ sum_arg: SumArg::Binary,
+ is_distinct: false,
+ expected_value: None,
+ },
+ TestCase {
+ name: "distinct sum",
+ data_type: DataType::Int64,
+ sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))),
+ sum_value_b: Precision::Absent,
+ sum_arg: SumArg::ColumnA,
+ is_distinct: true,
+ expected_value: None,
+ },
+ ] {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", case.data_type.clone(), true),
+ Field::new("b", case.data_type.clone(), true),
+ ]));
+
+ let statistics = Statistics {
+ num_rows: Precision::Absent,
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![
+ ColumnStatistics {
+ sum_value: case.sum_value_a,
+ ..Default::default()
+ },
+ ColumnStatistics {
+ sum_value: case.sum_value_b,
+ ..Default::default()
+ },
+ ],
+ };
+
+ let config = FileScanConfigBuilder::new(
+ ObjectStoreUrl::parse("test:///").unwrap(),
+ Arc::new(ParquetSource::new(Arc::clone(&schema))),
+ )
+ .with_file(PartitionedFile::new("x".to_string(), 100))
+ .with_statistics(statistics)
+ .build();
+
+ let source: Arc<dyn ExecutionPlan> =
DataSourceExec::from_data_source(config);
+ let schema = source.schema();
+
+ let (agg_args, alias): (Vec<Arc<dyn
datafusion_physical_expr::PhysicalExpr>>, _) =
+ match case.sum_arg {
+ SumArg::ColumnA => (vec![expressions::col("a", &schema)?],
"SUM(a)"),
+ SumArg::ColumnB => (vec![expressions::col("b", &schema)?],
"SUM(b)"),
+ SumArg::CastColumnA(cast_type) => (
+ vec![cast(expressions::col("a", &schema)?, &schema,
cast_type)?],
+ "SUM(CAST(a))",
+ ),
+ SumArg::Binary => (
+ vec![expressions::binary(
+ expressions::col("a", &schema)?,
+ Operator::Plus,
+ expressions::col("b", &schema)?,
+ &schema,
+ )?],
+ "SUM(a + b)",
+ ),
+ };
+
+ let sum_expr_builder = AggregateExprBuilder::new(sum_udaf(), agg_args)
+ .schema(Arc::clone(&schema))
+ .alias(alias);
+ let sum_expr_builder = if case.is_distinct {
+ sum_expr_builder.distinct()
+ } else {
+ sum_expr_builder
+ };
+ let sum_expr = sum_expr_builder.build()?;
+
+ let partial_agg = AggregateExec::try_new(
+ AggregateMode::Partial,
+ PhysicalGroupBy::default(),
+ vec![Arc::new(sum_expr.clone())],
+ vec![None],
+ source,
+ Arc::clone(&schema),
+ )?;
+
+ let final_agg = AggregateExec::try_new(
+ AggregateMode::Final,
+ PhysicalGroupBy::default(),
+ vec![Arc::new(sum_expr)],
+ vec![None],
+ Arc::new(partial_agg),
+ Arc::clone(&schema),
+ )?;
+
+ let conf = ConfigOptions::new();
+ let optimized =
+ AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;
+
+ if let Some(expected_value) = case.expected_value {
+ assert!(
+ optimized.is::<ProjectionExec>(),
+ "'{}': expected ProjectionExec",
+ case.name
+ );
+
+ let task_ctx = Arc::new(TaskContext::default());
+ let result = common::collect(optimized.execute(0,
task_ctx)?).await?;
+ assert_eq!(result.len(), 1, "'{}': expected 1 batch", case.name);
+ assert_eq!(
+ result[0].schema().field(0).data_type(),
+ &expected_value.data_type(),
+ "'{}': unexpected data type",
+ case.name
+ );
+ assert_eq!(
+ ScalarValue::try_from_array(result[0].column(0), 0)?,
+ expected_value,
+ "'{}': unexpected value",
+ case.name
+ );
+ } else {
+ assert!(
+ optimized.is::<AggregateExec>(),
+ "'{}': expected AggregateExec (not optimized)",
+ case.name
+ );
+ }
+ }
+
+ Ok(())
+}
diff --git a/datafusion/functions-aggregate/src/sum.rs
b/datafusion/functions-aggregate/src/sum.rs
index 8d1df28559..c124c7a1a0 100644
--- a/datafusion/functions-aggregate/src/sum.rs
+++ b/datafusion/functions-aggregate/src/sum.rs
@@ -29,6 +29,7 @@ use arrow::datatypes::{
};
use datafusion_common::hash_utils::RandomState;
use datafusion_common::internal_err;
+use datafusion_common::stats::Precision;
use datafusion_common::types::{
NativeType, logical_float64, logical_int8, logical_int16, logical_int32,
logical_int64, logical_uint8, logical_uint16, logical_uint32,
logical_uint64,
@@ -40,12 +41,13 @@ use datafusion_expr::function::{AccumulatorArgs,
StateFieldsArgs};
use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name};
use datafusion_expr::{
Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr,
GroupsAccumulator,
- Operator, ReversedUDAF, SetMonotonicity, Signature, TypeSignature,
+ Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs,
TypeSignature,
TypeSignatureClass, Volatility,
};
use
datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator;
use
datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator;
use datafusion_macros::user_doc;
+use datafusion_physical_expr::expressions::{CastExpr, Column};
use std::mem::size_of_val;
make_udaf_expr_and_func!(
@@ -410,6 +412,58 @@ impl AggregateUDFImpl for Sum {
// SUM(arg) + lit * COUNT(arg)
Ok(Some(sum_agg + (lit.clone() * count_agg)))
}
+
+ fn value_from_stats(&self, statistics_args: &StatisticsArgs) ->
Option<ScalarValue> {
+ if statistics_args.is_distinct {
+ return None;
+ }
+
+ let [expr] = statistics_args.exprs else {
+ return None;
+ };
+
+ let (col_expr, cast_type) = match expr.downcast_ref::<Column>() {
+ Some(col_expr) => (col_expr, None),
+ None => {
+ let cast_expr = expr.downcast_ref::<CastExpr>()?;
+ let col_expr = cast_expr.expr().downcast_ref::<Column>()?;
+ (col_expr, Some(cast_expr.cast_type()))
+ }
+ };
+
+ let col_stats = statistics_args
+ .statistics
+ .column_statistics
+ .get(col_expr.index())?;
+
+ // Replacing SUM with a literal is only valid for exact statistics.
+ // `cast_to_sum_type` also widens small integer stats to the SQL SUM
+ // return type, e.g. Int32 statistics become an Int64 SUM value.
+ let Precision::Exact(val) = col_stats.sum_value.cast_to_sum_type()
else {
+ return None;
+ };
+ if val.is_null() {
+ return None;
+ }
+
+ // SUM coercion can introduce a physical CAST around the input column
+ // (`SUM(Int32)` becomes `SUM(CAST(Int32 AS Int64))`). Only use the
+ // column's raw sum stats when the widened stats value matches that
+ // cast target and the aggregate return type.
+ if let Some(cast_type) = cast_type {
+ let value_type = val.data_type();
+ if cast_type != statistics_args.return_type || &value_type !=
cast_type {
+ return None;
+ }
+ return Some(val);
+ }
+
+ if &val.data_type() == statistics_args.return_type {
+ Some(val)
+ } else {
+ val.cast_to(statistics_args.return_type).ok()
+ }
+ }
}
/// This accumulator computes SUM incrementally
@@ -665,7 +719,7 @@ impl Accumulator for SlidingDistinctSumAccumulator {
mod tests {
use super::*;
use arrow::{
- array::Int64Array,
+ array::{Decimal128Array, Int64Array},
buffer::{NullBuffer, ScalarBuffer},
};
use std::sync::Arc;
@@ -709,4 +763,75 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn decimal_sum_accumulator_uses_widened_return_type() -> Result<()> {
+ let values: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![Some(99_999), Some(99_999)])
+ .with_precision_and_scale(5, 2)?,
+ );
+ let mut acc =
SumAccumulator::<Decimal128Type>::new(DataType::Decimal128(15, 2));
+
+ acc.update_batch(&[values])?;
+
+ assert_eq!(
+ acc.evaluate()?,
+ ScalarValue::Decimal128(Some(199_998), 15, 2)
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn sum_value_from_stats_widens_small_integer_sum() {
+ let statistics = datafusion_common::Statistics {
+ num_rows: Precision::Absent,
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![datafusion_common::ColumnStatistics {
+ sum_value: Precision::Exact(ScalarValue::Int32(Some(10))),
+ ..Default::default()
+ }],
+ };
+ let return_type = DataType::Int64;
+ let expr: Arc<dyn datafusion_physical_expr::PhysicalExpr> =
+ Arc::new(Column::new("a", 0));
+ let exprs = vec![expr];
+ let statistics_args = StatisticsArgs {
+ statistics: &statistics,
+ return_type: &return_type,
+ is_distinct: false,
+ exprs: &exprs,
+ };
+
+ assert_eq!(
+ Sum::new().value_from_stats(&statistics_args),
+ Some(ScalarValue::Int64(Some(10)))
+ );
+ }
+
+ #[test]
+ fn sum_value_from_stats_casts_decimal_sum_to_return_type() {
+ let statistics = datafusion_common::Statistics {
+ num_rows: Precision::Absent,
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![datafusion_common::ColumnStatistics {
+ sum_value:
Precision::Exact(ScalarValue::Decimal128(Some(12345), 5, 2)),
+ ..Default::default()
+ }],
+ };
+ let return_type = DataType::Decimal128(15, 2);
+ let expr: Arc<dyn datafusion_physical_expr::PhysicalExpr> =
+ Arc::new(Column::new("a", 0));
+ let exprs = vec![expr];
+ let statistics_args = StatisticsArgs {
+ statistics: &statistics,
+ return_type: &return_type,
+ is_distinct: false,
+ exprs: &exprs,
+ };
+
+ assert_eq!(
+ Sum::new().value_from_stats(&statistics_args),
+ Some(ScalarValue::Decimal128(Some(12345), 15, 2))
+ );
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
