Enter tokio runtime during other FFI calls, such as execute

Please see the discussion in the original post below.

Some users wish to spawn tasks during calls like execute and others. For pure rust implementations without FFI this isn't a problem. However when using the FFI layer, such as with datafusion-python, we can run into problems where we are not in the correct runtime for spawning tasks. Currently we enter the runtime during the calls to poll_next in the FFI_RecordBatchStream. We need to identify all of the other places we should enter the runtime when being called from foreign code.

This patch resolved the problem described in the discussion and this issue.

diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs
index 00602474d..d6b292450 100644
--- a/datafusion/ffi/src/execution_plan.rs
+++ b/datafusion/ffi/src/execution_plan.rs
@@ -112,6 +112,7 @@ unsafe extern "C" fn execute_fn_wrapper(
     let plan = &(*private_data).plan;
     let ctx = &(*private_data).context;
     let runtime = (*private_data).runtime.clone();
+    let _guard = runtime.as_ref().map(|rt| rt.enter());

     rresult!(plan
         .execute(partition, Arc::clone(ctx))

Discussed in #15691

^{Originally posted by westonpace April 11, 2025}
A contributor of ours (lance) hit a bit of a snag implementing a foreign table provider for datafusion python. It turns out the issue is that we spawn a background thread in ExecutionPlan::execute. The foreign table provider code re-attaches the tokio runtime when it polls the returned stream but it doesn't attach it for the call to execute. I think there are two ways this could be resolved and wanted to see if anyone had strong opinions one way or the other.

I could modify my execution plan implementation (in Lance) so that the background thread is not spawned until it is polled for the first time.
We could enter the foreign runtime before calling the execute function in datafusion

I guess the more general question is:

Is there an expectation that the ExecutionPlan streams do no background work until first polled?

Whoops, you can ignore that mention. I accidentally commented here instead of lance-format/lance#3953 😰

One possibility is to wrap execute...

let plan = plan.clone(); // Should be cheap since users almost always start with Arc<dyn ExecutionPlan>
let schema = plan.schema();
let data_stream = futures::stream::once(async move {
  // attach runtime here
  plan.execute(partition, context)
});
let data_stream = data_stream.try_flatten();
let data_stream = Box::pin(RecordBatchStreamAdapter::new(stream, schema));

This ensures that no tasks are spawned until the stream is polled regardless of how the individual ExecutionPlan nodes are implemented. I'd rather put the burden on the caller that has the need than put the burden on everyone building specialized nodes.

fix: don't panic on DF table provider if test has filter on indexed c…

@timsaucer

This ensures that no tasks are spawned until the stream is polled regardless of how the individual ExecutionPlan nodes are implemented. I'd rather put the burden on the caller that has the need than put the burden on everyone building specialized nodes.

this makes sense to me ☝ but I don't know if there other incoming changes that might superseed this change or other impact

Maybe @timsaucer has some input/thoughts here? 🙏

Navigation Menu

Repository navigation

Enter tokio runtime during other FFI calls, such as execute #16312

Discussed in #15691

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Enter tokio runtime during other FFI calls, such as execute #16312

Description

Discussed in #15691

Activity

westonpace commented on Jun 6, 2025

westonpace commented on Jun 6, 2025

renato2099 commented on Jul 3, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions