Provide session to the udtf call (review fixes)

askalt · askalt · commit ed923e592973 · 2026-02-09T14:04:34.000+03:00
This patch adds the passing of the current session to the UDTF call.
This helps implement session-dependent functions, for example, a function
that returns the list of registered tables.
diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md
@@ -207,13 +207,14 @@ cargo run --example dataframe -- dataframe
 
 #### Category: Single Process
 
-| Subcommand | File Path                                               | Description                                     |
-| ---------- | ------------------------------------------------------- | ----------------------------------------------- |
-| adv_udaf   | [`udf/advanced_udaf.rs`](examples/udf/advanced_udaf.rs) | Advanced User Defined Aggregate Function (UDAF) |
-| adv_udf    | [`udf/advanced_udf.rs`](examples/udf/advanced_udf.rs)   | Advanced User Defined Scalar Function (UDF)     |
-| adv_udwf   | [`udf/advanced_udwf.rs`](examples/udf/advanced_udwf.rs) | Advanced User Defined Window Function (UDWF)    |
-| async_udf  | [`udf/async_udf.rs`](examples/udf/async_udf.rs)         | Asynchronous User Defined Scalar Function       |
-| udaf       | [`udf/simple_udaf.rs`](examples/udf/simple_udaf.rs)     | Simple UDAF example                             |
-| udf        | [`udf/simple_udf.rs`](examples/udf/simple_udf.rs)       | Simple UDF example                              |
-| udtf       | [`udf/simple_udtf.rs`](examples/udf/simple_udtf.rs)     | Simple UDTF example                             |
-| udwf       | [`udf/simple_udwf.rs`](examples/udf/simple_udwf.rs)     | Simple UDWF example                             |
+| Subcommand      | File Path                                                   | Description                                     |
+| --------------- | ----------------------------------------------------------- | ----------------------------------------------- |
+| adv_udaf        | [`udf/advanced_udaf.rs`](examples/udf/advanced_udaf.rs)     | Advanced User Defined Aggregate Function (UDAF) |
+| adv_udf         | [`udf/advanced_udf.rs`](examples/udf/advanced_udf.rs)       | Advanced User Defined Scalar Function (UDF)     |
+| adv_udwf        | [`udf/advanced_udwf.rs`](examples/udf/advanced_udwf.rs)     | Advanced User Defined Window Function (UDWF)    |
+| async_udf       | [`udf/async_udf.rs`](examples/udf/async_udf.rs)             | Asynchronous User Defined Scalar Function       |
+| udaf            | [`udf/simple_udaf.rs`](examples/udf/simple_udaf.rs)         | Simple UDAF example                             |
+| udf             | [`udf/simple_udf.rs`](examples/udf/simple_udf.rs)           | Simple UDF example                              |
+| udtf            | [`udf/simple_udtf.rs`](examples/udf/simple_udtf.rs)         | Simple UDTF example                             |
+| udwf            | [`udf/simple_udwf.rs`](examples/udf/simple_udwf.rs)         | Simple UDWF example                             |
+| table_list_udtf | [`udf/table_list_udtf.rs`](examples/udf/table_list_udtf.rs) | Session-aware UDTF table list example           |
diff --git a/datafusion-examples/examples/udf/main.rs b/datafusion-examples/examples/udf/main.rs
@@ -50,6 +50,9 @@
 //!
 //! - `udwf`
 //!   (file: simple_udwf.rs, desc: Simple UDWF example)
+//!
+//! - `table_list_udtf`
+//!   (file: table_list_udtf.rs, desc: Session-aware UDTF table list example)
 
 mod advanced_udaf;
 mod advanced_udf;
diff --git a/datafusion-examples/examples/udf/table_list_udtf.rs b/datafusion-examples/examples/udf/table_list_udtf.rs
@@ -28,12 +28,12 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_common::{DataFusionError, plan_err};
-use futures::executor::block_on;
+use tokio::{runtime::Handle, task::block_in_place};
 
 const FUNCTION_NAME: &str = "table_list";
 
 // The example shows, how to create UDTF that depends on the session state.
-// There is `table_list` UDTF is defined which returns list of tables within session.
+// Defines a `table_list` UDTF that returns a list of tables within the provided session.
 
 pub async fn table_list_udtf() -> Result<()> {
     let ctx = SessionContext::new();
@@ -96,7 +96,10 @@ impl TableFunctionImpl for TableListUdtf {
                     continue;
                 };
                 for table_name in schema.table_names() {
-                    let Some(provider) = block_on(schema.table(&table_name))? else {
+                    let Some(provider) = block_in_place(|| {
+                        Handle::current().block_on(schema.table(&table_name))
+                    })?
+                    else {
                         continue;
                     };
                     catalogs.push(catalog_name.clone());
diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs
@@ -21,6 +21,7 @@ use std::sync::Arc;
 use abi_stable::StableAbi;
 use abi_stable::std_types::{RResult, RVec};
 use datafusion_catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
+use datafusion_common::DataFusionError;
 use datafusion_common::error::Result;
 use datafusion_execution::TaskContext;
 use datafusion_proto::logical_plan::from_proto::parse_exprs;
@@ -29,11 +30,13 @@ use datafusion_proto::logical_plan::{
     DefaultLogicalExtensionCodec, LogicalExtensionCodec,
 };
 use datafusion_proto::protobuf::LogicalExprList;
+use datafusion_session::Session;
 use prost::Message;
 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, rresult_return};
@@ -42,11 +45,18 @@ use crate::{df_result, rresult_return};
 #[repr(C)]
 #[derive(Debug, StableAbi)]
 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.
     pub call:
         unsafe extern "C" fn(udtf: &Self, args: RVec<u8>) -> FFIResult<FFI_TableProvider>,
 
+    /// Equivalent to the [`TableFunctionImpl::call_with_args`].
+    call_with_args: unsafe extern "C" fn(
+        udtf: &Self,
+        args: RVec<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
@@ -115,6 +125,48 @@ unsafe extern "C" fn call_fn_wrapper(
     ))
 }
 
+unsafe extern "C" fn call_with_args_wrapper(
+    udtf: &FFI_TableFunction,
+    args: RVec<u8>,
+    session: FFI_SessionRef,
+) -> FFIResult<FFI_TableProvider> {
+    let runtime = udtf.runtime();
+    let udtf_inner = udtf.inner();
+
+    let ctx: Arc<TaskContext> =
+        rresult_return!((&udtf.logical_codec.task_ctx_provider).try_into());
+    let codec: Arc<dyn LogicalExtensionCodec> = (&udtf.logical_codec).into();
+
+    let proto_filters = rresult_return!(LogicalExprList::decode(args.as_ref()));
+
+    let args = rresult_return!(parse_exprs(
+        proto_filters.expr.iter(),
+        ctx.as_ref(),
+        codec.as_ref()
+    ));
+
+    let mut foreign_session = None;
+    let session = rresult_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 = rresult_return!(udtf_inner.call_with_args(TableFunctionArgs {
+        args: &args,
+        session
+    }));
+    RResult::ROk(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());
@@ -170,6 +222,7 @@ impl FFI_TableFunction {
 
         Self {
             call: call_fn_wrapper,
+            call_with_args: call_with_args_wrapper,
             logical_codec,
             clone: clone_fn_wrapper,
             release: release_fn_wrapper,
@@ -209,12 +262,30 @@ impl From<FFI_TableFunction> for Arc<dyn TableFunctionImpl> {
 
 impl TableFunctionImpl for ForeignTableFunction {
     fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
+        let session =
+            FFI_SessionRef::new(args.session, None, self.0.logical_codec.clone());
         let codec: Arc<dyn LogicalExtensionCodec> = (&self.0.logical_codec).into();
         let expr_list = LogicalExprList {
             expr: serialize_exprs(args.args, codec.as_ref())?,
         };
         let filters_serialized = expr_list.encode_to_vec().into();
 
+        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();
+
         let table_provider = unsafe { (self.0.call)(&self.0, filters_serialized) };
 
         let table_provider = df_result!(table_provider)?;
@@ -340,7 +411,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 {
+            args: &[lit(6_u64), lit("one"), lit(2.0), lit(3_u64)],
+            session: &ctx.state(),
+        })?;
 
         let _ = ctx.register_table("test-table", table)?;
 
diff --git a/docs/source/library-user-guide/functions/adding-udfs.md b/docs/source/library-user-guide/functions/adding-udfs.md
@@ -1388,15 +1388,16 @@ in the CLI to read the metadata from a Parquet file.
 
 The simple UDTF used here takes a single `Int64` argument and returns a table with a single column with the value of the
 argument. To create a function in DataFusion, you need to implement the `TableFunctionImpl` trait. This trait has a
-single method, `call`, that takes a slice of `Expr`s and returns a `Result<Arc<dyn TableProvider>>`.
+single method, `call_with_args`, that takes a `TableFunctionArgs` struct and returns a `Result<Arc<dyn TableProvider>>`.
+Passed struct includes function arguments as a slice of `Expr`s.
 
-In the `call` method, you parse the input `Expr`s and return a `TableProvider`. You might also want to do some
+In the `call_with_args` method, you parse the input `Expr`s and return a `TableProvider`. You might also want to do some
 validation of the input `Expr`s, e.g. checking that the number of arguments is correct.
 
 ```rust
 use std::sync::Arc;
 use datafusion::common::{plan_err, ScalarValue, Result};
-use datafusion::catalog::{TableFunctionImpl, TableProvider};
+use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
 use datafusion::arrow::array::{ArrayRef, Int64Array};
 use datafusion::datasource::memory::MemTable;
 use arrow::record_batch::RecordBatch;
@@ -1438,7 +1439,7 @@ With the UDTF implemented, you can register it with the `SessionContext`:
 ```rust
 # use std::sync::Arc;
 # use datafusion::common::{plan_err, ScalarValue, Result};
-# use datafusion::catalog::{TableFunctionImpl, TableProvider};
+# use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
 # use datafusion::arrow::array::{ArrayRef, Int64Array};
 # use datafusion::datasource::memory::MemTable;
 # use arrow::record_batch::RecordBatch;
