From 17e8cbcb60d8875bfbb46c768cb989bbedfe2da1 Mon Sep 17 00:00:00 2001 From: Richard Date: Thu, 25 Jun 2026 09:52:12 -0400 Subject: [PATCH 1/8] introduce dictionarys as a supported group column type add schema support for dictionarys introduce high level GroupValuesColumn test introduce groupColumn trait test git issues introduce edge case/ regression section inital impl working implementation of dictionary for groupValuesCOlumns benchmarks show perf boost over groupvaluerows, TODO:dedupe items before inner append fix clippy errors & inline final builder add cache for arc ptr trim down test trim test LOC again trim PR revision 3 speed up low cardinlaity case working version introduce inter-batch caching break complex types into seperate parts fixed breaking test, re-allocate hashtable on each intern() call re-introduce cache remove mutex add cache to concat pointers to avoid un-needed allocations remove ptr caches and concat call reduce LOC revised PR comments add regression test to align with GroupValueRows & re-order overflow check add test to assert de-duplicated output dictionary tmp low card speed up optimize low-card case wip re-use allocations across calls add test final clean up --- .../group_values/multi_group_by/dictionary.rs | 626 ++++++++++++++++++ .../group_values/multi_group_by/mod.rs | 170 ++++- 2 files changed, 777 insertions(+), 19 deletions(-) create mode 100644 datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs new file mode 100644 index 0000000000000..e2830cd5aec51 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -0,0 +1,626 @@ +// 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 { + /// Deduplicated store of distinct values. + inner: Box, + /// 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, + /// 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, + /// 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, + /// Reusable hash buffer for the dictionary values array. + val_hashes: Vec, + _phantom: PhantomData, +} + +impl DictionaryGroupValuesColumn { + pub fn new(inner: Box, 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 = group_to_inner + .iter() + .map(|&slot| { + if values.is_null(slot) { + None + } else { + Some(K::Native::usize_as(slot)) + } + }) + .collect(); + Arc::new(DictionaryArray::::new(keys, values)) + } + + // https://github.com/apache/datafusion/issues/23127 + fn check_key_overflow(num_inner_slots: usize) -> Result<()> { + if !Self::key_type_fits(num_inner_slots) { + return exec_err!( + "Dictionary key type {:?} cannot represent {} distinct values", + K::DATA_TYPE, + num_inner_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 { + 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 { + 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 { + 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 GroupColumn + for DictionaryGroupValuesColumn +{ + 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::(); + 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::(); + 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(); + self.hash_values(dict_values); + self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx])? + } + }; + self.group_to_inner.push(inner_slot); + Self::check_key_overflow(self.inner.len()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let dict = array.as_dictionary::(); + 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()]; + 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::(); + 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(self.inner.len()) + } + + 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::() + + self.val_to_inner.capacity() * size_of::() + + self.val_hashes.capacity() * size_of::() + + self.null_array.get_array_memory_size() + + size_of::() + } + + fn build(self: Box) -> ArrayRef { + let values = self.inner.build(); + Self::into_dict(values, &self.group_to_inner) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + // `inner` is a trait object — the only way to extract its data is via `take_n`. + // Because group→inner slot mappings are non-contiguous, we drain all of `inner` + // at once, then re-append only the slots still referenced by the remaining groups. + let old_inner_len = self.inner.len(); + let all_inner_values = self.inner.take_n(old_inner_len); + + let emitted = + Self::into_dict(Arc::clone(&all_inner_values), &self.group_to_inner[..n]); + + let remaining = self.group_to_inner[n..].to_vec(); + + // Map each referenced old slot to a new contiguous index. + let mut old_to_new = vec![usize::MAX; old_inner_len]; + let mut new_to_old: Vec = Vec::new(); + for &old in &remaining { + if old_to_new[old] == usize::MAX { + old_to_new[old] = new_to_old.len(); + new_to_old.push(old); + } + } + + self.value_dedup = HashTable::new(); + self.value_dedup_size = 0; + self.null_inner_slot = None; + + for (new_slot, &old_slot) in new_to_old.iter().enumerate() { + if all_inner_values.is_null(old_slot) { + self.inner + .append_val(&self.null_array, 0) + .expect("append null failed in take_n"); + self.null_inner_slot = Some(new_slot); + } else { + self.inner + .append_val(&all_inner_values, old_slot) + .expect("append value failed in take_n"); + let single = all_inner_values.slice(old_slot, 1); + self.hash_values(&single); + self.value_dedup.insert_accounted( + (self.val_hashes[0], new_slot), + |&(entry_hash, _)| entry_hash, + &mut self.value_dedup_size, + ); + } + } + + self.group_to_inner = remaining.iter().map(|&old| old_to_new[old]).collect(); + Self::check_key_overflow(self.inner.len()).expect("key overflow in take_n"); + + emitted + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::aggregates::group_values::multi_group_by::bytes::ByteGroupValueBuilder; + use arrow::array::{ + Array, ArrayRef, BooleanBufferBuilder, DictionaryArray, Int32Array, StringArray, + UInt8Array, + }; + use arrow::compute::cast; + use arrow::datatypes::{DataType, Int32Type, UInt8Type}; + use datafusion_physical_expr::binary_map::OutputType; + use std::sync::Arc; + + fn utf8_col() -> DictionaryGroupValuesColumn { + let field = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &field, + ) + } + + fn dict_arr(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + Arc::new(DictionaryArray::::new( + Int32Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + + fn str_values(arr: &ArrayRef) -> Vec> { + let plain = cast(arr.as_ref(), &DataType::Utf8).unwrap(); + let strings = plain.as_any().downcast_ref::().unwrap(); + (0..strings.len()) + .map(|i| strings.is_valid(i).then(|| strings.value(i).to_owned())) + .collect() + } + + fn bool_vec(buf: &BooleanBufferBuilder) -> Vec { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + fn all_true(len: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(len); + buf.append_n(len, true); + buf + } + + // Null key and null-valued dict entry both map to the null group. + #[test] + fn null_key_and_null_value_in_dict() { + let mut col = utf8_col(); + // Row 0: null key, Row 1: key→null value, Row 2: key→"b" + let input = dict_arr(&[None, Some(0), Some(1)], &[None, Some("b")]); + for row in 0..3 { + col.append_val(&input, row).unwrap(); + } + + assert!(col.equal_to(0, &input, 1)); + assert!(col.equal_to(1, &input, 0)); + assert!(!col.equal_to(0, &input, 2)); + assert!(!col.equal_to(2, &input, 0)); + + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 2); + assert_eq!(str_values(&out), vec![None, None, Some("b".into())]); + } + + #[test] + fn take_n_remaps_slots_across_batches() { + use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; + use arrow::array::UInt64Array; + use arrow::datatypes::UInt64Type; + + let field = Field::new("", DataType::UInt64, true); + let mut col = DictionaryGroupValuesColumn::::new( + Box::new(PrimitiveGroupValueBuilder::::new( + DataType::UInt64, + )), + &field, + ); + + let u64_val = |arr: &ArrayRef, pos: usize| { + let dict = arr.as_dictionary::(); + dict.values() + .as_any() + .downcast_ref::() + .unwrap() + .value(dict.key(pos).unwrap()) + }; + + let batch1: ArrayRef = Arc::new(DictionaryArray::::new( + Int32Array::from(vec![Some(0), Some(1), None, Some(2), Some(0)]), + Arc::new(UInt64Array::from(vec![10u64, 20, 30])), + )); + col.vectorized_append(&batch1, &[0, 1, 2, 3, 4]).unwrap(); + + let emitted = col.take_n(3); + assert_eq!(u64_val(&emitted, 0), 10); + assert_eq!(u64_val(&emitted, 1), 20); + assert!(emitted.as_dictionary::().key(2).is_none()); + + let batch2: ArrayRef = Arc::new(DictionaryArray::::new( + Int32Array::from(vec![None, Some(0)]), + Arc::new(UInt64Array::from(vec![99u64])), + )); + col.vectorized_append(&batch2, &[0, 1]).unwrap(); + + let mut buf = all_true(2); + col.vectorized_equal_to(&[2, 3], &batch2, &[0, 1], &mut buf); + assert_eq!(bool_vec(&buf), vec![true, true]); + + let out = Box::new(col).build(); + assert_eq!(u64_val(&out, 0), 30); + assert_eq!(u64_val(&out, 1), 10); + assert!(out.as_dictionary::().key(2).is_none()); + assert_eq!(u64_val(&out, 3), 99); + } + + // Regression: https://github.com/apache/datafusion/issues/23127 + #[test] + fn key_type_overflow_returns_error() { + let field = Field::new("", DataType::Utf8, true); + let mut col = DictionaryGroupValuesColumn::::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &field, + ); + + let strs: Vec = (0..=255u16).map(|i| i.to_string()).collect(); + let str_refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + let full: ArrayRef = Arc::new(DictionaryArray::::new( + UInt8Array::from((0..=255u8).map(Some).collect::>()), + Arc::new(StringArray::from(str_refs)), + )); + col.vectorized_append(&full, &(0..256).collect::>()) + .unwrap(); + + let extra: ArrayRef = Arc::new(DictionaryArray::::new( + UInt8Array::from(vec![Some(0u8)]), + Arc::new(StringArray::from(vec![Some("overflow")])), + )); + assert!(col.append_val(&extra, 0).is_err()); + } + + // build_lookup_table must use the incoming batch's hashes, not + // stale ones left by the last vectorized_append call. + #[test] + fn vectorized_equal_to_uses_current_batch_hashes() { + let mut col = utf8_col(); + + let batch1 = dict_arr(&[Some(0)], &[Some("a"), Some("b")]); + col.vectorized_append(&batch1, &[0]).unwrap(); + + // values = ["z", "a"]; key 0 → "a" at val_idx 1. + // Stale hashes would probe val_idx 1 with hash("b") and miss. + let batch2 = dict_arr(&[Some(1)], &[Some("z"), Some("a")]); + let mut buf = all_true(1); + col.vectorized_equal_to(&[0], &batch2, &[0], &mut buf); + assert_eq!(bool_vec(&buf), vec![true]); + } + + #[test] + fn append_only_stores_referenced_values() { + let mut col = utf8_col(); + let values = Arc::new(StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + Some("f"), + Some("g"), + Some("h"), + Some("i"), + Some("j"), + ])); + let keys = Int32Array::from(vec![ + Some(0), // a + Some(2), // c + Some(7), // h + Some(0), + Some(2), + Some(7), + Some(7), + Some(0), + Some(2), + ]); + let input: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); + + col.vectorized_append(&input, &[0, 1, 2, 3, 4, 5, 6, 7, 8]) + .unwrap(); + + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 3); + assert_eq!( + str_values(&out), + vec![ + Some("a".into()), + Some("c".into()), + Some("h".into()), + Some("a".into()), + Some("c".into()), + Some("h".into()), + Some("h".into()), + Some("a".into()), + Some("c".into()), + ] + ); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 8b68152c477ac..f3870ce2fe496 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -20,6 +20,7 @@ mod boolean; mod bytes; pub mod bytes_view; +mod dictionary; pub mod primitive; pub mod row_backed; @@ -32,7 +33,6 @@ use crate::aggregates::group_values::multi_group_by::{ row_backed::RowsGroupColumn, }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; -use arrow::compute::cast; use arrow::datatypes::{ BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, @@ -46,7 +46,7 @@ use arrow::datatypes::{ }; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::EmitTo; use datafusion_physical_expr::binary_map::OutputType; @@ -972,7 +972,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Utf8View | DataType::BinaryView | DataType::Boolean - ) + ) || matches!(data_type, DataType::Dictionary(_,v ) if group_column_supported_type(v)) } /// Build a [`GroupColumn`] for a single schema field. @@ -994,7 +994,7 @@ fn make_group_column(field: &Field) -> Result> { let nullable = field.is_nullable(); let data_type = field.data_type(); let mut v: Vec> = Vec::with_capacity(1); - match *data_type { + match data_type { DataType::Int8 => instantiate_primitive!(v, nullable, Int8Type, data_type), DataType::Int16 => instantiate_primitive!(v, nullable, Int16Type, data_type), DataType::Int32 => instantiate_primitive!(v, nullable, Int32Type, data_type), @@ -1114,6 +1114,33 @@ fn make_group_column(field: &Field) -> Result> { v.push(Box::new(BooleanGroupValueBuilder::::new())); } } + DataType::Dictionary(key_dt, value_dt) => { + let new_field = Field::new("", *value_dt.clone(), true); + let inner = make_group_column(&new_field)?; + macro_rules! dict_col { + ($T:ty) => { + Box::new(dictionary::DictionaryGroupValuesColumn::<$T>::new( + inner, &new_field, + )) + }; + } + let col: Box = match key_dt.as_ref() { + DataType::Int8 => dict_col!(Int8Type), + DataType::Int16 => dict_col!(Int16Type), + DataType::Int32 => dict_col!(Int32Type), + DataType::Int64 => dict_col!(Int64Type), + DataType::UInt8 => dict_col!(UInt8Type), + DataType::UInt16 => dict_col!(UInt16Type), + DataType::UInt32 => dict_col!(UInt32Type), + DataType::UInt64 => dict_col!(UInt64Type), + _ => { + return not_impl_err!( + "Dictionary key type {key_dt} not supported in GroupValuesColumn" + ); + } + }; + v.push(col) + } // Generic fallback for nested types (Struct / List / LargeList / // FixedSizeList, recursively) that lack a type-specialized builder but // can be encoded by arrow's row format. This is what lets a mixed @@ -1162,7 +1189,7 @@ impl GroupValues for GroupValuesColumn { } fn emit(&mut self, emit_to: EmitTo) -> Result> { - let mut output = match emit_to { + let output = match emit_to { EmitTo::All => { // Replace the column builders with a fresh set so the // aggregator is immediately reusable after the drain. @@ -1252,20 +1279,6 @@ impl GroupValues for GroupValuesColumn { } }; - // TODO: Materialize dictionaries in group keys (#7647) - for (field, array) in self.schema.fields.iter().zip(&mut output) { - let expected = field.data_type(); - if let DataType::Dictionary(_, v) = expected { - let actual = array.data_type(); - if v.as_ref() != actual { - return Err(internal_datafusion_err!( - "Converted group rows expected dictionary of {v} got {actual}" - )); - } - *array = cast(array.as_ref(), expected)?; - } - } - Ok(output) } @@ -1624,6 +1637,19 @@ mod tests { DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth), DataType::Interval(arrow::datatypes::IntervalUnit::DayTime), DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int64)), + DataType::Dictionary( + Box::new(DataType::UInt16), + Box::new(DataType::LargeUtf8), + ), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Timestamp( + arrow::datatypes::TimeUnit::Nanosecond, + None, + )), + ), ]; for dt in &supported_cases { @@ -1651,6 +1677,11 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Float16)), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Decimal256(76, 10)), + ), ]; for dt in &unsupported_cases { @@ -1882,6 +1913,107 @@ mod tests { } } + // Regression for https://github.com/apache/datafusion/issues/23127: + // In a multi-column group-by there can be more than 128 groups while an + // Int8 dictionary column still has only a few distinct values. The toll + // row path must emit 129 groups without overflowing the Int8 key range. + // This also documents that the inner GroupColumn does NOT deduplicate dict + // values — each group gets its own slot, so values().len() == n_groups. + #[test] + fn multi_col_groupby_dict_many_groups_two_values() { + use arrow::array::{AsArray, DictionaryArray, Int16Array}; + use arrow::datatypes::Int16Type; + + let n_groups = 129_usize; + let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); + + // Each row has a unique label (forcing a new group) and alternates + // between the two dictionary values. Int16 keys are used so that + // 129 groups don't hit the Int8 overflow limit (i8::MAX = 127). + let labels: ArrayRef = Arc::new(StringArray::from( + (0..n_groups).map(|i| format!("g{i}")).collect::>(), + )); + let dict_keys = Int16Array::from( + (0..n_groups) + .map(|i| Some((i % 2) as i16)) + .collect::>(), + ); + let categories: ArrayRef = Arc::new(DictionaryArray::::new( + dict_keys, + Arc::clone(&dict_vocab), + )); + + let schema = Arc::new(Schema::new(vec![ + Field::new("label", DataType::Utf8, false), + Field::new( + "category", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + false, + ), + ])); + + let mut gv = GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + gv.intern(&[labels, categories], &mut vec![]).unwrap(); + let out = gv.emit(EmitTo::All).unwrap(); + + assert_eq!(out[0].len(), n_groups); + assert!(matches!( + out[1].data_type(), + DataType::Dictionary(k, v) + if k.as_ref() == &DataType::Int16 && v.as_ref() == &DataType::Utf8 + )); + // Both vectorized and streaming paths now deduplicate dict values. + assert_eq!(out[1].as_dictionary::().values().len(), 2); + } + + // Same as above but uses the non-vectorized (streaming) path via + // `GroupValuesColumn::`. The dict column's values array must be + // deduplicated — only 2 distinct entries ("cat" / "dog") regardless of + // how many groups were seen. + #[test] + fn multi_col_groupby_dict_many_groups_two_values_streaming() { + use arrow::array::{AsArray, DictionaryArray, Int16Array}; + use arrow::datatypes::Int16Type; + + let n_groups = 129_usize; + let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); + + let labels: ArrayRef = Arc::new(StringArray::from( + (0..n_groups).map(|i| format!("g{i}")).collect::>(), + )); + let dict_keys = Int16Array::from( + (0..n_groups) + .map(|i| Some((i % 2) as i16)) + .collect::>(), + ); + let categories: ArrayRef = Arc::new(DictionaryArray::::new( + dict_keys, + Arc::clone(&dict_vocab), + )); + + let schema = Arc::new(Schema::new(vec![ + Field::new("label", DataType::Utf8, false), + Field::new( + "category", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + false, + ), + ])); + + let mut gv = GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + gv.intern(&[labels, categories], &mut vec![]).unwrap(); + let out = gv.emit(EmitTo::All).unwrap(); + + assert_eq!(out[0].len(), n_groups); + assert!(matches!( + out[1].data_type(), + DataType::Dictionary(k, v) + if k.as_ref() == &DataType::Int16 && v.as_ref() == &DataType::Utf8 + )); + // append_val deduplicates: only 2 distinct values in the output values array. + assert_eq!(out[1].as_dictionary::().values().len(), 2); + } + #[test] fn test_intern_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); From 4ce2c9502cef412b297553ce22b65d838111a0ec Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Thu, 23 Jul 2026 11:13:03 -0400 Subject: [PATCH 2/8] add test, avoid extra hashes, fix i16 test to use i8 --- .../group_values/multi_group_by/dictionary.rs | 56 ++++++++++++-- .../group_values/multi_group_by/mod.rs | 73 +++---------------- 2 files changed, 59 insertions(+), 70 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index e2830cd5aec51..336454988e92a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -90,12 +90,15 @@ impl DictionaryGroupValuesColumn { } // https://github.com/apache/datafusion/issues/23127 - fn check_key_overflow(num_inner_slots: usize) -> Result<()> { - if !Self::key_type_fits(num_inner_slots) { + // 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, - num_inner_slots + non_null_slots ); } Ok(()) @@ -217,12 +220,13 @@ impl GroupColumn } Some(val_idx) => { let dict_values = dict.values(); - self.hash_values(dict_values); - self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx])? + 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(self.inner.len()) + self.check_key_overflow() } fn vectorized_equal_to( @@ -341,7 +345,7 @@ impl GroupColumn } } - Self::check_key_overflow(self.inner.len()) + self.check_key_overflow() } fn len(&self) -> usize { @@ -410,7 +414,7 @@ impl GroupColumn } self.group_to_inner = remaining.iter().map(|&old| old_to_new[old]).collect(); - Self::check_key_overflow(self.inner.len()).expect("key overflow in take_n"); + self.check_key_overflow().expect("key overflow in take_n"); emitted } @@ -558,6 +562,42 @@ mod tests { assert!(col.append_val(&extra, 0).is_err()); } + // A null value alongside 256 non-null values must not itself trigger an + // overflow: null groups emit a null key, not a key index. Adding a 257th + // non-null value after the null should be what triggers the error. + #[test] + fn null_does_not_count_toward_key_overflow() { + let field = Field::new("", DataType::Utf8, true); + let mut col = DictionaryGroupValuesColumn::::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &field, + ); + + // Fill all 256 UInt8 key slots (indices 0..=255) with distinct non-null values. + let strs: Vec = (0..=255u16).map(|i| i.to_string()).collect(); + let str_refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + let full: ArrayRef = Arc::new(DictionaryArray::::new( + UInt8Array::from((0..=255u8).map(Some).collect::>()), + Arc::new(StringArray::from(str_refs)), + )); + col.vectorized_append(&full, &(0..256).collect::>()) + .unwrap(); + + // Null does not consume a key index — appending it must succeed. + let null_arr: ArrayRef = Arc::new(DictionaryArray::::new( + UInt8Array::from(vec![None]), + Arc::new(StringArray::from(vec![Some("dummy")])), + )); + col.append_val(&null_arr, 0).unwrap(); + + // A 257th distinct non-null value now exceeds UInt8's capacity. + let extra: ArrayRef = Arc::new(DictionaryArray::::new( + UInt8Array::from(vec![Some(0u8)]), + Arc::new(StringArray::from(vec![Some("overflow")])), + )); + assert!(col.append_val(&extra, 0).is_err()); + } + // build_lookup_table must use the incoming batch's hashes, not // stale ones left by the last vectorized_append call. #[test] diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index f3870ce2fe496..7aeb97a0ea7af 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -1913,16 +1913,13 @@ mod tests { } } - // Regression for https://github.com/apache/datafusion/issues/23127: - // In a multi-column group-by there can be more than 128 groups while an - // Int8 dictionary column still has only a few distinct values. The toll - // row path must emit 129 groups without overflowing the Int8 key range. - // This also documents that the inner GroupColumn does NOT deduplicate dict - // values — each group gets its own slot, so values().len() == n_groups. + // https://github.com/apache/datafusion/issues/23127 + // validate DictionaryGroupColumn deduplicates values — only k distinct keys appear + // in the values array even when there are more than 128 groups total. #[test] fn multi_col_groupby_dict_many_groups_two_values() { - use arrow::array::{AsArray, DictionaryArray, Int16Array}; - use arrow::datatypes::Int16Type; + use arrow::array::{AsArray, DictionaryArray, Int8Array}; + use arrow::datatypes::Int8Type; let n_groups = 129_usize; let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); @@ -1933,12 +1930,12 @@ mod tests { let labels: ArrayRef = Arc::new(StringArray::from( (0..n_groups).map(|i| format!("g{i}")).collect::>(), )); - let dict_keys = Int16Array::from( + let dict_keys = Int8Array::from( (0..n_groups) - .map(|i| Some((i % 2) as i16)) + .map(|i| Some((i % 2) as i8)) .collect::>(), ); - let categories: ArrayRef = Arc::new(DictionaryArray::::new( + let categories: ArrayRef = Arc::new(DictionaryArray::::new( dict_keys, Arc::clone(&dict_vocab), )); @@ -1947,7 +1944,7 @@ mod tests { Field::new("label", DataType::Utf8, false), Field::new( "category", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), false, ), ])); @@ -1960,58 +1957,10 @@ mod tests { assert!(matches!( out[1].data_type(), DataType::Dictionary(k, v) - if k.as_ref() == &DataType::Int16 && v.as_ref() == &DataType::Utf8 + if k.as_ref() == &DataType::Int8 && v.as_ref() == &DataType::Utf8 )); // Both vectorized and streaming paths now deduplicate dict values. - assert_eq!(out[1].as_dictionary::().values().len(), 2); - } - - // Same as above but uses the non-vectorized (streaming) path via - // `GroupValuesColumn::`. The dict column's values array must be - // deduplicated — only 2 distinct entries ("cat" / "dog") regardless of - // how many groups were seen. - #[test] - fn multi_col_groupby_dict_many_groups_two_values_streaming() { - use arrow::array::{AsArray, DictionaryArray, Int16Array}; - use arrow::datatypes::Int16Type; - - let n_groups = 129_usize; - let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); - - let labels: ArrayRef = Arc::new(StringArray::from( - (0..n_groups).map(|i| format!("g{i}")).collect::>(), - )); - let dict_keys = Int16Array::from( - (0..n_groups) - .map(|i| Some((i % 2) as i16)) - .collect::>(), - ); - let categories: ArrayRef = Arc::new(DictionaryArray::::new( - dict_keys, - Arc::clone(&dict_vocab), - )); - - let schema = Arc::new(Schema::new(vec![ - Field::new("label", DataType::Utf8, false), - Field::new( - "category", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), - false, - ), - ])); - - let mut gv = GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); - gv.intern(&[labels, categories], &mut vec![]).unwrap(); - let out = gv.emit(EmitTo::All).unwrap(); - - assert_eq!(out[0].len(), n_groups); - assert!(matches!( - out[1].data_type(), - DataType::Dictionary(k, v) - if k.as_ref() == &DataType::Int16 && v.as_ref() == &DataType::Utf8 - )); - // append_val deduplicates: only 2 distinct values in the output values array. - assert_eq!(out[1].as_dictionary::().values().len(), 2); + assert_eq!(out[1].as_dictionary::().values().len(), 2); } #[test] From 43e0c0865f459c6b0d44cba811b8ad1b54f7b2bc Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Mon, 27 Jul 2026 10:43:14 -0400 Subject: [PATCH 3/8] address Pr comments and add test --- .../group_values/multi_group_by/dictionary.rs | 186 +++++++++++++++++- 1 file changed, 179 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index 336454988e92a..f5bb3573ce9da 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -93,8 +93,13 @@ impl DictionaryGroupValuesColumn { // 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) { + // Keys are raw slot indices; the null slot is excluded only when it is + // last (null groups emit None, not an index). + let inner_len = self.inner.len(); + let null_is_last = self.null_inner_slot.is_some_and(|s| s + 1 == inner_len); + let max_key_count = inner_len - null_is_last 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, @@ -194,6 +199,43 @@ impl DictionaryGroupValuesColumn { 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, + 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 GroupColumn @@ -241,6 +283,28 @@ impl GroupColumn let dict_values = dict.values(); let num_distinct = dict_values.len(); + // The lookup-table path pays O(num_distinct) upfront — hashing every + // dictionary value and probing value_dedup for each — so each row check + // becomes a single integer compare. That only pays for itself when the + // rows to check outnumber the distinct values (low-cardinality dicts + // under high repetition). For high-cardinality dicts with few rows to + // check (e.g. only null rows matched), the table build dominates: fall + // back to per-row value comparison instead. + // + // 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), @@ -253,7 +317,7 @@ impl GroupColumn 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, + // 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 @@ -369,7 +433,7 @@ impl GroupColumn fn take_n(&mut self, n: usize) -> ArrayRef { // `inner` is a trait object — the only way to extract its data is via `take_n`. - // Because group→inner slot mappings are non-contiguous, we drain all of `inner` + // Because group->inner slot mappings are non-contiguous, we drain all of `inner` // at once, then re-append only the slots still referenced by the remaining groups. let old_inner_len = self.inner.len(); let all_inner_values = self.inner.take_n(old_inner_len); @@ -393,6 +457,9 @@ impl GroupColumn self.value_dedup_size = 0; self.null_inner_slot = None; + // Hash all surviving values in one vectorized pass + self.hash_values(&all_inner_values); + for (new_slot, &old_slot) in new_to_old.iter().enumerate() { if all_inner_values.is_null(old_slot) { self.inner @@ -403,10 +470,8 @@ impl GroupColumn self.inner .append_val(&all_inner_values, old_slot) .expect("append value failed in take_n"); - let single = all_inner_values.slice(old_slot, 1); - self.hash_values(&single); self.value_dedup.insert_accounted( - (self.val_hashes[0], new_slot), + (self.val_hashes[old_slot], new_slot), |&(entry_hash, _)| entry_hash, &mut self.value_dedup_size, ); @@ -598,6 +663,113 @@ mod tests { assert!(col.append_val(&extra, 0).is_err()); } + // Regression: null mid-array means the last slot is non-null, so the max + // emitted key equals inner.len()-1, not inner.len()-2. + #[test] + fn key_overflow_null_slot_mid_array() { + use arrow::array::Int8Array; + use arrow::datatypes::Int8Type; + + let field = Field::new("", DataType::Utf8, true); + let mut col = DictionaryGroupValuesColumn::::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &field, + ); + + // 100 non-null values in slots 0..99 + let strs100: Vec = (0u8..100).map(|i| format!("s{i}")).collect(); + let refs100: Vec> = + strs100.iter().map(|s| Some(s.as_str())).collect(); + let b1: ArrayRef = Arc::new(DictionaryArray::::new( + Int8Array::from((0i8..100).map(Some).collect::>()), + Arc::new(StringArray::from(refs100)), + )); + col.vectorized_append(&b1, &(0usize..100).collect::>()) + .unwrap(); + + // null lands at slot 100 (mid-array) + let null_batch: ArrayRef = Arc::new(DictionaryArray::::new( + Int8Array::from(vec![None]), + Arc::new(StringArray::from(vec![Some("x")])), + )); + col.append_val(&null_batch, 0).unwrap(); + + // 28 more non-null values fill slots 101..128; slot 128 exceeds Int8::MAX + let strs28: Vec = (100u8..128).map(|i| format!("s{i}")).collect(); + let refs28: Vec> = strs28.iter().map(|s| Some(s.as_str())).collect(); + let b2: ArrayRef = Arc::new(DictionaryArray::::new( + Int8Array::from((0i8..28).map(Some).collect::>()), + Arc::new(StringArray::from(refs28)), + )); + assert!( + col.vectorized_append(&b2, &(0usize..28).collect::>()) + .is_err() + ); + } + + // Helpers shared by the three overflow boundary tests below. + fn int8_utf8_col() -> DictionaryGroupValuesColumn { + use arrow::datatypes::Int8Type; + let field = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &field, + ) + } + + fn int8_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + use arrow::array::Int8Array; + use arrow::datatypes::Int8Type; + Arc::new(DictionaryArray::::new( + Int8Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + + fn distinct_strs(start: usize, end: usize) -> ArrayRef { + let strs: Vec = (start..end).map(|i| format!("v{i}")).collect(); + let refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + int8_dict( + &(0..strs.len()).map(|i| Some(i as i8)).collect::>(), + &refs, + ) + } + + // null mid-array: last slot is non-null, so slot 128 overflows Int8 + #[test] + fn overflow_null_mid_pushes_over() { + let mut col = int8_utf8_col(); + // 100 non-null, then null at slot 100, then 28 more non-null (slots 101..128) + col.vectorized_append(&distinct_strs(0, 100), &(0..100).collect::>()) + .unwrap(); + col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + .unwrap(); + assert!( + col.vectorized_append(&distinct_strs(100, 128), &(0..28).collect::>()) + .is_err() + ); + } + + // null last: 128 non-null values fill slots 0..127 (Int8::MAX), null at 128 is ok + #[test] + fn overflow_null_last_at_max() { + let mut col = int8_utf8_col(); + col.vectorized_append(&distinct_strs(0, 128), &(0..128).collect::>()) + .unwrap(); + col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + .unwrap(); + } + + // null last: 127 non-null values, one below the limit + #[test] + fn overflow_null_last_just_below_max() { + let mut col = int8_utf8_col(); + col.vectorized_append(&distinct_strs(0, 127), &(0..127).collect::>()) + .unwrap(); + col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + .unwrap(); + } + // build_lookup_table must use the incoming batch's hashes, not // stale ones left by the last vectorized_append call. #[test] From 7af44a7d76a0831792b5785cb5b562403e44e354 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Mon, 27 Jul 2026 10:44:32 -0400 Subject: [PATCH 4/8] reduce comment bloat --- .../aggregates/group_values/multi_group_by/dictionary.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index f5bb3573ce9da..2320bf6be1bc1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -283,14 +283,6 @@ impl GroupColumn let dict_values = dict.values(); let num_distinct = dict_values.len(); - // The lookup-table path pays O(num_distinct) upfront — hashing every - // dictionary value and probing value_dedup for each — so each row check - // becomes a single integer compare. That only pays for itself when the - // rows to check outnumber the distinct values (low-cardinality dicts - // under high repetition). For high-cardinality dicts with few rows to - // check (e.g. only null rows matched), the table build dominates: fall - // back to per-row value comparison instead. - // // 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. From e799fe780f0205c869820425c50c02efaaa90355 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Mon, 27 Jul 2026 11:24:23 -0400 Subject: [PATCH 5/8] Fix clippy --- .../src/aggregates/group_values/multi_group_by/dictionary.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index 2320bf6be1bc1..1de129c039b1c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -290,7 +290,7 @@ impl GroupColumn self.equal_to_per_row( lhs_rows, dict_values, - &dict, + dict, rhs_rows, equal_to_results, ); From 2c76dd14d530f78049f4dcac561806ccc81fea15 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Fri, 31 Jul 2026 00:16:33 -0400 Subject: [PATCH 6/8] fix null handling & take_n --- .../group_values/multi_group_by/dictionary.rs | 110 +++++++++++++++--- .../group_values/multi_group_by/mod.rs | 4 +- 2 files changed, 97 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index 1de129c039b1c..cc9b078ce38e4 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -17,8 +17,10 @@ use crate::aggregates::group_values::multi_group_by::GroupColumn; use arrow::array::{ - Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, PrimitiveArray, + 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}; @@ -31,7 +33,7 @@ use std::sync::Arc; use crate::aggregates::AGGREGATION_HASH_SEED; -/// [`GroupColumn`] for dictionary-encoded columns. +/// [`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`, @@ -39,7 +41,7 @@ use crate::aggregates::AGGREGATION_HASH_SEED; pub struct DictionaryGroupValuesColumn { /// Deduplicated store of distinct values. inner: Box, - /// Single-element null array for appending null entries to `inner`. + /// 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, @@ -93,11 +95,12 @@ impl DictionaryGroupValuesColumn { // 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 only when it is - // last (null groups emit None, not an index). + // 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_is_last = self.null_inner_slot.is_some_and(|s| s + 1 == inner_len); - let max_key_count = inner_len - null_is_last as usize; + 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!( @@ -424,32 +427,61 @@ impl GroupColumn } fn take_n(&mut self, n: usize) -> ArrayRef { - // `inner` is a trait object — the only way to extract its data is via `take_n`. - // Because group->inner slot mappings are non-contiguous, we drain all of `inner` - // at once, then re-append only the slots still referenced by the remaining groups. let old_inner_len = self.inner.len(); let all_inner_values = self.inner.take_n(old_inner_len); - let emitted = - Self::into_dict(Arc::clone(&all_inner_values), &self.group_to_inner[..n]); + let mut emit_old_to_new = vec![usize::MAX; old_inner_len]; + let mut emit_new_to_old: Vec = 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 = 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])) + } + }) + .collect(); + let emitted: ArrayRef = + Arc::new(DictionaryArray::::new(emitted_keys, compact_emit_values)); + // Null deferred to last so null_inner_slot is always the highest index + // and check_key_overflow can subtract it without a false overflow. let remaining = self.group_to_inner[n..].to_vec(); - - // Map each referenced old slot to a new contiguous index. let mut old_to_new = vec![usize::MAX; old_inner_len]; let mut new_to_old: Vec = Vec::new(); + let mut null_old_slot: Option = None; for &old in &remaining { + if all_inner_values.is_null(old) { + if null_old_slot.is_none() { + null_old_slot = Some(old); + } + continue; + } if old_to_new[old] == usize::MAX { old_to_new[old] = new_to_old.len(); new_to_old.push(old); } } + if let Some(old) = null_old_slot { + old_to_new[old] = new_to_old.len(); + new_to_old.push(old); + } self.value_dedup = HashTable::new(); self.value_dedup_size = 0; self.null_inner_slot = None; - // Hash all surviving values in one vectorized pass self.hash_values(&all_inner_values); for (new_slot, &old_slot) in new_to_old.iter().enumerate() { @@ -762,6 +794,54 @@ mod tests { .unwrap(); } + // take_n compacts emitted values to only those referenced by the emitted groups. + #[test] + fn take_n_emits_compact_values() { + let mut col = utf8_col(); + // Four groups: [a, b, c, a] — three distinct values + let arr = dict_arr( + &[Some(0), Some(1), Some(2), Some(0)], + &[Some("a"), Some("b"), Some("c")], + ); + col.vectorized_append(&arr, &[0, 1, 2, 3]).unwrap(); + + // Emit 2 groups (a, b); "c" is only referenced by the remaining group. + let emitted = col.take_n(2); + + // Emitted values array must contain only "a" and "b", not "c". + assert_eq!(emitted.as_dictionary::().values().len(), 2); + assert_eq!( + str_values(&emitted), + vec![Some("a".into()), Some("b".into())] + ); + // Remaining group still resolves correctly. + let out = Box::new(col).build(); + assert_eq!(str_values(&out), vec![Some("c".into()), Some("a".into())]); + } + + // Regression: take_n must not panic when null appears before a non-null + // slot in remaining groups at Int8 key capacity (EmitTo::First boundary). + #[test] + fn take_n_key_limit_null_first_in_remaining() { + let mut col = int8_utf8_col(); + + // Fill Int8 capacity: 128 non-null values at slots 0..127. + col.vectorized_append(&distinct_strs(0, 128), &(0..128).collect::>()) + .unwrap(); + // Null at slot 128 (last) — still valid for Int8. + col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + .unwrap(); + // Re-append v0..v127 (reuses slots 0..127 via dedup). + // group_to_inner is now [0..127, 128(null), 0..127]. + col.vectorized_append(&distinct_strs(0, 128), &(0..128).collect::>()) + .unwrap(); + + // take_n(1) emits group 0 (v0). In remaining, null (old slot 128) appears + // before old slot 0 — without the null-last fix this panics. + let emitted = col.take_n(1); + assert_eq!(str_values(&emitted), vec![Some("v0".into())]); + } + // build_lookup_table must use the incoming batch's hashes, not // stale ones left by the last vectorized_append call. #[test] diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 7aeb97a0ea7af..6f4c212e92682 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -994,7 +994,7 @@ fn make_group_column(field: &Field) -> Result> { let nullable = field.is_nullable(); let data_type = field.data_type(); let mut v: Vec> = Vec::with_capacity(1); - match data_type { + match *data_type { DataType::Int8 => instantiate_primitive!(v, nullable, Int8Type, data_type), DataType::Int16 => instantiate_primitive!(v, nullable, Int16Type, data_type), DataType::Int32 => instantiate_primitive!(v, nullable, Int32Type, data_type), @@ -1114,7 +1114,7 @@ fn make_group_column(field: &Field) -> Result> { v.push(Box::new(BooleanGroupValueBuilder::::new())); } } - DataType::Dictionary(key_dt, value_dt) => { + DataType::Dictionary(ref key_dt, ref value_dt) => { let new_field = Field::new("", *value_dt.clone(), true); let inner = make_group_column(&new_field)?; macro_rules! dict_col { From 44ad7d11ae2da9793c265021c7fd9a56add00dc7 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Mon, 3 Aug 2026 23:45:45 -0400 Subject: [PATCH 7/8] fix merge conflicts, add test for dictionary overflow, trim LOC --- .../group_values/multi_group_by/dictionary.rs | 557 ++++++++---------- .../group_values/multi_group_by/mod.rs | 4 +- 2 files changed, 246 insertions(+), 315 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index cc9b078ce38e4..501b13d0cd183 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -77,36 +77,67 @@ impl DictionaryGroupValuesColumn { } } - fn into_dict(values: ArrayRef, group_to_inner: &[usize]) -> ArrayRef { + /// Build a `DictionaryArray` from `values` (all inner slots) and the + /// per-group slot mapping. The null inner slot, if any, is excluded from + /// the values array and its groups emit a null key — so it never consumes + /// a key index regardless of where it sits in `inner`. + fn into_dict( + values: ArrayRef, + group_to_inner: &[usize], + null_inner_slot: Option, + ) -> ArrayRef { + let Some(null_slot) = null_inner_slot else { + // Fast path: no null group — raw slot indices are valid keys. + let keys: PrimitiveArray = group_to_inner + .iter() + .map(|&slot| Some(K::Native::usize_as(slot))) + .collect(); + return Arc::new(DictionaryArray::::new(keys, values)); + }; + + // Build a compact remap: each non-null slot gets a contiguous key + // starting from 0; the null slot is skipped entirely. + let n = values.len(); + let mut remap = vec![0usize; n]; + let mut next = 0usize; + for (i, mapped) in remap.iter_mut().enumerate() { + if i != null_slot { + *mapped = next; + next += 1; + } + } + let keys: PrimitiveArray = group_to_inner .iter() .map(|&slot| { - if values.is_null(slot) { + if slot == null_slot { None } else { - Some(K::Native::usize_as(slot)) + Some(K::Native::usize_as(remap[slot])) } }) .collect(); - Arc::new(DictionaryArray::::new(keys, values)) + + // Compact values array: drop the null slot so key indices stay tight. + let compact_indices: Int64Array = (0..n) + .filter(|&i| i != null_slot) + .map(|i| i as i64) + .collect(); + let compact = + take(&*values, &compact_indices, None).expect("compact values in into_dict"); + Arc::new(DictionaryArray::::new(keys, compact)) } // 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. + // Null groups emit a null key (None), not a slot index, so the null inner + // slot never consumes a key index regardless of its position in inner. 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; + let non_null_count = self.inner.len() - self.null_inner_slot.is_some() as usize; + if !Self::key_type_fits(non_null_count) { return exec_err!( "Dictionary key type {:?} cannot represent {} distinct values", K::DATA_TYPE, - non_null_slots + non_null_count ); } Ok(()) @@ -422,8 +453,9 @@ impl GroupColumn } fn build(self: Box) -> ArrayRef { + let null_inner_slot = self.null_inner_slot; let values = self.inner.build(); - Self::into_dict(values, &self.group_to_inner) + Self::into_dict(values, &self.group_to_inner, null_inner_slot) } fn take_n(&mut self, n: usize) -> ArrayRef { @@ -433,6 +465,12 @@ impl GroupColumn let mut emit_old_to_new = vec![usize::MAX; old_inner_len]; let mut emit_new_to_old: Vec = Vec::new(); for &old in &self.group_to_inner[..n] { + // Null groups emit a null key (None) and need no slot in the + // values array, so excluding them keeps key indices tight and + // prevents overflow at key-type capacity. + if all_inner_values.is_null(old) { + continue; + } if emit_old_to_new[old] == usize::MAX { emit_old_to_new[old] = emit_new_to_old.len(); emit_new_to_old.push(old); @@ -518,30 +556,64 @@ mod tests { UInt8Array, }; use arrow::compute::cast; - use arrow::datatypes::{DataType, Int32Type, UInt8Type}; + use arrow::datatypes::{DataType, Int8Type, Int32Type, UInt8Type}; use datafusion_physical_expr::binary_map::OutputType; use std::sync::Arc; fn utf8_col() -> DictionaryGroupValuesColumn { - let field = Field::new("", DataType::Utf8, true); - DictionaryGroupValuesColumn::::new( + let f = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::new( Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), - &field, + &f, ) } - fn dict_arr(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + fn int8_col() -> DictionaryGroupValuesColumn { + let f = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &f, + ) + } + + fn uint8_col() -> DictionaryGroupValuesColumn { + let f = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &f, + ) + } + + fn i32_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { Arc::new(DictionaryArray::::new( Int32Array::from(keys.to_vec()), Arc::new(StringArray::from(values.to_vec())), )) } + fn i8_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + use arrow::array::Int8Array; + Arc::new(DictionaryArray::::new( + Int8Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + + fn u8_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + Arc::new(DictionaryArray::::new( + UInt8Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + fn str_values(arr: &ArrayRef) -> Vec> { let plain = cast(arr.as_ref(), &DataType::Utf8).unwrap(); - let strings = plain.as_any().downcast_ref::().unwrap(); - (0..strings.len()) - .map(|i| strings.is_valid(i).then(|| strings.value(i).to_owned())) + plain + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|v| v.map(|s| s.to_owned())) .collect() } @@ -555,356 +627,215 @@ mod tests { buf } - // Null key and null-valued dict entry both map to the null group. + // Builds an Int8-keyed dict of `end-start` distinct strings "v{start}".."v{end-1}". + fn distinct_i8_dict(start: usize, end: usize) -> ArrayRef { + let strs: Vec = (start..end).map(|i| format!("v{i}")).collect(); + let refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + i8_dict( + &(0..strs.len()).map(|i| Some(i as i8)).collect::>(), + &refs, + ) + } + + // Builds a UInt8-keyed dict of `count` distinct strings "u0".."u{count-1}". + fn distinct_u8_dict(count: usize) -> ArrayRef { + let strs: Vec = (0..count).map(|i| format!("u{i}")).collect(); + let refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + u8_dict( + &(0..count).map(|i| Some(i as u8)).collect::>(), + &refs, + ) + } + + #[test] + fn repeated_values_are_deduplicated_in_inner_store() { + let mut col = utf8_col(); + let arr = i32_dict( + &[Some(0), Some(1), Some(0), Some(1), Some(0)], + &[Some("a"), Some("b")], + ); + col.vectorized_append(&arr, &[0, 1, 2, 3, 4]).unwrap(); + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 2); + assert_eq!( + str_values(&out), + vec![ + Some("a".into()), + Some("b".into()), + Some("a".into()), + Some("b".into()), + Some("a".into()), + ] + ); + } + #[test] - fn null_key_and_null_value_in_dict() { + fn null_key_and_null_valued_entry_both_map_to_null_group() { let mut col = utf8_col(); - // Row 0: null key, Row 1: key→null value, Row 2: key→"b" - let input = dict_arr(&[None, Some(0), Some(1)], &[None, Some("b")]); + let input = i32_dict(&[None, Some(0), Some(1)], &[None, Some("b")]); for row in 0..3 { col.append_val(&input, row).unwrap(); } - assert!(col.equal_to(0, &input, 1)); - assert!(col.equal_to(1, &input, 0)); assert!(!col.equal_to(0, &input, 2)); - assert!(!col.equal_to(2, &input, 0)); - let out = Box::new(col).build(); - assert_eq!(out.as_dictionary::().values().len(), 2); + assert_eq!(out.as_dictionary::().values().len(), 1); assert_eq!(str_values(&out), vec![None, None, Some("b".into())]); } #[test] - fn take_n_remaps_slots_across_batches() { - use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; - use arrow::array::UInt64Array; - use arrow::datatypes::UInt64Type; - - let field = Field::new("", DataType::UInt64, true); - let mut col = DictionaryGroupValuesColumn::::new( - Box::new(PrimitiveGroupValueBuilder::::new( - DataType::UInt64, - )), - &field, + fn take_n_compacts_emitted_values_and_remaps_remaining_slots() { + let mut col = utf8_col(); + let b1 = i32_dict( + &[Some(0), Some(1), None, Some(2)], + &[Some("a"), Some("b"), Some("c")], ); + col.vectorized_append(&b1, &[0, 1, 2, 3]).unwrap(); - let u64_val = |arr: &ArrayRef, pos: usize| { - let dict = arr.as_dictionary::(); - dict.values() - .as_any() - .downcast_ref::() - .unwrap() - .value(dict.key(pos).unwrap()) - }; - - let batch1: ArrayRef = Arc::new(DictionaryArray::::new( - Int32Array::from(vec![Some(0), Some(1), None, Some(2), Some(0)]), - Arc::new(UInt64Array::from(vec![10u64, 20, 30])), - )); - col.vectorized_append(&batch1, &[0, 1, 2, 3, 4]).unwrap(); - - let emitted = col.take_n(3); - assert_eq!(u64_val(&emitted, 0), 10); - assert_eq!(u64_val(&emitted, 1), 20); - assert!(emitted.as_dictionary::().key(2).is_none()); + let emitted = col.take_n(2); + assert_eq!(emitted.as_dictionary::().values().len(), 2); + assert_eq!( + str_values(&emitted), + vec![Some("a".into()), Some("b".into())] + ); - let batch2: ArrayRef = Arc::new(DictionaryArray::::new( - Int32Array::from(vec![None, Some(0)]), - Arc::new(UInt64Array::from(vec![99u64])), - )); - col.vectorized_append(&batch2, &[0, 1]).unwrap(); + let b2 = i32_dict(&[None, Some(0)], &[Some("z")]); + col.vectorized_append(&b2, &[0, 1]).unwrap(); let mut buf = all_true(2); - col.vectorized_equal_to(&[2, 3], &batch2, &[0, 1], &mut buf); - assert_eq!(bool_vec(&buf), vec![true, true]); + col.vectorized_equal_to(&[0, 1], &b2, &[0, 1], &mut buf); + assert_eq!(bool_vec(&buf), vec![true, false]); let out = Box::new(col).build(); - assert_eq!(u64_val(&out, 0), 30); - assert_eq!(u64_val(&out, 1), 10); - assert!(out.as_dictionary::().key(2).is_none()); - assert_eq!(u64_val(&out, 3), 99); - } - - // Regression: https://github.com/apache/datafusion/issues/23127 - #[test] - fn key_type_overflow_returns_error() { - let field = Field::new("", DataType::Utf8, true); - let mut col = DictionaryGroupValuesColumn::::new( - Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), - &field, + assert_eq!( + str_values(&out), + vec![None, Some("c".into()), None, Some("z".into())] ); - - let strs: Vec = (0..=255u16).map(|i| i.to_string()).collect(); - let str_refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); - let full: ArrayRef = Arc::new(DictionaryArray::::new( - UInt8Array::from((0..=255u8).map(Some).collect::>()), - Arc::new(StringArray::from(str_refs)), - )); - col.vectorized_append(&full, &(0..256).collect::>()) - .unwrap(); - - let extra: ArrayRef = Arc::new(DictionaryArray::::new( - UInt8Array::from(vec![Some(0u8)]), - Arc::new(StringArray::from(vec![Some("overflow")])), - )); - assert!(col.append_val(&extra, 0).is_err()); } - // A null value alongside 256 non-null values must not itself trigger an - // overflow: null groups emit a null key, not a key index. Adding a 257th - // non-null value after the null should be what triggers the error. #[test] - fn null_does_not_count_toward_key_overflow() { - let field = Field::new("", DataType::Utf8, true); - let mut col = DictionaryGroupValuesColumn::::new( - Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), - &field, - ); - - // Fill all 256 UInt8 key slots (indices 0..=255) with distinct non-null values. - let strs: Vec = (0..=255u16).map(|i| i.to_string()).collect(); - let str_refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); - let full: ArrayRef = Arc::new(DictionaryArray::::new( - UInt8Array::from((0..=255u8).map(Some).collect::>()), - Arc::new(StringArray::from(str_refs)), - )); - col.vectorized_append(&full, &(0..256).collect::>()) + fn vectorized_equal_to_does_not_use_stale_hashes_from_prior_append() { + let mut col = utf8_col(); + col.vectorized_append(&i32_dict(&[Some(0)], &[Some("a"), Some("b")]), &[0]) .unwrap(); - - // Null does not consume a key index — appending it must succeed. - let null_arr: ArrayRef = Arc::new(DictionaryArray::::new( - UInt8Array::from(vec![None]), - Arc::new(StringArray::from(vec![Some("dummy")])), - )); - col.append_val(&null_arr, 0).unwrap(); - - // A 257th distinct non-null value now exceeds UInt8's capacity. - let extra: ArrayRef = Arc::new(DictionaryArray::::new( - UInt8Array::from(vec![Some(0u8)]), - Arc::new(StringArray::from(vec![Some("overflow")])), - )); - assert!(col.append_val(&extra, 0).is_err()); + let batch2 = i32_dict(&[Some(1)], &[Some("z"), Some("a")]); + let mut buf = all_true(1); + col.vectorized_equal_to(&[0], &batch2, &[0], &mut buf); + assert_eq!(bool_vec(&buf), vec![true]); } - // Regression: null mid-array means the last slot is non-null, so the max - // emitted key equals inner.len()-1, not inner.len()-2. #[test] - fn key_overflow_null_slot_mid_array() { - use arrow::array::Int8Array; - use arrow::datatypes::Int8Type; - - let field = Field::new("", DataType::Utf8, true); - let mut col = DictionaryGroupValuesColumn::::new( - Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), - &field, - ); + fn null_does_not_consume_a_key_slot_int8_null_first_mid_and_last() { + let rows128 = (0..128).collect::>(); - // 100 non-null values in slots 0..99 - let strs100: Vec = (0u8..100).map(|i| format!("s{i}")).collect(); - let refs100: Vec> = - strs100.iter().map(|s| Some(s.as_str())).collect(); - let b1: ArrayRef = Arc::new(DictionaryArray::::new( - Int8Array::from((0i8..100).map(Some).collect::>()), - Arc::new(StringArray::from(refs100)), - )); - col.vectorized_append(&b1, &(0usize..100).collect::>()) + let mut col = int8_col(); // null-last: 128 non-null + null — ok + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) .unwrap(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); - // null lands at slot 100 (mid-array) - let null_batch: ArrayRef = Arc::new(DictionaryArray::::new( - Int8Array::from(vec![None]), - Arc::new(StringArray::from(vec![Some("x")])), - )); - col.append_val(&null_batch, 0).unwrap(); - - // 28 more non-null values fill slots 101..128; slot 128 exceeds Int8::MAX - let strs28: Vec = (100u8..128).map(|i| format!("s{i}")).collect(); - let refs28: Vec> = strs28.iter().map(|s| Some(s.as_str())).collect(); - let b2: ArrayRef = Arc::new(DictionaryArray::::new( - Int8Array::from((0i8..28).map(Some).collect::>()), - Arc::new(StringArray::from(refs28)), - )); + let mut col = int8_col(); // null-first: null + 128 non-null — ok; 129th — error + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); assert!( - col.vectorized_append(&b2, &(0usize..28).collect::>()) + col.append_val(&i8_dict(&[Some(0)], &[Some("overflow")]), 0) .is_err() ); - } - // Helpers shared by the three overflow boundary tests below. - fn int8_utf8_col() -> DictionaryGroupValuesColumn { - use arrow::datatypes::Int8Type; - let field = Field::new("", DataType::Utf8, true); - DictionaryGroupValuesColumn::::new( - Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), - &field, - ) - } - - fn int8_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { - use arrow::array::Int8Array; - use arrow::datatypes::Int8Type; - Arc::new(DictionaryArray::::new( - Int8Array::from(keys.to_vec()), - Arc::new(StringArray::from(values.to_vec())), - )) - } - - fn distinct_strs(start: usize, end: usize) -> ArrayRef { - let strs: Vec = (start..end).map(|i| format!("v{i}")).collect(); - let refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); - int8_dict( - &(0..strs.len()).map(|i| Some(i as i8)).collect::>(), - &refs, - ) - } - - // null mid-array: last slot is non-null, so slot 128 overflows Int8 - #[test] - fn overflow_null_mid_pushes_over() { - let mut col = int8_utf8_col(); - // 100 non-null, then null at slot 100, then 28 more non-null (slots 101..128) - col.vectorized_append(&distinct_strs(0, 100), &(0..100).collect::>()) + let mut col = int8_col(); // null-mid: 100 + null + 28 = 128 total — ok; 129th — error + col.vectorized_append(&distinct_i8_dict(0, 100), &(0..100).collect::>()) .unwrap(); - col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(100, 128), &(0..28).collect::>()) .unwrap(); assert!( - col.vectorized_append(&distinct_strs(100, 128), &(0..28).collect::>()) + col.append_val(&i8_dict(&[Some(0)], &[Some("v128")]), 0) .is_err() ); - } - // null last: 128 non-null values fill slots 0..127 (Int8::MAX), null at 128 is ok - #[test] - fn overflow_null_last_at_max() { - let mut col = int8_utf8_col(); - col.vectorized_append(&distinct_strs(0, 128), &(0..128).collect::>()) - .unwrap(); - col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + let mut col = int8_col(); // build() null-first: 128 values (null excluded), null → None + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) .unwrap(); + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 128); + assert_eq!(str_values(&out)[0], None); + assert_eq!(str_values(&out)[1], Some("v0".into())); } - // null last: 127 non-null values, one below the limit #[test] - fn overflow_null_last_just_below_max() { - let mut col = int8_utf8_col(); - col.vectorized_append(&distinct_strs(0, 127), &(0..127).collect::>()) - .unwrap(); - col.append_val(&int8_dict(&[None], &[Some("x")]), 0) - .unwrap(); - } + fn null_does_not_consume_a_key_slot_uint8_null_first_and_last() { + let rows256 = (0..256).collect::>(); - // take_n compacts emitted values to only those referenced by the emitted groups. - #[test] - fn take_n_emits_compact_values() { - let mut col = utf8_col(); - // Four groups: [a, b, c, a] — three distinct values - let arr = dict_arr( - &[Some(0), Some(1), Some(2), Some(0)], - &[Some("a"), Some("b"), Some("c")], + let mut col = uint8_col(); // null-first: null + 256 non-null — ok; 257th — error + col.append_val(&u8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_u8_dict(256), &rows256) + .unwrap(); + assert!( + col.append_val(&u8_dict(&[Some(0)], &[Some("overflow")]), 0) + .is_err() ); - col.vectorized_append(&arr, &[0, 1, 2, 3]).unwrap(); - // Emit 2 groups (a, b); "c" is only referenced by the remaining group. - let emitted = col.take_n(2); - - // Emitted values array must contain only "a" and "b", not "c". - assert_eq!(emitted.as_dictionary::().values().len(), 2); - assert_eq!( - str_values(&emitted), - vec![Some("a".into()), Some("b".into())] - ); - // Remaining group still resolves correctly. + let mut col = uint8_col(); // build() null-first: 256 values (null excluded), last correct + col.append_val(&u8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_u8_dict(256), &rows256) + .unwrap(); let out = Box::new(col).build(); - assert_eq!(str_values(&out), vec![Some("c".into()), Some("a".into())]); + assert_eq!(out.as_dictionary::().values().len(), 256); + assert_eq!(str_values(&out)[0], None); + assert_eq!(str_values(&out)[256], Some("u255".into())); } - // Regression: take_n must not panic when null appears before a non-null - // slot in remaining groups at Int8 key capacity (EmitTo::First boundary). #[test] - fn take_n_key_limit_null_first_in_remaining() { - let mut col = int8_utf8_col(); + fn take_n_null_does_not_steal_key_slot_at_capacity() { + let rows128 = (0..128).collect::>(); - // Fill Int8 capacity: 128 non-null values at slots 0..127. - col.vectorized_append(&distinct_strs(0, 128), &(0..128).collect::>()) - .unwrap(); - // Null at slot 128 (last) — still valid for Int8. - col.append_val(&int8_dict(&[None], &[Some("x")]), 0) + // Int8 null-first + 128 non-null; emit all 129 — no panic, null → None + let mut col = int8_col(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) .unwrap(); - // Re-append v0..v127 (reuses slots 0..127 via dedup). - // group_to_inner is now [0..127, 128(null), 0..127]. - col.vectorized_append(&distinct_strs(0, 128), &(0..128).collect::>()) + let emitted = col.take_n(129); + assert!(emitted.as_dictionary::().key(0).is_none()); + assert_eq!(str_values(&emitted)[1], Some("v0".into())); + assert_eq!(str_values(&emitted)[128], Some("v127".into())); + + // UInt8 null-first + 256 non-null; emit all 257 — last must be "u255" not "u0" (wrap guard) + let mut col = uint8_col(); + col.append_val(&u8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_u8_dict(256), &(0..256).collect::>()) .unwrap(); - - // take_n(1) emits group 0 (v0). In remaining, null (old slot 128) appears - // before old slot 0 — without the null-last fix this panics. - let emitted = col.take_n(1); - assert_eq!(str_values(&emitted), vec![Some("v0".into())]); + let emitted = col.take_n(257); + assert!(emitted.as_dictionary::().key(0).is_none()); + assert_eq!(str_values(&emitted)[1], Some("u0".into())); + assert_eq!(str_values(&emitted)[256], Some("u255".into())); } - // build_lookup_table must use the incoming batch's hashes, not - // stale ones left by the last vectorized_append call. #[test] - fn vectorized_equal_to_uses_current_batch_hashes() { - let mut col = utf8_col(); + fn take_n_repeated_emissions_null_at_int8_capacity() { + let rows128 = (0..128).collect::>(); + let mut col = int8_col(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); - let batch1 = dict_arr(&[Some(0)], &[Some("a"), Some("b")]); - col.vectorized_append(&batch1, &[0]).unwrap(); + let first_half = col.take_n(64); + assert_eq!(str_values(&first_half)[0], Some("v0".into())); + assert_eq!(str_values(&first_half)[63], Some("v63".into())); + assert_eq!(first_half.as_dictionary::().values().len(), 64); - // values = ["z", "a"]; key 0 → "a" at val_idx 1. - // Stale hashes would probe val_idx 1 with hash("b") and miss. - let batch2 = dict_arr(&[Some(1)], &[Some("z"), Some("a")]); - let mut buf = all_true(1); - col.vectorized_equal_to(&[0], &batch2, &[0], &mut buf); - assert_eq!(bool_vec(&buf), vec![true]); - } + let second_half = col.take_n(64); + assert_eq!(str_values(&second_half)[0], Some("v64".into())); + assert_eq!(str_values(&second_half)[63], Some("v127".into())); - #[test] - fn append_only_stores_referenced_values() { - let mut col = utf8_col(); - let values = Arc::new(StringArray::from(vec![ - Some("a"), - Some("b"), - Some("c"), - Some("d"), - Some("e"), - Some("f"), - Some("g"), - Some("h"), - Some("i"), - Some("j"), - ])); - let keys = Int32Array::from(vec![ - Some(0), // a - Some(2), // c - Some(7), // h - Some(0), - Some(2), - Some(7), - Some(7), - Some(0), - Some(2), - ]); - let input: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); - - col.vectorized_append(&input, &[0, 1, 2, 3, 4, 5, 6, 7, 8]) - .unwrap(); + let null_group = col.take_n(1); + assert!(null_group.as_dictionary::().key(0).is_none()); let out = Box::new(col).build(); - assert_eq!(out.as_dictionary::().values().len(), 3); - assert_eq!( - str_values(&out), - vec![ - Some("a".into()), - Some("c".into()), - Some("h".into()), - Some("a".into()), - Some("c".into()), - Some("h".into()), - Some("h".into()), - Some("a".into()), - Some("c".into()), - ] - ); + assert_eq!(str_values(&out)[0], Some("v0".into())); + assert_eq!(str_values(&out)[127], Some("v127".into())); + assert_eq!(out.as_dictionary::().values().len(), 128); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 6f4c212e92682..76482eb209235 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -1925,8 +1925,8 @@ mod tests { let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); // Each row has a unique label (forcing a new group) and alternates - // between the two dictionary values. Int16 keys are used so that - // 129 groups don't hit the Int8 overflow limit (i8::MAX = 127). + // between the two dictionary values. Int8 keys are used; only 2 + // distinct values exist so the key type never overflows. let labels: ArrayRef = Arc::new(StringArray::from( (0..n_groups).map(|i| format!("g{i}")).collect::>(), )); From 63d4021b079355483bc9dbd2bcc3c8bcfb690b8f Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Tue, 4 Aug 2026 00:14:46 -0400 Subject: [PATCH 8/8] fix ci --- .../src/aggregates/group_values/multi_group_by/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 76482eb209235..b60f91d7c6363 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -1650,6 +1650,7 @@ mod tests { None, )), ), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Float16)), ]; for dt in &supported_cases { @@ -1677,7 +1678,6 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), - DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Float16)), DataType::Dictionary( Box::new(DataType::Int32), Box::new(DataType::Decimal256(76, 10)),