diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml index 79d74836..c2e853b3 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 d7c7131c..184c79e3 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 d61b94a9..43b1fdb6 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::()?; this.add_class::()?; this.add_class::()?; + this.add_class::()?; this.add_class::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index 3a059b14..b7084570 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>, + limit: Option, + filter: &Option, +) { + 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> { + 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 = 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, @@ -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
, + projection: Option>, + limit: Option, + filter: Option, +} + +#[pymethods] +impl PyTableRead { + /// Read the given splits into a list of PyArrow RecordBatches. + fn read(&self, py: Python<'_>, splits: &Bound<'_, PyAny>) -> PyResult>> { + 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::>().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, diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index 57f814e1..6d0f8f0a 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)