Repository navigation
Unsafe native partial aggregates survive child-triggered final fallback #5419
Description
Describe the bug
The unsafe-partial aggregate prepass checks aggregate expression support before child conversion is known. A final aggregate may pass those checks and then fall back because its shuffle child is not native. The corresponding native partial aggregate can survive even when the aggregate is explicitly ineligible for mixed Spark/Comet execution.
For decimal AVG, an empty native partial can then poison the Spark final aggregate's sum and produce an incorrect NULL result. The problem occurs with AQE both disabled and enabled.
Steps to reproduce
With Comet installed and enabled in a fresh local Spark session (for example, local[4]):
import tempfile
for key, value in {
"spark.sql.adaptive.enabled": "false", # Also reproduces with true.
"spark.sql.files.maxPartitionBytes": "1048576",
"spark.sql.parquet.filterPushdown": "false",
"spark.comet.scan.enabled": "false",
"spark.comet.convert.parquet.enabled": "true",
"spark.comet.shuffle.enabled": "false",
}.items():
spark.conf.set(key, value)
with tempfile.TemporaryDirectory() as tmp:
path = f"{tmp}/data"
(spark.range(8, numPartitions=4)
.selectExpr("id", "cast(200 as decimal(20,2)) amount")
.write.parquet(path))
spark.read.parquet(path).createOrReplaceTempView("decimal_avg_probe")
spark.sql("SELECT avg(amount) FROM decimal_avg_probe WHERE id = 1").show()Only one of the four scan partitions contains a matching row. The final aggregate falls back across the Spark shuffle, but the partial aggregate remains native. The query returns NULL instead of 200.000000.
Expected behavior
When the final aggregate actually falls back and its buffers are not declared safe for mixed execution, the feeding partial aggregate must also use Spark. This needs to happen before AQE materializes the partial output, and it must survive stage-only rule reapplication.
The result should be 200.000000. Supported mixed aggregates such as MIN/MAX and fully native aggregate chains should retain native execution; unrelated native filters and scans should not be reverted.
Additional context
Reproduced against Apache Comet main at 2699f59b71788e17a2714910e166a3f83deed937 using Spark 4.0.4. The existing expression-support prepass is still useful, but does not cover a final fallback caused by actual child eligibility. This also requires handling the intermediate PartialMerge/grouping and exchange chain for distinct aggregates.
Activity
andygrove commentedon Aug 29, 2026
Another route into this, found while reviewing #5537.
You do not need
spark.comet.shuffle.enabled=falseto reach the bad shape. With the ordinaryCometShuffleManagerandspark.comet.shuffle.mode=native,columnarShuffleFailureReasonsshort-circuits onisCometJVMShuffleMode(CometShuffleExchangeExec.scala:585), so an exchange whose native path fails has no columnar fallback either and reverts to a plain Spark exchange with the native partial still underneath it. Any unsupported native partitioning will do it, and unlike the decimal AVG case this one crashes instead of returning a wrong answer.with
spark.comet.shuffle.mode=nativeandspark.sql.shuffle.partitions=4:The array hash key is unsupported for native partitioning, so the plan comes out in exactly the shape this issue describes:
and it dies with:
collect_listproduces the same plan shape and fails slightly differently:Both reproduce with AQE on and off, on Spark 4.1.3 / Scala 2.13 / JDK 17.
sum(v)on the same query is fine, which matchessupportsMixedPartialFinal.The reason I am adding this here rather than opening a new issue is that #5537 has already built the fix. Its
preserveSparkAggregateBuffersinCometExecRulewalks down from a fallen-back exchange, reverts an incompatible native Partial back to Spark, and restores any native ancestors in between, which is what this issue asks for. It is gated behindisCometCelebornShuffleManagerEnabled, so the ordinary shuffle manager does not get it.I tried simply deleting that gate on the #5537 branch. Both queries above then pass, but three existing
CometAggregateSuitetests fail with "Expected 1 Comet aggregate operators, but found 0":group-by on variable length types,SUM decimal with DFandDecimal Avg with DF. All three assert a Comet Partial under a Spark Final withspark.comet.shuffle.enabled=false, andDecimal Avg with DFis covering the same decimal AVG shape this issue reports as a bug. So ungating is not a one-liner: someone needs to decide which of those assertions are pinning intended behaviour and which are pinning this bug, and #5509 probably belongs in that conversation too.sunchao commentedon Aug 29, 2026
Thanks for the additional reproducer. This is the same shuffle-fallback path that #5421 addresses. Its check runs after conversion and restores incompatible native partials when the final aggregate must run in Spark, including with the ordinary
CometShuffleManager.The current #5421 also handles the three test-expectation questions. It separates compatibility by direction: native
COUNTpartials can still feed Spark finals, while decimalSUMandAVGfall back across that boundary, with the affected expectations updated. So those failures do not all call for changing the expected native aggregate count to zero.The existing regressions cover
collect_list/collect_setwith shuffle enabled but native hash partitioning ineligible, under AQE on and off. Your exact array-key percentile query would provide useful additional coverage.fix: [branch-1.0] stop transition reversion from splitting unsafe agg…
fix: [branch-1.0] revert unsafe partial aggregates after final fallba…