Introduce Signature::Coercible (#12275)

* introduce signature float Signed-off-by: jayzhan211 <[email protected]> * fix test Signed-off-by: jayzhan211 <[email protected]> * change float to coercible Signed-off-by: jayzhan211 <[email protected]> * fmt Signed-off-by: jayzhan211 <[email protected]> * rm test Signed-off-by: jayzhan211 <[email protected]> * add comment Signed-off-by: jayzhan211 <[email protected]> * typo Signed-off-by: jayzhan211 <[email protected]> --------- Signed-off-by: jayzhan211 <[email protected]>
apache · Sep 2, 2024 · 8db30e2 · 8db30e2
1 parent ac74cd3
commit 8db30e2
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 18 deletions.
diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs
@@ -2428,7 +2428,8 @@ mod tests {
         let df: Vec<RecordBatch> = df.select(aggr_expr)?.collect().await?;
 
         assert_batches_sorted_eq!(
-            ["+-------------+----------+-----------------+---------------+--------+-----+------+----+------+",
+            [
+                "+-------------+----------+-----------------+---------------+--------+-----+------+----+------+",
                 "| first_value | last_val | approx_distinct | approx_median | median | max | min  | c2 | c3   |",
                 "+-------------+----------+-----------------+---------------+--------+-----+------+----+------+",
                 "|             |          |                 |               |        |     |      | 1  | -85  |",
@@ -2452,7 +2453,8 @@ mod tests {
                 "| -85         | 45       | 8               | -34           | 45     | 83  | -85  | 3  | -72  |",
                 "| -85         | 65       | 17              | -17           | 65     | 83  | -101 | 5  | -101 |",
                 "| -85         | 83       | 5               | -25           | 83     | 83  | -85  | 2  | -48  |",
-                "+-------------+----------+-----------------+---------------+--------+-----+------+----+------+"],
+                "+-------------+----------+-----------------+---------------+--------+-----+------+----+------+",
+            ],
             &df
         );
 

diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs
@@ -105,6 +105,11 @@ pub enum TypeSignature {
     Uniform(usize, Vec<DataType>),
     /// Exact number of arguments of an exact type
     Exact(Vec<DataType>),
+    /// The number of arguments that can be coerced to in order
+    /// For example, `Coercible(vec![DataType::Float64])` accepts
+    /// arguments like `vec![DataType::Int32]` or `vec![DataType::Float32]`
+    /// since i32 and f32 can be casted to f64
+    Coercible(Vec<DataType>),
     /// Fixed number of arguments of arbitrary types
     /// If a function takes 0 argument, its `TypeSignature` should be `Any(0)`
     Any(usize),
@@ -188,7 +193,7 @@ impl TypeSignature {
             TypeSignature::Numeric(num) => {
                 vec![format!("Numeric({})", num)]
             }
-            TypeSignature::Exact(types) => {
+            TypeSignature::Exact(types) | TypeSignature::Coercible(types) => {
                 vec![Self::join_types(types, ", ")]
             }
             TypeSignature::Any(arg_count) => {
@@ -300,6 +305,14 @@ impl Signature {
             volatility,
         }
     }
+    /// Target coerce types in order
+    pub fn coercible(target_types: Vec<DataType>, volatility: Volatility) -> Self {
+        Self {
+            type_signature: TypeSignature::Coercible(target_types),
+            volatility,
+        }
+    }
+
     /// A specified number of arguments of any type
     pub fn any(arg_count: usize, volatility: Volatility) -> Self {
         Signature {

diff --git a/datafusion/expr-common/src/type_coercion/aggregates.rs b/datafusion/expr-common/src/type_coercion/aggregates.rs
@@ -128,9 +128,11 @@ pub fn check_arg_count(
                 );
             }
         }
-        TypeSignature::UserDefined | TypeSignature::Numeric(_) => {
+        TypeSignature::UserDefined
+        | TypeSignature::Numeric(_)
+        | TypeSignature::Coercible(_) => {
             // User-defined signature is validated in `coerce_types`
-            // Numreic signature is validated in `get_valid_types`
+            // Numeric and Coercible signature is validated in `get_valid_types`
         }
         _ => {
             return internal_err!(

diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
@@ -175,7 +175,14 @@ fn try_coerce_types(
     let mut valid_types = valid_types;
 
     // Well-supported signature that returns exact valid types.
-    if !valid_types.is_empty() && matches!(type_signature, TypeSignature::UserDefined) {
+    if !valid_types.is_empty()
+        && matches!(
+            type_signature,
+            TypeSignature::UserDefined
+                | TypeSignature::Numeric(_)
+                | TypeSignature::Coercible(_)
+        )
+    {
         // exact valid types
         assert_eq!(valid_types.len(), 1);
         let valid_types = valid_types.swap_remove(0);
@@ -397,6 +404,30 @@ fn get_valid_types(
 
             vec![vec![valid_type; *number]]
         }
+        TypeSignature::Coercible(target_types) => {
+            if target_types.is_empty() {
+                return plan_err!(
+                    "The signature expected at least one argument but received {}",
+                    current_types.len()
+                );
+            }
+            if target_types.len() != current_types.len() {
+                return plan_err!(
+                    "The signature expected {} arguments but received {}",
+                    target_types.len(),
+                    current_types.len()
+                );
+            }
+
+            for (data_type, target_type) in current_types.iter().zip(target_types.iter())
+            {
+                if !can_cast_types(data_type, target_type) {
+                    return plan_err!("{data_type} is not coercible to {target_type}");
+                }
+            }
+
+            vec![target_types.to_owned()]
+        }
         TypeSignature::Uniform(number, valid_types) => valid_types
             .iter()
             .map(|valid_type| (0..*number).map(|_| valid_type.clone()).collect())

diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs
@@ -68,7 +68,10 @@ impl Stddev {
     /// Create a new STDDEV aggregate function
     pub fn new() -> Self {
         Self {
-            signature: Signature::numeric(1, Volatility::Immutable),
+            signature: Signature::coercible(
+                vec![DataType::Float64],
+                Volatility::Immutable,
+            ),
             alias: vec!["stddev_samp".to_string()],
         }
     }
@@ -88,11 +91,7 @@ impl AggregateUDFImpl for Stddev {
         &self.signature
     }
 
-    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        if !arg_types[0].is_numeric() {
-            return plan_err!("Stddev requires numeric input types");
-        }
-
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
         Ok(DataType::Float64)
     }
 

diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs
@@ -79,7 +79,10 @@ impl VarianceSample {
     pub fn new() -> Self {
         Self {
             aliases: vec![String::from("var_sample"), String::from("var_samp")],
-            signature: Signature::numeric(1, Volatility::Immutable),
+            signature: Signature::coercible(
+                vec![DataType::Float64],
+                Volatility::Immutable,
+            ),
         }
     }
 }
@@ -97,11 +100,7 @@ impl AggregateUDFImpl for VarianceSample {
         &self.signature
     }
 
-    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        if !arg_types[0].is_numeric() {
-            return plan_err!("Variance requires numeric input types");
-        }
-
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
         Ok(DataType::Float64)
     }