Refactor dictionary min/max logic and tests

Centralize dictionary batch handling for min/max operations.
Streamline min_max_batch_generic to initialize from the first
non-null element. Implement shared setup/assert helpers in
dictionary tests to reduce repetition while preserving test
coverage.

Author: kosiew

datafusion/functions-aggregate-common/src/min_max.rs

@@ -141,6 +141,16 @@ macro_rules! min_max_generic {
     }};
 }
 
+macro_rules! min_max_dictionary {
+    ($VALUE:expr, $DELTA:expr, wrap $KEY_TYPE:expr, $OP:ident) => {{
+        let winner = min_max_generic!($VALUE, $DELTA, $OP);
+        ScalarValue::Dictionary($KEY_TYPE.clone(), Box::new(winner))
+    }};
+    ($VALUE:expr, $DELTA:expr, $OP:ident) => {{
+        min_max_generic!($VALUE, $DELTA, $OP)
+    }};
+}
+
 // min/max of two scalar values of the same type
 macro_rules! min_max {
     ($VALUE:expr, $DELTA:expr, $OP:ident) => {{
@@ -417,22 +427,26 @@ macro_rules! min_max {
                 ScalarValue::Dictionary(key_type, lhs_inner),
                 ScalarValue::Dictionary(_, rhs_inner),
             ) => {
-                let winner = min_max_generic!(lhs_inner.as_ref(), rhs_inner.as_ref(), $OP);
-                ScalarValue::Dictionary(key_type.clone(), Box::new(winner))
+                min_max_dictionary!(
+                    lhs_inner.as_ref(),
+                    rhs_inner.as_ref(),
+                    wrap key_type,
+                    $OP
+                )
             }
 
             (
                 ScalarValue::Dictionary(_, lhs_inner),
                 rhs,
             ) => {
-                min_max_generic!(lhs_inner.as_ref(), rhs, $OP)
+                min_max_dictionary!(lhs_inner.as_ref(), rhs, $OP)
             }
 
             (
                 lhs,
                 ScalarValue::Dictionary(_, rhs_inner),
             ) => {
-                min_max_generic!(lhs, rhs_inner.as_ref(), $OP)
+                min_max_dictionary!(lhs, rhs_inner.as_ref(), $OP)
             }
 
             e => {
@@ -445,6 +459,17 @@ macro_rules! min_max {
     }};
 }
 
+fn dictionary_batch_extreme(
+    values: &ArrayRef,
+    extreme_fn: fn(&ArrayRef) -> Result<ScalarValue>,
+) -> Result<ScalarValue> {
+    let DataType::Dictionary(key_type, _) = values.data_type() else {
+        unreachable!("dictionary_batch_extreme requires dictionary arrays")
+    };
+    let inner = extreme_fn(values.as_any_dictionary().values())?;
+    Ok(ScalarValue::Dictionary(key_type.clone(), Box::new(inner)))
+}
+
 /// An accumulator to compute the maximum value
 #[derive(Debug, Clone)]
 pub struct MaxAccumulator {
@@ -788,32 +813,22 @@ pub fn min_batch(values: &ArrayRef) -> Result<ScalarValue> {
         DataType::FixedSizeList(_, _) => {
             min_max_batch_generic(values, Ordering::Greater)?
         }
-        DataType::Dictionary(key_type, _) => {
-            let dict_values = values.as_any_dictionary().values();
-            let inner = min_batch(dict_values)?;
-            ScalarValue::Dictionary(key_type.clone(), Box::new(inner))
-        }
+        DataType::Dictionary(_, _) => dictionary_batch_extreme(values, min_batch)?,
         _ => min_max_batch!(values, min),
     })
 }
 
 /// Generic min/max implementation for complex types
 fn min_max_batch_generic(array: &ArrayRef, ordering: Ordering) -> Result<ScalarValue> {
-    if array.len() == array.null_count() {
+    let mut non_null_indices = (0..array.len()).filter(|&i| !array.is_null(i));
+    let Some(first_idx) = non_null_indices.next() else {
         return ScalarValue::try_from(array.data_type());
-    }
-    let mut extreme = ScalarValue::try_from_array(array, 0)?;
-    for i in 1..array.len() {
+    };
+
+    let mut extreme = ScalarValue::try_from_array(array, first_idx)?;
+    for i in non_null_indices {
         let current = ScalarValue::try_from_array(array, i)?;
-        if current.is_null() {
-            continue;
-        }
-        if extreme.is_null() {
-            extreme = current;
-            continue;
-        }
-        let cmp = extreme.try_cmp(&current)?;
-        if cmp == ordering {
+        if extreme.try_cmp(&current)? == ordering {
             extreme = current;
         }
     }
@@ -870,11 +885,7 @@ pub fn max_batch(values: &ArrayRef) -> Result<ScalarValue> {
         DataType::List(_) => min_max_batch_generic(values, Ordering::Less)?,
         DataType::LargeList(_) => min_max_batch_generic(values, Ordering::Less)?,
         DataType::FixedSizeList(_, _) => min_max_batch_generic(values, Ordering::Less)?,
-        DataType::Dictionary(key_type, _) => {
-            let dict_values = values.as_any_dictionary().values();
-            let inner = max_batch(dict_values)?;
-            ScalarValue::Dictionary(key_type.clone(), Box::new(inner))
-        }
+        DataType::Dictionary(_, _) => dictionary_batch_extreme(values, max_batch)?,
         _ => min_max_batch!(values, max),
     })
 }

datafusion/functions-aggregate/src/min_max.rs

@@ -1267,94 +1267,79 @@ mod tests {
         ScalarValue::Dictionary(Box::new(key_type), Box::new(inner))
     }
 
-    #[test]
-    fn test_min_max_dictionary_without_coercion() -> Result<()> {
-        let values = StringArray::from(vec!["b", "c", "a", "d"]);
-        let keys = Int32Array::from(vec![Some(0), Some(1), Some(2), Some(3)]);
-        let dict_array =
-            DictionaryArray::try_new(keys, Arc::new(values) as ArrayRef).unwrap();
-        let dict_array_ref = Arc::new(dict_array) as ArrayRef;
-        let dict_type = dict_array_ref.data_type().clone();
+    fn utf8_dict_scalar(key_type: DataType, value: &str) -> ScalarValue {
+        dict_scalar(key_type, ScalarValue::Utf8(Some(value.to_string())))
+    }
+
+    fn string_dictionary_batch(
+        values: Vec<&str>,
+        keys: Vec<Option<i32>>,
+    ) -> ArrayRef {
+        let values = Arc::new(StringArray::from(values)) as ArrayRef;
+        Arc::new(DictionaryArray::try_new(Int32Array::from(keys), values).unwrap())
+            as ArrayRef
+    }
+
+    fn assert_dictionary_min_max(
+        dict_type: &DataType,
+        batches: &[ArrayRef],
+        expected_min: &str,
+        expected_max: &str,
+    ) -> Result<()> {
+        let key_type = match dict_type {
+            DataType::Dictionary(key_type, _) => key_type.as_ref().clone(),
+            other => panic!("expected dictionary type, got {other:?}"),
+        };
 
-        let mut min_acc = MinAccumulator::try_new(&dict_type)?;
-        min_acc.update_batch(&[Arc::clone(&dict_array_ref)])?;
-        let min_result = min_acc.evaluate()?;
+        let mut min_acc = MinAccumulator::try_new(dict_type)?;
+        for batch in batches {
+            min_acc.update_batch(&[Arc::clone(batch)])?;
+        }
         assert_eq!(
-            min_result,
-            dict_scalar(DataType::Int32, ScalarValue::Utf8(Some("a".to_string())))
+            min_acc.evaluate()?,
+            utf8_dict_scalar(key_type.clone(), expected_min)
         );
 
-        let mut max_acc = MaxAccumulator::try_new(&dict_type)?;
-        max_acc.update_batch(&[Arc::clone(&dict_array_ref)])?;
-        let max_result = max_acc.evaluate()?;
-        assert_eq!(
-            max_result,
-            dict_scalar(DataType::Int32, ScalarValue::Utf8(Some("d".to_string())))
-        );
+        let mut max_acc = MaxAccumulator::try_new(dict_type)?;
+        for batch in batches {
+            max_acc.update_batch(&[Arc::clone(batch)])?;
+        }
+        assert_eq!(max_acc.evaluate()?, utf8_dict_scalar(key_type, expected_max));
+
         Ok(())
     }
 
     #[test]
-    fn test_min_max_dictionary_with_nulls() -> Result<()> {
-        let values = StringArray::from(vec!["b", "c", "a"]);
-        let keys = Int32Array::from(vec![None, Some(0), None, Some(1), Some(2)]);
-        let dict_array =
-            DictionaryArray::try_new(keys, Arc::new(values) as ArrayRef).unwrap();
-        let dict_array_ref = Arc::new(dict_array) as ArrayRef;
+    fn test_min_max_dictionary_without_coercion() -> Result<()> {
+        let dict_array_ref = string_dictionary_batch(
+            vec!["b", "c", "a", "d"],
+            vec![Some(0), Some(1), Some(2), Some(3)],
+        );
         let dict_type = dict_array_ref.data_type().clone();
 
-        let mut min_acc = MinAccumulator::try_new(&dict_type)?;
-        min_acc.update_batch(&[Arc::clone(&dict_array_ref)])?;
-        let min_result = min_acc.evaluate()?;
-        assert_eq!(
-            min_result,
-            dict_scalar(DataType::Int32, ScalarValue::Utf8(Some("a".to_string())))
-        );
+        assert_dictionary_min_max(&dict_type, &[dict_array_ref], "a", "d")
+    }
 
-        let mut max_acc = MaxAccumulator::try_new(&dict_type)?;
-        max_acc.update_batch(&[Arc::clone(&dict_array_ref)])?;
-        let max_result = max_acc.evaluate()?;
-        assert_eq!(
-            max_result,
-            dict_scalar(DataType::Int32, ScalarValue::Utf8(Some("c".to_string())))
+    #[test]
+    fn test_min_max_dictionary_with_nulls() -> Result<()> {
+        let dict_array_ref = string_dictionary_batch(
+            vec!["b", "c", "a"],
+            vec![None, Some(0), None, Some(1), Some(2)],
         );
-        Ok(())
+        let dict_type = dict_array_ref.data_type().clone();
+
+        assert_dictionary_min_max(&dict_type, &[dict_array_ref], "a", "c")
     }
 
     #[test]
     fn test_min_max_dictionary_multi_batch() -> Result<()> {
         let dict_type =
             DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
+        let batch1 =
+            string_dictionary_batch(vec!["b", "c"], vec![Some(0), Some(1)]);
+        let batch2 =
+            string_dictionary_batch(vec!["a", "d"], vec![Some(0), Some(1)]);
 
-        let values1 = StringArray::from(vec!["b", "c"]);
-        let keys1 = Int32Array::from(vec![Some(0), Some(1)]);
-        let batch1 = Arc::new(
-            DictionaryArray::try_new(keys1, Arc::new(values1) as ArrayRef).unwrap(),
-        ) as ArrayRef;
-
-        let values2 = StringArray::from(vec!["a", "d"]);
-        let keys2 = Int32Array::from(vec![Some(0), Some(1)]);
-        let batch2 = Arc::new(
-            DictionaryArray::try_new(keys2, Arc::new(values2) as ArrayRef).unwrap(),
-        ) as ArrayRef;
-
-        let mut min_acc = MinAccumulator::try_new(&dict_type)?;
-        min_acc.update_batch(&[Arc::clone(&batch1)])?;
-        min_acc.update_batch(&[Arc::clone(&batch2)])?;
-        let min_result = min_acc.evaluate()?;
-        assert_eq!(
-            min_result,
-            dict_scalar(DataType::Int32, ScalarValue::Utf8(Some("a".to_string())))
-        );
-
-        let mut max_acc = MaxAccumulator::try_new(&dict_type)?;
-        max_acc.update_batch(&[Arc::clone(&batch1)])?;
-        max_acc.update_batch(&[Arc::clone(&batch2)])?;
-        let max_result = max_acc.evaluate()?;
-        assert_eq!(
-            max_result,
-            dict_scalar(DataType::Int32, ScalarValue::Utf8(Some("d".to_string())))
-        );
-        Ok(())
+        assert_dictionary_min_max(&dict_type, &[batch1, batch2], "a", "d")
     }
 }