apache / apache/datafusion-comet

Unsafe native partial aggregates survive child-triggered final fallback

Open
#5,419 2 comments 0 reactions 0 assignees View on GitHub
area:aggregation bug priority:critical
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 6h
Merged PRs (30d)
190

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]`):

```python
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.

Contributor guide

Open the contributing guide

Research direction

Start by running the provided decimal AVG reproduction with AQE both disabled and enabled, then trace the existing aggregate-expression support prepass and the child-eligibility fallback. Check how PartialMerge/grouping and exchange stages are reapplied, including distinct aggregates. Done means the query returns 200.000000, unsafe partials fall back with their finals, and supported native chains and unrelated scans or filters remain native.

Written by the indexing model from the issue text.

Assessment

Tech stack
scala
Domain
data, distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.