Add integration test that uses config values for scalar udf

Author: timsaucer

datafusion/ffi/src/tests/mod.rs

@@ -84,6 +84,8 @@ pub struct ForeignLibraryModule {
 
     pub create_nullary_udf: extern "C" fn() -> FFI_ScalarUDF,
 
+    pub create_timezone_udf: extern "C" fn() -> FFI_ScalarUDF,
+
     pub create_table_function:
         extern "C" fn(FFI_LogicalExtensionCodec) -> FFI_TableFunction,
 
@@ -157,6 +159,7 @@ pub fn get_foreign_library_module() -> ForeignLibraryModuleRef {
         create_table_factory: construct_table_provider_factory,
         create_scalar_udf: create_ffi_abs_func,
         create_nullary_udf: create_ffi_random_func,
+        create_timezone_udf: udf_udaf_udwf::create_timezone_func,
         create_table_function: create_ffi_table_func,
         create_sum_udaf: create_ffi_sum_func,
         create_stddev_udaf: create_ffi_stddev_func,

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -20,9 +20,10 @@ use std::sync::Arc;
 
 use arrow_schema::DataType;
 use datafusion_catalog::TableFunctionImpl;
+use datafusion_common::ScalarValue;
 use datafusion_expr::{
     AggregateUDF, ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
-    WindowUDF,
+    Volatility, WindowUDF,
 };
 use datafusion_functions::math::abs::AbsFunc;
 use datafusion_functions::math::random::RandomFunc;
@@ -78,6 +79,47 @@ pub(crate) extern "C" fn create_ffi_random_func() -> FFI_ScalarUDF {
     udf.into()
 }
 
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct TimeZoneUDF {
+    signature: Signature,
+}
+
+impl ScalarUDFImpl for TimeZoneUDF {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+    fn name(&self) -> &str {
+        "TimeZoneUDF"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(
+        &self,
+        _arg_types: &[DataType],
+    ) -> datafusion_common::Result<DataType> {
+        Ok(DataType::Utf8)
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: ScalarFunctionArgs,
+    ) -> datafusion_common::Result<ColumnarValue> {
+        let tz = args.config_options.execution.time_zone.clone();
+        Ok(ColumnarValue::Scalar(ScalarValue::from(tz)))
+    }
+}
+
+pub(crate) extern "C" fn create_timezone_func() -> FFI_ScalarUDF {
+    let udf: Arc<ScalarUDF> = Arc::new(ScalarUDF::from(TimeZoneUDF {
+        signature: Signature::uniform(1, vec![DataType::Utf8], Volatility::Stable),
+    }));
+
+    udf.into()
+}
+
 pub(crate) extern "C" fn create_ffi_table_func(
     codec: FFI_LogicalExtensionCodec,
 ) -> FFI_TableFunction {

datafusion/ffi/tests/ffi_udf.rs

@@ -19,15 +19,17 @@
 /// when the feature integration-tests is built
 #[cfg(feature = "integration-tests")]
 mod tests {
-    use std::sync::Arc;
-
+    use arrow::array::{Array, AsArray};
     use arrow::datatypes::DataType;
     use datafusion::common::record_batch;
     use datafusion::error::{DataFusionError, Result};
     use datafusion::logical_expr::{ScalarUDF, ScalarUDFImpl};
     use datafusion::prelude::{SessionContext, col};
+    use datafusion_execution::config::SessionConfig;
+    use datafusion_expr::lit;
     use datafusion_ffi::tests::create_record_batch;
     use datafusion_ffi::tests::utils::get_module;
+    use std::sync::Arc;
 
     /// This test validates that we can load an external module and use a scalar
     /// udf defined in it via the foreign function interface. In this case we are
@@ -100,4 +102,46 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_config_on_scalar_udf() -> Result<()> {
+        let module = get_module()?;
+
+        let ffi_udf =
+            module
+                .create_timezone_udf()
+                .ok_or(DataFusionError::NotImplemented(
+                    "External module failed to implement create_timezone_udf".to_string(),
+                ))?();
+        let foreign_udf: Arc<dyn ScalarUDFImpl> = (&ffi_udf).into();
+
+        let udf = ScalarUDF::new_from_shared_impl(foreign_udf);
+
+        let ctx = SessionContext::default();
+
+        let df = ctx
+            .read_empty()?
+            .select(vec![udf.call(vec![lit("a")]).alias("a")])?;
+
+        let result = df.collect().await?;
+        assert!(result[0].column(0).as_string::<i32>().is_null(0));
+
+        let mut config = SessionConfig::new();
+        config.options_mut().execution.time_zone = Some("AEST".into());
+
+        let ctx = SessionContext::new_with_config(config);
+
+        let df = ctx
+            .read_empty()?
+            .select(vec![udf.call(vec![lit("a")]).alias("a")])?;
+
+        let result = df.collect().await?;
+
+        assert!(result.len() == 1);
+        assert!(!result[0].column(0).as_string::<i32>().is_null(0));
+        let result = result[0].column(0).as_string::<i32>().value(0);
+        assert_eq!(result, "AEST");
+
+        Ok(())
+    }
 }