Skip to content

Commit

Permalink
fix: inconsistent scalar types in DistinctArrayAggAccumulator state (…
Browse files Browse the repository at this point in the history 
…apache#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]>
  • Loading branch information
@korowa @metesynnada
korowa and metesynnada authored Aug 24, 2023
1 parent 3f8c512 commit ea9144e
Show file tree
Hide file tree
Showing 2 changed files with 178 additions and 33 deletions.
182 changes: 155 additions & 27 deletions datafusion/physical-expr/src/aggregate/array_agg_distinct.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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(())
}
Expand All @@ -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,
Expand All @@ -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]
Expand All @@ -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]]
Expand Down Expand Up @@ -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,
))),
)
}
}
29 changes: 23 additions & 6 deletions datafusion/sqllogictest/test_files/aggregate.slt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit ea9144e

Please sign in to comment.