kumarUjjawal commented on code in PR #23187: URL: https://github.com/apache/datafusion/pull/23187#discussion_r3688362688
########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs: ########## @@ -0,0 +1,910 @@ +// 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::multi_group_by::GroupColumn; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, Int64Array, + PrimitiveArray, +}; +use arrow::compute::take; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field}; +use arrow::error::ArrowError; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::{DataFusionError, Result, exec_err}; +use datafusion_execution::memory_pool::proxy::HashTableAllocExt; +use hashbrown::hash_table::HashTable; +use std::marker::PhantomData; +use std::mem::size_of; +use std::sync::Arc; + +use crate::aggregates::AGGREGATION_HASH_SEED; + +/// [`GroupColumn`] for dictionary-encoded columns with key type `K`. +/// +/// `inner` holds one slot per distinct value seen across all batches. +/// `group_to_inner[group_idx]` maps each group to its slot in `inner`, +/// so groups with the same value share a slot rather than duplicating data. +pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> { + /// Deduplicated store of distinct values. + inner: Box<dyn GroupColumn>, + /// Unary null array (length 1) reused for every null appended to `inner`. + null_array: ArrayRef, + /// Maps each group index to its slot in `inner`. + group_to_inner: Vec<usize>, + /// Lookup table mapping `(value_hash, inner_slot)` for each non-null distinct value. + value_dedup: HashTable<(u64, usize)>, + /// Tracked allocation size of `value_dedup` for memory accounting via `size()`. + value_dedup_size: usize, + /// Slot in `inner` for the null group; `None` until the first null is seen. + null_inner_slot: Option<usize>, + /// Hash seed — must match `create_hashes` so hashes are consistent across calls. + random_state: RandomState, + /// Reusable scratch buffer mapping `val_idx → inner_slot` across batches. + val_to_inner: Vec<usize>, + /// Reusable hash buffer for the dictionary values array. + val_hashes: Vec<u64>, + _phantom: PhantomData<K>, +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> { + pub fn new(inner: Box<dyn GroupColumn>, field: &Field) -> Self { + let null_array = arrow::array::new_null_array(field.data_type(), 1); + Self { + inner, + null_array, + group_to_inner: Vec::new(), + value_dedup: HashTable::new(), + value_dedup_size: 0, + null_inner_slot: None, + random_state: AGGREGATION_HASH_SEED, + val_to_inner: Vec::default(), + val_hashes: Vec::default(), + _phantom: PhantomData, + } + } + + fn into_dict(values: ArrayRef, group_to_inner: &[usize]) -> ArrayRef { + let keys: PrimitiveArray<K> = group_to_inner + .iter() + .map(|&slot| { + if values.is_null(slot) { + None + } else { + Some(K::Native::usize_as(slot)) + } + }) + .collect(); + Arc::new(DictionaryArray::<K>::new(keys, values)) + } + + // https://github.com/apache/datafusion/issues/23127 + // Null groups emit a null key, not a key index into the values array, so the + // null inner slot does not consume a key index. + fn check_key_overflow(&self) -> Result<()> { + // Keys are raw slot indices. The null slot is excluded from key count + // only when it occupies the last position — any non-null slot above it + // still emits that slot's raw index as a key. + let inner_len = self.inner.len(); Review Comment: Null first plus the maximum non-null values fails, while the same values with null last succeeds. ########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs: ########## @@ -0,0 +1,910 @@ +// 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::multi_group_by::GroupColumn; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, Int64Array, + PrimitiveArray, +}; +use arrow::compute::take; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field}; +use arrow::error::ArrowError; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::{DataFusionError, Result, exec_err}; +use datafusion_execution::memory_pool::proxy::HashTableAllocExt; +use hashbrown::hash_table::HashTable; +use std::marker::PhantomData; +use std::mem::size_of; +use std::sync::Arc; + +use crate::aggregates::AGGREGATION_HASH_SEED; + +/// [`GroupColumn`] for dictionary-encoded columns with key type `K`. +/// +/// `inner` holds one slot per distinct value seen across all batches. +/// `group_to_inner[group_idx]` maps each group to its slot in `inner`, +/// so groups with the same value share a slot rather than duplicating data. +pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> { + /// Deduplicated store of distinct values. + inner: Box<dyn GroupColumn>, + /// Unary null array (length 1) reused for every null appended to `inner`. + null_array: ArrayRef, + /// Maps each group index to its slot in `inner`. + group_to_inner: Vec<usize>, + /// Lookup table mapping `(value_hash, inner_slot)` for each non-null distinct value. + value_dedup: HashTable<(u64, usize)>, + /// Tracked allocation size of `value_dedup` for memory accounting via `size()`. + value_dedup_size: usize, + /// Slot in `inner` for the null group; `None` until the first null is seen. + null_inner_slot: Option<usize>, + /// Hash seed — must match `create_hashes` so hashes are consistent across calls. + random_state: RandomState, + /// Reusable scratch buffer mapping `val_idx → inner_slot` across batches. + val_to_inner: Vec<usize>, + /// Reusable hash buffer for the dictionary values array. + val_hashes: Vec<u64>, + _phantom: PhantomData<K>, +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> { + pub fn new(inner: Box<dyn GroupColumn>, field: &Field) -> Self { + let null_array = arrow::array::new_null_array(field.data_type(), 1); + Self { + inner, + null_array, + group_to_inner: Vec::new(), + value_dedup: HashTable::new(), + value_dedup_size: 0, + null_inner_slot: None, + random_state: AGGREGATION_HASH_SEED, + val_to_inner: Vec::default(), + val_hashes: Vec::default(), + _phantom: PhantomData, + } + } + + fn into_dict(values: ArrayRef, group_to_inner: &[usize]) -> ArrayRef { + let keys: PrimitiveArray<K> = group_to_inner + .iter() + .map(|&slot| { + if values.is_null(slot) { + None + } else { + Some(K::Native::usize_as(slot)) + } + }) + .collect(); + Arc::new(DictionaryArray::<K>::new(keys, values)) + } + + // https://github.com/apache/datafusion/issues/23127 + // Null groups emit a null key, not a key index into the values array, so the + // null inner slot does not consume a key index. + fn check_key_overflow(&self) -> Result<()> { + // Keys are raw slot indices. The null slot is excluded from key count + // only when it occupies the last position — any non-null slot above it + // still emits that slot's raw index as a key. + let inner_len = self.inner.len(); + let null_slot_excluded = self.null_inner_slot.is_some_and(|s| s + 1 == inner_len); + let max_key_count = inner_len - null_slot_excluded as usize; + if !Self::key_type_fits(max_key_count) { + let non_null_slots = inner_len - self.null_inner_slot.is_some() as usize; + return exec_err!( + "Dictionary key type {:?} cannot represent {} distinct values", + K::DATA_TYPE, + non_null_slots + ); + } + Ok(()) + } + + fn key_type_fits(num_values: usize) -> bool { + let max: usize = match K::DATA_TYPE { + DataType::Int8 => i8::MAX as usize, + DataType::Int16 => i16::MAX as usize, + DataType::Int32 => i32::MAX as usize, + DataType::Int64 => i64::MAX as usize, + DataType::UInt8 => u8::MAX as usize, + DataType::UInt16 => u16::MAX as usize, + DataType::UInt32 => u32::MAX as usize, + DataType::UInt64 => usize::MAX, + _ => return false, + }; + num_values == 0 || num_values - 1 <= max + } + + fn hash_values(&mut self, values: &ArrayRef) { + self.val_hashes.clear(); + self.val_hashes.resize(values.len(), 0); + create_hashes( + std::slice::from_ref(values), + &self.random_state, + &mut self.val_hashes, + ) + .unwrap(); + } + + fn find_or_insert_value( + &mut self, + dict_values: &ArrayRef, + val_idx: usize, + hash: u64, + ) -> Result<usize> { + let inner = &*self.inner; + let existing = self + .value_dedup + .find(hash, |&(entry_hash, slot)| { + entry_hash == hash && inner.equal_to(slot, dict_values, val_idx) + }) + .map(|&(_, slot)| slot); + + match existing { + Some(slot) => Ok(slot), + None => { + let slot = self.inner.len(); + self.inner.append_val(dict_values, val_idx)?; + self.value_dedup.insert_accounted( + (hash, slot), + |&(entry_hash, _)| entry_hash, + &mut self.value_dedup_size, + ); + Ok(slot) + } + } + } + + fn find_or_insert_null(&mut self) -> Result<usize> { + if let Some(slot) = self.null_inner_slot { + return Ok(slot); + } + let slot = self.inner.len(); + self.inner.append_val(&self.null_array, 0)?; + self.null_inner_slot = Some(slot); + Ok(slot) + } + + fn build_lookup_table( + &self, + dict_values: &ArrayRef, + val_hashes: &[u64], + ) -> Vec<usize> { + let num_distinct = dict_values.len(); + let mut table = vec![usize::MAX; num_distinct + 1]; + let inner = &*self.inner; + for val_idx in 0..num_distinct { + if dict_values.is_null(val_idx) { + table[val_idx] = self.null_inner_slot.unwrap_or(usize::MAX); + } else { + let hash = val_hashes[val_idx]; + if let Some(&(_, slot)) = + self.value_dedup.find(hash, |&(entry_hash, slot)| { + entry_hash == hash && inner.equal_to(slot, dict_values, val_idx) + }) + { + table[val_idx] = slot; + } + } + } + table[num_distinct] = self.null_inner_slot.unwrap_or(usize::MAX); + table + } + + /// Per-row fallback for `vectorized_equal_to` used when the number of rows + /// to check is smaller than the dictionary cardinality, making the O(D) + /// lookup-table build more expensive than direct value comparison. + /// + /// `#[cold]` + `#[inline(never)]` keeps this code out of the hot + /// lookup-table loops in `vectorized_equal_to` so LLVM can pipeline them. + #[cold] + #[inline(never)] + fn equal_to_per_row( + &self, + lhs_rows: &[usize], + dict_values: &ArrayRef, + dict: &DictionaryArray<K>, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let group_to_inner = self.group_to_inner.as_slice(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + let lhs_slot = group_to_inner[lhs_row]; + let equal = match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_slot, &self.null_array, 0), + Some(val_idx) if dict_values.is_null(val_idx) => { + self.inner.equal_to(lhs_slot, &self.null_array, 0) + } + Some(val_idx) => self.inner.equal_to(lhs_slot, dict_values, val_idx), + }; + if !equal { + equal_to_results.set_bit(idx, false); + } + } + } +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn + for DictionaryGroupValuesColumn<K> +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + let lhs_slot = self.group_to_inner[lhs_row]; + let dict = array.as_dictionary::<K>(); + match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_slot, &self.null_array, 0), + Some(val_idx) if dict.values().is_null(val_idx) => { + self.inner.equal_to(lhs_slot, &self.null_array, 0) + } + Some(val_idx) => self.inner.equal_to(lhs_slot, dict.values(), val_idx), + } + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let dict = array.as_dictionary::<K>(); + let inner_slot = match dict.key(row) { + None => self.find_or_insert_null()?, + Some(val_idx) if dict.values().is_null(val_idx) => { + self.find_or_insert_null()? + } + Some(val_idx) => { + let dict_values = dict.values(); + let single = dict_values.slice(val_idx, 1); + self.hash_values(&single); + self.find_or_insert_value(dict_values, val_idx, self.val_hashes[0])? + } + }; + self.group_to_inner.push(inner_slot); + self.check_key_overflow() + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let dict = array.as_dictionary::<K>(); + let dict_keys = dict.keys(); + let dict_values = dict.values(); + let num_distinct = dict_values.len(); + + // The fallback is in a separate #[cold] function so its code does not + // appear inline here and cannot prevent LLVM from pipelining / unrolling + // the hot lookup-table loops below. + if rhs_rows.len() < num_distinct { + self.equal_to_per_row( + lhs_rows, + dict_values, + dict, + rhs_rows, + equal_to_results, + ); + return; + } + + let mut val_hashes = vec![0u64; dict_values.len()]; + create_hashes( + std::slice::from_ref(dict_values), + &self.random_state, + &mut val_hashes, + ) + .unwrap(); + let lookup = self.build_lookup_table(dict_values, &val_hashes); + + let group_to_inner = self.group_to_inner.as_slice(); + + if dict_keys.null_count() == 0 { + // No null keys : skip the get_bit guard: we only ever write false, + // so overwriting an already-false bit is a no-op. + let raw_keys = dict_keys.values(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + let rhs_slot = lookup[raw_keys[rhs_row].as_usize()]; + if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot { + equal_to_results.set_bit(idx, false); + } + } + } else { + let null_buf = dict_keys.nulls().unwrap(); + let raw_keys = dict_keys.values(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if equal_to_results.get_bit(idx) { + let val_idx = if null_buf.is_null(rhs_row) { + num_distinct + } else { + raw_keys[rhs_row].as_usize() + }; + let rhs_slot = lookup[val_idx]; + if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot { + equal_to_results.set_bit(idx, false); + } + } + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let dict = array.as_dictionary::<K>(); + let dict_keys = dict.keys(); + let dict_values = dict.values(); + let num_distinct = dict_values.len(); + + self.hash_values(dict_values); + self.val_to_inner.clear(); + self.val_to_inner.resize(num_distinct, usize::MAX); + + self.group_to_inner.try_reserve(rows.len()).map_err(|e| { + DataFusionError::ArrowError( + Box::new(ArrowError::MemoryError(e.to_string())), + None, + ) + })?; + + if dict_keys.null_count() == 0 { + let raw_keys = dict_keys.values(); + for &row in rows { + let val_idx = raw_keys[row].as_usize(); + if self.val_to_inner[val_idx] == usize::MAX { + // A non-null key can still point to a null value in the values array. + self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { + self.find_or_insert_null()? + } else { + self.find_or_insert_value( + dict_values, + val_idx, + self.val_hashes[val_idx], + )? + }; + } + self.group_to_inner.push(self.val_to_inner[val_idx]); + } + } else { + let raw_keys = dict_keys.values(); + let null_buf = dict_keys.nulls().unwrap(); + for &row in rows { + let slot = if null_buf.is_null(row) { + self.find_or_insert_null()? + } else { + let val_idx = raw_keys[row].as_usize(); + if self.val_to_inner[val_idx] == usize::MAX { + self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { + self.find_or_insert_null()? + } else { + self.find_or_insert_value( + dict_values, + val_idx, + self.val_hashes[val_idx], + )? + }; + } + self.val_to_inner[val_idx] + }; + self.group_to_inner.push(slot); + } + } + + self.check_key_overflow() + } + + fn len(&self) -> usize { + self.group_to_inner.len() + } + + fn size(&self) -> usize { + self.inner.size() + + self.value_dedup_size + + self.group_to_inner.capacity() * size_of::<usize>() + + self.val_to_inner.capacity() * size_of::<usize>() + + self.val_hashes.capacity() * size_of::<u64>() + + self.null_array.get_array_memory_size() + + size_of::<Self>() + } + + fn build(self: Box<Self>) -> ArrayRef { + let values = self.inner.build(); + Self::into_dict(values, &self.group_to_inner) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let old_inner_len = self.inner.len(); + let all_inner_values = self.inner.take_n(old_inner_len); + + let mut emit_old_to_new = vec![usize::MAX; old_inner_len]; + let mut emit_new_to_old: Vec<usize> = Vec::new(); + for &old in &self.group_to_inner[..n] { + if emit_old_to_new[old] == usize::MAX { + emit_old_to_new[old] = emit_new_to_old.len(); + emit_new_to_old.push(old); + } + } + let emit_indices = + Int64Array::from_iter(emit_new_to_old.iter().map(|&i| i as i64)); + let compact_emit_values = + take(&*all_inner_values, &emit_indices, None).expect("take emit values"); + let emitted_keys: PrimitiveArray<K> = self.group_to_inner[..n] + .iter() + .map(|&old| { + if all_inner_values.is_null(old) { + None + } else { + Some(K::Native::usize_as(emit_old_to_new[old])) Review Comment: Null participates in emitted slot ordering. At key capacity, the final value gets index 128 for Int8 or 256 for UInt8. Int8 panics; UInt8 wraps and returns the wrong value. Move null last or exclude it. Add repeated boundary-emission tests. ########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs: ########## @@ -0,0 +1,910 @@ +// 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::multi_group_by::GroupColumn; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, Int64Array, + PrimitiveArray, +}; +use arrow::compute::take; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field}; +use arrow::error::ArrowError; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::{DataFusionError, Result, exec_err}; +use datafusion_execution::memory_pool::proxy::HashTableAllocExt; +use hashbrown::hash_table::HashTable; +use std::marker::PhantomData; +use std::mem::size_of; +use std::sync::Arc; + +use crate::aggregates::AGGREGATION_HASH_SEED; + +/// [`GroupColumn`] for dictionary-encoded columns with key type `K`. +/// +/// `inner` holds one slot per distinct value seen across all batches. +/// `group_to_inner[group_idx]` maps each group to its slot in `inner`, +/// so groups with the same value share a slot rather than duplicating data. +pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> { + /// Deduplicated store of distinct values. + inner: Box<dyn GroupColumn>, + /// Unary null array (length 1) reused for every null appended to `inner`. + null_array: ArrayRef, + /// Maps each group index to its slot in `inner`. + group_to_inner: Vec<usize>, + /// Lookup table mapping `(value_hash, inner_slot)` for each non-null distinct value. + value_dedup: HashTable<(u64, usize)>, + /// Tracked allocation size of `value_dedup` for memory accounting via `size()`. + value_dedup_size: usize, + /// Slot in `inner` for the null group; `None` until the first null is seen. + null_inner_slot: Option<usize>, + /// Hash seed — must match `create_hashes` so hashes are consistent across calls. + random_state: RandomState, + /// Reusable scratch buffer mapping `val_idx → inner_slot` across batches. + val_to_inner: Vec<usize>, + /// Reusable hash buffer for the dictionary values array. + val_hashes: Vec<u64>, + _phantom: PhantomData<K>, +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> { + pub fn new(inner: Box<dyn GroupColumn>, field: &Field) -> Self { + let null_array = arrow::array::new_null_array(field.data_type(), 1); + Self { + inner, + null_array, + group_to_inner: Vec::new(), + value_dedup: HashTable::new(), + value_dedup_size: 0, + null_inner_slot: None, + random_state: AGGREGATION_HASH_SEED, + val_to_inner: Vec::default(), + val_hashes: Vec::default(), + _phantom: PhantomData, + } + } + + fn into_dict(values: ArrayRef, group_to_inner: &[usize]) -> ArrayRef { + let keys: PrimitiveArray<K> = group_to_inner + .iter() + .map(|&slot| { + if values.is_null(slot) { + None + } else { + Some(K::Native::usize_as(slot)) + } + }) + .collect(); + Arc::new(DictionaryArray::<K>::new(keys, values)) + } + + // https://github.com/apache/datafusion/issues/23127 + // Null groups emit a null key, not a key index into the values array, so the + // null inner slot does not consume a key index. + fn check_key_overflow(&self) -> Result<()> { + // Keys are raw slot indices. The null slot is excluded from key count + // only when it occupies the last position — any non-null slot above it + // still emits that slot's raw index as a key. + let inner_len = self.inner.len(); + let null_slot_excluded = self.null_inner_slot.is_some_and(|s| s + 1 == inner_len); + let max_key_count = inner_len - null_slot_excluded as usize; + if !Self::key_type_fits(max_key_count) { + let non_null_slots = inner_len - self.null_inner_slot.is_some() as usize; + return exec_err!( + "Dictionary key type {:?} cannot represent {} distinct values", + K::DATA_TYPE, + non_null_slots + ); + } + Ok(()) + } + + fn key_type_fits(num_values: usize) -> bool { + let max: usize = match K::DATA_TYPE { + DataType::Int8 => i8::MAX as usize, + DataType::Int16 => i16::MAX as usize, + DataType::Int32 => i32::MAX as usize, + DataType::Int64 => i64::MAX as usize, + DataType::UInt8 => u8::MAX as usize, + DataType::UInt16 => u16::MAX as usize, + DataType::UInt32 => u32::MAX as usize, + DataType::UInt64 => usize::MAX, + _ => return false, + }; + num_values == 0 || num_values - 1 <= max + } + + fn hash_values(&mut self, values: &ArrayRef) { + self.val_hashes.clear(); + self.val_hashes.resize(values.len(), 0); + create_hashes( + std::slice::from_ref(values), + &self.random_state, + &mut self.val_hashes, + ) + .unwrap(); + } + + fn find_or_insert_value( + &mut self, + dict_values: &ArrayRef, + val_idx: usize, + hash: u64, + ) -> Result<usize> { + let inner = &*self.inner; + let existing = self + .value_dedup + .find(hash, |&(entry_hash, slot)| { + entry_hash == hash && inner.equal_to(slot, dict_values, val_idx) + }) + .map(|&(_, slot)| slot); + + match existing { + Some(slot) => Ok(slot), + None => { + let slot = self.inner.len(); + self.inner.append_val(dict_values, val_idx)?; + self.value_dedup.insert_accounted( + (hash, slot), + |&(entry_hash, _)| entry_hash, + &mut self.value_dedup_size, + ); + Ok(slot) + } + } + } + + fn find_or_insert_null(&mut self) -> Result<usize> { + if let Some(slot) = self.null_inner_slot { + return Ok(slot); + } + let slot = self.inner.len(); + self.inner.append_val(&self.null_array, 0)?; + self.null_inner_slot = Some(slot); + Ok(slot) + } + + fn build_lookup_table( + &self, + dict_values: &ArrayRef, + val_hashes: &[u64], + ) -> Vec<usize> { + let num_distinct = dict_values.len(); + let mut table = vec![usize::MAX; num_distinct + 1]; + let inner = &*self.inner; + for val_idx in 0..num_distinct { + if dict_values.is_null(val_idx) { + table[val_idx] = self.null_inner_slot.unwrap_or(usize::MAX); + } else { + let hash = val_hashes[val_idx]; + if let Some(&(_, slot)) = + self.value_dedup.find(hash, |&(entry_hash, slot)| { + entry_hash == hash && inner.equal_to(slot, dict_values, val_idx) + }) + { + table[val_idx] = slot; + } + } + } + table[num_distinct] = self.null_inner_slot.unwrap_or(usize::MAX); + table + } + + /// Per-row fallback for `vectorized_equal_to` used when the number of rows + /// to check is smaller than the dictionary cardinality, making the O(D) + /// lookup-table build more expensive than direct value comparison. + /// + /// `#[cold]` + `#[inline(never)]` keeps this code out of the hot + /// lookup-table loops in `vectorized_equal_to` so LLVM can pipeline them. + #[cold] + #[inline(never)] + fn equal_to_per_row( + &self, + lhs_rows: &[usize], + dict_values: &ArrayRef, + dict: &DictionaryArray<K>, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let group_to_inner = self.group_to_inner.as_slice(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + let lhs_slot = group_to_inner[lhs_row]; + let equal = match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_slot, &self.null_array, 0), + Some(val_idx) if dict_values.is_null(val_idx) => { + self.inner.equal_to(lhs_slot, &self.null_array, 0) + } + Some(val_idx) => self.inner.equal_to(lhs_slot, dict_values, val_idx), + }; + if !equal { + equal_to_results.set_bit(idx, false); + } + } + } +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn + for DictionaryGroupValuesColumn<K> +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + let lhs_slot = self.group_to_inner[lhs_row]; + let dict = array.as_dictionary::<K>(); + match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_slot, &self.null_array, 0), + Some(val_idx) if dict.values().is_null(val_idx) => { + self.inner.equal_to(lhs_slot, &self.null_array, 0) + } + Some(val_idx) => self.inner.equal_to(lhs_slot, dict.values(), val_idx), + } + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let dict = array.as_dictionary::<K>(); + let inner_slot = match dict.key(row) { + None => self.find_or_insert_null()?, + Some(val_idx) if dict.values().is_null(val_idx) => { + self.find_or_insert_null()? + } + Some(val_idx) => { + let dict_values = dict.values(); + let single = dict_values.slice(val_idx, 1); + self.hash_values(&single); + self.find_or_insert_value(dict_values, val_idx, self.val_hashes[0])? + } + }; + self.group_to_inner.push(inner_slot); + self.check_key_overflow() + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let dict = array.as_dictionary::<K>(); + let dict_keys = dict.keys(); + let dict_values = dict.values(); + let num_distinct = dict_values.len(); + + // The fallback is in a separate #[cold] function so its code does not + // appear inline here and cannot prevent LLVM from pipelining / unrolling + // the hot lookup-table loops below. + if rhs_rows.len() < num_distinct { + self.equal_to_per_row( + lhs_rows, + dict_values, + dict, + rhs_rows, + equal_to_results, + ); + return; + } + + let mut val_hashes = vec![0u64; dict_values.len()]; + create_hashes( + std::slice::from_ref(dict_values), + &self.random_state, + &mut val_hashes, + ) + .unwrap(); + let lookup = self.build_lookup_table(dict_values, &val_hashes); + + let group_to_inner = self.group_to_inner.as_slice(); + + if dict_keys.null_count() == 0 { + // No null keys : skip the get_bit guard: we only ever write false, + // so overwriting an already-false bit is a no-op. + let raw_keys = dict_keys.values(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + let rhs_slot = lookup[raw_keys[rhs_row].as_usize()]; + if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot { + equal_to_results.set_bit(idx, false); + } + } + } else { + let null_buf = dict_keys.nulls().unwrap(); + let raw_keys = dict_keys.values(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if equal_to_results.get_bit(idx) { + let val_idx = if null_buf.is_null(rhs_row) { + num_distinct + } else { + raw_keys[rhs_row].as_usize() + }; + let rhs_slot = lookup[val_idx]; + if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot { + equal_to_results.set_bit(idx, false); + } + } + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let dict = array.as_dictionary::<K>(); + let dict_keys = dict.keys(); + let dict_values = dict.values(); + let num_distinct = dict_values.len(); + + self.hash_values(dict_values); + self.val_to_inner.clear(); + self.val_to_inner.resize(num_distinct, usize::MAX); + + self.group_to_inner.try_reserve(rows.len()).map_err(|e| { + DataFusionError::ArrowError( + Box::new(ArrowError::MemoryError(e.to_string())), + None, + ) + })?; + + if dict_keys.null_count() == 0 { + let raw_keys = dict_keys.values(); + for &row in rows { + let val_idx = raw_keys[row].as_usize(); + if self.val_to_inner[val_idx] == usize::MAX { + // A non-null key can still point to a null value in the values array. + self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { + self.find_or_insert_null()? + } else { + self.find_or_insert_value( + dict_values, + val_idx, + self.val_hashes[val_idx], + )? + }; + } + self.group_to_inner.push(self.val_to_inner[val_idx]); + } + } else { + let raw_keys = dict_keys.values(); + let null_buf = dict_keys.nulls().unwrap(); + for &row in rows { + let slot = if null_buf.is_null(row) { + self.find_or_insert_null()? + } else { + let val_idx = raw_keys[row].as_usize(); + if self.val_to_inner[val_idx] == usize::MAX { + self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { + self.find_or_insert_null()? + } else { + self.find_or_insert_value( + dict_values, + val_idx, + self.val_hashes[val_idx], + )? + }; + } + self.val_to_inner[val_idx] + }; + self.group_to_inner.push(slot); + } + } + + self.check_key_overflow() + } + + fn len(&self) -> usize { + self.group_to_inner.len() + } + + fn size(&self) -> usize { + self.inner.size() + + self.value_dedup_size + + self.group_to_inner.capacity() * size_of::<usize>() + + self.val_to_inner.capacity() * size_of::<usize>() + + self.val_hashes.capacity() * size_of::<u64>() + + self.null_array.get_array_memory_size() + + size_of::<Self>() + } + + fn build(self: Box<Self>) -> ArrayRef { + let values = self.inner.build(); + Self::into_dict(values, &self.group_to_inner) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let old_inner_len = self.inner.len(); Review Comment: Every partial emit still drains, hashes, and rebuilds all surviving values, retaining the O(G² / batch_size) cost and high peak memory. Also, Arrow take retains full backing storage for Utf8View, BinaryView, and nested dictionaries. -- 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]
