rluvaton commented on issue #24704: URL: https://github.com/apache/datafusion/issues/24704#issuecomment-5665631916
**Recap over the blocked implementation I implemented and experimented with.** **Once we agree on the design we can continue with outlining the plan for migration, impl, etc.** # Benchmark results ## implementation with blocked vec backed by mmap: (results when having mmap implementation for blocked vec over stack allocated items, and at the time of writing not the latest commit implementation) - external_aggr which tracks memory usage and perf: https://github.com/apache/datafusion/pull/24928#issuecomment-5603717435 - h2o_medium benchmarks: https://github.com/apache/datafusion/pull/24928#issuecomment-5603527366 - clickbench_partitioned: https://github.com/apache/datafusion/pull/24928#issuecomment-5603510552 - tpcds_sf1 - https://github.com/apache/datafusion/pull/24928#issuecomment-5603469619 - tpcds_sf10 - https://github.com/apache/datafusion/pull/24928#issuecomment-5603464928 - tpch_sf1 - https://github.com/apache/datafusion/pull/24928#issuecomment-5603444733 the performance is not as good as the current implementation on main in some cases but not that bad (it is faster in some cases). # Goals 1. Make the migrations for the users as easy as possible 2. Have close enough performance to the current implementation 3. Try to avoid breaking changes 4. Make the implementation for blocked aggregate simple # Findings for the Blocked aggregate impl I've implemented full blocked approach to see: - every gotchas before our users got to encounter - every design flaw that we want to avoid before commiting to the design - every performance issue from this implementation Below are my finding: ## We cannot use the existing `GroupsAccumulator` trait Given that we still want to have the `EmitTo::All` we can't use the existing trait since emitting all in blocked aggregate Should output Vec of Blocks rather than 1 very large block. emit all is used when spilling or when doing early emit (like partial passthrough), and having it produce 1 large block would force a copy even when blocks are needed. ## Having a temporary blocked aggregation stream will result in a huge code Unlike the recent split of the aggregate streams, having a temporary blocked aggregate streams will result in a huge extra code that we have to sync until fully migrating. we can avoid this by having adapters to the user public API (like `GroupsAccumulator` and `Accumulator`) but know that the groups accumulator adapter have a significant performance hit. and replacing the implementation of the stream in a single pr, the amount of change needed to the streams themselves are not big at all. ## Having `EmitTo::First(n)` for the `BlockedGroupsAccumulator` is recommended while complicating the implementation in the helpers, this allows to have an adapter between the existing `GroupsAccumulator` and the `BlockedGroupsAccumulator`. ## Having an abstraction around `BlocksIndex` Having the abstraction around `BlocksIndex` allows us to change between implementation without forcing us onto specific implementation. this is important since if we have the `mmap` implementation for `BlockedVec` having the flat index (like treating `BlockedVec` as a flat Vec) would be more performant versus having 2 fields (`block_index` and `index_in_block`) or having it in a single `u64` and require shifting and masking (which would force us to have batch size as power of 2 which is not the requirement in the conf) ## Most of the blocked aggregate expressions and group by types need a small number of building blocks After implementing a lot of aggregate expression, both in this repo and in my company code I came to the conclusion that you can implement most of them without having to deal with manual blocking and others gotchas and those are: 1. Blocked `Vec`: 1. stack allocated items 2. heap allocated items - this is a pain with allocated size tracking 2. Blocked `BooleanBuilder` 3. Blocked `NullBufferBuilder` 4. Blocked `OffsetBufferBuilder` (this is separate from `BlockedVec` since having `take_n` requires shifting the offsets) 5. Blocked byte buffer (unliked the regular blocked vec, this does not manage the blocks itself, but the users are - this is for byte and view byte array types) 6. Blocked `Rows` Optional building blocks that abstract away the block management so eac - `BlockedCustomHeapAllocatedInputBuilder` that the Blocked vec over heap allocated items can use it behind the scenes (which is what I did) this should not force the type to be allowed to create on it's own (for example using the wrapper to implement blocked arrow `Rows`, which you can't initialize without the row converter) ## Tracking memory in blocked helpers can be a pain for heap allocated items If you implement `allocated_size` function for the blocked helpers (which you should have since you don't expose the blocks themselves) for helpers that deal with items on the stack like `BlockedOffsetBufferBuilder` or `BlockedVec` for stack allocated items, this is easy to implement and will be materialized to avoid iterating over all the blocks each time. however for heap allocated items, this is annoying, since you need to allow the user to specify a way to get the allocated size for each item (which is easy). but you can't allow to have `IndexMut` implemented since it might change the size of the touched item ## Don't force blocked helpers to pre-allocate `block_size` this should keep the decision to the user like before, since it might be expensive and not wanted due to memory bloat ## Supported nested in multi group by force all the group column impl in the multi group by have block size optionally being dynamic and controlled from outside For example, grouping by `string` and `list<int>` requires the `int` inside the list to be laid out as blocked but not fixed one since each list is different size. I've implemented the `GroupValues` trait with `const IS_FIXED_BLOCK_SIZE: bool` generic which pass to the blocked building blocks helpers (like `BlockedVec`) so the performance of calculating blocking or not would not be affected for the other case when it is false the helpers that all the implementation are using are not handling pushing between blocks this is proven to be easy to do and quite simple ## Having adapter to `GroupsAccumulator` is easy and possible but have high overhead Having an adapter will allow for gradual implementation from the users and us into the `BlockedGroupsAccumulator` however this mean that every emit_all will require multiple calls to the underlying `GroupsAccumulator` with `EmitTo::First(batch_size)`, causing shifting ## Having nested Vec harm performance due to the extra indirection and add cost for the low cardinality case Implementation of `BlockedVec` that implement it using `Vec<Vec<T>>`/`VecDeque<Vec<T>>` force index calculation (in case you store it in a single `u64` and not 2 fields) reduce cache locality and hurt overall performance. instead we can implement it using `mmap` which complicate the implementation but will allow for better performance and less indirection since we will treat it as single flat vec. this must be behind feature flag that will be disabled for wasm since it is not available there, and we should provide an API to update custom user allocators that track the allocations like: 1. OOM guard (that was presented in one of the meetups) 2. tracking allocator as form of testing (like we have in our company code in our memory tests) ## Having to spill force a lot of spill files/expensive copy in case we emit in blocks and we don't have sort across multiple blocks Currently, when we have memory pressure we call `EmitTo::All` and we sort a huge batch and then spill it in parts, this is good, however in the new API, emit_all should returns blocks and not a single huge one like explained above, this means that if we keep the existing code that for every block it sort and spill it independently, it will result in a lot of spill files, which will result in a really large performance hit (several benchmarks were affected by this) doing concat on the blocks and going to the existing impl defeat the purpose of having blocked impl, and also add performance penalty as well. instead what I did is supporting in arrow sort (I've coping the original arrow sort code to my branch) and allow to sort across multiple blocks while keeping them separated. this reduce the amount of spill files and avoid the copy. ## `GroupValues` being public add another breaking change vector Because `GroupValues` is public we would introduce a breaking change when changing it into blocked api. # Traits `BlockedGroupsAccumulator`: ```rust pub trait BlockedGroupsAccumulator: Send + std::any::Any { fn batch_size(&self) -> usize; fn update_batch( &mut self, values: &[ArrayRef], // Get blocks index rather than flat index to avoid extra computation when inserting and searching (BlocksIndex can be implemented as flat usize though) group_indices: &[BlocksIndex], opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()>; /// Same as `GroupsAccumulator::evaluate` but returning vec of blocks and different EmitTo /// For /// - `BlockedEmitTo::All` it should return `Vec<Block>` /// - `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>>; // Same as `GroupsAccumulator::evaluate_preserving` but with blocked indices fn evaluate_preserving( &mut self, _selection: BlockedGroupSelection<'_>, ) -> Result<ArrayRef> { not_impl_err!("Preserving grouped evaluation is not implemented") } /// Returns `true` if [`Self::evaluate_preserving`] is implemented. fn supports_evaluate_preserving(&self) -> bool { false } /// Same as `GroupsAccumulator::state` but returning vec of blocks and different EmitTo /// For /// - `BlockedEmitTo::All` it should return `Vec<Block>` /// - `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 state(&mut self, emit_to: BlockedEmitTo) -> Result<Vec<Vec<ArrayRef>>>; /// Same as `GroupsAccumulator::state_preserving` but with blocked indices fn state_preserving( &mut self, _selection: BlockedGroupSelection<'_>, ) -> Result<Vec<ArrayRef>> { not_impl_err!("Preserving grouped state is not implemented") } /// Returns `true` if [`Self::state_preserving`] is implemented. fn supports_state_preserving(&self) -> bool { false } fn merge_batch( &mut self, values: &[ArrayRef], // Get blocks index rather than flat index to avoid extra computation when inserting and searching (BlocksIndex can be implemented as flat usize though) group_indices: &[BlocksIndex], total_num_groups: usize, ) -> Result<()>; /// Same as `GroupsAccumulator::convert_to_state` fn convert_to_state( &self, values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> Result<Vec<ArrayRef>>; /// Same as `GroupsAccumulator::size` fn size(&self) -> usize; } ``` `BlockedGroupColumn`: ```rust pub trait BlockedGroupColumn<const IS_FIXED_BLOCK_SIZE: bool>: Send + Sync { fn batch_size(&self) -> usize; /// Same as `GroupColumn::equal_to` but with BlocksIndex rather than flat index /// Get blocks index rather than flat index to avoid extra computation when inserting and searching (BlocksIndex can be implemented as flat usize though) fn equal_to(&self, lhs_row: BlocksIndex, array: &ArrayRef, rhs_row: usize) -> bool; /// Appends the row at `row` in `array` to this builder fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()>; /// Same as `GroupColumn::vectorized_equal_to`, but with BlocksIndex rather than flat index fn vectorized_equal_to( &self, // Get blocks index rather than flat index to avoid extra computation when inserting and searching (BlocksIndex can be implemented as flat usize though) lhs_rows: &[BlocksIndex], array: &ArrayRef, rhs_rows: &[usize], equal_to_results: &mut BooleanBufferBuilder, ); /// The vectorized version `append_val` fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()>; /// Returns the number of rows stored in this builder fn len(&self) -> usize; /// true if len == 0 fn is_empty(&self) -> bool { self.len() == 0 } /// Returns the number of bytes used by this [`BlockedGroupColumn`] fn size(&self) -> usize; /// Same as `GroupColumn::values_preserving` but with BlocksIndex rather than flat index in the group selection fn values_preserving(&self, selection: BlockedGroupSelection<'_>) -> Result<ArrayRef>; /// Builds a new array from the first `n` stored rows, shifting the /// remaining rows to the start of the builder /// /// n must be smaller than block size and len, and greater than 0 /// /// the `adjusted_block_size_iter` argument is for nested implementation and how many items in each block should exists after the take. fn take_n(&mut self, n: usize, // TODO - need to figure out how to do that in rust since Clone requires `Sized` but we want to be able to clone the // iterator as we need to have the same block sizing for both the values and the nulls for some implementations adjusted_block_size_iter: Option<Box<dyn Iterator<Item = usize> + Clone>>, ) -> ArrayRef; /// Take next block, if no blocks return `None` fn take_next_block(&mut self) -> Option<ArrayRef>; /// Builds a new array from all of the stored rows /// /// returns blocks of array, each block is of size `block_size` except the last one fn take_all(self: Box<Self>) -> Vec<ArrayRef>; /// When the block sizing is externally managed (when `IS_FIXED_BLOCK_SIZE` is `false`) this call signals that a new block has started. fn start_new_block(&mut self); } ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
