Add streaming_groupby for stateful streaming aggregation - #21924
Conversation
PointKernel
added
feature request
|
Results of the groupby cardinality benchmark, the input is fixed at 20M elements: num_aggregations = 8 cc @devavret |
|
I have an idea for how you can work with non-fixed width keys. I'll explain here with just keys and not values
For each batch, while inserting, we mark the new keys we saw in this batch. Then we'll compact and store it in a way that can be referred to while inserting subsequent batches. We'll end up with several set of key tables held inside the groupby object which should hopefully be less memory than if we had sparsely populated a key table of the same size as the single large hash set. So we start with a pre-allocated large hash set and similarly sized arrays for values. Along with that, we allocate two more vectors: key_table, key_idx. These refer to the batch we first saw the key in and the index of the key row in the compacted batch. So suppose we allocate an empty hash set with 8 empty slots like so:
and we receive our first batch of keys
and they end up in the hash set like so:
while inserting, we remember where each inserted key ended up so we store the slot index in a key batch sized vector
Normally we use this to gather and return the unique keys
but here we'll also use it to mark the location of these keys in the map
Now the next batch of keys come in
While inserting in the hash set, we have enough information to either refer to this batch itself or to keys seen in batches before. The row comparator would have the following logic After insertion the new batch looks like this
And the hash set and associated vectors look like this
|
@devavret Thanks for sharing your proposal. My main question is about how the comparison would actually be performed. cudf row operators are designed for either self-comparison or comparing two tables. How would we compare a row in the new batch against a row from a previous batch? Are you suggesting storing all preprocessed keys from previous batches so that we can then use a two-table comparator between the new batch and a specific previous batch? There’s also a subtle issue around atomically updating the associated vectors. For example, if thread A wins the hash table CAS and inserts into slot S, there’s a window before it writes key_table[S] and key_idx[S]. During that time, thread B may probe slot S, see it as occupied, and read (partial) uninitialized data from those arrays. The root cause is that the hash set slot and the associated metadata arrays live in separate memory locations where cuco’s atomic CAS only protects the slot itself, not the side arrays. There’s no straightforward way to atomically update all three together without either packing them into a single value (e.g., a struct/tuple stored as the slot) or introducing fences/spin-waits, both of which would hurt performance. |
Yes, we'd need to store preprocessed keys but only the unique ones. So each batch gets preprocessed once before inserting, then after it has been inserted and compacted, preprocess again and store. I'm suggesting we generalize this so the comparison can happen across arbitrary number of tables. I feel like there must be a way to wrap the current two table row operator to achieve this. The first row is from the table being inserted. The batch idx helps us point to the other table and the row idx tells us which row in the other batch to compare to. So the row comparator has access to current table and a store containing deduplicated previously seen batches.
Right, I'm not suggesting the batchIdx and kyIdx be updated immediately upon batch insertion. The steps I described are like so:
|
|
/ok to test fc91494 |
|
Thanks @devavret. The implementation has been updated to support variable-width keys following your design described in #21924 (comment). The main difference is in how we identify keys in the hash set. Instead of inserting batch-local indices and using slot positions for companion vector lookups, we maintain a We use this approach because cuco's comparator receives stored values, not slot positions, so there is no way to index companion vectors by slot. The tradeoff is that the encoded index space is consumed by the full batch size (not just distinct keys), so cumulative batch rows must stay within For your reference, I collected the latest results of the same benchmark launched in #21924 (comment): and The performance gap between normal groupby and streaming groupby has been significantly reduced, with both approaches now demonstrating comparable performance in high-cardinality scenarios for both small and large numbers of aggregations. The remaining gap exists only in low-cardinality cases, where streaming groupby has not yet incorporated shared memory optimizations. Additionally, peak memory usage has decreased from 229 MB to 152 MB across all cases, meaning streaming groupby now consumes less peak memory than normal groupby. |
|
@devavret would you please complete your review? |
|
@ttnghia would you please share your review? |
|
Does streaming aggregate with 1 iteration produce exactly the same output as the normal aggregation? If so, should we considering deprecate the normal aggreation pipeline, and adjust the implementation of streaming aggregation accordingly? I.e., instead of calling |
| * @throws std::overflow_error if accumulated rows plus batch size exceeds `max_groups` | ||
| * @throws cudf::logic_error if distinct keys exceed `max_groups` | ||
| */ | ||
| void aggregate(table_view const& data, rmm::cuda_stream_view stream = cudf::get_default_stream()); |
There was a problem hiding this comment.
request: can we have an API tell us the unique number of keys so far. seems like it's stored anyway and would be trivial to surface.
There was a problem hiding this comment.
bump ^. Realistically, it's not possible to know the exact cardinailty beforehand and engines are supposed to make a guess. So when it does grow beyond initial mad_distinct_keys, we want to create a new streaming_groupby with more capacity, merge the current one into it and destroy. Currently, this means this aggregate() call lives in a try{} with the resizing done in catch{} block. It would be nicer to not have to rely on exceptions for this runtime behaviour.
There was a problem hiding this comment.
Also, I'm wondering whether this should store a stream instead of taking a new one each time. Do we need to wait on the previous stream before calling aggregate with a new batch and new stream?
There was a problem hiding this comment.
Okay, I see that you only get that information AFTER checking for new distinct keys in the current batch. Which means you need to do some work before you can say the current batch will overflow or not.
There was a problem hiding this comment.
This is already exposed as distinct_keys.
Thank you @ttnghia, this is a great question. Currently the streaming aggregate with 1 iteration produces the same result, yes. But the big difference is that streaming groupby can't use the shared memory optimization, and there isn't a clear way to extend streaming aggregations to support that. So I expect that both streaming and non-streaming will exist for the foreseeable future. Would you please share another pass of review feedback? |
- Move transient-encoding overflow check from constructor to per-batch in aggregate - Preserve key schema when no groups are materialized - Add non-ASCII UTF-8 string key test - Note in dictionary sum test why streaming path is skipped - Drop unused cudf/copying.hpp include in group_max benchmark
3701ea9 to
8bf5ceb
Compare
|
/merge |
962d15b
into
NVIDIA:release/26.06
8 participants
Description
Closes #18182
This PR adds
streaming_groupby, a stateful groupby that accumulates partial aggregates across batches using a single persistent hash table.Users specify
max_groups, the maximum number of distinct groups expected, and all main data structures are allocated once and reused without resizing. The hash table stores asize_typegroup ID per slot.The ID is global across the stream: each distinct group, in the order it is first seen, is assigned a stable ID in
[0, distinct_count)that is shared by the result table (used as the row index) and by the companion array (used as the lookup index). The actual keys live in a list of per-batch compacted key tables. The companion array, of lengthmax_groups, holds a{batch_id, row_id}pair for each group ID, pointing back to where that group's representative key is stored. Equality probes resolve a slot ID through the companion array into the correct preprocessed batch table and compare via an n-table row comparator.Each batch is processed in two steps. The first step calls
insert_and_findagainst the hash set. Winners write a transient valuemax_groups + batch_idxinto their slot, which is distinguishable from any real group ID since real IDs live in[0, max_groups). Existing slots already hold a final ID and are returned as-is. A side flag array marks which rows won their slot, and a slot-offset array records where each row landed for cheap revisits. The newly inserted rows are stream-compacted and gathered into a fresh compacted key table that is appended to the per-batch list. The second step walks only the new keys, atomically rewrites their transient slot values to stable global IDs starting at the current distinct count, and writes the matching{batch_id, row_id}entries into the companion array. A final reread converts any remaining transient slot reads into global IDs, so every row in the batch ends up mapped to its stable group ID. Aggregations are updated atomically into a single result table indexed directly by these IDs.Merging reprobes the other object's compacted keys against this hash table to recover their target group IDs in this object's ID space, then atomically combines the matching result rows. Finalization concatenates the per-batch compacted key tables to produce the distinct-keys output, slices the result table to
[0, distinct_count)— no gather is needed, since the global IDs are already the row indices — and runs the compound-aggregation finalizers to produce user-facing columns. The internal state is left intact, so furtheraggregatecalls remain valid.Certain trade-offs are intentional. For example, the streaming groupby is designed to deep-copy all distinct keys locally. This enables batch-based processing: once a batch has completed the add step, its input data can be released, which helps reduce memory usage.
Additionally, the current code path does not support shared memory. As a result, inputs with very low cardinality can suffer from poor runtime performance due to high atomic contention, since many updates target the same key or memory location. This is an accepted trade-off. In practice, downstream users can run a standard groupby to estimate cardinality; if it is low, they can concatenate all input data and use a regular groupby instead, which typically yields better performance.
Checklist