Fix Type Coercion for UDF Arguments

datafusion/common/Cargo.toml

@@ -64,7 +64,7 @@ log = { workspace = true }
 object_store = { workspace = true, optional = true }
 parquet = { workspace = true, optional = true, default-features = true }
 paste = "1.0.15"
-pyo3 = { version = "0.23.3", optional = true }
+pyo3 = { version = "0.23.4", optional = true }
 recursive = { workspace = true, optional = true }
 sqlparser = { workspace = true }
 tokio = { workspace = true }

datafusion/expr-common/src/signature.rs

@@ -498,7 +498,7 @@ impl Signature {
         }
     }
 
-    /// A specified number of numeric arguments
+    /// A specified number of string arguments
     pub fn string(arg_count: usize, volatility: Volatility) -> Self {
         Self {
             type_signature: TypeSignature::String(arg_count),

datafusion/expr/src/type_coercion/functions.rs

@@ -53,14 +53,14 @@ pub fn data_types_with_scalar_udf(
     let type_signature = &signature.type_signature;
 
     if current_types.is_empty() {
-        if type_signature.supports_zero_argument() {
-            return Ok(vec![]);
+        return if type_signature.supports_zero_argument() {
+            Ok(vec![])
         } else if type_signature.used_to_support_zero_arguments() {
             // Special error to help during upgrade: https://github.com/apache/datafusion/issues/13763
-            return plan_err!("{} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.", func.name());
+            plan_err!("{} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.", func.name())
         } else {
-            return plan_err!("{} does not support zero arguments.", func.name());
-        }
+            plan_err!("{} does not support zero arguments.", func.name())
+        };
     }
 
     let valid_types =
@@ -91,14 +91,14 @@ pub fn data_types_with_aggregate_udf(
     let type_signature = &signature.type_signature;
 
     if current_types.is_empty() {
-        if type_signature.supports_zero_argument() {
-            return Ok(vec![]);
+        return if type_signature.supports_zero_argument() {
+            Ok(vec![])
         } else if type_signature.used_to_support_zero_arguments() {
             // Special error to help during upgrade: https://github.com/apache/datafusion/issues/13763
-            return plan_err!("{} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.", func.name());
+            plan_err!("{} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.", func.name())
         } else {
-            return plan_err!("{} does not support zero arguments.", func.name());
-        }
+            plan_err!("{} does not support zero arguments.", func.name())
+        };
     }
 
     let valid_types =
@@ -128,14 +128,14 @@ pub fn data_types_with_window_udf(
     let type_signature = &signature.type_signature;
 
     if current_types.is_empty() {
-        if type_signature.supports_zero_argument() {
-            return Ok(vec![]);
+        return if type_signature.supports_zero_argument() {
+            Ok(vec![])
         } else if type_signature.used_to_support_zero_arguments() {
             // Special error to help during upgrade: https://github.com/apache/datafusion/issues/13763
-            return plan_err!("{} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.", func.name());
+            plan_err!("{} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.", func.name())
         } else {
-            return plan_err!("{} does not support zero arguments.", func.name());
-        }
+            plan_err!("{} does not support zero arguments.", func.name())
+        };
     }
 
     let valid_types =
@@ -165,20 +165,20 @@ pub fn data_types(
     let type_signature = &signature.type_signature;
 
     if current_types.is_empty() {
-        if type_signature.supports_zero_argument() {
-            return Ok(vec![]);
+        return if type_signature.supports_zero_argument() {
+            Ok(vec![])
         } else if type_signature.used_to_support_zero_arguments() {
             // Special error to help during upgrade: https://github.com/apache/datafusion/issues/13763
-            return plan_err!(
+            plan_err!(
                 "signature {:?} does not support zero arguments. Use TypeSignature::Nullary for zero arguments.",
                 type_signature
-            );
+            )
         } else {
-            return plan_err!(
+            plan_err!(
                 "signature {:?} does not support zero arguments.",
                 type_signature
-            );
-        }
+            )
+        };
     }
 
     let valid_types = get_valid_types(type_signature, current_types)?;
@@ -584,23 +584,12 @@ fn get_valid_types(
                 match target_type_class {
                     TypeSignatureClass::Native(native_type) => {
                         let target_type = native_type.native();
-                        if &logical_type == target_type {
-                            return target_type.default_cast_for(current_type);
-                        }
-
-                        if logical_type == NativeType::Null {
-                            return target_type.default_cast_for(current_type);
-                        }
-
-                        if target_type.is_integer() && logical_type.is_integer() {
-                            return target_type.default_cast_for(current_type);
-                        }
-
-                        internal_err!(
-                            "Expect {} but received {}",
-                            target_type_class,
-                            current_type
-                        )
+                        // if target_type == &NativeType::String {
+                        //     ...
+                        // } else {
+                        //     target_type.default_cast_for(current_type)
+                        // }
+                        target_type.default_cast_for(current_type)
                     }
                     // Not consistent with Postgres and DuckDB but to avoid regression we implicit cast string to timestamp
                     TypeSignatureClass::Timestamp
@@ -881,7 +870,7 @@ fn coerced_from<'a>(
             Null | Timestamp(_, None) | Date32 | Utf8 | LargeUtf8,
         ) => Some(type_into.clone()),
         (Interval(_), Utf8 | LargeUtf8) => Some(type_into.clone()),
-        // We can go into a Utf8View from a Utf8 or LargeUtf8
+        // We can go into a Utf8View from an Utf8 or LargeUtf8
         (Utf8View, Utf8 | LargeUtf8 | Null) => Some(type_into.clone()),
         // Any type can be coerced into strings
         (Utf8 | LargeUtf8, _) => Some(type_into.clone()),
@@ -931,12 +920,17 @@ fn coerced_from<'a>(
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
 
-    use crate::Volatility;
-
-    use super::*;
-    use arrow::datatypes::Field;
-    use datafusion_common::assert_contains;
+    use crate::type_coercion::functions::{
+        can_coerce_from, coerced_from, data_types, get_valid_types, maybe_data_types,
+    };
+    use crate::TypeSignature;
+    use arrow::datatypes::{DataType, Field, TimeUnit};
+    use datafusion_common::{assert_contains, Result};
+    use datafusion_expr_common::signature::{
+        Signature, Volatility, FIXED_SIZE_LIST_WILDCARD,
+    };
 
     #[test]
     fn test_string_conversion() {
@@ -1124,7 +1118,7 @@ mod tests {
             Volatility::Stable,
         );
 
-        let coerced_data_types = data_types("test", &current_types, &signature).unwrap();
+        let coerced_data_types = data_types("test", &current_types, &signature)?;
         assert_eq!(coerced_data_types, current_types);
 
         // make sure it can't coerce to a different size
@@ -1140,7 +1134,7 @@ mod tests {
             vec![DataType::FixedSizeList(Arc::clone(&inner), 2)],
             Volatility::Stable,
         );
-        let coerced_data_types = data_types("test", &current_types, &signature).unwrap();
+        let coerced_data_types = data_types("test", &current_types, &signature)?;
         assert_eq!(coerced_data_types, current_types);
 
         Ok(())

datafusion/functions/src/string/ascii.rs

@@ -19,9 +19,12 @@ use crate::utils::make_scalar_function;
 use arrow::array::{ArrayAccessor, ArrayIter, ArrayRef, AsArray, Int32Array};
 use arrow::datatypes::DataType;
 use arrow::error::ArrowError;
+use datafusion_common::types::logical_string;
 use datafusion_common::{internal_err, Result};
-use datafusion_expr::{ColumnarValue, Documentation};
-use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass,
+    Volatility,
+};
 use datafusion_macros::user_doc;
 use std::any::Any;
 use std::sync::Arc;
@@ -61,7 +64,10 @@ impl Default for AsciiFunc {
 impl AsciiFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::string(1, Volatility::Immutable),
+            signature: Signature::coercible(
+                vec![TypeSignatureClass::Native(logical_string())],
+                Volatility::Immutable,
+            ),
         }
     }
 }
@@ -130,7 +136,7 @@ pub fn ascii(args: &[ArrayRef]) -> Result<ArrayRef> {
             let string_array = args[0].as_string_view();
             Ok(calculate_ascii(string_array)?)
         }
-        _ => internal_err!("Unsupported data type"),
+        other => internal_err!("Unsupported data type for ascii: {:?}", other),
     }
 }
 
@@ -143,7 +149,7 @@ mod tests {
     use datafusion_common::{Result, ScalarValue};
     use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
 
-    macro_rules! test_ascii {
+    macro_rules! test_ascii_invoke {
         ($INPUT:expr, $EXPECTED:expr) => {
             test_function!(
                 AsciiFunc::new(),
@@ -175,11 +181,11 @@ mod tests {
     }
 
     #[test]
-    fn test_functions() -> Result<()> {
-        test_ascii!(Some(String::from("x")), Ok(Some(120)));
-        test_ascii!(Some(String::from("a")), Ok(Some(97)));
-        test_ascii!(Some(String::from("")), Ok(Some(0)));
-        test_ascii!(None, Ok(None));
+    fn test_ascii_invoke() -> Result<()> {
+        test_ascii_invoke!(Some(String::from("x")), Ok(Some(120)));
+        test_ascii_invoke!(Some(String::from("a")), Ok(Some(97)));
+        test_ascii_invoke!(Some(String::from("")), Ok(Some(0)));
+        test_ascii_invoke!(None, Ok(None));
         Ok(())
     }
 }

datafusion/functions/src/string/bit_length.rs

@@ -20,9 +20,12 @@ use arrow::datatypes::DataType;
 use std::any::Any;
 
 use crate::utils::utf8_to_int_type;
+use datafusion_common::types::logical_string;
 use datafusion_common::{exec_err, Result, ScalarValue};
-use datafusion_expr::{ColumnarValue, Documentation, Volatility};
-use datafusion_expr::{ScalarUDFImpl, Signature};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass,
+    Volatility,
+};
 use datafusion_macros::user_doc;
 
 #[user_doc(
@@ -55,7 +58,10 @@ impl Default for BitLengthFunc {
 impl BitLengthFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::string(1, Volatility::Immutable),
+            signature: Signature::coercible(
+                vec![TypeSignatureClass::Native(logical_string())],
+                Volatility::Immutable,
+            ),
         }
     }
 }

datafusion/functions/src/string/contains.rs

@@ -21,10 +21,12 @@ use arrow::compute::contains as arrow_contains;
 use arrow::datatypes::DataType;
 use arrow::datatypes::DataType::{Boolean, LargeUtf8, Utf8, Utf8View};
 use datafusion_common::exec_err;
+use datafusion_common::types::logical_string;
 use datafusion_common::DataFusionError;
 use datafusion_common::Result;
 use datafusion_expr::{
-    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature,
+    TypeSignatureClass, Volatility,
 };
 use datafusion_macros::user_doc;
 use std::any::Any;
@@ -59,7 +61,16 @@ impl Default for ContainsFunc {
 impl ContainsFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::string(2, Volatility::Immutable),
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::String(2),
+                    TypeSignature::Coercible(vec![
+                        TypeSignatureClass::Native(logical_string()),
+                        TypeSignatureClass::Native(logical_string()),
+                    ]),
+                ],
+                Volatility::Immutable,
+            ),
         }
     }
 }

datafusion/functions/src/string/ends_with.rs

@@ -22,9 +22,12 @@ use arrow::array::ArrayRef;
 use arrow::datatypes::DataType;
 
 use crate::utils::make_scalar_function;
+use datafusion_common::types::logical_string;
 use datafusion_common::{internal_err, Result};
-use datafusion_expr::{ColumnarValue, Documentation, Volatility};
-use datafusion_expr::{ScalarUDFImpl, Signature};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature,
+    TypeSignatureClass, Volatility,
+};
 use datafusion_macros::user_doc;
 
 #[user_doc(
@@ -62,7 +65,16 @@ impl Default for EndsWithFunc {
 impl EndsWithFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::string(2, Volatility::Immutable),
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::String(2),
+                    TypeSignature::Coercible(vec![
+                        TypeSignatureClass::Native(logical_string()),
+                        TypeSignatureClass::Native(logical_string()),
+                    ]),
+                ],
+                Volatility::Immutable,
+            ),
         }
     }
 }

datafusion/functions/src/string/octet_length.rs

@@ -20,9 +20,12 @@ use arrow::datatypes::DataType;
 use std::any::Any;
 
 use crate::utils::utf8_to_int_type;
+use datafusion_common::types::logical_string;
 use datafusion_common::{exec_err, Result, ScalarValue};
-use datafusion_expr::{ColumnarValue, Documentation, Volatility};
-use datafusion_expr::{ScalarUDFImpl, Signature};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass,
+    Volatility,
+};
 use datafusion_macros::user_doc;
 
 #[user_doc(
@@ -55,7 +58,10 @@ impl Default for OctetLengthFunc {
 impl OctetLengthFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::string(1, Volatility::Immutable),
+            signature: Signature::coercible(
+                vec![TypeSignatureClass::Native(logical_string())],
+                Volatility::Immutable,
+            ),
         }
     }
 }

datafusion/functions/src/string/starts_with.rs

@@ -23,9 +23,12 @@ use arrow::datatypes::DataType;
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 
 use crate::utils::make_scalar_function;
+use datafusion_common::types::logical_string;
 use datafusion_common::{internal_err, Result, ScalarValue};
-use datafusion_expr::{ColumnarValue, Documentation, Expr, Like};
-use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
+use datafusion_expr::{
+    ColumnarValue, Documentation, Expr, Like, ScalarUDFImpl, Signature, TypeSignature,
+    TypeSignatureClass, Volatility,
+};
 use datafusion_macros::user_doc;
 
 /// Returns true if string starts with prefix.
@@ -64,7 +67,16 @@ impl Default for StartsWithFunc {
 impl StartsWithFunc {
     pub fn new() -> Self {
         Self {
-            signature: Signature::string(2, Volatility::Immutable),
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::String(2),
+                    TypeSignature::Coercible(vec![
+                        TypeSignatureClass::Native(logical_string()),
+                        TypeSignatureClass::Native(logical_string()),
+                    ]),
+                ],
+                Volatility::Immutable,
+            ),
         }
     }
 }