Change ScalarValue::List to store ArrayRef

datafusion/common/src/utils.rs

@@ -19,10 +19,12 @@
 
 use crate::{DataFusionError, Result, ScalarValue};
 use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::buffer::OffsetBuffer;
 use arrow::compute;
 use arrow::compute::{partition, SortColumn, SortOptions};
-use arrow::datatypes::{SchemaRef, UInt32Type};
+use arrow::datatypes::{Field, SchemaRef, UInt32Type};
 use arrow::record_batch::RecordBatch;
+use arrow_array::ListArray;
 use sqlparser::ast::Ident;
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::Parser;
@@ -334,6 +336,18 @@ pub fn longest_consecutive_prefix<T: Borrow<usize>>(
     count
 }
 
+/// Wrap an array into a single element `ListArray`.
+/// For example `[1, 2, 3]` would be converted into `[[1, 2, 3]]`
+pub fn wrap_into_list_array(arr: ArrayRef) -> ListArray {
+    let offsets = OffsetBuffer::from_lengths([arr.len()]);
+    ListArray::new(
+        Arc::new(Field::new("item", arr.data_type().to_owned(), true)),
+        offsets,
+        arr,
+        None,
+    )
+}
+
 /// An extension trait for smart pointers. Provides an interface to get a
 /// raw pointer to the data (with metadata stripped away).
 ///

datafusion/core/tests/sql/aggregates.rs

@@ -47,25 +47,23 @@ async fn csv_query_array_agg_distinct() -> Result<()> {
     let column = actual[0].column(0);
     assert_eq!(column.len(), 1);
 
-    if let ScalarValue::List(Some(mut v), _) = ScalarValue::try_from_array(column, 0)? {
-        // workaround lack of Ord of ScalarValue
-        let cmp = |a: &ScalarValue, b: &ScalarValue| {
-            a.partial_cmp(b).expect("Can compare ScalarValues")
-        };
-        v.sort_by(cmp);
-        assert_eq!(
-            *v,
-            vec![
-                ScalarValue::UInt32(Some(1)),
-                ScalarValue::UInt32(Some(2)),
-                ScalarValue::UInt32(Some(3)),
-                ScalarValue::UInt32(Some(4)),
-                ScalarValue::UInt32(Some(5))
-            ]
-        );
-    } else {
-        unreachable!();
-    }
+    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")
+    };
+    scalars.sort_by(cmp);
+    assert_eq!(
+        scalars,
+        vec![
+            ScalarValue::UInt32(Some(1)),
+            ScalarValue::UInt32(Some(2)),
+            ScalarValue::UInt32(Some(3)),
+            ScalarValue::UInt32(Some(4)),
+            ScalarValue::UInt32(Some(5))
+        ]
+    );
 
     Ok(())
 }

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -30,7 +30,10 @@ use arrow::{
     error::ArrowError,
     record_batch::RecordBatch,
 };
-use datafusion_common::tree_node::{RewriteRecursion, TreeNode, TreeNodeRewriter};
+use datafusion_common::{
+    cast::{as_large_list_array, as_list_array},
+    tree_node::{RewriteRecursion, TreeNode, TreeNodeRewriter},
+};
 use datafusion_common::{
     exec_err, internal_err, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue,
 };
@@ -392,8 +395,11 @@ impl<'a> ConstEvaluator<'a> {
                         "Could not evaluate the expression, found a result of length {}",
                         a.len()
                     )
+                } else if as_list_array(&a).is_ok() || as_large_list_array(&a).is_ok() {
+                    Ok(ScalarValue::List(a))
                 } else {
-                    Ok(ScalarValue::try_from_array(&a, 0)?)
+                    // Non-ListArray
+                    ScalarValue::try_from_array(&a, 0)
                 }
             }
             ColumnarValue::Scalar(s) => Ok(s),

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

@@ -22,8 +22,11 @@ use crate::expressions::format_state_name;
 use crate::{AggregateExpr, PhysicalExpr};
 use arrow::array::ArrayRef;
 use arrow::datatypes::{DataType, Field};
+use arrow_array::Array;
+use datafusion_common::cast::as_list_array;
+use datafusion_common::utils::wrap_into_list_array;
+use datafusion_common::Result;
 use datafusion_common::ScalarValue;
-use datafusion_common::{internal_err, DataFusionError, Result};
 use datafusion_expr::Accumulator;
 use std::any::Any;
 use std::sync::Arc;
@@ -102,7 +105,7 @@ impl PartialEq<dyn Any> for ArrayAgg {
 
 #[derive(Debug)]
 pub(crate) struct ArrayAggAccumulator {
-    values: Vec<ScalarValue>,
+    values: Vec<ArrayRef>,
     datatype: DataType,
 }
 
@@ -117,50 +120,60 @@ impl ArrayAggAccumulator {
 }
 
 impl Accumulator for ArrayAggAccumulator {
+    // Append value like Int64Array(1,2,3)
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         if values.is_empty() {
             return Ok(());
         }
         assert!(values.len() == 1, "array_agg can only take 1 param!");
-        let arr = &values[0];
-        (0..arr.len()).try_for_each(|index| {
-            let scalar = ScalarValue::try_from_array(arr, index)?;
-            self.values.push(scalar);
-            Ok(())
-        })
+        let val = values[0].clone();
+        self.values.push(val);
+        Ok(())
     }
 
+    // Append value like ListArray(Int64Array(1,2,3), Int64Array(4,5,6))
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
         if states.is_empty() {
             return Ok(());
         }
         assert!(states.len() == 1, "array_agg states must be singleton!");
-        let arr = &states[0];
-        (0..arr.len()).try_for_each(|index| {
-            let scalar = ScalarValue::try_from_array(arr, index)?;
-            if let ScalarValue::List(Some(values), _) = scalar {
-                self.values.extend(values);
-                Ok(())
-            } else {
-                internal_err!("array_agg state must be list!")
-            }
-        })
+
+        let list_arr = as_list_array(&states[0])?;
+        for arr in list_arr.iter().flatten() {
+            self.values.push(arr);
+        }
+        Ok(())
     }
 
     fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(vec![self.evaluate()?])
     }
 
     fn evaluate(&self) -> Result<ScalarValue> {
-        Ok(ScalarValue::new_list(
-            Some(self.values.clone()),
-            self.datatype.clone(),
-        ))
+        // Transform Vec<ListArr> to ListArr
+
+        let element_arrays: Vec<&dyn Array> =
+            self.values.iter().map(|a| a.as_ref()).collect();
+
+        if element_arrays.is_empty() {
+            let arr = ScalarValue::new_list(&[], &self.datatype);
+            return Ok(ScalarValue::List(arr));
+        }
+
+        let concated_array = arrow::compute::concat(&element_arrays)?;
+        let list_array = wrap_into_list_array(concated_array);
+
+        Ok(ScalarValue::List(Arc::new(list_array)))
     }
 
     fn size(&self) -> usize {
-        std::mem::size_of_val(self) + ScalarValue::size_of_vec(&self.values)
-            - std::mem::size_of_val(&self.values)
+        std::mem::size_of_val(self)
+            + (std::mem::size_of::<ArrayRef>() * self.values.capacity())
+            + self
+                .values
+                .iter()
+                .map(|arr| arr.get_array_memory_size())
+                .sum::<usize>()
             + self.datatype.size()
             - std::mem::size_of_val(&self.datatype)
     }
@@ -176,72 +189,78 @@ mod tests {
     use arrow::array::Int32Array;
     use arrow::datatypes::*;
     use arrow::record_batch::RecordBatch;
+    use arrow_array::Array;
+    use arrow_array::ListArray;
+    use arrow_buffer::OffsetBuffer;
+    use datafusion_common::DataFusionError;
     use datafusion_common::Result;
 
     #[test]
     fn array_agg_i32() -> Result<()> {
         let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
 
-        let list = 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)),
-            ]),
-            DataType::Int32,
-        );
+        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));
 
         generic_test_op!(a, DataType::Int32, ArrayAgg, list, DataType::Int32)
     }
 
     #[test]
     fn array_agg_nested() -> Result<()> {
-        let l1 = ScalarValue::new_list(
-            Some(vec![
-                ScalarValue::new_list(
-                    Some(vec![
-                        ScalarValue::from(1i32),
-                        ScalarValue::from(2i32),
-                        ScalarValue::from(3i32),
-                    ]),
-                    DataType::Int32,
-                ),
-                ScalarValue::new_list(
-                    Some(vec![ScalarValue::from(4i32), ScalarValue::from(5i32)]),
-                    DataType::Int32,
-                ),
-            ]),
-            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        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 l2 = ScalarValue::new_list(
-            Some(vec![
-                ScalarValue::new_list(
-                    Some(vec![ScalarValue::from(6i32)]),
-                    DataType::Int32,
-                ),
-                ScalarValue::new_list(
-                    Some(vec![ScalarValue::from(7i32), ScalarValue::from(8i32)]),
-                    DataType::Int32,
-                ),
-            ]),
-            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        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 l3 = ScalarValue::new_list(
-            Some(vec![ScalarValue::new_list(
-                Some(vec![ScalarValue::from(9i32)]),
-                DataType::Int32,
-            )]),
-            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        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 = ScalarValue::new_list(
-            Some(vec![l1.clone(), l2.clone(), l3.clone()]),
-            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
+        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();