Skip to content

Commit

Permalink
Revert "Support List for Array aggregate order and distinct (apache#9234
Browse files Browse the repository at this point in the history 
)"

This reverts commit fc84a63.
  • Loading branch information
@erratic-pattern
erratic-pattern committed Mar 14, 2024
1 parent 1a4dc00 commit 5965d67
Show file tree
Hide file tree
Showing 6 changed files with 124 additions and 195 deletions.
74 changes: 25 additions & 49 deletions datafusion/common/src/scalar/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ use arrow::{
UInt16Type, UInt32Type, UInt64Type, UInt8Type, DECIMAL128_MAX_PRECISION,
},
};
use arrow_array::cast::as_list_array;
use arrow_array::{ArrowNativeTypeOp, Scalar};

pub use struct_builder::ScalarStructBuilder;
Expand Down Expand Up @@ -2137,67 +2138,28 @@ impl ScalarValue {

/// Retrieve ScalarValue for each row in `array`
///
/// Example 1: Array (ScalarValue::Int32)
/// Example
/// ```
/// use datafusion_common::ScalarValue;
/// use arrow::array::ListArray;
/// use arrow::datatypes::{DataType, Int32Type};
///
/// // Equivalent to [[1,2,3], [4,5]]
/// let list_arr = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
/// Some(vec![Some(1), Some(2), Some(3)]),
/// None,
/// Some(vec![Some(4), Some(5)])
/// ]);
///
/// // Convert the array into Scalar Values for each row
/// let scalar_vec = ScalarValue::convert_array_to_scalar_vec(&list_arr).unwrap();
///
/// let expected = vec![
/// vec![
/// vec![
/// ScalarValue::Int32(Some(1)),
/// ScalarValue::Int32(Some(2)),
/// ScalarValue::Int32(Some(3)),
/// ],
/// vec![
/// ScalarValue::Int32(Some(4)),
/// ScalarValue::Int32(Some(5)),
/// ],
/// ];
///
/// assert_eq!(scalar_vec, expected);
/// ```
///
/// Example 2: Nested array (ScalarValue::List)
/// ```
/// use datafusion_common::ScalarValue;
/// use arrow::array::ListArray;
/// use arrow::datatypes::{DataType, Int32Type};
/// use datafusion_common::utils::array_into_list_array;
/// use std::sync::Arc;
///
/// let list_arr = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
/// Some(vec![Some(1), Some(2), Some(3)]),
/// Some(vec![Some(4), Some(5)])
/// ]);
///
/// // Wrap into another layer of list, we got nested array as [ [[1,2,3], [4,5]] ]
/// let list_arr = array_into_list_array(Arc::new(list_arr));
///
/// // Convert the array into Scalar Values for each row, we got 1D arrays in this example
/// let scalar_vec = ScalarValue::convert_array_to_scalar_vec(&list_arr).unwrap();
///
/// let l1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
/// Some(vec![Some(1), Some(2), Some(3)]),
/// ]);
/// let l2 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
/// Some(vec![Some(4), Some(5)]),
/// ]);
///
/// let expected = vec![
/// vec![
/// ScalarValue::List(Arc::new(l1)),
/// ScalarValue::List(Arc::new(l2)),
/// ],
/// vec![],
/// vec![ScalarValue::Int32(Some(4)), ScalarValue::Int32(Some(5))]
/// ];
///
/// assert_eq!(scalar_vec, expected);
Expand All @@ -2206,13 +2168,27 @@ impl ScalarValue {
let mut scalars = Vec::with_capacity(array.len());

for index in 0..array.len() {
let nested_array = array.as_list::<i32>().value(index);
let scalar_values = (0..nested_array.len())
.map(|i| ScalarValue::try_from_array(&nested_array, i))
.collect::<Result<Vec<_>>>()?;
let scalar_values = match array.data_type() {
DataType::List(_) => {
let list_array = as_list_array(array);
match list_array.is_null(index) {
true => Vec::new(),
false => {
let nested_array = list_array.value(index);
ScalarValue::convert_array_to_scalar_vec(&nested_array)?
.into_iter()
.flatten()
.collect()
}
}
}
_ => {
let scalar = ScalarValue::try_from_array(array, index)?;
vec![scalar]
}
};
scalars.push(scalar_values);
}

Ok(scalars)
}

Expand Down
2 changes: 1 addition & 1 deletion datafusion/core/tests/sql/aggregates.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,9 @@ async fn csv_query_array_agg_distinct() -> Result<()> {
// We should have 1 row containing a list
let column = actual[0].column(0);
assert_eq!(column.len(), 1);

let scalar_vec = ScalarValue::convert_array_to_scalar_vec(&column)?;
let mut scalars = scalar_vec[0].clone();

// workaround lack of Ord of ScalarValue
let cmp = |a: &ScalarValue, b: &ScalarValue| {
a.partial_cmp(b).expect("Can compare ScalarValues")
Expand Down
173 changes: 92 additions & 81 deletions datafusion/physical-expr/src/aggregate/array_agg_distinct.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ use std::sync::Arc;

use arrow::array::ArrayRef;
use arrow::datatypes::{DataType, Field};
use arrow_array::cast::AsArray;

use crate::aggregate::utils::down_cast_any_ref;
use crate::expressions::format_state_name;
Expand Down Expand Up @@ -139,10 +138,9 @@ impl Accumulator for DistinctArrayAggAccumulator {
assert_eq!(values.len(), 1, "batch input should only include 1 column!");

let array = &values[0];

for i in 0..array.len() {
let scalar = ScalarValue::try_from_array(&array, i)?;
self.values.insert(scalar);
let scalar_vec = ScalarValue::convert_array_to_scalar_vec(array)?;
for scalars in scalar_vec {
self.values.extend(scalars);
}

Ok(())
Expand All @@ -153,12 +151,7 @@ impl Accumulator for DistinctArrayAggAccumulator {
return Ok(());
}

let array = &states[0];

assert_eq!(array.len(), 1, "state array should only include 1 row!");
// Unwrap outer ListArray then do update batch
let inner_array = array.as_list::<i32>().value(0);
self.update_batch(&[inner_array])
self.update_batch(states)
}

fn evaluate(&mut self) -> Result<ScalarValue> {
Expand Down Expand Up @@ -189,54 +182,46 @@ mod tests {
use arrow_array::ListArray;
use arrow_buffer::OffsetBuffer;
use datafusion_common::internal_err;
use datafusion_common::utils::array_into_list_array;

// arrow::compute::sort cann't sort ListArray directly, so we need to sort the inner primitive array and wrap it back into ListArray.
fn sort_list_inner(arr: ScalarValue) -> ScalarValue {
let arr = match arr {
ScalarValue::List(arr) => arr.value(0),
_ => {
panic!("Expected ScalarValue::List, got {:?}", arr)
}
};

// arrow::compute::sort can't sort nested ListArray directly, so we compare the scalar values pair-wise.
fn compare_list_contents(
expected: Vec<ScalarValue>,
actual: ScalarValue,
) -> Result<()> {
let array = actual.to_array()?;
let list_array = array.as_list::<i32>();
let inner_array = list_array.value(0);
let mut actual_scalars = vec![];
for index in 0..inner_array.len() {
let sv = ScalarValue::try_from_array(&inner_array, index)?;
actual_scalars.push(sv);
}

if actual_scalars.len() != expected.len() {
return internal_err!(
"Expected and actual list lengths differ: expected={}, actual={}",
expected.len(),
actual_scalars.len()
);
}
let arr = arrow::compute::sort(&arr, None).unwrap();
let list_arr = array_into_list_array(arr);
ScalarValue::List(Arc::new(list_arr))
}

let mut seen = vec![false; expected.len()];
for v in expected {
let mut found = false;
for (i, sv) in actual_scalars.iter().enumerate() {
if sv == &v {
seen[i] = true;
found = true;
break;
fn compare_list_contents(expected: ScalarValue, actual: ScalarValue) -> Result<()> {
let actual = sort_list_inner(actual);

match (&expected, &actual) {
(ScalarValue::List(arr1), ScalarValue::List(arr2)) => {
if arr1.eq(arr2) {
Ok(())
} else {
internal_err!(
"Actual value {:?} not found in expected values {:?}",
actual,
expected
)
}
}
if !found {
return internal_err!(
"Expected value {:?} not found in actual values {:?}",
v,
actual_scalars
);
_ => {
internal_err!("Expected scalar lists as inputs")
}
}

Ok(())
}

fn check_distinct_array_agg(
input: ArrayRef,
expected: Vec<ScalarValue>,
expected: ScalarValue,
datatype: DataType,
) -> Result<()> {
let schema = Schema::new(vec![Field::new("a", datatype.clone(), false)]);
Expand All @@ -249,13 +234,14 @@ mod tests {
true,
));
let actual = aggregate(&batch, agg)?;

compare_list_contents(expected, actual)
}

fn check_merge_distinct_array_agg(
input1: ArrayRef,
input2: ArrayRef,
expected: Vec<ScalarValue>,
expected: ScalarValue,
datatype: DataType,
) -> Result<()> {
let schema = Schema::new(vec![Field::new("a", datatype.clone(), false)]);
Expand All @@ -276,20 +262,23 @@ mod tests {
accum1.merge_batch(&[array])?;

let actual = accum1.evaluate()?;

compare_list_contents(expected, actual)
}

#[test]
fn distinct_array_agg_i32() -> Result<()> {
let col: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 4, 5, 2]));

let expected = vec![
ScalarValue::Int32(Some(1)),
ScalarValue::Int32(Some(2)),
ScalarValue::Int32(Some(4)),
ScalarValue::Int32(Some(5)),
ScalarValue::Int32(Some(7)),
];
let expected =
ScalarValue::List(Arc::new(
ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
Some(1),
Some(2),
Some(4),
Some(5),
Some(7),
])]),
));

check_distinct_array_agg(col, expected, DataType::Int32)
}
Expand All @@ -299,15 +288,18 @@ mod tests {
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 expected = 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)),
];
let expected =
ScalarValue::List(Arc::new(
ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
Some(1),
Some(2),
Some(3),
Some(4),
Some(5),
Some(7),
Some(8),
])]),
));

check_merge_distinct_array_agg(col1, col2, expected, DataType::Int32)
}
Expand Down Expand Up @@ -359,16 +351,23 @@ mod tests {
let l2 = ScalarValue::List(Arc::new(l2));
let l3 = ScalarValue::List(Arc::new(l3));

// Duplicate l1 and l3 in the input array and check that it is deduped in the output.
let array = ScalarValue::iter_to_array(vec![
l1.clone(),
l2.clone(),
l3.clone(),
l3.clone(),
l1.clone(),
])
.unwrap();
let expected = vec![l1, l2, l3];
// Duplicate l1 in the input array and check that it is deduped in the output.
let array = ScalarValue::iter_to_array(vec![l1.clone(), l2, l3, l1]).unwrap();

let expected =
ScalarValue::List(Arc::new(
ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
Some(1),
Some(2),
Some(3),
Some(4),
Some(5),
Some(6),
Some(7),
Some(8),
Some(9),
])]),
));

check_distinct_array_agg(
array,
Expand Down Expand Up @@ -427,10 +426,22 @@ mod tests {
let l3 = ScalarValue::List(Arc::new(l3));

// 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.clone()]).unwrap();
let input2 = ScalarValue::iter_to_array(vec![l1.clone(), l3.clone()]).unwrap();

let expected = vec![l1, l2, l3];
let input1 = ScalarValue::iter_to_array(vec![l1.clone(), l2]).unwrap();
let input2 = ScalarValue::iter_to_array(vec![l1, l3]).unwrap();

let expected =
ScalarValue::List(Arc::new(
ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
Some(1),
Some(2),
Some(3),
Some(4),
Some(5),
Some(6),
Some(7),
Some(8),
])]),
));

check_merge_distinct_array_agg(input1, input2, expected, DataType::Int32)
}
Expand Down
Loading

0 comments on commit 5965d67

Please sign in to comment.