[python] Speed up LeRobot worker startup with lazy fallback splits

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()
@@ -143,7 +145,11 @@ 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().with_projection(
+                self._projection).new_scan().plan().splits()
+        self._set_splits(splits)
         self._process_id = process_id
 
     def _index_scanner(self, predicate):
@@ -171,6 +177,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):

paimon-python/pypaimon/tests/lerobot_native_reader_test.py

@@ -256,3 +256,91 @@ def init_with_auth(scan, *args, **kwargs):
             assert calls
     finally:
         current.close()
+
+
+@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": "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]}))
+    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()
+
+
+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"})
+    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()