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 3e6c43e6 [python] Expose streaming native table reader (#870)
3e6c43e6 is described below
commit 3e6c43e6220fdf0c42a260bdd76c4217a983dae3
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Sep 18 14:14:36 2026 +0800
[python] Expose streaming native table reader (#870)
---
.../python/python/pypaimon_rust/datafusion.pyi | 8 ++
bindings/python/src/context.rs | 1 +
bindings/python/src/read.rs | 100 ++++++++++++++++-----
bindings/python/tests/test_read.py | 14 +++
4 files changed, 100 insertions(+), 23 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index c8abd01b..e9aadcf6 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -56,7 +56,15 @@ class TableScan:
def with_row_position_shard(self, index: int, count: int) -> "TableScan":
...
def plan(self) -> Plan: ...
+class RecordBatchReader:
+ def __iter__(self) -> "RecordBatchReader": ...
+ def __next__(self) -> pyarrow.RecordBatch: ...
+ def read_next_batch(self) -> Optional[pyarrow.RecordBatch]: ...
+
class TableRead:
+ def read_arrow(self, splits: Sequence[Split]) -> RecordBatchReader:
+ """Lazily read splits as PyArrow RecordBatches."""
+ ...
def read(self, splits: Sequence[Split]) -> List[pyarrow.RecordBatch]: ...
class ReadBuilder:
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 9fea4b9f..1e285fd1 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -542,6 +542,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::read::PyTableScan>()?;
this.add_class::<crate::read::PyPlan>()?;
this.add_class::<crate::read::PyTableRead>()?;
+ this.add_class::<crate::read::PyRecordBatchReader>()?;
this.add_class::<crate::read::PySplit>()?;
this.add_class::<crate::schema::PyTableSchema>()?;
this.add_class::<crate::schema::PyDataField>()?;
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index 40eab8b2..0a5ab881 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -16,14 +16,14 @@
// under the License.
use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
use arrow::pyarrow::ToPyArrow;
use futures::TryStreamExt;
use paimon::spec::Predicate;
-use paimon::table::{DataSplit, IncrementalScanMode, RowRange, Table};
+use paimon::table::{ArrowRecordBatchStream, DataSplit, IncrementalScanMode,
RowRange, Table};
use paimon_datafusion::runtime::runtime;
-use pyo3::exceptions::{PyTypeError, PyValueError};
+use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};
@@ -380,32 +380,51 @@ pub struct PyTableRead {
case_sensitive: bool,
}
+impl PyTableRead {
+ fn read_stream(
+ &self,
+ py: Python<'_>,
+ splits: &Bound<'_, PyAny>,
+ ) -> PyResult<ArrowRecordBatchStream> {
+ let splits = extract_splits(splits)?;
+ py.detach(|| {
+ let mut builder = self.table.new_read_builder();
+ apply_read_config(
+ &mut builder,
+ &self.projection,
+ self.limit,
+ &self.filter,
+ self.case_sensitive,
+ )?;
+ // Validate config (e.g. projection) before the empty-splits fast
+ // path so an invalid projection fails consistently regardless of
+ // how many splits are passed.
+ let read = builder.new_read().map_err(to_py_err)?;
+ read.to_arrow(&splits).map_err(to_py_err)
+ })
+ }
+}
+
#[pymethods]
impl PyTableRead {
+ /// Lazily read the given splits as an iterator of PyArrow RecordBatches.
+ fn read_arrow(
+ &self,
+ py: Python<'_>,
+ splits: &Bound<'_, PyAny>,
+ ) -> PyResult<PyRecordBatchReader> {
+ Ok(PyRecordBatchReader {
+ stream: Mutex::new(self.read_stream(py, splits)?),
+ })
+ }
+
/// Read the given splits into a list of PyArrow RecordBatches.
fn read(&self, py: Python<'_>, splits: &Bound<'_, PyAny>) ->
PyResult<Vec<Py<PyAny>>> {
- let splits = extract_splits(splits)?;
+ let stream = self.read_stream(py, splits)?;
let rt = runtime();
let batches = py.detach(|| {
- rt.block_on(async {
- let mut builder = self.table.new_read_builder();
- apply_read_config(
- &mut builder,
- &self.projection,
- self.limit,
- &self.filter,
- self.case_sensitive,
- )?;
- // Validate config (e.g. projection) before the empty-splits
fast
- // path so an invalid projection fails consistently regardless
of
- // how many splits are passed.
- let read = builder.new_read().map_err(to_py_err)?;
- if splits.is_empty() {
- return Ok(Vec::new());
- }
- let stream = read.to_arrow(&splits).map_err(to_py_err)?;
- stream.try_collect::<Vec<_>>().await.map_err(to_py_err)
- })
+ rt.block_on(stream.try_collect::<Vec<_>>())
+ .map_err(to_py_err)
})?;
batches
.iter()
@@ -414,6 +433,41 @@ impl PyTableRead {
}
}
+#[pyclass(name = "RecordBatchReader", module = "pypaimon_rust.datafusion")]
+pub struct PyRecordBatchReader {
+ stream: Mutex<ArrowRecordBatchStream>,
+}
+
+impl PyRecordBatchReader {
+ 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)
+ })?;
+ batch
+ .map(|batch| Ok(batch.to_pyarrow(py)?.unbind()))
+ .transpose()
+ }
+}
+
+#[pymethods]
+impl PyRecordBatchReader {
+ fn read_next_batch(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
+ self.next_batch(py)
+ }
+
+ fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+ slf
+ }
+
+ fn __next__(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
+ self.next_batch(py)
+ }
+}
+
#[pyclass(name = "Plan", module = "pypaimon_rust.datafusion")]
pub struct PyPlan {
splits: Vec<DataSplit>,
diff --git a/bindings/python/tests/test_read.py
b/bindings/python/tests/test_read.py
index 35f019e8..448cc47a 100644
--- a/bindings/python/tests/test_read.py
+++ b/bindings/python/tests/test_read.py
@@ -593,10 +593,22 @@ def test_read_returns_expected_rows():
assert t.sort_by("id").to_pydict() == {"id": [1, 2, 3], "name": ["a",
"b", "c"]}
+def test_read_arrow_streams_expected_rows():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ builder = table.new_read_builder().with_projection(["id"])
+ splits = builder.new_scan().plan().splits()
+ reader = builder.new_read().read_arrow(splits)
+ batches = list(reader)
+ assert pa.Table.from_batches(batches).to_pydict() == {"id": [1, 2, 3]}
+ assert reader.read_next_batch() is None
+
+
def test_read_empty_splits():
with tempfile.TemporaryDirectory() as warehouse:
table = _make_table_with_data(warehouse)
assert table.new_read_builder().new_read().read([]) == []
+ assert list(table.new_read_builder().new_read().read_arrow([])) == []
def test_read_empty_splits_still_validates_projection():
@@ -607,6 +619,8 @@ def test_read_empty_splits_still_validates_projection():
read =
table.new_read_builder().with_projection(["does_not_exist"]).new_read()
with pytest.raises(Exception):
read.read([])
+ with pytest.raises(Exception):
+ read.read_arrow([])
def test_read_non_split_raises_typeerror():
