perf: let a partial aggregate skip aggregation for AVG

docs/source/user-guide/latest/tuning/operators.md

@@ -172,7 +172,7 @@ pruning does not require that setting.
 Set `spark.comet.exec.aggregate.skipPartial.enabled=true` to let Comet bypass partial hash
 aggregation for high-cardinality grouping when it is not reducing the number of rows enough. This
 experimental optimization is disabled by default. It currently applies only to fused native
-shuffle-writer plans whose partial aggregates are grouping-only or single-argument `COUNT`.
+shuffle-writer plans whose partial aggregates are grouping-only, single-argument `COUNT`, or `AVG` over non-decimal inputs.
 Low-cardinality inputs continue to aggregate normally. The SQL metric
 `rows bypassing partial aggregation` shows whether skipping occurred.
 
@@ -186,7 +186,7 @@ enabling it for a workload.
 
 Eligibility is conservative for the whole fused native plan: any unsupported partial accumulator,
 Spark `PartialMerge`, or mixed-mode aggregate disables skipping in that plan. Multi-argument
-`COUNT` and other accumulators are not admitted. Distribution-required grouping-only stages
+`COUNT`, `AVG` over decimal inputs, and other accumulators are not admitted. Distribution-required grouping-only stages
 still fully deduplicate, and non-native-shuffle plans retain ordinary aggregation.
 
 To experiment with the thresholds, also enable `spark.comet.exec.respectDataFusionConfigs`,

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};
@@ -1035,6 +1036,14 @@ 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)
+                                // Only a DOUBLE result, which the planner builds as Comet's
+                                // Float64 AVG, can convert rows to states. A decimal result
+                                // is an `AvgDecimal`, which cannot.
+                                || 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.
@@ -2713,7 +2722,12 @@ 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::data_type::{
+        data_type_info::DatatypeStruct, DataTypeInfo, DecimalInfo,
+    };
+    use datafusion_comet_proto::spark_expression::{
+        AggExpr, Avg, Count, DataType as ProtoDataType, Expr, Sum,
+    };
     use datafusion_comet_proto::spark_operator::{HashAggregate, ShuffleWriter};
     use std::cell::Cell;
     use std::future::Future;
@@ -2730,6 +2744,27 @@ mod tests {
             expr_struct: Some(AggExprStruct::Sum(Sum::default())),
             ..Default::default()
         };
+        let proto_type = |type_id: DataTypeId| ProtoDataType {
+            type_id: type_id as i32,
+            ..Default::default()
+        };
+        let decimal = ProtoDataType {
+            type_info: Some(Box::new(DataTypeInfo {
+                datatype_struct: Some(DatatypeStruct::Decimal(DecimalInfo {
+                    precision: 38,
+                    scale: 10,
+                })),
+            })),
+            ..proto_type(DataTypeId::Decimal)
+        };
+        let avg = |child: Option<Expr>, datatype: Option<ProtoDataType>| AggExpr {
+            expr_struct: Some(AggExprStruct::Avg(Avg {
+                child,
+                datatype,
+                ..Default::default()
+            })),
+            ..Default::default()
+        };
         let partial = HashAggregate {
             grouping_exprs: vec![Expr::default()],
             agg_exprs: vec![count.clone()],
@@ -2775,7 +2810,14 @@ mod tests {
                 ..partial.clone()
             },
             HashAggregate {
-                agg_exprs: vec![count.clone(), count],
+                agg_exprs: vec![
+                    count.clone(),
+                    avg(Some(Expr::default()), Some(proto_type(DataTypeId::Double))),
+                ],
+                ..partial.clone()
+            },
+            HashAggregate {
+                agg_exprs: vec![count.clone(), count.clone()],
                 ..partial.clone()
             },
         ] {
@@ -2798,6 +2840,32 @@ mod tests {
                 agg_exprs: vec![AggExpr::default()],
                 ..partial.clone()
             },
+            // AVG without an input, without a result type, or with a result type other than
+            // DOUBLE is not admitted. A decimal result is an `AvgDecimal`, which cannot convert
+            // rows to partial states.
+            HashAggregate {
+                agg_exprs: vec![avg(None, Some(proto_type(DataTypeId::Double)))],
+                ..partial.clone()
+            },
+            HashAggregate {
+                agg_exprs: vec![avg(Some(Expr::default()), None)],
+                ..partial.clone()
+            },
+            HashAggregate {
+                agg_exprs: vec![avg(Some(Expr::default()), Some(decimal.clone()))],
+                ..partial.clone()
+            },
+            HashAggregate {
+                agg_exprs: vec![count.clone(), avg(Some(Expr::default()), Some(decimal))],
+                ..partial.clone()
+            },
+            HashAggregate {
+                agg_exprs: vec![avg(
+                    Some(Expr::default()),
+                    Some(proto_type(DataTypeId::Float)),
+                )],
+                ..partial.clone()
+            },
             HashAggregate {
                 agg_exprs: vec![AggExpr {
                     expr_struct: Some(AggExprStruct::Count(Count {

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<Vec<ArrayRef>> {
-        not_impl_err!("Input batch conversion to state not implemented")
+        let values = values[0].as_primitive::<T>();
+        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::<T>::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::<arrow::datatypes::Float64Type, _>::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(())
+    }
 }

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -718,8 +718,9 @@ 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 " +
+          "and whose aggregate functions, if any, are all single-argument COUNT or AVG over " +
+          "non-decimal inputs. 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.")
       .booleanConf

spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala

@@ -3431,6 +3431,59 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("skip partial aggregation converts AVG rows to partial states") {
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "input").toUri.toString
+      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+        // Mostly distinct keys within each task, with some keys repeating across tasks so the
+        // final aggregate merges converted states with each other. Values include nulls.
+        spark
+          .range(0L, 16384L, 1L, 8)
+          .selectExpr(
+            "id % 12000 AS k",
+            "CASE WHEN id % 5 = 0 THEN NULL ELSE CAST(id % 17 AS INT) END AS i",
+            "CASE WHEN id % 7 = 0 THEN NULL ELSE CAST(id % 11 AS DOUBLE) / 3 END AS d",
+            "CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id % 13 AS DECIMAL(10, 2)) / 7 END AS small",
+            "CASE WHEN id % 4 = 0 THEN NULL ELSE CAST(id % 19 AS DECIMAL(38, 10)) / 7 END AS big")
+          .write
+          .parquet(path)
+      }
+      withParquetTable(path, "skip_partial_avg") {
+        withSQLConf(
+          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+          SQLConf.SHUFFLE_PARTITIONS.key -> "8",
+          CometConf.COMET_BATCH_SIZE.key -> "128",
+          CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+          CometConf.COMET_SHUFFLE_MODE.key -> "native",
+          CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED.key -> "true",
+          CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
+          "spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold" -> "100",
+          "spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "0.8") {
+          def skipped(query: String): Long = {
+            val (_, plan) = checkSparkAnswerAndOperator(query)
+            stripAQEPlan(plan).collect { case aggregate: CometHashAggregateExec =>
+              aggregate.metrics.get("skipped_aggregation_rows").map(_.value).getOrElse(0L)
+            }.sum
+          }
+          def avgQuery(expressions: String): String =
+            s"SELECT k, $expressions FROM skip_partial_avg GROUP BY k"
+
+          // AVG over integral and floating-point inputs stops aggregating.
+          for (expressions <- Seq("avg(i)", "avg(d)", "avg(i), avg(d), count(i)")) {
+            assert(skipped(avgQuery(expressions)) > 0L, s"Expected skipping for $expressions")
+          }
+          // AVG over a decimal input that Comet aggregates as a decimal cannot convert rows to
+          // partial states, so it never stops aggregating.
+          for (expressions <- Seq("avg(big)", "avg(i), avg(big)")) {
+            assert(skipped(avgQuery(expressions)) == 0L, s"Unexpected skipping for $expressions")
+          }
+          // Whether or not a small decimal skips, the answer matches Spark's.
+          skipped(avgQuery("avg(small)"))
+        }
+      }
+    }
+  }
+
   test("first/last") {
     withSQLConf(
       SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",