Intermediate result blocked approach to aggregation memory management - #15591
Intermediate result blocked approach to aggregation memory management#15591Rachelint wants to merge 102 commits into
Conversation
|
Hi @Rachelint I think I have a alternative proposal that seems relatively easy to implement. |
Really thanks. This design in pr indeed still introduces quite a few code changes... I tried to not modify anythings about
But I found this way will introduce too many extra cost... Maybe we place the |
cc37eba to
f690940
Compare
95c6a36 to
a4c6f42
Compare
2100a5b to
0ee951c
Compare
|
Has finished development(and test) of all needed common structs!
|
c51d409 to
2863809
Compare
|
It is very close, just need to add more tests! |
31d660d to
2b8dd1e
Compare
|
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 |
@adriangb hello, is it possible to authorize me to trigger benchmark through bot? |
|
run benchmarks clickbench_partitioned |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing intermeidate-result-blocked-approach (5869167) to a27f030 (merge-base) diff using: clickbench_partitioned 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 |
| /// - [`EmitTo::NextBlock`]: pops a single block. | ||
| pub fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<T>> { | ||
| match emit_to { | ||
| EmitTo::All => self.inner.pop_block().ok_or_else(|| { |
There was a problem hiding this comment.
this doesn't seem to match the above description:
[
EmitTo::All]: drains every block via repeatedpop_blockand
concatenates the results into a singleVec<T>.
There was a problem hiding this comment.
Maybe a bit confusing now due to:
EmitTo::FirstandEmitTo::Allonly meaningful inflat modeEmitTo::NextBlockonly meaningful inblocked mode
Here (vec_block_store.rs) maybe more suitable to return error when encountering EmitTo::First and EmitTo::All?
| return None; | ||
| } | ||
|
|
||
| let block = mem::take(&mut self.inner[self.cursor]); |
There was a problem hiding this comment.
You are initializing block even when not needed, which might not reduce memory usage.
why not using VecDeque and using pop_front instead of pop block and having all of this
There was a problem hiding this comment.
VecDeque is used at the beginning, but the index op in VecDeque actually not really trivial as Vec, and switch to Vec later suggested by reviews.
But I agree with that VecDeque may be more suitable, due to:
- its main target is saving memory (although Vec will not occupy much)
- the logic can become simple, and I think as a initial pr, it should keep as simple as possible.
| /// [`GroupsAccumulator::supports_blocked_groups`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator::supports_blocked_groups | ||
| /// | ||
| #[derive(Debug)] | ||
| pub struct BlockedBlockStore<B: Block> { |
There was a problem hiding this comment.
I've tried to copy the code ArrowBytesMap to BlockedArrowBytesMap where the Block is:
#[derive(Debug)]
struct BufferAndOffsetBlock<O: OffsetSizeTrait> {
/// In progress arrow `Buffer` containing all values
buffer: BufferBuilder<u8>,
/// Offsets into `buffer` for each distinct value. These offsets as used
/// directly to create the final `GenericBinaryArray`. The `i`th string is
/// stored in the range `offsets[i]..offsets[i+1]` in `buffer`. Null values
/// are stored as a zero length string.
offsets: Vec<O>,
}
impl<O: OffsetSizeTrait> Default for BufferAndOffsetBlock<O> {
fn default() -> Self {
Self::new(INITIAL_BUFFER_CAPACITY)
}
}
impl<O: OffsetSizeTrait> Block for BufferAndOffsetBlock<O> {
type T = Vec<u8>;
fn new(capacity: usize) -> Self {
Self {
buffer: BufferBuilder::new(capacity),
offsets: vec![O::default()], // first offset is always 0
}
}
fn fill_default_value(&mut self, fill_len: usize, default_value: Self::T) {
for _ in 0..fill_len {
self.buffer.append_slice(&default_value);
}
let mut last_offset = *self.offsets.last().unwrap();
let len_in_bytes = O::usize_as(default_value.len());
for _ in 0..fill_len {
self.offsets.push(last_offset + len_in_bytes);
last_offset = last_offset + len_in_bytes;
}
}
fn len(&self) -> usize {
self.offsets.len() - 1
}
}but then I realized the fill_default_value will not serve any benefit when called from resize since I will have to update the offsets (zero cost, O(1)) and add to the buffer - expensive since it requires shifting the bytes (unless the default value is 0)
There was a problem hiding this comment.
Methods in Block trait is mainly abstracted from logic in prim_op.rs, and may only useful for accumulators.
Maybe we should define some trait different for groups case?
For me, how to define the suitable traits is actually a main block point...
|
So I took this pr and added support locally for binary and it fixed the issue we had with huge batches coming out of aggregate, however I have concerns because ideally It should be enabled in all cases in order for it to be benefitial - for us it's partial/final with both spill/emit early support but there.
|
I think this is the key observation -- I think there is no way we can add the extra indirection and not slow down performance for small aggregates. For higher cardinality, I think it may not be as bad as we are going to be doing lots of random memory acesses anyways as the table doesn't fit in the processor caches. So I am guessing what we will need is some way to use direct indexing for small hash tables, but when the table grows above a certain size switch to 2 part indexes But that is just a theory and I have not tested it yet |
|
Another complexity is supporting |
Yeah I think this should be possible -- have two paths, direct indexing (for a single block under the threshold) / the block-based approach that is chosen per batch. I think some other things that could be tried:
|
|
I’m actually more concerned about the implementation complexity, especially for complex types—the work required just to extract the trait was already quite painful. I’m somewhat less concerned about the performance impact. Based on the benchmarks that have been run throughout this PR, the regression appears relatively modest, ranging from roughly neutral to around 10%. |
| /// [`BlockStore`] lets flat and blocked group state share the same accumulation | ||
| /// flow while using different physical layouts. Implementations should keep | ||
| /// block lookup cheap because it is used by per-row accumulator update paths. | ||
| pub trait BlockStore<B: Block>: |
There was a problem hiding this comment.
@Dandandan @ariel-miculas @2010YOUY01 @alamb
I think the greatest difficulty currently being faced is that how to abstract the common trait (BlockStore here) to help conveniently impl blocked accumulator and group values
Current one can work and zero cost when we disable blocked mode, but I think it seems too complex...
| if let Some(emit_to) = self.group_ordering.oom_emit_to(n) | ||
| && let Some(batch) = self.emit(emit_to, false)? | ||
| { | ||
| return Ok(Some(ExecutionState::ProducingOutput(batch))); |
There was a problem hiding this comment.
Serialize and spill the batch to disk.
| /// Internally uses a `Vec<B>` with a `start` offset to track the first active | ||
| /// block. When blocks are popped via [`Self::pop_block`], the block is swapped | ||
| /// out in O(1) using `mem::replace` and the `start` cursor advances, avoiding | ||
| /// the O(n) shift cost of `Vec::remove(0)`. |
There was a problem hiding this comment.
This avoids the O(n) shift but it doesn't free the memory, I'm wondering whether a VecDeque would work better here, so inner: VecDeque instead of Vec
The memory usage of Vec may not be large (8bytes(pointer) * num_blocks).
But index op of VecDeque is very expansive compared with Vec, so I think it may be worth to waste few memory to get better performance?
| /// - [`EmitTo::NextBlock`]: pops a single block. | ||
| pub fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<T>> { | ||
| match emit_to { | ||
| EmitTo::All => self.inner.pop_block().ok_or_else(|| { |
There was a problem hiding this comment.
Maybe a bit confusing now due to:
EmitTo::FirstandEmitTo::Allonly meaningful inflat modeEmitTo::NextBlockonly meaningful inblocked mode
Here (vec_block_store.rs) maybe more suitable to return error when encountering EmitTo::First and EmitTo::All?
| return None; | ||
| } | ||
|
|
||
| let block = mem::take(&mut self.inner[self.cursor]); |
There was a problem hiding this comment.
VecDeque is used at the beginning, but the index op in VecDeque actually not really trivial as Vec, and switch to Vec later suggested by reviews.
But I agree with that VecDeque may be more suitable, due to:
- its main target is saving memory (although Vec will not occupy much)
- the logic can become simple, and I think as a initial pr, it should keep as simple as possible.
| /// [`GroupsAccumulator::supports_blocked_groups`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator::supports_blocked_groups | ||
| /// | ||
| #[derive(Debug)] | ||
| pub struct BlockedBlockStore<B: Block> { |
There was a problem hiding this comment.
Methods in Block trait is mainly abstracted from logic in prim_op.rs, and may only useful for accumulators.
Maybe we should define some trait different for groups case?
For me, how to define the suitable traits is actually a main block point...
To quantify the implementation effort for options 1 and 3: we would need to migrate all implementations of
Maybe it's a good idea to experimenting the hardest one first 🤔 |
Which issue does this PR close?
Rationale for this change
As mentioned in #7065 , we use a single
Vecto manageaggregation intermediate resultsboth inGroupAccumulatorandGroupValues.It is simple but not efficient enough in high-cardinality aggregation, because when
Vecis not large enough, we need to allocate a newVecand copy all data from the old one.So this pr introduces a
blocked approachto manage theaggregation intermediate results. We will never resize theVecin the approach, and instead we split the data to blocks, when the capacity is not enough, we just allocate a new block. Detail can see #7065What changes are included in this PR?
PrimitiveGroupsAccumulatorandGroupValuesPrimitiveas the exampleAre these changes tested?
Test by exist tests. And new unit tests, new fuzzy tests.
Are there any user-facing changes?
Two functions are added to
GroupValuesandGroupAccumulatortrait.But as you can see, there are default implementations for them, and users can choose to really support the blocked approach when wanting a better performance for their
udafs.