rebase_main

Author: coderfender

datafusion/ffi/src/udtf.rs

@@ -19,23 +19,24 @@ use std::any::Any;
 use std::ffi::c_void;
 use std::sync::Arc;
 
-use datafusion_catalog::{TableFunctionImpl, TableProvider};
+use datafusion_catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
+use datafusion_common::DataFusionError;
 use datafusion_common::error::Result;
 use datafusion_execution::TaskContext;
-use datafusion_expr::Expr;
 use datafusion_proto::logical_plan::from_proto::parse_exprs;
 use datafusion_proto::logical_plan::to_proto::serialize_exprs;
 use datafusion_proto::logical_plan::{
     DefaultLogicalExtensionCodec, LogicalExtensionCodec,
 };
 use datafusion_proto::protobuf::LogicalExprList;
+use datafusion_session::Session;
 use prost::Message;
-
 use stabby::vec::Vec as SVec;
 use tokio::runtime::Handle;
 
 use crate::execution::FFI_TaskContextProvider;
 use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::session::{FFI_SessionRef, ForeignSession};
 use crate::table_provider::FFI_TableProvider;
 use crate::util::FFIResult;
 use crate::{df_result, sresult_return};
@@ -44,11 +45,22 @@ use crate::{df_result, sresult_return};
 #[repr(C)]
 #[derive(Debug)]
 pub struct FFI_TableFunction {
-    /// Equivalent to the `call` function of the TableFunctionImpl.
+    /// Equivalent to the [`TableFunctionImpl::call`].
     /// The arguments are Expr passed as protobuf encoded bytes.
+    #[deprecated(
+        since = "53.0.0",
+        note = "See TableFunctionImpl::call deprecation note"
+    )]
     pub call:
         unsafe extern "C" fn(udtf: &Self, args: SVec<u8>) -> FFIResult<FFI_TableProvider>,
 
+    /// Equivalent to the [`TableFunctionImpl::call_with_args`].
+    call_with_args: unsafe extern "C" fn(
+        udtf: &Self,
+        args: SVec<u8>,
+        session: FFI_SessionRef,
+    ) -> FFIResult<FFI_TableProvider>,
+
     pub logical_codec: FFI_LogicalExtensionCodec,
 
     /// Used to create a clone on the provider of the udtf. This should
@@ -117,6 +129,47 @@ unsafe extern "C" fn call_fn_wrapper(
     ))
 }
 
+unsafe extern "C" fn call_with_args_wrapper(
+    udtf: &FFI_TableFunction,
+    args: SVec<u8>,
+    session: FFI_SessionRef,
+) -> FFIResult<FFI_TableProvider> {
+    let runtime = udtf.runtime();
+    let udtf_inner = udtf.inner();
+
+    let ctx: Arc<TaskContext> =
+        sresult_return!((&udtf.logical_codec.task_ctx_provider).try_into());
+    let codec: Arc<dyn LogicalExtensionCodec> = (&udtf.logical_codec).into();
+
+    let proto_filters = sresult_return!(LogicalExprList::decode(args.as_ref()));
+
+    let args = sresult_return!(parse_exprs(
+        proto_filters.expr.iter(),
+        ctx.as_ref(),
+        codec.as_ref()
+    ));
+
+    let mut foreign_session = None;
+    let session = sresult_return!(
+        session
+            .as_local()
+            .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+            .unwrap_or_else(|| {
+                foreign_session = Some(ForeignSession::try_from(&session)?);
+                Ok(foreign_session.as_ref().unwrap())
+            })
+    );
+    let table_provider = sresult_return!(
+        udtf_inner.call_with_args(TableFunctionArgs::new(&args, session))
+    );
+    FFIResult::Ok(FFI_TableProvider::new_with_ffi_codec(
+        table_provider,
+        false,
+        runtime,
+        udtf.logical_codec.clone(),
+    ))
+}
+
 unsafe extern "C" fn release_fn_wrapper(udtf: &mut FFI_TableFunction) {
     unsafe {
         debug_assert!(!udtf.private_data.is_null());
@@ -177,7 +230,9 @@ impl FFI_TableFunction {
         let private_data = Box::new(TableFunctionPrivateData { udtf, runtime });
 
         Self {
+            #[expect(deprecated)]
             call: call_fn_wrapper,
+            call_with_args: call_with_args_wrapper,
             logical_codec,
             clone: clone_fn_wrapper,
             release: release_fn_wrapper,
@@ -216,13 +271,35 @@ impl From<FFI_TableFunction> for Arc<dyn TableFunctionImpl> {
 }
 
 impl TableFunctionImpl for ForeignTableFunction {
-    fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
+        let session = FFI_SessionRef::new(
+            args.session(),
+            self.0.runtime(),
+            self.0.logical_codec.clone(),
+        );
+        let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
+        let expr_list = LogicalExprList {
+            expr: serialize_exprs(args.exprs(), codec.as_ref())?,
+        };
+        let filters_serialized = expr_list.encode_to_vec().into_iter().collect();
+
+        let table_provider =
+            unsafe { (self.0.call_with_args)(&self.0, filters_serialized, session) };
+
+        let table_provider = df_result!(table_provider)?;
+        let table_provider: Arc<dyn TableProvider> = (&table_provider).into();
+
+        Ok(table_provider)
+    }
+
+    fn call(&self, args: &[datafusion_expr::Expr]) -> Result<Arc<dyn TableProvider>> {
         let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
         let expr_list = LogicalExprList {
             expr: serialize_exprs(args, codec.as_ref())?,
         };
         let filters_serialized = expr_list.encode_to_vec().into_iter().collect();
 
+        #[expect(deprecated)]
         let table_provider = unsafe { (self.0.call)(&self.0, filters_serialized) };
 
         let table_provider = df_result!(table_provider)?;
@@ -243,16 +320,22 @@ mod tests {
     use datafusion::logical_expr::ptr_eq::arc_ptr_eq;
     use datafusion::prelude::{SessionContext, lit};
     use datafusion::scalar::ScalarValue;
+    use datafusion_catalog::TableFunctionArgs;
     use datafusion_execution::TaskContextProvider;
+    use datafusion_expr::Expr;
 
     use super::*;
 
     #[derive(Debug)]
     struct TestUDTF {}
 
     impl TableFunctionImpl for TestUDTF {
-        fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+        fn call_with_args(
+            &self,
+            args: TableFunctionArgs,
+        ) -> Result<Arc<dyn TableProvider>> {
             let args = args
+                .exprs()
                 .iter()
                 .map(|arg| {
                     if let Expr::Literal(scalar, _) = arg {
@@ -342,7 +425,10 @@ mod tests {
 
         let foreign_udf: Arc<dyn TableFunctionImpl> = local_udtf.into();
 
-        let table = foreign_udf.call(&[lit(6_u64), lit("one"), lit(2.0), lit(3_u64)])?;
+        let table = foreign_udf.call_with_args(TableFunctionArgs::new(
+            &[lit(6_u64), lit("one"), lit(2.0), lit(3_u64)],
+            &ctx.state(),
+        ))?;
 
         let _ = ctx.register_table("test-table", table)?;
 

datafusion/ffi/tests/ffi_catalog.rs

@@ -24,7 +24,6 @@ mod tests {
     use std::sync::Arc;
 
     use datafusion::catalog::{CatalogProvider, CatalogProviderList};
-    use datafusion_common::DataFusionError;
     use datafusion_ffi::tests::utils::get_module;
 
     #[tokio::test]