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-23863-c4a539bd8ef52308726ff3e9e655ee0f51786c96 in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
