refactor: Split hash aggregation logic into separated streams#22729
refactor: Split hash aggregation logic into separated streams#227292010YOUY01 wants to merge 5 commits into
Conversation
| impl Stream for PartialFinalHashAggregateStream { | ||
| type Item = Result<RecordBatch>; | ||
|
|
||
| fn poll_next( |
There was a problem hiding this comment.
The state machines are identical for now, but in follow-up work, such as skipping partial aggregation for high-cardinality inputs, their control flows will diverge. I think separating them improves clarity, as discussed in #22710.
Some duplication is inevitable, but that is the trade-off.
There was a problem hiding this comment.
Given these are two separate structs, what is the rationale for keepign them in the same module?
datafusion/physical-plan/src/aggregates/hash_aggregate.rs
Perhaps it would be clearer if we put them in their own modules (could be a follow on PR)
like
datafusion/physical-plan/src/aggregates/streams/partial_final.rs
datafusion/physical-plan/src/aggregates/streams/initial_partial.rs
|
cc @Dandandan @ariel-miculas @alamb, who have expressed interest before. |
|
I'm curious about the high-level vision: is the plan to close #15591 in favor of this new approach? I would like the redesign of hash aggregation to take into account the memory constraints imposed by the finite memory pool, i.e. how does the implementation perform under OOM conditions.
Otherwise we'll end up with the same issues that exist now. E.g. EmitTo::First(n) wasn't designed for emitting a large portion of the existing groups, so it over-allocated when used for emitting early in partial aggregation OOM case. |
Yes, the goal is to support blocked state management. The existing challenge is that the current implementation is hard to extend and review. I want to clean things up through this refactor first, and then apply the actual change.
All of these issues are symptoms of managing state in a large contiguous |
Current one seems the refactor of original But in #15591, code changes in I totally agree with we should refactor code in I think maybe #15591 and refactor can be pushed forward in parallel, they don't have very much in common. |
That's true, the actual code change is about the same. The major thing get reduced is the cognitive overhead when maintaining this feature, I'll try to explain better: I think the real pain point is:
The purpose of this refactor is to make that boundary explicit. With this split, the key blocked state-management changes (
Yes I agree. I was thinking maybe #15591 can be updated after this refactor, if we can ensure the blocked |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you @2010YOUY01 -- this looks like a good idea, but I am slightly confused
- The PR says it "splits" logic but I actually see two entirely new streams (and no reduction / refactoring of the existing streams)
- I am not sure how the two new streams are related to the existing stream
I am also worried about so much new code (I realize that some duplication is inevitable, but this seems like it adds 2 new additional copies)
Also, the PR descriptio says
Todo in this PR:
Add a temporary configuration enable_migration_aggregate to turn off this path
But it is unchecked -- do you still plan to do that?
| impl Stream for PartialFinalHashAggregateStream { | ||
| type Item = Result<RecordBatch>; | ||
|
|
||
| fn poll_next( |
There was a problem hiding this comment.
Given these are two separate structs, what is the rationale for keepign them in the same module?
datafusion/physical-plan/src/aggregates/hash_aggregate.rs
Perhaps it would be clearer if we put them in their own modules (could be a follow on PR)
like
datafusion/physical-plan/src/aggregates/streams/partial_final.rs
datafusion/physical-plan/src/aggregates/streams/initial_partial.rs
| /// is for the final stage. | ||
| /// | ||
| /// See [`InitialPartialHashAggregateStream`] for details. | ||
| pub(crate) struct PartialFinalHashAggregateStream { |
There was a problem hiding this comment.
A minor nit is I found this naming confusing
- PartialFinalHashAggregateStream -- the
Partialis beforeFinal InitialPartialHashAggregateStream-- thePartialis afterInitial
Would it be more consistent to call it PartialInitialHashAggregateStream
There was a problem hiding this comment.
I think it might be clearer to call them PartialHashAggregateStream and FinalHashAggregateStream 🤔. I'll update the names.
| InitialPartialHash(InitialPartialHashAggregateStream), | ||
| PartialFinalHash(PartialFinalHashAggregateStream), | ||
| GroupedHash(GroupedHashAggregateStream), |
There was a problem hiding this comment.
I am a little unclear what the difference betweeen InitialPartialHash PartialFinalHash and GroupedHash
Is it possible to leave some comments?
There was a problem hiding this comment.
Also, I am not clear that the relationship between these different strema types are
For example is GroupedHash the same as InitialPartialHash?
There was a problem hiding this comment.
Also, why os GroupedHashAggregateStreamnot changed 🤔
There was a problem hiding this comment.
Added comments in 3184e00 for 1. input/output schema 2. Incremental migration strategy for why GroupsHashAggregateStream hasn't changed.
For example is
GroupedHashthe same as InitialPartialHash?
GroupedHash is a reused path for InitialPartial and ParitalFinal and many others, this might explain why changing things inside GroupedHash is very hard.
| //! | ||
| //! `AggregateExec` keeps finite-memory, ordered, limit, grouping-set, | ||
| //! `partial state -> partial state`, and single-stage aggregation on | ||
| //! `GroupedHashAggregateStream` for now. |
There was a problem hiding this comment.
What does it mean "for now" ? That seems like a comment about the PR rather than about the state of the code (and thus may be confusing once this PR is merged)
|
I love the idea of refactoring to reduce the cognative load, BTW |
@alamb Thank you for the review! The idea is incremental migration: we have to keep some duplication during the migration, but eventually the old // Before
AggregateExec::execute() {
// One multiplexed stream handles all paths.
GroupedHashAggregateStream::new(...)
}
// Incremental migration process
AggregateExec::execute() {
match self.plan_stream() {
// Migrated path 1/5
StreamPlan::Partial => PartialHashAggregateStream::new(...),
// Migrated path 2/5
StreamPlan::Final => FinalHashAggregateStream::new(...),
// Paths not migrated yet still use the old implementation.
_ => GroupedHashAggregateStream::new(...),
}
}
// After all paths are migrated
AggregateExec::execute() {
match self.plan_stream() {
StreamPlan::Partial => PartialHashAggregateStream::new(...),
StreamPlan::Final => FinalHashAggregateStream::new(...),
// Migrated path 3/5
// ...
// No fallback remains; GroupedHashAggregateStream can be deleted.
}
}The reasons for this approach are:
Added in 4a2b907 |
github-actions
Bot
added
sqllogictest
I see -- the plan for eventually removing |
|
run benchmarks |
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you @2010YOUY01 -- this looks good to me as long as the benchmarks show no performance regression
my only concern is that we will get half way through the migration and leave the code in an even worse state (even more complicated) but I think your comments have made it about as clear as possible
Thank you for pushing this along
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing split-aggr-refactor-only (09462b9) to 48adae4 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing split-aggr-refactor-only (09462b9) to 48adae4 (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing split-aggr-refactor-only (09462b9) to 48adae4 (merge-base) diff using: tpch File an issue against this benchmark runner |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
|
Benchmark for this request failed. Last 20 lines of output: Click to expandFile an issue against this benchmark runner |
|
run benchmarks |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing split-aggr-refactor-only (bd75229) to 48adae4 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing split-aggr-refactor-only (bd75229) to 48adae4 (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing split-aggr-refactor-only (bd75229) to 48adae4 (merge-base) diff using: tpch File an issue against this benchmark runner |
|
I expect very high-cardinality aggregations to be slower for now, because I plan to address that in a follow-up PR: I would change the default for the configuration added in 4a2b907 to |
Make sense. Currently , it is actually very very hard to understand the state machine in
Sounds good, #15591 still not apply the blocked logic into splling path and others, and I think really hard if we do it in the old logic... |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
5 participants
Which issue does this PR close?
Rationale for this change
See issues.
This PR split out partial and final aggregate strem from
GroupsHashAggregateStreamTo fully migrate hash aggregation, we have to
I think they should be leave to follow up PRs
Todo in this PR:
enable_migration_aggregateto turn off this pathSince it should be a regression if the above features are not added, it also helps if to prevent potential regressions from the migration of other aggregate streams.
What changes are included in this PR?
Split out the streams from
GroupsHashAggregateStreamAre these changes tested?
Are there any user-facing changes?