andygrove opened a new issue, #6857:
URL: https://github.com/apache/datafusion-comet/issues/6857
### Describe the bug
`count(DISTINCT x)` next to another aggregate comes back too large when the
native aggregate that deduplicates the distinct keys runs short of memory.
Spark plans `SELECT g, sum(v), count(DISTINCT d) FROM t GROUP BY g` as four
aggregates. The second one is a `PartialMerge` keyed by `(g, d)`. It has to
leave exactly one row per `(g, d)` pair, because the stage above it counts
those rows.
Comet runs a `PartialMerge` as a DataFusion `AggregateMode::Partial`
aggregate (`native/core/src/execution/planner.rs`, the
`ProtoAggregateMode::PartialMerge => DFAggregateMode::Partial` arm, with
`MergeAsPartial` wrappers). In `Partial` mode, DataFusion's
`GroupedHashAggregateStream` handles memory pressure with
`OutOfMemoryMode::EmitEarly`: it emits the groups it holds and starts again.
That's fine for a real partial aggregate, but here the same `(g, d)` key leaves
the stage more than once, and the distinct count includes every copy.
Nothing fails. `sum(v)` is still right, because partial sums add up, but
`count(DISTINCT d)` is wrong.
### Steps to reproduce
On `main` at dc880d6d26 with Spark 4.1.3, in a `CometTestBase` suite (2g
off-heap):
```scala
spark
.range(0, 400000, 1, 1)
.selectExpr("id % 100 AS g", "id % 200000 AS d", "id AS v")
.write
.parquet(path)
```
Then, with these settings:
```
spark.sql.adaptive.enabled=false
spark.sql.shuffle.partitions=1
spark.comet.batchSize=1024
spark.comet.exec.memoryPool.fraction=0.002
spark.comet.exec.respectDataFusionConfigs=true
spark.comet.datafusion.execution.sort_spill_reservation_bytes=65536
```
```sql
SELECT g, sum(v), count(DISTINCT d) FROM t GROUP BY g
```
Spark returns `count(DISTINCT d) = 2000` for every group. Comet returns 4000
for all 100 groups, and every `sum(v)` matches.
The metrics show where the extra rows come from. The `PartialMerge`
aggregate keyed by `(g, d)` reports `output_rows` = 400,000, twice the 200,000
distinct pairs. No aggregate reports a spill. With the default pool, the same
stage reports 200,000 rows and the answer is right.
### Expected behavior
The same counts as Spark. Under memory pressure, a `PartialMerge` that
deduplicates keys should spill or fail. It should never emit the same key twice.
### Additional context
- It isn't specific to `sum`. Since #4727, `collect_list` in a
distinct-combined plan runs natively too. Under the same settings, `SELECT _1,
collect_list(_3), count(DISTINCT _2) FROM t GROUP BY _1`, with a 32,768-row
table where `_1 = i % 4096`, `_2 = i % 8192` and `_3` is a string of about 260
characters, returns 4 or 5 instead of 2. On 1.1.0 that plan fell back to Spark
(`hasNativeArrayBufferAgg`), so the `collect_list` case is new on `main`.
- The `sum` case goes through the native `PartialMerge` support from #4003,
and `branch-1.0` and `branch-1.1` plan it the same way, so 1.0.x and 1.1.0 are
probably affected too. I haven't run it on those branches.
- `CometBaseAggregate.doConvert` already serializes an empty `PartialMerge`
(no aggregate expressions) as native `Final`, because it "must fully
deduplicate its keys". A `PartialMerge` with aggregate expressions needs the
same guarantee, but it runs as native `Partial`.
- I found this while reviewing #5903. That PR doesn't change it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
