feat: Enable adaptive partial aggregation for eligible native shuffle plans

[peterxcli]

Which issue does this PR close?

Closes #5717.

Rationale for this change

High-cardinality partial aggregation can spend substantial CPU maintaining a hash table while emitting almost as many rows as it consumes. Re-enable DataFusion's adaptive bypass only for supported native-shuffle plans, rather than lowering the threshold globally.

DataFusion 55 requires convert_to_state, but several Comet accumulators still return an unsupported-operation error. Spark's grouping-only DISTINCT stages also need an explicit mode to preserve full deduplication.

What changes are included in this PR?

• Serialize distribution-required grouping-only aggregates as native Final to preserve DISTINCT semantics.
• Admit adaptive skipping only for native shuffle-writer plans whose partial aggregates are grouping-only or single-argument COUNT. An unsupported partial accumulator, PartialMerge, or mixed mode disables skipping throughout the fused plan. Non-native-shuffle roots remain disabled.
• Apply the eligibility guard after DataFusion configuration pass-through, so testing overrides cannot bypass safety. Eligible plans retain the adaptive default and explicit disablement remains available.
• Expose skipped_aggregation_rows as a SQL metric and document the conservative scope.

No new configuration, dependency, or execution wrapper is introduced.

How are these changes tested?

• Full CometAggregateSuite: 93 passed, zero failures, 2 ignored. The two new focused regressions also passed independently.
• Native skip_partial_eligibility_is_fail_closed: 1 passed.
• Native release build, Rust formatting check, Scala Spotless formatting, and diff whitespace check passed.
• Separate 2.4-million-row correctness fixture: 22 exact comparisons against Spark, zero errors, covering COUNT, nullable/filtered COUNT, DISTINCT, empty/null inputs, SUM, and mixed PartialMerge.
• The checked-in Spark regression asserts actual skipping for eligible high-cardinality COUNT and no skipping for low cardinality, explicit disablement, unsupported accumulators/modes, and JVM shuffle.

Release-mode Spark A/B on the DataFusion 55 base (75fdddc9285ec61c0cd326977c61dd41fca39a8b): 16,777,216 Parquet rows, 8 writers, 32 reducers, local[8], Spark 4.1.3, native direct-read shuffle, AQE/exchange reuse disabled, one warmup per mode, alternating pairs. Build/test processes had finished before timing.

Workload	Pairs	Disabled median	Guarded enabled median	Speedup (wall-time reduction)
High-cardinality COUNT	7	664 ms	479 ms	27.9%
Low-cardinality control	7	199 ms	185 ms	7.0% observed; no skipping, no optimization gain claimed
High-cardinality COUNT, 128 MiB native memory	3	945 ms	858 ms	9.2% observed; limited sample

Wall-time reduction is (disabled median - enabled median) / disabled median × 100%.

High-cardinality COUNT improved 27.9%, winning all seven pairs; enabled runs bypassed 15,925,248 rows each, versus zero when disabled. The low-cardinality control bypassed zero rows in both modes, so no optimization gain is claimed there.

All eight constrained-memory executions (including warmups) returned correct results with 64 aggregate spills and approximately 273 MB spilled each. Enabled executions still bypassed 15,925,248 rows, exercising skipping and spilling together. These are workstation measurements, not a universal speedup or coverage of every memory pool, AQE, or distributed deployment. Additional benchmark context is recorded in #5717.

Focused reproduction from the repository root:

./mvnw test -Dtest=none '-Dsuites=org.apache.comet.exec.CometAggregateSuite skip partial aggregation' -Dscalastyle.skip=true -Djni.dir="$PWD/native/target/release"

Native regression from native/ (on macOS, expose the installed JVM's lib/server via DYLD_LIBRARY_PATH):

RUSTFLAGS='-Ctarget-cpu=native' cargo test --release --locked -p datafusion-comet skip_partial_eligibility --lib

[sunchao]

Correctness

Reviewed 8d915424c1cd91e2f53c23a8e183f4fd028c3adb against 75fdddc9285ec61c0cd326977c61dd41fca39a8b.

The previous session configuration disabled adaptive partial aggregation globally. This change enables it only for native shuffle-writer roots when every partial aggregate in the fused plan is grouping-only or single-argument COUNT. The recursive guard rejects unsupported accumulators, PartialMerge, mixed modes and unknown modes, and runs after configuration pass-through so an override cannot enable unsupported state conversion.

The DISTINCT mode change matches the maintained Spark 3.5 and 4.0 planning paths: local grouping-only stages have no required child distribution, while the final and pre-distinct merge stages require distribution and must fully deduplicate. Serializing those latter stages as native Final preserves that obligation, including when Spark can satisfy distribution without inserting another exchange. COUNT's locked DataFusion 55 implementation emits non-null Int64 0/1 states using both argument validity and FILTER validity. Downstream count-state merging matches Spark's additive buffer semantics. Accumulated groups drain before later input bypasses aggregation. Ordered and global aggregation paths do not gain this bypass. I found no additional correctness issue in those paths. Maintained Spark 3.4 and 4.1 sources were unavailable, so this source comparison does not qualify those versions.

The redundant string interpolation in the new regression test is fixed. The refreshed current-head checks at 2026-09-05T22:08:12.363Z show 15 successes, seven skips and six checks still running, with no failures. The Scala syntactic lint job passed. The completed syntactic lint log confirms checkout of synthetic merge b27feaae6ecffe8c3f3e373139f1f67ff0f1da97, which has parents c348f775d5b20d91d1989756b72f62f366d67d32 and the reviewed head. Its four additional cast files make it a different tree, while all five files changed here, Cargo.lock and the lint configuration match the reviewed head. The native eligibility test passed on the preceding CI merge with identical native code and dependencies. That is historical supporting evidence, not a completed current-head native run. The full JVM suite and 22 Spark comparisons remain author-reported. I did not run local product tests and found no remaining actionable P1/P2 issue.

Performance

The new plan walk occurs once during native session creation. It does not add a per-row Comet wrapper. Eligible unordered partials retain DataFusion's default probe at 100,000 input rows and a groups/rows ratio greater than 0.8. A threshold of 1.1 creates no probe, preserving explicit disablement and the unsupported path. On memory pressure the unordered partial stream emits states early. This should be distinguished from final-stage disk spills.

The author's guarded workload reports seven alternating high-cardinality pairs improving median wall time from 664 to 479 ms, with 15,925,248 bypassed rows. The low-cardinality control bypassed zero rows, and the 128 MiB case supplied a smaller spilling sample. These are useful bounded workstation measurements, not independently reproduced or distributed/AQE qualification.

Design

A single conservative decision for the fused plan is easy to follow and avoids assigning different session settings to individual operators. Inspecting children even below Final prevents an unsupported nested partial from being hidden by its parent. The grouping-only mode correction addresses the semantic reason global enablement was unsafe, while the post-override guard makes accumulator support an invariant. The current boundary is intentionally narrow. Broadening it should require working state conversion and runtime evidence for each additional accumulator and execution context.

Abstraction & complexity

The helper, explicit mode selection, metric and focused admission tests fit the existing planner and metric interfaces. There is no new execution layer or configuration surface to maintain. The tests distinguish actual bypass from merely setting a threshold, and exercise unsupported modes and JVM shuffle. I found no unnecessary abstraction or actionable simplification.

[sunchao]

Correctness

Follow-up on a8f5969d3caf97338828212baf733cee2bddddaf against 17f54da8ca5cb0ad5dbe8357b6e037fef8a0db2c, after review 5123205955. All five authored patches match the previously reviewed contribution. The 148-file increment matches main's advancement, and the eligibility guard, session configuration, aggregate planner, Scala aggregate serialization and aggregate metrics blocks are unchanged. I found no new or remaining actionable P1/P2 issue in this update and the relevant inherited interactions.

The maintained Spark 3.5/4.0 planning contracts still support the grouping-only Final mapping for stages requiring distribution. The locked DataFusion 55 COUNT state conversion still combines argument/FILTER validity into non-null Int64 states, and the partial stream drains accumulated states before forwarding later input. Unsupported accumulators and PartialMerge remain excluded after configuration overrides. The native writer still supplies the ShuffleWriter root. Scan plan-data injection preserves it. Inherited AQE input traversal and nested shuffle hashing do not alter these eligibility or grouping-mode decisions. Maintained Spark 3.4/4.1 branches remain unavailable, so the source comparison does not qualify those versions.

Both new aggregation regressions passed in the current Spark 3.4, 3.5, 4.0 and 4.1 execution jobs. These check post-shuffle DISTINCT correctness and actual bypass/admission, including the low-cardinality, explicit-disable, unsupported aggregate, mixed-mode and JVM-shuffle controls. The native fail-closed eligibility test also passed. That run reports 1,252 passed and five skipped tests. All inspected jobs ran merge bd0cc3640207384b11ba11994f606cd820675efd, whose parents are the assigned BASE/HEAD and whose tracked tree equals HEAD. Producer, artifact metadata and the four Spark consumers agree on native artifact 10087734700 and its SHA256 digest.

At 2026-09-09T05:54:26.032357+00:00, checks were 64 successful, nine skipped and one failed, with none running. The failed Spark SQL 4.1 job reports SQLAppStatusListenerMemoryLeakSuite's noLiveData() assertion. Immediately preceding activity from that test logs that the Comet extension is disabled because of the shuffle-manager configuration. This does not demonstrate a failure in adaptive native aggregation, but the listener failure remains unresolved and CI is still red. No local product tests or builds were run. GitHub reports the existing approval attached to this head. This COMMENT preserves that approval.

Performance

The merge leaves the one-time plan walk and locked adaptive probe/state-conversion paths unchanged. I found no added per-row Comet allocation or repeated work in this update. The author's guarded benchmark remains the available measurement: seven alternating high-cardinality COUNT pairs, median 664 to 479 ms, with 15,925,248 bypassed rows. The low-cardinality control bypassed zero rows, and the constrained-memory run supplied three measured pairs. Those are historical workstation results. Main also advanced runtime, shuffle and other components, so feature-block equivalence does not establish a fresh speedup or an identical benchmark environment. No new timing, distributed/AQE or deployed-build qualification is claimed.

Design

The update preserves the conservative whole-plan decision and the independent grouping-only mode correction. Final nodes still allow recursive inspection of nested partials, and configuration overrides cannot enlarge accumulator eligibility. The inherited writer and scan-data changes preserve the native plan boundary used by the guard. I found no changed design issue requiring an additional layer or a broader allowlist.

Abstraction & complexity

The helper, metric and admission tests are unchanged and remain within existing planner interfaces. The merge introduces no additional abstraction for this feature. I found no actionable simplification in the reviewed update.

[peterxcli]

reports SQLAppStatusListenerMemoryLeakSuite's noLiveData() assertio ... This COMMENT preserves that approval.

@sunchao could you please help me to retry the ci? thanks! the failed job must share the same reason as #5023

[sunchao]

Correctness

Discussion follow-up on a8f5969d3caf97338828212baf733cee2bddddaf against 17f54da8ca5cb0ad5dbe8357b6e037fef8a0db2c, after review 5150301288. Head and base are unchanged. I re-read the full five-file contribution and verified its source against the prior packet. The incremental diff is empty. I found no new or remaining verified P1/P2. The existing approval stands.

I checked the controls behind the documentation comment. Overrides require spark.comet.exec.respectDataFusionConfigs=true, a development/testing option that defaults to false. The keys are spark.comet.datafusion.execution.skip_partial_aggregation_probe_rows_threshold and spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold. Locked DataFusion 55 defaults are 100,000 input rows per partial-aggregation partition and a groups/input-rows ratio greater than 0.8. Setting the ratio to 1.1 disables bypass. Both aggregate implementations omit the probe for ratios at least 1.0. Eligible plans preserve explicit overrides, while unsupported plans force 1.1 after pass-through. The tuning page currently documents eligibility and the metric but omits these controls. That omission is already covered by the existing thread.

The refreshed maintained Spark 3.5/4.0 references still support the unchanged grouping-only Final mapping for distribution-required stages and COUNT's zero-initialized, additive Long state. DataFusion's unchanged COUNT conversion includes input/FILTER validity and emits non-null 0/1 states. The update adds no type, null, ANSI/overflow, error or fallback change. Maintained Spark 3.4/4.1 source gaps remain.

Validation and CI

At 2026-09-09T18:56:35Z, checks remain 64 successful, nine skipped and one failed. The same completed jobs and run attempt retain the passing native eligibility test and both aggregation regressions on Spark 3.4/3.5/4.0/4.1. I reconciled their job identities and saved log provenance. Merge bd0cc3640207384b11ba11994f606cd820675efd has the assigned base/head parents and the same tracked tree as HEAD. This reuses the previous CI execution evidence, with no new local build, test or benchmark.

The failed listener job still reports SQLAppStatusListenerMemoryLeakSuite's noLiveData() assertion, with Comet disabled in the preceding test activity. The author's retry request links #5023, which reports a macOS scan SIGSEGV and has a comment attributing it to libhdfs JNI lifetime handling. That report does not establish a common cause for this Linux listener assertion. CI remains red and the listener failure remains unresolved.

Performance

There is no execution change since the previous review. The eligibility walk remains once per native context. Lower probe thresholds can trade hash-table work for more emitted partial states, so the existing author measurements do not establish a universal tuning value. No fresh timing, distributed/AQE or deployed-build benefit is claimed.

Design

The documented controls should reflect the existing post-override safety boundary: tuning can change when an eligible plan skips, but cannot admit unsupported accumulators or modes. I verified that boundary against the configuration serialization and native session setup. No changed design issue requires another control or execution path.

Abstraction & complexity

The helper, metric, mode mapping and admission tests are unchanged. The existing DataFusion pass-through already supplies the tuning mechanism. I found no new abstraction issue or actionable simplification in this follow-up.

[sunchao]

Reviewed 7ec7cfcd against 17f54da8. The documentation now addresses the tuning and disablement request: it explains the configuration gate, the 100,000-row and 0.8 defaults, and disabling bypass with 1.1. I checked the examples against the configuration path and confirmed that eligibility checks still run after overrides.

This update changes only the documentation; execution code, regression tests, and dependencies are unchanged from a8f5969d. No new or remaining verified P1/P2 findings.

Current check metadata reports 56 successes and 10 skips. Current job logs were not reverified; the fetched merge preview includes newer main changes and DataFusion 55.1, while the reviewed head locks 55.0. Earlier verified native and Spark 3.4–4.1 regression runs support the unchanged implementation. No local tests or fresh performance benchmark were run.

[andygrove]

Thanks for the careful approach here. Mapping the distribution-required grouping-only stage to native Final addresses the actual cause of the old q97 corruption, and the guard in configure_skip_partial_aggregation is fail-closed. I traced the plan shapes it admits (grouping-only and single-argument Count partials under a native ShuffleWriter) and could not construct one that produces wrong results, including when the Spark Final aggregate falls back while the shuffle stays native. CI is green across all Spark versions.

A few comments inline. The test hardening and the config question are worth addressing before merge. The third is about a pre-existing gap that this PR's own comment documents, so I would just like a tracking issue for it.

[andygrove]

This test compares results but does not confirm that the stage 1 partial actually bypassed aggregation. If the probe never fires, the test passes without exercising the DISTINCT dedup path it is meant to protect. It might be worth collecting the CometHashAggregateExec nodes and asserting skipped_aggregation_rows is greater than zero on the Partial stage, the same way the eligibility test does.

[andygrove]

The only documented way to turn this off is spark.comet.exec.respectDataFusionConfigs, which is a TESTING category option, plus a raw DataFusion key. Since this feature was disabled once before for a TPC-DS wrong-result (#788, apache/datafusion#11850) and is now on by default, would you consider a first-class boolean such as spark.comet.exec.aggregate.skipPartial.enabled (default true) that the guard in configure_skip_partial_aggregation checks alongside the plan walk? That gives operators a supported kill switch without pointing them at a testing flag.

[andygrove]

This comment says a distribution-required stage must fully deduplicate its keys, including the pre-distinct PartialMerge case. For the non-empty PartialMerge stage, Comet still plans a DataFusion Partial with MergeAsPartial accumulators, and DataFusion's Partial mode uses EmitEarly under memory pressure, which can re-emit a key that was already emitted. Does that mean stage 3's distinct count(k) can overcount when stage 2 spills? I think this is pre-existing and separate from this PR, and the guard correctly refuses to skip in that shape, but could you open a tracking issue and link it from this comment so the obligation is recorded?