feat: add close method to RecordBatchStream for resource management

Author: kosiew

python/datafusion/record_batch.py

@@ -54,6 +54,9 @@ class RecordBatchStream:
 
     These are typically the result of a
     :py:func:`~datafusion.dataframe.DataFrame.execute_stream` operation.
+
+    Call :py:meth:`close` when finished consuming the stream to avoid
+    lingering background tasks.
     """
 
     def __init__(self, record_batch_stream: df_internal.RecordBatchStream) -> None:
@@ -82,6 +85,16 @@ def __iter__(self) -> typing_extensions.Self:
         """Iterator function."""
         return self
 
+    def close(self) -> None:
+        """Close the stream and release associated resources.
+
+        This drains any remaining batches and allows the underlying
+        :class:`SessionContext` to be released. Call this when you are
+        done consuming the stream to avoid leaving tasks running in the
+        background.
+        """
+        self.rbs.close()
+
 
 def to_record_batch_stream(df: DataFrame) -> RecordBatchStream:
     """Convert a DataFrame into a RecordBatchStream.

python/tests/test_dataframe.py

@@ -1598,6 +1598,18 @@ def test_stream_keeps_context_alive():
     assert table.equals(expected)
 
 
+def test_record_batch_stream_close(df):
+    stream = df.execute_stream()
+
+    batch = next(stream)
+    assert batch is not None
+
+    stream.close()
+
+    with pytest.raises(RuntimeError):
+        next(stream)
+
+
 def test_empty_to_arrow_table(df):
     # Convert empty datafusion dataframe to pyarrow Table
     pyarrow_table = df.limit(0).to_arrow_table()

src/record_batch.rs

@@ -24,7 +24,7 @@ use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::execution::context::SessionContext;
 use datafusion::physical_plan::SendableRecordBatchStream;
 use futures::StreamExt;
-use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration};
+use pyo3::exceptions::{PyRuntimeError, PyStopAsyncIteration, PyStopIteration};
 use pyo3::prelude::*;
 use pyo3::{pyclass, pymethods, PyObject, PyResult, Python};
 use tokio::sync::Mutex;
@@ -66,25 +66,33 @@ pub(crate) fn record_batches_to_pyarrow(
 
 #[pyclass(name = "RecordBatchStream", module = "datafusion", subclass)]
 pub struct PyRecordBatchStream {
-    stream: Arc<Mutex<SendableRecordBatchStream>>,
-    // Hold on to the session context to ensure the underlying execution
-    // environment remains alive for the duration of the stream
-    _ctx: Arc<SessionContext>,
+    /// The underlying stream.  Wrapped in an [`Option`] so it can be
+    /// taken during [`close`](Self::close) and [`Drop`].
+    stream: Option<Arc<Mutex<SendableRecordBatchStream>>>,
+    /// Hold on to the session context to ensure the underlying execution
+    /// environment remains alive for the duration of the stream.  This is
+    /// also wrapped in an [`Option`] so it can be released once the stream
+    /// is closed.
+    _ctx: Option<Arc<SessionContext>>,
 }
 
 impl PyRecordBatchStream {
     pub fn new(stream: SendableRecordBatchStream, ctx: Arc<SessionContext>) -> Self {
         Self {
-            stream: Arc::new(Mutex::new(stream)),
-            _ctx: ctx,
+            stream: Some(Arc::new(Mutex::new(stream))),
+            _ctx: Some(ctx),
         }
     }
 }
 
 #[pymethods]
 impl PyRecordBatchStream {
     fn next(&mut self, py: Python) -> PyResult<PyRecordBatch> {
-        let stream = self.stream.clone();
+        let stream = self
+            .stream
+            .as_ref()
+            .ok_or_else(|| PyRuntimeError::new_err("stream is closed"))?
+            .clone();
         wait_for_future(py, next_stream(stream, true))?
     }
 
@@ -93,7 +101,11 @@ impl PyRecordBatchStream {
     }
 
     fn __anext__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let stream = self.stream.clone();
+        let stream = self
+            .stream
+            .as_ref()
+            .ok_or_else(|| PyRuntimeError::new_err("stream is closed"))?
+            .clone();
         pyo3_async_runtimes::tokio::future_into_py(py, next_stream(stream, false))
     }
 
@@ -104,6 +116,21 @@ impl PyRecordBatchStream {
     fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
         slf
     }
+
+    /// Close the underlying stream.
+    ///
+    /// This drains any remaining record batches and releases the
+    /// [`SessionContext`] associated with the stream.  Users should call
+    /// this method when they are finished consuming the stream to avoid
+    /// leaving background tasks running.
+    fn close(&mut self, py: Python) -> PyResult<()> {
+        if let Some(stream) = self.stream.take() {
+            wait_for_future(py, close_stream(stream))?;
+        }
+        // Drop the context once the stream is drained
+        self._ctx.take();
+        Ok(())
+    }
 }
 
 /// Polls the next batch from a `SendableRecordBatchStream`, converting the `Option<Result<_>>` form.
@@ -132,3 +159,24 @@ async fn next_stream(
         Err(e) => Err(PyDataFusionError::from(e))?,
     }
 }
+
+/// Drain the remaining batches from a [`SendableRecordBatchStream`].
+///
+/// Errors are ignored as the stream is being closed.
+async fn close_stream(stream: Arc<Mutex<SendableRecordBatchStream>>) {
+    let mut stream = stream.lock().await;
+    while let Ok(Some(_)) = poll_next_batch(&mut stream).await {}
+}
+
+impl Drop for PyRecordBatchStream {
+    fn drop(&mut self) {
+        if let Some(stream) = self.stream.take() {
+            // Drain the stream while the context is still alive
+            Python::with_gil(|py| {
+                let _ = wait_for_future(py, close_stream(stream));
+            });
+        }
+        // Drop the context after the stream has been fully drained
+        self._ctx.take();
+    }
+}