This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 79bdc93798 [python] Use native planning for LeRobot frame lookups
(#10701)
79bdc93798 is described below
commit 79bdc9379885c2cba19f2284c94a26fac37b876c
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Oct 21 14:15:28 2026 +0800
[python] Use native planning for LeRobot frame lookups (#10701)
---
docs/docs/pypaimon/lerobot.md | 5 +
.../pypaimon/multimodal/lerobot/reader.py | 25 +-
.../pypaimon/tests/lerobot_native_reader_test.py | 258 +++++++++++++++++++++
3 files changed, 287 insertions(+), 1 deletion(-)
diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md
index 53c32451a5..45120edb2f 100644
--- a/docs/docs/pypaimon/lerobot.md
+++ b/docs/docs/pypaimon/lerobot.md
@@ -248,6 +248,11 @@ on `index`; payloads remain lazy. Indexed videos prefetch
metadata and target GO
uncached PyAV reads on demand. Unindexed videos use the TorchCodec/PyAV scan
path. Set
`video_backend` to force either decoder.
+With `pypaimon-rust` installed, LeRobot defaults to Rust frame lookup and
table reads
+unless native read/planning options are explicitly configured. Set
+`table = table.copy({"read.native.enabled": "false"})` to use Python.
+Reads stay on the selected snapshot; unavailable or unsupported native paths
use Python.
+
Set `table = table.copy({"read.video.max-open-decoders": "8"})` before
constructing the dataset to limit cached decoders per camera per worker
(default: 16).
Each open decoder uses memory and a file handle, so usage grows with cameras
and workers.
diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py
b/paimon-python/pypaimon/multimodal/lerobot/reader.py
index a63945fd86..fe53da6e46 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/reader.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py
@@ -30,15 +30,38 @@ class _PaimonTableFrameReader:
def __init__(self, frames_table, *, columns):
self._table, self.snapshot_id, splits = _indexed_read_table(
frames_table, columns)
+ options = self._table.options.options
+ if not any(options.contains(option) for option in (
+ CoreOptions.READ_NATIVE_ENABLED,
CoreOptions.SCAN_NATIVE_PLAN_ENABLED)):
+ from pypaimon.execution.native.runtime import
native_runtime_available
+ if native_runtime_available():
+ self._table = self._table.copy_without_time_travel({
+ CoreOptions.READ_NATIVE_ENABLED.key(): "true"})
snapshot = self._table.snapshot_manager().get_snapshot_by_id(
self.snapshot_id)
+ 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)
def read_indices(self, indices, columns):
- splits, needs_filter = self._locator.locate(indices)
builder = self._table.new_read_builder().with_projection(columns)
+ if (self._table.options.native_plan_enabled()
+ or self._table.options.native_read_enabled()):
+ from pypaimon.execution.native.runtime import
native_runtime_available
+ if native_runtime_available():
+ # Keep lookup predicates in the native plan and residual read.
+ builder.with_filter(_index_predicate(self._table, indices))
+ scan = builder.new_scan()
+ if scan._native_plan_supported():
+ # Authorized scans must retain the complete planning path.
+ if scan._query_auth_fn is not None or
self._table.options.query_auth_enabled:
+ return
builder.new_read().to_arrow(scan.plan().splits())
+ plan = scan._try_native_plan()
+ if plan is not None:
+ return builder.new_read().to_arrow(plan.splits())
+ builder =
self._table.new_read_builder().with_projection(columns)
+ splits, needs_filter = self._locator.locate(indices)
if needs_filter:
builder = builder.with_filter(
_index_predicate(self._table, indices))
diff --git a/paimon-python/pypaimon/tests/lerobot_native_reader_test.py
b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py
new file mode 100644
index 0000000000..ef413f2a64
--- /dev/null
+++ b/paimon-python/pypaimon/tests/lerobot_native_reader_test.py
@@ -0,0 +1,258 @@
+# 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 multiprocessing
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+import pypaimon.multimodal as pm
+from pypaimon.multimodal.lerobot.reader import _PaimonTableFrameReader
+from pypaimon.execution.native.runtime import native_runtime_available
+from pypaimon.execution.native.read.plan import native_plan
+
+pytest.importorskip("torch")
+
+
+def _read_in_worker(reader, output):
+ try:
+ assert reader._table.options.native_read_enabled()
+ with patch.object(reader._locator, "locate",
side_effect=AssertionError("Python frame lookup")):
+ output.put(reader.read_indices([10, 40],
["index"])["index"].to_pylist())
+ finally:
+ reader.close()
+
+
[email protected]
+def frames(tmp_path):
+ conn = pm.connect(options={"warehouse": str(tmp_path)})
+ table = conn.create_table("frames", data=pa.table({
+ "index": [10, 20, 30], "value": [1, 2, 3],
+ }), 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")
+ return table
+
+
+def reader(table, **options):
+ return _PaimonTableFrameReader(
+ table.raw_table.copy(options), columns=["index", "value"])
+
+
[email protected](not native_runtime_available(), reason="native reader
required")
[email protected]("option", [None, "read.native.enabled",
"scan.native-plan.enabled"])
+def test_native_frame_lookup_keeps_residual_and_snapshot(frames, option):
+ # The index covers the original rows, not this append.
+ frames.add(pa.table({"index": [40, 50], "value": [4, 5]}))
+ current = reader(frames, **({option: "true"} if option else {}))
+ if option is None:
+ assert current._table.options.native_read_enabled()
+ frames.add(pa.table({"index": [60], "value": [6]}))
+ native_plans = []
+
+ def checked_plan(*args, **kwargs):
+ plan = native_plan(*args, **kwargs)
+ native_plans.append(plan)
+ return plan
+
+ try:
+ with patch.object(current._locator, "locate",
+ side_effect=AssertionError("Python frame lookup")), \
+ patch("pypaimon.execution.native.read.plan.native_plan",
side_effect=checked_plan):
+ for indices, expected in [([40, 20, 40], [20, 40]),
+ ([60, 99], []), ([], [])]:
+ result = current.read_indices(indices, ["value", "index"])
+ assert result.column_names == ["value", "index"]
+ assert sorted(result["index"].to_pylist()) == expected
+ assert len(native_plans) == 3
+ assert all(plan.snapshot_id == current.snapshot_id for plan in
native_plans)
+ finally:
+ current.close()
+
+
[email protected](not native_runtime_available(), reason="native reader
required")
+def test_native_frame_lookup_preserves_tag_after_pickle(frames):
+ frames.raw_table.create_tag("training")
+ current = reader(frames, **{
+ "scan.tag-name": "training",
+ })
+ restored = pickle.loads(pickle.dumps(current))
+ frames.add(pa.table({"index": [40], "value": [4]}))
+ frames.raw_table.replace_tag("training")
+ try:
+ assert current.snapshot_id == restored.snapshot_id
+ assert restored.read_indices([10, 40], ["index"])["index"].to_pylist()
== [10]
+ finally:
+ current.close()
+ restored.close()
+
+
[email protected](not native_runtime_available(), reason="native reader
required")
+def test_native_frame_lookup_in_spawn_worker(frames):
+ current = reader(frames)
+ frames.add(pa.table({"index": [40], "value": [4]}))
+ context = multiprocessing.get_context("spawn")
+ output = context.Queue()
+ process = context.Process(target=_read_in_worker, args=(current, output))
+ try:
+ process.start()
+ assert output.get(timeout=30) == [10]
+ process.join(timeout=30)
+ assert process.exitcode == 0
+ finally:
+ if process.is_alive():
+ process.terminate()
+ process.join()
+ output.close()
+ current.close()
+
+
+def test_missing_native_runtime_keeps_cached_python_lookup(frames):
+ current = reader(frames, **{"read.native.enabled": "true"})
+ try:
+ with
patch("pypaimon.execution.native.runtime.native_runtime_available",
return_value=False), \
+ patch.object(current._locator, "locate",
wraps=current._locator.locate) as locate:
+ assert current.read_indices([30], ["value"])["value"].to_pylist()
== [3]
+ locate.assert_called_once_with([30])
+ finally:
+ current.close()
+
+
[email protected](not native_runtime_available(), reason="native reader
required")
+def test_native_plan_failure_reuses_scanner_and_keeps_predicate(frames):
+ from pypaimon.globalindex import DataEvolutionGlobalIndexScanner
+ frames.add(pa.table({"index": [40, 50], "value": [4, 5]}))
+ current = reader(frames, **{
+ "scan.native-plan.enabled": "true", "read.native.enabled": "false",
+ })
+ try:
+ with patch("pypaimon.execution.native.read.plan.native_plan",
+ side_effect=RuntimeError("unavailable")), \
+ patch.object(current._locator, "locate",
wraps=current._locator.locate) as locate, \
+ patch.object(DataEvolutionGlobalIndexScanner, "create",
+ wraps=DataEvolutionGlobalIndexScanner.create) as
create:
+ for indices, expected in [([20], [2]), ([40], [4]), ([10, 99],
[1])]:
+ assert current.read_indices(indices,
["value"])["value"].to_pylist() == expected
+ assert locate.call_count == 3
+ assert create.call_count == 1
+ scanner = current._locator._scanner
+ with patch.object(scanner, "close", wraps=scanner.close) as close:
+ current.close()
+ close.assert_called_once()
+ finally:
+ current.close()
+
+
[email protected]_plan
[email protected]_read
[email protected]("options", [
+ {"read.native.enabled": "false"},
+ {"scan.native-plan.enabled": "false"},
+])
+def test_explicit_python_lookup_is_preserved(frames, options):
+ current = reader(frames, **options)
+ try:
+ assert not current._table.options.native_read_enabled()
+ with patch.object(current._locator, "locate",
wraps=current._locator.locate) as locate:
+ assert current.read_indices([20], ["value"])["value"].to_pylist()
== [2]
+ locate.assert_called_once_with([20])
+ finally:
+ current.close()
+
+
[email protected]_plan
[email protected]_read
[email protected]("available", [False, True])
+def test_auto_native_uses_private_table_copy(frames, available):
+ original = dict(frames.raw_table.options.options.to_map())
+ with patch("pypaimon.execution.native.runtime.native_runtime_available",
return_value=available):
+ current = reader(frames)
+ try:
+ assert current._table.options.native_read_enabled() == available
+ assert frames.raw_table.options.options.to_map() == original
+ if not available:
+ with patch.object(current._locator, "locate",
wraps=current._locator.locate) as locate:
+ assert current.read_indices([30],
["value"])["value"].to_pylist() == [3]
+ locate.assert_called_once_with([30])
+ finally:
+ current.close()
+
+
+def test_unsupported_native_plan_reuses_python_index_scanner(frames):
+ from pypaimon.globalindex import DataEvolutionGlobalIndexScanner
+ frames.add(pa.table({"index": [40], "value": [4]}))
+ with patch("pypaimon.execution.native.runtime.native_runtime_available",
return_value=True):
+ current = reader(frames)
+ try:
+ assert current._table.options.native_read_enabled()
+ with
patch("pypaimon.execution.native.runtime.native_runtime_available",
return_value=True), \
+
patch("pypaimon.execution.native.read.plan.NativeTableScanMixin._native_plan_supported",
+ return_value=False), \
+ patch("pypaimon.execution.native.read.plan.native_plan",
+ side_effect=AssertionError("unsupported native plan")), \
+ patch.object(DataEvolutionGlobalIndexScanner, "create",
+ wraps=DataEvolutionGlobalIndexScanner.create) as
create:
+ for indices, expected in [([20], [2]), ([40], [4]), ([10, 99],
[1])]:
+ result = current.read_indices(indices, ["value"])
+ assert result["value"].to_pylist() == expected
+ assert create.call_count == 1
+ scanner = current._locator._scanner
+ with patch.object(scanner, "close", wraps=scanner.close) as close:
+ current.close()
+ close.assert_called_once()
+ finally:
+ current.close()
+
+
[email protected](not native_runtime_available(), reason="native reader
required")
[email protected]("denied", [False, True])
+def test_native_lookup_preserves_query_authorization(frames, denied):
+ from pypaimon.read.table_scan import TableScan
+ current = reader(frames)
+ calls = []
+
+ def authorize(*args, **kwargs):
+ calls.append(True)
+ if denied:
+ raise PermissionError("denied")
+ return None
+
+ original_init = TableScan.__init__
+
+ def init_with_auth(scan, *args, **kwargs):
+ original_init(scan, *args, **kwargs)
+ scan._query_auth_fn = authorize
+
+ try:
+ with patch.object(TableScan, "__init__", init_with_auth), \
+ patch.object(current._locator, "locate",
side_effect=AssertionError("authorization bypass")), \
+ patch("pypaimon.execution.native.read.plan.native_plan",
+ side_effect=RuntimeError("unavailable")):
+ if denied:
+ with pytest.raises(PermissionError, match="denied"):
+ current.read_indices([20], ["value"])
+ else:
+ assert current.read_indices([20],
["value"])["value"].to_pylist() == [2]
+ assert calls
+ finally:
+ current.close()
