ariel-miculas commented on code in PR #22481: URL: https://github.com/apache/datafusion/pull/22481#discussion_r3302380666
########## datafusion/physical-plan/src/aggregates/group_values/single_group_by/flat_primitive.rs: ########## @@ -0,0 +1,292 @@ +// 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 crate::aggregates::group_values::GroupValues; +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, cast::AsArray, +}; +use arrow::datatypes::DataType; +use datafusion_common::Result; +use datafusion_expr::EmitTo; +use num_traits::AsPrimitive; +use std::mem::size_of; +use std::sync::Arc; + +/// Sentinel value indicating an unoccupied slot in the flat array. +const EMPTY: u32 = u32::MAX; + +/// A [`GroupValues`] implementation using direct array indexing for integer-typed +/// GROUP BY columns with a known, bounded value range. +/// +/// Instead of hashing each key and probing a hash table, this computes +/// `index = key_as_u64.wrapping_sub(offset)` to directly index into a flat array, +/// yielding O(1) group lookups with no hashing, no collisions, and excellent +/// cache locality for small ranges. +/// +/// Inspired by the `ArrayMap` used for perfect hash joins (see `joins/array_map.rs`). +pub struct GroupValuesFlatPrimitive<T: ArrowPrimitiveType> +where + T::Native: AsPrimitive<u64>, +{ + data_type: DataType, + /// Maps `key - offset` -> group_id (u32). EMPTY means no group assigned yet. + slots: Vec<u32>, + /// The minimum value (as u64 via two's complement) used as the base offset. + offset: u64, + /// The group index assigned to NULL values, if any. + null_group: Option<usize>, + /// Ordered group values for emit. `values[group_id]` is the original key. + values: Vec<T::Native>, +} + +impl<T: ArrowPrimitiveType> GroupValuesFlatPrimitive<T> +where + T::Native: AsPrimitive<u64>, +{ + /// Creates a new flat-indexed GroupValues. + /// + /// `min_val` and `max_val` define the key range (as u64 via two's complement cast). + /// The allocated flat array has `max_val - min_val + 1` slots. + pub fn new(data_type: DataType, min_val: u64, max_val: u64) -> Self { + assert!(PrimitiveArray::<T>::is_compatible(&data_type)); + let range = max_val.wrapping_sub(min_val); + let size = (range as usize) + 1; + Self { + data_type, + slots: vec![EMPTY; size], + offset: min_val, + null_group: None, + values: Vec::new(), + } + } + + /// Returns the range (number of slots) of this flat map. + pub fn range(&self) -> usize { + self.slots.len() + } +} + +impl<T: ArrowPrimitiveType> GroupValues for GroupValuesFlatPrimitive<T> +where + T::Native: AsPrimitive<u64> + Default + Copy + Send, +{ + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + assert_eq!(cols.len(), 1); + groups.clear(); + + let arr = cols[0].as_primitive::<T>(); + + if arr.null_count() == 0 { + for i in 0..arr.len() { + // SAFETY: null_count == 0 guarantees all values are valid + let key: u64 = unsafe { arr.value_unchecked(i) }.as_(); + let idx = key.wrapping_sub(self.offset) as usize; + debug_assert!(idx < self.slots.len()); + + let slot = unsafe { self.slots.get_unchecked_mut(idx) }; + let group_id = if *slot == EMPTY { + let g = self.values.len() as u32; + *slot = g; + self.values.push(unsafe { arr.value_unchecked(i) }); + g as usize + } else { + *slot as usize + }; + groups.push(group_id); + } + } else { + for i in 0..arr.len() { + let group_id = if arr.is_null(i) { + *self.null_group.get_or_insert_with(|| { + let g = self.values.len(); + self.values.push(Default::default()); + g + }) + } else { + let key: u64 = unsafe { arr.value_unchecked(i) }.as_(); + let idx = key.wrapping_sub(self.offset) as usize; + debug_assert!(idx < self.slots.len()); + + let slot = unsafe { self.slots.get_unchecked_mut(idx) }; + if *slot == EMPTY { + let g = self.values.len() as u32; + *slot = g; + self.values.push(unsafe { arr.value_unchecked(i) }); + g as usize + } else { + *slot as usize + } + }; + groups.push(group_id); + } + } + Ok(()) + } + + fn size(&self) -> usize { + self.slots.len() * size_of::<u32>() + + self.values.capacity() * size_of::<T::Native>() + } + + fn is_empty(&self) -> bool { + self.values.is_empty() + } + + fn len(&self) -> usize { + self.values.len() + } + + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + let array: PrimitiveArray<T> = match emit_to { + EmitTo::All => { + self.slots.fill(EMPTY); + let null_idx = self.null_group.take(); + let values = std::mem::take(&mut self.values); + build_primitive::<T>(values, null_idx) + } + EmitTo::First(n) => { + // Shift all slot references down by n, remove those < n + for slot in self.slots.iter_mut() { + if *slot != EMPTY { + let g = *slot as usize; + if g < n { + *slot = EMPTY; + } else { + *slot = (g - n) as u32; + } + } + } + let null_group = match &mut self.null_group { + Some(v) if *v >= n => { + *v -= n; + None + } + Some(_) => self.null_group.take(), + None => None, + }; + let mut split = self.values.split_off(n); + std::mem::swap(&mut self.values, &mut split); Review Comment: I'm trying to replace this pattern (split_off + swap) with `split_vec_min_alloc` https://github.com/apache/datafusion/pull/22416 to avoid unnecessary allocations -- 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]
