Skip to content

Commit 9fd697c

Browse files
authored
Move array_agg unit tests to slt (#10402)
1 parent ede3de8 commit 9fd697c

File tree

2 files changed

+24
-119
lines changed

2 files changed

+24
-119
lines changed

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

Lines changed: 0 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -186,122 +186,3 @@ impl Accumulator for ArrayAggAccumulator {
186186
- std::mem::size_of_val(&self.datatype)
187187
}
188188
}
189-
190-
#[cfg(test)]
191-
mod tests {
192-
use super::*;
193-
use crate::expressions::col;
194-
use crate::expressions::tests::aggregate;
195-
use arrow::array::Int32Array;
196-
use arrow::datatypes::*;
197-
use arrow::record_batch::RecordBatch;
198-
use arrow_array::ListArray;
199-
use arrow_buffer::OffsetBuffer;
200-
use datafusion_common::DataFusionError;
201-
202-
macro_rules! test_op {
203-
($ARRAY:expr, $DATATYPE:expr, $OP:ident, $EXPECTED:expr) => {
204-
test_op!($ARRAY, $DATATYPE, $OP, $EXPECTED, $EXPECTED.data_type())
205-
};
206-
($ARRAY:expr, $DATATYPE:expr, $OP:ident, $EXPECTED:expr, $EXPECTED_DATATYPE:expr) => {{
207-
let schema = Schema::new(vec![Field::new("a", $DATATYPE, true)]);
208-
209-
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![$ARRAY])?;
210-
211-
let agg = Arc::new(<$OP>::new(
212-
col("a", &schema)?,
213-
"bla".to_string(),
214-
$EXPECTED_DATATYPE,
215-
true,
216-
));
217-
let actual = aggregate(&batch, agg)?;
218-
let expected = ScalarValue::from($EXPECTED);
219-
220-
assert_eq!(expected, actual);
221-
222-
Ok(()) as Result<(), DataFusionError>
223-
}};
224-
}
225-
226-
#[test]
227-
fn array_agg_i32() -> Result<()> {
228-
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
229-
230-
let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
231-
Some(1),
232-
Some(2),
233-
Some(3),
234-
Some(4),
235-
Some(5),
236-
])]);
237-
let list = ScalarValue::List(Arc::new(list));
238-
239-
test_op!(a, DataType::Int32, ArrayAgg, list, DataType::Int32)
240-
}
241-
242-
#[test]
243-
fn array_agg_nested() -> Result<()> {
244-
let a1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
245-
Some(1),
246-
Some(2),
247-
Some(3),
248-
])]);
249-
let a2 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
250-
Some(4),
251-
Some(5),
252-
])]);
253-
let l1 = ListArray::new(
254-
Arc::new(Field::new("item", a1.data_type().to_owned(), true)),
255-
OffsetBuffer::from_lengths([a1.len() + a2.len()]),
256-
arrow::compute::concat(&[&a1, &a2])?,
257-
None,
258-
);
259-
260-
let a1 =
261-
ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(6)])]);
262-
let a2 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
263-
Some(7),
264-
Some(8),
265-
])]);
266-
let l2 = ListArray::new(
267-
Arc::new(Field::new("item", a1.data_type().to_owned(), true)),
268-
OffsetBuffer::from_lengths([a1.len() + a2.len()]),
269-
arrow::compute::concat(&[&a1, &a2])?,
270-
None,
271-
);
272-
273-
let a1 =
274-
ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(9)])]);
275-
let l3 = ListArray::new(
276-
Arc::new(Field::new("item", a1.data_type().to_owned(), true)),
277-
OffsetBuffer::from_lengths([a1.len()]),
278-
arrow::compute::concat(&[&a1])?,
279-
None,
280-
);
281-
282-
let list = ListArray::new(
283-
Arc::new(Field::new("item", l1.data_type().to_owned(), true)),
284-
OffsetBuffer::from_lengths([l1.len() + l2.len() + l3.len()]),
285-
arrow::compute::concat(&[&l1, &l2, &l3])?,
286-
None,
287-
);
288-
let list = ScalarValue::List(Arc::new(list));
289-
let l1 = ScalarValue::List(Arc::new(l1));
290-
let l2 = ScalarValue::List(Arc::new(l2));
291-
let l3 = ScalarValue::List(Arc::new(l3));
292-
293-
let array = ScalarValue::iter_to_array(vec![l1, l2, l3]).unwrap();
294-
295-
test_op!(
296-
array,
297-
DataType::List(Arc::new(Field::new_list(
298-
"item",
299-
Field::new("item", DataType::Int32, true),
300-
true,
301-
))),
302-
ArrayAgg,
303-
list,
304-
DataType::List(Arc::new(Field::new("item", DataType::Int32, true,)))
305-
)
306-
}
307-
}

datafusion/sqllogictest/test_files/aggregate.slt

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1925,6 +1925,30 @@ SELECT ARRAY_AGG([1]);
19251925
----
19261926
[[1]]
19271927

1928+
# array_agg_i32
1929+
statement ok
1930+
create table t (c1 int) as values (1), (2), (3), (4), (5);
1931+
1932+
query ?
1933+
select array_agg(c1) from t;
1934+
----
1935+
[1, 2, 3, 4, 5]
1936+
1937+
statement ok
1938+
drop table t;
1939+
1940+
# array_agg_nested
1941+
statement ok
1942+
create table t as values (make_array([1, 2, 3], [4, 5])), (make_array([6], [7, 8])), (make_array([9]));
1943+
1944+
query ?
1945+
select array_agg(column1) from t;
1946+
----
1947+
[[[1, 2, 3], [4, 5]], [[6], [7, 8]], [[9]]]
1948+
1949+
statement ok
1950+
drop table t;
1951+
19281952
# variance_single_value
19291953
query RRRR
19301954
select var(sq.column1), var_pop(sq.column1), stddev(sq.column1), stddev_pop(sq.column1) from (values (1.0)) as sq;

0 commit comments

Comments
 (0)