Implement FFI task context and task context provider (#18918)

## Which issue does this PR close?

Addresses part of #18671 but
does not close it.

## Rationale for this change

We have an issue in the current FFI code in that we will need access to
the `TaskContext` to do function encoding and decoding to pass logical
expressions across the boundary. We use the proto crate for this
operations.

We cannot encode and decode functions that have been registered since
our current code creates a default session context. This causes problems
in crates where we are using UDFs as inputs to either aggregations,
window functions, or table filters.

With this change we keep a _weak_ reference to a new trait, a
`TaskContextProvider`. By keeping a weak trait we make sure we do not
create a circular dependency between function that is internally holding
on to the provider, which holds the task context, which holds the
function.

## What changes are included in this PR?

- Introduce the `TaskContextProvider` trait
- Implement FFI versions of `TaskContext` and `TaskContextProvider`.

This PR does _not_ use these structures in the current code. That is
coming as part of a later PR in an effort to keep the size of the PRs
small for effective code review.

## Are these changes tested?

Unit tests are added. Coverage report:

<img width="525" height="218" alt="Screenshot 2025-11-25 at 2 36 57 PM"
src="https://github.com/user-attachments/assets/4e1da78c-8058-41dc-bf87-e4ccb3d2b894"
/>

## Are there any user-facing changes?

No.

Author: timsaucer

Cargo.lock

@@ -1860,6 +1860,12 @@ impl FunctionRegistry for SessionContext {
     }
 }
 
+impl datafusion_execution::TaskContextProvider for SessionContext {
+    fn task_ctx(&self) -> Arc<TaskContext> {
+        SessionContext::task_ctx(self)
+    }
+}
+
 /// Create a new task context instance from SessionContext
 impl From<&SessionContext> for TaskContext {
     fn from(session: &SessionContext) -> Self {

datafusion/core/src/execution/context/mod.rs

@@ -1988,6 +1988,12 @@ impl FunctionRegistry for SessionState {
     }
 }
 
+impl datafusion_execution::TaskContextProvider for SessionState {
+    fn task_ctx(&self) -> Arc<TaskContext> {
+        SessionState::task_ctx(self)
+    }
+}
+
 impl OptimizerConfig for SessionState {
     fn query_execution_start_time(&self) -> DateTime<Utc> {
         self.execution_props.query_execution_start_time

datafusion/core/src/execution/session_state.rs

@@ -47,4 +47,4 @@ pub mod registry {
 pub use disk_manager::DiskManager;
 pub use registry::FunctionRegistry;
 pub use stream::{RecordBatchStream, SendableRecordBatchStream};
-pub use task::TaskContext;
+pub use task::{TaskContext, TaskContextProvider};

datafusion/execution/src/lib.rs

@@ -211,6 +211,11 @@ impl FunctionRegistry for TaskContext {
     }
 }
 
+/// Produce the [`TaskContext`].
+pub trait TaskContextProvider {
+    fn task_ctx(&self) -> Arc<TaskContext>;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion/execution/src/task.rs

@@ -48,6 +48,7 @@ async-ffi = { version = "0.5.0", features = ["abi_stable"] }
 async-trait = { workspace = true }
 datafusion = { workspace = true, default-features = false }
 datafusion-common = { workspace = true }
+datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-functions-aggregate-common = { workspace = true }
 datafusion-physical-expr = { workspace = true }

datafusion/ffi/Cargo.toml

@@ -170,6 +170,29 @@ your unit tests you should override this with
 `crate::mock_foreign_marker_id` to force your test to create the foreign
 variant of your struct.
 
+## Task Context Provider
+
+Many of the FFI structs in this crate contain a `FFI_TaskContextProvider`. The
+purpose of this struct is to _weakly_ hold a reference to a method to
+access the current `TaskContext`. The reason we need this accessor is because
+we use the `datafusion-proto` crate to serialize and deserialize data across
+the FFI boundary. In particular, we need to serialize and deserialize
+functions using a `TaskContext`, which implements `FunctionRegistry`.
+
+This becomes difficult because we may need to register multiple user defined
+functions, table or catalog providers, etc with a `Session`, and each of these
+will need the `TaskContext` to perform the processing. For this reason we
+cannot simply include the `TaskContext` at the time of registration because
+it would not have knowledge of anything registered afterward.
+
+The `FFI_TaskContextProvider` is built from a trait that provides a method
+to get the current `TaskContext`. `FFI_TaskContextProvider` only holds a
+`Weak` reference to the `TaskContextProvider`, because otherwise we could
+create a circular dependency at runtime. It is imperative that if you use
+these methods that your provider remains valid for the lifetime of the
+calls. The `FFI_TaskContextProvider` is implemented on `SessionContext`
+and it is easy to implement on any struct that implements `Session`.
+
 [apache datafusion]: https://datafusion.apache.org/
 [api docs]: http://docs.rs/datafusion-ffi/latest
 [rust abi]: https://doc.rust-lang.org/reference/abi.html

datafusion/ffi/README.md

@@ -0,0 +1,22 @@
+// 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.
+
+mod task_ctx;
+pub mod task_ctx_provider;
+
+pub use task_ctx::FFI_TaskContext;
+pub use task_ctx_provider::FFI_TaskContextProvider;

datafusion/ffi/src/execution/mod.rs

@@ -0,0 +1,270 @@
+// 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;
+use std::sync::Arc;
+
+use abi_stable::pmr::ROption;
+use abi_stable::std_types::{RHashMap, RString};
+use abi_stable::StableAbi;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::runtime_env::RuntimeEnv;
+use datafusion_execution::TaskContext;
+use datafusion_expr::{
+    AggregateUDF, AggregateUDFImpl, ScalarUDF, ScalarUDFImpl, WindowUDF, WindowUDFImpl,
+};
+
+use crate::session_config::FFI_SessionConfig;
+use crate::udaf::FFI_AggregateUDF;
+use crate::udf::FFI_ScalarUDF;
+use crate::udwf::FFI_WindowUDF;
+
+/// A stable struct for sharing [`TaskContext`] across FFI boundaries.
+#[repr(C)]
+#[derive(Debug, StableAbi)]
+#[allow(non_camel_case_types)]
+pub struct FFI_TaskContext {
+    /// Return the session ID.
+    pub session_id: unsafe extern "C" fn(&Self) -> RString,
+
+    /// Return the task ID.
+    pub task_id: unsafe extern "C" fn(&Self) -> ROption<RString>,
+
+    /// Return the session configuration.
+    pub session_config: unsafe extern "C" fn(&Self) -> FFI_SessionConfig,
+
+    /// Returns a hashmap of names to scalar functions.
+    pub scalar_functions: unsafe extern "C" fn(&Self) -> RHashMap<RString, FFI_ScalarUDF>,
+
+    /// Returns a hashmap of names to aggregate functions.
+    pub aggregate_functions:
+        unsafe extern "C" fn(&Self) -> RHashMap<RString, FFI_AggregateUDF>,
+
+    /// Returns a hashmap of names to window functions.
+    pub window_functions: unsafe extern "C" fn(&Self) -> RHashMap<RString, FFI_WindowUDF>,
+
+    /// Release the memory of the private data when it is no longer being used.
+    pub release: unsafe extern "C" fn(arg: &mut Self),
+
+    /// Internal data. This is only to be accessed by the provider of the plan.
+    /// The foreign library should never attempt to access this data.
+    pub private_data: *mut c_void,
+
+    /// Utility to identify when FFI objects are accessed locally through
+    /// the foreign interface. See [`crate::get_library_marker_id`] and
+    /// the crate's `README.md` for more information.
+    pub library_marker_id: extern "C" fn() -> usize,
+}
+
+struct TaskContextPrivateData {
+    ctx: Arc<TaskContext>,
+}
+
+impl FFI_TaskContext {
+    unsafe fn inner(&self) -> &Arc<TaskContext> {
+        let private_data = self.private_data as *const TaskContextPrivateData;
+        &(*private_data).ctx
+    }
+}
+
+unsafe extern "C" fn session_id_fn_wrapper(ctx: &FFI_TaskContext) -> RString {
+    let ctx = ctx.inner();
+    ctx.session_id().into()
+}
+
+unsafe extern "C" fn task_id_fn_wrapper(ctx: &FFI_TaskContext) -> ROption<RString> {
+    let ctx = ctx.inner();
+    ctx.task_id().map(|s| s.as_str().into()).into()
+}
+
+unsafe extern "C" fn session_config_fn_wrapper(
+    ctx: &FFI_TaskContext,
+) -> FFI_SessionConfig {
+    let ctx = ctx.inner();
+    ctx.session_config().into()
+}
+
+unsafe extern "C" fn scalar_functions_fn_wrapper(
+    ctx: &FFI_TaskContext,
+) -> RHashMap<RString, FFI_ScalarUDF> {
+    let ctx = ctx.inner();
+    ctx.scalar_functions()
+        .iter()
+        .map(|(name, udf)| (name.to_owned().into(), Arc::clone(udf).into()))
+        .collect()
+}
+
+unsafe extern "C" fn aggregate_functions_fn_wrapper(
+    ctx: &FFI_TaskContext,
+) -> RHashMap<RString, FFI_AggregateUDF> {
+    let ctx = ctx.inner();
+    ctx.aggregate_functions()
+        .iter()
+        .map(|(name, udaf)| {
+            (
+                name.to_owned().into(),
+                FFI_AggregateUDF::from(Arc::clone(udaf)),
+            )
+        })
+        .collect()
+}
+
+unsafe extern "C" fn window_functions_fn_wrapper(
+    ctx: &FFI_TaskContext,
+) -> RHashMap<RString, FFI_WindowUDF> {
+    let ctx = ctx.inner();
+    ctx.window_functions()
+        .iter()
+        .map(|(name, udf)| (name.to_owned().into(), FFI_WindowUDF::from(Arc::clone(udf))))
+        .collect()
+}
+
+unsafe extern "C" fn release_fn_wrapper(ctx: &mut FFI_TaskContext) {
+    let private_data = Box::from_raw(ctx.private_data as *mut TaskContextPrivateData);
+    drop(private_data);
+}
+
+impl Drop for FFI_TaskContext {
+    fn drop(&mut self) {
+        unsafe { (self.release)(self) }
+    }
+}
+
+impl From<Arc<TaskContext>> for FFI_TaskContext {
+    fn from(ctx: Arc<TaskContext>) -> Self {
+        let private_data = Box::new(TaskContextPrivateData { ctx });
+
+        FFI_TaskContext {
+            session_id: session_id_fn_wrapper,
+            task_id: task_id_fn_wrapper,
+            session_config: session_config_fn_wrapper,
+            scalar_functions: scalar_functions_fn_wrapper,
+            aggregate_functions: aggregate_functions_fn_wrapper,
+            window_functions: window_functions_fn_wrapper,
+            release: release_fn_wrapper,
+            private_data: Box::into_raw(private_data) as *mut c_void,
+            library_marker_id: crate::get_library_marker_id,
+        }
+    }
+}
+
+impl From<FFI_TaskContext> for Arc<TaskContext> {
+    fn from(ffi_ctx: FFI_TaskContext) -> Self {
+        unsafe {
+            if (ffi_ctx.library_marker_id)() == crate::get_library_marker_id() {
+                return Arc::clone(ffi_ctx.inner());
+            }
+
+            let task_id = (ffi_ctx.task_id)(&ffi_ctx).map(|s| s.to_string()).into();
+            let session_id = (ffi_ctx.session_id)(&ffi_ctx).into();
+            let session_config = (ffi_ctx.session_config)(&ffi_ctx);
+            let session_config =
+                SessionConfig::try_from(&session_config).unwrap_or_default();
+
+            let scalar_functions = (ffi_ctx.scalar_functions)(&ffi_ctx)
+                .into_iter()
+                .map(|kv_pair| {
+                    let udf = <Arc<dyn ScalarUDFImpl>>::from(&kv_pair.1);
+
+                    (
+                        kv_pair.0.into_string(),
+                        Arc::new(ScalarUDF::new_from_shared_impl(udf)),
+                    )
+                })
+                .collect();
+            let aggregate_functions = (ffi_ctx.aggregate_functions)(&ffi_ctx)
+                .into_iter()
+                .map(|kv_pair| {
+                    let udaf = <Arc<dyn AggregateUDFImpl>>::from(&kv_pair.1);
+
+                    (
+                        kv_pair.0.into_string(),
+                        Arc::new(AggregateUDF::new_from_shared_impl(udaf)),
+                    )
+                })
+                .collect();
+            let window_functions = (ffi_ctx.window_functions)(&ffi_ctx)
+                .into_iter()
+                .map(|kv_pair| {
+                    let udwf = <Arc<dyn WindowUDFImpl>>::from(&kv_pair.1);
+
+                    (
+                        kv_pair.0.into_string(),
+                        Arc::new(WindowUDF::new_from_shared_impl(udwf)),
+                    )
+                })
+                .collect();
+
+            let runtime = Arc::new(RuntimeEnv::default());
+
+            Arc::new(TaskContext::new(
+                task_id,
+                session_id,
+                session_config,
+                scalar_functions,
+                aggregate_functions,
+                window_functions,
+                runtime,
+            ))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use datafusion::prelude::SessionContext;
+    use datafusion_common::Result;
+    use datafusion_execution::TaskContext;
+
+    use crate::execution::FFI_TaskContext;
+
+    #[test]
+    fn ffi_task_ctx_round_trip() -> Result<()> {
+        let session_ctx = SessionContext::new();
+        let original = session_ctx.task_ctx();
+        let mut ffi_task_ctx = FFI_TaskContext::from(Arc::clone(&original));
+        ffi_task_ctx.library_marker_id = crate::mock_foreign_marker_id;
+
+        let foreign_task_ctx: Arc<TaskContext> = ffi_task_ctx.into();
+
+        // TaskContext doesn't implement Eq (nor should it) so check some of the
+        // data is round tripping correctly.
+
+        assert_eq!(
+            original.scalar_functions(),
+            foreign_task_ctx.scalar_functions()
+        );
+        assert_eq!(
+            original.aggregate_functions(),
+            foreign_task_ctx.aggregate_functions()
+        );
+        assert_eq!(
+            original.window_functions(),
+            foreign_task_ctx.window_functions()
+        );
+        assert_eq!(original.task_id(), foreign_task_ctx.task_id());
+        assert_eq!(original.session_id(), foreign_task_ctx.session_id());
+        assert_eq!(
+            format!("{:?}", original.session_config()),
+            format!("{:?}", foreign_task_ctx.session_config())
+        );
+
+        Ok(())
+    }
+}