functions-aggregate: Support dictionary scalar coercion for min/max

Author: brancz

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

@@ -432,8 +432,15 @@ pub struct MaxAccumulator {
 impl MaxAccumulator {
     /// new max accumulator
     pub fn try_new(datatype: &DataType) -> Result<Self> {
+        // Unwrap dictionary types to match what max_batch() returns: it
+        // recurses into the dictionary's values array, so the resulting
+        // ScalarValue is always the value type, not the dictionary type.
+        let dt = match datatype {
+            DataType::Dictionary(_, v) => v.as_ref(),
+            other => other,
+        };
         Ok(Self {
-            max: ScalarValue::try_from(datatype)?,
+            max: ScalarValue::try_from(dt)?,
         })
     }
 }
@@ -473,8 +480,15 @@ pub struct MinAccumulator {
 impl MinAccumulator {
     /// new min accumulator
     pub fn try_new(datatype: &DataType) -> Result<Self> {
+        // Unwrap dictionary types to match what min_batch() returns: it
+        // recurses into the dictionary's values array, so the resulting
+        // ScalarValue is always the value type, not the dictionary type.
+        let dt = match datatype {
+            DataType::Dictionary(_, v) => v.as_ref(),
+            other => other,
+        };
         Ok(Self {
-            min: ScalarValue::try_from(datatype)?,
+            min: ScalarValue::try_from(dt)?,
         })
     }
 }

datafusion/functions-aggregate/src/min_max.rs

@@ -392,8 +392,12 @@ pub struct SlidingMaxAccumulator {
 impl SlidingMaxAccumulator {
     /// new max accumulator
     pub fn try_new(datatype: &DataType) -> Result<Self> {
+        let dt = match datatype {
+            DataType::Dictionary(_, v) => v.as_ref(),
+            other => other,
+        };
         Ok(Self {
-            max: ScalarValue::try_from(datatype)?,
+            max: ScalarValue::try_from(dt)?,
             moving_max: MovingMax::<ScalarValue>::new(),
         })
     }
@@ -679,8 +683,12 @@ pub struct SlidingMinAccumulator {
 
 impl SlidingMinAccumulator {
     pub fn try_new(datatype: &DataType) -> Result<Self> {
+        let dt = match datatype {
+            DataType::Dictionary(_, v) => v.as_ref(),
+            other => other,
+        };
         Ok(Self {
-            min: ScalarValue::try_from(datatype)?,
+            min: ScalarValue::try_from(dt)?,
             moving_min: MovingMin::<ScalarValue>::new(),
         })
     }
@@ -1270,4 +1278,82 @@ mod tests {
         assert_eq!(max_result, ScalarValue::Utf8(Some("🦀".to_string())));
         Ok(())
     }
+
+    #[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;
+
+        // Pass the raw Dictionary type — no get_min_max_result_type() unwrap.
+        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, ScalarValue::Utf8(Some("a".to_string())));
+
+        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, ScalarValue::Utf8(Some("d".to_string())));
+        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;
+
+        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, ScalarValue::Utf8(Some("a".to_string())));
+
+        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, ScalarValue::Utf8(Some("c".to_string())));
+        Ok(())
+    }
+
+    #[test]
+    fn test_min_max_dictionary_multi_batch() -> Result<()> {
+        let dict_type =
+            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
+
+        // First batch.
+        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;
+
+        // Second batch with a new min and max.
+        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, 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, ScalarValue::Utf8(Some("d".to_string())));
+        Ok(())
+    }
 }