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 e602ec58 [core] Support copying stream scan state in Python bindings
(#1187)
e602ec58 is described below
commit e602ec58358684a9d4dafedebd83fa743a38719d
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 21 16:22:14 2026 +0800
[core] Support copying stream scan state in Python bindings (#1187)
---
.github/workflows/ci.yml | 6 +-
.../python/python/pypaimon_rust/datafusion.pyi | 3 +
bindings/python/src/read.rs | 149 +++++++++++++++++----
bindings/python/tests/test_read.py | 18 ++-
bindings/python/tests/test_stream_table_scan.py | 84 ++++++++++++
crates/paimon/src/table/read_builder.rs | 8 +-
crates/paimon/src/table/stream_table_scan.rs | 4 +-
crates/paimon/tests/scan_planning_parity_test.rs | 49 +++++++
crates/paimon/tests/stream_table_scan_test.rs | 67 +++++++++
9 files changed, 348 insertions(+), 40 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d568322f..3b75498a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -187,6 +187,11 @@ jobs:
- ubuntu-latest
- macos-latest
- windows-latest
+ env:
+ # Core and REST-server tests share a runner. Keep backtrace line numbers
+ # without full debug symbols or one-off incremental build caches.
+ CARGO_PROFILE_TEST_DEBUG: line-tables-only
+ CARGO_INCREMENTAL: "0"
steps:
- uses: actions/checkout@v7
@@ -196,7 +201,6 @@ jobs:
run: |
# Full-feature test binaries can exhaust the runner's paging file.
"CARGO_BUILD_JOBS=2" >> $env:GITHUB_ENV
- "CARGO_PROFILE_TEST_DEBUG=line-tables-only" >> $env:GITHUB_ENV
- name: Test
run: cargo test --locked -p paimon --all-targets --features
fulltext,vortex
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index c5b66362..4d4dcefb 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -186,6 +186,9 @@ class TableScan:
class StreamTableScan:
"""Stateful continuous scan; checkpoint is the next snapshot to consume."""
+ def __copy__(self) -> StreamTableScan:
+ """Copy scan configuration and progress without persisting a
consumer."""
+ ...
def with_bucket_filter(self, filter: Optional[Callable[[int], bool]] =
None) -> StreamTableScan: ...
def with_consumer_id(self, consumer_id: str) -> StreamTableScan: ...
def with_shard(self, index: int, count: int) -> StreamTableScan: ...
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index ed7759f3..30ffdf52 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -545,6 +545,43 @@ struct PyBucketFilterState {
error: Option<Py<PyBaseException>>,
}
+impl PyStreamTableScan {
+ fn install_bucket_filter(&mut self) {
+ if self.bucket_filter.lock().unwrap().filter.is_none() {
+ self.scan.with_bucket_filter(|_| Ok(true));
+ return;
+ }
+ let state = Arc::clone(&self.bucket_filter);
+ self.scan.with_bucket_filter(move |bucket| {
+ Python::attach(|py| {
+ // Release the lock before invoking Python: callbacks may run
+ // GC, which traverses this same state.
+ let filter = state
+ .lock()
+ .unwrap()
+ .filter
+ .as_ref()
+ .map(|filter| filter.clone_ref(py));
+ let Some(filter) = filter else {
+ return Ok(true);
+ };
+ filter
+ .call1(py, (bucket,))
+ .and_then(|result| result.is_truthy(py))
+ .map_err(|error| {
+ let error = error.into_value(py);
+ let previous_error =
state.lock().unwrap().error.replace(error);
+ drop(previous_error);
+ paimon::Error::UnexpectedError {
+ message: "Python bucket filter failed".into(),
+ source: None,
+ }
+ })
+ })
+ });
+ }
+}
+
#[pymethods]
impl PyStreamTableScan {
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
@@ -563,6 +600,28 @@ impl PyStreamTableScan {
drop(references);
}
+ /// Copy scan progress with independently owned Python callback references.
+ fn __copy__(&self, py: Python<'_>) -> Self {
+ let filter = self
+ .bucket_filter
+ .lock()
+ .unwrap()
+ .filter
+ .as_ref()
+ .map(|filter| filter.clone_ref(py));
+ let mut copy = Self {
+ scan: self.scan.clone(),
+ bucket_filter: Arc::new(Mutex::new(PyBucketFilterState {
+ filter,
+ error: None,
+ })),
+ };
+ // Each wrapper owns one Python reference and its core callback
captures
+ // that wrapper's state. Sharing the Arc would duplicate GC traversal.
+ copy.install_bucket_filter();
+ copy
+ }
+
#[pyo3(signature = (filter = None))]
fn with_bucket_filter(
mut slf: PyRefMut<'_, Self>,
@@ -577,34 +636,7 @@ impl PyStreamTableScan {
let previous_filter =
std::mem::replace(&mut slf.bucket_filter.lock().unwrap().filter,
filter);
drop(previous_filter);
- let state = Arc::clone(&slf.bucket_filter);
- slf.scan.with_bucket_filter(move |bucket| {
- Python::attach(|py| {
- // Release the lock before invoking Python: callbacks may run
- // GC, which traverses this same state.
- let filter = state
- .lock()
- .unwrap()
- .filter
- .as_ref()
- .map(|filter| filter.clone_ref(py));
- let Some(filter) = filter else {
- return Ok(true);
- };
- filter
- .call1(py, (bucket,))
- .and_then(|result| result.is_truthy(py))
- .map_err(|error| {
- let error = error.into_value(py);
- let previous_error =
state.lock().unwrap().error.replace(error);
- drop(previous_error);
- paimon::Error::UnexpectedError {
- message: "Python bucket filter failed".into(),
- source: None,
- }
- })
- })
- });
+ slf.install_bucket_filter();
Ok(slf)
}
@@ -1210,6 +1242,67 @@ mod tests {
use super::*;
+ #[test]
+ fn stream_scan_copy_keeps_callback_ownership_and_progress_independent() {
+ Python::attach(|py| {
+ for bound_method in [false, true] {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "stream_copy"),
+ "memory:/stream-copy".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let scan =
+ Py::new(py,
PyReadBuilder::new(table).new_stream_scan().unwrap()).unwrap();
+ let globals = PyDict::new(py);
+ globals.set_item("scan", scan).unwrap();
+ globals.set_item("bound_method", bound_method).unwrap();
+ py.run(
+ c"
+import copy
+import gc
+import weakref
+
+class Owner:
+ def __init__(self):
+ self.scan = scan.with_bucket_filter(self.select if bound_method else
self)
+ self.scan.restore(1)
+ self.copied = copy.copy(self.scan)
+
+ def select(self, bucket):
+ return True
+
+ def __call__(self, bucket):
+ return self.select(bucket)
+
+owner = Owner()
+assert owner.scan is not owner.copied
+owner.copied.restore(2)
+assert owner.scan.checkpoint() == 1
+assert owner.copied.checkpoint() == 2
+owner.scan.restore(3)
+assert owner.copied.checkpoint() == 2
+reference = weakref.ref(owner)
+del scan
+gc.collect()
+assert reference() is owner
+del owner
+gc.collect()
+assert reference() is None
+",
+ Some(&globals),
+ None,
+ )
+ .unwrap();
+ }
+ });
+ }
+
#[test]
fn table_read_execution_keywords_validate_before_empty_splits() {
Python::attach(|py| {
diff --git a/bindings/python/tests/test_read.py
b/bindings/python/tests/test_read.py
index 4b56326e..316b6434 100644
--- a/bindings/python/tests/test_read.py
+++ b/bindings/python/tests/test_read.py
@@ -1880,20 +1880,26 @@ def
test_append_row_position_slices_and_shards_are_native_readable(row_tracking)
assert shards == list(range(9))
-def test_append_row_positions_follow_stats_pruned_file_order_before_limit():
[email protected]("row_tracking", [False, True])
[email protected]("limit,expected", [(None, [4, 5, 6, 7]), (2, [4, 5])])
+def
test_append_row_positions_follow_stats_pruned_file_order_before_limit(row_tracking,
limit, expected):
with tempfile.TemporaryDirectory() as warehouse:
- table = _make_append_position_table(warehouse)
+ table = _make_append_position_table(warehouse, row_tracking)
builder = (
table.new_read_builder()
.with_projection(["id"])
.with_filter({"method": "greaterOrEqual", "field": "id",
"literals": [3]})
- .with_limit(2)
)
+ if limit is not None:
+ builder.with_limit(limit)
plan = builder.new_scan().with_row_position_slice(1, 5).plan()
+ # Stats remove the first file before positions are assigned. LIMIT
+ # must not truncate the candidate files before applying the slice.
+ full_selection =
table.new_read_builder().with_projection(["id"]).new_read().read(plan.splits())
+ assert pa.Table.from_batches(full_selection).column("id").to_pylist()
== [4, 5, 6, 7]
actual = pa.Table.from_batches(builder.new_read().read(plan.splits()))
- # with_limit is a planning hint in the Rust API. The complete selected
- # range must survive planning; PyPaimon enforces the final two-row
limit.
- assert actual.column("id").to_pylist() == [4, 5, 6, 7]
+ # The configured reader enforces the final LIMIT after selection.
+ assert actual.column("id").to_pylist() == expected
def test_append_incremental_row_positions_use_combined_delta_order():
diff --git a/bindings/python/tests/test_stream_table_scan.py
b/bindings/python/tests/test_stream_table_scan.py
index 07948428..15ca6624 100644
--- a/bindings/python/tests/test_stream_table_scan.py
+++ b/bindings/python/tests/test_stream_table_scan.py
@@ -17,6 +17,7 @@
"""Stateful core streaming through the Python wrapper."""
+import copy
import gc
import weakref
@@ -228,3 +229,86 @@ def test_snapshot_reader_is_an_internal_component():
assert not hasattr(datafusion, 'SnapshotReader')
assert not hasattr(datafusion.Table, 'new_snapshot_reader')
assert not hasattr(datafusion.ReadBuilder, 'new_snapshot_reader')
+
+
[email protected]('follow_up', [False, True])
+def test_copy_has_independent_progress_and_bucket_callback(table, follow_up):
+ builder = table.new_read_builder()
+ original = builder.new_stream_scan().with_bucket_filter(lambda bucket:
bucket % 2 == 0)
+ original.restore(1 if follow_up else None)
+ copied = copy.copy(original)
+ assert copied is not original
+ assert copied.checkpoint() == original.checkpoint()
+ copied.with_bucket_filter(lambda bucket: bucket % 2 == 1)
+ first = original.plan()
+ assert copied.checkpoint() == (1 if follow_up else None)
+ second = copied.plan()
+ expected = 100 if follow_up else 200
+ assert sorted(_rows(builder, first) + _rows(builder, second)) == [
+ (i, expected + i) for i in range(32)]
+ original.restore(1)
+ assert copied.checkpoint() == (2 if follow_up else 3)
+ assert original.checkpoint() == 1
+
+
[email protected]('bound_method', [False, True])
+def test_copy_owns_independent_callback_references_and_collects_cycles(table,
bound_method):
+ class Owner:
+ def __init__(self):
+ self.scan = table.new_stream_scan().with_bucket_filter(self.select
if bound_method else self)
+ self.copied = copy.copy(self.scan)
+
+ def select(self, bucket):
+ gc.collect()
+ return True
+
+ def __call__(self, bucket):
+ return self.select(bucket)
+
+ owner = Owner()
+ reference = weakref.ref(owner)
+ gc.collect()
+ assert owner.scan.plan().snapshot_id() == 2
+ assert owner.copied.plan().snapshot_id() == 2
+ del owner
+ gc.collect()
+ assert reference() is None
+
+
+def test_replacing_original_callback_does_not_clear_the_copy(table):
+ class Owner:
+ def select(self, bucket):
+ return True
+
+ owner = Owner()
+ reference = weakref.ref(owner)
+ original = table.new_stream_scan().with_bucket_filter(owner.select)
+ copied = copy.copy(original)
+ del owner
+ original.with_bucket_filter(None)
+ gc.collect()
+ assert reference() is not None
+ assert copied.plan().snapshot_id() == 2
+ copied.with_bucket_filter(None)
+ gc.collect()
+ assert reference() is None
+
+
+def
test_copy_preserves_callback_error_identity_without_sharing_configuration(table):
+ failure = LookupError('copied bucket callback failure')
+
+ def select(bucket):
+ raise failure
+
+ original = table.new_stream_scan().with_bucket_filter(select)
+ copied = copy.copy(original)
+ with pytest.raises(LookupError) as caught:
+ copied.plan()
+ assert caught.value is failure
+ assert copied.checkpoint() is None
+ copied.with_bucket_filter(None)
+ assert copied.plan().snapshot_id() == 2
+ with pytest.raises(LookupError) as caught:
+ original.plan()
+ assert caught.value is failure
+ assert original.checkpoint() is None
diff --git a/crates/paimon/src/table/read_builder.rs
b/crates/paimon/src/table/read_builder.rs
index 275203cf..590ebe2b 100644
--- a/crates/paimon/src/table/read_builder.rs
+++ b/crates/paimon/src/table/read_builder.rs
@@ -538,14 +538,14 @@ impl<'a> PaimonReadBuilder<'a> {
}
}
- /// Push a row-limit hint down to scan planning.
+ /// Configure a row limit for scan planning and reads.
///
/// This allows paimon-core scan planning to generate fewer splits when the
/// current scan state keeps split-level `merged_row_count()` conservative.
///
- /// Data-evolution [`TableRead`] enforces the limit before BLOB resolution.
- /// Other read paths still treat it only as a planning hint, so callers or
- /// query engines must enforce the final LIMIT themselves.
+ /// A [`TableRead`] created by this builder also enforces the final limit
+ /// across splits after filtering and row selection. Data-evolution reads
+ /// enforce it before BLOB resolution.
pub fn with_limit(&mut self, limit: usize) -> &mut Self {
self.limit = Some(limit);
self
diff --git a/crates/paimon/src/table/stream_table_scan.rs
b/crates/paimon/src/table/stream_table_scan.rs
index 2052d3ba..88b52f3a 100644
--- a/crates/paimon/src/table/stream_table_scan.rs
+++ b/crates/paimon/src/table/stream_table_scan.rs
@@ -29,7 +29,9 @@ use crate::{Error, Result};
/// otherwise. Empty follow-up plans and irrelevant commits are skipped, as in
/// Java. `None` means there is no available next snapshot; clients can poll.
/// Progress is persisted only by an explicit `notify_checkpoint_complete`.
-#[derive(Debug)]
+/// Clones retain the current progress and read configuration; subsequent
+/// planning, restoration and configuration changes are independent.
+#[derive(Debug, Clone)]
pub struct StreamTableScan {
scan: TableScan<'static>,
next_snapshot_id: Option<i64>,
diff --git a/crates/paimon/tests/scan_planning_parity_test.rs
b/crates/paimon/tests/scan_planning_parity_test.rs
index c5fb07d7..1411a850 100644
--- a/crates/paimon/tests/scan_planning_parity_test.rs
+++ b/crates/paimon/tests/scan_planning_parity_test.rs
@@ -167,6 +167,55 @@ async fn read_column(builder: &ReadBuilder<'_>, plan:
&Plan, column: usize) -> V
.collect()
}
+#[tokio::test]
+async fn append_stats_pruning_then_positions_then_reader_limit() {
+ for row_tracking in [false, true] {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("row-tracking.enabled", row_tracking.to_string())
+ .option("source.split.target-size", "1b")
+ .build()
+ .unwrap();
+ let path = "memory:/planning_parity/append_position_limit";
+ let (io, table) = memory_table(path, TableSchema::new(0, &schema));
+ setup_dirs(&io, path).await;
+ persist_table_schema(&io, path, table.schema()).await;
+ for start in [0, 3, 6] {
+ append_ids(&table, start, start + 3).await;
+ }
+ let predicate = PredicateBuilder::new(table.schema().fields())
+ .greater_or_equal("id", Datum::Int(3))
+ .unwrap();
+ for limit in [None, Some(2)] {
+ let mut builder = table.new_read_builder();
+ builder.with_filter(predicate.clone());
+ if let Some(limit) = limit {
+ builder.with_limit(limit);
+ }
+ let plan = builder
+ .new_scan()
+ .with_row_position_slice(1, 5)
+ .unwrap()
+ .plan()
+ .await
+ .unwrap();
+ assert_eq!(plan.snapshot_id(), Some(3));
+ // An unlimited reader proves that scan planning kept the complete
+ // slice after stats pruned the first three-row file.
+ assert_eq!(read_ids(&table, &plan).await, vec![4, 5, 6, 7]);
+ assert_eq!(
+ read_column(&builder, &plan, 0).await,
+ if limit.is_some() {
+ vec![4, 5]
+ } else {
+ vec![4, 5, 6, 7]
+ }
+ );
+ }
+ }
+}
+
async fn delete_ids(table: &Table, ids: &[i64]) {
let builder = table.new_write_builder();
let mut delete = builder.new_delete().unwrap();
diff --git a/crates/paimon/tests/stream_table_scan_test.rs
b/crates/paimon/tests/stream_table_scan_test.rs
index ab6a7575..ea853eee 100644
--- a/crates/paimon/tests/stream_table_scan_test.rs
+++ b/crates/paimon/tests/stream_table_scan_test.rs
@@ -133,6 +133,73 @@ async fn stream_scan_owns_table_and_read_configuration() {
);
}
+#[tokio::test]
+async fn cloned_scan_keeps_independent_progress_watermarks_and_configuration()
{
+ let table = table("memory:/stream-cloned-state", "none").await;
+ write_batch(&table, &make_batch(vec![1], vec![10])).await;
+ patch_snapshot(&table, 1, "watermark", 100.into()).await;
+ write_batch(&table, &make_batch(vec![1], vec![20])).await;
+ patch_snapshot(&table, 2, "watermark", 200.into()).await;
+ let mut original = table.new_stream_scan().unwrap();
+ original.with_consumer_id("original").unwrap();
+ original.restore(Some(1)).unwrap();
+ assert_eq!(
+ values(&table, original.plan().await.unwrap().unwrap()).await,
+ vec![(1, 10)]
+ );
+ let mut copy = original.clone();
+ assert_eq!(copy.checkpoint(), Some(2));
+ assert_eq!(copy.watermark(), Some(100));
+
+ copy.with_consumer_id("copy").unwrap();
+ copy.with_bucket_filter(|_| Ok(false));
+ assert!(copy.plan().await.unwrap().is_none());
+ assert_eq!(copy.checkpoint(), Some(3));
+ assert_eq!(copy.watermark(), Some(200));
+ assert_eq!(original.checkpoint(), Some(2));
+ assert_eq!(original.watermark(), Some(100));
+ assert_eq!(
+ values(&table, original.plan().await.unwrap().unwrap()).await,
+ vec![(1, 20)]
+ );
+
+ original.restore(Some(1)).unwrap();
+ assert_eq!(copy.checkpoint(), Some(3));
+ assert_eq!(copy.watermark(), Some(200));
+ copy.with_bucket_filter(|_| Ok(true));
+ copy.restore(Some(1)).unwrap();
+ assert_eq!(
+ values(&table, copy.plan().await.unwrap().unwrap()).await,
+ vec![(1, 10)]
+ );
+ assert_eq!(copy.watermark(), Some(100));
+ assert_eq!(original.checkpoint(), Some(1));
+ assert_eq!(original.watermark(), Some(200));
+
+ original.with_bucket_filter(|_| Ok(false));
+ assert!(original.plan().await.unwrap().is_none());
+ assert_eq!(original.checkpoint(), Some(3));
+ assert_eq!(copy.checkpoint(), Some(2));
+ original
+ .notify_checkpoint_complete(original.checkpoint())
+ .await
+ .unwrap();
+ copy.notify_checkpoint_complete(copy.checkpoint())
+ .await
+ .unwrap();
+ let consumers = table.consumer_manager();
+ assert_eq!(consumers.get("original").await.unwrap(), Some(3));
+ assert_eq!(consumers.get("copy").await.unwrap(), Some(2));
+
+ // Resetting one copy restores its own acknowledged consumer checkpoint.
+ copy.restore(None).unwrap();
+ assert_eq!(copy.plan().await.unwrap().unwrap().snapshot_id(), Some(2));
+ assert_eq!(copy.checkpoint(), Some(3));
+ assert_eq!(copy.watermark(), Some(200));
+ assert_eq!(original.checkpoint(), Some(3));
+ assert_eq!(consumers.get("original").await.unwrap(), Some(3));
+}
+
#[tokio::test]
async fn
initial_shards_merge_all_key_versions_and_follow_up_shards_keep_events() {
for (engine, dv, bucket, producer, expected) in [
