Refactor stream handling to improve Python signal checking and add utility for polling record batch streams

Author: kosiew

src/dataframe.rs

@@ -43,17 +43,16 @@ use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods};
 use rayon::prelude::*;
-use tokio::time::{sleep, Duration};
 
 use crate::catalog::PyTable;
-use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError};
+use crate::errors::{py_datafusion_err, PyDataFusionError};
 use crate::expr::sort_expr::to_sort_expressions;
 use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{
     get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_and_wait, validate_pycapsule,
-    wait_for_future,
+    wait_for_future, wait_for_stream_next,
 };
 use crate::{
     errors::PyDataFusionResult,
@@ -1018,30 +1017,10 @@ impl Iterator for ArrowStreamReader {
     type Item = Result<RecordBatch, ArrowError>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(1_000);
-        let rt = &get_tokio_runtime().0;
-        let fut = self.stream.next();
-
-        let result = Python::with_gil(|py| {
-            py.allow_threads(|| {
-                rt.block_on(async {
-                    tokio::pin!(fut);
-                    loop {
-                        tokio::select! {
-                            res = &mut fut => break res,
-                            _ = sleep(INTERVAL_CHECK_SIGNALS) => {
-                                if let Err(err) = Python::with_gil(|py| py.check_signals()) {
-                                    break Some(Err(to_datafusion_err(err)));
-                                }
-                            }
-                        }
-                    }
-                })
-            })
-        });
+        let result = Python::with_gil(|py| wait_for_stream_next(py, &mut self.stream));
 
         match result {
-            Some(Ok(batch)) => {
+            Ok(Some(batch)) => {
                 let batch = if self.project {
                     match record_batch_into_schema(batch, self.schema.as_ref()) {
                         Ok(b) => b,
@@ -1052,8 +1031,8 @@ impl Iterator for ArrowStreamReader {
                 };
                 Some(Ok(batch))
             }
-            Some(Err(e)) => Some(Err(ArrowError::from(e))),
-            None => None,
+            Ok(None) => None,
+            Err(e) => Some(Err(ArrowError::from(e))),
         }
     }
 }

src/record_batch.rs

@@ -18,7 +18,7 @@
 use std::sync::Arc;
 
 use crate::errors::PyDataFusionError;
-use crate::utils::wait_for_future;
+use crate::utils::{wait_for_future, wait_for_stream_next};
 use datafusion::arrow::pyarrow::ToPyArrow;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::physical_plan::SendableRecordBatchStream;
@@ -59,17 +59,14 @@ impl PyRecordBatchStream {
     }
 }
 
-pub(crate) async fn pull_next_batch(
-    stream: &mut SendableRecordBatchStream,
-) -> Option<datafusion::common::Result<RecordBatch>> {
-    stream.next().await
-}
-
 #[pymethods]
 impl PyRecordBatchStream {
     fn next(&mut self, py: Python) -> PyResult<PyRecordBatch> {
-        let stream = self.stream.clone();
-        wait_for_future(py, next_stream(stream, true))?
+        let mut stream = wait_for_future(py, self.stream.lock())?;
+        match wait_for_stream_next(py, &mut stream).map_err(PyDataFusionError::from)? {
+            Some(batch) => Ok(batch.into()),
+            None => Err(PyStopIteration::new_err("stream exhausted")),
+        }
     }
 
     fn __next__(&mut self, py: Python) -> PyResult<PyRecordBatch> {
@@ -95,7 +92,7 @@ async fn next_stream(
     sync: bool,
 ) -> PyResult<PyRecordBatch> {
     let mut stream = stream.lock().await;
-    match pull_next_batch(&mut stream).await {
+    match stream.next().await {
         Some(Ok(batch)) => Ok(batch.into()),
         Some(Err(e)) => Err(PyDataFusionError::from(e))?,
         None => {

src/utils.rs

@@ -20,9 +20,11 @@ use crate::{
     errors::{to_datafusion_err, PyDataFusionError, PyDataFusionResult},
     TokioRuntime,
 };
+use datafusion::{arrow::record_batch::RecordBatch, physical_plan::SendableRecordBatchStream};
 use datafusion::{
     common::ScalarValue, execution::context::SessionContext, logical_expr::Volatility,
 };
+use futures::StreamExt;
 use pyo3::prelude::*;
 use pyo3::{exceptions::PyValueError, types::PyCapsule};
 use std::{future::Future, sync::OnceLock, time::Duration};
@@ -84,6 +86,22 @@ where
     })
 }
 
+/// Poll `SendableRecordBatchStream::next` while checking Python signals.
+///
+/// This utility mirrors [`wait_for_future`] for streams and converts any
+/// Python interruptions into a [`DataFusionError`].
+pub fn wait_for_stream_next(
+    py: Python,
+    stream: &mut SendableRecordBatchStream,
+) -> datafusion::common::Result<Option<RecordBatch>> {
+    match wait_for_future(py, stream.next()) {
+        Ok(Some(Ok(batch))) => Ok(Some(batch)),
+        Ok(Some(Err(e))) => Err(e),
+        Ok(None) => Ok(None),
+        Err(err) => Err(to_datafusion_err(err)),
+    }
+}
+
 pub fn spawn_and_wait<F, T>(py: Python, fut: F) -> PyDataFusionResult<T>
 where
     F: Future<Output = datafusion::common::Result<T>> + Send + 'static,