Allow type coersion of zero input arrays to nullary (#15487)

* ScalarUDF via FFI would break with nullary type inputs. We could not coerce empty vector of arguments to nullary due to check. This also exposes return_type_from_args instead of just return_type that was causing some UDFs to fail. Added unit tests and moved around FFI test modules a little.

* Add license text

* Correct error in documentation

* Error message changed in test due to updated scalar coercion

* Perform check of user defined types when looking for empty argument types

* Updated error messages during unit test

Author: timsaucer

datafusion/expr/src/type_coercion/functions.rs

@@ -49,7 +49,7 @@ pub fn data_types_with_scalar_udf(
     let signature = func.signature();
     let type_signature = &signature.type_signature;
 
-    if current_types.is_empty() {
+    if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
         if type_signature.supports_zero_argument() {
             return Ok(vec![]);
         } else if type_signature.used_to_support_zero_arguments() {
@@ -87,7 +87,7 @@ pub fn data_types_with_aggregate_udf(
     let signature = func.signature();
     let type_signature = &signature.type_signature;
 
-    if current_types.is_empty() {
+    if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
         if type_signature.supports_zero_argument() {
             return Ok(vec![]);
         } else if type_signature.used_to_support_zero_arguments() {
@@ -124,7 +124,7 @@ pub fn data_types_with_window_udf(
     let signature = func.signature();
     let type_signature = &signature.type_signature;
 
-    if current_types.is_empty() {
+    if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
         if type_signature.supports_zero_argument() {
             return Ok(vec![]);
         } else if type_signature.used_to_support_zero_arguments() {
@@ -161,7 +161,7 @@ pub fn data_types(
 ) -> Result<Vec<DataType>> {
     let type_signature = &signature.type_signature;
 
-    if current_types.is_empty() {
+    if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
         if type_signature.supports_zero_argument() {
             return Ok(vec![]);
         } else if type_signature.used_to_support_zero_arguments() {

datafusion/ffi/src/tests/mod.rs

@@ -37,12 +37,13 @@ use datafusion::{
     common::record_batch,
 };
 use sync_provider::create_sync_table_provider;
-use udf_udaf_udwf::create_ffi_abs_func;
+use udf_udaf_udwf::{create_ffi_abs_func, create_ffi_random_func};
 
 mod async_provider;
 pub mod catalog;
 mod sync_provider;
 mod udf_udaf_udwf;
+pub mod utils;
 
 #[repr(C)]
 #[derive(StableAbi)]
@@ -60,6 +61,8 @@ pub struct ForeignLibraryModule {
     /// Create a scalar UDF
     pub create_scalar_udf: extern "C" fn() -> FFI_ScalarUDF,
 
+    pub create_nullary_udf: extern "C" fn() -> FFI_ScalarUDF,
+
     pub version: extern "C" fn() -> u64,
 }
 
@@ -105,6 +108,7 @@ pub fn get_foreign_library_module() -> ForeignLibraryModuleRef {
         create_catalog: create_catalog_provider,
         create_table: construct_table_provider,
         create_scalar_udf: create_ffi_abs_func,
+        create_nullary_udf: create_ffi_random_func,
         version: super::version,
     }
     .leak_into_prefix()

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -16,7 +16,10 @@
 // under the License.
 
 use crate::udf::FFI_ScalarUDF;
-use datafusion::{functions::math::abs::AbsFunc, logical_expr::ScalarUDF};
+use datafusion::{
+    functions::math::{abs::AbsFunc, random::RandomFunc},
+    logical_expr::ScalarUDF,
+};
 
 use std::sync::Arc;
 
@@ -25,3 +28,9 @@ pub(crate) extern "C" fn create_ffi_abs_func() -> FFI_ScalarUDF {
 
     udf.into()
 }
+
+pub(crate) extern "C" fn create_ffi_random_func() -> FFI_ScalarUDF {
+    let udf: Arc<ScalarUDF> = Arc::new(RandomFunc::new().into());
+
+    udf.into()
+}

datafusion/ffi/src/tests/utils.rs

@@ -0,0 +1,87 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::tests::ForeignLibraryModuleRef;
+use abi_stable::library::RootModule;
+use datafusion::error::{DataFusionError, Result};
+use std::path::Path;
+
+/// Compute the path to the library. It would be preferable to simply use
+/// abi_stable::library::development_utils::compute_library_path however
+/// our current CI pipeline has a `ci` profile that we need to use to
+/// find the library.
+pub fn compute_library_path<M: RootModule>(
+    target_path: &Path,
+) -> std::io::Result<std::path::PathBuf> {
+    let debug_dir = target_path.join("debug");
+    let release_dir = target_path.join("release");
+    let ci_dir = target_path.join("ci");
+
+    let debug_path = M::get_library_path(&debug_dir.join("deps"));
+    let release_path = M::get_library_path(&release_dir.join("deps"));
+    let ci_path = M::get_library_path(&ci_dir.join("deps"));
+
+    let all_paths = vec![
+        (debug_dir.clone(), debug_path),
+        (release_dir, release_path),
+        (ci_dir, ci_path),
+    ];
+
+    let best_path = all_paths
+        .into_iter()
+        .filter(|(_, path)| path.exists())
+        .filter_map(|(dir, path)| path.metadata().map(|m| (dir, m)).ok())
+        .filter_map(|(dir, meta)| meta.modified().map(|m| (dir, m)).ok())
+        .max_by_key(|(_, date)| *date)
+        .map(|(dir, _)| dir)
+        .unwrap_or(debug_dir);
+
+    Ok(best_path)
+}
+
+pub fn get_module() -> Result<ForeignLibraryModuleRef> {
+    let expected_version = crate::version();
+
+    let crate_root = Path::new(env!("CARGO_MANIFEST_DIR"));
+    let target_dir = crate_root
+        .parent()
+        .expect("Failed to find crate parent")
+        .parent()
+        .expect("Failed to find workspace root")
+        .join("target");
+
+    // Find the location of the library. This is specific to the build environment,
+    // so you will need to change the approach here based on your use case.
+    // let target: &std::path::Path = "../../../../target/".as_ref();
+    let library_path =
+        compute_library_path::<ForeignLibraryModuleRef>(target_dir.as_path())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?
+            .join("deps");
+
+    // Load the module
+    let module = ForeignLibraryModuleRef::load_from_directory(&library_path)
+        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+    assert_eq!(
+        module
+            .version()
+            .expect("Unable to call version on FFI module")(),
+        expected_version
+    );
+
+    Ok(module)
+}

datafusion/ffi/src/udf.rs renamed to datafusion/ffi/src/udf/mod.rs

@@ -29,14 +29,20 @@ use arrow::{
 };
 use datafusion::{
     error::DataFusionError,
-    logical_expr::type_coercion::functions::data_types_with_scalar_udf,
+    logical_expr::{
+        type_coercion::functions::data_types_with_scalar_udf, ReturnInfo, ReturnTypeArgs,
+    },
 };
 use datafusion::{
     error::Result,
     logical_expr::{
         ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
     },
 };
+use return_info::FFI_ReturnInfo;
+use return_type_args::{
+    FFI_ReturnTypeArgs, ForeignReturnTypeArgs, ForeignReturnTypeArgsOwned,
+};
 
 use crate::{
     arrow_wrappers::{WrappedArray, WrappedSchema},
@@ -45,6 +51,9 @@ use crate::{
     volatility::FFI_Volatility,
 };
 
+pub mod return_info;
+pub mod return_type_args;
+
 /// A stable struct for sharing a [`ScalarUDF`] across FFI boundaries.
 #[repr(C)]
 #[derive(Debug, StableAbi)]
@@ -66,6 +75,14 @@ pub struct FFI_ScalarUDF {
         arg_types: RVec<WrappedSchema>,
     ) -> RResult<WrappedSchema, RString>,
 
+    /// Determines the return info of the underlying [`ScalarUDF`]. Either this
+    /// or return_type may be implemented on a UDF.
+    pub return_type_from_args: unsafe extern "C" fn(
+        udf: &Self,
+        args: FFI_ReturnTypeArgs,
+    )
+        -> RResult<FFI_ReturnInfo, RString>,
+
     /// Execute the underlying [`ScalarUDF`] and return the result as a `FFI_ArrowArray`
     /// within an AbiStable wrapper.
     pub invoke_with_args: unsafe extern "C" fn(
@@ -123,6 +140,23 @@ unsafe extern "C" fn return_type_fn_wrapper(
     rresult!(return_type)
 }
 
+unsafe extern "C" fn return_type_from_args_fn_wrapper(
+    udf: &FFI_ScalarUDF,
+    args: FFI_ReturnTypeArgs,
+) -> RResult<FFI_ReturnInfo, RString> {
+    let private_data = udf.private_data as *const ScalarUDFPrivateData;
+    let udf = &(*private_data).udf;
+
+    let args: ForeignReturnTypeArgsOwned = rresult_return!((&args).try_into());
+    let args_ref: ForeignReturnTypeArgs = (&args).into();
+
+    let return_type = udf
+        .return_type_from_args((&args_ref).into())
+        .and_then(FFI_ReturnInfo::try_from);
+
+    rresult!(return_type)
+}
+
 unsafe extern "C" fn coerce_types_fn_wrapper(
     udf: &FFI_ScalarUDF,
     arg_types: RVec<WrappedSchema>,
@@ -209,6 +243,7 @@ impl From<Arc<ScalarUDF>> for FFI_ScalarUDF {
             short_circuits,
             invoke_with_args: invoke_with_args_fn_wrapper,
             return_type: return_type_fn_wrapper,
+            return_type_from_args: return_type_from_args_fn_wrapper,
             coerce_types: coerce_types_fn_wrapper,
             clone: clone_fn_wrapper,
             release: release_fn_wrapper,
@@ -281,6 +316,16 @@ impl ScalarUDFImpl for ForeignScalarUDF {
         result.and_then(|r| (&r.0).try_into().map_err(DataFusionError::from))
     }
 
+    fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result<ReturnInfo> {
+        let args: FFI_ReturnTypeArgs = args.try_into()?;
+
+        let result = unsafe { (self.udf.return_type_from_args)(&self.udf, args) };
+
+        let result = df_result!(result);
+
+        result.and_then(|r| r.try_into())
+    }
+
     fn invoke_with_args(&self, invoke_args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         let ScalarFunctionArgs {
             args,

datafusion/ffi/src/udf/return_info.rs

@@ -0,0 +1,53 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use abi_stable::StableAbi;
+use arrow::{datatypes::DataType, ffi::FFI_ArrowSchema};
+use datafusion::{error::DataFusionError, logical_expr::ReturnInfo};
+
+use crate::arrow_wrappers::WrappedSchema;
+
+/// A stable struct for sharing a [`ReturnInfo`] across FFI boundaries.
+#[repr(C)]
+#[derive(Debug, StableAbi)]
+#[allow(non_camel_case_types)]
+pub struct FFI_ReturnInfo {
+    return_type: WrappedSchema,
+    nullable: bool,
+}
+
+impl TryFrom<ReturnInfo> for FFI_ReturnInfo {
+    type Error = DataFusionError;
+
+    fn try_from(value: ReturnInfo) -> Result<Self, Self::Error> {
+        let return_type = WrappedSchema(FFI_ArrowSchema::try_from(value.return_type())?);
+        Ok(Self {
+            return_type,
+            nullable: value.nullable(),
+        })
+    }
+}
+
+impl TryFrom<FFI_ReturnInfo> for ReturnInfo {
+    type Error = DataFusionError;
+
+    fn try_from(value: FFI_ReturnInfo) -> Result<Self, Self::Error> {
+        let return_type = DataType::try_from(&value.return_type.0)?;
+
+        Ok(ReturnInfo::new(return_type, value.nullable))
+    }
+}