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 f091bc2 [python] Add TableRead.read(splits) returning Arrow batches
(#428)
f091bc2 is described below
commit f091bc225f28a87b9a9f38e0ce8235aa13e15a02
Author: Junrui Lee <[email protected]>
AuthorDate: Wed Jul 1 21:19:18 2026 +0800
[python] Add TableRead.read(splits) returning Arrow batches (#428)
---
bindings/python/Cargo.toml | 1 +
.../python/python/pypaimon_rust/datafusion.pyi | 4 +
bindings/python/src/context.rs | 1 +
bindings/python/src/read.rs | 99 +++++++++++++++++++---
bindings/python/tests/test_read.py | 79 +++++++++++++++++
5 files changed, 173 insertions(+), 11 deletions(-)
diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml
index 79d7483..c2e853b 100644
--- a/bindings/python/Cargo.toml
+++ b/bindings/python/Cargo.toml
@@ -30,6 +30,7 @@ doc = false
arrow = { workspace = true, features = ["pyarrow"] }
datafusion = { workspace = true }
datafusion-ffi = { workspace = true }
+futures = "0.3"
paimon = { path = "../../crates/paimon", features = ["storage-all"] }
paimon-datafusion = { path = "../../crates/integrations/datafusion", features
= ["fulltext"] }
pyo3 = { version = "0.28", features = ["abi3-py310"] }
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index d7c7131..184c79e 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -47,11 +47,15 @@ class Plan:
class TableScan:
def plan(self) -> Plan: ...
+class TableRead:
+ def read(self, splits: Sequence[Split]) -> List[pyarrow.RecordBatch]: ...
+
class ReadBuilder:
def with_projection(self, columns: List[str]) -> "ReadBuilder": ...
def with_limit(self, limit: int) -> "ReadBuilder": ...
def with_filter(self, predicate: dict) -> "ReadBuilder": ...
def new_scan(self) -> TableScan: ...
+ def new_read(self) -> "TableRead": ...
class Table:
def identifier(self) -> str: ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index d61b94a..43b1fdb 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -272,6 +272,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::read::PyReadBuilder>()?;
this.add_class::<crate::read::PyTableScan>()?;
this.add_class::<crate::read::PyPlan>()?;
+ this.add_class::<crate::read::PyTableRead>()?;
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 3a059b1..b708457 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -17,16 +17,57 @@
use std::sync::Arc;
+use arrow::pyarrow::ToPyArrow;
+use futures::TryStreamExt;
use paimon::spec::Predicate;
use paimon::table::{DataSplit, Table};
use paimon_datafusion::runtime::runtime;
-use pyo3::exceptions::PyValueError;
+use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};
use crate::error::to_py_err;
use crate::predicate::dict_to_predicate;
+/// Apply projection/limit/filter from a config snapshot onto a core
ReadBuilder.
+/// Shared by PyTableScan::plan and PyTableRead::read so scan and read stay
consistent.
+fn apply_read_config(
+ builder: &mut paimon::table::ReadBuilder<'_>,
+ projection: &Option<Vec<String>>,
+ limit: Option<usize>,
+ filter: &Option<Predicate>,
+) {
+ if let Some(projection) = projection {
+ let cols: Vec<&str> = projection.iter().map(String::as_str).collect();
+ builder.with_projection(&cols);
+ }
+ if let Some(limit) = limit {
+ builder.with_limit(limit);
+ }
+ if let Some(filter) = filter {
+ builder.with_filter(filter.clone());
+ }
+}
+
+/// Extract a sequence of Python `Split` objects into core `DataSplit`s.
Accepts
+/// any iterable (list/tuple/generator). Runs under the GIL since it touches
+/// Python objects. A non-iterable argument or a non-`Split` element raises
+/// `TypeError`.
+fn extract_splits(splits: &Bound<'_, PyAny>) -> PyResult<Vec<DataSplit>> {
+ let iter = splits
+ .try_iter()
+ .map_err(|_| PyTypeError::new_err("read() expects a sequence of Split
objects"))?;
+ let mut out = Vec::new();
+ for item in iter {
+ let item = item?;
+ let split: PyRef<PySplit> = item
+ .extract()
+ .map_err(|_| PyTypeError::new_err("read() expects a sequence of
Split objects"))?;
+ out.push(split.inner.clone());
+ }
+ Ok(out)
+}
+
#[pyclass(name = "ReadBuilder", module = "pypaimon_rust.datafusion")]
pub struct PyReadBuilder {
table: Arc<Table>,
@@ -79,6 +120,15 @@ impl PyReadBuilder {
filter: self.filter.clone(),
}
}
+
+ fn new_read(&self) -> PyTableRead {
+ PyTableRead {
+ table: Arc::clone(&self.table),
+ projection: self.projection.clone(),
+ limit: self.limit,
+ filter: self.filter.clone(),
+ }
+ }
}
#[pyclass(name = "TableScan", module = "pypaimon_rust.datafusion")]
@@ -96,16 +146,7 @@ impl PyTableScan {
let splits = py.detach(|| {
rt.block_on(async {
let mut builder = self.table.new_read_builder();
- if let Some(projection) = &self.projection {
- let cols: Vec<&str> =
projection.iter().map(String::as_str).collect();
- builder.with_projection(&cols);
- }
- if let Some(limit) = self.limit {
- builder.with_limit(limit);
- }
- if let Some(filter) = &self.filter {
- builder.with_filter(filter.clone());
- }
+ apply_read_config(&mut builder, &self.projection, self.limit,
&self.filter);
let plan = builder.new_scan().plan().await.map_err(to_py_err)?;
Ok::<_, PyErr>(plan.splits().to_vec())
})
@@ -114,6 +155,42 @@ impl PyTableScan {
}
}
+#[pyclass(name = "TableRead", module = "pypaimon_rust.datafusion")]
+pub struct PyTableRead {
+ table: Arc<Table>,
+ projection: Option<Vec<String>>,
+ limit: Option<usize>,
+ filter: Option<Predicate>,
+}
+
+#[pymethods]
+impl PyTableRead {
+ /// 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 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);
+ // 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)
+ })
+ })?;
+ batches
+ .iter()
+ .map(|batch| Ok(batch.to_pyarrow(py)?.unbind()))
+ .collect()
+ }
+}
+
#[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 57f814e..6d0f8f0 100644
--- a/bindings/python/tests/test_read.py
+++ b/bindings/python/tests/test_read.py
@@ -18,6 +18,7 @@
import pickle
import tempfile
+import pyarrow as pa
import pytest
from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
@@ -280,3 +281,81 @@ def test_filter_overwrite():
.with_filter({"method": "equal", "field": "dt", "literals":
["p2"]})
.new_scan().plan().splits())
assert overwritten == only_p2
+
+
+def test_read_returns_expected_rows():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ b = table.new_read_builder()
+ splits = b.new_scan().plan().splits()
+ batches = b.new_read().read(splits)
+ t = pa.Table.from_batches(batches)
+ assert t.num_rows == 3
+ assert t.sort_by("id").to_pydict() == {"id": [1, 2, 3], "name": ["a",
"b", "c"]}
+
+
+def test_read_empty_splits():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ assert table.new_read_builder().new_read().read([]) == []
+
+
+def test_read_empty_splits_still_validates_projection():
+ # An invalid projection must fail regardless of split count; the
empty-splits
+ # fast path should not bypass config validation.
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ read =
table.new_read_builder().with_projection(["does_not_exist"]).new_read()
+ with pytest.raises(Exception):
+ read.read([])
+
+
+def test_read_non_split_raises_typeerror():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ with pytest.raises(TypeError):
+ table.new_read_builder().new_read().read([object()])
+
+
+def test_read_pickled_split():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ b = table.new_read_builder()
+ splits = b.new_scan().plan().splits()
+ pickled = [pickle.loads(pickle.dumps(s)) for s in splits]
+ batches = b.new_read().read(pickled)
+ t = pa.Table.from_batches(batches)
+ assert t.num_rows == 3
+ assert sorted(t.column("id").to_pylist()) == [1, 2, 3]
+
+
+def test_read_projection_applied():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ b = table.new_read_builder().with_projection(["id"])
+ splits = b.new_scan().plan().splits()
+ batches = b.new_read().read(splits)
+ t = pa.Table.from_batches(batches)
+ assert t.schema.names == ["id"]
+ assert sorted(t.column("id").to_pylist()) == [1, 2, 3]
+
+
+def test_read_is_config_snapshot():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ b = table.new_read_builder().with_projection(["id"])
+ splits = b.new_scan().plan().splits()
+ read = b.new_read()
+ b.with_projection(["name"]) # mutate builder AFTER new_read()
+ t = pa.Table.from_batches(read.read(splits))
+ assert t.schema.names == ["id"]
+
+
+def test_read_with_filter_smoke():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ b = table.new_read_builder().with_filter(
+ {"method": "greaterOrEqual", "field": "id", "literals": [1]})
+ splits = b.new_scan().plan().splits()
+ batches = b.new_read().read(splits)
+ assert isinstance(batches, list)
