andygrove opened a new issue, #6880:
URL: https://github.com/apache/datafusion-comet/issues/6880
### Describe the bug
#6867 fixes #6857 by running a native `PartialMerge` as DataFusion `Single`,
which spills instead of emitting a key twice. When the `PartialMerge`'s input
is sorted on a prefix of its grouping keys, DataFusion 55.2 doesn't run
`SingleHashAggregateStream`: `should_use_single_hash_stream` requires
`InputOrderMode::Linear`, so the aggregate falls back to the legacy
`GroupedHashAggregateStream`. Under memory pressure, that stream's spill asks
for more than its share of the pool to sort the batch it is about to write, and
the task fails:
```
Failed to reserve memory for sort during spill: Resources exhausted:
Additional allocation failed for GroupedHashAggregateStream[0] (col_0) ...
Error: Failed to acquire 745281 bytes where this consumer already holds 8396
bytes and the fair limit is 419430 bytes, 5 registered (0 bytes overcommitted)
```
Before #6867, this shape ran as DataFusion `Partial`. Its
`OrderedPartialAggregateStream` emits all its groups when a reservation fails,
so it returned the wrong `count(DISTINCT)` of #6857. A task failure is better
than a wrong answer, but the stage should spill.
The shape needs the sort and both aggregates in one native plan, with no
exchange between the `Partial` and the `PartialMerge`. That happens when the
child's hash partitioning on `g` already satisfies the clustered distribution
on `(g, d)`: after a native sort-merge join on `g`, for example, or in the
repro below.
### Steps to reproduce
On the #6867 branch (2ce73e59e1), in a `CometTestBase` suite (2g off-heap):
```scala
spark
.range(0, 125000, 1, 1)
.selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS k",
"id AS v")
.write
.parquet(path)
withSQLConf(
SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName,
CometConf.COMET_BATCH_SIZE.key -> "1024",
CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
"spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536",
CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (2.0 / 2048).toString)
{
spark.read
.parquet(path)
.repartition(1, col("g"))
.sortWithinPartitions("g")
.groupBy("g")
.agg(expr("max(v)"), count_distinct(col("k")))
.collect()
}
```
It fails with the error above. Of the pools from 2 to 48 MiB that I tried
(pool fraction `n / 2048`), only 24 and 48 MiB passed. A larger pool doesn't
make it go away, because the sort reservation grows with it: about 0.75 MB at 2
MiB, 2.5 MB at 8 and 12 MiB, 4.9 MB at 14 to 20 MiB and 9.6 MB at 32 MiB,
between 1.0 and 1.8 times the stream's fair share.
### Expected behavior
The `PartialMerge` spills and reads its spill files back, as it does when
its input is unsorted.
### Additional context
- `SpillReplayPool` doesn't help. The refused request comes from the spill
itself, while the stream holds only 8 KB, not from reading the spill files back.
- DataFusion `main` replaces this fallback with
`OrderedSingleAggregateStream`
(`datafusion/physical-plan/src/aggregates/ordered_single_stream.rs`), so this
should be rechecked when Comet moves to DataFusion 56 (#6410).
- If it has to be fixed before then, Comet could run a `PartialMerge` with
sorted input as DataFusion `Final`. Its `OrderedFinalAggregateStream` spills,
and the `MergeAsPartial` struct output works under `Final` too. `Final` reads
the states right after the grouping keys, which holds for a pure
`PartialMerge`, and it spills with its input schema, which is native state when
the input is sorted by a native operator.
--
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]
