Implement set_agg aggregate function

ballista/rust/core/proto/ballista.proto

@@ -169,6 +169,7 @@ enum AggregateFunction {
   COUNT = 4;
   APPROX_DISTINCT = 5;
   ARRAY_AGG = 6;
+  SET_AGG = 7;
 }
 
 message AggregateExprNode {

ballista/rust/core/src/serde/logical_plan/to_proto.rs

@@ -1129,6 +1129,7 @@ impl TryInto<protobuf::LogicalExprNode> for &Expr {
                         protobuf::AggregateFunction::ApproxDistinct
                     }
                     AggregateFunction::ArrayAgg => protobuf::AggregateFunction::ArrayAgg,
+                    AggregateFunction::SetAgg => protobuf::AggregateFunction::SetAgg,
                     AggregateFunction::Min => protobuf::AggregateFunction::Min,
                     AggregateFunction::Max => protobuf::AggregateFunction::Max,
                     AggregateFunction::Sum => protobuf::AggregateFunction::Sum,
@@ -1364,6 +1365,7 @@ impl From<&AggregateFunction> for protobuf::AggregateFunction {
             AggregateFunction::Count => Self::Count,
             AggregateFunction::ApproxDistinct => Self::ApproxDistinct,
             AggregateFunction::ArrayAgg => Self::ArrayAgg,
+            AggregateFunction::SetAgg => Self::SetAgg,
         }
     }
 }

ballista/rust/core/src/serde/mod.rs

@@ -119,6 +119,7 @@ impl From<protobuf::AggregateFunction> for AggregateFunction {
                 AggregateFunction::ApproxDistinct
             }
             protobuf::AggregateFunction::ArrayAgg => AggregateFunction::ArrayAgg,
+            protobuf::AggregateFunction::SetAgg => AggregateFunction::SetAgg,
         }
     }
 }

datafusion/src/physical_plan/aggregates.rs

@@ -64,6 +64,8 @@ pub enum AggregateFunction {
     ApproxDistinct,
     /// array_agg
     ArrayAgg,
+    /// set_agg
+    SetAgg,
 }
 
 impl fmt::Display for AggregateFunction {
@@ -84,6 +86,7 @@ impl FromStr for AggregateFunction {
             "sum" => AggregateFunction::Sum,
             "approx_distinct" => AggregateFunction::ApproxDistinct,
             "array_agg" => AggregateFunction::ArrayAgg,
+            "set_agg" => AggregateFunction::SetAgg,
             _ => {
                 return Err(DataFusionError::Plan(format!(
                     "There is no built-in function named {}",
@@ -122,6 +125,11 @@ pub fn return_type(
             coerced_data_types[0].clone(),
             true,
         )))),
+        AggregateFunction::SetAgg => Ok(DataType::List(Box::new(Field::new(
+            "item",
+            coerced_data_types[0].clone(),
+            true,
+        )))),
     }
 }
 
@@ -192,6 +200,11 @@ pub fn create_aggregate_expr(
             name,
             coerced_exprs_types[0].clone(),
         )),
+        (AggregateFunction::SetAgg, _) => Arc::new(expressions::SetAgg::new(
+            coerced_phy_exprs[0].clone(),
+            name,
+            coerced_exprs_types[0].clone(),
+        )),
         (AggregateFunction::Min, _) => Arc::new(expressions::Min::new(
             coerced_phy_exprs[0].clone(),
             name,
@@ -245,7 +258,8 @@ pub fn signature(fun: &AggregateFunction) -> Signature {
     match fun {
         AggregateFunction::Count
         | AggregateFunction::ApproxDistinct
-        | AggregateFunction::ArrayAgg => Signature::any(1, Volatility::Immutable),
+        | AggregateFunction::ArrayAgg
+        | AggregateFunction::SetAgg => Signature::any(1, Volatility::Immutable),
         AggregateFunction::Min | AggregateFunction::Max => {
             let valid = STRINGS
                 .iter()

datafusion/src/physical_plan/coercion_rule/aggregate_rule.rs

@@ -59,6 +59,7 @@ pub(crate) fn coerce_types(
             Ok(input_types.to_vec())
         }
         AggregateFunction::ArrayAgg => Ok(input_types.to_vec()),
+        AggregateFunction::SetAgg => Ok(input_types.to_vec()),
         AggregateFunction::Min | AggregateFunction::Max => {
             // min and max support the dictionary data type
             // unpack the dictionary to get the value

datafusion/src/physical_plan/expressions/mod.rs

@@ -50,6 +50,7 @@ mod nth_value;
 mod nullif;
 mod rank;
 mod row_number;
+mod set_agg;
 mod sum;
 mod try_cast;
 
@@ -84,6 +85,7 @@ pub use nth_value::NthValue;
 pub use nullif::{nullif_func, SUPPORTED_NULLIF_TYPES};
 pub use rank::{dense_rank, percent_rank, rank};
 pub use row_number::RowNumber;
+pub use set_agg::SetAgg;
 pub(crate) use sum::is_sum_support_arg_type;
 pub use sum::{sum_return_type, Sum};
 pub use try_cast::{try_cast, TryCastExpr};