feat: enter Tokio runtime integration before scan for TableProvider

Author: kosiew

docs/source/user-guide/io/table_provider.rst

@@ -56,3 +56,8 @@ to the ``SessionContext``.
     ctx.register_table_provider("my_table", provider)
 
     ctx.table("my_table").show()
+
+Using ``wait_for_future`` ensures the embedded Tokio runtime is active when
+DataFusion invokes ``TableProvider::scan``. Your implementation of ``scan`` can
+safely call ``tokio::spawn_blocking`` by entering the runtime with
+``datafusion.enter_tokio_runtime()`` when required.

src/utils.rs

@@ -26,7 +26,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
 use std::future::Future;
 use std::sync::OnceLock;
-use tokio::runtime::Runtime;
+use tokio::runtime::{EnterGuard, Runtime};
 
 /// Utility to get the Tokio Runtime from Python
 #[inline]
@@ -47,14 +47,23 @@ pub(crate) fn get_global_ctx() -> &'static SessionContext {
     CTX.get_or_init(SessionContext::new)
 }
 
+/// Enter the Tokio runtime and return the guard
+#[inline]
+pub(crate) fn enter_tokio_runtime() -> EnterGuard<'static> {
+    get_tokio_runtime().0.enter()
+}
+
 /// Utility to collect rust futures with GIL released
 pub fn wait_for_future<F>(py: Python, f: F) -> F::Output
 where
     F: Future + Send,
     F::Output: Send,
 {
     let runtime: &Runtime = &get_tokio_runtime().0;
-    py.allow_threads(|| runtime.block_on(f))
+    py.allow_threads(|| {
+        let _guard = enter_tokio_runtime();
+        runtime.block_on(f)
+    })
 }
 
 pub(crate) fn parse_volatility(value: &str) -> PyDataFusionResult<Volatility> {