diff --git a/src/utils.rs b/src/utils.rs index 0d72eaf75..807517b4c 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -20,20 +20,16 @@ use crate::TokioRuntime; use datafusion::logical_expr::Volatility; use pyo3::prelude::*; use std::future::Future; +use std::sync::{Arc, OnceLock}; use tokio::runtime::Runtime; /// Utility to get the Tokio Runtime from Python -pub(crate) fn get_tokio_runtime(py: Python) -> PyRef { - let datafusion = py.import_bound("datafusion._internal").unwrap(); - let tmp = datafusion.getattr("runtime").unwrap(); - match tmp.extract::>() { - Ok(runtime) => runtime, - Err(_e) => { - let rt = TokioRuntime(tokio::runtime::Runtime::new().unwrap()); - let obj: Bound<'_, TokioRuntime> = Py::new(py, rt).unwrap().into_bound(py); - obj.extract().unwrap() - } - } +pub(crate) fn get_tokio_runtime(_: Python) -> Arc { + static RUNTIME: OnceLock> = OnceLock::new(); + RUNTIME.get_or_init(|| { + let rt = TokioRuntime(tokio::runtime::Runtime::new().unwrap()); + Arc::new(rt) + }).clone() } /// Utility to collect rust futures with GIL released