Move array_agg unit tests to slt (apache#10402)

andygrove · May 7, 2024 · 9fd697c · 9fd697c
1 parent ede3de8
commit 9fd697c
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 119 deletions.
diff --git a/datafusion/physical-expr/src/aggregate/array_agg.rs b/datafusion/physical-expr/src/aggregate/array_agg.rs
@@ -186,122 +186,3 @@ impl Accumulator for ArrayAggAccumulator {
             - std::mem::size_of_val(&self.datatype)
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::expressions::col;
-    use crate::expressions::tests::aggregate;
-    use arrow::array::Int32Array;
-    use arrow::datatypes::*;
-    use arrow::record_batch::RecordBatch;
-    use arrow_array::ListArray;
-    use arrow_buffer::OffsetBuffer;
-    use datafusion_common::DataFusionError;
-
-    macro_rules! test_op {
-        ($ARRAY:expr, $DATATYPE:expr, $OP:ident, $EXPECTED:expr) => {
-            test_op!($ARRAY, $DATATYPE, $OP, $EXPECTED, $EXPECTED.data_type())
-        };
-        ($ARRAY:expr, $DATATYPE:expr, $OP:ident, $EXPECTED:expr, $EXPECTED_DATATYPE:expr) => {{
-            let schema = Schema::new(vec![Field::new("a", $DATATYPE, true)]);
-
-            let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![$ARRAY])?;
-
-            let agg = Arc::new(<$OP>::new(
-                col("a", &schema)?,
-                "bla".to_string(),
-                $EXPECTED_DATATYPE,
-                true,
-            ));
-            let actual = aggregate(&batch, agg)?;
-            let expected = ScalarValue::from($EXPECTED);
-
-            assert_eq!(expected, actual);
-
-            Ok(()) as Result<(), DataFusionError>
-        }};
-    }
-
-    #[test]
-    fn array_agg_i32() -> Result<()> {
-        let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
-
-        let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
-            Some(1),
-            Some(2),
-            Some(3),
-            Some(4),
-            Some(5),
-        ])]);
-        let list = ScalarValue::List(Arc::new(list));
-
-        test_op!(a, DataType::Int32, ArrayAgg, list, DataType::Int32)
-    }
-
-    #[test]
-    fn array_agg_nested() -> Result<()> {
-        let a1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
-            Some(1),
-            Some(2),
-            Some(3),
-        ])]);
-        let a2 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
-            Some(4),
-            Some(5),
-        ])]);
-        let l1 = ListArray::new(
-            Arc::new(Field::new("item", a1.data_type().to_owned(), true)),
-            OffsetBuffer::from_lengths([a1.len() + a2.len()]),
-            arrow::compute::concat(&[&a1, &a2])?,
-            None,
-        );
-
-        let a1 =
-            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(6)])]);
-        let a2 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
-            Some(7),
-            Some(8),
-        ])]);
-        let l2 = ListArray::new(
-            Arc::new(Field::new("item", a1.data_type().to_owned(), true)),
-            OffsetBuffer::from_lengths([a1.len() + a2.len()]),
-            arrow::compute::concat(&[&a1, &a2])?,
-            None,
-        );
-
-        let a1 =
-            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(9)])]);
-        let l3 = ListArray::new(
-            Arc::new(Field::new("item", a1.data_type().to_owned(), true)),
-            OffsetBuffer::from_lengths([a1.len()]),
-            arrow::compute::concat(&[&a1])?,
-            None,
-        );
-
-        let list = ListArray::new(
-            Arc::new(Field::new("item", l1.data_type().to_owned(), true)),
-            OffsetBuffer::from_lengths([l1.len() + l2.len() + l3.len()]),
-            arrow::compute::concat(&[&l1, &l2, &l3])?,
-            None,
-        );
-        let list = ScalarValue::List(Arc::new(list));
-        let l1 = ScalarValue::List(Arc::new(l1));
-        let l2 = ScalarValue::List(Arc::new(l2));
-        let l3 = ScalarValue::List(Arc::new(l3));
-
-        let array = ScalarValue::iter_to_array(vec![l1, l2, l3]).unwrap();
-
-        test_op!(
-            array,
-            DataType::List(Arc::new(Field::new_list(
-                "item",
-                Field::new("item", DataType::Int32, true),
-                true,
-            ))),
-            ArrayAgg,
-            list,
-            DataType::List(Arc::new(Field::new("item", DataType::Int32, true,)))
-        )
-    }
-}
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -1925,6 +1925,30 @@ SELECT ARRAY_AGG([1]);
 ----
 [[1]]
 
+# array_agg_i32
+statement ok
+create table t (c1 int) as values (1), (2), (3), (4), (5);
+
+query ?
+select array_agg(c1) from t;
+----
+[1, 2, 3, 4, 5]
+
+statement ok
+drop table t;
+
+# array_agg_nested
+statement ok
+create table t as values (make_array([1, 2, 3], [4, 5])), (make_array([6], [7, 8])), (make_array([9]));
+
+query ?
+select array_agg(column1) from t;
+----
+[[[1, 2, 3], [4, 5]], [[6], [7, 8]], [[9]]]
+
+statement ok
+drop table t;
+
 # variance_single_value
 query RRRR
 select var(sq.column1), var_pop(sq.column1), stddev(sq.column1), stddev_pop(sq.column1) from (values (1.0)) as sq;