refactor: replace spawn_stream and spawn_streams with spawn_future for consistency

Author: kosiew

src/context.rs

@@ -45,7 +45,7 @@ use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
-use crate::utils::{get_global_ctx, spawn_stream, validate_pycapsule, wait_for_future};
+use crate::utils::{get_global_ctx, spawn_future, validate_pycapsule, wait_for_future};
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
@@ -1131,7 +1131,7 @@ impl PySessionContext {
     ) -> PyDataFusionResult<PyRecordBatchStream> {
         let ctx: TaskContext = TaskContext::from(&self.ctx.state());
         let plan = plan.plan.clone();
-        let stream = spawn_stream(py, async move { plan.execute(part, Arc::new(ctx)) })?;
+        let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?;
         Ok(PyRecordBatchStream::new(stream))
     }
 }

src/dataframe.rs

@@ -51,8 +51,8 @@ use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::{poll_next_batch, PyRecordBatchStream};
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{
-    get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_stream, spawn_streams,
-    validate_pycapsule, wait_for_future,
+    get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_future, validate_pycapsule,
+    wait_for_future,
 };
 use crate::{
     errors::PyDataFusionResult,
@@ -967,7 +967,7 @@ impl PyDataFrame {
         requested_schema: Option<Bound<'py, PyCapsule>>,
     ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {
         let df = self.df.as_ref().clone();
-        let streams = spawn_streams(py, async move { df.execute_stream_partitioned().await })?;
+        let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?;
 
         let mut schema: Schema = self.df.schema().to_owned().into();
         let mut projection: Option<SchemaRef> = None;
@@ -1020,13 +1020,13 @@ impl PyDataFrame {
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {
         let df = self.df.as_ref().clone();
-        let stream = spawn_stream(py, async move { df.execute_stream().await })?;
+        let stream = spawn_future(py, async move { df.execute_stream().await })?;
         Ok(PyRecordBatchStream::new(stream))
     }
 
     fn execute_stream_partitioned(&self, py: Python) -> PyResult<Vec<PyRecordBatchStream>> {
         let df = self.df.as_ref().clone();
-        let streams = spawn_streams(py, async move { df.execute_stream_partitioned().await })?;
+        let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?;
         Ok(streams.into_iter().map(PyRecordBatchStream::new).collect())
     }
 

src/utils.rs

@@ -21,8 +21,7 @@ use crate::{
     TokioRuntime,
 };
 use datafusion::{
-    common::ScalarValue, execution::context::SessionContext, execution::SendableRecordBatchStream,
-    logical_expr::Volatility,
+    common::ScalarValue, execution::context::SessionContext, logical_expr::Volatility,
 };
 use pyo3::prelude::*;
 use pyo3::{exceptions::PyValueError, types::PyCapsule};
@@ -114,27 +113,6 @@ where
     Ok(inner_result?)
 }
 
-/// 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
-/// while respecting Python signal handling.
-pub(crate) fn spawn_streams<F>(
-    py: Python,
-    fut: F,
-) -> PyDataFusionResult<Vec<SendableRecordBatchStream>>
-where
-    F: Future<Output = datafusion::common::Result<Vec<SendableRecordBatchStream>>> + Send + 'static,
-{
-    spawn_future(py, fut)
-}
-
 pub(crate) fn parse_volatility(value: &str) -> PyDataFusionResult<Volatility> {
     Ok(match value {
         "immutable" => Volatility::Immutable,