diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 07586e63b..53bc83e0c 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -67,6 +67,9 @@ class ReadBuilder: ... def with_limit(self, limit: int) -> "ReadBuilder": ... def with_filter(self, predicate: dict) -> "ReadBuilder": ... + def with_row_ranges(self, ranges: Sequence[tuple[int, int]]) -> "ReadBuilder": + """Set Data Evolution row ranges. Empty selects no rows; format tables are unsupported.""" + ... def new_scan(self) -> TableScan: ... def new_read(self) -> "TableRead": ... diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index 49217a432..d5d63ecef 100644 --- a/bindings/python/src/read.rs +++ b/bindings/python/src/read.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use arrow::pyarrow::ToPyArrow; use futures::TryStreamExt; use paimon::spec::Predicate; -use paimon::table::{DataSplit, Table}; +use paimon::table::{DataSplit, RowRange, Table}; use paimon_datafusion::runtime::runtime; use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::prelude::*; @@ -62,8 +62,7 @@ fn find_time_travel_selector(opts: &HashMap) -> Option<(&str, &s .find_map(|&name| opts.get(name).map(|v| (name, v.as_str()))) } -/// 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. +/// Apply common scan/read config onto a core ReadBuilder. fn apply_read_config( builder: &mut paimon::table::ReadBuilder<'_>, projection: &Option>, @@ -110,6 +109,7 @@ pub struct PyReadBuilder { projection: Option>, limit: Option, filter: Option, + row_ranges: Option>, case_sensitive: bool, } @@ -120,6 +120,7 @@ impl PyReadBuilder { projection: None, limit: None, filter: None, + row_ranges: None, case_sensitive: true, } } @@ -162,6 +163,7 @@ impl PyReadBuilder { projection: None, limit: None, filter: None, + row_ranges: None, case_sensitive: true, }) } @@ -203,12 +205,32 @@ impl PyReadBuilder { Ok(slf) } + /// Set inclusive row ID ranges for Data Evolution scan planning. + /// Planned splits carry the ranges used by readers. + fn with_row_ranges( + mut slf: PyRefMut<'_, Self>, + ranges: Vec<(i64, i64)>, + ) -> PyResult> { + let mut row_ranges = Vec::with_capacity(ranges.len()); + for (from, to) in ranges { + if from > to { + return Err(PyValueError::new_err(format!( + "row range start {from} exceeds end {to}" + ))); + } + row_ranges.push(RowRange::new(from, to)); + } + slf.row_ranges = Some(row_ranges); + Ok(slf) + } + fn new_scan(&self) -> PyTableScan { PyTableScan { table: Arc::clone(&self.table), projection: self.projection.clone(), limit: self.limit, filter: self.filter.clone(), + row_ranges: self.row_ranges.clone(), case_sensitive: self.case_sensitive, } } @@ -230,6 +252,7 @@ pub struct PyTableScan { projection: Option>, limit: Option, filter: Option, + row_ranges: Option>, case_sensitive: bool, } @@ -247,6 +270,9 @@ impl PyTableScan { &self.filter, self.case_sensitive, )?; + if let Some(row_ranges) = &self.row_ranges { + builder.with_row_ranges(row_ranges.clone()); + } let plan = builder.new_scan().plan().await.map_err(to_py_err)?; Ok::<_, PyErr>(plan.splits().to_vec()) }) diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index c0ea73675..0f4600a46 100644 --- a/bindings/python/tests/test_read.py +++ b/bindings/python/tests/test_read.py @@ -66,6 +66,49 @@ def test_with_limit(): assert plan is not None +def test_with_row_ranges(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + ctx.sql("CREATE SCHEMA paimon.rdb") + ctx.sql("""CREATE TABLE paimon.rdb.de (id INT, name STRING) WITH ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true')""") + ctx.sql("""INSERT INTO paimon.rdb.de (id, name) + VALUES (1, 'a'), (2, 'b'), (3, 'c')""") + table = PaimonCatalog({"warehouse": warehouse}).get_table("rdb.de") + builder = table.new_read_builder().with_row_ranges([(0, 1)]) + plan = builder.new_scan().plan() + batches = builder.new_read().read(plan.splits()) + assert pa.Table.from_batches(batches).column("id").to_pylist() == [1, 2] + + restored_splits = [pickle.loads(pickle.dumps(split)) for split in plan.splits()] + restored_batches = builder.new_read().read(restored_splits) + assert pa.Table.from_batches(restored_batches).column("id").to_pylist() == [1, 2] + + empty_builder = table.new_read_builder().with_row_ranges([]) + empty_plan = empty_builder.new_scan().plan() + assert empty_plan.splits() == [] + assert empty_builder.new_read().read(empty_plan.splits()) == [] + + with pytest.raises(ValueError, match="start 2 exceeds end 1"): + table.new_read_builder().with_row_ranges([(2, 1)]) + + +def test_format_table_rejects_row_ranges(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + ctx.sql("CREATE SCHEMA paimon.rdb") + ctx.sql("""CREATE TABLE paimon.rdb.ft (id INT) WITH ( + 'type' = 'format-table', + 'file.format' = 'parquet')""") + table = PaimonCatalog({"warehouse": warehouse}).get_table("rdb.ft") + + with pytest.raises(NotImplementedError, match="not supported for format tables"): + table.new_read_builder().with_row_ranges([]).new_scan().plan() + + def test_plan_len(): with tempfile.TemporaryDirectory() as warehouse: table = _make_table_with_data(warehouse) diff --git a/crates/paimon/src/table/format_read_builder.rs b/crates/paimon/src/table/format_read_builder.rs index 5d6c1b84d..4360d041d 100644 --- a/crates/paimon/src/table/format_read_builder.rs +++ b/crates/paimon/src/table/format_read_builder.rs @@ -39,6 +39,7 @@ pub(crate) struct FormatReadBuilder<'a> { partition_filter: Option, data_predicates: Vec, limit: Option, + row_ranges: Option>, case_sensitive: bool, parquet_read_budget: Option>, } @@ -52,6 +53,7 @@ impl<'a> FormatReadBuilder<'a> { partition_filter: None, data_predicates: Vec::new(), limit: None, + row_ranges: None, case_sensitive: true, parquet_read_budget: None, } @@ -98,7 +100,8 @@ impl<'a> FormatReadBuilder<'a> { false } - pub(crate) fn with_row_ranges(&mut self, _ranges: Vec) -> &mut Self { + pub(crate) fn with_row_ranges(&mut self, ranges: Vec) -> &mut Self { + self.row_ranges = Some(ranges); self } @@ -119,7 +122,7 @@ impl<'a> FormatReadBuilder<'a> { Vec::new(), None, self.limit, - None, + self.row_ranges.clone(), ) } diff --git a/crates/paimon/src/table/format_table_scan.rs b/crates/paimon/src/table/format_table_scan.rs index df2a48e02..f2f9a075c 100644 --- a/crates/paimon/src/table/format_table_scan.rs +++ b/crates/paimon/src/table/format_table_scan.rs @@ -24,7 +24,7 @@ use crate::spec::{ Datum, PartitionComputer, Predicate, PredicateOperator, }; use crate::table::partition_filter::PartitionFilter; -use crate::table::source::DataSplitBuilder; +use crate::table::source::{DataSplitBuilder, RowRange}; use chrono::NaiveDate; #[derive(Debug, Clone)] @@ -32,6 +32,7 @@ pub(crate) struct FormatTableScan<'a> { table: &'a Table, partition_filter: Option, limit: Option, + row_ranges: Option>, } impl<'a> FormatTableScan<'a> { @@ -39,14 +40,21 @@ impl<'a> FormatTableScan<'a> { table: &'a Table, partition_filter: Option, limit: Option, + row_ranges: Option>, ) -> Self { Self { table, partition_filter, limit, + row_ranges, } } + pub(crate) fn with_row_ranges(mut self, ranges: Vec) -> Self { + self.row_ranges = Some(ranges); + self + } + pub(crate) async fn plan(&self) -> crate::Result { self.ensure_query_auth_allowed()?; self.plan_inner(None).await @@ -64,6 +72,11 @@ impl<'a> FormatTableScan<'a> { } async fn plan_inner(&self, trace: Option<&mut ScanTrace>) -> crate::Result { + if self.row_ranges.is_some() { + return Err(crate::Error::Unsupported { + message: "Row ranges are not supported for format tables".to_string(), + }); + } let core_options = CoreOptions::new(self.table.schema().options()); let format_extension = supported_format_table_extension(core_options.file_format())?; let schema_id = self.table.schema().id(); diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index 94c89c2bd..cce6838ec 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -210,7 +210,8 @@ impl<'a> ReadBuilder<'a> { } } - /// Set row ID ranges `[from, to]` (inclusive) for filtering in data evolution mode. + /// Set Data Evolution row ID ranges `[from, to]` (inclusive). + /// An empty vector selects no rows. Format tables are not supported. pub fn with_row_ranges(&mut self, ranges: Vec) -> &mut Self { match &mut self.0 { ReadBuilderKind::Paimon(builder) => { @@ -399,13 +400,9 @@ impl<'a> PaimonReadBuilder<'a> { ) } - /// Set row ID ranges `[from, to]` (inclusive) for filtering in data evolution mode. + /// Set row ID ranges `[from, to]` (inclusive). An empty vector selects no rows. pub fn with_row_ranges(&mut self, ranges: Vec) -> &mut Self { - self.row_ranges = if ranges.is_empty() { - None - } else { - Some(ranges) - }; + self.row_ranges = Some(ranges); self } @@ -780,6 +777,50 @@ mod tests { ) } + #[test] + fn test_with_empty_row_ranges_is_preserved() { + let table = simple_table(); + let mut builder = table.new_read_builder(); + builder.with_row_ranges(Vec::new()); + + assert_eq!(paimon_builder(&builder).row_ranges, Some(Vec::new())); + } + + #[tokio::test] + async fn test_format_table_rejects_row_ranges() { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .option("type", "format-table") + .option("file.format", "parquet") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "format_t"), + "memory:/format_t".to_string(), + TableSchema::new(0, &schema), + None, + ); + + let mut builder = table.new_read_builder(); + builder.with_row_ranges(Vec::new()); + let error = builder.new_scan().plan().await.unwrap_err(); + assert!( + matches!(error, crate::Error::Unsupported { ref message } if message.contains("format tables")) + ); + + let error = table + .new_read_builder() + .new_scan() + .with_row_ranges(Vec::new()) + .plan() + .await + .unwrap_err(); + assert!( + matches!(error, crate::Error::Unsupported { ref message } if message.contains("format tables")) + ); + } + fn dv_pk_table(table_path: &str, merge_engine: &str) -> Table { let file_io = FileIOBuilder::new("file").build().unwrap(); let table_schema = TableSchema::new( diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index ded0cab02..6e3e69ea3 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -796,6 +796,7 @@ impl<'a> TableScan<'a> { table, partition_filter, limit, + row_ranges, ))) } else { Self(TableScanKind::Paimon(PaimonTableScan::new( @@ -821,7 +822,9 @@ impl<'a> TableScan<'a> { TableScanKind::Paimon(scan) => { Self(TableScanKind::Paimon(scan.with_row_ranges(ranges))) } - TableScanKind::Format(scan) => Self(TableScanKind::Format(scan)), + TableScanKind::Format(scan) => { + Self(TableScanKind::Format(scan.with_row_ranges(ranges))) + } } } @@ -964,13 +967,9 @@ impl<'a> PaimonTableScan<'a> { /// Set row ranges for scan-time filtering. /// /// This replaces any existing row_ranges. Typically used to inject - /// results from global index lookups (e.g. full-text search). + /// results from global index lookups. An empty vector selects no rows. pub fn with_row_ranges(mut self, ranges: Vec) -> Self { - self.row_ranges = if ranges.is_empty() { - None - } else { - Some(ranges) - }; + self.row_ranges = Some(ranges); self } @@ -2476,6 +2475,15 @@ mod tests { ) } + #[test] + fn test_scan_with_empty_row_ranges_is_preserved() { + let table = limit_test_table(); + let scan = PaimonTableScan::new(&table, None, Vec::new(), None, None, None) + .with_row_ranges(Vec::new()); + + assert_eq!(scan.row_ranges, Some(Vec::new())); + } + fn limit_test_split(file_name: &str, row_count: i64) -> DataSplit { let mut file = test_data_file_meta(Vec::new(), Vec::new(), Vec::new(), row_count); file.file_name = file_name.to_string();