fix: inconsistent scalar types in DistinctArrayAggAccumulator state (#7385)

* fix: inconsistent types in array_agg_distinct merge_batch

* Apply suggestions from code review

Co-authored-by: Metehan Yıldırım <[email protected]>

* filtering NULLs & validating sqllogictest output

---------

Co-authored-by: Metehan Yıldırım <[email protected]>

Author: korowa

datafusion/physical-expr/src/aggregate/array_agg_distinct.rs

@@ -28,8 +28,7 @@ use std::collections::HashSet;
 use crate::aggregate::utils::down_cast_any_ref;
 use crate::expressions::format_state_name;
 use crate::{AggregateExpr, PhysicalExpr};
-use datafusion_common::Result;
-use datafusion_common::ScalarValue;
+use datafusion_common::{internal_err, DataFusionError, Result, ScalarValue};
 use datafusion_expr::Accumulator;
 
 /// Expression for a ARRAY_AGG(DISTINCT) aggregation.
@@ -135,23 +134,36 @@ impl Accumulator for DistinctArrayAggAccumulator {
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         assert_eq!(values.len(), 1, "batch input should only include 1 column!");
 
-        let arr = &values[0];
-        for i in 0..arr.len() {
-            self.values.insert(ScalarValue::try_from_array(arr, i)?);
-        }
-        Ok(())
+        let array = &values[0];
+        (0..array.len()).try_for_each(|i| {
+            if !array.is_null(i) {
+                self.values.insert(ScalarValue::try_from_array(array, i)?);
+            }
+            Ok(())
+        })
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
         if states.is_empty() {
             return Ok(());
         }
 
-        for array in states {
-            for j in 0..array.len() {
-                self.values.insert(ScalarValue::try_from_array(array, j)?);
+        assert_eq!(
+            states.len(),
+            1,
+            "array_agg_distinct states must contain single array"
+        );
+
+        let array = &states[0];
+        (0..array.len()).try_for_each(|i| {
+            let scalar = ScalarValue::try_from_array(array, i)?;
+            if let ScalarValue::List(Some(values), _) = scalar {
+                self.values.extend(values);
+                Ok(())
+            } else {
+                internal_err!("array_agg_distinct state must be list")
             }
-        }
+        })?;
 
         Ok(())
     }
@@ -174,12 +186,35 @@ impl Accumulator for DistinctArrayAggAccumulator {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::aggregate::utils::get_accum_scalar_values_as_arrays;
     use crate::expressions::col;
     use crate::expressions::tests::aggregate;
     use arrow::array::{ArrayRef, Int32Array};
     use arrow::datatypes::{DataType, Schema};
     use arrow::record_batch::RecordBatch;
 
+    fn compare_list_contents(expected: ScalarValue, actual: ScalarValue) -> Result<()> {
+        match (expected, actual) {
+            (ScalarValue::List(Some(mut e), _), ScalarValue::List(Some(mut a), _)) => {
+                // workaround lack of Ord of ScalarValue
+                let cmp = |a: &ScalarValue, b: &ScalarValue| {
+                    a.partial_cmp(b).expect("Can compare ScalarValues")
+                };
+
+                e.sort_by(cmp);
+                a.sort_by(cmp);
+                // Check that the inputs are the same
+                assert_eq!(e, a);
+            }
+            _ => {
+                return Err(DataFusionError::Internal(
+                    "Expected scalar lists as inputs".to_string(),
+                ));
+            }
+        }
+        Ok(())
+    }
+
     fn check_distinct_array_agg(
         input: ArrayRef,
         expected: ScalarValue,
@@ -195,24 +230,34 @@ mod tests {
         ));
         let actual = aggregate(&batch, agg)?;
 
-        match (expected, actual) {
-            (ScalarValue::List(Some(mut e), _), ScalarValue::List(Some(mut a), _)) => {
-                // workaround lack of Ord of ScalarValue
-                let cmp = |a: &ScalarValue, b: &ScalarValue| {
-                    a.partial_cmp(b).expect("Can compare ScalarValues")
-                };
+        compare_list_contents(expected, actual)
+    }
 
-                e.sort_by(cmp);
-                a.sort_by(cmp);
-                // Check that the inputs are the same
-                assert_eq!(e, a);
-            }
-            _ => {
-                unreachable!()
-            }
-        }
+    fn check_merge_distinct_array_agg(
+        input1: ArrayRef,
+        input2: ArrayRef,
+        expected: ScalarValue,
+        datatype: DataType,
+    ) -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", datatype.clone(), false)]);
+        let agg = Arc::new(DistinctArrayAgg::new(
+            col("a", &schema)?,
+            "bla".to_string(),
+            datatype,
+        ));
 
-        Ok(())
+        let mut accum1 = agg.create_accumulator()?;
+        let mut accum2 = agg.create_accumulator()?;
+
+        accum1.update_batch(&[input1])?;
+        accum2.update_batch(&[input2])?;
+
+        let state = get_accum_scalar_values_as_arrays(accum2.as_ref())?;
+        accum1.merge_batch(&state)?;
+
+        let actual = accum1.evaluate()?;
+
+        compare_list_contents(expected, actual)
     }
 
     #[test]
@@ -233,6 +278,27 @@ mod tests {
         check_distinct_array_agg(col, out, DataType::Int32)
     }
 
+    #[test]
+    fn merge_distinct_array_agg_i32() -> Result<()> {
+        let col1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 4, 5, 2]));
+        let col2: ArrayRef = Arc::new(Int32Array::from(vec![1, 3, 7, 8, 4]));
+
+        let out = ScalarValue::new_list(
+            Some(vec![
+                ScalarValue::Int32(Some(1)),
+                ScalarValue::Int32(Some(2)),
+                ScalarValue::Int32(Some(3)),
+                ScalarValue::Int32(Some(4)),
+                ScalarValue::Int32(Some(5)),
+                ScalarValue::Int32(Some(7)),
+                ScalarValue::Int32(Some(8)),
+            ]),
+            DataType::Int32,
+        );
+
+        check_merge_distinct_array_agg(col1, col2, out, DataType::Int32)
+    }
+
     #[test]
     fn distinct_array_agg_nested() -> Result<()> {
         // [[1, 2, 3], [4, 5]]
@@ -296,4 +362,66 @@ mod tests {
             ))),
         )
     }
+
+    #[test]
+    fn merge_distinct_array_agg_nested() -> Result<()> {
+        // [[1, 2], [3, 4]]
+        let l1 = ScalarValue::new_list(
+            Some(vec![
+                ScalarValue::new_list(
+                    Some(vec![ScalarValue::from(1i32), ScalarValue::from(2i32)]),
+                    DataType::Int32,
+                ),
+                ScalarValue::new_list(
+                    Some(vec![ScalarValue::from(3i32), ScalarValue::from(4i32)]),
+                    DataType::Int32,
+                ),
+            ]),
+            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        );
+
+        // [[5]]
+        let l2 = ScalarValue::new_list(
+            Some(vec![ScalarValue::new_list(
+                Some(vec![ScalarValue::from(5i32)]),
+                DataType::Int32,
+            )]),
+            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        );
+
+        // [[6, 7], [8]]
+        let l3 = ScalarValue::new_list(
+            Some(vec![
+                ScalarValue::new_list(
+                    Some(vec![ScalarValue::from(6i32), ScalarValue::from(7i32)]),
+                    DataType::Int32,
+                ),
+                ScalarValue::new_list(
+                    Some(vec![ScalarValue::from(8i32)]),
+                    DataType::Int32,
+                ),
+            ]),
+            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        );
+
+        let expected = ScalarValue::new_list(
+            Some(vec![l1.clone(), l2.clone(), l3.clone()]),
+            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        );
+
+        // Duplicate l1 in the input array and check that it is deduped in the output.
+        let input1 = ScalarValue::iter_to_array(vec![l1.clone(), l2]).unwrap();
+        let input2 = ScalarValue::iter_to_array(vec![l1, l3]).unwrap();
+
+        check_merge_distinct_array_agg(
+            input1,
+            input2,
+            expected,
+            DataType::List(Arc::new(Field::new_list(
+                "item",
+                Field::new("item", DataType::Int32, true),
+                true,
+            ))),
+        )
+    }
 }

datafusion/sqllogictest/test_files/aggregate.slt

@@ -1271,14 +1271,31 @@ NULL 4 29 1.260869565217 123 -117 23
 NULL 5 -194 -13.857142857143 118 -101 14
 NULL NULL 781 7.81 125 -117 100
 
-# TODO this querys output is non determinisitic (the order of the elements
-# differs run to run
+# TODO: array_agg_distinct output is non-determinisitic -- rewrite with array_sort(list_sort)
+#       unnest is also not available, so manually unnesting via CROSS JOIN
+# additional count(1) forces array_agg_distinct instead of array_agg over aggregated by c2 data
 #
 # csv_query_array_agg_distinct
-# query T
-# SELECT array_agg(distinct c2) FROM aggregate_test_100
-# ----
-# [4, 2, 3, 5, 1]
+query III
+WITH indices AS (
+  SELECT 1 AS idx UNION ALL
+  SELECT 2 AS idx UNION ALL
+  SELECT 3 AS idx UNION ALL
+  SELECT 4 AS idx UNION ALL
+  SELECT 5 AS idx
+)
+SELECT data.arr[indices.idx] as element, array_length(data.arr) as array_len, dummy
+FROM (
+  SELECT array_agg(distinct c2) as arr, count(1) as dummy FROM aggregate_test_100
+) data
+  CROSS JOIN indices
+ORDER BY 1
+----
+1 5 100
+2 5 100
+3 5 100
+4 5 100
+5 5 100
 
 # aggregate_time_min_and_max
 query TT