fix: [branch-1.1] keep one row per key in a native PartialMerge that runs short of memory (#6867)

[andygrove]

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 (feat: add native distinct-combined collect_list aggregate support #4727). I started from this branch's file and applied only the fix: keep one row per key in a native PartialMerge that runs short of memory #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 fix: count pool overcommit as untracked memory in the native memory usage log #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 (feat: add native distinct-combined collect_list aggregate support #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.

[kazuyukitanimura]

docs/source/contributor-guide/memory_management.md (pool-stack diagram, and the section "Final aggregates reading their spill files back"). Thanks for the careful
port. The SpillReplayPool description says it covers final aggregates only. With this change it also covers the SingleHashAggregateStream that Comet now uses for
a PartialMerge, as the spill_replay.rs header says. Could you update both places in this PR? Also please open a matching fix for main, since #6867 did not touch
the page and the two branches should agree.

[andygrove]

docs/source/contributor-guide/memory_management.md (pool-stack diagram, and the section "Final aggregates reading their spill files back"). Thanks for the careful port. The SpillReplayPool description says it covers final aggregates only. With this change it also covers the SingleHashAggregateStream that Comet now uses for a PartialMerge, as the spill_replay.rs header says. Could you update both places in this PR? Also please open a matching fix for main, since #6867 did not touch the page and the two branches should agree.

I went ahead and merged since this fixes a correctness bug. The suggested doc edit is for the contributors guide which we do not maintain on release branches. It only gets published from the main branch.