first draft

Signed-off-by: jayzhan211 <[email protected]>

Author: jayzhan211

datafusion/core/src/physical_planner.rs

@@ -85,7 +85,6 @@ use datafusion_expr::{
     DescribeTable, DmlStatement, Extension, Filter, RecursiveQuery, StringifiedPlan,
     WindowFrame, WindowFrameBound, WriteOp,
 };
-use datafusion_functions_aggregate::array_agg::array_agg_udaf;
 use datafusion_physical_expr::expressions::Literal;
 use datafusion_physical_expr::LexOrdering;
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
@@ -1840,11 +1839,11 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                 .unwrap_or(sqlparser::ast::NullTreatment::RespectNulls)
                 == NullTreatment::IgnoreNulls;
 
+            // TODO: Remove this after array_agg are all udafs
             let (agg_expr, filter, order_by) = match func_def {
-                AggregateFunctionDefinition::BuiltIn(
-                    datafusion_expr::AggregateFunction::ArrayAgg,
-                ) if !distinct && order_by.is_none() => {
-                    let sort_exprs = order_by.clone().unwrap_or(vec![]);
+                AggregateFunctionDefinition::UDF(udf)
+                    if udf.name() == "ARRAY_AGG" && (*distinct || order_by.is_some()) =>
+                {
                     let physical_sort_exprs = match order_by {
                         Some(exprs) => Some(create_physical_sort_exprs(
                             exprs,
@@ -1855,16 +1854,15 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                     };
                     let ordering_reqs: Vec<PhysicalSortExpr> =
                         physical_sort_exprs.clone().unwrap_or(vec![]);
-                    let agg_expr = udaf::create_aggregate_expr(
-                        &array_agg_udaf(),
+                    let fun = aggregates::AggregateFunction::ArrayAgg;
+                    let agg_expr = aggregates::create_aggregate_expr(
+                        &fun,
+                        *distinct,
                         &physical_args,
-                        args,
-                        &sort_exprs,
                         &ordering_reqs,
                         physical_input_schema,
                         name,
                         ignore_nulls,
-                        *distinct,
                     )?;
                     (agg_expr, filter, physical_sort_exprs)
                 }
@@ -1916,6 +1914,7 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                     (agg_expr, filter, physical_sort_exprs)
                 }
             };
+
             Ok((agg_expr, filter, order_by))
         }
         other => internal_err!("Invalid aggregate expression '{other:?}'"),

datafusion/expr/src/function.rs

@@ -106,9 +106,6 @@ pub struct StateFieldsArgs<'a> {
     /// The input type of the aggregate function.
     pub input_type: &'a DataType,
 
-    /// If the input type is nullable.
-    pub input_nullable: bool,
-
     /// The return type of the aggregate function.
     pub return_type: &'a DataType,
 

datafusion/functions-aggregate/src/array_agg.rs

@@ -22,16 +22,12 @@ use arrow::datatypes::DataType;
 use arrow_schema::Field;
 
 use datafusion_common::cast::as_list_array;
-use datafusion_common::utils::array_into_list_array;
+use datafusion_common::utils::array_into_list_array_nullable;
 use datafusion_common::Result;
 use datafusion_common::ScalarValue;
-use datafusion_expr::expr::AggregateFunction;
-use datafusion_expr::expr::AggregateFunctionDefinition;
-use datafusion_expr::function::AccumulatorArgs;
-use datafusion_expr::simplify::SimplifyInfo;
+use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
 use datafusion_expr::AggregateUDFImpl;
-use datafusion_expr::Expr;
 use datafusion_expr::{Accumulator, Signature, Volatility};
 use std::sync::Arc;
 
@@ -84,47 +80,17 @@ impl AggregateUDFImpl for ArrayAgg {
         ))))
     }
 
-    fn state_fields(
-        &self,
-        args: datafusion_expr::function::StateFieldsArgs,
-    ) -> Result<Vec<Field>> {
+    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
         Ok(vec![Field::new_list(
             format_state_name(args.name, "array_agg"),
             Field::new("item", args.input_type.clone(), true),
-            args.input_nullable,
+            true,
         )])
     }
 
     fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
         Ok(Box::new(ArrayAggAccumulator::try_new(acc_args.input_type)?))
     }
-
-    fn reverse_expr(&self) -> datafusion_expr::ReversedUDAF {
-        datafusion_expr::ReversedUDAF::Identical
-    }
-
-    fn simplify(
-        &self,
-    ) -> Option<datafusion_expr::function::AggregateFunctionSimplification> {
-        let simplify = |aggregate_function: AggregateFunction, _: &dyn SimplifyInfo| {
-            if aggregate_function.order_by.is_some() || aggregate_function.distinct {
-                Ok(Expr::AggregateFunction(AggregateFunction {
-                    func_def: AggregateFunctionDefinition::BuiltIn(
-                        datafusion_expr::aggregate_function::AggregateFunction::ArrayAgg,
-                    ),
-                    args: aggregate_function.args,
-                    distinct: aggregate_function.distinct,
-                    filter: aggregate_function.filter,
-                    order_by: aggregate_function.order_by,
-                    null_treatment: aggregate_function.null_treatment,
-                }))
-            } else {
-                Ok(Expr::AggregateFunction(aggregate_function))
-            }
-        };
-
-        Some(Box::new(simplify))
-    }
 }
 
 #[derive(Debug)]
@@ -150,8 +116,11 @@ impl Accumulator for ArrayAggAccumulator {
             return Ok(());
         }
         assert!(values.len() == 1, "array_agg can only take 1 param!");
-        let val = values[0].clone();
-        self.values.push(val);
+
+        let val = Arc::clone(&values[0]);
+        if val.len() > 0 {
+            self.values.push(val);
+        }
         Ok(())
     }
 
@@ -175,17 +144,15 @@ impl Accumulator for ArrayAggAccumulator {
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
         // 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));
+            return Ok(ScalarValue::new_null_list(self.datatype.clone(), true, 1));
         }
 
         let concated_array = arrow::compute::concat(&element_arrays)?;
-        let list_array = array_into_list_array(concated_array);
+        let list_array = array_into_list_array_nullable(concated_array);
 
         Ok(ScalarValue::List(Arc::new(list_array)))
     }

datafusion/functions-aggregate/src/first_last.rs

@@ -440,7 +440,6 @@ impl AggregateUDFImpl for LastValue {
         let StateFieldsArgs {
             name,
             input_type,
-            input_nullable: _,
             return_type: _,
             ordering_fields,
             is_distinct: _,

datafusion/functions-aggregate/src/lib.rs

@@ -56,8 +56,8 @@
 pub mod macros;
 
 pub mod approx_distinct;
-pub mod correlation;
 pub mod array_agg;
+pub mod correlation;
 pub mod count;
 pub mod covariance;
 pub mod first_last;
@@ -92,8 +92,8 @@ pub mod expr_fn {
     pub use super::approx_median::approx_median;
     pub use super::approx_percentile_cont::approx_percentile_cont;
     pub use super::approx_percentile_cont_with_weight::approx_percentile_cont_with_weight;
-    pub use super::average::avg;
     pub use super::array_agg::array_agg;
+    pub use super::average::avg;
     pub use super::bit_and_or_xor::bit_and;
     pub use super::bit_and_or_xor::bit_or;
     pub use super::bit_and_or_xor::bit_xor;

datafusion/functions-array/src/planner.rs

@@ -19,8 +19,9 @@
 
 use datafusion_common::{utils::list_ndims, DFSchema, Result};
 use datafusion_expr::{
+    expr::AggregateFunctionDefinition,
     planner::{ExprPlanner, PlannerResult, RawBinaryExpr, RawFieldAccessExpr},
-    sqlparser, AggregateFunction, Expr, ExprSchemable, GetFieldAccess,
+    sqlparser, Expr, ExprSchemable, GetFieldAccess,
 };
 use datafusion_functions::expr_fn::get_field;
 use datafusion_functions_aggregate::nth_value::nth_value_udaf;
@@ -153,8 +154,9 @@ impl ExprPlanner for FieldAccessPlanner {
 }
 
 fn is_array_agg(agg_func: &datafusion_expr::expr::AggregateFunction) -> bool {
-    agg_func.func_def
-        == datafusion_expr::expr::AggregateFunctionDefinition::BuiltIn(
-            AggregateFunction::ArrayAgg,
-        )
+    if let AggregateFunctionDefinition::UDF(udf) = &agg_func.func_def {
+        return udf.name() == "ARRAY_AGG";
+    }
+
+    false
 }

datafusion/physical-expr-common/src/aggregate/mod.rs

@@ -88,7 +88,6 @@ pub fn create_aggregate_expr(
         ordering_fields,
         is_distinct,
         input_type: input_exprs_types[0].clone(),
-        input_nullable: input_phy_exprs[0].nullable(schema)?,
     }))
 }
 
@@ -264,7 +263,6 @@ pub struct AggregateFunctionExpr {
     ordering_fields: Vec<Field>,
     is_distinct: bool,
     input_type: DataType,
-    input_nullable: bool,
 }
 
 impl AggregateFunctionExpr {
@@ -293,7 +291,6 @@ impl AggregateExpr for AggregateFunctionExpr {
         let args = StateFieldsArgs {
             name: &self.name,
             input_type: &self.input_type,
-            input_nullable: self.input_nullable,
             return_type: &self.data_type,
             ordering_fields: &self.ordering_fields,
             is_distinct: self.is_distinct,
@@ -303,11 +300,7 @@ impl AggregateExpr for AggregateFunctionExpr {
     }
 
     fn field(&self) -> Result<Field> {
-        Ok(Field::new(
-            &self.name,
-            self.data_type.clone(),
-            self.input_nullable,
-        ))
+        Ok(Field::new(&self.name, self.data_type.clone(), true))
     }
 
     fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {

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

@@ -30,7 +30,7 @@ use std::sync::Arc;
 
 use arrow::datatypes::Schema;
 
-use datafusion_common::{not_impl_err, Result};
+use datafusion_common::{internal_err, not_impl_err, Result};
 use datafusion_expr::AggregateFunction;
 
 use crate::expressions::{self};

datafusion/physical-plan/src/aggregates/no_grouping.rs

@@ -218,6 +218,7 @@ fn aggregate_batch(
                 Some(filter) => Cow::Owned(batch_filter(&batch, filter)?),
                 None => Cow::Borrowed(&batch),
             };
+
             // 1.3
             let values = &expr
                 .iter()

datafusion/proto/src/physical_plan/to_proto.rs

@@ -23,10 +23,9 @@ use datafusion::datasource::file_format::parquet::ParquetSink;
 use datafusion::physical_expr::window::{NthValueKind, SlidingAggregateWindowExpr};
 use datafusion::physical_expr::{PhysicalSortExpr, ScalarFunctionExpr};
 use datafusion::physical_plan::expressions::{
-    BinaryExpr, CaseExpr, CastExpr, Column, CumeDist, DistinctArrayAgg,
-    InListExpr, IsNotNullExpr, IsNullExpr, Literal, Max, Min, NegativeExpr, NotExpr,
-    NthValue, Ntile, OrderSensitiveArrayAgg, Rank, RankType, RowNumber, TryCastExpr,
-    WindowShift,
+    BinaryExpr, CaseExpr, CastExpr, Column, CumeDist, DistinctArrayAgg, InListExpr,
+    IsNotNullExpr, IsNullExpr, Literal, Max, Min, NegativeExpr, NotExpr, NthValue, Ntile,
+    OrderSensitiveArrayAgg, Rank, RankType, RowNumber, TryCastExpr, WindowShift,
 };
 use datafusion::physical_plan::udaf::AggregateFunctionExpr;
 use datafusion::physical_plan::windows::{BuiltInWindowExpr, PlainAggregateWindowExpr};
@@ -245,9 +244,7 @@ fn aggr_expr_to_aggr_fn(expr: &dyn AggregateExpr) -> Result<AggrFn> {
     let mut distinct = false;
 
     // TODO: remove
-    let inner = if aggr_expr.downcast_ref::<ArrayAgg>().is_some() {
-        protobuf::AggregateFunction::ArrayAgg
-    } else if aggr_expr.downcast_ref::<DistinctArrayAgg>().is_some() {
+    let inner = if aggr_expr.downcast_ref::<DistinctArrayAgg>().is_some() {
         distinct = true;
         protobuf::AggregateFunction::ArrayAgg
     } else if aggr_expr.downcast_ref::<OrderSensitiveArrayAgg>().is_some() {