diff --git a/datafusion/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs index 2af3bbfd5d2e7..29d5c66afa50c 100644 --- a/datafusion/functions-aggregate/src/min_max.rs +++ b/datafusion/functions-aggregate/src/min_max.rs @@ -56,6 +56,77 @@ fn get_min_max_result_type(input_types: &[DataType]) -> Result> { } } +/// Return the physical input type accepted by min/max. +/// +/// Dictionary-encoded Utf8 can be consumed directly by the string +/// accumulators. Other dictionary and run-end encoded inputs retain the +/// existing behavior and are coerced to their logical value type. +fn get_min_max_input_type(input_types: &[DataType]) -> Result> { + if input_types.len() != 1 { + return exec_err!( + "min/max was called with {} arguments. It requires only 1.", + input_types.len() + ); + } + + match &input_types[0] { + DataType::Dictionary(_, value_type) if value_type.as_ref() == &DataType::Utf8 => { + Ok(input_types.to_vec()) + } + _ => get_min_max_result_type(input_types), + } +} + +/// Adapts the single-group accumulator when type coercion preserves a +/// dictionary input. The inner accumulator compares dictionary values +/// directly, while aggregate state and output use the logical value type. +#[derive(Debug)] +struct DictionaryInputAccumulator { + inner: A, +} + +impl DictionaryInputAccumulator { + fn new(inner: A) -> Self { + Self { inner } + } +} + +impl Accumulator for DictionaryInputAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.inner.update_batch(values) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + self.inner.merge_batch(states) + } + + fn state(&mut self) -> Result> { + Ok(self.inner.state()?.iter().map(logical_scalar).collect()) + } + + fn evaluate(&mut self) -> Result { + self.inner.evaluate().map(|value| logical_scalar(&value)) + } + + fn size(&self) -> usize { + size_of_val(self) - size_of_val(&self.inner) + self.inner.size() + } +} + +fn has_dictionary_input(args: &AccumulatorArgs<'_>) -> bool { + matches!( + args.expr_fields.first().map(|field| field.data_type()), + Some(DataType::Dictionary(_, _)) + ) +} + +fn logical_scalar(value: &ScalarValue) -> ScalarValue { + match value { + ScalarValue::Dictionary(_, value) => logical_scalar(value), + value => value.clone(), + } +} + #[user_doc( doc_section(label = "General Functions"), description = "Returns the maximum value in the specified column.", @@ -201,7 +272,7 @@ impl FromColumnStatistics for Max { if let Precision::Exact(ref val) = col_stats.max_value && !val.is_null() { - return Some(val.clone()); + return Some(logical_scalar(val)); } None } @@ -217,7 +288,7 @@ impl AggregateUDFImpl for Max { } fn return_type(&self, arg_types: &[DataType]) -> Result { - Ok(arg_types[0].to_owned()) + Ok(get_min_max_result_type(arg_types)?[0].clone()) } fn return_field(&self, arg_fields: &[FieldRef]) -> Result { @@ -225,9 +296,15 @@ impl AggregateUDFImpl for Max { } fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - Ok(Box::new(MaxAccumulator::try_new( - acc_args.return_field.data_type(), - )?)) + let data_type = acc_args.return_field.data_type(); + if has_dictionary_input(&acc_args) { + let input_type = acc_args.expr_fields[0].data_type(); + Ok(Box::new(DictionaryInputAccumulator::new( + MaxAccumulator::try_new(input_type)?, + ))) + } else { + Ok(Box::new(MaxAccumulator::try_new(data_type)?)) + } } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { @@ -297,7 +374,7 @@ impl AggregateUDFImpl for Max { } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - get_min_max_result_type(arg_types) + get_min_max_input_type(arg_types) } fn reverse_expr(&self) -> datafusion_expr::ReversedUDAF { datafusion_expr::ReversedUDAF::Identical @@ -350,6 +427,10 @@ impl Accumulator for SlidingMaxAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { for idx in 0..values[0].len() { let val = ScalarValue::try_from_array(&values[0], idx)?; + let val = match val { + ScalarValue::Dictionary(_, value) => *value, + value => value, + }; if !val.is_null() { self.moving_max.push(val); } @@ -433,7 +514,7 @@ impl FromColumnStatistics for Min { if let Precision::Exact(ref val) = col_stats.min_value && !val.is_null() { - return Some(val.clone()); + return Some(logical_scalar(val)); } None } @@ -449,7 +530,7 @@ impl AggregateUDFImpl for Min { } fn return_type(&self, arg_types: &[DataType]) -> Result { - Ok(arg_types[0].to_owned()) + Ok(get_min_max_result_type(arg_types)?[0].clone()) } fn return_field(&self, arg_fields: &[FieldRef]) -> Result { @@ -457,9 +538,15 @@ impl AggregateUDFImpl for Min { } fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - Ok(Box::new(MinAccumulator::try_new( - acc_args.return_field.data_type(), - )?)) + let data_type = acc_args.return_field.data_type(); + if has_dictionary_input(&acc_args) { + let input_type = acc_args.expr_fields[0].data_type(); + Ok(Box::new(DictionaryInputAccumulator::new( + MinAccumulator::try_new(input_type)?, + ))) + } else { + Ok(Box::new(MinAccumulator::try_new(data_type)?)) + } } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { @@ -533,7 +620,7 @@ impl AggregateUDFImpl for Min { } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - get_min_max_result_type(arg_types) + get_min_max_input_type(arg_types) } fn reverse_expr(&self) -> datafusion_expr::ReversedUDAF { @@ -587,6 +674,10 @@ impl Accumulator for SlidingMinAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { for idx in 0..values[0].len() { let val = ScalarValue::try_from_array(&values[0], idx)?; + let val = match val { + ScalarValue::Dictionary(_, value) => *value, + value => value, + }; if !val.is_null() { self.moving_min.push(val); } @@ -1078,6 +1169,7 @@ mod tests { let expr: Arc = Arc::new(expressions::Column::new("v", 0)); let expr_field = Arc::new(Field::new("v", data_type.clone(), true)); + let return_type = udf.return_type(std::slice::from_ref(&data_type))?; let args = || AccumulatorArgs { schema: &schema, expr_fields: std::slice::from_ref(&expr_field), @@ -1086,7 +1178,7 @@ mod tests { name: udf.name(), ignore_nulls: false, is_reversed: false, - return_field: Arc::new(Field::new(udf.name(), data_type.clone(), true)), + return_field: Arc::new(Field::new(udf.name(), return_type.clone(), true)), order_bys: &[], }; @@ -1112,10 +1204,27 @@ mod tests { second.merge_batch(&state, &merge_group_indices, total_num_groups)?; let result = second.evaluate(EmitTo::All)?; - assert_eq!(result.data_type(), &data_type); + assert_eq!(result.data_type(), &return_type); Ok(result) } + #[test] + fn dictionary_min_max_groups_accumulator_returns_plain_values() -> Result<()> { + let values = Arc::new(StringArray::from(vec!["z", "b", "a", "c"])) as ArrayRef; + let dictionary = Arc::new(DictionaryArray::try_new( + Int32Array::from(vec![0, 1, 2, 3]), + values, + )?) as ArrayRef; + let groups = [0, 1, 0, 1]; + + let min = evaluate_grouped(&Min::new(), &dictionary, &groups, 2)?; + assert_eq!(min.as_string::(), &StringArray::from(vec!["a", "b"])); + + let max = evaluate_grouped(&Max::new(), &dictionary, &groups, 2)?; + assert_eq!(max.as_string::(), &StringArray::from(vec!["z", "c"])); + Ok(()) + } + #[test] fn interval_min_max_groups_accumulator() -> Result<()> { // Group 0 and 1 get values from both halves of the input, group 2 only @@ -1494,6 +1603,14 @@ mod tests { let result = fun.coerce_types(input_type); assert_eq!(*input_type, result.unwrap()); } + + let dictionary = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + assert_eq!( + fun.coerce_types(std::slice::from_ref(&dictionary)).unwrap(), + vec![dictionary.clone()] + ); + assert_eq!(fun.return_type(&[dictionary]).unwrap(), DataType::Utf8); } } @@ -1561,6 +1678,78 @@ mod tests { Ok(()) } + #[test] + fn sliding_min_max_dictionary_returns_plain_values() -> Result<()> { + let values = Arc::new(StringArray::from(vec!["b", "a", "c"])) as ArrayRef; + let dictionary = Arc::new(DictionaryArray::try_new( + Int32Array::from(vec![0, 1, 2]), + values, + )?) as ArrayRef; + + let mut min = SlidingMinAccumulator::try_new(&DataType::Utf8)?; + min.update_batch(std::slice::from_ref(&dictionary))?; + assert_eq!(min.evaluate()?, ScalarValue::Utf8(Some("a".to_string()))); + min.retract_batch(&[dictionary.slice(0, 2)])?; + assert_eq!(min.evaluate()?, ScalarValue::Utf8(Some("c".to_string()))); + + let mut max = SlidingMaxAccumulator::try_new(&DataType::Utf8)?; + max.update_batch(std::slice::from_ref(&dictionary))?; + assert_eq!(max.evaluate()?, ScalarValue::Utf8(Some("c".to_string()))); + Ok(()) + } + + #[test] + fn dictionary_input_uses_public_accumulator_contract() -> Result<()> { + let values = Arc::new(StringArray::from(vec!["b", "a", "c"])) as ArrayRef; + let dictionary = Arc::new(DictionaryArray::try_new( + Int32Array::from(vec![0, 1, 2, 1]), + values, + )?) as ArrayRef; + let dictionary_type = dictionary.data_type().clone(); + let schema = Schema::new(vec![Field::new("v", dictionary_type.clone(), true)]); + let expr: Arc = + Arc::new(expressions::Column::new("v", 0)); + let expr_field = Arc::new(Field::new("v", dictionary_type, true)); + + let min = Min::new(); + let max = Max::new(); + for (udf, expected) in [ + ( + &min as &dyn AggregateUDFImpl, + ScalarValue::Utf8(Some("a".to_string())), + ), + ( + &max as &dyn AggregateUDFImpl, + ScalarValue::Utf8(Some("c".to_string())), + ), + ] { + let return_field = udf.return_field(std::slice::from_ref(&expr_field))?; + assert_eq!(return_field.data_type(), &DataType::Utf8); + let args = || AccumulatorArgs { + schema: &schema, + expr_fields: std::slice::from_ref(&expr_field), + exprs: std::slice::from_ref(&expr), + is_distinct: false, + name: udf.name(), + ignore_nulls: false, + is_reversed: false, + return_field: Arc::clone(&return_field), + order_bys: &[], + }; + + let mut accumulator = udf.accumulator(args())?; + accumulator.update_batch(&[Arc::clone(&dictionary)])?; + assert_eq!(accumulator.evaluate()?, expected); + + let state = accumulator.state()?; + assert_eq!(state, vec![expected.clone()]); + let mut merged = udf.accumulator(args())?; + merged.merge_batch(&[ScalarValue::iter_to_array(state)?])?; + assert_eq!(merged.evaluate()?, expected); + } + Ok(()) + } + fn dict_scalar(key_type: DataType, inner: ScalarValue) -> ScalarValue { ScalarValue::Dictionary(Box::new(key_type), Box::new(inner)) } diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index dc6544440b2b6..ac4fdae552c2e 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -4,9 +4,12 @@ use arrow::array::{ Array, ArrayRef, AsArray, BinaryBuilder, BinaryViewBuilder, BooleanArray, - LargeBinaryBuilder, LargeStringBuilder, StringBuilder, StringViewBuilder, + DictionaryArray, LargeBinaryBuilder, LargeStringBuilder, StringArray, StringBuilder, + StringViewBuilder, }; -use arrow::datatypes::DataType; +use arrow::compute::cast; +use arrow::datatypes::{ArrowDictionaryKeyType, DataType}; +use arrow::downcast_dictionary_array; use datafusion_common::hash_map::Entry; use datafusion_common::{HashMap, Result, internal_err}; use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; @@ -143,6 +146,55 @@ impl MinMaxBytesAccumulator { assert_eq!(&self.inner.data_type, result.data_type()); Ok(result) } + + fn update_dictionary( + &mut self, + dictionary: &DictionaryArray, + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + cmp: F, + ) -> Result<()> + where + K: ArrowDictionaryKeyType, + F: FnMut(&[u8], &[u8]) -> bool + Send + Sync, + { + fn str_to_bytes<'a>( + it: impl Iterator>, + ) -> impl Iterator> { + it.map(|s| s.map(str::as_bytes)) + } + + fn apply_filter<'a, T: 'a>( + iter: impl Iterator> + 'a, + filter: Option<&'a BooleanArray>, + ) -> impl Iterator> + 'a { + iter.enumerate().map(move |(index, value)| { + if filter + .is_none_or(|filter| filter.is_valid(index) && filter.value(index)) + { + value + } else { + None + } + }) + } + + self.inner.update_batch( + apply_filter( + str_to_bytes( + dictionary + .downcast_dict::() + .expect("dictionary value type was checked") + .into_iter(), + ), + opt_filter, + ), + group_indices, + total_num_groups, + cmp, + ) + } } impl GroupsAccumulator for MinMaxBytesAccumulator { @@ -155,10 +207,17 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { ) -> Result<()> { let array = &values[0]; assert_eq!(array.len(), group_indices.len()); - assert_eq!(array.data_type(), &self.inner.data_type); - - // apply filter if needed - let array = apply_filter_as_nulls(array, opt_filter)?; + let input_type = match array.data_type() { + DataType::Dictionary(_, value_type) => value_type.as_ref(), + data_type => data_type, + }; + if input_type != &self.inner.data_type { + return internal_err!( + "Expected {} input for MinMaxBytesAccumulator, got {}", + self.inner.data_type, + array.data_type() + ); + } // dispatch to appropriate kernel / specialized implementation fn string_min(a: &[u8], b: &[u8]) -> bool { @@ -193,6 +252,31 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { it.map(|s| s.map(|s| s.as_bytes())) } + if matches!(array.data_type(), DataType::Dictionary(_, _)) { + if self.inner.data_type != DataType::Utf8 { + return internal_err!( + "Unexpected dictionary value type for MinMaxBytesAccumulator: {}", + self.inner.data_type + ); + } + let cmp: fn(&[u8], &[u8]) -> bool = + if self.is_min { string_min } else { string_max }; + let array = array.as_ref(); + return downcast_dictionary_array! { + array => self.update_dictionary( + array, + group_indices, + opt_filter, + total_num_groups, + cmp, + ), + _ => unreachable!("dictionary type was checked"), + }; + } + + // apply filter if needed + let array = apply_filter_as_nulls(array, opt_filter)?; + match (self.is_min, &self.inner.data_type) { // Utf8/LargeUtf8/Utf8View Min (true, &DataType::Utf8) => self.inner.update_batch( @@ -342,7 +426,12 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { ) -> Result> { // Min/max do not change the values as they are their own states // apply the filter by combining with the null mask, if any - let output = apply_filter_as_nulls(&values[0], opt_filter)?; + let values = if matches!(values[0].data_type(), DataType::Dictionary(_, _)) { + cast(&values[0], &self.inner.data_type)? + } else { + Arc::clone(&values[0]) + }; + let output = apply_filter_as_nulls(&values, opt_filter)?; Ok(vec![output]) } fn size(&self) -> usize { @@ -540,7 +629,8 @@ impl MinMaxBytesState { #[cfg(test)] mod tests { use super::*; - use arrow::array::StringArray; + use arrow::array::{DictionaryArray, Int8Array, StringArray, UInt64Array}; + use arrow::datatypes::{Int8Type, UInt64Type}; fn expected_size(state: &MinMaxBytesState) -> usize { state.min_max.capacity() * size_of::>>() @@ -638,4 +728,64 @@ mod tests { ); Ok(()) } + + #[test] + fn dictionary_updates_filters_and_plain_state_merge() -> Result<()> { + let values = Arc::new(StringArray::from(vec![ + Some("z"), + None, + Some("b"), + Some("ignored"), + ])) as ArrayRef; + let first_values = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), Some(1), None, Some(2), Some(3)]), + values, + )?) as ArrayRef; + let first_groups = [0, 0, 1, 1, 2]; + let filter = BooleanArray::from(vec![true, true, true, true, false]); + + let converted = MinMaxBytesAccumulator::new_min(DataType::Utf8) + .convert_to_state(std::slice::from_ref(&first_values), Some(&filter))?; + assert_eq!(converted[0].data_type(), &DataType::Utf8); + + let mut max = MinMaxBytesAccumulator::new_max(DataType::Utf8); + max.update_batch( + &[Arc::clone(&first_values)], + &first_groups, + Some(&filter), + 3, + )?; + assert_eq!( + max.evaluate(EmitTo::All)?.as_string::(), + &StringArray::from(vec![Some("z"), Some("b"), None]) + ); + + let mut first = MinMaxBytesAccumulator::new_min(DataType::Utf8); + first.update_batch(&[first_values], &first_groups, Some(&filter), 3)?; + let state = first.state(EmitTo::All)?; + assert_eq!(state[0].data_type(), &DataType::Utf8); + assert_eq!( + state[0].as_string::(), + &StringArray::from(vec![Some("z"), Some("b"), None]) + ); + + // A different key width and value mapping must be interpreted using this + // batch's dictionary, rather than relying on key ordering or identity. + let values = Arc::new(StringArray::from(vec!["a", "y", "c"])) as ArrayRef; + let second_values = Arc::new(DictionaryArray::::try_new( + UInt64Array::from(vec![Some(1), Some(0), Some(2)]), + values, + )?) as ArrayRef; + let mut second = MinMaxBytesAccumulator::new_min(DataType::Utf8); + second.update_batch(&[second_values], &[0, 1, 2], None, 3)?; + second.merge_batch(&state, &[0, 1, 2], 3)?; + + let result = second.evaluate(EmitTo::All)?; + assert_eq!(result.data_type(), &DataType::Utf8); + assert_eq!( + result.as_string::(), + &StringArray::from(vec![Some("y"), Some("a"), Some("c")]) + ); + Ok(()) + } }