From 4146cf9b5463444e4574be8b90365ee4ecd087ad Mon Sep 17 00:00:00 2001 From: Eduard Karacharov Date: Tue, 22 Aug 2023 11:24:51 +0300 Subject: [PATCH 1/3] fix: inconsistent types in array_agg_distinct merge_batch --- .../src/aggregate/array_agg_distinct.rs | 169 +++++++++++++++--- .../sqllogictest/test_files/aggregate.slt | 4 + 2 files changed, 151 insertions(+), 22 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs b/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs index 2d7a6e5b0e4d..0a0b512ccd23 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs +++ b/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. @@ -147,11 +146,21 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } - for array in states { - for j in 0..array.len() { - self.values.insert(ScalarValue::try_from_array(array, j)?); + assert!( + states.len() == 1, + "array_agg_distinct states must contain single array" + ); + + let state = &states[0]; + (0..state.len()).try_for_each(|i| { + let scalar = ScalarValue::try_from_array(state, 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 +183,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 +227,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 +275,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 +359,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, + ))), + ) + } } diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index e881acf5755b..fc70ad74c612 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1280,6 +1280,10 @@ NULL NULL 781 7.81 125 -117 100 # ---- # [4, 2, 3, 5, 1] +# additional count(1) forces array_agg_distinct instead of array_agg over aggregated by c2 data +statement ok +SELECT array_agg(distinct c2), count(1) FROM aggregate_test_100 + # aggregate_time_min_and_max query TT select min(t), max(t) from (select '00:00:00' as t union select '00:00:01' union select '00:00:02') From b9a0666ba686e5409504b2cf11ce968dfabeea41 Mon Sep 17 00:00:00 2001 From: Eduard Karacharov <13005055+korowa@users.noreply.github.com> Date: Wed, 23 Aug 2023 20:44:32 +0300 Subject: [PATCH 2/3] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Metehan Yıldırım <100111937+metesynnada@users.noreply.github.com> --- .../physical-expr/src/aggregate/array_agg_distinct.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs b/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs index 0a0b512ccd23..f941dd11015c 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs @@ -146,14 +146,15 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } - assert!( - states.len() == 1, + assert_eq!( + states.len(), + 1, "array_agg_distinct states must contain single array" ); - let state = &states[0]; - (0..state.len()).try_for_each(|i| { - let scalar = ScalarValue::try_from_array(state, i)?; + 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(()) From ace0eec8a7a73e7bebdcc6b1192f9147cd16c304 Mon Sep 17 00:00:00 2001 From: Eduard Karacharov Date: Thu, 24 Aug 2023 10:25:40 +0300 Subject: [PATCH 3/3] filtering NULLs & validating sqllogictest output --- .../src/aggregate/array_agg_distinct.rs | 12 ++++--- .../sqllogictest/test_files/aggregate.slt | 33 +++++++++++++------ 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs b/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs index f941dd11015c..422eecd20155 100644 --- a/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs +++ b/datafusion/physical-expr/src/aggregate/array_agg_distinct.rs @@ -134,11 +134,13 @@ 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<()> { diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index fc70ad74c612..c89362a2f293 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1271,18 +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] - -# additional count(1) forces array_agg_distinct instead of array_agg over aggregated by c2 data -statement ok -SELECT array_agg(distinct c2), count(1) FROM aggregate_test_100 +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