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-5723-d03cb7b9ccba8cb0e82b2672a56af806d05466b3 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit cf639fdec71e0337b42bf0068adbb3d1813edde3 Author: Peter Lee <[email protected]> AuthorDate: Sat Sep 12 20:10:49 2026 +0000 feat: Enable adaptive partial aggregation for eligible native shuffle plans (#5723) * Enable adaptive partial aggregation for eligible native shuffle plans * Fix redundant interpolation in skip-partial regression * docs: explain adaptive aggregation tuning and disablement --- docs/source/user-guide/latest/tuning.md | 40 +++++ native/core/src/execution/jni_api.rs | 171 +++++++++++++++++++-- .../apache/spark/sql/comet/CometMetricNode.scala | 3 + .../org/apache/spark/sql/comet/operators.scala | 9 ++ .../apache/comet/exec/CometAggregateSuite.scala | 104 ++++++++++++- 5 files changed, 313 insertions(+), 14 deletions(-) diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index ceec8b1a15..25b805d1cc 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -31,6 +31,46 @@ It is recommended that `COMET_WORKER_THREADS` be set to the number of executor c in some environments, such as Kubernetes, where the number of cores allocated to a pod will already be equal to the number of executor cores. +## Adaptive Partial Aggregation + +For high-cardinality grouping, Comet can bypass partial hash aggregation when it is not +reducing the number of rows enough. This currently applies only to fused native shuffle-writer +plans whose partial aggregates are grouping-only or single-argument `COUNT`. Low-cardinality +inputs continue to aggregate normally. The SQL metric `rows bypassing partial aggregation` +shows whether skipping occurred. + +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 +still fully deduplicate, and non-native-shuffle plans retain ordinary aggregation. +The DataFusion testing configuration override does not bypass these safety checks. + +DataFusion 55 defaults to probing after 100,000 input rows per partial aggregation +partition and skipping when the number of groups divided by input rows exceeds `0.8`. +To experiment with these thresholds, enable `spark.comet.exec.respectDataFusionConfigs`, +a development and testing option that defaults to `false`. For example, the following +SQL settings pass through the default threshold values, which you can adjust: + +```sql +SET spark.comet.exec.respectDataFusionConfigs=true; +SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold=100000; +SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold=0.8; +``` + +A lower row threshold allows an earlier decision; a lower ratio threshold makes +skipping more likely. Skipping can increase the number of partial states emitted +and the amount of shuffle data, so measure the effect on your workload. + +To disable skipping, keep `spark.comet.exec.respectDataFusionConfigs=true` and set +the ratio threshold above the maximum possible groups/input-rows ratio: + +```sql +SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold=1.1; +``` + +These settings only tune eligible plans. Unsupported accumulators and modes remain +disabled even when configuration overrides are enabled. + ## Memory Tuning It is necessary to specify how much memory Comet can use in addition to memory already allocated to Spark. In some diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 65a2d68ec1..692b0d1bcc 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -31,7 +31,7 @@ use std::collections::HashSet; use arrow::array::{Array, RecordBatch, UInt32Array}; use arrow::compute::{take, TakeOptions}; use arrow::datatypes::DataType as ArrowDataType; -use datafusion::common::{DataFusionError, Result as DataFusionResult, ScalarValue}; +use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -41,7 +41,8 @@ use datafusion::{ physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream}, prelude::{SessionConfig, SessionContext}, }; -use datafusion_comet_proto::spark_operator::{Operator, ShuffleScan}; +use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct; +use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan}; use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl}; use datafusion_spark::function::array::array_contains::SparkArrayContains; use datafusion_spark::function::array::repeat::SparkArrayRepeat; @@ -551,6 +552,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( max_temp_directory_size, task_cpus as usize, &spark_config, + &spark_plan, )?; let plan_creation_time = start.elapsed(); @@ -668,6 +670,42 @@ pub extern "system" fn Java_org_apache_comet_Native_setShufflePartitionPusher( }) } +/// Only admit the validated native-shuffle path. A session belongs to one fused Spark plan, +/// so an unsafe partial aggregate disables skipping for the whole plan, including its children. +/// This deliberately gives up some opportunities rather than changing execution contexts per op. +fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operator) { + fn supported(plan: &Operator) -> bool { + let supported_aggregate = match &plan.op_struct { + Some(OpStruct::HashAgg(agg)) => match AggregateMode::try_from(agg.mode) { + // Final never skips. Still inspect its children below. + Ok(AggregateMode::Final) => true, + Ok(AggregateMode::Partial) => { + agg.expr_modes + .iter() + .all(|mode| *mode == AggregateMode::Partial as i32) + && agg.agg_exprs.iter().all(|expr| { + matches!(&expr.expr_struct, Some(AggExprStruct::Count(count)) + if count.children.len() == 1) + }) + } + // PartialMerge is represented as native Partial, but consumes states, not rows. + _ => false, + }, + _ => true, + }; + supported_aggregate && plan.children.iter().all(supported) + } + + if !matches!(&plan.op_struct, Some(OpStruct::ShuffleWriter(_))) || !supported(plan) { + // Enforce safety after config pass-through: a testing override cannot make unsupported + // accumulators convertible. DF 55 removed supports_convert_to_state(). + config + .options_mut() + .execution + .skip_partial_aggregation_probe_ratio_threshold = 1.1; + } +} + /// Configure DataFusion session context. fn prepare_datafusion_session_context( batch_size: usize, @@ -676,6 +714,7 @@ fn prepare_datafusion_session_context( max_temp_directory_size: u64, task_cpus: usize, spark_config: &HashMap<String, String>, + spark_plan: &Operator, ) -> CometResult<SessionContext> { let paths = local_dirs.into_iter().map(PathBuf::from).collect(); let disk_manager = DiskManagerBuilder::default() @@ -689,17 +728,7 @@ fn prepare_datafusion_session_context( // This DataFusion context is within the scope of an executing Spark Task. We want to set // its internal parallelism to the number of CPUs allocated to Spark Tasks. This can be // modified by changing spark.task.cpus in the Spark config. - .with_batch_size(batch_size) - // DataFusion partial aggregates can emit duplicate rows so we disable the - // skip partial aggregation feature because this is not compatible with Spark's - // use of partial aggregates. - .set( - "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", - // this is the threshold of number of groups / number of rows and the - // maximum value is 1.0, so we set the threshold a little higher just - // to be safe - &ScalarValue::Float64(Some(1.1)), - ); + .with_batch_size(batch_size); // Translate the Comet-namespaced row-level pushdown flag into the equivalent // DataFusion session options. `pushdown_filters` enables the parquet reader's @@ -726,6 +755,8 @@ fn prepare_datafusion_session_context( } } + configure_skip_partial_aggregation(&mut session_config, spark_plan); + let runtime = rt_config.build()?; let mut session_ctx = SessionContext::new_with_config_rt(session_config, Arc::new(runtime)); @@ -1590,6 +1621,120 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( mod tests { use super::*; use datafusion::execution::memory_pool::{MemoryConsumer, UnboundedMemoryPool}; + use datafusion_comet_proto::spark_expression::{AggExpr, Count, Expr, Sum}; + use datafusion_comet_proto::spark_operator::{HashAggregate, ShuffleWriter}; + + #[test] + fn skip_partial_eligibility_is_fail_closed() { + let count = AggExpr { + expr_struct: Some(AggExprStruct::Count(Count { + children: vec![Expr::default()], + })), + ..Default::default() + }; + let sum = AggExpr { + expr_struct: Some(AggExprStruct::Sum(Sum::default())), + ..Default::default() + }; + let partial = HashAggregate { + grouping_exprs: vec![Expr::default()], + agg_exprs: vec![count.clone()], + mode: AggregateMode::Partial as i32, + ..Default::default() + }; + let writer = |agg: HashAggregate| Operator { + op_struct: Some(OpStruct::ShuffleWriter(ShuffleWriter::default())), + children: vec![Operator { + op_struct: Some(OpStruct::HashAgg(agg)), + ..Default::default() + }], + ..Default::default() + }; + let ratio = |plan: &Operator, requested: f64| { + let mut config = SessionConfig::new(); + config + .options_mut() + .execution + .skip_partial_aggregation_probe_rows_threshold = 37; + config + .options_mut() + .execution + .skip_partial_aggregation_probe_ratio_threshold = requested; + configure_skip_partial_aggregation(&mut config, plan); + assert_eq!( + config + .options() + .execution + .skip_partial_aggregation_probe_rows_threshold, + 37 + ); + config + .options() + .execution + .skip_partial_aggregation_probe_ratio_threshold + }; + + for agg in [ + partial.clone(), + HashAggregate { + agg_exprs: vec![], + ..partial.clone() + }, + HashAggregate { + agg_exprs: vec![count.clone(), count], + ..partial.clone() + }, + ] { + let plan = writer(agg); + assert_eq!(ratio(&plan, 0.8), 0.8); + assert_eq!(ratio(&plan, 0.5), 0.5); + assert_eq!(ratio(&plan, 1.1), 1.1); + // Non-native shuffle / standalone native blocks stay disabled. + assert_eq!(ratio(&plan.children[0], 0.8), 1.1); + } + + for agg in [ + HashAggregate { + agg_exprs: vec![sum], + ..partial.clone() + }, + HashAggregate { + agg_exprs: vec![AggExpr::default()], + ..partial.clone() + }, + HashAggregate { + agg_exprs: vec![AggExpr { + expr_struct: Some(AggExprStruct::Count(Count { + children: vec![Expr::default(), Expr::default()], + })), + ..Default::default() + }], + ..partial.clone() + }, + HashAggregate { + mode: AggregateMode::PartialMerge as i32, + ..partial.clone() + }, + HashAggregate { + expr_modes: vec![AggregateMode::PartialMerge as i32], + ..partial.clone() + }, + HashAggregate { + mode: 99, + ..partial.clone() + }, + ] { + let plan = writer(agg); + assert_eq!(ratio(&plan, 0.8), 1.1); + // An eligible sibling or a Final parent must not hide the unsafe child. + let mut nested = writer(HashAggregate { + mode: AggregateMode::Final as i32, + ..partial.clone() + }); + nested.children[0].children = plan.children; + assert_eq!(ratio(&nested, 0.8), 1.1); + } + } fn entry_count(thread_id: u64) -> usize { get_thread_memory_pools() diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 0308779ce0..336cbf1a6e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -262,6 +262,9 @@ object CometMetricNode { def aggregateMetrics(sc: SparkContext): Map[String, SQLMetric] = { Map( + "skipped_aggregation_rows" -> SQLMetrics.createMetric( + sc, + "rows bypassing partial aggregation"), "spill_count" -> SQLMetrics.createMetric(sc, "number of spills"), "spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "total spilled bytes"), "spilled_rows" -> SQLMetrics.createMetric(sc, "number of spilled rows"), diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 3aa6b7220d..dfb362595a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1734,6 +1734,15 @@ trait CometBaseAggregate { if (aggregateExpressions.isEmpty) { val hashAggBuilder = OperatorOuterClass.HashAggregate.newBuilder() hashAggBuilder.addAllGroupingExprs(groupingExprs.map(_.get).asJava) + // Spark has no expression mode to serialize here. An empty aggregate with a required child + // distribution must fully deduplicate its keys (Final, or a pre-distinct PartialMerge), so + // use native Final to keep skip-partial disabled. + val mode = if (aggregate.requiredChildDistributionExpressions.isDefined) { + CometAggregateMode.Final + } else { + CometAggregateMode.Partial + } + hashAggBuilder.setModeValue(mode.getNumber) buildAggOp( builder, hashAggBuilder, 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 23c6b53509..4b3e7f4705 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{CometListenerBusUtils, SparkConf} import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.Cast -import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} +import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning import org.apache.spark.sql.comet.CometHashAggregateExec @@ -1623,6 +1623,108 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("skip partial aggregation preserves post-shuffle distinct") { + val writers = 8 + val rowsPerWriter = 300L + val overlap = 60L + val rows = writers * rowsPerWriter + + withTempDir { dir => + val path = new Path(dir.toURI.toString, "input") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0L, rows, 1L, writers) + .selectExpr(s"id - spark_partition_id() * $overlap AS k") + .write + .parquet(path.toUri.toString) + } + + withParquetTable(path.toUri.toString, "skip_partial_distinct") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> writers.toString, + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + 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") { + checkSparkAnswerAndOperator( + "SELECT count(*) FROM (SELECT DISTINCT k FROM skip_partial_distinct)") + checkSparkAnswerAndOperator("SELECT count(DISTINCT k) FROM skip_partial_distinct") + } + } + } + } + + test("skip partial aggregation admits only supported native shuffle plans") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "input").toUri.toString + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0L, 16384L, 1L, 8) + .selectExpr( + "id AS k", + "CASE WHEN id % 5 = 0 THEN NULL ELSE id % 17 END AS v", + "CASE WHEN id % 7 = 0 THEN NULL ELSE id % 11 END AS w") + .write + .parquet(path) + } + withParquetTable(path, "skip_partial_eligibility") { + 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_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold" -> "100") { + def aggregates(query: String): Seq[CometHashAggregateExec] = { + val (_, plan) = checkSparkAnswerAndOperator(query) + val result = stripAQEPlan(plan).collect { case aggregate: CometHashAggregateExec => + aggregate + } + assert(result.nonEmpty) + result + } + def skipped(aggregate: CometHashAggregateExec): Long = + aggregate.metrics.get("skipped_aggregation_rows").map(_.value).getOrElse(0L) + + val countQuery = "SELECT sum(n) FROM " + + "(SELECT k, count(*) n FROM skip_partial_eligibility GROUP BY k)" + // No ratio override: the eligible plan uses DataFusion's adaptive default. + assert(aggregates(countQuery).map(skipped).sum > 0L) + assert( + aggregates("SELECT sum(n) FROM " + + "(SELECT k % 2, count(*) n FROM skip_partial_eligibility GROUP BY k % 2)") + .map(skipped) + .sum == 0L) + + withSQLConf( + "spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "1.1") { + assert(aggregates(countQuery).map(skipped).sum == 0L) + } + withSQLConf( + "spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "0.8") { + for (expression <- Seq("sum(v)", "count(v, w)", "count(*) + sum(v)")) { + val result = aggregates( + "SELECT sum(n) FROM " + + s"(SELECT k, $expression n FROM skip_partial_eligibility GROUP BY k)") + assert(result.map(skipped).sum == 0L, s"Unexpected skipping for $expression") + } + val mixed = + aggregates("SELECT count(DISTINCT k), count(*) FROM skip_partial_eligibility") + .filter(_.modes.contains(PartialMerge)) + assert(mixed.nonEmpty, "Expected a PartialMerge stage in the DISTINCT plan") + assert(mixed.map(skipped).sum == 0L) + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + assert(aggregates(countQuery).map(skipped).sum == 0L) + } + } + } + } + } + } + test("first/last") { withSQLConf( SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
