apache / apache/datafusion-comet
Unsafe native partial aggregates survive child-triggered final fallback
- 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
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