Refactor spawn_future and spawn_stream functions for improved flexibility and code reuse

Author: kosiew

src/utils.rs

@@ -85,17 +85,27 @@ where
     })
 }
 
-/// Spawn a [`SendableRecordBatchStream`] on the Tokio runtime and wait for completion
+/// Spawn a [`Future`] on the Tokio runtime and wait for completion
 /// while respecting Python signal handling.
-pub(crate) fn spawn_stream<F>(py: Python, fut: F) -> PyDataFusionResult<SendableRecordBatchStream>
+pub(crate) fn spawn_future<F, T>(py: Python, fut: F) -> PyDataFusionResult<T>
 where
-    F: Future<Output = datafusion::common::Result<SendableRecordBatchStream>> + Send + 'static,
+    F: Future<Output = datafusion::common::Result<T>> + Send + 'static,
+    T: Send + 'static,
 {
     let rt = &get_tokio_runtime().0;
-    let handle: JoinHandle<datafusion::common::Result<SendableRecordBatchStream>> = rt.spawn(fut);
+    let handle: JoinHandle<datafusion::common::Result<T>> = rt.spawn(fut);
     Ok(wait_for_future(py, async {
         handle.await.map_err(to_datafusion_err)
-    })???)
+    })??)
+}
+
+/// Spawn a [`SendableRecordBatchStream`] on the Tokio runtime and wait for completion
+/// while respecting Python signal handling.
+pub(crate) fn spawn_stream<F>(py: Python, fut: F) -> PyDataFusionResult<SendableRecordBatchStream>
+where
+    F: Future<Output = datafusion::common::Result<SendableRecordBatchStream>> + Send + 'static,
+{
+    spawn_future(py, fut)
 }
 
 /// Spawn a partitioned [`SendableRecordBatchStream`] on the Tokio runtime and wait for completion
@@ -107,12 +117,7 @@ pub(crate) fn spawn_streams<F>(
 where
     F: Future<Output = datafusion::common::Result<Vec<SendableRecordBatchStream>>> + Send + 'static,
 {
-    let rt = &get_tokio_runtime().0;
-    let handle: JoinHandle<datafusion::common::Result<Vec<SendableRecordBatchStream>>> =
-        rt.spawn(fut);
-    Ok(wait_for_future(py, async {
-        handle.await.map_err(to_datafusion_err)
-    })???)
+    spawn_future(py, fut)
 }
 
 pub(crate) fn parse_volatility(value: &str) -> PyDataFusionResult<Volatility> {