Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

perf(hash agg): use get_mut instead of pop+put pattern #8691

Merged
merged 3 commits into from
Mar 21, 2023
Merged

perf(hash agg): use get_mut instead of pop+put pattern #8691

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
57c3140
use `get_mut` instead of `pop`+`put`
stdrc Mar 21, 2023
f848ac5
fix comment
stdrc Mar 21, 2023
3fdc03a
may be better
stdrc Mar 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions src/stream/src/executor/hash_agg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

use std::collections::{HashMap, HashSet};
use std::marker::PhantomData;
use std::ptr::NonNull;
use std::sync::Arc;

use futures::{stream, StreamExt, TryStreamExt};
Expand Down Expand Up @@ -436,11 +437,16 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {

// Calculate current outputs, concurrently.
let futs = keys_in_batch.into_iter().map(|key| {
// Pop out the agg group temporarily.
let mut agg_group = vars
.agg_group_cache
.pop(&key)
.expect("changed group must have corresponding AggGroup");
// Get agg group of the key.
let agg_group = {
let mut cache_ptr: NonNull<_> = (&mut vars.agg_group_cache).into();
// SAFETY: `key`s in `keys_in_batch` are unique by nature, because they're
// from `group_change_set` which is a set.
let cache = unsafe { cache_ptr.as_mut() };
cache
.get_mut(&key)
.expect("changed group must have corresponding AggGroup")
};
async {
let curr_outputs = agg_group.get_outputs(&this.storages).await?;
Ok::<_, StreamExecutorError>((key, agg_group, curr_outputs))
Expand All @@ -452,7 +458,7 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
.try_collect()
.await?;

for (key, mut agg_group, curr_outputs) in outputs_in_batch {
for (key, agg_group, curr_outputs) in outputs_in_batch {
let AggChangesInfo {
n_appended_ops,
result_row,
Expand All @@ -477,9 +483,6 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
this.result_table.insert(result_row);
}
}

// Put the agg group back into the agg group cache.
vars.agg_group_cache.put(key, agg_group);
}

let columns = builders
Expand Down