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 5946150 [python] Expose DataFrame-style scan/split planning API (#415)
5946150 is described below
commit 59461507f70b2449c6b5c04bf673096803f0e8cc
Author: Junrui Lee <[email protected]>
AuthorDate: Sat Jun 27 22:36:47 2026 +0800
[python] Expose DataFrame-style scan/split planning API (#415)
---
bindings/python/Cargo.toml | 1 +
.../python/python/pypaimon_rust/datafusion.pyi | 17 +++
bindings/python/src/context.rs | 4 +
bindings/python/src/lib.rs | 1 +
bindings/python/src/read.rs | 160 +++++++++++++++++++++
bindings/python/src/table.rs | 6 +
bindings/python/tests/test_read.py | 87 +++++++++++
crates/paimon/src/table/source.rs | 21 ++-
8 files changed, 296 insertions(+), 1 deletion(-)
diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml
index 0c3b487..12f0cbd 100644
--- a/bindings/python/Cargo.toml
+++ b/bindings/python/Cargo.toml
@@ -33,4 +33,5 @@ datafusion-ffi = { workspace = true }
paimon = { path = "../../crates/paimon", features = ["storage-all"] }
paimon-datafusion = { path = "../../crates/integrations/datafusion", features
= ["fulltext"] }
pyo3 = { version = "0.28", features = ["abi3-py310"] }
+serde_json = "1.0"
tokio = { workspace = true }
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 172c1fe..cd12609 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -36,10 +36,27 @@ class TableSchema:
def options(self) -> Dict[str, str]: ...
def comment(self) -> Optional[str]: ...
+class Split:
+ def __init__(self, state: bytes) -> None: ...
+ def row_count(self) -> int: ...
+
+class Plan:
+ def splits(self) -> List[Split]: ...
+ def __len__(self) -> int: ...
+
+class TableScan:
+ def plan(self) -> Plan: ...
+
+class ReadBuilder:
+ def with_projection(self, columns: List[str]) -> "ReadBuilder": ...
+ def with_limit(self, limit: int) -> "ReadBuilder": ...
+ def new_scan(self) -> TableScan: ...
+
class Table:
def identifier(self) -> str: ...
def location(self) -> str: ...
def schema(self) -> TableSchema: ...
+ def new_read_builder(self) -> ReadBuilder: ...
class PaimonCatalog:
def __init__(self, catalog_options: Dict[str, str]) -> None: ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 2eb9424..0880302 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -269,6 +269,10 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
let this = PyModule::new(py, "datafusion")?;
this.add_class::<PaimonCatalog>()?;
this.add_class::<crate::table::PyTable>()?;
+ this.add_class::<crate::read::PyReadBuilder>()?;
+ this.add_class::<crate::read::PyTableScan>()?;
+ this.add_class::<crate::read::PyPlan>()?;
+ this.add_class::<crate::read::PySplit>()?;
this.add_class::<crate::schema::PyTableSchema>()?;
this.add_class::<crate::schema::PyDataField>()?;
this.add_class::<PyPythonScalarUDFObject>()?;
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index d0d2002..c4db29a 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -20,6 +20,7 @@ use pyo3::prelude::*;
mod blob;
mod context;
mod error;
+mod read;
mod schema;
mod table;
mod udf;
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
new file mode 100644
index 0000000..4e650d1
--- /dev/null
+++ b/bindings/python/src/read.rs
@@ -0,0 +1,160 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use paimon::table::{DataSplit, Table};
+use paimon_datafusion::runtime::runtime;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyBytes;
+
+use crate::error::to_py_err;
+
+#[pyclass(name = "ReadBuilder", module = "pypaimon_rust.datafusion")]
+pub struct PyReadBuilder {
+ table: Arc<Table>,
+ projection: Option<Vec<String>>,
+ limit: Option<usize>,
+}
+
+impl PyReadBuilder {
+ pub fn new(table: Arc<Table>) -> Self {
+ Self {
+ table,
+ projection: None,
+ limit: None,
+ }
+ }
+}
+
+#[pymethods]
+impl PyReadBuilder {
+ fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec<String>) ->
PyRefMut<'_, Self> {
+ slf.projection = Some(columns);
+ slf
+ }
+
+ fn with_limit(mut slf: PyRefMut<'_, Self>, limit: usize) -> PyRefMut<'_,
Self> {
+ slf.limit = Some(limit);
+ slf
+ }
+
+ fn new_scan(&self) -> PyTableScan {
+ PyTableScan {
+ table: Arc::clone(&self.table),
+ projection: self.projection.clone(),
+ limit: self.limit,
+ }
+ }
+}
+
+#[pyclass(name = "TableScan", module = "pypaimon_rust.datafusion")]
+pub struct PyTableScan {
+ table: Arc<Table>,
+ projection: Option<Vec<String>>,
+ limit: Option<usize>,
+}
+
+#[pymethods]
+impl PyTableScan {
+ fn plan(&self, py: Python<'_>) -> PyResult<PyPlan> {
+ let rt = runtime();
+ 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);
+ }
+ let plan = builder.new_scan().plan().await.map_err(to_py_err)?;
+ Ok::<_, PyErr>(plan.splits().to_vec())
+ })
+ })?;
+ Ok(PyPlan { splits })
+ }
+}
+
+#[pyclass(name = "Plan", module = "pypaimon_rust.datafusion")]
+pub struct PyPlan {
+ splits: Vec<DataSplit>,
+}
+
+#[pymethods]
+impl PyPlan {
+ fn splits(&self) -> Vec<PySplit> {
+ self.splits
+ .iter()
+ .cloned()
+ .map(|inner| PySplit { inner })
+ .collect()
+ }
+
+ fn __len__(&self) -> usize {
+ self.splits.len()
+ }
+}
+
+#[pyclass(name = "Split", module = "pypaimon_rust.datafusion")]
+pub struct PySplit {
+ pub(crate) inner: DataSplit,
+}
+
+impl PySplit {
+ fn to_bytes(&self) -> PyResult<Vec<u8>> {
+ serde_json::to_vec(&self.inner)
+ .map_err(|e| PyValueError::new_err(format!("failed to serialize
split: {e}")))
+ }
+
+ fn from_bytes(bytes: &[u8]) -> PyResult<DataSplit> {
+ serde_json::from_slice(bytes)
+ .map_err(|e| PyValueError::new_err(format!("failed to deserialize
split: {e}")))
+ }
+}
+
+#[pymethods]
+impl PySplit {
+ /// Physical row count: sum of data-file row counts (not a logical result
count).
+ fn row_count(&self) -> i64 {
+ self.inner.row_count()
+ }
+
+ /// Reduce to `Split(bytes)` for pickle/copy. The bytes are an opaque,
+ /// implementation-detail encoding; only same/compatible-version round-trip
+ /// is guaranteed.
+ fn __reduce__<'py>(
+ slf: &Bound<'py, Self>,
+ py: Python<'py>,
+ ) -> PyResult<(Py<PyAny>, (Py<PyBytes>,))> {
+ let bytes = slf.borrow().to_bytes()?;
+ let cls = slf.get_type().unbind().into_any();
+ Ok((cls, (PyBytes::new(py, &bytes).unbind(),)))
+ }
+
+ /// Reconstruct a split from opaque bytes produced by pickling. Direct
+ /// construction without those bytes is unsupported; obtain splits from
+ /// `ReadBuilder.new_scan().plan()`.
+ #[new]
+ fn new(state: &Bound<'_, PyBytes>) -> PyResult<Self> {
+ Ok(Self {
+ inner: Self::from_bytes(state.as_bytes())?,
+ })
+ }
+}
diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs
index 0a0c35c..263c821 100644
--- a/bindings/python/src/table.rs
+++ b/bindings/python/src/table.rs
@@ -19,6 +19,7 @@ use std::sync::Arc;
use pyo3::prelude::*;
+use crate::read::PyReadBuilder;
use crate::schema::PyTableSchema;
#[pyclass(name = "Table", module = "pypaimon_rust.datafusion")]
@@ -46,4 +47,9 @@ impl PyTable {
fn schema(&self) -> PyTableSchema {
PyTableSchema::new(self.inner.schema().clone())
}
+
+ /// Create a [`PyReadBuilder`] for DataFrame-style scan planning.
+ fn new_read_builder(&self) -> PyReadBuilder {
+ PyReadBuilder::new(Arc::clone(&self.inner))
+ }
}
diff --git a/bindings/python/tests/test_read.py
b/bindings/python/tests/test_read.py
new file mode 100644
index 0000000..7b3442f
--- /dev/null
+++ b/bindings/python/tests/test_read.py
@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pickle
+import tempfile
+
+from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
+
+
+def _make_table_with_data(warehouse):
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": warehouse})
+ ctx.sql("CREATE SCHEMA paimon.rdb")
+ ctx.sql("CREATE TABLE paimon.rdb.t (id INT, name STRING)")
+ ctx.sql("INSERT INTO paimon.rdb.t VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+ catalog = PaimonCatalog({"warehouse": warehouse})
+ return catalog.get_table("rdb.t")
+
+
+def test_read_builder_chain_exists():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ builder = table.new_read_builder()
+ scan = builder.with_projection(["id"]).with_limit(2).new_scan()
+ # plan() returns a Plan; deeper assertions are in later tasks.
+ plan = scan.plan()
+ assert plan is not None
+
+
+def test_new_read_builder_plan():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ plan = table.new_read_builder().new_scan().plan()
+ assert len(plan.splits()) >= 1
+
+
+def test_with_projection():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ plan =
table.new_read_builder().with_projection(["id"]).new_scan().plan()
+ assert plan is not None
+
+
+def test_with_limit():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ # limit is a planning hint; assert only that planning succeeds.
+ plan = table.new_read_builder().with_limit(1).new_scan().plan()
+ assert plan is not None
+
+
+def test_plan_len():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ plan = table.new_read_builder().new_scan().plan()
+ assert len(plan) == len(plan.splits())
+
+
+def test_plan_without_filter_succeeds():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ plan = table.new_read_builder().new_scan().plan()
+ assert len(plan.splits()) >= 1
+
+
+def test_split_pickle_roundtrip():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_table_with_data(warehouse)
+ splits = table.new_read_builder().new_scan().plan().splits()
+ assert len(splits) >= 1
+ split = splits[0]
+ restored = pickle.loads(pickle.dumps(split))
+ assert restored.row_count() == split.row_count()
diff --git a/crates/paimon/src/table/source.rs
b/crates/paimon/src/table/source.rs
index 5cc310f..229afa0 100644
--- a/crates/paimon/src/table/source.rs
+++ b/crates/paimon/src/table/source.rs
@@ -413,7 +413,7 @@ impl PartitionBucket {
/// Input split for reading: partition + bucket + list of data files and
optional deletion files.
///
/// Reference:
[org.apache.paimon.table.source.DataSplit](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/source/DataSplit.java)
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataSplit {
snapshot_id: i64,
partition: BinaryRow,
@@ -777,6 +777,25 @@ mod tests {
.unwrap()
}
+ #[test]
+ fn data_split_serde_json_round_trip() {
+ let split = DataSplit::builder()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("file:/tmp/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![])
+ .build()
+ .unwrap();
+
+ let bytes = serde_json::to_vec(&split).expect("serialize");
+ let restored: DataSplit =
serde_json::from_slice(&bytes).expect("deserialize");
+ assert_eq!(restored.snapshot_id(), split.snapshot_id());
+ assert_eq!(restored.bucket(), split.bucket());
+ assert_eq!(restored.bucket_path(), split.bucket_path());
+ }
+
/// Raw convertible split without deletion files: physical sum is exact.
#[test]
fn test_merged_row_count_raw_convertible_sums_physical_rows() {
