From be64cb098f101152d3e04cbfdd8bc759a862d00c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 2 Nov 2026 10:38:53 -0700 Subject: [PATCH 1/2] feat: convert input rows to partial states in Comet's double average A partial aggregate that stops aggregating sends each row on as a partial state, which DataFusion builds with GroupsAccumulator::convert_to_state. Comet's Float64 AvgGroupsAccumulator, which runs Spark's average of a decimal as avg(UnscaledValue(x)), returned NotImplemented. Each row's state is now its value and a count of one, or a zero sum and count for a null or filtered-out row, what state() emits for such a group. (cherry picked from commit 5000f3a837ca1d24ff9cd8233ad4c5d2bea4aa50) --- native/spark-expr/src/agg_funcs/avg.rs | 55 ++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 3 deletions(-) diff --git a/native/spark-expr/src/agg_funcs/avg.rs b/native/spark-expr/src/agg_funcs/avg.rs index 2b2422ce775..c44d30eda2a 100644 --- a/native/spark-expr/src/agg_funcs/avg.rs +++ b/native/spark-expr/src/agg_funcs/avg.rs @@ -334,12 +334,33 @@ where ]) } + /// Each row's own partial state, as `state` emits it: the value and a count of one, or a + /// zero sum and count for a null or filtered-out row. fn convert_to_state( &self, - _values: &[ArrayRef], - _opt_filter: Option<&arrow::array::BooleanArray>, + values: &[ArrayRef], + opt_filter: Option<&arrow::array::BooleanArray>, ) -> Result> { - not_impl_err!("Input batch conversion to state not implemented") + let values = values[0].as_primitive::(); + let mut sums = Vec::with_capacity(values.len()); + let mut counts = Vec::with_capacity(values.len()); + for idx in 0..values.len() { + let included = values.is_valid(idx) + && opt_filter.is_none_or(|filter| filter.is_valid(idx) && filter.value(idx)); + if included { + sums.push(values.value(idx)); + counts.push(1); + } else { + sums.push(T::default_value()); + counts.push(0); + } + } + let sums = PrimitiveArray::::new(sums.into(), None) + .with_data_type(self.return_data_type.clone()); + Ok(vec![ + Arc::new(sums) as ArrayRef, + Arc::new(Int64Array::new(counts.into(), None)) as ArrayRef, + ]) } fn size(&self) -> usize { @@ -365,4 +386,32 @@ mod tests { assert_eq!(partial.evaluate()?, ScalarValue::Float64(None)); Ok(()) } + + /// Converting rows straight to partial states, as a partial aggregate that stops aggregating + /// does, gives each row the state that aggregating it alone into a group would. + #[test] + fn convert_to_state_matches_one_row_groups() -> Result<()> { + use arrow::array::BooleanArray; + let values: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.5), + None, + Some(-2.0), + Some(4.0), + ])); + let filter = BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]); + let new_accumulator = || { + AvgGroupsAccumulator::::new( + &Float64, + |sum: f64, count: i64| Ok(sum / count as f64), + ) + }; + + let converted = + new_accumulator().convert_to_state(&[Arc::clone(&values)], Some(&filter))?; + + let mut grouped = new_accumulator(); + grouped.update_batch(&[values], &[0, 1, 2, 3], Some(&filter), 4)?; + assert_eq!(converted, grouped.state(EmitTo::All)?); + Ok(()) + } } From d7261c8357fa03117034d6398babd8012b55a784 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 2 Nov 2026 12:52:57 -0700 Subject: [PATCH 2/2] feat: let a partial aggregate with a double AVG stop aggregating configure_skip_partial_aggregation only admitted single-argument COUNT. With the double AvgGroupsAccumulator now converting rows to partial states, a single-argument AVG whose result is a double is admitted too. A decimal AVG runs as AvgDecimal, whose convert_to_state is not implemented, so it stays excluded. --- native/core/src/execution/jni_api.rs | 35 ++++++++++++++++++- .../scala/org/apache/comet/CometConf.scala | 9 ++--- .../comet/exec/CometAggregateSuite.scala | 12 ++++++- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d915abfed0b..bc2bc835334 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -61,6 +61,7 @@ use datafusion::{ }; use datafusion_comet_common::decode_string_arrays; use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct; +use datafusion_comet_proto::spark_expression::data_type::DataTypeId; use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan}; use datafusion_comet_shuffle::decoded_batch_coalescer::DecodedBatchCoalescer; use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl}; @@ -1034,6 +1035,12 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato && agg.agg_exprs.iter().all(|expr| { matches!(&expr.expr_struct, Some(AggExprStruct::Count(count)) if count.children.len() == 1) + // A decimal average's accumulator can't convert rows to states. + || matches!(&expr.expr_struct, Some(AggExprStruct::Avg(avg)) + if avg.child.is_some() + && avg.datatype.as_ref().is_some_and(|datatype| { + datatype.type_id == DataTypeId::Double as i32 + })) }) } // PartialMerge consumes states, not rows. @@ -2712,7 +2719,7 @@ mod tests { use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::ExecutionPlan; use datafusion_comet_proto::spark_expression; - use datafusion_comet_proto::spark_expression::{AggExpr, Count, Expr, Sum}; + use datafusion_comet_proto::spark_expression::{AggExpr, Avg, Count, Expr, Sum}; use datafusion_comet_proto::spark_operator::{HashAggregate, ShuffleWriter}; use std::cell::Cell; use std::future::Future; @@ -2729,6 +2736,17 @@ mod tests { expr_struct: Some(AggExprStruct::Sum(Sum::default())), ..Default::default() }; + let avg = |child: Option, type_id: DataTypeId| AggExpr { + expr_struct: Some(AggExprStruct::Avg(Avg { + child, + datatype: Some(spark_expression::DataType { + type_id: type_id as i32, + type_info: None, + }), + ..Default::default() + })), + ..Default::default() + }; let partial = HashAggregate { grouping_exprs: vec![Expr::default()], agg_exprs: vec![count.clone()], @@ -2773,6 +2791,13 @@ mod tests { agg_exprs: vec![], ..partial.clone() }, + HashAggregate { + agg_exprs: vec![ + count.clone(), + avg(Some(Expr::default()), DataTypeId::Double), + ], + ..partial.clone() + }, HashAggregate { agg_exprs: vec![count.clone(), count], ..partial.clone() @@ -2797,6 +2822,14 @@ mod tests { agg_exprs: vec![AggExpr::default()], ..partial.clone() }, + HashAggregate { + agg_exprs: vec![avg(None, DataTypeId::Double)], + ..partial.clone() + }, + HashAggregate { + agg_exprs: vec![avg(Some(Expr::default()), DataTypeId::Decimal)], + ..partial.clone() + }, HashAggregate { agg_exprs: vec![AggExpr { expr_struct: Some(AggExprStruct::Count(Count { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 824118f3efe..742662f4006 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -710,10 +710,11 @@ object CometConf extends ShimCometConf { "Experimental opt-in: let a native partial aggregate stop aggregating once its input " + "looks mostly distinct, and send the rest of the task's rows to the shuffle " + "unaggregated. Only applies to partial aggregates that feed Comet native shuffle " + - "and whose aggregate functions, if any, are all single-argument COUNT. The check " + - "starts after the first 100,000 input rows of a task, and once aggregation stops " + - "it does not resume, so a task whose keys repeat after a mostly distinct start can " + - s"shuffle many times more rows than it would with this disabled. $TUNING_GUIDE.") + "and whose aggregate functions, if any, are all single-argument COUNT or AVG of a " + + "non-decimal value. The check starts after the first 100,000 input rows of a task, " + + "and once aggregation stops it does not resume, so a task whose keys repeat after a " + + "mostly distinct start can shuffle many times more rows than it would with this " + + s"disabled. $TUNING_GUIDE.") .booleanConf .createWithDefault(false) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 2b9156346c5..ca319df0767 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -3254,7 +3254,17 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } withSQLConf( "spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "0.8") { - for (expression <- Seq("sum(v)", "count(v, w)", "count(*) + sum(v)")) { + // A non-decimal AVG converts rows to partial states, including null rows. + assert( + aggregates("SELECT k, avg(v), avg(w) FROM skip_partial_eligibility GROUP BY k") + .map(skipped) + .sum > 0L) + for (expression <- Seq( + "sum(v)", + "count(v, w)", + "count(*) + sum(v)", + // A decimal AVG's accumulator can't convert rows to partial states. + "avg(CAST(v AS DECIMAL(30, 2)))")) { val result = aggregates( "SELECT sum(n) FROM " + s"(SELECT k, $expression n FROM skip_partial_eligibility GROUP BY k)")