feat: Enable adaptive partial aggregation for eligible native shuffle plans

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

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()

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"),

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,