Unsafe native partial aggregates survive child-triggered final fallback

[sunchao]

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.

[andygrove]

Another route into this, found while reviewing #5537.

You do not need spark.comet.shuffle.enabled=false to reach the bad shape. With the ordinary CometShuffleManager and spark.comet.shuffle.mode=native, columnarShuffleFailureReasons short-circuits on isCometJVMShuffleMode (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.

spark.range(0, 18, 1, 4).selectExpr("id % 3 AS k", "id % 5 AS v").createOrReplaceTempView("t")

with spark.comet.shuffle.mode=native and spark.sql.shuffle.partitions=4:

SELECT array(k) AS ak, percentile(v, 0.5) FROM t GROUP BY array(k)

The array hash key is unsupported for native partitioning, so the plan comes out in exactly the shape this issue describes:

ObjectHashAggregate(keys=[_groupingexpression#25], functions=[percentile(v#2L, 0.5, 1, 0, 0, false)])
+- Exchange hashpartitioning(_groupingexpression#25, 4), ENSURE_REQUIREMENTS
   +- CometHashAggregate [v#2L, _groupingexpression#25], [Partial], [partial_percentile(v#2L, 0.5, 1, 0, 0, false)]

and it dies with:

java.io.EOFException
  at java.io.DataInputStream.readFully(DataInputStream.java:203)
  at org.apache.spark.sql.catalyst.expressions.aggregate.TypedAggregateWithHashMapAsBuffer.deserialize(interfaces.scala:688)
  at org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate.merge(interfaces.scala:589)
  at org.apache.spark.sql.execution.aggregate.AggregationIterator$$anonfun$1.$anonfun$applyOrElse$3(AggregationIterator.scala:201)

collect_list produces the same plan shape and fails slightly differently:

java.lang.NullPointerException
  at org.apache.spark.sql.catalyst.expressions.aggregate.Collect.deserialize(collect.scala:86)
  at org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate.merge(interfaces.scala:589)

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 matches supportsMixedPartialFinal.

The reason I am adding this here rather than opening a new issue is that #5537 has already built the fix. Its preserveSparkAggregateBuffers in CometExecRule walks 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 behind isCometCelebornShuffleManagerEnabled, 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 CometAggregateSuite tests fail with "Expected 1 Comet aggregate operators, but found 0": group-by on variable length types, SUM decimal with DF and Decimal Avg with DF. All three assert a Comet Partial under a Spark Final with spark.comet.shuffle.enabled=false, and Decimal Avg with DF is 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]

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 COUNT partials can still feed Spark finals, while decimal SUM and AVG fall 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_set with shuffle enabled but native hash partitioning ineligible, under AQE on and off. Your exact array-key percentile query would provide useful additional coverage.