Skip to content

Commit c19ebbd

Browse files
andygroveandygrove
committed
fix: make adaptive partial aggregation opt-in with spark.comet.exec.aggregate.skipPartial.enabled
DataFusion decides once per task, after the first 100,000 input rows, whether to stop partial aggregation, and never checks again, so a task whose keys repeat after a mostly distinct start can shuffle many times more rows than in 1.0.0. Put the bypass behind a new config that defaults to false, which restores the 1.0.0 behavior. Closes #6466.
1 parent b38a60b commit c19ebbd

7 files changed

Lines changed: 104 additions & 38 deletions

File tree

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

Lines changed: 19 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -61,43 +61,41 @@ with the setting disabled to distinguish reduced hash-probe work from reader I/O
6161

6262
## Adaptive Partial Aggregation
6363

64-
For high-cardinality grouping, Comet can bypass partial hash aggregation when it is not
65-
reducing the number of rows enough. This currently applies only to fused native shuffle-writer
66-
plans whose partial aggregates are grouping-only or single-argument `COUNT`. Low-cardinality
67-
inputs continue to aggregate normally. The SQL metric `rows bypassing partial aggregation`
68-
shows whether skipping occurred.
64+
Set `spark.comet.exec.aggregate.skipPartial.enabled=true` to let Comet bypass partial hash
65+
aggregation for high-cardinality grouping when it is not reducing the number of rows enough. This
66+
experimental optimization is disabled by default. It currently applies only to fused native
67+
shuffle-writer plans whose partial aggregates are grouping-only or single-argument `COUNT`.
68+
Low-cardinality inputs continue to aggregate normally. The SQL metric
69+
`rows bypassing partial aggregation` shows whether skipping occurred.
70+
71+
DataFusion makes the decision separately in each task. It starts checking after the first 100,000
72+
input rows, and as soon as the number of groups divided by the number of input rows exceeds `0.8`,
73+
it stops aggregating and sends the rest of the task's rows to the shuffle as they are. It does not
74+
check again, so a task whose keys repeat after a mostly distinct start, such as several snapshot
75+
files of the same keys packed into one split, can shuffle many times more rows than it would with
76+
skipping disabled. Compare the shuffle write metrics with the setting enabled and disabled before
77+
enabling it for a workload.
6978

7079
Eligibility is conservative for the whole fused native plan: any unsupported partial accumulator,
7180
Spark `PartialMerge`, or mixed-mode aggregate disables skipping in that plan. Multi-argument
7281
`COUNT` and other accumulators are not admitted. Distribution-required grouping-only stages
7382
still fully deduplicate, and non-native-shuffle plans retain ordinary aggregation.
74-
The DataFusion testing configuration override does not bypass these safety checks.
7583

76-
DataFusion 55 defaults to probing after 100,000 input rows per partial aggregation
77-
partition and skipping when the number of groups divided by input rows exceeds `0.8`.
78-
To experiment with these thresholds, enable `spark.comet.exec.respectDataFusionConfigs`,
84+
To experiment with the thresholds, also enable `spark.comet.exec.respectDataFusionConfigs`,
7985
a development and testing option that defaults to `false`. For example, the following
8086
SQL settings pass through the default threshold values, which you can adjust:
8187

8288
```sql
89+
SET spark.comet.exec.aggregate.skipPartial.enabled=true;
8390
SET spark.comet.exec.respectDataFusionConfigs=true;
8491
SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold=100000;
8592
SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold=0.8;
8693
```
8794

8895
A lower row threshold allows an earlier decision; a lower ratio threshold makes
89-
skipping more likely. Skipping can increase the number of partial states emitted
90-
and the amount of shuffle data, so measure the effect on your workload.
91-
92-
To disable skipping, keep `spark.comet.exec.respectDataFusionConfigs=true` and set
93-
the ratio threshold above the maximum possible groups/input-rows ratio:
94-
95-
```sql
96-
SET spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold=1.1;
97-
```
98-
99-
These settings only tune eligible plans. Unsupported accumulators and modes remain
100-
disabled even when configuration overrides are enabled.
96+
skipping more likely. These settings only tune eligible plans. They cannot enable skipping
97+
while `spark.comet.exec.aggregate.skipPartial.enabled` is `false`, or for unsupported
98+
accumulators and modes.
10199

102100
## Local TopK Fusion
103101

‎native/core/src/execution/jni_api.rs‎

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,8 @@ use crate::execution::tracing::{
119119

120120
use crate::execution::memory_pools::logging_pool::LoggingMemoryPool;
121121
use crate::execution::spark_config::{
122-
SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED,
122+
SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY,
123+
COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXPLAIN_NATIVE_ENABLED,
123124
COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED,
124125
COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES,
125126
};
@@ -806,10 +807,15 @@ pub extern "system" fn Java_org_apache_comet_Native_setShufflePartitionPusher(
806807
})
807808
}
808809

809-
/// Only admit the validated native-shuffle path. A session belongs to one fused Spark plan,
810-
/// so an unsafe partial aggregate disables skipping for the whole plan, including its children.
811-
/// This deliberately gives up some opportunities rather than changing execution contexts per op.
812-
fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operator) {
810+
/// Skipping is opt-in (`spark.comet.exec.aggregate.skipPartial.enabled`): once DataFusion's probe
811+
/// decides to skip, it never aggregates again, so a task whose keys repeat after a mostly distinct
812+
/// start shuffles every later row (#6466).
813+
///
814+
/// When enabled, only admit the validated native-shuffle path. A session belongs to one fused
815+
/// Spark plan, so an unsafe partial aggregate disables skipping for the whole plan, including its
816+
/// children. This deliberately gives up some opportunities rather than changing execution
817+
/// contexts per op.
818+
fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operator, enabled: bool) {
813819
fn supported(plan: &Operator) -> bool {
814820
let supported_aggregate = match &plan.op_struct {
815821
Some(OpStruct::HashAgg(agg)) => match AggregateMode::try_from(agg.mode) {
@@ -832,9 +838,12 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato
832838
supported_aggregate && plan.children.iter().all(supported)
833839
}
834840

835-
if !matches!(&plan.op_struct, Some(OpStruct::ShuffleWriter(_))) || !supported(plan) {
836-
// Enforce safety after config pass-through: a testing override cannot make unsupported
837-
// accumulators convertible. DF 55 removed supports_convert_to_state().
841+
let eligible =
842+
enabled && matches!(&plan.op_struct, Some(OpStruct::ShuffleWriter(_))) && supported(plan);
843+
if !eligible {
844+
// Enforce this after config pass-through: a testing override can neither turn skipping
845+
// on nor make unsupported accumulators convertible. DF 55 removed
846+
// supports_convert_to_state().
838847
config
839848
.options_mut()
840849
.execution
@@ -893,7 +902,11 @@ fn prepare_datafusion_session_context(
893902
}
894903
}
895904

896-
configure_skip_partial_aggregation(&mut session_config, spark_plan);
905+
configure_skip_partial_aggregation(
906+
&mut session_config,
907+
spark_plan,
908+
spark_config.get_bool(COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED),
909+
);
897910

898911
let runtime = rt_config.build()?;
899912

@@ -2041,7 +2054,7 @@ mod tests {
20412054
}],
20422055
..Default::default()
20432056
};
2044-
let ratio = |plan: &Operator, requested: f64| {
2057+
let ratio = |plan: &Operator, requested: f64, enabled: bool| {
20452058
let mut config = SessionConfig::new();
20462059
config
20472060
.options_mut()
@@ -2051,7 +2064,7 @@ mod tests {
20512064
.options_mut()
20522065
.execution
20532066
.skip_partial_aggregation_probe_ratio_threshold = requested;
2054-
configure_skip_partial_aggregation(&mut config, plan);
2067+
configure_skip_partial_aggregation(&mut config, plan, enabled);
20552068
assert_eq!(
20562069
config
20572070
.options()
@@ -2077,11 +2090,13 @@ mod tests {
20772090
},
20782091
] {
20792092
let plan = writer(agg);
2080-
assert_eq!(ratio(&plan, 0.8), 0.8);
2081-
assert_eq!(ratio(&plan, 0.5), 0.5);
2082-
assert_eq!(ratio(&plan, 1.1), 1.1);
2093+
assert_eq!(ratio(&plan, 0.8, true), 0.8);
2094+
assert_eq!(ratio(&plan, 0.5, true), 0.5);
2095+
assert_eq!(ratio(&plan, 1.1, true), 1.1);
2096+
// Skipping is opt-in, and a DataFusion override cannot turn it on.
2097+
assert_eq!(ratio(&plan, 0.8, false), 1.1);
20832098
// Non-native shuffle / standalone native blocks stay disabled.
2084-
assert_eq!(ratio(&plan.children[0], 0.8), 1.1);
2099+
assert_eq!(ratio(&plan.children[0], 0.8, true), 1.1);
20852100
}
20862101

20872102
for agg in [
@@ -2116,14 +2131,14 @@ mod tests {
21162131
},
21172132
] {
21182133
let plan = writer(agg);
2119-
assert_eq!(ratio(&plan, 0.8), 1.1);
2134+
assert_eq!(ratio(&plan, 0.8, true), 1.1);
21202135
// An eligible sibling or a Final parent must not hide the unsafe child.
21212136
let mut nested = writer(HashAggregate {
21222137
mode: AggregateMode::Final as i32,
21232138
..partial.clone()
21242139
});
21252140
nested.children[0].children = plan.children;
2126-
assert_eq!(ratio(&nested, 0.8), 1.1);
2141+
assert_eq!(ratio(&nested, 0.8, true), 1.1);
21272142
}
21282143
}
21292144

‎native/core/src/execution/spark_config.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirec
2424
pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory";
2525
pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str =
2626
"spark.comet.parquet.rowFilterPushdown.enabled";
27+
pub(crate) const COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: &str =
28+
"spark.comet.exec.aggregate.skipPartial.enabled";
2729
pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores";
2830

2931
/// Comet configs read through this trait must be resolved by the JVM first:

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -437,6 +437,20 @@ object CometConf extends ShimCometConf {
437437
.booleanConf
438438
.createWithDefault(false)
439439

440+
val COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: ConfigEntry[Boolean] =
441+
conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregate.skipPartial.enabled")
442+
.category(CATEGORY_EXEC)
443+
.doc(
444+
"Experimental opt-in: let a native partial aggregate stop aggregating once its input " +
445+
"looks mostly distinct, and send the rest of the task's rows to the shuffle " +
446+
"unaggregated. Only applies to partial aggregates that feed Comet native shuffle " +
447+
"and whose aggregate functions, if any, are all single-argument COUNT. The check " +
448+
"starts after the first 100,000 input rows of a task, and once aggregation stops " +
449+
"it does not resume, so a task whose keys repeat after a mostly distinct start can " +
450+
s"shuffle many times more rows than it would with this disabled. $TUNING_GUIDE.")
451+
.booleanConf
452+
.createWithDefault(false)
453+
440454
val COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED: ConfigEntry[Boolean] =
441455
conf(s"$COMET_EXEC_CONFIG_PREFIX.join.dynamicFilter.enabled")
442456
.category(CATEGORY_EXEC)

‎spark/src/main/scala/org/apache/comet/CometExecIterator.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -729,6 +729,7 @@ object CometExecIterator extends Logging {
729729
Seq[ConfigEntry[_]](
730730
CometConf.COMET_DEBUG_ENABLED,
731731
CometConf.COMET_DEBUG_MEMORY_ENABLED,
732+
CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED,
732733
CometConf.COMET_EXPLAIN_NATIVE_ENABLED,
733734
CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE,
734735
CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED,

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2407,6 +2407,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
24072407
SQLConf.SHUFFLE_PARTITIONS.key -> writers.toString,
24082408
CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
24092409
CometConf.COMET_SHUFFLE_MODE.key -> "native",
2410+
CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED.key -> "true",
24102411
CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
24112412
"spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold" -> "100",
24122413
"spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "0.8") {
@@ -2418,6 +2419,33 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
24182419
}
24192420
}
24202421

2422+
test("skip partial aggregation is disabled by default") {
2423+
// https://github.com/apache/datafusion-comet/issues/6466. DataFusion stops aggregating for
2424+
// the rest of a task once its first 100,000 rows look mostly distinct, so a task whose keys
2425+
// repeat after a distinct start shuffled every later row.
2426+
withTempPath { dir =>
2427+
val path = dir.getCanonicalPath
2428+
withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
2429+
// One file, read by one task. Keys 0 to 199999 cycle 10 times, so the first 100,000
2430+
// rows are all distinct, but every key repeats 10 times within the task.
2431+
spark.range(0, 2000000, 1, 1).selectExpr("id % 200000 AS k").write.parquet(path)
2432+
}
2433+
withSQLConf(
2434+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
2435+
SQLConf.FILES_MAX_PARTITION_BYTES.key -> (1L << 30).toString,
2436+
SQLConf.SHUFFLE_PARTITIONS.key -> "4",
2437+
CometConf.COMET_SHUFFLE_MODE.key -> "native") {
2438+
val df = spark.read.parquet(path).groupBy("k").count()
2439+
df.collect()
2440+
val written = collect(df.queryExecution.executedPlan) {
2441+
case e: CometShuffleExchangeExec =>
2442+
e.metrics("shuffleRecordsWritten").value
2443+
}.sum
2444+
assert(written == 200000, s"the partial aggregate shuffled $written rows")
2445+
}
2446+
}
2447+
}
2448+
24212449
test("skip partial aggregation admits only supported native shuffle plans") {
24222450
withTempDir { dir =>
24232451
val path = new Path(dir.toURI.toString, "input").toUri.toString
@@ -2438,6 +2466,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
24382466
CometConf.COMET_BATCH_SIZE.key -> "128",
24392467
CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
24402468
CometConf.COMET_SHUFFLE_MODE.key -> "native",
2469+
CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED.key -> "true",
24412470
CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
24422471
"spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold" -> "100") {
24432472
def aggregates(query: String): Seq[CometHashAggregateExec] = {
@@ -2455,6 +2484,12 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
24552484
"(SELECT k, count(*) n FROM skip_partial_eligibility GROUP BY k)"
24562485
// No ratio override: the eligible plan uses DataFusion's adaptive default.
24572486
assert(aggregates(countQuery).map(skipped).sum > 0L)
2487+
// The Comet config gates skipping, whatever the DataFusion thresholds say.
2488+
withSQLConf(
2489+
CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED.key -> "false",
2490+
"spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "0.8") {
2491+
assert(aggregates(countQuery).map(skipped).sum == 0L)
2492+
}
24582493
assert(
24592494
aggregates("SELECT sum(n) FROM " +
24602495
"(SELECT k % 2, count(*) n FROM skip_partial_eligibility GROUP BY k % 2)")

‎spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ class CometExecSuite extends CometTestBase {
9191
val flags = Seq(
9292
CometConf.COMET_DEBUG_ENABLED,
9393
CometConf.COMET_DEBUG_MEMORY_ENABLED,
94+
CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED,
9495
CometConf.COMET_EXPLAIN_NATIVE_ENABLED,
9596
CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED,
9697
CometConf.COMET_TRACING_ENABLED)

0 commit comments

Comments
 (0)