Repository navigation
fix: [branch-1.2] spill a native PartialMerge whose input is sorted on some of its keys (#7222) - #7286
Merged
Conversation
…n some of its keys (apache#7222) * fix: spill a native PartialMerge whose input is sorted on some of its keys Comet runs a Spark PartialMerge as DataFusion Single so that it spills rather than emit a key twice (apache#6857). When the input is sorted on a prefix of the grouping keys, DataFusion 55.2 runs Single with the legacy GroupedHashAggregateStream, whose spill asks for more than its share of the pool to sort the batch it writes, and the task fails. Plan that PartialMerge as DataFusion Final instead. Final runs an OrderedFinalAggregateStream there, which spills and reads its spill files back, and MergeAsPartial evaluates to its state under Final as it does under Single. Final reads each state from the columns after the keys and spills its groups in the input's types, so the planner switches only when the input is the grouping keys followed by native-typed states. A Spark partial's binary collect buffers keep the Single plan. Closes apache#6880. * refactor: decide the PartialMerge aggregate mode before building it Read the input order from the child's equivalence properties, as AggregateExec does, instead of building the aggregate as Single and again as Final. Simplify the state type guard, note when the switch can be rechecked, and cover input sorted on every key in the planner test. * address review: lazy sort check, document fair-limit exemption, cover changing g and a sort-merge join * address review: walk state types once, share planner test helpers * test: cover the PartialMerge inputs Final rejects and run the join case under AQE Add planner cases for swapped keys, states after one key and a column after the states, which each keep the PartialMerge as Single. Check that the join test's PartialMerge reaches the sort-merge join without an exchange, and keep AQE on by disabling the suite's runtime broadcast threshold. Name the DataFusion 56 stream as it is on main. * address review: keep Single for a PartialMerge over fully sorted input * address review: shorten the PartialMerge planning comment (cherry picked from commit d8ead26) Adapted for branch-1.2: - Dropped the comment on `limited_to_its_share` in `fair_pool.rs`, the comment on `ORDERED_FINAL_AGGREGATE` in `spill_replay.rs`, and the paragraph in `memory_management.md`. They describe the OrderedFinalAggregateStream exemption from the fair limit, which came in with apache#7096 and is not on this branch. - The sort-merge join test uses a 32 MiB pool instead of 16 MiB. Without apache#7096 a consumer's fair limit is the pool divided by the number of consumers, and at 16 MiB the sort below the join fails before the PartialMerge runs.
github-actions
Bot
added
bug
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Backport of #7222 to
branch-1.2. #7222 closes #6880 onmain.Part of #7204.
Rationale for this change
Since #6867, a native
PartialMergeruns as a single-mode aggregate, which DataFusion 55.2 cannot spill when its input is sorted on a prefix of the grouping keys. Acount(DISTINCT k)over input sorted on the group key then fails the task withAdditional allocation failed for GroupedHashAggregateStream. #7222 runs such aPartialMergeas a final aggregate, which does have a spill path.branch-1.2has the failure, so it needs the fix before 1.2.0 ships.branch-1.1also has #6867, through #7020, so it has the same failure. Its backport is not part of this PR.What changes are included in this PR?
git cherry-pick -x d8ead26b1eof the squash commit of #7222. The planner changes inplanner.rs,planner/aggregate.rsandmerge_as_partial.rsapplied without conflicts. Adaptations:limited_to_its_shareinfair_pool.rs, the comment onORDERED_FINAL_AGGREGATEinspill_replay.rs, and the paragraph inmemory_management.md. They describe theOrderedFinalAggregateStreamexemption from the fair limit, which came in with fix: make fair_unified's per-consumer limit bind when tasks run concurrently #7096 and is not on this branch. The header comment inspill_replay.rsapplies as it is.Not enough memory to continue external sortbefore thePartialMergeruns. At 32 MiB thePartialMergespills, and the test still asserts that it does.How are these changes tested?
On this branch, with a
libcometbuilt from the pick:cargo check --all-targetsinnative/passes, and the new Rust testpartial_merge_over_sorted_input_runs_as_finalpasses.CometAggregateSuitepasses: 164 tests, none failed.libcometbuilt from unmodifiedbranch-1.2, both new tests fail withAdditional allocation failed for GroupedHashAggregateStream. So the tests catch the bug, and the branch has it.I tested the pick before #7222 merged, then re-picked the squash commit. The re-pick has the same content, and
cargo check --all-targetspasses on it. I did not rerun the JVM suites on the re-pick, and I did not run the other suites that #7222 ran.