Move Median to functions-aggregate and Introduce Numeric signature (#…

…10644) * introduce median udaf Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * rm agg median Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * rm old median Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * introduce numeric signature Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * address comment Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * fix doc Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * add proto roundtrip Signed-off-by: jayzhan211 <jayzhan211@gmail.com> --------- Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
apache · May 26, 2024 · 26b44f4 · 26b44f4
1 parent e40f50b
commit 26b44f4
Show file tree

Hide file tree

Showing 28 changed files with 258 additions and 133 deletions.
diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs
@@ -50,10 +50,11 @@ use datafusion_common::{
 };
 use datafusion_expr::lit;
 use datafusion_expr::{
- avg, count, max, median, min, stddev, utils::COUNT_STAR_EXPANSION,
+ avg, count, max, min, stddev, utils::COUNT_STAR_EXPANSION,
  TableProviderFilterPushDown, UNNAMED_TABLE,
 };
 use datafusion_expr::{case, is_null, sum};
+use datafusion_functions_aggregate::expr_fn::median;
 
 use async_trait::async_trait;
 

diff --git a/datafusion/expr/src/aggregate_function.rs b/datafusion/expr/src/aggregate_function.rs
@@ -43,8 +43,6 @@ pub enum AggregateFunction {
  Max,
  /// Average
  Avg,
- /// Median
- Median,
  /// Approximate distinct function
  ApproxDistinct,
  /// Aggregation into an array
@@ -114,7 +112,6 @@ impl AggregateFunction {
  Min => "MIN",
  Max => "MAX",
  Avg => "AVG",
- Median => "MEDIAN",
  ApproxDistinct => "APPROX_DISTINCT",
  ArrayAgg => "ARRAY_AGG",
  FirstValue => "FIRST_VALUE",
@@ -168,7 +165,6 @@ impl FromStr for AggregateFunction {
  "count" => AggregateFunction::Count,
  "max" => AggregateFunction::Max,
  "mean" => AggregateFunction::Avg,
- "median" => AggregateFunction::Median,
  "min" => AggregateFunction::Min,
  "sum" => AggregateFunction::Sum,
  "array_agg" => AggregateFunction::ArrayAgg,
@@ -275,9 +271,7 @@ impl AggregateFunction {
  AggregateFunction::ApproxPercentileContWithWeight => {
  Ok(coerced_data_types[0].clone())
  }
- AggregateFunction::ApproxMedian | AggregateFunction::Median => {
- Ok(coerced_data_types[0].clone())
- }
+ AggregateFunction::ApproxMedian => Ok(coerced_data_types[0].clone()),
  AggregateFunction::Grouping => Ok(DataType::Int32),
  AggregateFunction::FirstValue
  | AggregateFunction::LastValue
@@ -335,7 +329,6 @@ impl AggregateFunction {
  | AggregateFunction::VariancePop
  | AggregateFunction::Stddev
  | AggregateFunction::StddevPop
- | AggregateFunction::Median
  | AggregateFunction::ApproxMedian
  | AggregateFunction::FirstValue
  | AggregateFunction::LastValue => {

diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
@@ -296,18 +296,6 @@ pub fn approx_distinct(expr: Expr) -> Expr {
  ))
 }
 
-/// Calculate the median for `expr`.
-pub fn median(expr: Expr) -> Expr {
- Expr::AggregateFunction(AggregateFunction::new(
- aggregate_function::AggregateFunction::Median,
- vec![expr],
- false,
- None,
- None,
- None,
- ))
-}
-
 /// Calculate an approximation of the median for `expr`.
 pub fn approx_median(expr: Expr) -> Expr {
  Expr::AggregateFunction(AggregateFunction::new(

diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs
@@ -82,6 +82,9 @@ pub struct AccumulatorArgs<'a> {
 
  /// The number of arguments the aggregate function takes.
  pub args_num: usize,
+
+ /// The name of the expression
+ pub name: &'a str,
 }
 
 /// [`StateFieldsArgs`] contains information about the fields that an

diff --git a/datafusion/expr/src/signature.rs b/datafusion/expr/src/signature.rs
@@ -119,6 +119,9 @@ pub enum TypeSignature {
  OneOf(Vec<TypeSignature>),
  /// Specifies Signatures for array functions
  ArraySignature(ArrayFunctionSignature),
+ /// Fixed number of arguments of numeric types.
+ /// See <https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#method.is_numeric> to know which type is considered numeric
+ Numeric(usize),
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -178,6 +181,9 @@ impl TypeSignature {
  .collect::<Vec<String>>()
  .join(", ")]
  }
+ TypeSignature::Numeric(num) => {
+ vec![format!("Numeric({})", num)]
+ }
  TypeSignature::Exact(types) => {
  vec![Self::join_types(types, ", ")]
  }
@@ -259,6 +265,14 @@ impl Signature {
  volatility,
  }
  }
+
+ pub fn numeric(num: usize, volatility: Volatility) -> Self {
+ Self {
+ type_signature: TypeSignature::Numeric(num),
+ volatility,
+ }
+ }
+
  /// An arbitrary number of arguments of any type.
  pub fn variadic_any(volatility: Volatility) -> Self {
  Self {

diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs
@@ -332,7 +332,7 @@ impl TreeNode for Expr {
  Ok(Expr::AggregateFunction(AggregateFunction::new_udf(
  fun,
  new_args,
- false,
+ distinct,
  new_filter,
  new_order_by,
  null_treatment,

diff --git a/datafusion/expr/src/type_coercion/aggregates.rs b/datafusion/expr/src/type_coercion/aggregates.rs
@@ -283,9 +283,9 @@ pub fn coerce_types(
  }
  Ok(input_types.to_vec())
  }
- AggregateFunction::Median
- | AggregateFunction::FirstValue
- | AggregateFunction::LastValue => Ok(input_types.to_vec()),
+ AggregateFunction::FirstValue | AggregateFunction::LastValue => {
+  Ok(input_types.to_vec())
+ }
  AggregateFunction::NthValue => Ok(input_types.to_vec()),
  AggregateFunction::Grouping => Ok(vec![input_types[0].clone()]),
  AggregateFunction::StringAgg => {
@@ -355,6 +355,10 @@ pub fn check_arg_count(
  );
  }
  }
+ TypeSignature::UserDefined | TypeSignature::Numeric(_) => {
+ // User-defined signature is validated in `coerce_types`
+ // Numreic signature is validated in `get_valid_types`
+ }
  _ => {
  return internal_err!(
  "Aggregate functions do not support this {signature:?}"

diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
@@ -322,6 +322,38 @@ fn get_valid_types(
  .iter()
  .map(|valid_type| current_types.iter().map(|_| valid_type.clone()).collect())
  .collect(),
+ TypeSignature::Numeric(number) => {
+ if *number < 1 {
+ return plan_err!(
+ "The signature expected at least one argument but received {}",
+ current_types.len()
+ );
+ }
+ if *number != current_types.len() {
+ return plan_err!(
+ "The signature expected {} arguments but received {}",
+ number,
+ current_types.len()
+ );
+ }
+
+ let mut valid_type = current_types.first().unwrap().clone();
+ for t in current_types.iter().skip(1) {
+ if let Some(coerced_type) =
+ comparison_binary_numeric_coercion(&valid_type, t)
+ {
+ valid_type = coerced_type;
+ } else {
+ return plan_err!(
+ "{} and {} are not coercible to a common numeric type",
+ valid_type,
+ t
+ );
+ }
+ }
+
+ vec![vec![valid_type; *number]]
+ }
  TypeSignature::Uniform(number, valid_types) => valid_types
  .iter()
  .map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())

diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs
@@ -193,9 +193,10 @@ impl AggregateUDF {
  self.inner.create_groups_accumulator()
  }
 
- pub fn coerce_types(&self, _args: &[DataType]) -> Result<Vec<DataType>> {
+ pub fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
  not_impl_err!("coerce_types not implemented for {:?} yet", self.name())
  }
+
  /// Do the function rewrite
  ///
  /// See [`AggregateUDFImpl::simplify`] for more details.

diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml
@@ -39,6 +39,7 @@ path = "src/lib.rs"
 
 [dependencies]
 arrow = { workspace = true }
+arrow-schema = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }

diff --git a/datafusion/functions-aggregate/src/lib.rs b/datafusion/functions-aggregate/src/lib.rs
@@ -57,6 +57,7 @@ pub mod macros;
 
 pub mod covariance;
 pub mod first_last;
+pub mod median;
 
 use datafusion_common::Result;
 use datafusion_execution::FunctionRegistry;
@@ -68,6 +69,7 @@ use std::sync::Arc;
 pub mod expr_fn {
  pub use super::covariance::covar_samp;
  pub use super::first_last::first_value;
+ pub use super::median::median;
 }
 
 /// Registers all enabled packages with a [`FunctionRegistry`]
@@ -76,6 +78,7 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
  first_last::first_value_udaf(),
  covariance::covar_samp_udaf(),
  covariance::covar_pop_udaf(),
+ median::median_udaf(),
  ];
 
  functions.into_iter().try_for_each(|udf| {