Work in progress on integration testing of udwf

Author: timsaucer

datafusion/ffi/src/tests/mod.rs

@@ -31,6 +31,8 @@ use crate::{catalog_provider::FFI_CatalogProvider, udtf::FFI_TableFunction};
 
 use crate::udaf::FFI_AggregateUDF;
 
+use crate::udwf::FFI_WindowUDF;
+
 use super::{table_provider::FFI_TableProvider, udf::FFI_ScalarUDF};
 use arrow::array::RecordBatch;
 use async_provider::create_async_table_provider;
@@ -40,7 +42,7 @@ use datafusion::{
 };
 use sync_provider::create_sync_table_provider;
 use udf_udaf_udwf::{
-    create_ffi_abs_func, create_ffi_random_func, create_ffi_stddev_func,
+    create_ffi_abs_func, create_ffi_random_func, create_ffi_rank_func, create_ffi_stddev_func,
     create_ffi_sum_func, create_ffi_table_func,
 };
 
@@ -76,6 +78,8 @@ pub struct ForeignLibraryModule {
     /// Createa  grouping UDAF using stddev
     pub create_stddev_udaf: extern "C" fn() -> FFI_AggregateUDF,
 
+    pub create_rank_udwf: extern "C" fn() -> FFI_WindowUDF,
+
     pub version: extern "C" fn() -> u64,
 }
 
@@ -125,6 +129,7 @@ pub fn get_foreign_library_module() -> ForeignLibraryModuleRef {
         create_table_function: create_ffi_table_func,
         create_sum_udaf: create_ffi_sum_func,
         create_stddev_udaf: create_ffi_stddev_func,
+        create_rank_udwf: create_ffi_rank_func,
         version: super::version,
     }
     .leak_into_prefix()

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -15,13 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::{udaf::FFI_AggregateUDF, udf::FFI_ScalarUDF, udtf::FFI_TableFunction};
+use crate::{udaf::FFI_AggregateUDF, udf::FFI_ScalarUDF, udtf::FFI_TableFunction, udwf::FFI_WindowUDF};
 use datafusion::{
     catalog::TableFunctionImpl,
     functions::math::{abs::AbsFunc, random::RandomFunc},
     functions_aggregate::{stddev::Stddev, sum::Sum},
     functions_table::generate_series::RangeFunc,
-    logical_expr::{AggregateUDF, ScalarUDF},
+    functions_window::rank::Rank,
+    logical_expr::{AggregateUDF, ScalarUDF, WindowUDF},
 };
 
 use std::sync::Arc;
@@ -55,3 +56,9 @@ pub(crate) extern "C" fn create_ffi_stddev_func() -> FFI_AggregateUDF {
 
     udaf.into()
 }
+
+pub(crate) extern "C" fn create_ffi_rank_func() -> FFI_WindowUDF {
+    let udwf: Arc<WindowUDF> = Arc::new(Rank::new("rank_demo".to_string(), datafusion::functions_window::rank::RankType::Basic).into());
+
+    udwf.into()
+}

datafusion/ffi/src/udwf/mod.rs

@@ -87,8 +87,6 @@ pub struct FFI_WindowUDF {
         arg_types: RVec<WrappedSchema>,
     ) -> RResult<RVec<WrappedSchema>, RString>,
 
-    pub schema: unsafe extern "C" fn(udwf: &Self) -> WrappedSchema,
-
     pub sort_options: ROption<FFI_SortOptions>,
 
     /// Used to create a clone on the provider of the udf. This should
@@ -108,19 +106,13 @@ unsafe impl Sync for FFI_WindowUDF {}
 
 pub struct WindowUDFPrivateData {
     pub udf: Arc<WindowUDF>,
-    pub schema: SchemaRef,
 }
 
 impl FFI_WindowUDF {
     unsafe fn inner(&self) -> &Arc<WindowUDF> {
         let private_data = self.private_data as *const WindowUDFPrivateData;
         &(*private_data).udf
     }
-
-    unsafe fn inner_schema(&self) -> SchemaRef {
-        let private_data = self.private_data as *const WindowUDFPrivateData;
-        Arc::clone(&(*private_data).schema)
-    }
 }
 
 unsafe extern "C" fn partition_evaluator_fn_wrapper(
@@ -167,12 +159,6 @@ unsafe extern "C" fn coerce_types_fn_wrapper(
     rresult!(vec_datatype_to_rvec_wrapped(&return_types))
 }
 
-unsafe extern "C" fn schema_fn_wrapper(udwf: &FFI_WindowUDF) -> WrappedSchema {
-    let schema = udwf.inner_schema();
-
-    schema.into()
-}
-
 unsafe extern "C" fn release_fn_wrapper(udwf: &mut FFI_WindowUDF) {
     let private_data = Box::from_raw(udwf.private_data as *mut WindowUDFPrivateData);
     drop(private_data);
@@ -187,7 +173,6 @@ unsafe extern "C" fn clone_fn_wrapper(udwf: &FFI_WindowUDF) -> FFI_WindowUDF {
     // });
     let private_data = Box::new(WindowUDFPrivateData {
         udf: Arc::clone(udwf.inner()),
-        schema: udwf.inner_schema(),
     });
 
     FFI_WindowUDF {
@@ -196,7 +181,6 @@ unsafe extern "C" fn clone_fn_wrapper(udwf: &FFI_WindowUDF) -> FFI_WindowUDF {
         volatility: udwf.volatility.clone(),
         partition_evaluator: partition_evaluator_fn_wrapper,
         sort_options: udwf.sort_options.clone(),
-        schema: schema_fn_wrapper,
         coerce_types: coerce_types_fn_wrapper,
         field: field_fn_wrapper,
         clone: clone_fn_wrapper,
@@ -211,22 +195,21 @@ impl Clone for FFI_WindowUDF {
     }
 }
 
-impl FFI_WindowUDF {
-    pub fn new(udf: Arc<WindowUDF>, schema: SchemaRef) -> Self {
+impl From<Arc<WindowUDF>> for FFI_WindowUDF {
+    fn from(udf: Arc<WindowUDF>) -> Self {
         let name = udf.name().into();
         let aliases = udf.aliases().iter().map(|a| a.to_owned().into()).collect();
         let volatility = udf.signature().volatility.into();
         let sort_options = udf.sort_options().map(|v| (&v).into()).into();
 
-        let private_data = Box::new(WindowUDFPrivateData { udf, schema });
+        let private_data = Box::new(WindowUDFPrivateData { udf });
 
         Self {
             name,
             aliases,
             volatility,
             partition_evaluator: partition_evaluator_fn_wrapper,
             sort_options,
-            schema: schema_fn_wrapper,
             coerce_types: coerce_types_fn_wrapper,
             field: field_fn_wrapper,
             clone: clone_fn_wrapper,
@@ -307,8 +290,7 @@ impl WindowUDFImpl for ForeignWindowUDF {
         args: datafusion::logical_expr::function::PartitionEvaluatorArgs,
     ) -> Result<Box<dyn PartitionEvaluator>> {
         let evaluator = unsafe {
-            let schema = (self.udf.schema)(&self.udf);
-            let args = FFI_PartitionEvaluatorArgs::new(args, schema.into())?;
+            let args = FFI_PartitionEvaluatorArgs::try_from(args)?;
             (self.udf.partition_evaluator)(&self.udf, args)
         };
 

datafusion/ffi/src/udwf/partition_evaluator_args.rs

@@ -15,18 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 use abi_stable::{std_types::RVec, StableAbi};
 use arrow::{
-    datatypes::{DataType, SchemaRef},
+    datatypes::{DataType, Field, Schema, SchemaRef},
     error::ArrowError,
     ffi::FFI_ArrowSchema,
 };
 use datafusion::{
     error::{DataFusionError, Result},
     logical_expr::function::PartitionEvaluatorArgs,
-    physical_plan::PhysicalExpr,
+    physical_plan::{expressions::Column, PhysicalExpr},
     prelude::SessionContext,
 };
 use datafusion_proto::{
@@ -51,11 +51,40 @@ pub struct FFI_PartitionEvaluatorArgs {
     schema: WrappedSchema,
 }
 
-impl FFI_PartitionEvaluatorArgs {
-    pub fn new(
+
+impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
+
+    type Error = DataFusionError;
+    fn try_from(
         args: PartitionEvaluatorArgs,
-        schema: SchemaRef,
     ) -> Result<Self, DataFusionError> {
+        // This is a bit of a hack. Since PartitionEvaluatorArgs does not carry a schema
+        // around, and instead passes the data types directly we are unable to decode the
+        // protobuf PhysicalExpr correctly. In evaluating the code the only place these
+        // appear to be really used are the Column data types. So here we will find all
+        // of the required columns and create a schema that has empty fields except for
+        // the ones we require. Ideally we would enhance PartitionEvaluatorArgs to just
+        // pass along the schema, but that is a larger breaking change.
+        let required_columns: HashMap<usize, (&str, &DataType)> = args.input_exprs().iter().zip(args.input_types())
+            .filter_map(|(expr, data_type)| {
+                expr.as_any().downcast_ref::<Column>().map(|column| (column.index(), (column.name(), data_type)))
+            })
+            .collect();
+
+        let max_column = required_columns.keys().max().unwrap_or(&0).to_owned();
+        let fields: Vec<_> = (0..max_column).into_iter()
+            .map(|idx| {
+                match required_columns.get(&idx) {
+                    Some((name, data_type)) => {
+                        Field::new(*name, (*data_type).clone(), true)
+                    }
+                    None => {
+                        Field::new(format!("ffi_partition_evaluator_col_{idx}"), DataType::Null, true)
+                    }
+                }
+            }).collect();
+        let schema = Arc::new(Schema::new(fields));
+
         let codec = DefaultPhysicalExtensionCodec {};
         let input_exprs = serialize_physical_exprs(args.input_exprs(), &codec)?
             .into_iter()

datafusion/ffi/tests/ffi_integration.rs

@@ -20,11 +20,15 @@
 #[cfg(feature = "integration-tests")]
 mod tests {
     use datafusion::error::{DataFusionError, Result};
-    use datafusion::prelude::SessionContext;
+    use datafusion::prelude::{col, SessionContext};
     use datafusion_ffi::catalog_provider::ForeignCatalogProvider;
     use datafusion_ffi::table_provider::ForeignTableProvider;
     use datafusion_ffi::tests::create_record_batch;
     use datafusion_ffi::tests::utils::get_module;
+    use datafusion::logical_expr::{ScalarUDF, WindowUDF};
+    use datafusion_ffi::udf::ForeignScalarUDF;
+    use datafusion_ffi::udwf::ForeignWindowUDF;
+    use std::path::Path;
     use std::sync::Arc;
 
     /// It is important that this test is in the `tests` directory and not in the
@@ -95,4 +99,41 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_rank_udwf() -> Result<()> {
+        let module = get_module()?;
+
+        let ffi_rank_func =
+            module
+                .create_rank_udwf()
+                .ok_or(DataFusionError::NotImplemented(
+                    "External table provider failed to implement create_scalar_udf"
+                        .to_string(),
+                ))?();
+        let foreign_rank_func: ForeignWindowUDF = (&ffi_rank_func).try_into()?;
+
+        let udwf: WindowUDF = foreign_rank_func.into();
+
+        let ctx = SessionContext::default();
+        let df = ctx.read_batch(create_record_batch(-5, 5))?;
+
+        let df = df
+            .with_column("rank_a", udwf.call(vec![]))?;
+
+        let result = df.collect().await?;
+
+        let expected = record_batch!(
+            ("a", Int32, vec![-5, -4, -3, -2, -1]),
+            ("b", Float64, vec![-5., -4., -3., -2., -1.]),
+            ("abs_a", Int32, vec![5, 4, 3, 2, 1]),
+            ("abs_b", Float64, vec![5., 4., 3., 2., 1.])
+        )?;
+
+        assert!(result.len() == 1);
+        assert!(result[0] == expected);
+
+        Ok(())
+
+    }
 }