diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 058a6a00cd7..9f25ba28ace 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/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`, diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index a01789dead3..2a7689f1dcf 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}; @@ -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, datatype: Option| 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 { 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(()) + } } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a80cdd76f1b..a9b4f10edc8 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/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 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 a5696212145..993b70852c7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/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",