zhuqi-lucas commented on code in PR #23187: URL: https://github.com/apache/datafusion/pull/23187#discussion_r3655293391
########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs: ########## @@ -0,0 +1,666 @@ +// 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, PrimitiveArray, +}; +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. +/// +/// `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>, + /// Single-element null array for appending null entries 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<()> { + let non_null_slots = self.inner.len() - self.null_inner_slot.is_some() as usize; + if !Self::key_type_fits(non_null_slots) { + 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 + } +} + +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(); + + let mut val_hashes = vec![0u64; dict_values.len()]; Review Comment: I tracked down the remaining regression (`size_65536_card_65536_null_0.10`): `vectorized_equal_to` is O(dictionary cardinality) instead of O(rows to check) — every call hashes **all** dict values and probes `value_dedup` for each of them to build the full lookup table, plus a `vec![usize::MAX; D+1]` allocation (512KB at D=65536). Why only this config regresses: `intern` only calls `vectorized_equal_to` for rows whose hash matched an existing group. With unique values and no nulls there are almost no matches, so the path rarely runs (hence `null_0.00` is fine at 1.04). But with 10% nulls, **every batch's ~6.5K null rows hash-match the existing null group**, so every batch pays O(65536) hashing + the table build just to check ~6.5K rows. Suggestion: fall back to per-row value comparison when the rows to check are fewer than the distinct values — no hashing, no O(D) work. The lookup-table path stays as-is for the low-cardinality/high-repetition shape it is designed for (its 1.8× wins are untouched — those configs take the table branch either way). Verified locally on this branch (criterion, Apple Silicon): | config | before | after | |---|---|---| | `size_65536_card_65536_null_0.10` (the regression) | 9.78ms | **8.04ms (−18%)** | | `size_8192_card_8192_null_0.10` | 1.41ms | **0.92ms (−34%)** | | low-cardinality configs | — | unchanged within noise | `multi_group_by` module tests all pass. ```suggestion // The lookup-table strategy below pays O(num_distinct) up front — // hashing every dictionary value and probing `value_dedup` for each — // so that each row check is a single integer compare. That only pays // for itself when the rows to check outnumber the distinct values // (low-cardinality dictionaries under high repetition). For // high-cardinality dictionaries checking few rows (e.g. a batch whose // only hash matches against existing groups are its null rows), the // table build dominates: delegate per row to the inner builder's // value comparison instead — no hashing, no O(num_distinct) work. if rhs_rows.len() < num_distinct { 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); } } return; } let mut val_hashes = vec![0u64; dict_values.len()]; ``` -- 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]
