This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new f33dee15 fix: make Python record batch reader cancellable (#885)
f33dee15 is described below
commit f33dee1568118897a6c2ee39223d15c7107128a1
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Sep 20 16:44:48 2026 +0800
fix: make Python record batch reader cancellable (#885)
---
.../python/python/pypaimon_rust/datafusion.pyi | 3 +
bindings/python/src/read.rs | 108 +++++++++++++++++++--
bindings/python/tests/test_read.py | 13 +++
3 files changed, 115 insertions(+), 9 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 8bf6b3cd..d179df9d 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -60,6 +60,9 @@ class RecordBatchReader:
def __iter__(self) -> "RecordBatchReader": ...
def __next__(self) -> pyarrow.RecordBatch: ...
def read_next_batch(self) -> Optional[pyarrow.RecordBatch]: ...
+ def close(self) -> None:
+ """Stop an in-flight read and release the underlying native stream."""
+ ...
class TableRead:
def read_arrow(self, splits: Sequence[Split]) -> RecordBatchReader:
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index 42c89165..df706c0b 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -16,9 +16,11 @@
// under the License.
use std::collections::HashMap;
+use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use arrow::pyarrow::ToPyArrow;
+use arrow::record_batch::RecordBatch;
use futures::TryStreamExt;
use paimon::spec::Predicate;
use paimon::table::{ArrowRecordBatchStream, DataSplit, IncrementalScanMode,
RowRange, Table};
@@ -26,6 +28,7 @@ use paimon_datafusion::runtime::runtime;
use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};
+use tokio::sync::Notify;
use crate::error::to_py_err;
use crate::predicate::dict_to_predicate;
@@ -438,7 +441,9 @@ impl PyTableRead {
splits: &Bound<'_, PyAny>,
) -> PyResult<PyRecordBatchReader> {
Ok(PyRecordBatchReader {
- stream: Mutex::new(self.read_stream(py, splits)?),
+ stream: Mutex::new(Some(self.read_stream(py, splits)?)),
+ closed: AtomicBool::new(false),
+ close_notify: Notify::new(),
})
}
@@ -459,22 +464,56 @@ impl PyTableRead {
#[pyclass(name = "RecordBatchReader", module = "pypaimon_rust.datafusion")]
pub struct PyRecordBatchReader {
- stream: Mutex<ArrowRecordBatchStream>,
+ stream: Mutex<Option<ArrowRecordBatchStream>>,
+ closed: AtomicBool,
+ close_notify: Notify,
}
impl PyRecordBatchReader {
+ fn next_record_batch(&self) -> PyResult<Option<RecordBatch>> {
+ let mut stream = self
+ .stream
+ .lock()
+ .map_err(|_| PyRuntimeError::new_err("native record batch reader
lock poisoned"))?;
+ if self.closed.load(Ordering::Acquire) {
+ stream.take();
+ return Ok(None);
+ }
+ let result = match stream.as_mut() {
+ Some(stream) => runtime()
+ .block_on(async {
+ tokio::select! {
+ biased;
+ _ = self.close_notify.notified() => Ok(None),
+ batch = stream.try_next() => batch,
+ }
+ })
+ .map_err(to_py_err),
+ None => Ok(None),
+ };
+ if self.closed.load(Ordering::Acquire) {
+ stream.take();
+ }
+ result
+ }
+
fn next_batch(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
- let batch = py.detach(|| {
- let mut stream = self
- .stream
- .lock()
- .map_err(|_| PyRuntimeError::new_err("native record batch
reader lock poisoned"))?;
- runtime().block_on(stream.try_next()).map_err(to_py_err)
- })?;
+ let batch = py.detach(|| self.next_record_batch())?;
batch
.map(|batch| Ok(batch.to_pyarrow(py)?.unbind()))
.transpose()
}
+
+ fn close_reader(&self) {
+ self.closed.store(true, Ordering::Release);
+ // notify_one stores a permit when next_record_batch has not started
+ // polling yet, avoiding a lost wake-up between its closed check and
+ // the select.
+ self.close_notify.notify_one();
+ if let Ok(mut stream) = self.stream.try_lock() {
+ stream.take();
+ }
+ }
}
#[pymethods]
@@ -483,6 +522,11 @@ impl PyRecordBatchReader {
self.next_batch(py)
}
+ /// Stop an in-flight read and release the underlying stream. Idempotent.
+ fn close(&self) {
+ self.close_reader();
+ }
+
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
@@ -586,3 +630,49 @@ impl PySplit {
})
}
}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::{mpsc, Arc};
+ use std::task::Poll;
+ use std::time::Duration;
+
+ use super::*;
+
+ #[test]
+ fn record_batch_reader_close_interrupts_pending_next() {
+ let (polled_tx, polled_rx) = mpsc::channel();
+ let mut announced = false;
+ let stream: ArrowRecordBatchStream =
Box::pin(futures::stream::poll_fn(move |_cx| {
+ if !announced {
+ announced = true;
+ polled_tx.send(()).unwrap();
+ }
+ Poll::<Option<paimon::Result<RecordBatch>>>::Pending
+ }));
+ let reader = Arc::new(PyRecordBatchReader {
+ stream: Mutex::new(Some(stream)),
+ closed: AtomicBool::new(false),
+ close_notify: Notify::new(),
+ });
+ let worker_reader = Arc::clone(&reader);
+ let (result_tx, result_rx) = mpsc::channel();
+ let worker = std::thread::spawn(move || {
+ result_tx
+ .send(matches!(worker_reader.next_record_batch(), Ok(None)))
+ .unwrap();
+ });
+
+ polled_rx
+ .recv_timeout(Duration::from_secs(5))
+ .expect("stream was not polled");
+ reader.close_reader();
+
+ assert!(result_rx
+ .recv_timeout(Duration::from_secs(5))
+ .expect("pending next was not interrupted"));
+ worker.join().unwrap();
+ assert!(matches!(reader.next_record_batch(), Ok(None)));
+ reader.close_reader();
+ }
+}
diff --git a/bindings/python/tests/test_read.py
b/bindings/python/tests/test_read.py
index e772f73d..fa836077 100644
--- a/bindings/python/tests/test_read.py
+++ b/bindings/python/tests/test_read.py
@@ -617,6 +617,19 @@ def test_read_arrow_streams_expected_rows():
assert reader.read_next_batch() is None
+def test_read_arrow_close_is_idempotent_and_stops_iteration():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ builder = table.new_read_builder()
+ splits = builder.new_scan().plan().splits()
+ reader = builder.new_read().read_arrow(splits)
+
+ reader.close()
+ assert reader.read_next_batch() is None
+ assert list(reader) == []
+ reader.close()
+
+
def test_read_empty_splits():
with tempfile.TemporaryDirectory() as warehouse:
table = _make_table_with_data(warehouse)
