inject tokio runtime

Author: ccciudatu

src/dataframe.rs

@@ -714,12 +714,6 @@ impl Shard {
         &CODEC
     }
 }
-#[pymethods]
-impl Shard {
-    pub fn stream(&self) -> PyResult<PyRecordBatchStream> {
-        shard_stream(self.serialized_plan.as_ref())
-    }
-}
 
 #[pyclass(get_all)]
 #[derive(Debug, Clone)]
@@ -795,7 +789,7 @@ async fn split_physical_plan(df: &DataFrame, num_shards: usize) -> Result<Distri
 }
 
 #[pyfunction]
-pub fn shard_stream(serialized_shard_plan: &[u8]) -> PyResult<PyRecordBatchStream> {
+pub fn shard_stream(serialized_shard_plan: &[u8], py: Python) -> PyResult<PyRecordBatchStream> {
     deltalake::ensure_initialized();
     let registry = MemoryFunctionRegistry::default();
     let runtime = RuntimeEnvBuilder::new().build()?;
@@ -805,8 +799,13 @@ pub fn shard_stream(serialized_shard_plan: &[u8]) -> PyResult<PyRecordBatchStrea
         .map_err(py_datafusion_err)?;
     let plan = node.try_into_physical_plan(&registry, &runtime, &codec)?;
     println!("Shard plan: {}", displayable(plan.as_ref()).one_line());
-    let ctx = TaskContext::default();
-    execute_stream(plan, Arc::new(ctx)).map(PyRecordBatchStream::new).map_err(py_datafusion_err)
+    let stream_with_runtime = get_tokio_runtime().0.spawn(async move {
+        execute_stream(plan, Arc::new(TaskContext::default()))
+    });
+    wait_for_future(py, stream_with_runtime)
+        .map_err(py_datafusion_err)?
+        .map(PyRecordBatchStream::new)
+        .map_err(py_datafusion_err)
 }
 
 /// Print DataFrame