Refactor min_max handling and reduce test duplication

kosiew · kosiew · commit f2330b913be5 · 2026-04-02T22:09:30.000+08:00
Normalize dictionary scalars before min_max type match.
Remove separate dictionary batch scan for a shared generic path.
Consolidate repeated logic with a single update_extreme helper.
Reduce test duplication in min_max.rs by reusing accumulators
and consolidating dictionary builders. Clean up tests in
basic.rs by extracting RecordBatch construction and
compressing final-plan assertions.
diff --git a/datafusion/core/tests/sql/aggregates/basic.rs b/datafusion/core/tests/sql/aggregates/basic.rs
@@ -445,6 +445,20 @@ async fn count_distinct_dictionary_mixed_values() -> Result<()> {
 
 #[tokio::test]
 async fn min_max_dictionary_uses_planned_dictionary_path() -> Result<()> {
+    fn dictionary_batch(
+        schema: SchemaRef,
+        keys: Vec<Option<i32>>,
+        values: Vec<&str>,
+    ) -> Result<RecordBatch> {
+        Ok(RecordBatch::try_new(
+            schema,
+            vec![Arc::new(DictionaryArray::new(
+                Int32Array::from(keys),
+                Arc::new(StringArray::from(values)),
+            ))],
+        )?)
+    }
+
     let ctx =
         SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2));
 
@@ -456,20 +470,13 @@ async fn min_max_dictionary_uses_planned_dictionary_path() -> Result<()> {
         true,
     )]));
 
-    let batch1 = RecordBatch::try_new(
+    let batch1 = dictionary_batch(
         schema.clone(),
-        vec![Arc::new(DictionaryArray::new(
-            Int32Array::from(vec![Some(1), Some(1), None]),
-            Arc::new(StringArray::from(vec!["a", "z", "zz_unused"])),
-        ))],
-    )?;
-    let batch2 = RecordBatch::try_new(
-        schema.clone(),
-        vec![Arc::new(DictionaryArray::new(
-            Int32Array::from(vec![Some(0), Some(1)]),
-            Arc::new(StringArray::from(vec!["a", "d"])),
-        ))],
+        vec![Some(1), Some(1), None],
+        vec!["a", "z", "zz_unused"],
     )?;
+    let batch2 =
+        dictionary_batch(schema.clone(), vec![Some(0), Some(1)], vec!["a", "d"])?;
     let provider = MemTable::try_new(schema, vec![vec![batch1], vec![batch2]])?;
     ctx.register_table("t", Arc::new(provider))?;
 
@@ -479,10 +486,9 @@ async fn min_max_dictionary_uses_planned_dictionary_path() -> Result<()> {
     let physical_plan = df.clone().create_physical_plan().await?;
     let formatted_plan = format!("{}", displayable(physical_plan.as_ref()).indent(true));
     assert!(formatted_plan.contains("AggregateExec: mode=Partial, gby=[]"));
-    assert!(
-        formatted_plan.contains("AggregateExec: mode=Final, gby=[]")
-            || formatted_plan.contains("AggregateExec: mode=FinalPartitioned, gby=[]")
-    );
+    assert!(["Final", "FinalPartitioned"].iter().any(|mode| {
+        formatted_plan.contains(&format!("AggregateExec: mode={mode}, gby=[]"))
+    }));
 
     let results = df.collect().await?;
 
diff --git a/datafusion/functions-aggregate-common/src/min_max.rs b/datafusion/functions-aggregate-common/src/min_max.rs
@@ -144,7 +144,18 @@ macro_rules! min_max_generic {
 // min/max of two scalar values of the same type
 macro_rules! min_max {
     ($VALUE:expr, $DELTA:expr, $OP:ident) => {{
-        Ok(match ($VALUE, $DELTA) {
+        let value = $VALUE;
+        let delta = $DELTA;
+        let wrap_key_type = match (value, delta) {
+            (ScalarValue::Dictionary(key_type, _), ScalarValue::Dictionary(_, _)) => {
+                Some(key_type.as_ref())
+            }
+            _ => None,
+        };
+        let value = unwrap_dictionary_scalar(value);
+        let delta = unwrap_dictionary_scalar(delta);
+
+        let result = match (value, delta) {
             (ScalarValue::Null, ScalarValue::Null) => ScalarValue::Null,
             (
                 lhs @ ScalarValue::Decimal32(lhsv, lhsp, lhss),
@@ -413,65 +424,51 @@ macro_rules! min_max {
                 min_max_generic!(lhs, rhs, $OP)
             }
 
-            (
-                ScalarValue::Dictionary(key_type, lhs_inner),
-                ScalarValue::Dictionary(_, rhs_inner),
-            ) => {
-                wrap_dictionary_scalar(
-                    key_type.as_ref(),
-                    min_max_generic!(lhs_inner.as_ref(), rhs_inner.as_ref(), $OP),
-                )
-            }
-
-            (
-                ScalarValue::Dictionary(_, lhs_inner),
-                rhs,
-            ) => {
-                min_max_generic!(lhs_inner.as_ref(), rhs, $OP)
-            }
-
-            (
-                lhs,
-                ScalarValue::Dictionary(_, rhs_inner),
-            ) => {
-                min_max_generic!(lhs, rhs_inner.as_ref(), $OP)
-            }
-
             e => {
                 return internal_err!(
                     "MIN/MAX is not expected to receive scalars of incompatible types {:?}",
                     e
                 )
             }
+        };
+
+        Ok(match wrap_key_type {
+            Some(key_type) => wrap_dictionary_scalar(key_type, result),
+            None => result,
         })
     }};
 }
 
-fn dictionary_batch_extreme(
-    values: &ArrayRef,
-    ordering: Ordering,
-) -> Result<ScalarValue> {
-    let mut extreme: Option<ScalarValue> = None;
-
-    for i in 0..values.len() {
-        let current = ScalarValue::try_from_array(values, i)?;
-        if current.is_null() {
-            continue;
-        }
-
-        match &extreme {
-            Some(existing) if existing.try_cmp(&current)? != ordering => {}
-            _ => extreme = Some(current),
-        }
+fn unwrap_dictionary_scalar(value: &ScalarValue) -> &ScalarValue {
+    match value {
+        ScalarValue::Dictionary(_, inner) => inner.as_ref(),
+        _ => value,
     }
-
-    extreme.map_or_else(|| ScalarValue::try_from(values.data_type()), Ok)
 }
 
 fn wrap_dictionary_scalar(key_type: &DataType, value: ScalarValue) -> ScalarValue {
     ScalarValue::Dictionary(Box::new(key_type.clone()), Box::new(value))
 }
 
+fn update_extreme(
+    extreme: &mut Option<ScalarValue>,
+    current: ScalarValue,
+    ordering: Ordering,
+) -> Result<()> {
+    if current.is_null() {
+        return Ok(());
+    }
+
+    if !matches!(
+        extreme,
+        Some(existing) if existing.try_cmp(&current)? != ordering
+    ) {
+        *extreme = Some(current);
+    }
+
+    Ok(())
+}
+
 /// An accumulator to compute the maximum value
 #[derive(Debug, Clone)]
 pub struct MaxAccumulator {
@@ -815,32 +812,24 @@ pub fn min_batch(values: &ArrayRef) -> Result<ScalarValue> {
         DataType::FixedSizeList(_, _) => {
             min_max_batch_generic(values, Ordering::Greater)?
         }
-        DataType::Dictionary(_, _) => {
-            dictionary_batch_extreme(values, Ordering::Greater)?
-        }
+        DataType::Dictionary(_, _) => min_max_batch_generic(values, Ordering::Greater)?,
         _ => min_max_batch!(values, min),
     })
 }
 
 /// Generic min/max implementation for complex types
 fn min_max_batch_generic(array: &ArrayRef, ordering: Ordering) -> Result<ScalarValue> {
-    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, 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.try_cmp(&current)? == ordering {
-            extreme = current;
-        }
+    let mut extreme = None;
+
+    for i in 0..array.len() {
+        update_extreme(
+            &mut extreme,
+            ScalarValue::try_from_array(array, i)?,
+            ordering,
+        )?;
     }
 
-    Ok(extreme)
+    extreme.map_or_else(|| ScalarValue::try_from(array.data_type()), Ok)
 }
 
 /// dynamically-typed max(array) -> ScalarValue
@@ -892,7 +881,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(_, _) => dictionary_batch_extreme(values, Ordering::Less)?,
+        DataType::Dictionary(_, _) => min_max_batch_generic(values, Ordering::Less)?,
         _ => min_max_batch!(values, max),
     })
 }
diff --git a/datafusion/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs
@@ -1255,22 +1255,23 @@ mod tests {
 
     #[test]
     fn test_min_max_dictionary_after_coercion() -> Result<()> {
-        let values = StringArray::from(vec!["b", "c", "a", "🦀", "d"]);
-        let keys = Int32Array::from(vec![Some(0), Some(1), Some(2), None, Some(4)]);
-        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_array_ref = string_dictionary_batch(
+            vec!["b", "c", "a", "🦀", "d"],
+            &[Some(0), Some(1), Some(2), None, Some(4)],
+        );
         let rt_type =
             get_min_max_result_type(&[dict_array_ref.data_type().clone()])?[0].clone();
 
-        let mut min_acc = MinAccumulator::try_new(&rt_type)?;
-        min_acc.update_batch(&[Arc::clone(&dict_array_ref)])?;
-        let min_result = min_acc.evaluate()?;
+        let min_result = evaluate_dictionary_accumulator(
+            MinAccumulator::try_new(&rt_type)?,
+            &[Arc::clone(&dict_array_ref)],
+        )?;
         assert_eq!(min_result, ScalarValue::Utf8(Some("a".to_string())));
 
-        let mut max_acc = MaxAccumulator::try_new(&rt_type)?;
-        max_acc.update_batch(&[Arc::clone(&dict_array_ref)])?;
-        let max_result = max_acc.evaluate()?;
+        let max_result = evaluate_dictionary_accumulator(
+            MaxAccumulator::try_new(&rt_type)?,
+            &[dict_array_ref],
+        )?;
         assert_eq!(max_result, ScalarValue::Utf8(Some("d".to_string())));
         Ok(())
     }
@@ -1283,18 +1284,11 @@ mod tests {
         dict_scalar(key_type, ScalarValue::Utf8(Some(value.to_string())))
     }
 
-    fn string_dictionary_batch(values: &[&str], keys: &[Option<i32>]) -> ArrayRef {
-        let values = Arc::new(StringArray::from(values.to_vec())) as ArrayRef;
-        Arc::new(
-            DictionaryArray::try_new(Int32Array::from(keys.to_vec()), values).unwrap(),
-        ) as ArrayRef
-    }
-
-    fn optional_string_dictionary_batch(
-        values: &[Option<&str>],
-        keys: &[Option<i32>],
-    ) -> ArrayRef {
-        let values = Arc::new(StringArray::from(values.to_vec())) as ArrayRef;
+    fn string_dictionary_batch<T>(values: Vec<T>, keys: &[Option<i32>]) -> ArrayRef
+    where
+        StringArray: From<Vec<T>>,
+    {
+        let values = Arc::new(StringArray::from(values)) as ArrayRef;
         Arc::new(
             DictionaryArray::try_new(Int32Array::from(keys.to_vec()), values).unwrap(),
         ) as ArrayRef
@@ -1339,7 +1333,7 @@ mod tests {
     #[test]
     fn test_min_max_dictionary_without_coercion() -> Result<()> {
         let dict_array_ref = string_dictionary_batch(
-            &["b", "c", "a", "d"],
+            vec!["b", "c", "a", "d"],
             &[Some(0), Some(1), Some(2), Some(3)],
         );
         let dict_type = dict_array_ref.data_type().clone();
@@ -1350,7 +1344,7 @@ mod tests {
     #[test]
     fn test_min_max_dictionary_with_nulls() -> Result<()> {
         let dict_array_ref = string_dictionary_batch(
-            &["b", "c", "a"],
+            vec!["b", "c", "a"],
             &[None, Some(0), None, Some(1), Some(2)],
         );
         let dict_type = dict_array_ref.data_type().clone();
@@ -1360,17 +1354,19 @@ mod tests {
 
     #[test]
     fn test_min_max_dictionary_ignores_unreferenced_values() -> Result<()> {
-        let dict_array_ref =
-            string_dictionary_batch(&["a", "z", "zz_unused"], &[Some(1), Some(1), None]);
+        let dict_array_ref = string_dictionary_batch(
+            vec!["a", "z", "zz_unused"],
+            &[Some(1), Some(1), None],
+        );
         let dict_type = dict_array_ref.data_type().clone();
 
         assert_dictionary_min_max(&dict_type, &[dict_array_ref], "z", "z")
     }
 
     #[test]
     fn test_min_max_dictionary_ignores_referenced_null_values() -> Result<()> {
-        let dict_array_ref = optional_string_dictionary_batch(
-            &[Some("b"), None, Some("a"), Some("d")],
+        let dict_array_ref = string_dictionary_batch(
+            vec![Some("b"), None, Some("a"), Some("d")],
             &[Some(0), Some(1), Some(2), Some(3)],
         );
         let dict_type = dict_array_ref.data_type().clone();
@@ -1382,8 +1378,8 @@ mod tests {
     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(&["b", "c"], &[Some(0), Some(1)]);
-        let batch2 = string_dictionary_batch(&["a", "d"], &[Some(0), Some(1)]);
+        let batch1 = string_dictionary_batch(vec!["b", "c"], &[Some(0), Some(1)]);
+        let batch2 = string_dictionary_batch(vec!["a", "d"], &[Some(0), Some(1)]);
 
         assert_dictionary_min_max(&dict_type, &[batch1, batch2], "a", "d")
     }
