andygrove opened a new pull request, #6867:
URL: https://github.com/apache/datafusion-comet/pull/6867
## Which issue does this PR close?
Closes #6857.
## Rationale for this change
Comet ran Spark's `PartialMerge` aggregate as DataFusion `Partial`, with
`MergeAsPartial` wrappers that turn its input into merges. When DataFusion's
`Partial` runs short of memory, it emits the groups it holds and starts again,
so a key can leave the stage more than once. In Spark's single-distinct plan,
the stage above the `PartialMerge` keyed by `(g, d)` counts its rows, so
`count(DISTINCT d)` counts every copy. In the issue's repro it returns 4000
instead of 2000, and nothing fails.
DataFusion's other merging modes don't fix this on their own:
- `PartialReduce` merges states into states, but it may emit a group more
than once too. On DataFusion `main` it emits early when it runs short of memory
(apache/datafusion#24785). In 55.2 it picks a stream that can't spill, because
Comet's pools report an unknown memory limit.
- `Final` spills, but `FinalHashAggregateStream` writes its spill files with
its input schema. When a Spark partial feeds the native `PartialMerge`, that
input still holds Spark's binary collect buffers while the spilled state is a
native list, so the spill fails with "column types must match schema types".
## What changes are included in this PR?
- The planner runs a `PartialMerge` as DataFusion `Single`. `Single` reads
the state columns as its input, which `MergeAsPartial` already redirects to
`merge_batch` (decoding Spark's buffers where needed). It spills with the state
schema and never emits a key twice. The mixed `{Partial, PartialMerge}` stage
above it still runs as `Partial`, since the `Final` after it merges whatever it
emits.
- `MergeAsPartial` evaluates to its state packed in a struct, and a
projection unpacks the structs into one column per state field, so the output
schema is the same as the `Partial` plan's. The inner accumulator is still
built with the inner expression's return field, because DataFusion's `max`
picks its accumulator by return type.
- `SpillReplayPool` (#6254) also matches `SingleHashAggregateStream`, which
merges and replays its spill files the same way `FinalHashAggregateStream`
does. Without it, a spilled `PartialMerge` failed reading its spill files back
at a 2 MiB pool.
## How are these changes tested?
- Three new `CometAggregateSuite` tests, each asserting that the native
`PartialMerge` spilled:
- The issue's `sum(v), count(DISTINCT d)` repro, which returns 4000
instead of 2000 without the fix.
- `collect_list` next to `count(DISTINCT)`, over a native partial and over
a Spark partial whose buffers the `PartialMerge` decodes. Both return 5 or 6
instead of 2 without the fix. The strings are 128 characters rather than the
issue's 260. With 260, the `PartialMerge` now spills correctly, but the `Final`
above it, which now receives complete lists, can't read back its 2 MB spill
chunks in the 4 MiB pool. That is #6363, not this change.
- A `PartialMerge` that reads several spill files back at a 2 MiB pool,
which fails without the `SpillReplayPool` change.
- A planner test checks that a `PartialMerge` plans as `Single` under the
projection, emits the same schema as the `Partial` plan, and merges `max` and
`count` states. It fails if the return-field override is removed.
- Ran locally: `CometAggregateSuite`, the `expressions/aggregate/` SQL file
tests, `CometFuzzAggregateSuite`, `CometFloatSemanticsSuite`,
`CometExecRuleSuite`, `CometNativeShuffleSuite`, `CometShuffleSuite`,
`DisableAQECometShuffleSuite`, `CometExecSuite` and `cargo test -p
datafusion-comet --lib`. The new tests also pass on Spark 3.4 and 3.5.
--
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]
