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..359c4dfcbd949 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -0,0 +1,401 @@ +// 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::compute::concat; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field}; +use datafusion_common::Result; +use hashbrown::HashMap; +use std::marker::PhantomData; +use std::sync::{Arc, Mutex}; + +/// Use the cache path in `vectorized_equal_to` only when the number of unique lhs group IDs +/// is at most this fraction of the total comparisons (1/10 = 10%). +const CACHE_USE_THRESHOLD: usize = 10; + +pub struct DictionaryGroupValuesColumn { + inner: Box, + null_array: ArrayRef, + // Mutex is required because `vectorized_equal_to` takes `&self` (GroupColumn trait + // constraint) but needs to populate this cache on the hot path. + key_to_group_cache: Mutex), bool>>, + last_values: Option, + cached_combined: Option, + _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, + last_values: None, + key_to_group_cache: Mutex::new(HashMap::new()), + cached_combined: None, + _phantom: PhantomData, + } + } + + /// Returns `concat([values, null_sentinel])`, rebuilding only when `values` changes. + #[inline] + fn get_combined(&mut self, values: &ArrayRef) -> Result { + if !self + .last_values + .as_ref() + .is_some_and(|v| Arc::ptr_eq(v, values)) + { + self.cached_combined = + Some(concat(&[values.as_ref(), self.null_array.as_ref()])?); + self.last_values = Some(Arc::clone(values)); + // The incoming dictionary values buffer changed, so cached comparisons are stale. + self.key_to_group_cache.lock().unwrap().clear(); + } + Ok(Arc::clone(self.cached_combined.as_ref().unwrap())) + } + + fn into_dict(values: ArrayRef) -> ArrayRef { + let num_values = values.len(); + assert!( + Self::valid_bounds::(num_values), + "Dictionary key type {:?} cannot hold {} values", + K::DATA_TYPE, + num_values + ); + let keys: PrimitiveArray = (0..num_values) + .map(|i| (!values.is_null(i)).then(|| K::Native::usize_as(i))) + .collect(); + Arc::new(DictionaryArray::::new(keys, values)) + } + + fn valid_bounds(num_values: usize) -> bool { + let max: usize = match T::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 + } +} + +impl GroupColumn + for DictionaryGroupValuesColumn +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + let dict = array.as_dictionary::(); + match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_row, &self.null_array, 0), + Some(key) => self.inner.equal_to(lhs_row, dict.values(), key), + } + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let dict = array.as_dictionary::(); + match dict.key(row) { + None => self.inner.append_val(&self.null_array, 0), + Some(key) => self.inner.append_val(dict.values(), key), + } + } + + 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_values = dict.values(); + + // When lhs group IDs are highly repeated relative to the number of comparisons, + // caching (group_id, dict_key) -> bool avoids redundant inner comparisons. + let unique_lhs_row_count = + lhs_rows.iter().collect::>().len(); + if unique_lhs_row_count <= rhs_rows.len() / CACHE_USE_THRESHOLD { + let mut comparison_cache = self.key_to_group_cache.lock().unwrap(); + for (position, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows).enumerate() + { + if !equal_to_results.get_bit(position) { + continue; + } + let dict_key = dict.key(rhs_row); + let is_equal = *comparison_cache + .entry((lhs_row, dict_key)) + .or_insert_with(|| match dict_key { + None => self.inner.equal_to(lhs_row, &self.null_array, 0), + Some(key) => self.inner.equal_to(lhs_row, dict_values, key), + }); + if !is_equal { + equal_to_results.set_bit(position, false); + } + } + } else { + let dict_keys = dict.keys(); + if dict_keys.null_count() == 0 { + let value_indices: Vec = rhs_rows + .iter() + .map(|&row_index| dict_keys.value(row_index).as_usize()) + .collect(); + self.inner.vectorized_equal_to( + lhs_rows, + dict_values, + &value_indices, + equal_to_results, + ); + } else { + for (position, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows).enumerate() + { + if !equal_to_results.get_bit(position) { + continue; + } + let is_equal = match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_row, &self.null_array, 0), + Some(key) => self.inner.equal_to(lhs_row, dict_values, key), + }; + if !is_equal { + equal_to_results.set_bit(position, false); + } + } + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let dict = array.as_dictionary::(); + let dict_keys = dict.keys(); + + if dict_keys.null_count() == 0 { + let value_indices: Vec = rows + .iter() + .map(|&row_index| dict_keys.value(row_index).as_usize()) + .collect(); + self.inner.vectorized_append(dict.values(), &value_indices) + } else { + // `get_combined` appends the null sentinel as the final element, so any + // null dictionary key maps to that last index. + let combined_with_null_sentinel = self.get_combined(dict.values())?; + let null_sentinel_index = combined_with_null_sentinel.len() - 1; + let value_indices: Vec = rows + .iter() + .map(|&row_index| dict.key(row_index).unwrap_or(null_sentinel_index)) + .collect(); + self.inner + .vectorized_append(&combined_with_null_sentinel, &value_indices) + } + } + + fn len(&self) -> usize { + self.inner.len() + } + + fn size(&self) -> usize { + use std::mem::size_of; + let cache_bytes = self.key_to_group_cache.lock().unwrap().capacity() + * (size_of::<(usize, Option)>() + size_of::() + 1); + self.inner.size() + cache_bytes + } + + fn build(self: Box) -> ArrayRef { + Self::into_dict(self.inner.build()) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + // Drop cache entries for the group IDs being emitted; the remaining + // entries still reference valid (not-yet-taken) group IDs. + self.key_to_group_cache + .get_mut() + .unwrap() + .retain(|(group_id, _), _| *group_id >= n); + Self::into_dict(self.inner.take_n(n)) + } +} + +#[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, + }; + use arrow::compute::cast; + use arrow::datatypes::{DataType, Int32Type}; + 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: &[&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 assert_is_dict_utf8(arr: &ArrayRef) { + assert!( + matches!(arr.data_type(), + DataType::Dictionary(k, v) + if k.as_ref() == &DataType::Int32 && v.as_ref() == &DataType::Utf8), + "expected Dictionary(Int32, Utf8), got {:?}", + arr.data_type() + ); + } + + fn true_buf(len: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(len); + buf.append_n(len, true); + buf + } + + fn buf_to_vec(buf: &BooleanBufferBuilder) -> Vec { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + #[test] + fn utf8_dict_null_handling() { + let mut col = utf8_col(); + let input: ArrayRef = Arc::new(DictionaryArray::::new( + Int32Array::from(vec![None, Some(0), Some(1)]), + Arc::new(StringArray::from(vec![None::<&str>, Some("b")])), + )); + for row in 0..3 { + col.append_val(&input, row).unwrap(); + } + + assert!(col.equal_to(0, &input, 0) && col.equal_to(1, &input, 1)); + assert!(col.equal_to(0, &input, 1)); // null == null + assert!(!col.equal_to(0, &input, 2) && !col.equal_to(2, &input, 0)); + + let mut buf = true_buf(3); + col.vectorized_equal_to(&[0, 1, 2], &input, &[2, 2, 1], &mut buf); + assert_eq!(buf_to_vec(&buf), vec![false, false, false]); + + let mut buf = true_buf(3); + buf.set_bit(0, false); + col.vectorized_equal_to(&[0, 1, 2], &input, &[0, 1, 2], &mut buf); + assert_eq!(buf_to_vec(&buf), vec![false, true, true]); + + assert_eq!(col.take_n(0).len(), 0); + let taken = col.take_n(3); + assert_is_dict_utf8(&taken); + assert_eq!(str_values(&taken), vec![None, None, Some("b".into())]); + + let input2 = dict_arr(&[Some(0), None, Some(1), Some(0), None], &["cat", "dog"]); + col.vectorized_append(&input2, &[0, 1, 2, 3, 4]).unwrap(); + let taken2 = col.take_n(5); + assert_is_dict_utf8(&taken2); + assert_eq!( + str_values(&taken2), + vec![ + Some("cat".into()), + None, + Some("dog".into()), + Some("cat".into()), + None, + ] + ); + + let empty = Box::new(col).build(); + assert_is_dict_utf8(&empty); + assert_eq!(empty.len(), 0); + } + + #[test] + fn u64_primitive_inner_multi_batch() { + 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 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 mut buf = true_buf(5); + col.vectorized_equal_to(&[0, 1, 2, 3, 4], &batch1, &[0, 1, 2, 3, 0], &mut buf); + assert_eq!(buf_to_vec(&buf), vec![true, true, true, true, true]); + + let mut buf = true_buf(2); + col.vectorized_equal_to(&[0, 4], &batch1, &[1, 1], &mut buf); + assert_eq!(buf_to_vec(&buf), vec![false, false]); + + // Helper: resolve key at `pos` through the dictionary and return the u64 value. + 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 taken1 = col.take_n(3); + assert_eq!(u64_val(&taken1, 0), 10); + assert_eq!(u64_val(&taken1, 1), 20); + assert!(taken1.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 = true_buf(2); + col.vectorized_equal_to(&[2, 3], &batch2, &[0, 1], &mut buf); + assert_eq!(buf_to_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); + } +} 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 f275d777c3279..8bd1112559d2a 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; use std::mem::{self, size_of}; @@ -30,7 +31,6 @@ use crate::aggregates::group_values::multi_group_by::{ bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; -use arrow::compute::cast; use arrow::datatypes::{ BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Field, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Schema, SchemaRef, @@ -41,7 +41,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; @@ -955,7 +955,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. @@ -977,7 +977,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), @@ -1067,6 +1067,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) + } _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), } debug_assert_eq!( @@ -1107,7 +1134,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. @@ -1197,20 +1224,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) } @@ -1259,7 +1272,9 @@ enum Nulls { mod tests { use std::{collections::HashMap, sync::Arc}; - use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray}; + use arrow::array::{ + ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray, StringViewArray, + }; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; use datafusion_common::utils::proxy::HashTableAllocExt; @@ -1309,6 +1324,19 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Microsecond), DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond), DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + 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 { @@ -1337,6 +1365,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 { @@ -1466,7 +1499,7 @@ mod tests { // `emit(EmitTo::First(4))` calls can `take_n` without panicking. // The hashmap entries below reference group indices 0..=11, so the // single column builder needs at least 12 rows to back them. - let seed: ArrayRef = Arc::new(arrow::array::Int32Array::from(vec![0_i32; 12])); + let seed: ArrayRef = Arc::new(Int32Array::from(vec![0_i32; 12])); for row in 0..12 { group_values.group_values[0] .append_val(&seed, row)