Refactor record batch streaming to use poll_next_batch for improved error handling

Author: kosiew

src/dataframe.rs

@@ -48,7 +48,7 @@ use crate::catalog::PyTable;
 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::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,
@@ -390,11 +390,11 @@ impl Iterator for DataFrameStreamReader {
         // respecting Python signal handling (e.g. ``KeyboardInterrupt``).
         // This mirrors the behaviour of other synchronous wrappers and
         // prevents blocking indefinitely when a Python interrupt is raised.
-        let fut = self.stream.next();
+        let fut = poll_next_batch(&mut self.stream);
         let result = Python::with_gil(|py| wait_for_future(py, fut));
 
         match result {
-            Ok(Some(Ok(batch))) => {
+            Ok(Ok(Some(batch))) => {
                 let batch = if let Some(ref schema) = self.projection {
                     match record_batch_into_schema(batch, schema.as_ref()) {
                         Ok(b) => b,
@@ -405,8 +405,8 @@ impl Iterator for DataFrameStreamReader {
                 };
                 Some(Ok(batch))
             }
-            Ok(Some(Err(e))) => Some(Err(ArrowError::ExternalError(Box::new(e)))),
-            Ok(None) => None,
+            Ok(Ok(None)) => None,
+            Ok(Err(e)) => Some(Err(ArrowError::ExternalError(Box::new(e)))),
             Err(e) => Some(Err(ArrowError::ExternalError(Box::new(e)))),
         }
     }

src/record_batch.rs

@@ -84,15 +84,20 @@ impl PyRecordBatchStream {
     }
 }
 
+pub(crate) async fn poll_next_batch(
+    stream: &mut SendableRecordBatchStream,
+) -> datafusion::error::Result<Option<RecordBatch>> {
+    stream.next().await.transpose()
+}
+
 async fn next_stream(
     stream: Arc<Mutex<SendableRecordBatchStream>>,
     sync: bool,
 ) -> PyResult<PyRecordBatch> {
     let mut stream = stream.lock().await;
-    match stream.next().await {
-        Some(Ok(batch)) => Ok(batch.into()),
-        Some(Err(e)) => Err(PyDataFusionError::from(e))?,
-        None => {
+    match poll_next_batch(&mut stream).await {
+        Ok(Some(batch)) => Ok(batch.into()),
+        Ok(None) => {
             // Depending on whether the iteration is sync or not, we raise either a
             // StopIteration or a StopAsyncIteration
             if sync {
@@ -101,5 +106,6 @@ async fn next_stream(
                 Err(PyStopAsyncIteration::new_err("stream exhausted"))
             }
         }
+        Err(e) => Err(PyDataFusionError::from(e))?,
     }
 }