Initial commit of UDWF via FFI

Author: timsaucer

datafusion/ffi/src/lib.rs

@@ -37,6 +37,7 @@ pub mod table_source;
 pub mod udaf;
 pub mod udf;
 pub mod udtf;
+pub mod udwf;
 pub mod util;
 pub mod volatility;
 

datafusion/ffi/src/udwf/mod.rs

@@ -0,0 +1,373 @@
+// 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 std::{ffi::c_void, sync::Arc};
+
+use abi_stable::{
+    std_types::{ROption, RResult, RString, RVec},
+    StableAbi,
+};
+use arrow::datatypes::Schema;
+use arrow::{
+    compute::SortOptions,
+    datatypes::{DataType, SchemaRef},
+};
+use datafusion::{
+    error::DataFusionError,
+    logical_expr::{
+        function::WindowUDFFieldArgs,
+        type_coercion::functions::data_types_with_window_udf, PartitionEvaluator,
+    },
+};
+use datafusion::{
+    error::Result,
+    logical_expr::{Signature, WindowUDF, WindowUDFImpl},
+};
+use partition_evaluator::{FFI_PartitionEvaluator, ForeignPartitionEvaluator};
+use partition_evaluator_args::{
+    FFI_PartitionEvaluatorArgs, ForeignPartitionEvaluatorArgs,
+};
+mod partition_evaluator;
+mod partition_evaluator_args;
+mod range;
+
+use crate::{
+    arrow_wrappers::WrappedSchema,
+    df_result, rresult, rresult_return,
+    util::{rvec_wrapped_to_vec_datatype, vec_datatype_to_rvec_wrapped},
+    volatility::FFI_Volatility,
+};
+
+/// A stable struct for sharing a [`WindowUDF`] across FFI boundaries.
+#[repr(C)]
+#[derive(Debug, StableAbi)]
+#[allow(non_camel_case_types)]
+pub struct FFI_WindowUDF {
+    /// FFI equivalent to the `name` of a [`WindowUDF`]
+    pub name: RString,
+
+    /// FFI equivalent to the `aliases` of a [`WindowUDF`]
+    pub aliases: RVec<RString>,
+
+    /// FFI equivalent to the `volatility` of a [`WindowUDF`]
+    pub volatility: FFI_Volatility,
+
+    pub partition_evaluator:
+        unsafe extern "C" fn(
+            udwf: &Self,
+            args: FFI_PartitionEvaluatorArgs,
+        ) -> RResult<FFI_PartitionEvaluator, RString>,
+
+    pub field: unsafe extern "C" fn(
+        udwf: &Self,
+        input_types: RVec<WrappedSchema>,
+        display_name: RString,
+    ) -> RResult<WrappedSchema, RString>,
+
+    /// Performs type coersion. To simply this interface, all UDFs are treated as having
+    /// user defined signatures, which will in turn call coerce_types to be called. This
+    /// call should be transparent to most users as the internal function performs the
+    /// appropriate calls on the underlying [`WindowUDF`]
+    pub coerce_types: unsafe extern "C" fn(
+        udf: &Self,
+        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
+    /// only need to be called by the receiver of the udf.
+    pub clone: unsafe extern "C" fn(udf: &Self) -> Self,
+
+    /// Release the memory of the private data when it is no longer being used.
+    pub release: unsafe extern "C" fn(udf: &mut Self),
+
+    /// Internal data. This is only to be accessed by the provider of the udf.
+    /// A [`ForeignWindowUDF`] should never attempt to access this data.
+    pub private_data: *mut c_void,
+}
+
+unsafe impl Send for FFI_WindowUDF {}
+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(
+    udwf: &FFI_WindowUDF,
+    args: FFI_PartitionEvaluatorArgs,
+) -> RResult<FFI_PartitionEvaluator, RString> {
+    let inner = udwf.inner();
+
+    let args = rresult_return!(ForeignPartitionEvaluatorArgs::try_from(args));
+
+    let evaluator = rresult_return!(inner.partition_evaluator_factory((&args).into()));
+
+    RResult::ROk(evaluator.into())
+}
+
+unsafe extern "C" fn field_fn_wrapper(
+    udwf: &FFI_WindowUDF,
+    input_types: 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 field = rresult_return!(
+        inner.field(WindowUDFFieldArgs::new(&input_types, display_name.as_str()))
+    );
+
+    let schema = Arc::new(Schema::new(vec![field]));
+
+    RResult::ROk(WrappedSchema::from(schema))
+}
+
+unsafe extern "C" fn coerce_types_fn_wrapper(
+    udwf: &FFI_WindowUDF,
+    arg_types: RVec<WrappedSchema>,
+) -> RResult<RVec<WrappedSchema>, RString> {
+    let inner = udwf.inner();
+
+    let arg_types = rresult_return!(rvec_wrapped_to_vec_datatype(&arg_types));
+
+    let return_types = rresult_return!(data_types_with_window_udf(&arg_types, inner));
+
+    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);
+}
+
+unsafe extern "C" fn clone_fn_wrapper(udwf: &FFI_WindowUDF) -> FFI_WindowUDF {
+    // let private_data = udf.private_data as *const WindowUDFPrivateData;
+    // let udf_data = &(*private_data);
+
+    // let private_data = Box::new(WindowUDFPrivateData {
+    //     udf: Arc::clone(&udf_data.udf),
+    // });
+    let private_data = Box::new(WindowUDFPrivateData {
+        udf: Arc::clone(udwf.inner()),
+        schema: udwf.inner_schema(),
+    });
+
+    FFI_WindowUDF {
+        name: udwf.name.clone(),
+        aliases: udwf.aliases.clone(),
+        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,
+        release: release_fn_wrapper,
+        private_data: Box::into_raw(private_data) as *mut c_void,
+    }
+}
+
+impl Clone for FFI_WindowUDF {
+    fn clone(&self) -> Self {
+        unsafe { (self.clone)(self) }
+    }
+}
+
+impl FFI_WindowUDF {
+    pub fn new(udf: Arc<WindowUDF>, schema: SchemaRef) -> 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 });
+
+        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,
+            release: release_fn_wrapper,
+            private_data: Box::into_raw(private_data) as *mut c_void,
+        }
+    }
+}
+
+impl Drop for FFI_WindowUDF {
+    fn drop(&mut self) {
+        unsafe { (self.release)(self) }
+    }
+}
+
+/// This struct is used to access an UDF provided by a foreign
+/// library across a FFI boundary.
+///
+/// The ForeignWindowUDF is to be used by the caller of the UDF, so it has
+/// no knowledge or access to the private data. All interaction with the UDF
+/// must occur through the functions defined in FFI_WindowUDF.
+#[derive(Debug)]
+pub struct ForeignWindowUDF {
+    name: String,
+    aliases: Vec<String>,
+    udf: FFI_WindowUDF,
+    signature: Signature,
+}
+
+unsafe impl Send for ForeignWindowUDF {}
+unsafe impl Sync for ForeignWindowUDF {}
+
+impl TryFrom<&FFI_WindowUDF> for ForeignWindowUDF {
+    type Error = DataFusionError;
+
+    fn try_from(udf: &FFI_WindowUDF) -> Result<Self, Self::Error> {
+        let name = udf.name.to_owned().into();
+        let signature = Signature::user_defined((&udf.volatility).into());
+
+        let aliases = udf.aliases.iter().map(|s| s.to_string()).collect();
+
+        Ok(Self {
+            name,
+            udf: udf.clone(),
+            aliases,
+            signature,
+        })
+    }
+}
+
+impl WindowUDFImpl for ForeignWindowUDF {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        unsafe {
+            let arg_types = vec_datatype_to_rvec_wrapped(arg_types)?;
+            let result_types = df_result!((self.udf.coerce_types)(&self.udf, arg_types))?;
+            Ok(rvec_wrapped_to_vec_datatype(&result_types)?)
+        }
+    }
+
+    fn partition_evaluator(
+        &self,
+        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())?;
+            (self.udf.partition_evaluator)(&self.udf, args)
+        };
+
+        df_result!(evaluator).map(|evaluator| {
+            Box::new(ForeignPartitionEvaluator::from(evaluator))
+                as Box<dyn PartitionEvaluator>
+        })
+    }
+
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<arrow::datatypes::Field> {
+        unsafe {
+            let input_types = vec_datatype_to_rvec_wrapped(field_args.input_types())?;
+            let schema = df_result!((self.udf.field)(
+                &self.udf,
+                input_types,
+                field_args.name().into()
+            ))?;
+            let schema: SchemaRef = schema.into();
+
+            match schema.fields().is_empty() {
+                true => Err(DataFusionError::Execution(
+                    "Unable to retrieve field in WindowUDF via FFI".to_string(),
+                )),
+                false => Ok(schema.field(0).to_owned()),
+            }
+        }
+    }
+
+    fn sort_options(&self) -> Option<SortOptions> {
+        let options: Option<&FFI_SortOptions> = self.udf.sort_options.as_ref().into();
+        options.map(|s| s.into())
+    }
+}
+
+#[repr(C)]
+#[derive(Debug, StableAbi, Clone)]
+#[allow(non_camel_case_types)]
+pub struct FFI_SortOptions {
+    pub descending: bool,
+    pub nulls_first: bool,
+}
+
+impl From<&SortOptions> for FFI_SortOptions {
+    fn from(value: &SortOptions) -> Self {
+        Self {
+            descending: value.descending,
+            nulls_first: value.nulls_first,
+        }
+    }
+}
+
+impl From<&FFI_SortOptions> for SortOptions {
+    fn from(value: &FFI_SortOptions) -> Self {
+        Self {
+            descending: value.descending,
+            nulls_first: value.nulls_first,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {}