From f18595791fcd217eae117b29f6a68c330190ae70 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 21 Oct 2026 02:39:58 -0700 Subject: [PATCH 1/3] [python] Rebuild native LeRobot fallback splits lazily in workers --- .../pypaimon/multimodal/lerobot/reader.py | 9 ++++++- .../tests/lerobot_native_reader_test.py | 25 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py index fe53da6e4612..ff2dee72a637 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/reader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -143,7 +143,10 @@ def _ensure_process(self): return self._scanner = None self._scanner_initialized = False - self._set_splits(self._splits) + splits = self._splits + if splits is None: + splits = self._table.new_read_builder().new_scan().plan().splits() + self._set_splits(splits) self._process_id = process_id def _index_scanner(self, predicate): @@ -171,6 +174,10 @@ def __getstate__(self): state["_process_id"] = None state["_split_ranges"] = None state["_split_range_index"] = None + if (self._table.options.native_read_enabled() + or self._table.options.native_plan_enabled()): + # Native reads do not need fallback splits in the worker payload. + state["_splits"] = None return state def __del__(self): diff --git a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py index f5940360cb10..c8e7d932463c 100644 --- a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py +++ b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py @@ -256,3 +256,28 @@ def init_with_auth(scan, *args, **kwargs): assert calls finally: current.close() + + +def test_pickled_native_locator_rebuilds_pinned_fallback_once(frames): + from pypaimon.globalindex import DataEvolutionGlobalIndexScanner + frames.add(pa.table({"index": [40], "value": [4]})) + current = reader(frames, **{ + "scan.native-plan.enabled": "true", "read.native.enabled": "false"}) + restored = pickle.loads(pickle.dumps(current)) + assert restored._locator._splits is None + frames.add(pa.table({"index": [60], "value": [6]})) + try: + with patch("pypaimon.execution.native.read.plan.native_plan", + side_effect=RuntimeError("native unavailable")), \ + patch.object(DataEvolutionGlobalIndexScanner, "create", + wraps=DataEvolutionGlobalIndexScanner.create) as create: + for unused in range(3): + result = restored.read_indices([10, 40, 60], ["value", "index"]) + assert result.column_names == ["value", "index"] + assert sorted(result["index"].to_pylist()) == [10, 40] + assert sorted(result["value"].to_pylist()) == [1, 4] + assert create.call_count == 1 + assert restored._locator._splits is not None + finally: + current.close() + restored.close() From 4b9339eb66fc9f243670350d13ade2f105a3801e Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 21 Oct 2026 07:47:39 -0700 Subject: [PATCH 2/3] Test native read options and Python locator serialization --- .../tests/lerobot_native_reader_test.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py index c8e7d932463c..bc864d243d02 100644 --- a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py +++ b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py @@ -258,11 +258,13 @@ def init_with_auth(scan, *args, **kwargs): current.close() -def test_pickled_native_locator_rebuilds_pinned_fallback_once(frames): +@pytest.mark.parametrize("option", ["scan.native-plan.enabled", "read.native.enabled"]) +def test_pickled_native_locator_rebuilds_pinned_fallback_once(frames, option): from pypaimon.globalindex import DataEvolutionGlobalIndexScanner frames.add(pa.table({"index": [40], "value": [4]})) current = reader(frames, **{ - "scan.native-plan.enabled": "true", "read.native.enabled": "false"}) + "scan.native-plan.enabled": "false", "read.native.enabled": "false", + option: "true"}) restored = pickle.loads(pickle.dumps(current)) assert restored._locator._splits is None frames.add(pa.table({"index": [60], "value": [6]})) @@ -281,3 +283,15 @@ def test_pickled_native_locator_rebuilds_pinned_fallback_once(frames): finally: current.close() restored.close() + + +def test_pickled_python_locator_keeps_fallback_splits(frames): + current = reader(frames, **{ + "scan.native-plan.enabled": "false", "read.native.enabled": "false"}) + restored = pickle.loads(pickle.dumps(current)) + try: + assert restored._locator._splits is not None + assert restored.read_indices([20], ["value"])["value"].to_pylist() == [2] + finally: + current.close() + restored.close() From 3f525d72afd69f5dfedad556e6db1e15a1c78ee4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 21 Oct 2026 23:56:37 +0800 Subject: [PATCH 3/3] [python] Preserve LeRobot fallback projection --- .../pypaimon/multimodal/lerobot/reader.py | 9 ++-- .../tests/lerobot_native_reader_test.py | 49 +++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py index ff2dee72a637..f85269095d0e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/reader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -42,7 +42,8 @@ def __init__(self, frames_table, *, columns): self._table = self._table._copy_with_snapshot(snapshot) self.num_rows = snapshot.next_row_id self.file_io = self._table.file_io - self._locator = _FrameLocator(self._table, snapshot, splits) + self._locator = _FrameLocator( + self._table, snapshot, splits, projection=columns) def read_indices(self, indices, columns): builder = self._table.new_read_builder().with_projection(columns) @@ -74,9 +75,10 @@ def close(self): class _FrameLocator: """Locate LeRobot frame rows in one fixed Paimon snapshot.""" - def __init__(self, table, snapshot, splits): + def __init__(self, table, snapshot, splits, *, projection): self._table = table self._snapshot = snapshot + self._projection = list(projection) self._scanner = None self._scanner_initialized = False self._process_id = os.getpid() @@ -145,7 +147,8 @@ def _ensure_process(self): self._scanner_initialized = False splits = self._splits if splits is None: - splits = self._table.new_read_builder().new_scan().plan().splits() + splits = self._table.new_read_builder().with_projection( + self._projection).new_scan().plan().splits() self._set_splits(splits) self._process_id = process_id diff --git a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py index bc864d243d02..c8f7a3ad17bd 100644 --- a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py +++ b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py @@ -285,6 +285,55 @@ def test_pickled_native_locator_rebuilds_pinned_fallback_once(frames, option): restored.close() +def test_pickled_native_locator_rebuilds_original_projection(tmp_path): + from pypaimon.catalog.table_query_auth import TableQueryAuthResult + from pypaimon.read.table_scan import TableScan + + conn = pm.connect(options={"warehouse": str(tmp_path)}) + table = conn.create_table("projected_frames", data=pa.table({ + "index": [10, 20, 30], + "value": [1, 2, 3], + "secret": [4, 5, 6], + }), options={ + "row-tracking.enabled": "true", + "data-evolution.enabled": "true", + "global-index.enabled": "true", + "bucket": "-1", + }) + table.raw_table.create_global_index("index", index_type="btree") + selected = [] + + def authorize(columns): + selected.append(columns) + if "secret" in columns: + raise PermissionError("secret column denied") + return TableQueryAuthResult(None, None) + + original_init = TableScan.__init__ + + def init_with_auth(scan, *args, **kwargs): + original_init(scan, *args, **kwargs) + scan._query_auth_fn = authorize + + with patch.object(TableScan, "__init__", init_with_auth): + current = reader(table, **{ + "scan.native-plan.enabled": "false", + "read.native.enabled": "true", + }) + restored = pickle.loads(pickle.dumps(current)) + try: + assert restored._locator._splits is None + with patch.object(TableScan, "__init__", init_with_auth), \ + patch("pypaimon.execution.native.runtime.native_runtime_available", + return_value=False): + result = restored.read_indices([20], ["value"]) + assert result["value"].to_pylist() == [2] + assert selected == [["index", "value"], ["index", "value"]] + finally: + current.close() + restored.close() + + def test_pickled_python_locator_keeps_fallback_splits(frames): current = reader(frames, **{ "scan.native-plan.enabled": "false", "read.native.enabled": "false"})