Support Signature::Any(0), and add tests for zero arguments

apache · Dec 21, 2024 · 3f0822a · 3f0822a
1 parent 7ba33c7
commit 3f0822a
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
@@ -1119,6 +1119,45 @@ async fn create_scalar_function_from_sql_statement() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_valid_zero_argument_signatures() {
+    let signatures = vec![
+        Signature::exact(vec![], Volatility::Immutable),
+        Signature::any(0, Volatility::Immutable),
+        Signature::variadic(vec![], Volatility::Immutable),
+        Signature::variadic_any(Volatility::Immutable),
+        Signature::uniform(0, vec![], Volatility::Immutable),
+        Signature::coercible(vec![], Volatility::Immutable),
+        Signature::comparable(0, Volatility::Immutable),
+        Signature::nullary(Volatility::Immutable),
+    ];
+    for signature in signatures {
+        let ctx = SessionContext::new();
+        let udf = ScalarFunctionWrapper {
+            name: "bad_signature".to_string(),
+            expr: lit(1),
+            signature,
+            return_type: DataType::Int32,
+        };
+        ctx.register_udf(ScalarUDF::from(udf));
+        let results = ctx
+            .sql("select bad_signature()")
+            .await
+            .unwrap()
+            .collect()
+            .await
+            .unwrap();
+        let expected = vec![
+            "+-----------------+",
+            "| bad_signature() |",
+            "+-----------------+",
+            "| 1               |",
+            "+-----------------+",
+        ];
+        assert_batches_eq!(expected, &results);
+    }
+}
+
 /// Saves whatever is passed to it as a scalar function
 #[derive(Debug, Default)]
 struct RecordingFunctionFactory {

diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs
@@ -343,6 +343,12 @@ impl TypeSignature {
     pub fn supports_zero_argument(&self) -> bool {
         match &self {
             TypeSignature::Exact(vec) => vec.is_empty(),
+            TypeSignature::Any(0) => true,
+            TypeSignature::Variadic(vec) => vec.is_empty(),
+            TypeSignature::VariadicAny => true,
+            TypeSignature::Uniform(0, _) => true,
+            TypeSignature::Coercible(vec) => vec.is_empty(),
+            TypeSignature::Comparable(0) => true,
             TypeSignature::Nullary => true,
             TypeSignature::OneOf(types) => types
                 .iter()