andygrove opened a new pull request, #7020: URL: https://github.com/apache/datafusion-comet/pull/7020
## Which issue does this PR close? Closes #6857 on `branch-1.1`. #6867 already closed it on `main`. ## Rationale for this change This backports #6867 to `branch-1.1`. Comet runs a Spark `PartialMerge` aggregate as a DataFusion `Partial` aggregate. When `Partial` runs short of memory it emits the groups it holds and starts again, which is fine for a real partial aggregate. But Spark's single-distinct plan counts the rows of the `PartialMerge` in the stage above it, so a key that leaves the stage twice is counted twice. `SELECT g, sum(v), count(DISTINCT d) FROM t GROUP BY g` then returns too large a `count(DISTINCT d)`, with no error (`priority:critical`). The `PartialMerge` and `Partial` mapping in the planner is the same on this branch, and the two tests below fail here without the change. This is not a 1.1.0 regression, and it needs memory pressure on the aggregate. I am proposing it because the query shape is common and the result is silently wrong, in the same way as #6270 and #6700. It is also the largest of the backports so far, so I'd like a closer look at the hand port in `merge_as_partial.rs` than a clean cherry-pick would need. ## What changes are included in this PR? One cherry-pick with `-x`, resolved by hand where this branch differs. See #6867 for the design: the `PartialMerge` runs as DataFusion `Single`, which spills. `MergeAsPartial` evaluates to a struct of its state, and a projection unpacks the struct into the columns the `Partial` plan emitted. `SpillReplayPool` also covers `SingleHashAggregateStream`. The adaptations are recorded in the commit message: - `merge_as_partial.rs` is ported by hand onto this branch's file, which has no `PartialMergeStateDecoder` or `collect_list` state fields (#4727). I started from this branch's file and applied only the #6867 changes. The only difference from `main`'s diff is the decoder-related lines, which are not here. - The planner has no `ordered_by_grouping_keys` sort, which is a later change. So the planner hunk keeps `main`'s unpacking and nothing after it, and the tests for that sort are not taken. The `PartialMerge` planner test is otherwise identical, minus a `..Default::default()` that clippy flags once the missing field is gone. - `spill_replay.rs` has no `unwrap_spill_replay`, because #6271 is not on this branch. That only moved the conflict context. - `CometAggregateSuite`: the `collect_list` `PartialMerge` test is not taken, because this branch falls back to Spark for `collect_list` next to a distinct aggregate (#4727 is not here). I ran it to confirm: it fails with "Expected a native PartialMerge aggregate", which is the fallback and not a wrong result. - The suite's `functions` import line is unchanged, since the dropped test was the only user of the extra imports. ## How are these changes tested? Run locally on this branch with JDK 17: - On the default Spark 4.1 profile, `CometAggregateSuite`, the `aggregate` fixtures in `CometSqlFileTestSuite` and `CometExecSuite` pass (327 tests). That includes the two new `PartialMerge` tests. Scalastyle and spotless ran in the same build. - On Spark 3.5 and Scala 2.12, the sources and tests compile. - `cargo test -p datafusion-comet --lib` for the planner, `merge_as_partial` and memory pool tests passes (73 tests, including the new planner test). `cargo fmt --check` and workspace `cargo clippy --all-targets -D warnings` are clean. My local toolchain is Rust 1.97, and CI uses stable. - To check that the tests catch the bug on this branch, I changed the planner to run `PartialMerge` as `Partial` again, without the unpacking, rebuilt `libcomet`, and ran the two new tests. Both failed with `Results do not match`. The `md5` of the library Maven copied matched the rebuilt one both times. By CI's path routing, which I checked with `dev/ci/compute-changes.py`, this PR runs the all-profiles Linux build, macOS, the Spark SQL jobs for 3.5, 4.0 and 4.1, and the four Iceberg versions, so no `run-*` label is needed. -- 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]
