Rebase due to UDF changes upstream

Author: timsaucer

datafusion/ffi/src/tests/mod.rs

@@ -42,8 +42,8 @@ use datafusion::{
 };
 use sync_provider::create_sync_table_provider;
 use udf_udaf_udwf::{
-    create_ffi_abs_func, create_ffi_random_func, create_ffi_rank_func, create_ffi_stddev_func,
-    create_ffi_sum_func, create_ffi_table_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,
 };
 
 mod async_provider;

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -15,7 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::{udaf::FFI_AggregateUDF, udf::FFI_ScalarUDF, udtf::FFI_TableFunction, udwf::FFI_WindowUDF};
+use crate::{
+    udaf::FFI_AggregateUDF, udf::FFI_ScalarUDF, udtf::FFI_TableFunction,
+    udwf::FFI_WindowUDF,
+};
 use datafusion::{
     catalog::TableFunctionImpl,
     functions::math::{abs::AbsFunc, random::RandomFunc},
@@ -58,7 +61,13 @@ pub(crate) extern "C" fn create_ffi_stddev_func() -> FFI_AggregateUDF {
 }
 
 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());
+    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

@@ -26,11 +26,12 @@ use arrow::{
     compute::SortOptions,
     datatypes::{DataType, SchemaRef},
 };
+use arrow_schema::{Field, FieldRef};
 use datafusion::{
     error::DataFusionError,
     logical_expr::{
-        function::WindowUDFFieldArgs,
-        type_coercion::functions::data_types_with_window_udf, PartitionEvaluator,
+        function::WindowUDFFieldArgs, type_coercion::functions::fields_with_window_udf,
+        PartitionEvaluator,
     },
 };
 use datafusion::{
@@ -45,6 +46,7 @@ mod partition_evaluator;
 mod partition_evaluator_args;
 mod range;
 
+use crate::util::{rvec_wrapped_to_vec_fieldref, vec_fieldref_to_rvec_wrapped};
 use crate::{
     arrow_wrappers::WrappedSchema,
     df_result, rresult, rresult_return,
@@ -130,16 +132,17 @@ unsafe extern "C" fn partition_evaluator_fn_wrapper(
 
 unsafe extern "C" fn field_fn_wrapper(
     udwf: &FFI_WindowUDF,
-    input_types: RVec<WrappedSchema>,
+    input_fields: RVec<WrappedSchema>,
     display_name: RString,
 ) -> RResult<WrappedSchema, RString> {
     let inner = udwf.inner();
 
-    let input_types = rresult_return!(rvec_wrapped_to_vec_datatype(&input_types));
+    let input_fields = rresult_return!(rvec_wrapped_to_vec_fieldref(&input_fields));
 
-    let field = rresult_return!(
-        inner.field(WindowUDFFieldArgs::new(&input_types, display_name.as_str()))
-    );
+    let field = rresult_return!(inner.field(WindowUDFFieldArgs::new(
+        &input_fields,
+        display_name.as_str()
+    )));
 
     let schema = Arc::new(Schema::new(vec![field]));
 
@@ -152,9 +155,17 @@ unsafe extern "C" fn coerce_types_fn_wrapper(
 ) -> RResult<RVec<WrappedSchema>, RString> {
     let inner = udwf.inner();
 
-    let arg_types = rresult_return!(rvec_wrapped_to_vec_datatype(&arg_types));
+    let arg_fields = rresult_return!(rvec_wrapped_to_vec_datatype(&arg_types))
+        .into_iter()
+        .map(|dt| Field::new("f", dt, false))
+        .map(Arc::new)
+        .collect::<Vec<_>>();
 
-    let return_types = rresult_return!(data_types_with_window_udf(&arg_types, inner));
+    let return_fields = rresult_return!(fields_with_window_udf(&arg_fields, inner));
+    let return_types = return_fields
+        .into_iter()
+        .map(|f| f.data_type().to_owned())
+        .collect::<Vec<_>>();
 
     rresult!(vec_datatype_to_rvec_wrapped(&return_types))
 }
@@ -300,9 +311,9 @@ impl WindowUDFImpl for ForeignWindowUDF {
         })
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<arrow::datatypes::Field> {
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
         unsafe {
-            let input_types = vec_datatype_to_rvec_wrapped(field_args.input_types())?;
+            let input_types = vec_fieldref_to_rvec_wrapped(field_args.input_fields())?;
             let schema = df_result!((self.udf.field)(
                 &self.udf,
                 input_types,
@@ -314,7 +325,7 @@ impl WindowUDFImpl for ForeignWindowUDF {
                 true => Err(DataFusionError::Execution(
                     "Unable to retrieve field in WindowUDF via FFI".to_string(),
                 )),
-                false => Ok(schema.field(0).to_owned()),
+                false => Ok(schema.field(0).to_owned().into()),
             }
         }
     }

datafusion/ffi/src/udwf/partition_evaluator.rs

@@ -17,6 +17,7 @@
 
 use std::{ffi::c_void, ops::Range};
 
+use crate::{arrow_wrappers::WrappedArray, df_result, rresult, rresult_return};
 use abi_stable::{
     std_types::{RResult, RString, RVec},
     StableAbi,
@@ -29,10 +30,11 @@ use datafusion::{
 };
 use prost::Message;
 
-use crate::{arrow_wrappers::WrappedArray, df_result, rresult, rresult_return};
-
 use super::range::FFI_Range;
 
+/// A stable struct for sharing [`PartitionEvaluator`] across FFI boundaries.
+/// For an explanation of each field, see the corresponding function
+/// defined in [`PartitionEvaluator`].
 #[repr(C)]
 #[derive(Debug, StableAbi)]
 #[allow(non_camel_case_types)]
@@ -108,7 +110,8 @@ unsafe extern "C" fn evaluate_all_fn_wrapper(
         .collect::<Result<Vec<ArrayRef>>>();
     let values_arrays = rresult_return!(values_arrays);
 
-    let return_array = (inner.evaluate_all(&values_arrays, num_rows))
+    let return_array = inner
+        .evaluate_all(&values_arrays, num_rows)
         .and_then(|array| WrappedArray::try_from(&array).map_err(DataFusionError::from));
 
     rresult!(return_array)
@@ -148,7 +151,8 @@ unsafe extern "C" fn evaluate_all_with_rank_fn_wrapper(
         .map(Range::from)
         .collect::<Vec<_>>();
 
-    let return_array = (inner.evaluate_all_with_rank(num_rows, &ranks_in_partition))
+    let return_array = inner
+        .evaluate_all_with_rank(num_rows, &ranks_in_partition)
         .and_then(|array| WrappedArray::try_from(&array).map_err(DataFusionError::from));
 
     rresult!(return_array)

datafusion/ffi/src/udwf/partition_evaluator_args.rs

@@ -17,12 +17,14 @@
 
 use std::{collections::HashMap, sync::Arc};
 
+use crate::arrow_wrappers::WrappedSchema;
 use abi_stable::{std_types::RVec, StableAbi};
 use arrow::{
     datatypes::{DataType, Field, Schema, SchemaRef},
     error::ArrowError,
     ffi::FFI_ArrowSchema,
 };
+use arrow_schema::FieldRef;
 use datafusion::{
     error::{DataFusionError, Result},
     logical_expr::function::PartitionEvaluatorArgs,
@@ -38,51 +40,52 @@ use datafusion_proto::{
 };
 use prost::Message;
 
-use crate::arrow_wrappers::WrappedSchema;
-
+/// A stable struct for sharing [`PartitionEvaluatorArgs`] across FFI boundaries.
+/// For an explanation of each field, see the corresponding function
+/// defined in [`PartitionEvaluatorArgs`].
 #[repr(C)]
 #[derive(Debug, StableAbi)]
 #[allow(non_camel_case_types)]
 pub struct FFI_PartitionEvaluatorArgs {
     input_exprs: RVec<RVec<u8>>,
-    input_types: RVec<WrappedSchema>,
+    input_fields: RVec<WrappedSchema>,
     is_reversed: bool,
     ignore_nulls: bool,
     schema: WrappedSchema,
 }
 
-
 impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
-
     type Error = DataFusionError;
-    fn try_from(
-        args: PartitionEvaluatorArgs,
-    ) -> Result<Self, DataFusionError> {
+    fn try_from(args: PartitionEvaluatorArgs) -> 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)))
+        let required_columns: HashMap<usize, (&str, &DataType)> = args
+            .input_exprs()
+            .iter()
+            .zip(args.input_fields())
+            .filter_map(|(expr, field)| {
+                expr.as_any()
+                    .downcast_ref::<Column>()
+                    .map(|column| (column.index(), (column.name(), field.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 fields: Vec<_> = (0..max_column)
+            .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 {};
@@ -91,8 +94,8 @@ impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
             .map(|expr_node| expr_node.encode_to_vec().into())
             .collect();
 
-        let input_types = args
-            .input_types()
+        let input_fields = args
+            .input_fields()
             .iter()
             .map(|input_type| FFI_ArrowSchema::try_from(input_type).map(WrappedSchema))
             .collect::<Result<Vec<_>, ArrowError>>()?
@@ -102,7 +105,7 @@ impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
 
         Ok(Self {
             input_exprs,
-            input_types,
+            input_fields,
             schema,
             is_reversed: args.is_reversed(),
             ignore_nulls: args.ignore_nulls(),
@@ -116,7 +119,7 @@ impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
 /// PartitionEvaluatorArgs can then reference.
 pub struct ForeignPartitionEvaluatorArgs {
     input_exprs: Vec<Arc<dyn PhysicalExpr>>,
-    input_types: Vec<DataType>,
+    input_fields: Vec<FieldRef>,
     is_reversed: bool,
     ignore_nulls: bool,
 }
@@ -142,14 +145,14 @@ impl TryFrom<FFI_PartitionEvaluatorArgs> for ForeignPartitionEvaluatorArgs {
             })
             .collect::<Result<Vec<_>>>()?;
 
-        let input_types = input_exprs
+        let input_fields = input_exprs
             .iter()
-            .map(|expr| expr.data_type(&schema))
+            .map(|expr| expr.return_field(&schema))
             .collect::<Result<Vec<_>>>()?;
 
         Ok(Self {
             input_exprs,
-            input_types,
+            input_fields,
             is_reversed: value.is_reversed,
             ignore_nulls: value.ignore_nulls,
         })
@@ -160,7 +163,7 @@ impl<'a> From<&'a ForeignPartitionEvaluatorArgs> for PartitionEvaluatorArgs<'a>
     fn from(value: &'a ForeignPartitionEvaluatorArgs) -> Self {
         PartitionEvaluatorArgs::new(
             &value.input_exprs,
-            &value.input_types,
+            &value.input_fields,
             value.is_reversed,
             value.ignore_nulls,
         )

datafusion/ffi/src/udwf/range.rs

@@ -19,6 +19,9 @@ use std::ops::Range;
 
 use abi_stable::StableAbi;
 
+/// A stable struct for sharing [`Range`] across FFI boundaries.
+/// For an explanation of each field, see the corresponding function
+/// defined in [`Range`].
 #[repr(C)]
 #[derive(Debug, StableAbi)]
 #[allow(non_camel_case_types)]

datafusion/ffi/tests/ffi_integration.rs

@@ -20,15 +20,11 @@
 #[cfg(feature = "integration-tests")]
 mod tests {
     use datafusion::error::{DataFusionError, Result};
-    use datafusion::prelude::{col, SessionContext};
+    use datafusion::prelude::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
@@ -99,41 +95,4 @@ 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(())
-
-    }
 }