rluvaton commented on code in PR #25877: URL: https://github.com/apache/datafusion/pull/25877#discussion_r4177014748
########## datafusion/physical-plan/src/aggregates/group_values/blocked/primitive.rs: ########## @@ -0,0 +1,412 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, + cast::AsArray, +}; +use arrow::datatypes::DataType; +use datafusion_common::Result; +use datafusion_common::hash_utils::RandomState; +use datafusion_expr::{BlockedEmitTo, BlocksIndex}; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::BlockedVec; +use hashbrown::hash_table::HashTable; + +use super::BlockedGroupValues; +use crate::aggregates::group_values::HashValue; + +/// A [`BlockedGroupValues`] storing a single column of primitive values. +/// +/// Like `GroupValuesPrimitive`, but the values are stored in a [`BlockedVec`] +/// and the hash table entries hold `(group_index, value)` instead of +/// `(group_index, hash)`, so probing never reads the (blocked) values. +pub(crate) struct BlockedGroupValuesPrimitive<T: ArrowPrimitiveType> { + /// The data type of the output array + data_type: DataType, + /// Stores `(group_index, value)` based on the hash of the value + /// + /// Storing the value instead of its hash means comparing keys never goes + /// through `values`, and rehashing is cheap for primitive values. + map: HashTable<(BlocksIndex, T::Native)>, + /// The group index of the null value if any + null_group: Option<BlocksIndex>, + /// The values for each group index + values: BlockedVec<T::Native>, + /// The random state used to generate hashes + random_state: RandomState, +} + +impl<T: ArrowPrimitiveType> BlockedGroupValuesPrimitive<T> { + pub(crate) fn new(data_type: DataType, block_size: usize) -> Self { + assert!(PrimitiveArray::<T>::is_compatible(&data_type)); + Self { + data_type, + map: HashTable::with_capacity(128), + values: BlockedVec::new(block_size), + null_group: None, + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + } + } + + /// Builds the output array of one block whose first group is `start` + fn build_block(&self, values: Vec<T::Native>, start: usize) -> ArrayRef { + let len = values.len(); + let block_size = self.values.block_size(); + let nulls = self + .null_group + .map(|null_group| null_group.flat(block_size)) + .filter(|null_group| (start..start + len).contains(null_group)) + .map(|null_group| { + let null_idx = null_group - start; + let mut buffer = NullBufferBuilder::new(len); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(len - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }); + Arc::new( + PrimitiveArray::<T>::new(values.into(), nulls) + .with_data_type(self.data_type.clone()), + ) + } + + /// Updates the null group after the first `n` groups were emitted + fn shift_null_group(&mut self, n: usize) { + let block_size = self.values.block_size(); + self.null_group = self + .null_group + .and_then(|null_group| null_group.flat(block_size).checked_sub(n)) + .map(|flat| BlocksIndex::from_flat(flat, block_size)); + } +} + +impl<T: ArrowPrimitiveType> BlockedGroupValues for BlockedGroupValuesPrimitive<T> +where + T::Native: HashValue, +{ + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<BlocksIndex>) -> Result<()> { + assert_eq!(cols.len(), 1); + groups.clear(); + + for v in cols[0].as_primitive::<T>() { + let group_id = match v { + None => *self + .null_group + .get_or_insert_with(|| self.values.push(Default::default())), + Some(key) => { + // Fold equivalence-class duplicates (e.g. `-0.0` → `+0.0`) + // so the bit-equal `is_eq` matches and the stored value is + // the canonical representative. + let key = key.canonicalize(); + let state = &self.random_state; + let hash = key.hash(state); + let insert = self.map.entry( + hash, + |&(_, k)| k.is_eq(key), + |&(_, k)| k.hash(state), + ); + + match insert { + hashbrown::hash_table::Entry::Occupied(o) => o.get().0, + hashbrown::hash_table::Entry::Vacant(v) => { + let g = self.values.push(key); + v.insert((g, key)); + g + } + } + } + }; + groups.push(group_id) + } + Ok(()) + } + + fn size(&self) -> usize { + self.map.capacity() * size_of::<(BlocksIndex, T::Native)>() + + self.values.allocated_size() + } + + fn is_empty(&self) -> bool { + self.values.is_empty() + } + + fn len(&self) -> usize { + self.values.len() + } + + fn emit(&mut self, emit_to: BlockedEmitTo) -> Result<Vec<Vec<ArrayRef>>> { + let blocks = match emit_to { + BlockedEmitTo::All => { + self.map.clear(); + let mut start = 0; + let blocks = self + .values + .take_all() + .into_iter() + .map(|block| { + let len = block.len(); + let array = self.build_block(block, start); + start += len; + vec![array] + }) + .collect(); + self.null_group = None; + blocks + } + BlockedEmitTo::NextBlock => { Review Comment: But you are right, just not in this PR -- 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]
