Repository navigation
Conversation
evaluate_to_columns wrapped all evaluated blocks as a single entry, producing batches with one column per block once there was more than one block.
|
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 |
There was a problem hiding this comment.
Ignore the implementation itself, but just look at the API itself, this can be changed to change into blocks or flat depending on some threshold
jayzhan211
left a comment
There was a problem hiding this comment.
I'd suggest reversing the order: first get an implementation that measurably fixes the problem in #24704, then shape the API around it. An API designed before we know what the fast, memory-efficient implementation needs tends to get in its way, and it's much harder to change once it's public.
The current draft shows the risk. Every aggregation now goes through the blocked traits, and most of it runs through adapters over flat state, which is slower than main and uses more memory:
group_values/blocked.rs:195andblocked_groups_accumulator.rs:312: emit All through the adapters is O(G²/B). With datafusion-cli, 1 partition,GROUP BY concat('k', v)plussum(v)at 4M groups takes 0.21 s on main and 30.2 s here. A primitive key withsumis 0.1 s vs 0.4 s.aggregates/spill.rs:220: one spill file per block.aggregate_memory_spill.sltseesspill_countgo from 7 to about 800, and thecount(DISTINCT)cases at L95/L106 now fail withResourcesExhausted.
Proposal:
- Pick one target from #24704, e.g. peak memory / spill-free limit and emit time for high-cardinality
GROUP BY <primitive>withsum/count/min/max. - Implement blocked storage natively for just that case, behind a gate (blocked only when the group values and all accumulators support it, flat path unchanged otherwise), with no adapters.
- Show before/after numbers: peak memory, spill count, and ClickBench/TPC-H timings with no regressions elsewhere.
- Then extract the API from what that implementation actually needed, and extend it to more accumulators and group-value types.
That keeps each step reviewable and makes sure the API we commit to is one that performs well.
| is_reversed: bool, | ||
| input_fields: Vec<FieldRef>, | ||
| is_nullable: bool, | ||
| batch_size: Option<usize>, |
There was a problem hiding this comment.
The block size shouldn't be set at plan time on AggregateFunctionExpr. The table uses the runtime context.session_config().batch_size(), so a plan executed with a different TaskContext fails at the assert in aggregate_hash_table/common.rs:100 with Internal("... left: 4096, right: 8192: Block size mismatch ..."). Suggest dropping the field and the builder method, and having the caller pass the runtime block size in:
pub fn blocked_groups_accumulator_supported(&self, block_size: usize) -> bool {pub fn create_blocked_groups_accumulator(
&self,
block_size: usize,
) -> Result<Box<dyn BlockedGroupsAccumulator>> {| } | ||
|
|
||
| pub trait BlockedGroupsAccumulator: Send + std::any::Any { | ||
| fn batch_size(&self) -> usize; |
There was a problem hiding this comment.
nit: this is named batch_size(), but everything else (BlockedVec::block_size, BlockedGroupValues::block_size) says block size. Suggest fn block_size(&self) -> usize;, and also renaming BlockedAccumulatorArgs::batch_size.
| /// Create an accumulator for `agg_expr` -- a [`BlockedGroupsAccumulator`] if | ||
| /// that is supported by the aggregate, or a | ||
| /// [`BlockedGroupsAccumulatorAdapter`] if not. | ||
| pub(in crate::aggregates) fn create_blocked_group_accumulator( |
There was a problem hiding this comment.
Should the blocked API be opt-in rather than the route for every aggregation? A flat GroupsAccumulator/GroupValues can't serve BlockedEmitTo cheaply: the adapters loop EmitTo::First(block_size), and each call shifts the whole table. Suggest taking the blocked path only when the group values and every accumulator support it natively, and keeping the flat path unchanged otherwise. That also keeps the adapters out of the public shape.
There was a problem hiding this comment.
this mean I have to duplicate the hash aggregate streams code and I wanted to avoid this
GroupsAccumulator should not be used and should be removed in the end, we will not release until all GroupValues has been converted
| /// Requirements: | ||
| /// 1. `n` must be smaller than block_size | ||
| /// 2. `n` is not 0 | ||
| First(usize), |
There was a problem hiding this comment.
The First(n) precondition n < block_size leaks to callers: common_ordered.rs:481 already rebuilds EmitTo::First(n) as NextBlocks plus a First remainder. Either allow any n, or define that decomposition once, next to the enum:
impl BlockedEmitTo {
/// Splits a flat `EmitTo` into block-sized emits.
pub fn from_emit_to(emit_to: EmitTo, block_size: usize) -> Vec<Self> {
match emit_to {
EmitTo::All => vec![Self::All],
EmitTo::First(n) => {
let mut emits = vec![Self::NextBlock; n / block_size];
if n % block_size != 0 {
emits.push(Self::First(n % block_size));
}
emits
}
}
}
}| /// - `BlockedEmitTo::NextBlock` it should return single item vector with the block or empty vec in case of no blocks | ||
| /// - `BlockedEmitTo::First(n)` it should return single item vector with the first n rows in the first block. n must be smaller than block size and length | ||
| /// | ||
| fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result<Vec<ArrayRef>>; |
There was a problem hiding this comment.
Return shapes:
NextBlockyields 0 or 1 blocks, so returning aVechides that;Option<ArrayRef>(or a dedicatedevaluate_next_block) would say so.evaluate_preserving(L348) andstate_preservingreturn one flat array rather than blocks, which is inconsistent withevaluateandstate.- If the plan is an iterator for emit All (the TODO in
group_values/blocked.rs), it's cheaper to put it in the trait now than to change it later.
There was a problem hiding this comment.
evaluate_preserving is like that since it provide the indices, I can split the indices by block size and return Vec
Return shapes:
NextBlockyields 0 or 1 blocks, so returning aVechides that;Option<ArrayRef>(or a dedicatedevaluate_next_block) would say so.
I thought about it originally but in order to not explode the code with functions to implement I thought this will be a better way, what do you think?
evaluate_preserving(L348) andstate_preservingreturn one flat array rather than blocks, which is inconsistent withevaluateandstate.
evaluate_preserving/state_preserving is like that since it provide the indices, I can split the indices by block size and return Vec if that is what you mean
- If the plan is an iterator for emit All (the TODO in
group_values/blocked.rs), it's cheaper to put it in the trait now than to change it later.
the iterator that emit all is just an idea, this will be useful when materializing will use more memory than needed
There was a problem hiding this comment.
having iterator create challenges regarding memory size
| } | ||
| } | ||
|
|
||
| pub trait BlockedGroupsAccumulator: Send + std::any::Any { |
There was a problem hiding this comment.
This trait repeats GroupsAccumulator method for method, and BlockedGroupSelection (L37) repeats GroupSelection; only the index type changes. Could the two share a generic index type, or could the blocked methods be defaulted on GroupsAccumulator? That would avoid two parallel public traits that have to stay in sync.
There was a problem hiding this comment.
BlockedGroupsAccumulator would have to extend GroupsAccumulator and I don't want that since every function call now will be confusion to what is being called and also the plan is to remove GroupsAccumulator from the way I see it and not keep it, so adding dependencies to old implementation would just couple more the implementations
| /// margin, `index_in_block < block_size` (the batch size) and the number of blocks is | ||
| /// `groups / block_size`. The API keeps taking and returning `usize`. | ||
| #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)] | ||
| pub struct BlocksIndex { |
There was a problem hiding this comment.
BlocksIndex puts a lot of public surface in expr-common: sub_flat, sub_flat_checked (L199), gte_flat, prev_block, add, increment. Only group-values implementations need these. Suggest keeping the public type to new, block_index, index_in_block and the flat conversions, and moving the rest next to their only user, or marking them #[doc(hidden)].
Yes, but this is to make it easier to review, I can copy the entire aggregation code and replace there with the blocked but it is harder to review
I've did not implement in this PR the bytes/bytes view group by, but it is implemented in later pr
I'm aware of that and I added a comment and fixed that problem in later PR by adding
Yes, but I've done this way to make it easier to review, I can copy the entire aggregation code and replace there with the blocked but it is harder to review since you have no clear way to see what I actually changed
This pr is after I've done all of this (but without the flat approach), the pr that implement it entirely is:
|
|
@rluvaton Do you think it's possible to break #24928 into several PRs, where each one includes not just the API but also part of the implementation? Ideally each PR would be safe to merge into main on its own with no regression, assuming the follow-ups might not land for several releases. If that's not workable, I think we'd need a feature branch for this. On a feature branch we could merge the API first and then the rest incrementally. |
|
@jayzhan211 Ok, so instead of having an adapter and harm performance I created another PR with alternative approach which support both and have less performance hit for unsupported cases |
Created alternative PR in:
outlined the difference between those 2 there
Which issue does this PR close?
Part of:
Rationale for this change
See issue
What changes are included in this PR?
It contain
BlockedGroupsAccumulatortrait, helperBlockedVec, supportcountin blocked so you see the example usage, change the entire aggregate to work with blocked (this is possible due to the added adapters), implement BlockedGroupValues for primitive so you will see how it is being usedWhat is the testing strategy for this PR?
Existing tests
Are there any user-facing changes?
yes