From 67d9b6569f701c466e00a51e65d599f26721093e Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 01:56:50 -0700 Subject: [PATCH 01/16] [python] Support logical frame readers in LeRobot datasets --- docs/docs/pypaimon/lerobot.md | 22 ++ paimon-python/pypaimon/multimodal/__init__.py | 6 +- .../pypaimon/multimodal/lerobot/__init__.py | 6 +- .../pypaimon/multimodal/lerobot/dataset.py | 335 +++++++++++++++--- .../pypaimon/tests/multimodal_lerobot_test.py | 99 ++++++ 5 files changed, 410 insertions(+), 58 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index b1e487fff1f1..6671cd70c923 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -226,3 +226,25 @@ Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force either decoder. + +For logical frames assembled from multiple tables, implement one batched +reader instead of copying them into a LeRobot table group: + +```python +from pypaimon.multimodal import LeRobotFrameReader, PaimonLeRobotDataset + +class Frames(LeRobotFrameReader): + schema = logical_frame_schema + + def read_indices(self, indices, columns): + return resolve_frame_rows(indices, columns) # pyarrow.Table + +dataset = PaimonLeRobotDataset.from_reader(Frames(), metadata) +``` + +`metadata` contains `info`, `episodes`, `tasks`, and optional `stats`, +`subtasks`, `repo_id`, and `revision`. The reader must bind its source tables +to stable snapshots and resolve all requested indices in one Arrow table. +Image values are encoded bytes. The Dataset retains episode selection, delta +windows, padding, transforms, and Torch conversion. Logical video readers are +not yet supported. diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index 00d59d5c7184..e9249457b332 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -29,7 +29,10 @@ Hdf5File, Hdf5LoadResult, ) -from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset +from pypaimon.multimodal.lerobot.dataset import ( + LeRobotFrameReader, + PaimonLeRobotDataset, +) from pypaimon.multimodal.rosbag import ( RosbagLoadResult, RosbagSource, @@ -63,6 +66,7 @@ "BlobStore", "Hdf5File", "Hdf5LoadResult", + "LeRobotFrameReader", "MultimodalConnection", "MultimodalTable", "NoSuchKey", diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index 25e196dbb873..48bcf48afcff 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -17,11 +17,15 @@ """LeRobot Dataset v3 integration for multimodal Paimon tables.""" from pypaimon.multimodal.lerobot.api import load_from_lerobot -from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset +from pypaimon.multimodal.lerobot.dataset import ( + LeRobotFrameReader, + PaimonLeRobotDataset, +) from pypaimon.multimodal.lerobot.writer import PaimonLeRobotWriter __all__ = [ + "LeRobotFrameReader", "PaimonLeRobotDataset", "PaimonLeRobotWriter", "load_from_lerobot", diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 7a13731e5fcd..100a18a08015 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -24,7 +24,9 @@ import operator import os import sys +from abc import ABC, abstractmethod from collections import OrderedDict +from collections.abc import Mapping from functools import partial import pyarrow as pa @@ -75,6 +77,27 @@ }) +class LeRobotFrameReader(ABC): + """Batch reader for logical LeRobot frame rows. + + Implementations may resolve one logical row from multiple Paimon tables. + Image values must be encoded bytes. Readers used by DataLoader workers + must be picklable. + """ + + @property + @abstractmethod + def schema(self): + """Return the logical frame schema as :class:`pyarrow.Schema`.""" + + @abstractmethod + def read_indices(self, indices, columns): + """Return requested rows as a :class:`pyarrow.Table`.""" + + def close(self): + """Release reader resources.""" + + class PaimonLeRobotDataset: """Map-style LeRobot reader backed by indexed Paimon reads. @@ -98,12 +121,62 @@ def __init__( blob_parallelism=16, video_backend=None, return_uint8=False): - if sys.version_info < (3, 10): - raise RuntimeError( - "PaimonLeRobotDataset requires Python 3.10 or newer; " - "install and run 'pypaimon[lerobot]' on a supported Python " - "version.") - raw_table, self.meta = _load_dataset(table, tag_name) + _require_dataset_python() + raw_table, metadata = _load_dataset(table, tag_name) + info = self._init_dataset( + metadata, + tag_name, + episodes, + image_transforms, + delta_timestamps, + tolerance_s, + blob_parallelism, + video_backend, + return_uint8, + ) + self._init_reader(raw_table, info) + + @classmethod + def from_reader( + cls, + reader, + metadata, + *, + episodes=None, + image_transforms=None, + delta_timestamps=None, + tolerance_s=1e-4, + return_uint8=False): + """Build a dataset over a logical, possibly multi-table reader.""" + _require_dataset_python() + dataset = cls.__new__(cls) + metadata = _reader_metadata(metadata) + info = dataset._init_dataset( + metadata, + None, + episodes, + image_transforms, + delta_timestamps, + tolerance_s, + 1, + None, + return_uint8, + ) + dataset._init_external_reader(reader, info) + return dataset + + def _init_dataset( + self, + metadata, + tag_name, + episodes, + image_transforms, + delta_timestamps, + tolerance_s, + blob_parallelism, + video_backend, + return_uint8): + self.meta = metadata self.tag_name = tag_name self.repo_id = self.meta.repo_id self.image_transforms = image_transforms @@ -125,7 +198,7 @@ def __init__( info = self._init_metadata() self._init_episodes(episodes) - self._init_reader(raw_table, info) + return info def _init_metadata(self): info = dict(_metadata_member(self.meta, "info", {})) @@ -199,27 +272,8 @@ def _init_episodes(self, episodes): def _init_reader(self, raw_table, info): target_schema = _target_schema(raw_table) - table_fields = set(target_schema.names) - tasks = _metadata_member(self.meta, "tasks") - subtasks = _metadata_member(self.meta, "subtasks") - _validate_component_metadata( - self._features, self._total_tasks, tasks, subtasks) - source_schema = _schema_from_info(info) - _validate_lerobot_schema(source_schema, target_schema, self.repo_id) - validation_context = _build_frame_validation_context( - self.meta, - self._episode_ranges, - self._fps, - tasks, - subtasks, - source_schema.field("timestamp").type, - ) - projection = list(self._features) - missing = set(projection) - table_fields - if missing: - raise ValueError( - "Paimon table is missing LeRobot fields: %s" - % sorted(missing)) + projection, validation_context, subtasks = \ + self._init_frame_contract(target_schema, info, True) self._read_table, self._snapshot_id, splits = _indexed_read_table( raw_table, projection) @@ -232,6 +286,7 @@ def _init_reader(self, raw_table, info): self._projection = projection self._frame_locator = _FrameLocator( self._read_table, snapshot, splits) + self._external_reader = None self._validation_context = validation_context self._file_io = self._read_table.file_io self._video_collators = [ @@ -246,6 +301,60 @@ def _init_reader(self, raw_table, info): ) for key in self._video_keys ] + self._init_delta_projection(validation_context, subtasks) + + def _init_external_reader(self, reader, info): + if self._video_keys: + raise ValueError( + "LeRobotFrameReader does not support video features.") + schema = getattr(reader, "schema", None) + if not isinstance(schema, pa.Schema): + raise TypeError( + "LeRobotFrameReader.schema must be a pyarrow.Schema.") + projection, validation_context, subtasks = \ + self._init_frame_contract(schema, info, False) + if not callable(getattr(reader, "read_indices", None)): + raise TypeError( + "LeRobotFrameReader must define read_indices().") + self._read_table = None + self._snapshot_id = None + self._projection = projection + self._frame_locator = None + self._external_reader = reader + self._validation_context = validation_context + self._file_io = None + self._video_collators = [] + self._init_delta_projection(validation_context, subtasks) + + def _init_frame_contract(self, target_schema, info, validate_metadata): + tasks = _metadata_member(self.meta, "tasks") + subtasks = _metadata_member(self.meta, "subtasks") + _validate_component_metadata( + self._features, self._total_tasks, tasks, subtasks) + source_schema = _schema_from_info(info) + if validate_metadata: + _validate_lerobot_schema( + source_schema, target_schema, self.repo_id) + else: + _validate_reader_schema( + source_schema, target_schema, self.repo_id) + validation_context = _build_frame_validation_context( + self.meta, + self._episode_ranges, + self._fps, + tasks, + subtasks, + source_schema.field("timestamp").type, + ) + projection = list(self._features) + missing = set(projection) - set(target_schema.names) + if missing: + raise ValueError( + "LeRobot frame schema is missing fields: %s" + % sorted(missing)) + return projection, validation_context, subtasks + + def _init_delta_projection(self, validation_context, subtasks): self._task_names = validation_context["task_names"] self._subtask_names = validation_context["subtask_names"] self._delta_projection = None @@ -307,9 +416,12 @@ def __getitems__(self, indices): if position not in unique_frame_index_set }) lookup_indices = sorted(unique_frame_index_set.union(delta_indices)) - splits, needs_filter = self._frame_locator.locate(lookup_indices) - rows = self._read_rows( - lookup_indices, self._projection, splits, needs_filter) + if getattr(self, "_external_reader", None) is None: + splits, needs_filter = self._frame_locator.locate(lookup_indices) + rows = self._read_rows( + lookup_indices, self._projection, splits, needs_filter) + else: + rows = self._read_rows(lookup_indices, self._projection) base_rows = { index: rows[index] for index in unique_frame_indices } @@ -323,28 +435,36 @@ def __getitems__(self, indices): _attach_task_labels( base_rows, self._task_names, self._subtask_names) row_groups = [base_rows, delta_rows] - image_sources = _image_blob_sources( - row_groups, self._image_keys) - for attempt in range(_IMAGE_READ_ATTEMPTS): - if attempt: - _restore_image_blob_sources(image_sources) - try: - _resolve_image_blobs( - self._file_io, - row_groups, - self._image_keys, - self.blob_parallelism, - ) - _decode_image_rows( - row_groups, - self._image_keys, - self._features, - self.return_uint8, - ) - break - except OSError: - if attempt + 1 == _IMAGE_READ_ATTEMPTS: - raise + if getattr(self, "_external_reader", None) is None: + image_sources = _image_blob_sources( + row_groups, self._image_keys) + for attempt in range(_IMAGE_READ_ATTEMPTS): + if attempt: + _restore_image_blob_sources(image_sources) + try: + _resolve_image_blobs( + self._file_io, + row_groups, + self._image_keys, + self.blob_parallelism, + ) + _decode_image_rows( + row_groups, + self._image_keys, + self._features, + self.return_uint8, + ) + break + except OSError: + if attempt + 1 == _IMAGE_READ_ATTEMPTS: + raise + else: + _decode_image_rows( + row_groups, + self._image_keys, + self._features, + self.return_uint8, + ) _decode_video_rows( row_groups, getattr(self, "_video_collators", ())) @@ -382,6 +502,15 @@ def __getitems__(self, indices): def close(self): first_error = None + reader = getattr(self, "_external_reader", None) + self._external_reader = None + if reader is not None: + try: + close = getattr(reader, "close", None) + if callable(close): + close() + except Exception as error: + first_error = error locator = getattr(self, "_frame_locator", None) if locator is not None: try: @@ -407,6 +536,16 @@ def _read_rows( self, indices, projection, splits=None, needs_filter=True): if not indices: return {} + external_reader = getattr(self, "_external_reader", None) + if external_reader is not None: + return _read_external_rows( + external_reader, + projection, + indices, + self._validation_context, + self.tolerance_s, + self._features, + ) return _read_rows_by_index( self._read_table, projection, @@ -623,9 +762,39 @@ def shapes(self): } def get_task_index(self, task): - if task not in self.tasks.index: + if hasattr(self.tasks, "loc"): + if task not in self.tasks.index: + return None + return int(self.tasks.loc[task].task_index) + try: + return list(self.tasks).index(task) + except ValueError: return None - return int(self.tasks.loc[task].task_index) + + +def _reader_metadata(metadata): + info = _metadata_member(metadata, "info") + if not isinstance(info, Mapping): + raise TypeError("LeRobot reader metadata must contain an info map.") + info = dict(info) + features = info.get("features") + if isinstance(features, Mapping): + info["features"] = { + name: dict(feature) for name, feature in features.items() + } + for feature in info["features"].values(): + if "shape" in feature: + feature["shape"] = tuple(feature["shape"]) + stats = _metadata_member(metadata, "stats") + return _PaimonLeRobotMetadata( + str(_metadata_member(metadata, "repo_id", "logical-reader")), + _metadata_member(metadata, "revision"), + info, + _numpy_stats(stats) if stats is not None else None, + _metadata_member(metadata, "episodes"), + _metadata_member(metadata, "tasks"), + _metadata_member(metadata, "subtasks"), + ) def _load_dataset(table, tag_name): @@ -729,7 +898,8 @@ def _numpy_stats(value): def _metadata_member(metadata, name, default=None): - value = getattr(metadata, name, None) + value = metadata.get(name) if isinstance(metadata, Mapping) \ + else getattr(metadata, name, None) return default if value is None else value @@ -935,6 +1105,51 @@ def _index_predicate(table, indices): "index", indices) +def _validate_reader_schema(source_schema, reader_schema, source): + for source_field in source_schema: + target_index = reader_schema.get_field_index(source_field.name) + if target_index < 0: + continue + target_type = reader_schema.field(target_index).type + if source_field.type != target_type: + raise ValueError( + "LeRobot feature %s from %s expects %s, found %s." + % (source_field.name, source, source_field.type, + target_type)) + + +def _read_external_rows( + reader, projection, indices, validation_context, tolerance_s, + features): + values = reader.read_indices(tuple(indices), tuple(projection)) + if not isinstance(values, pa.Table): + raise TypeError( + "LeRobotFrameReader.read_indices() must return a pyarrow.Table.") + missing = set(projection) - set(values.column_names) + if missing: + raise ValueError( + "LeRobotFrameReader result is missing fields: %s" + % sorted(missing)) + rows = _arrow_rows(values.select(projection), features) + expected = set(indices) + result = {} + for row in rows: + index = _control_index(row, "index", -1) + if index not in expected or index in result: + raise ValueError( + "LeRobotFrameReader returned an unexpected or duplicate " + "index: %d." % index) + _validate_control_row( + index, row, validation_context, tolerance_s) + result[index] = row + missing = expected - set(result) + if missing: + raise RuntimeError( + "LeRobotFrameReader did not return indices %s." + % sorted(missing)) + return result + + def _read_rows_by_index( table, projection, indices, validation_context, tolerance_s, features, splits=None, needs_filter=True): @@ -1428,6 +1643,14 @@ def _normalize_index(index, size): return index +def _require_dataset_python(): + if sys.version_info < (3, 10): + raise RuntimeError( + "PaimonLeRobotDataset requires Python 3.10 or newer; " + "install and run 'pypaimon[lerobot]' on a supported Python " + "version.") + + def _positive_int(value, name): try: value = operator.index(value) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 97b180d39e65..c999346b8897 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -325,6 +325,105 @@ def test_dataset_requires_supported_python(self): pmm.PaimonLeRobotDataset(Mock()) load.assert_not_called() + def test_dataset_reads_one_batch_from_logical_frame_reader(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + + info = { + "codebase_version": "v3.0", + "total_frames": 3, + "total_episodes": 1, + "total_tasks": 1, + "fps": 10, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "frame_index": {"dtype": "int64", "shape": [1]}, + "timestamp": {"dtype": "float32", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + "observation.state": {"dtype": "float32", "shape": [2]}, + "action": {"dtype": "float32", "shape": [1]}, + "camera.image": {"dtype": "image", "shape": [2, 2, 3]}, + }, + } + rows = { + index: { + "index": index, + "episode_index": 0, + "frame_index": index, + "timestamp": index / 10, + "task_index": 0, + "observation.state": [index, index + 1], + "action": float(index), + "camera.image": _image_bytes( + np.full((2, 2, 3), index, dtype=np.uint8), None), + } + for index in range(3) + } + + class Reader(pmm.LeRobotFrameReader): + + def __init__(self): + self.calls = [] + self.closed = False + + @property + def schema(self): + return _schema_from_info(info) + + def read_indices(self, indices, columns): + self.calls.append((indices, columns)) + return pa.Table.from_pylist([ + {name: rows[index][name] for name in columns} + for index in indices + ], schema=self.schema) + + def close(self): + self.closed = True + + reader = Reader() + metadata = { + "repo_id": "logical/multi-table", + "revision": "dataset-version-12", + "info": info, + "episodes": [{ + "episode_index": 0, + "dataset_from_index": 0, + "dataset_to_index": 3, + "length": 3, + "tasks": ["pick"], + }], + "tasks": ["pick"], + "stats": {"action": {"mean": [1.0]}}, + } + dataset = pmm.PaimonLeRobotDataset.from_reader( + reader, + metadata, + delta_timestamps={"action": [-0.1, 0.0, 0.1]}, + ) + + sample, _ = dataset.__getitems__([1, 2]) + + self.assertEqual([((0, 1, 2), tuple(info["features"]))], + reader.calls) + self.assertIsNone(dataset.tag_name) + self.assertEqual("dataset-version-12", dataset.meta.revision) + self.assertEqual((2,), dataset.features["observation.state"]["shape"]) + self.assertEqual([1.0], dataset.meta.stats["action"]["mean"].tolist()) + self.assertEqual(0, dataset.meta.get_task_index("pick")) + self.assertEqual("pick", sample["task"]) + torch.testing.assert_close( + sample["observation.state"], torch.tensor([1.0, 2.0])) + torch.testing.assert_close( + sample["action"], torch.tensor([0.0, 1.0, 2.0])) + self.assertEqual([3, 2, 2], list(sample["camera.image"].shape)) + self.assertEqual([False, False, False], + sample["action_is_pad"].tolist()) + dataset.close() + self.assertTrue(reader.closed) + def test_metadata_json_preserves_nested_values(self): values = { "name": "机器人", From 1a74bdad000d88d7ce9ace52a0d674bdfb61290c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 02:07:08 -0700 Subject: [PATCH 02/16] [python] Clarify logical LeRobot frame readers --- docs/docs/pypaimon/lerobot.md | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 6671cd70c923..fb5c8a6ab110 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -227,8 +227,7 @@ on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force either decoder. -For logical frames assembled from multiple tables, implement one batched -reader instead of copying them into a LeRobot table group: +Use a frame reader when one logical frame is assembled from multiple tables: ```python from pypaimon.multimodal import LeRobotFrameReader, PaimonLeRobotDataset @@ -242,9 +241,7 @@ class Frames(LeRobotFrameReader): dataset = PaimonLeRobotDataset.from_reader(Frames(), metadata) ``` -`metadata` contains `info`, `episodes`, `tasks`, and optional `stats`, -`subtasks`, `repo_id`, and `revision`. The reader must bind its source tables -to stable snapshots and resolve all requested indices in one Arrow table. -Image values are encoded bytes. The Dataset retains episode selection, delta -windows, padding, transforms, and Torch conversion. Logical video readers are -not yet supported. +`LeRobotFrameReader` is a logical row contract, not a physical `frames` table. +It may query any tables. `metadata` separately supplies `info`, `episodes`, +`tasks`, and optional `stats` and `subtasks`. Logical video readers are not yet +supported. From e7f484257a44977170a94889be4736f8c3c3d30d Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 02:13:17 -0700 Subject: [PATCH 03/16] [python] Clarify multi-table frame reader example --- docs/docs/pypaimon/lerobot.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index fb5c8a6ab110..761749a056e3 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -232,13 +232,13 @@ Use a frame reader when one logical frame is assembled from multiple tables: ```python from pypaimon.multimodal import LeRobotFrameReader, PaimonLeRobotDataset -class Frames(LeRobotFrameReader): +class MultiTableFrameReader(LeRobotFrameReader): schema = logical_frame_schema def read_indices(self, indices, columns): return resolve_frame_rows(indices, columns) # pyarrow.Table -dataset = PaimonLeRobotDataset.from_reader(Frames(), metadata) +dataset = PaimonLeRobotDataset.from_reader(MultiTableFrameReader(), metadata) ``` `LeRobotFrameReader` is a logical row contract, not a physical `frames` table. From 40fa3ffb17f0fff48e7dba1e091f6b8207e4baed Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 02:13:57 -0700 Subject: [PATCH 04/16] [python] Simplify custom frame reader example --- docs/docs/pypaimon/lerobot.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 761749a056e3..821ba1d26c65 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -232,13 +232,13 @@ Use a frame reader when one logical frame is assembled from multiple tables: ```python from pypaimon.multimodal import LeRobotFrameReader, PaimonLeRobotDataset -class MultiTableFrameReader(LeRobotFrameReader): +class CustomFrameReader(LeRobotFrameReader): schema = logical_frame_schema def read_indices(self, indices, columns): return resolve_frame_rows(indices, columns) # pyarrow.Table -dataset = PaimonLeRobotDataset.from_reader(MultiTableFrameReader(), metadata) +dataset = PaimonLeRobotDataset.from_reader(CustomFrameReader(), metadata) ``` `LeRobotFrameReader` is a logical row contract, not a physical `frames` table. From afa9c388eb060a494efab61a21301ca5e9a265fd Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 02:36:32 -0700 Subject: [PATCH 05/16] [python] Unify LeRobot dataset readers --- docs/docs/pypaimon/lerobot.md | 9 +- paimon-python/pypaimon/multimodal/__init__.py | 6 +- .../pypaimon/multimodal/lerobot/__init__.py | 6 +- .../pypaimon/multimodal/lerobot/dataset.py | 368 ++++-------------- .../pypaimon/multimodal/lerobot/reader.py | 218 +++++++++++ .../pypaimon/tests/multimodal_lerobot_test.py | 7 +- 6 files changed, 299 insertions(+), 315 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/lerobot/reader.py diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 821ba1d26c65..4041069aa29a 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -238,10 +238,13 @@ class CustomFrameReader(LeRobotFrameReader): def read_indices(self, indices, columns): return resolve_frame_rows(indices, columns) # pyarrow.Table -dataset = PaimonLeRobotDataset.from_reader(CustomFrameReader(), metadata) +dataset = PaimonLeRobotDataset( + reader=CustomFrameReader(), + metadata=metadata, +) ``` `LeRobotFrameReader` is a logical row contract, not a physical `frames` table. It may query any tables. `metadata` separately supplies `info`, `episodes`, -`tasks`, and optional `stats` and `subtasks`. Logical video readers are not yet -supported. +`tasks`, and optional `stats` and `subtasks`. Descriptor-backed media readers +also expose `file_io`. diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index e9249457b332..e0fb41749e0c 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -29,10 +29,8 @@ Hdf5File, Hdf5LoadResult, ) -from pypaimon.multimodal.lerobot.dataset import ( - LeRobotFrameReader, - PaimonLeRobotDataset, -) +from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset +from pypaimon.multimodal.lerobot.reader import LeRobotFrameReader from pypaimon.multimodal.rosbag import ( RosbagLoadResult, RosbagSource, diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index 48bcf48afcff..363eb86c0533 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -17,10 +17,8 @@ """LeRobot Dataset v3 integration for multimodal Paimon tables.""" from pypaimon.multimodal.lerobot.api import load_from_lerobot -from pypaimon.multimodal.lerobot.dataset import ( - LeRobotFrameReader, - PaimonLeRobotDataset, -) +from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset +from pypaimon.multimodal.lerobot.reader import LeRobotFrameReader from pypaimon.multimodal.lerobot.writer import PaimonLeRobotWriter diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 100a18a08015..e82a0b9b3487 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -22,16 +22,13 @@ import json import math import operator -import os import sys -from abc import ABC, abstractmethod from collections import OrderedDict from collections.abc import Mapping from functools import partial import pyarrow as pa -from pypaimon.common.options.core_options import CoreOptions from pypaimon.multimodal.lerobot.metadata import ( _companion_table_identifiers, _restore_pandas_metadata, @@ -39,6 +36,10 @@ _validate_tag_name, ) from pypaimon.multimodal.lerobot.loader import _DECLARED_NUMERIC_RANGES +from pypaimon.multimodal.lerobot.reader import ( + LeRobotFrameReader, + _PaimonTableFrameReader, +) from pypaimon.multimodal.lerobot.schema import ( _feature_shape, _require_v3, @@ -47,7 +48,6 @@ ) from pypaimon.multimodal.table import _target_schema, _time_travel_table from pypaimon.multimodal.video import VideoFrameCollator -from pypaimon.read.query_auth_split import QueryAuthSplit _TORCH_DTYPE_NAMES = { @@ -77,27 +77,6 @@ }) -class LeRobotFrameReader(ABC): - """Batch reader for logical LeRobot frame rows. - - Implementations may resolve one logical row from multiple Paimon tables. - Image values must be encoded bytes. Readers used by DataLoader workers - must be picklable. - """ - - @property - @abstractmethod - def schema(self): - """Return the logical frame schema as :class:`pyarrow.Schema`.""" - - @abstractmethod - def read_indices(self, indices, columns): - """Return requested rows as a :class:`pyarrow.Table`.""" - - def close(self): - """Release reader resources.""" - - class PaimonLeRobotDataset: """Map-style LeRobot reader backed by indexed Paimon reads. @@ -111,8 +90,10 @@ class PaimonLeRobotDataset: def __init__( self, - table, + table=None, *, + reader=None, + metadata=None, tag_name=None, episodes=None, image_transforms=None, @@ -122,7 +103,20 @@ def __init__( video_backend=None, return_uint8=False): _require_dataset_python() - raw_table, metadata = _load_dataset(table, tag_name) + if (table is None) == (reader is None): + raise ValueError("Provide exactly one of table or reader.") + if reader is not None: + if metadata is None: + raise ValueError("metadata is required with reader.") + if tag_name is not None: + raise ValueError( + "tag_name is managed by a custom reader and must be None.") + metadata = _reader_metadata(metadata) + raw_table = None + else: + if metadata is not None: + raise ValueError("metadata is only accepted with reader.") + raw_table, metadata = _load_dataset(table, tag_name) info = self._init_dataset( metadata, tag_name, @@ -134,36 +128,10 @@ def __init__( video_backend, return_uint8, ) - self._init_reader(raw_table, info) - - @classmethod - def from_reader( - cls, - reader, - metadata, - *, - episodes=None, - image_transforms=None, - delta_timestamps=None, - tolerance_s=1e-4, - return_uint8=False): - """Build a dataset over a logical, possibly multi-table reader.""" - _require_dataset_python() - dataset = cls.__new__(cls) - metadata = _reader_metadata(metadata) - info = dataset._init_dataset( - metadata, - None, - episodes, - image_transforms, - delta_timestamps, - tolerance_s, - 1, - None, - return_uint8, - ) - dataset._init_external_reader(reader, info) - return dataset + if reader is None: + self._init_paimon_reader(raw_table, info) + else: + self._init_reader(reader, info, False) def _init_dataset( self, @@ -270,28 +238,42 @@ def _init_episodes(self, episodes): if self._delta_indices and self._episode_ranges is None: raise ValueError("delta_timestamps requires episode metadata.") - def _init_reader(self, raw_table, info): + def _init_paimon_reader(self, raw_table, info): target_schema = _target_schema(raw_table) projection, validation_context, subtasks = \ self._init_frame_contract(target_schema, info, True) - - self._read_table, self._snapshot_id, splits = _indexed_read_table( - raw_table, projection) - snapshot = self._read_table.snapshot_manager().get_snapshot_by_id( - self._snapshot_id) - if snapshot.next_row_id != self._total_frames: - raise ValueError( - "Paimon table has %d rows but metadata declares %d frames." - % (snapshot.next_row_id, self._total_frames)) + reader = _PaimonTableFrameReader( + raw_table, projection, self._total_frames) + self._set_reader( + reader, projection, validation_context, subtasks) + + def _init_reader(self, reader, info, validate_metadata): + if not isinstance(reader, LeRobotFrameReader): + raise TypeError("reader must be a LeRobotFrameReader.") + schema = getattr(reader, "schema", None) + if not isinstance(schema, pa.Schema): + raise TypeError( + "LeRobotFrameReader.schema must be a pyarrow.Schema.") + projection, validation_context, subtasks = \ + self._init_frame_contract(schema, info, validate_metadata) + self._set_reader( + reader, projection, validation_context, subtasks) + + def _set_reader( + self, reader, projection, validation_context, subtasks): + self._reader = reader + self._snapshot_id = getattr(reader, "snapshot_id", None) + self._read_table = getattr(reader, "_table", None) + self._frame_locator = getattr(reader, "_locator", None) self._projection = projection - self._frame_locator = _FrameLocator( - self._read_table, snapshot, splits) - self._external_reader = None self._validation_context = validation_context - self._file_io = self._read_table.file_io + self._file_io = getattr(reader, "file_io", None) + if self._video_keys and self._file_io is None: + raise ValueError( + "A video-backed LeRobotFrameReader must expose file_io.") self._video_collators = [ VideoFrameCollator( - self._read_table, + reader, video_column=key, decoder_factory=partial( _open_video_decoder, backend=self.video_backend), @@ -303,29 +285,6 @@ def _init_reader(self, raw_table, info): ] self._init_delta_projection(validation_context, subtasks) - def _init_external_reader(self, reader, info): - if self._video_keys: - raise ValueError( - "LeRobotFrameReader does not support video features.") - schema = getattr(reader, "schema", None) - if not isinstance(schema, pa.Schema): - raise TypeError( - "LeRobotFrameReader.schema must be a pyarrow.Schema.") - projection, validation_context, subtasks = \ - self._init_frame_contract(schema, info, False) - if not callable(getattr(reader, "read_indices", None)): - raise TypeError( - "LeRobotFrameReader must define read_indices().") - self._read_table = None - self._snapshot_id = None - self._projection = projection - self._frame_locator = None - self._external_reader = reader - self._validation_context = validation_context - self._file_io = None - self._video_collators = [] - self._init_delta_projection(validation_context, subtasks) - def _init_frame_contract(self, target_schema, info, validate_metadata): tasks = _metadata_member(self.meta, "tasks") subtasks = _metadata_member(self.meta, "subtasks") @@ -416,12 +375,7 @@ def __getitems__(self, indices): if position not in unique_frame_index_set }) lookup_indices = sorted(unique_frame_index_set.union(delta_indices)) - if getattr(self, "_external_reader", None) is None: - splits, needs_filter = self._frame_locator.locate(lookup_indices) - rows = self._read_rows( - lookup_indices, self._projection, splits, needs_filter) - else: - rows = self._read_rows(lookup_indices, self._projection) + rows = self._read_rows(lookup_indices, self._projection) base_rows = { index: rows[index] for index in unique_frame_indices } @@ -435,7 +389,7 @@ def __getitems__(self, indices): _attach_task_labels( base_rows, self._task_names, self._subtask_names) row_groups = [base_rows, delta_rows] - if getattr(self, "_external_reader", None) is None: + if self._file_io is not None: image_sources = _image_blob_sources( row_groups, self._image_keys) for attempt in range(_IMAGE_READ_ATTEMPTS): @@ -502,27 +456,20 @@ def __getitems__(self, indices): def close(self): first_error = None - reader = getattr(self, "_external_reader", None) - self._external_reader = None - if reader is not None: - try: - close = getattr(reader, "close", None) - if callable(close): - close() - except Exception as error: - first_error = error - locator = getattr(self, "_frame_locator", None) - if locator is not None: - try: - locator.close() - except Exception as error: - first_error = error for collator in getattr(self, "_video_collators", ()): try: collator.close() except Exception as error: if first_error is None: first_error = error + reader = getattr(self, "_reader", None) + self._reader = None + if reader is not None: + try: + reader.close() + except Exception as error: + if first_error is None: + first_error = error if first_error is not None: raise first_error @@ -532,29 +479,16 @@ def __del__(self): except Exception: pass - def _read_rows( - self, indices, projection, splits=None, needs_filter=True): + def _read_rows(self, indices, projection): if not indices: return {} - external_reader = getattr(self, "_external_reader", None) - if external_reader is not None: - return _read_external_rows( - external_reader, - projection, - indices, - self._validation_context, - self.tolerance_s, - self._features, - ) - return _read_rows_by_index( - self._read_table, + return _read_reader_rows( + self._reader, projection, indices, self._validation_context, self.tolerance_s, self._features, - splits, - needs_filter, ) def set_image_transforms(self, image_transforms): @@ -597,115 +531,6 @@ def __repr__(self): self.num_frames, list(self.features))) -class _FrameLocator: - """Locate LeRobot frame rows in one fixed Paimon snapshot.""" - - def __init__(self, table, snapshot, splits): - self._table = table - self._snapshot = snapshot - self._scanner = None - self._scanner_initialized = False - self._process_id = os.getpid() - self._set_splits(splits) - - def _set_splits(self, splits): - from pypaimon.read.datasource.torch_dataset import ( - SplitRangeIndex, - row_ranges_for_split, - ) - - self._splits = splits - self._split_ranges = [ - row_ranges_for_split(split) for split in splits - ] - self._split_range_index = SplitRangeIndex(self._split_ranges) - - def locate(self, indices): - """Return narrowed splits and whether rows still need filtering.""" - self._ensure_process() - predicate = _index_predicate(self._table, indices) - try: - scanner = self._index_scanner(predicate) - except Exception as error: - raise RuntimeError( - "Failed to open the Paimon global index for LeRobot frame " - "lookups.") from error - if scanner is None: - raise RuntimeError( - "PaimonLeRobotDataset requires a readable global index on " - "the frame 'index' column.") - try: - evaluation = scanner.scan_with_coverage(predicate) - if evaluation is None: - raise RuntimeError( - "The Paimon global index could not evaluate the LeRobot " - "frame index predicate.") - unindexed = scanner.unindexed_ranges( - predicate, - search_mode=self._table.options.scalar_index_search_mode(), - contributing_field_ids=evaluation.contributing_field_ids, - ) - ranges = evaluation.result.results().to_range_list() + unindexed - from pypaimon.read.datasource.torch_dataset import ( - select_indexed_splits, - ) - from pypaimon.utils.range import Range - return select_indexed_splits( - self._splits, - self._split_ranges, - self._split_range_index, - Range.sort_and_merge_overlap(ranges, True), - ), bool(unindexed) - except RuntimeError: - raise - except Exception as error: - raise RuntimeError( - "Failed to query the Paimon global index for LeRobot " - "frames.") from error - - def _ensure_process(self): - process_id = os.getpid() - if process_id == self._process_id: - return - self._scanner = None - self._scanner_initialized = False - self._set_splits(self._splits) - self._process_id = process_id - - def _index_scanner(self, predicate): - if not self._scanner_initialized: - from pypaimon.globalindex import DataEvolutionGlobalIndexScanner - self._scanner = DataEvolutionGlobalIndexScanner.create( - self._table, - predicate=predicate, - snapshot=self._snapshot, - ) - self._scanner_initialized = True - return self._scanner - - def close(self): - scanner = self._scanner - self._scanner = None - self._scanner_initialized = False - if scanner is not None and self._process_id == os.getpid(): - scanner.close() - - def __getstate__(self): - state = self.__dict__.copy() - state["_scanner"] = None - state["_scanner_initialized"] = False - state["_process_id"] = None - state["_split_ranges"] = None - state["_split_range_index"] = None - return state - - def __del__(self): - try: - self.close() - except Exception: - pass - - class _PaimonLeRobotMetadata: def __init__( @@ -1075,36 +900,6 @@ def _delta_indices(delta_timestamps, fps, tolerance_s, features): return result -def _indexed_read_table(raw_table, projection): - read_table = raw_table.copy({ - CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true" - }) - plan = read_table.new_read_builder().with_projection( - projection).new_scan().plan() - splits = plan.splits() - if any( - isinstance(split, QueryAuthSplit) - and ( - getattr(split.auth_result, "filter", None) - or getattr(split.auth_result, "column_masking", None) - ) - for split in splits): - raise ValueError( - "PaimonLeRobotDataset does not support query authorization " - "filters or column masking.") - if plan.snapshot_id is None: - raise ValueError("Paimon LeRobot frames table has no snapshot.") - if read_table.options.scan_tag_name() is None: - read_table = _time_travel_table( - read_table, snapshot_id=plan.snapshot_id) - return read_table, plan.snapshot_id, splits - - -def _index_predicate(table, indices): - return table.new_read_builder().new_predicate_builder().is_in( - "index", indices) - - def _validate_reader_schema(source_schema, reader_schema, source): for source_field in source_schema: target_index = reader_schema.get_field_index(source_field.name) @@ -1118,7 +913,7 @@ def _validate_reader_schema(source_schema, reader_schema, source): target_type)) -def _read_external_rows( +def _read_reader_rows( reader, projection, indices, validation_context, tolerance_s, features): values = reader.read_indices(tuple(indices), tuple(projection)) @@ -1150,33 +945,6 @@ def _read_external_rows( return result -def _read_rows_by_index( - table, projection, indices, validation_context, tolerance_s, features, - splits=None, needs_filter=True): - builder = table.new_read_builder().with_projection(projection) - if needs_filter: - builder = builder.with_filter(_index_predicate(table, indices)) - if splits is None: - splits = builder.new_scan().plan().splits() - rows = _arrow_rows(builder.new_read().to_arrow(splits), features) - expected = set(indices) - result = {} - for row in rows: - index = _control_index(row, "index", -1) - if index not in expected or index in result: - raise ValueError( - "Paimon BTree returned an unexpected or duplicate LeRobot " - "index: %d." % index) - _validate_control_row(index, row, validation_context, tolerance_s) - result[index] = row - missing = expected - set(result) - if missing: - raise RuntimeError( - "Paimon index lookup did not return LeRobot indices %s." - % sorted(missing)) - return result - - def _arrow_rows(table, features): """Convert indexed Arrow results without expanding tensors to lists.""" rows = [{} for unused in range(table.num_rows)] diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py new file mode 100644 index 000000000000..e2eac62b8ef7 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -0,0 +1,218 @@ +# 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. + +"""Frame readers used by :class:`PaimonLeRobotDataset`.""" + +import os +from abc import ABC, abstractmethod + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.table import _target_schema, _time_travel_table +from pypaimon.read.query_auth_split import QueryAuthSplit + + +class LeRobotFrameReader(ABC): + """Batch reader for logical LeRobot frame rows. + + Implementations may resolve one logical row from multiple Paimon tables. + Image values must be encoded bytes or BLOB descriptors. Video values must + be serialized ``VideoFrameDescriptor`` instances. Descriptor-backed + readers expose their resolved ``file_io``. + """ + + @property + @abstractmethod + def schema(self): + """Return the logical frame schema as :class:`pyarrow.Schema`.""" + + @abstractmethod + def read_indices(self, indices, columns): + """Return requested rows as a :class:`pyarrow.Table`.""" + + def close(self): + """Release reader resources.""" + + +class _PaimonTableFrameReader(LeRobotFrameReader): + """Default frame reader for one indexed Paimon frames table.""" + + def __init__(self, raw_table, projection, total_frames): + self._schema = _target_schema(raw_table) + self._table, self.snapshot_id, splits = _indexed_read_table( + raw_table, projection) + snapshot = self._table.snapshot_manager().get_snapshot_by_id( + self.snapshot_id) + if snapshot.next_row_id != total_frames: + raise ValueError( + "Paimon table has %d rows but metadata declares %d frames." + % (snapshot.next_row_id, total_frames)) + self.file_io = self._table.file_io + self._locator = _FrameLocator(self._table, snapshot, splits) + + @property + def schema(self): + return self._schema + + def read_indices(self, indices, columns): + splits, needs_filter = self._locator.locate(indices) + builder = self._table.new_read_builder().with_projection(columns) + if needs_filter: + builder = builder.with_filter( + _index_predicate(self._table, indices)) + return builder.new_read().to_arrow(splits) + + def close(self): + self._locator.close() + + +class _FrameLocator: + """Locate LeRobot frame rows in one fixed Paimon snapshot.""" + + def __init__(self, table, snapshot, splits): + self._table = table + self._snapshot = snapshot + self._scanner = None + self._scanner_initialized = False + self._process_id = os.getpid() + self._set_splits(splits) + + def _set_splits(self, splits): + from pypaimon.read.datasource.torch_dataset import ( + SplitRangeIndex, + row_ranges_for_split, + ) + + self._splits = splits + self._split_ranges = [ + row_ranges_for_split(split) for split in splits + ] + self._split_range_index = SplitRangeIndex(self._split_ranges) + + def locate(self, indices): + """Return narrowed splits and whether rows still need filtering.""" + self._ensure_process() + predicate = _index_predicate(self._table, indices) + try: + scanner = self._index_scanner(predicate) + except Exception as error: + raise RuntimeError( + "Failed to open the Paimon global index for LeRobot frame " + "lookups.") from error + if scanner is None: + raise RuntimeError( + "PaimonLeRobotDataset requires a readable global index on " + "the frame 'index' column.") + try: + evaluation = scanner.scan_with_coverage(predicate) + if evaluation is None: + raise RuntimeError( + "The Paimon global index could not evaluate the LeRobot " + "frame index predicate.") + unindexed = scanner.unindexed_ranges( + predicate, + search_mode=self._table.options.scalar_index_search_mode(), + contributing_field_ids=evaluation.contributing_field_ids, + ) + ranges = evaluation.result.results().to_range_list() + unindexed + from pypaimon.read.datasource.torch_dataset import ( + select_indexed_splits, + ) + from pypaimon.utils.range import Range + return select_indexed_splits( + self._splits, + self._split_ranges, + self._split_range_index, + Range.sort_and_merge_overlap(ranges, True), + ), bool(unindexed) + except RuntimeError: + raise + except Exception as error: + raise RuntimeError( + "Failed to query the Paimon global index for LeRobot " + "frames.") from error + + def _ensure_process(self): + process_id = os.getpid() + if process_id == self._process_id: + return + self._scanner = None + self._scanner_initialized = False + self._set_splits(self._splits) + self._process_id = process_id + + def _index_scanner(self, predicate): + if not self._scanner_initialized: + from pypaimon.globalindex import DataEvolutionGlobalIndexScanner + self._scanner = DataEvolutionGlobalIndexScanner.create( + self._table, + predicate=predicate, + snapshot=self._snapshot, + ) + self._scanner_initialized = True + return self._scanner + + def close(self): + scanner = self._scanner + self._scanner = None + self._scanner_initialized = False + if scanner is not None and self._process_id == os.getpid(): + scanner.close() + + def __getstate__(self): + state = self.__dict__.copy() + state["_scanner"] = None + state["_scanner_initialized"] = False + state["_process_id"] = None + state["_split_ranges"] = None + state["_split_range_index"] = None + return state + + def __del__(self): + try: + self.close() + except Exception: + pass + + +def _indexed_read_table(raw_table, projection): + read_table = raw_table.copy({ + CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true" + }) + plan = read_table.new_read_builder().with_projection( + projection).new_scan().plan() + splits = plan.splits() + if any( + isinstance(split, QueryAuthSplit) + and ( + getattr(split.auth_result, "filter", None) + or getattr(split.auth_result, "column_masking", None) + ) + for split in splits): + raise ValueError( + "PaimonLeRobotDataset does not support query authorization " + "filters or column masking.") + if plan.snapshot_id is None: + raise ValueError("Paimon LeRobot frames table has no snapshot.") + if read_table.options.scan_tag_name() is None: + read_table = _time_travel_table( + read_table, snapshot_id=plan.snapshot_id) + return read_table, plan.snapshot_id, splits + + +def _index_predicate(table, indices): + return table.new_read_builder().new_predicate_builder().is_in( + "index", indices) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index c999346b8897..5abefd92a559 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -398,9 +398,9 @@ def close(self): "tasks": ["pick"], "stats": {"action": {"mean": [1.0]}}, } - dataset = pmm.PaimonLeRobotDataset.from_reader( - reader, - metadata, + dataset = pmm.PaimonLeRobotDataset( + reader=reader, + metadata=metadata, delta_timestamps={"action": [-0.1, 0.0, 0.1]}, ) @@ -3046,7 +3046,6 @@ def counted_plan(scan): self.assertIs(scanner, dataset._frame_locator._scanner) self.assertEqual(1, read.call_count) self.assertEqual([0, 1, 3, 4], read.call_args.args[0]) - self.assertFalse(read.call_args.args[3]) self.assertEqual(1, fetch.call_count) self.assertEqual(3, fetch.call_args.args[3]) self.assertEqual("place", last["task"]) From 0882fa6ce989b758bf054bc1f26f67b9254d148b Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 02:48:13 -0700 Subject: [PATCH 06/16] [python] Align custom LeRobot source construction --- docs/docs/pypaimon/lerobot.md | 18 +-- paimon-python/pypaimon/multimodal/__init__.py | 4 +- .../pypaimon/multimodal/lerobot/__init__.py | 4 +- .../pypaimon/multimodal/lerobot/dataset.py | 124 ++++++++---------- .../lerobot/{reader.py => dataset_source.py} | 24 +++- .../pypaimon/tests/multimodal_lerobot_test.py | 20 +-- 6 files changed, 98 insertions(+), 96 deletions(-) rename paimon-python/pypaimon/multimodal/lerobot/{reader.py => dataset_source.py} (92%) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 4041069aa29a..d7512c63f991 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -227,24 +227,20 @@ on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force either decoder. -Use a frame reader when one logical frame is assembled from multiple tables: +Use a dataset source when one logical frame is assembled from multiple tables: ```python -from pypaimon.multimodal import LeRobotFrameReader, PaimonLeRobotDataset +from pypaimon.multimodal import LeRobotDatasetSource, PaimonLeRobotDataset -class CustomFrameReader(LeRobotFrameReader): +class CustomDatasetSource(LeRobotDatasetSource): + metadata = dataset_metadata schema = logical_frame_schema def read_indices(self, indices, columns): return resolve_frame_rows(indices, columns) # pyarrow.Table -dataset = PaimonLeRobotDataset( - reader=CustomFrameReader(), - metadata=metadata, -) +dataset = PaimonLeRobotDataset(CustomDatasetSource()) ``` -`LeRobotFrameReader` is a logical row contract, not a physical `frames` table. -It may query any tables. `metadata` separately supplies `info`, `episodes`, -`tasks`, and optional `stats` and `subtasks`. Descriptor-backed media readers -also expose `file_io`. +`LeRobotDatasetSource` supplies metadata and logical frame rows. It may query +any tables. Descriptor-backed media sources also expose `file_io`. diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index e0fb41749e0c..069e237bb6da 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -30,7 +30,7 @@ Hdf5LoadResult, ) from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset -from pypaimon.multimodal.lerobot.reader import LeRobotFrameReader +from pypaimon.multimodal.lerobot.dataset_source import LeRobotDatasetSource from pypaimon.multimodal.rosbag import ( RosbagLoadResult, RosbagSource, @@ -64,7 +64,7 @@ "BlobStore", "Hdf5File", "Hdf5LoadResult", - "LeRobotFrameReader", + "LeRobotDatasetSource", "MultimodalConnection", "MultimodalTable", "NoSuchKey", diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index 363eb86c0533..dde2858e1bf0 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -18,12 +18,12 @@ from pypaimon.multimodal.lerobot.api import load_from_lerobot from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset -from pypaimon.multimodal.lerobot.reader import LeRobotFrameReader +from pypaimon.multimodal.lerobot.dataset_source import LeRobotDatasetSource from pypaimon.multimodal.lerobot.writer import PaimonLeRobotWriter __all__ = [ - "LeRobotFrameReader", + "LeRobotDatasetSource", "PaimonLeRobotDataset", "PaimonLeRobotWriter", "load_from_lerobot", diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index e82a0b9b3487..b1f34f007044 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -36,9 +36,9 @@ _validate_tag_name, ) from pypaimon.multimodal.lerobot.loader import _DECLARED_NUMERIC_RANGES -from pypaimon.multimodal.lerobot.reader import ( - LeRobotFrameReader, - _PaimonTableFrameReader, +from pypaimon.multimodal.lerobot.dataset_source import ( + LeRobotDatasetSource, + _PaimonTableDatasetSource, ) from pypaimon.multimodal.lerobot.schema import ( _feature_shape, @@ -80,8 +80,8 @@ class PaimonLeRobotDataset: """Map-style LeRobot reader backed by indexed Paimon reads. - LeRobot metadata is resolved from the Paimon table group and remains - available through :attr:`meta`. + The input is a Paimon table group or :class:`LeRobotDatasetSource`. + Resolved LeRobot metadata remains available through :attr:`meta`. Set ``return_uint8=True`` to keep 8-bit visual frames in their decoded ``torch.uint8`` representation instead of normalizing them to float32. @@ -90,10 +90,8 @@ class PaimonLeRobotDataset: def __init__( self, - table=None, + table, *, - reader=None, - metadata=None, tag_name=None, episodes=None, image_transforms=None, @@ -103,19 +101,15 @@ def __init__( video_backend=None, return_uint8=False): _require_dataset_python() - if (table is None) == (reader is None): - raise ValueError("Provide exactly one of table or reader.") - if reader is not None: - if metadata is None: - raise ValueError("metadata is required with reader.") + if isinstance(table, LeRobotDatasetSource): if tag_name is not None: raise ValueError( - "tag_name is managed by a custom reader and must be None.") - metadata = _reader_metadata(metadata) + "tag_name is managed by a custom source and must be None.") + source = table + metadata = _source_metadata(source.metadata) raw_table = None else: - if metadata is not None: - raise ValueError("metadata is only accepted with reader.") + source = None raw_table, metadata = _load_dataset(table, tag_name) info = self._init_dataset( metadata, @@ -128,10 +122,10 @@ def __init__( video_backend, return_uint8, ) - if reader is None: - self._init_paimon_reader(raw_table, info) + if source is None: + self._init_paimon_source(raw_table, metadata, info) else: - self._init_reader(reader, info, False) + self._init_source(source, info, False) def _init_dataset( self, @@ -238,42 +232,40 @@ def _init_episodes(self, episodes): if self._delta_indices and self._episode_ranges is None: raise ValueError("delta_timestamps requires episode metadata.") - def _init_paimon_reader(self, raw_table, info): + def _init_paimon_source(self, raw_table, metadata, info): target_schema = _target_schema(raw_table) projection, validation_context, subtasks = \ self._init_frame_contract(target_schema, info, True) - reader = _PaimonTableFrameReader( - raw_table, projection, self._total_frames) - self._set_reader( - reader, projection, validation_context, subtasks) - - def _init_reader(self, reader, info, validate_metadata): - if not isinstance(reader, LeRobotFrameReader): - raise TypeError("reader must be a LeRobotFrameReader.") - schema = getattr(reader, "schema", None) + source = _PaimonTableDatasetSource( + raw_table, metadata, projection, self._total_frames) + self._set_source( + source, projection, validation_context, subtasks) + + def _init_source(self, source, info, validate_metadata): + schema = getattr(source, "schema", None) if not isinstance(schema, pa.Schema): raise TypeError( - "LeRobotFrameReader.schema must be a pyarrow.Schema.") + "LeRobotDatasetSource.schema must be a pyarrow.Schema.") projection, validation_context, subtasks = \ self._init_frame_contract(schema, info, validate_metadata) - self._set_reader( - reader, projection, validation_context, subtasks) - - def _set_reader( - self, reader, projection, validation_context, subtasks): - self._reader = reader - self._snapshot_id = getattr(reader, "snapshot_id", None) - self._read_table = getattr(reader, "_table", None) - self._frame_locator = getattr(reader, "_locator", None) + self._set_source( + source, projection, validation_context, subtasks) + + def _set_source( + self, source, projection, validation_context, subtasks): + self._source = source + self._snapshot_id = getattr(source, "snapshot_id", None) + self._read_table = getattr(source, "_table", None) + self._frame_locator = getattr(source, "_locator", None) self._projection = projection self._validation_context = validation_context - self._file_io = getattr(reader, "file_io", None) + self._file_io = getattr(source, "file_io", None) if self._video_keys and self._file_io is None: raise ValueError( - "A video-backed LeRobotFrameReader must expose file_io.") + "A video-backed LeRobotDatasetSource must expose file_io.") self._video_collators = [ VideoFrameCollator( - reader, + source, video_column=key, decoder_factory=partial( _open_video_decoder, backend=self.video_backend), @@ -295,7 +287,7 @@ def _init_frame_contract(self, target_schema, info, validate_metadata): _validate_lerobot_schema( source_schema, target_schema, self.repo_id) else: - _validate_reader_schema( + _validate_source_schema( source_schema, target_schema, self.repo_id) validation_context = _build_frame_validation_context( self.meta, @@ -462,11 +454,11 @@ def close(self): except Exception as error: if first_error is None: first_error = error - reader = getattr(self, "_reader", None) - self._reader = None - if reader is not None: + source = getattr(self, "_source", None) + self._source = None + if source is not None: try: - reader.close() + source.close() except Exception as error: if first_error is None: first_error = error @@ -482,8 +474,8 @@ def __del__(self): def _read_rows(self, indices, projection): if not indices: return {} - return _read_reader_rows( - self._reader, + return _read_source_rows( + self._source, projection, indices, self._validation_context, @@ -597,10 +589,10 @@ def get_task_index(self, task): return None -def _reader_metadata(metadata): +def _source_metadata(metadata): info = _metadata_member(metadata, "info") if not isinstance(info, Mapping): - raise TypeError("LeRobot reader metadata must contain an info map.") + raise TypeError("LeRobot source metadata must contain an info map.") info = dict(info) features = info.get("features") if isinstance(features, Mapping): @@ -612,7 +604,7 @@ def _reader_metadata(metadata): feature["shape"] = tuple(feature["shape"]) stats = _metadata_member(metadata, "stats") return _PaimonLeRobotMetadata( - str(_metadata_member(metadata, "repo_id", "logical-reader")), + str(_metadata_member(metadata, "repo_id", "logical-source")), _metadata_member(metadata, "revision"), info, _numpy_stats(stats) if stats is not None else None, @@ -900,30 +892,30 @@ def _delta_indices(delta_timestamps, fps, tolerance_s, features): return result -def _validate_reader_schema(source_schema, reader_schema, source): - for source_field in source_schema: - target_index = reader_schema.get_field_index(source_field.name) +def _validate_source_schema(expected_schema, actual_schema, source): + for expected_field in expected_schema: + target_index = actual_schema.get_field_index(expected_field.name) if target_index < 0: continue - target_type = reader_schema.field(target_index).type - if source_field.type != target_type: + target_type = actual_schema.field(target_index).type + if expected_field.type != target_type: raise ValueError( "LeRobot feature %s from %s expects %s, found %s." - % (source_field.name, source, source_field.type, + % (expected_field.name, source, expected_field.type, target_type)) -def _read_reader_rows( - reader, projection, indices, validation_context, tolerance_s, +def _read_source_rows( + source, projection, indices, validation_context, tolerance_s, features): - values = reader.read_indices(tuple(indices), tuple(projection)) + values = source.read_indices(tuple(indices), tuple(projection)) if not isinstance(values, pa.Table): raise TypeError( - "LeRobotFrameReader.read_indices() must return a pyarrow.Table.") + "LeRobotDatasetSource.read_indices() must return a pyarrow.Table.") missing = set(projection) - set(values.column_names) if missing: raise ValueError( - "LeRobotFrameReader result is missing fields: %s" + "LeRobotDatasetSource result is missing fields: %s" % sorted(missing)) rows = _arrow_rows(values.select(projection), features) expected = set(indices) @@ -932,7 +924,7 @@ def _read_reader_rows( index = _control_index(row, "index", -1) if index not in expected or index in result: raise ValueError( - "LeRobotFrameReader returned an unexpected or duplicate " + "LeRobotDatasetSource returned an unexpected or duplicate " "index: %d." % index) _validate_control_row( index, row, validation_context, tolerance_s) @@ -940,7 +932,7 @@ def _read_reader_rows( missing = expected - set(result) if missing: raise RuntimeError( - "LeRobotFrameReader did not return indices %s." + "LeRobotDatasetSource did not return indices %s." % sorted(missing)) return result diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/dataset_source.py similarity index 92% rename from paimon-python/pypaimon/multimodal/lerobot/reader.py rename to paimon-python/pypaimon/multimodal/lerobot/dataset_source.py index e2eac62b8ef7..d7115a3dfeaf 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/reader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset_source.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Frame readers used by :class:`PaimonLeRobotDataset`.""" +"""Logical frame sources used by :class:`PaimonLeRobotDataset`.""" import os from abc import ABC, abstractmethod @@ -25,15 +25,20 @@ from pypaimon.read.query_auth_split import QueryAuthSplit -class LeRobotFrameReader(ABC): - """Batch reader for logical LeRobot frame rows. +class LeRobotDatasetSource(ABC): + """Data and metadata source for logical LeRobot frame rows. Implementations may resolve one logical row from multiple Paimon tables. Image values must be encoded bytes or BLOB descriptors. Video values must be serialized ``VideoFrameDescriptor`` instances. Descriptor-backed - readers expose their resolved ``file_io``. + sources expose their resolved ``file_io``. """ + @property + @abstractmethod + def metadata(self): + """Return LeRobot metadata for this source.""" + @property @abstractmethod def schema(self): @@ -47,10 +52,11 @@ def close(self): """Release reader resources.""" -class _PaimonTableFrameReader(LeRobotFrameReader): - """Default frame reader for one indexed Paimon frames table.""" +class _PaimonTableDatasetSource(LeRobotDatasetSource): + """Default source for one indexed Paimon frames table.""" - def __init__(self, raw_table, projection, total_frames): + def __init__(self, raw_table, metadata, projection, total_frames): + self._metadata = metadata self._schema = _target_schema(raw_table) self._table, self.snapshot_id, splits = _indexed_read_table( raw_table, projection) @@ -63,6 +69,10 @@ def __init__(self, raw_table, projection, total_frames): self.file_io = self._table.file_io self._locator = _FrameLocator(self._table, snapshot, splits) + @property + def metadata(self): + return self._metadata + @property def schema(self): return self._schema diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 5abefd92a559..b345bf7763dd 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -325,7 +325,7 @@ def test_dataset_requires_supported_python(self): pmm.PaimonLeRobotDataset(Mock()) load.assert_not_called() - def test_dataset_reads_one_batch_from_logical_frame_reader(self): + def test_dataset_reads_one_batch_from_logical_source(self): try: import torch except ImportError as error: @@ -363,12 +363,17 @@ def test_dataset_reads_one_batch_from_logical_frame_reader(self): for index in range(3) } - class Reader(pmm.LeRobotFrameReader): + class Source(pmm.LeRobotDatasetSource): - def __init__(self): + def __init__(self, metadata): + self._metadata = metadata self.calls = [] self.closed = False + @property + def metadata(self): + return self._metadata + @property def schema(self): return _schema_from_info(info) @@ -383,7 +388,6 @@ def read_indices(self, indices, columns): def close(self): self.closed = True - reader = Reader() metadata = { "repo_id": "logical/multi-table", "revision": "dataset-version-12", @@ -398,16 +402,16 @@ def close(self): "tasks": ["pick"], "stats": {"action": {"mean": [1.0]}}, } + source = Source(metadata) dataset = pmm.PaimonLeRobotDataset( - reader=reader, - metadata=metadata, + source, delta_timestamps={"action": [-0.1, 0.0, 0.1]}, ) sample, _ = dataset.__getitems__([1, 2]) self.assertEqual([((0, 1, 2), tuple(info["features"]))], - reader.calls) + source.calls) self.assertIsNone(dataset.tag_name) self.assertEqual("dataset-version-12", dataset.meta.revision) self.assertEqual((2,), dataset.features["observation.state"]["shape"]) @@ -422,7 +426,7 @@ def close(self): self.assertEqual([False, False, False], sample["action_is_pad"].tolist()) dataset.close() - self.assertTrue(reader.closed) + self.assertTrue(source.closed) def test_metadata_json_preserves_nested_values(self): values = { From bcf02aa268f46536e8c7e1dee5cc8c5ea7b10b0b Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 03:12:13 -0700 Subject: [PATCH 07/16] [python] Make LeRobot reader the extension point --- docs/docs/pypaimon/lerobot.md | 24 +- paimon-python/pypaimon/multimodal/__init__.py | 8 +- .../pypaimon/multimodal/lerobot/__init__.py | 8 +- .../pypaimon/multimodal/lerobot/dataset.py | 268 +++++++++++++----- .../lerobot/{dataset_source.py => reader.py} | 48 +--- .../pypaimon/tests/multimodal_lerobot_test.py | 44 +-- 6 files changed, 257 insertions(+), 143 deletions(-) rename paimon-python/pypaimon/multimodal/lerobot/{dataset_source.py => reader.py} (82%) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index d7512c63f991..5ea1d7449cb5 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -227,20 +227,28 @@ on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force either decoder. -Use a dataset source when one logical frame is assembled from multiple tables: +Subclass `PaimonDatasetReader` when one logical frame is assembled from +multiple tables: ```python -from pypaimon.multimodal import LeRobotDatasetSource, PaimonLeRobotDataset +from pypaimon.multimodal import PaimonDatasetReader, PaimonLeRobotDataset -class CustomDatasetSource(LeRobotDatasetSource): - metadata = dataset_metadata - schema = logical_frame_schema +class CustomDatasetReader(PaimonDatasetReader): + def __init__(self, version, **kwargs): + self._version = version + super().__init__(version.lerobot_metadata, **kwargs) + + @property + def schema(self): + return logical_frame_schema def read_indices(self, indices, columns): return resolve_frame_rows(indices, columns) # pyarrow.Table -dataset = PaimonLeRobotDataset(CustomDatasetSource()) +reader = CustomDatasetReader(version, delta_timestamps=delta_timestamps) +dataset = PaimonLeRobotDataset(reader) ``` -`LeRobotDatasetSource` supplies metadata and logical frame rows. It may query -any tables. Descriptor-backed media sources also expose `file_io`. +`PaimonDatasetReader` may query any tables. It batches logical row reads and +reuses the standard Episode, delta-window, media, and Torch handling. Readers +returning media descriptors must set `file_io` before calling `super().__init__`. diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index 069e237bb6da..25cfc784c5d7 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -29,8 +29,10 @@ Hdf5File, Hdf5LoadResult, ) -from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset -from pypaimon.multimodal.lerobot.dataset_source import LeRobotDatasetSource +from pypaimon.multimodal.lerobot.dataset import ( + PaimonDatasetReader, + PaimonLeRobotDataset, +) from pypaimon.multimodal.rosbag import ( RosbagLoadResult, RosbagSource, @@ -64,11 +66,11 @@ "BlobStore", "Hdf5File", "Hdf5LoadResult", - "LeRobotDatasetSource", "MultimodalConnection", "MultimodalTable", "NoSuchKey", "ObjectInfo", + "PaimonDatasetReader", "PaimonLeRobotDataset", "PutObjectResult", "RosbagLoadResult", diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index dde2858e1bf0..482184a52b99 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -17,13 +17,15 @@ """LeRobot Dataset v3 integration for multimodal Paimon tables.""" from pypaimon.multimodal.lerobot.api import load_from_lerobot -from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset -from pypaimon.multimodal.lerobot.dataset_source import LeRobotDatasetSource +from pypaimon.multimodal.lerobot.dataset import ( + PaimonDatasetReader, + PaimonLeRobotDataset, +) from pypaimon.multimodal.lerobot.writer import PaimonLeRobotWriter __all__ = [ - "LeRobotDatasetSource", + "PaimonDatasetReader", "PaimonLeRobotDataset", "PaimonLeRobotWriter", "load_from_lerobot", diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index b1f34f007044..6d517c2c3632 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -23,6 +23,7 @@ import math import operator import sys +from abc import ABC, abstractmethod from collections import OrderedDict from collections.abc import Mapping from functools import partial @@ -36,10 +37,7 @@ _validate_tag_name, ) from pypaimon.multimodal.lerobot.loader import _DECLARED_NUMERIC_RANGES -from pypaimon.multimodal.lerobot.dataset_source import ( - LeRobotDatasetSource, - _PaimonTableDatasetSource, -) +from pypaimon.multimodal.lerobot.reader import _PaimonTableFrameReader from pypaimon.multimodal.lerobot.schema import ( _feature_shape, _require_v3, @@ -77,10 +75,10 @@ }) -class PaimonLeRobotDataset: - """Map-style LeRobot reader backed by indexed Paimon reads. +class PaimonDatasetReader(ABC): + """Read-side implementation for Paimon-backed LeRobot datasets. - The input is a Paimon table group or :class:`LeRobotDatasetSource`. + Subclasses provide a logical Arrow schema and batched ``read_indices``. Resolved LeRobot metadata remains available through :attr:`meta`. Set ``return_uint8=True`` to keep 8-bit visual frames in their decoded @@ -90,7 +88,7 @@ class PaimonLeRobotDataset: def __init__( self, - table, + meta, *, tag_name=None, episodes=None, @@ -99,18 +97,10 @@ def __init__( tolerance_s=1e-4, blob_parallelism=16, video_backend=None, - return_uint8=False): + return_uint8=False, + _resolved_meta=False): _require_dataset_python() - if isinstance(table, LeRobotDatasetSource): - if tag_name is not None: - raise ValueError( - "tag_name is managed by a custom source and must be None.") - source = table - metadata = _source_metadata(source.metadata) - raw_table = None - else: - source = None - raw_table, metadata = _load_dataset(table, tag_name) + metadata = meta if _resolved_meta else _reader_metadata(meta) info = self._init_dataset( metadata, tag_name, @@ -122,10 +112,31 @@ def __init__( video_backend, return_uint8, ) - if source is None: - self._init_paimon_source(raw_table, metadata, info) - else: - self._init_source(source, info, False) + schema = self.schema + if not isinstance(schema, pa.Schema): + raise TypeError( + "PaimonDatasetReader.schema must be a pyarrow.Schema.") + projection, validation_context, subtasks = \ + self._init_frame_contract( + schema, info, self._validate_physical_metadata()) + rows = self._open_frame_rows(projection) + self._set_frame_rows( + rows, projection, validation_context, subtasks) + + @property + @abstractmethod + def schema(self): + """Return the logical frame schema as :class:`pyarrow.Schema`.""" + + @abstractmethod + def read_indices(self, indices, columns): + """Return requested logical frame rows as a :class:`pyarrow.Table`.""" + + def _validate_physical_metadata(self): + return False + + def _open_frame_rows(self, projection): + return None def _init_dataset( self, @@ -232,40 +243,22 @@ def _init_episodes(self, episodes): if self._delta_indices and self._episode_ranges is None: raise ValueError("delta_timestamps requires episode metadata.") - def _init_paimon_source(self, raw_table, metadata, info): - target_schema = _target_schema(raw_table) - projection, validation_context, subtasks = \ - self._init_frame_contract(target_schema, info, True) - source = _PaimonTableDatasetSource( - raw_table, metadata, projection, self._total_frames) - self._set_source( - source, projection, validation_context, subtasks) - - def _init_source(self, source, info, validate_metadata): - schema = getattr(source, "schema", None) - if not isinstance(schema, pa.Schema): - raise TypeError( - "LeRobotDatasetSource.schema must be a pyarrow.Schema.") - projection, validation_context, subtasks = \ - self._init_frame_contract(schema, info, validate_metadata) - self._set_source( - source, projection, validation_context, subtasks) - - def _set_source( - self, source, projection, validation_context, subtasks): - self._source = source - self._snapshot_id = getattr(source, "snapshot_id", None) - self._read_table = getattr(source, "_table", None) - self._frame_locator = getattr(source, "_locator", None) + def _set_frame_rows( + self, rows, projection, validation_context, subtasks): + self._frame_rows = rows + access = rows if rows is not None else self + self._snapshot_id = getattr(access, "snapshot_id", None) + self._read_table = getattr(access, "_table", None) + self._frame_locator = getattr(access, "_locator", None) self._projection = projection self._validation_context = validation_context - self._file_io = getattr(source, "file_io", None) + self._file_io = getattr(access, "file_io", None) if self._video_keys and self._file_io is None: raise ValueError( - "A video-backed LeRobotDatasetSource must expose file_io.") + "A video-backed PaimonDatasetReader must expose file_io.") self._video_collators = [ VideoFrameCollator( - source, + access, video_column=key, decoder_factory=partial( _open_video_decoder, backend=self.video_backend), @@ -287,7 +280,7 @@ def _init_frame_contract(self, target_schema, info, validate_metadata): _validate_lerobot_schema( source_schema, target_schema, self.repo_id) else: - _validate_source_schema( + _validate_reader_schema( source_schema, target_schema, self.repo_id) validation_context = _build_frame_validation_context( self.meta, @@ -341,6 +334,26 @@ def num_episodes(self): def __len__(self): return self.num_frames + @property + def absolute_to_relative_idx(self): + if self._selected_ranges is None: + return None + result = {} + relative = 0 + for begin, end in self._selected_ranges: + for absolute in range(begin, end): + result[absolute] = relative + relative += 1 + return result + + def get_item(self, index): + """Return one fully assembled frame.""" + return self[index] + + def get_items(self, indices): + """Return fully assembled frames for one batch.""" + return self.__getitems__(indices) + def __getitem__(self, index): if isinstance(index, slice): return self.__getitems__(range(*index.indices(len(self)))) @@ -454,11 +467,11 @@ def close(self): except Exception as error: if first_error is None: first_error = error - source = getattr(self, "_source", None) - self._source = None - if source is not None: + rows = getattr(self, "_frame_rows", None) + self._frame_rows = None + if rows is not None: try: - source.close() + rows.close() except Exception as error: if first_error is None: first_error = error @@ -474,8 +487,9 @@ def __del__(self): def _read_rows(self, indices, projection): if not indices: return {} - return _read_source_rows( - self._source, + rows = self._frame_rows if self._frame_rows is not None else self + return _read_reader_rows( + rows, projection, indices, self._validation_context, @@ -523,6 +537,121 @@ def __repr__(self): self.num_frames, list(self.features))) +class _PaimonTableDatasetReader(PaimonDatasetReader): + + def __init__( + self, + table, + *, + tag_name=None, + episodes=None, + image_transforms=None, + delta_timestamps=None, + tolerance_s=1e-4, + blob_parallelism=16, + video_backend=None, + return_uint8=False): + self._raw_table, meta = _load_dataset(table, tag_name) + super().__init__( + meta, + tag_name=tag_name, + episodes=episodes, + image_transforms=image_transforms, + delta_timestamps=delta_timestamps, + tolerance_s=tolerance_s, + blob_parallelism=blob_parallelism, + video_backend=video_backend, + return_uint8=return_uint8, + _resolved_meta=True, + ) + + @property + def schema(self): + return _target_schema(self._raw_table) + + def read_indices(self, indices, columns): + return self._frame_rows.read_indices(indices, columns) + + def _validate_physical_metadata(self): + return True + + def _open_frame_rows(self, projection): + return _PaimonTableFrameReader( + self._raw_table, projection, self._total_frames) + + +class PaimonLeRobotDataset: + """Map-style Dataset facade backed by :class:`PaimonDatasetReader`.""" + + def __init__( + self, + table, + *, + tag_name=None, + episodes=None, + image_transforms=None, + delta_timestamps=None, + tolerance_s=1e-4, + blob_parallelism=16, + video_backend=None, + return_uint8=False): + _require_dataset_python() + if isinstance(table, PaimonDatasetReader): + if ( + tag_name is not None + or episodes is not None + or image_transforms is not None + or delta_timestamps is not None + or tolerance_s != 1e-4 + or blob_parallelism != 16 + or video_backend is not None + or return_uint8 + ): + raise ValueError( + "Configure Dataset options on PaimonDatasetReader.") + self.reader = table + else: + self.reader = _PaimonTableDatasetReader( + table, + tag_name=tag_name, + episodes=episodes, + image_transforms=image_transforms, + delta_timestamps=delta_timestamps, + tolerance_s=tolerance_s, + blob_parallelism=blob_parallelism, + video_backend=video_backend, + return_uint8=return_uint8, + ) + + def __len__(self): + return len(self.reader) + + def __getitem__(self, index): + return self.reader[index] + + def __getitems__(self, indices): + return self.reader.get_items(indices) + + def set_image_transforms(self, image_transforms): + self.reader.set_image_transforms(image_transforms) + + def clear_image_transforms(self): + self.reader.clear_image_transforms() + + def close(self): + self.reader.close() + + def __getattr__(self, name): + reader = self.__dict__.get("reader") + if reader is None: + raise AttributeError(name) + return getattr(reader, name) + + def __repr__(self): + return repr(self.reader).replace( + self.reader.__class__.__name__, self.__class__.__name__, 1) + + class _PaimonLeRobotMetadata: def __init__( @@ -589,10 +718,11 @@ def get_task_index(self, task): return None -def _source_metadata(metadata): +def _reader_metadata(metadata): info = _metadata_member(metadata, "info") if not isinstance(info, Mapping): - raise TypeError("LeRobot source metadata must contain an info map.") + raise TypeError( + "PaimonDatasetReader metadata must contain an info map.") info = dict(info) features = info.get("features") if isinstance(features, Mapping): @@ -604,7 +734,7 @@ def _source_metadata(metadata): feature["shape"] = tuple(feature["shape"]) stats = _metadata_member(metadata, "stats") return _PaimonLeRobotMetadata( - str(_metadata_member(metadata, "repo_id", "logical-source")), + str(_metadata_member(metadata, "repo_id", "custom-reader")), _metadata_member(metadata, "revision"), info, _numpy_stats(stats) if stats is not None else None, @@ -892,7 +1022,7 @@ def _delta_indices(delta_timestamps, fps, tolerance_s, features): return result -def _validate_source_schema(expected_schema, actual_schema, source): +def _validate_reader_schema(expected_schema, actual_schema, source): for expected_field in expected_schema: target_index = actual_schema.get_field_index(expected_field.name) if target_index < 0: @@ -905,17 +1035,17 @@ def _validate_source_schema(expected_schema, actual_schema, source): target_type)) -def _read_source_rows( - source, projection, indices, validation_context, tolerance_s, +def _read_reader_rows( + reader, projection, indices, validation_context, tolerance_s, features): - values = source.read_indices(tuple(indices), tuple(projection)) + values = reader.read_indices(tuple(indices), tuple(projection)) if not isinstance(values, pa.Table): raise TypeError( - "LeRobotDatasetSource.read_indices() must return a pyarrow.Table.") + "PaimonDatasetReader.read_indices() must return a pyarrow.Table.") missing = set(projection) - set(values.column_names) if missing: raise ValueError( - "LeRobotDatasetSource result is missing fields: %s" + "PaimonDatasetReader result is missing fields: %s" % sorted(missing)) rows = _arrow_rows(values.select(projection), features) expected = set(indices) @@ -924,7 +1054,7 @@ def _read_source_rows( index = _control_index(row, "index", -1) if index not in expected or index in result: raise ValueError( - "LeRobotDatasetSource returned an unexpected or duplicate " + "PaimonDatasetReader returned an unexpected or duplicate " "index: %d." % index) _validate_control_row( index, row, validation_context, tolerance_s) @@ -932,7 +1062,7 @@ def _read_source_rows( missing = expected - set(result) if missing: raise RuntimeError( - "LeRobotDatasetSource did not return indices %s." + "PaimonDatasetReader did not return indices %s." % sorted(missing)) return result diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset_source.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py similarity index 82% rename from paimon-python/pypaimon/multimodal/lerobot/dataset_source.py rename to paimon-python/pypaimon/multimodal/lerobot/reader.py index d7115a3dfeaf..5f67dd1676d9 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset_source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -15,49 +15,19 @@ # specific language governing permissions and limitations # under the License. -"""Logical frame sources used by :class:`PaimonLeRobotDataset`.""" +"""Indexed frame-row reader used by :class:`PaimonDatasetReader`.""" import os -from abc import ABC, abstractmethod from pypaimon.common.options.core_options import CoreOptions -from pypaimon.multimodal.table import _target_schema, _time_travel_table +from pypaimon.multimodal.table import _time_travel_table from pypaimon.read.query_auth_split import QueryAuthSplit -class LeRobotDatasetSource(ABC): - """Data and metadata source for logical LeRobot frame rows. +class _PaimonTableFrameReader: + """Read logical frame rows from one indexed Paimon table.""" - Implementations may resolve one logical row from multiple Paimon tables. - Image values must be encoded bytes or BLOB descriptors. Video values must - be serialized ``VideoFrameDescriptor`` instances. Descriptor-backed - sources expose their resolved ``file_io``. - """ - - @property - @abstractmethod - def metadata(self): - """Return LeRobot metadata for this source.""" - - @property - @abstractmethod - def schema(self): - """Return the logical frame schema as :class:`pyarrow.Schema`.""" - - @abstractmethod - def read_indices(self, indices, columns): - """Return requested rows as a :class:`pyarrow.Table`.""" - - def close(self): - """Release reader resources.""" - - -class _PaimonTableDatasetSource(LeRobotDatasetSource): - """Default source for one indexed Paimon frames table.""" - - def __init__(self, raw_table, metadata, projection, total_frames): - self._metadata = metadata - self._schema = _target_schema(raw_table) + def __init__(self, raw_table, projection, total_frames): self._table, self.snapshot_id, splits = _indexed_read_table( raw_table, projection) snapshot = self._table.snapshot_manager().get_snapshot_by_id( @@ -69,14 +39,6 @@ def __init__(self, raw_table, metadata, projection, total_frames): self.file_io = self._table.file_io self._locator = _FrameLocator(self._table, snapshot, splits) - @property - def metadata(self): - return self._metadata - - @property - def schema(self): - return self._schema - def read_indices(self, indices, columns): splits, needs_filter = self._locator.locate(indices) builder = self._table.new_read_builder().with_projection(columns) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index b345bf7763dd..16c226c75743 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -97,6 +97,16 @@ av = None +class _ManualDatasetReader(pmm.PaimonDatasetReader): + + @property + def schema(self): + raise NotImplementedError + + def read_indices(self, indices, columns): + raise NotImplementedError + + def _replaced_contract(field, old, new): description = field.metadata[b"description"].decode("utf-8") if old not in description: @@ -325,7 +335,7 @@ def test_dataset_requires_supported_python(self): pmm.PaimonLeRobotDataset(Mock()) load.assert_not_called() - def test_dataset_reads_one_batch_from_logical_source(self): + def test_dataset_reads_one_batch_from_custom_reader(self): try: import torch except ImportError as error: @@ -363,16 +373,12 @@ def test_dataset_reads_one_batch_from_logical_source(self): for index in range(3) } - class Source(pmm.LeRobotDatasetSource): + class Reader(pmm.PaimonDatasetReader): - def __init__(self, metadata): - self._metadata = metadata + def __init__(self, metadata, **kwargs): self.calls = [] self.closed = False - - @property - def metadata(self): - return self._metadata + super().__init__(metadata, **kwargs) @property def schema(self): @@ -386,6 +392,7 @@ def read_indices(self, indices, columns): ], schema=self.schema) def close(self): + super().close() self.closed = True metadata = { @@ -402,16 +409,19 @@ def close(self): "tasks": ["pick"], "stats": {"action": {"mean": [1.0]}}, } - source = Source(metadata) - dataset = pmm.PaimonLeRobotDataset( - source, + reader = Reader( + metadata, delta_timestamps={"action": [-0.1, 0.0, 0.1]}, ) + dataset = pmm.PaimonLeRobotDataset(reader) + self.assertIsInstance(dataset.reader, pmm.PaimonDatasetReader) + self.assertIsNone(dataset.reader.absolute_to_relative_idx) + self.assertTrue(repr(dataset).startswith("PaimonLeRobotDataset(")) sample, _ = dataset.__getitems__([1, 2]) self.assertEqual([((0, 1, 2), tuple(info["features"]))], - source.calls) + reader.calls) self.assertIsNone(dataset.tag_name) self.assertEqual("dataset-version-12", dataset.meta.revision) self.assertEqual((2,), dataset.features["observation.state"]["shape"]) @@ -426,7 +436,7 @@ def close(self): self.assertEqual([False, False, False], sample["action_is_pad"].tolist()) dataset.close() - self.assertTrue(source.closed) + self.assertTrue(reader.closed) def test_metadata_json_preserves_nested_values(self): values = { @@ -735,7 +745,7 @@ def jpeg(mode, values): "L", np.full((4, 5), 80 + index, np.uint8)), }) - dataset = object.__new__(pmm.PaimonLeRobotDataset) + dataset = object.__new__(_ManualDatasetReader) dataset._total_frames = 2 dataset.episodes = None dataset._selected_ranges = None @@ -795,7 +805,7 @@ def test_dataset_retries_image_fetch_and_decode_together(self): "observation.image": descriptor, }] - dataset = object.__new__(pmm.PaimonLeRobotDataset) + dataset = object.__new__(_ManualDatasetReader) dataset._total_frames = 1 dataset.episodes = None dataset._selected_ranges = None @@ -3039,8 +3049,8 @@ def counted_plan(scan): return original_plan(scan) with patch.object(TableScan, "plan", new=counted_plan), patch.object( - dataset, "_read_rows", - wraps=dataset._read_rows) as read, patch( + dataset.reader, "_read_rows", + wraps=dataset.reader._read_rows) as read, patch( "pypaimon.multimodal.blob_read.fetch_blob_bodies", wraps=fetch_blob_bodies) as fetch: last, first = dataset.__getitems__([4, 0]) From fbf5b95e98881b1891e1e369498c6e3029a8282e Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 03:26:22 -0700 Subject: [PATCH 08/16] [python] Describe custom LeRobot readers generically --- docs/docs/pypaimon/lerobot.md | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 5ea1d7449cb5..537ec6e548d4 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -227,8 +227,7 @@ on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force either decoder. -Subclass `PaimonDatasetReader` when one logical frame is assembled from -multiple tables: +Subclass `PaimonDatasetReader` for a custom logical frame layout: ```python from pypaimon.multimodal import PaimonDatasetReader, PaimonLeRobotDataset @@ -249,6 +248,6 @@ reader = CustomDatasetReader(version, delta_timestamps=delta_timestamps) dataset = PaimonLeRobotDataset(reader) ``` -`PaimonDatasetReader` may query any tables. It batches logical row reads and -reuses the standard Episode, delta-window, media, and Torch handling. Readers -returning media descriptors must set `file_io` before calling `super().__init__`. +`PaimonDatasetReader` batches logical row reads and reuses the standard +Episode, delta-window, media, and Torch handling. Readers returning media +descriptors must set `file_io` before calling `super().__init__`. From 91d6ff19991c6ec07767221a870c23235e4f5f8c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 03:32:16 -0700 Subject: [PATCH 09/16] [python] Make reader methods the primary dataset API --- .../pypaimon/multimodal/lerobot/dataset.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 6d517c2c3632..fc9deb43d7a9 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -348,18 +348,10 @@ def absolute_to_relative_idx(self): def get_item(self, index): """Return one fully assembled frame.""" - return self[index] + return self.get_items([index])[0] def get_items(self, indices): """Return fully assembled frames for one batch.""" - return self.__getitems__(indices) - - def __getitem__(self, index): - if isinstance(index, slice): - return self.__getitems__(range(*index.indices(len(self)))) - return self.__getitems__([index])[0] - - def __getitems__(self, indices): dataset_indices = [ _normalize_index(index, len(self)) for index in indices ] @@ -459,6 +451,14 @@ def __getitems__(self, indices): result.append(item) return result + def __getitem__(self, index): + if isinstance(index, slice): + return self.get_items(range(*index.indices(len(self)))) + return self.get_item(index) + + def __getitems__(self, indices): + return self.get_items(indices) + def close(self): first_error = None for collator in getattr(self, "_video_collators", ()): From c3e5fc8c480ac5ed0e02a5ef6f3daf8251b917c9 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 04:29:58 -0700 Subject: [PATCH 10/16] [python] Clarify custom dataset reader contract --- .../pypaimon/multimodal/lerobot/dataset.py | 13 +++++++++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 4 ---- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index fc9deb43d7a9..78ad21f96bb3 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -79,13 +79,17 @@ class PaimonDatasetReader(ABC): """Read-side implementation for Paimon-backed LeRobot datasets. Subclasses provide a logical Arrow schema and batched ``read_indices``. - Resolved LeRobot metadata remains available through :attr:`meta`. + Resolved LeRobot metadata remains available through :attr:`meta`. Readers + must be picklable for DataLoader workers. Readers returning BLOB or video + descriptors must set ``file_io`` before calling ``super().__init__``. Set ``return_uint8=True`` to keep 8-bit visual frames in their decoded ``torch.uint8`` representation instead of normalizing them to float32. Higher-bit-depth images retain the existing float32 behavior. """ + file_io = None + def __init__( self, meta, @@ -130,7 +134,12 @@ def schema(self): @abstractmethod def read_indices(self, indices, columns): - """Return requested logical frame rows as a :class:`pyarrow.Table`.""" + """Return one row per requested absolute index as a PyArrow Table. + + ``indices`` are unique; result order is unrestricted. Returned rows + must contain every requested column without missing or duplicate + indices. + """ def _validate_physical_metadata(self): return False diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 16c226c75743..320fc9eb0437 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -355,7 +355,6 @@ def test_dataset_reads_one_batch_from_custom_reader(self): "task_index": {"dtype": "int64", "shape": [1]}, "observation.state": {"dtype": "float32", "shape": [2]}, "action": {"dtype": "float32", "shape": [1]}, - "camera.image": {"dtype": "image", "shape": [2, 2, 3]}, }, } rows = { @@ -367,8 +366,6 @@ def test_dataset_reads_one_batch_from_custom_reader(self): "task_index": 0, "observation.state": [index, index + 1], "action": float(index), - "camera.image": _image_bytes( - np.full((2, 2, 3), index, dtype=np.uint8), None), } for index in range(3) } @@ -432,7 +429,6 @@ def close(self): sample["observation.state"], torch.tensor([1.0, 2.0])) torch.testing.assert_close( sample["action"], torch.tensor([0.0, 1.0, 2.0])) - self.assertEqual([3, 2, 2], list(sample["camera.image"].shape)) self.assertEqual([False, False, False], sample["action_is_pad"].tolist()) dataset.close() From c401dfd5a6b78d0d7fd1f3534fc67a3b4a01529b Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 04:52:07 -0700 Subject: [PATCH 11/16] [python] Keep pickle protocol on dataset facade --- .../pypaimon/multimodal/lerobot/dataset.py | 2 ++ .../pypaimon/tests/multimodal_lerobot_test.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 78ad21f96bb3..30e71e082094 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -651,6 +651,8 @@ def close(self): self.reader.close() def __getattr__(self, name): + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) reader = self.__dict__.get("reader") if reader is None: raise AttributeError(name) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 320fc9eb0437..7bfe16093f13 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -107,6 +107,15 @@ def read_indices(self, indices, columns): raise NotImplementedError +class _PickleDatasetReader(_ManualDatasetReader): + + def __init__(self, value): + self.value = value + + def __getstate__(self): + return {"value": self.value} + + def _replaced_contract(field, old, new): description = field.metadata[b"description"].decode("utf-8") if old not in description: @@ -434,6 +443,15 @@ def close(self): dataset.close() self.assertTrue(reader.closed) + def test_dataset_does_not_proxy_pickle_protocol(self): + dataset = pmm.PaimonLeRobotDataset(_PickleDatasetReader(7)) + with self.assertRaises(AttributeError): + dataset.__getattr__("__getstate__") + restored = pickle.loads(pickle.dumps(dataset)) + + self.assertIsInstance(restored.reader, _PickleDatasetReader) + self.assertEqual(7, restored.reader.value) + def test_metadata_json_preserves_nested_values(self): values = { "name": "机器人", From cc3a39f62c048921fc972ea4264aa94c71e21df3 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 05:09:38 -0700 Subject: [PATCH 12/16] [python] Separate frame reading from metadata validation --- docs/docs/pypaimon/lerobot.md | 21 +++++++++++-------- .../pypaimon/multimodal/lerobot/dataset.py | 12 +++++++---- .../pypaimon/multimodal/lerobot/reader.py | 9 +++----- 3 files changed, 23 insertions(+), 19 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 537ec6e548d4..35737d9e60e1 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -233,21 +233,24 @@ Subclass `PaimonDatasetReader` for a custom logical frame layout: from pypaimon.multimodal import PaimonDatasetReader, PaimonLeRobotDataset class CustomDatasetReader(PaimonDatasetReader): - def __init__(self, version, **kwargs): - self._version = version - super().__init__(version.lerobot_metadata, **kwargs) + def __init__(self, metadata, source, **kwargs): + self._source = source + self.file_io = getattr(source, "file_io", None) + super().__init__(metadata, **kwargs) @property def schema(self): - return logical_frame_schema + return self._source.schema def read_indices(self, indices, columns): - return resolve_frame_rows(indices, columns) # pyarrow.Table + return self._source.read_indices(indices, columns) -reader = CustomDatasetReader(version, delta_timestamps=delta_timestamps) +reader = CustomDatasetReader( + metadata, source, delta_timestamps=delta_timestamps +) dataset = PaimonLeRobotDataset(reader) ``` -`PaimonDatasetReader` batches logical row reads and reuses the standard -Episode, delta-window, media, and Torch handling. Readers returning media -descriptors must set `file_io` before calling `super().__init__`. +`source` exposes `schema`, `read_indices`, and optional `file_io`. +`PaimonDatasetReader` reuses the standard Episode, delta-window, media, and +Torch handling. diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 30e71e082094..7f910e16bef4 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -560,7 +560,7 @@ def __init__( blob_parallelism=16, video_backend=None, return_uint8=False): - self._raw_table, meta = _load_dataset(table, tag_name) + self._frames_table, meta = _load_dataset(table, tag_name) super().__init__( meta, tag_name=tag_name, @@ -576,7 +576,7 @@ def __init__( @property def schema(self): - return _target_schema(self._raw_table) + return _target_schema(self._frames_table) def read_indices(self, indices, columns): return self._frame_rows.read_indices(indices, columns) @@ -585,8 +585,12 @@ def _validate_physical_metadata(self): return True def _open_frame_rows(self, projection): - return _PaimonTableFrameReader( - self._raw_table, projection, self._total_frames) + rows = _PaimonTableFrameReader(self._frames_table, projection) + if rows.num_rows != self._total_frames: + raise ValueError( + "Paimon table has %d rows but metadata declares %d frames." + % (rows.num_rows, self._total_frames)) + return rows class PaimonLeRobotDataset: diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py index 5f67dd1676d9..a032011912d1 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/reader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -27,15 +27,12 @@ class _PaimonTableFrameReader: """Read logical frame rows from one indexed Paimon table.""" - def __init__(self, raw_table, projection, total_frames): + def __init__(self, frames_table, projection): self._table, self.snapshot_id, splits = _indexed_read_table( - raw_table, projection) + frames_table, projection) snapshot = self._table.snapshot_manager().get_snapshot_by_id( self.snapshot_id) - if snapshot.next_row_id != total_frames: - raise ValueError( - "Paimon table has %d rows but metadata declares %d frames." - % (snapshot.next_row_id, total_frames)) + self.num_rows = snapshot.next_row_id self.file_io = self._table.file_io self._locator = _FrameLocator(self._table, snapshot, splits) From 276bbc7e87916397c555a633eca19f3634d3784f Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 05:31:06 -0700 Subject: [PATCH 13/16] [python] Make custom reader inputs explicit --- docs/docs/pypaimon/lerobot.md | 12 +++--- .../pypaimon/multimodal/lerobot/dataset.py | 39 +++++++------------ .../pypaimon/multimodal/lerobot/reader.py | 4 +- .../pypaimon/tests/multimodal_lerobot_test.py | 17 ++++---- 4 files changed, 29 insertions(+), 43 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 35737d9e60e1..cb6fa3aae688 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -235,12 +235,12 @@ from pypaimon.multimodal import PaimonDatasetReader, PaimonLeRobotDataset class CustomDatasetReader(PaimonDatasetReader): def __init__(self, metadata, source, **kwargs): self._source = source - self.file_io = getattr(source, "file_io", None) - super().__init__(metadata, **kwargs) - - @property - def schema(self): - return self._source.schema + super().__init__( + metadata, + schema=source.schema, + file_io=getattr(source, "file_io", None), + **kwargs, + ) def read_indices(self, indices, columns): return self._source.read_indices(indices, columns) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 7f910e16bef4..a999c42704d3 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -78,23 +78,21 @@ class PaimonDatasetReader(ABC): """Read-side implementation for Paimon-backed LeRobot datasets. - Subclasses provide a logical Arrow schema and batched ``read_indices``. - Resolved LeRobot metadata remains available through :attr:`meta`. Readers - must be picklable for DataLoader workers. Readers returning BLOB or video - descriptors must set ``file_io`` before calling ``super().__init__``. + Subclasses provide batched ``read_indices``. Resolved LeRobot metadata + remains available through :attr:`meta`. Readers must be picklable for + DataLoader workers. Set ``return_uint8=True`` to keep 8-bit visual frames in their decoded ``torch.uint8`` representation instead of normalizing them to float32. Higher-bit-depth images retain the existing float32 behavior. """ - file_io = None - def __init__( self, meta, *, - tag_name=None, + schema, + file_io=None, episodes=None, image_transforms=None, delta_timestamps=None, @@ -105,9 +103,13 @@ def __init__( _resolved_meta=False): _require_dataset_python() metadata = meta if _resolved_meta else _reader_metadata(meta) + if not isinstance(schema, pa.Schema): + raise TypeError( + "PaimonDatasetReader schema must be a pyarrow.Schema.") + self.schema = schema + self.file_io = file_io info = self._init_dataset( metadata, - tag_name, episodes, image_transforms, delta_timestamps, @@ -116,10 +118,6 @@ def __init__( video_backend, return_uint8, ) - schema = self.schema - if not isinstance(schema, pa.Schema): - raise TypeError( - "PaimonDatasetReader.schema must be a pyarrow.Schema.") projection, validation_context, subtasks = \ self._init_frame_contract( schema, info, self._validate_physical_metadata()) @@ -127,11 +125,6 @@ def __init__( self._set_frame_rows( rows, projection, validation_context, subtasks) - @property - @abstractmethod - def schema(self): - """Return the logical frame schema as :class:`pyarrow.Schema`.""" - @abstractmethod def read_indices(self, indices, columns): """Return one row per requested absolute index as a PyArrow Table. @@ -150,7 +143,6 @@ def _open_frame_rows(self, projection): def _init_dataset( self, metadata, - tag_name, episodes, image_transforms, delta_timestamps, @@ -159,7 +151,6 @@ def _init_dataset( video_backend, return_uint8): self.meta = metadata - self.tag_name = tag_name self.repo_id = self.meta.repo_id self.image_transforms = image_transforms self.delta_timestamps = delta_timestamps @@ -561,9 +552,10 @@ def __init__( video_backend=None, return_uint8=False): self._frames_table, meta = _load_dataset(table, tag_name) + self.tag_name = tag_name super().__init__( meta, - tag_name=tag_name, + schema=_target_schema(self._frames_table), episodes=episodes, image_transforms=image_transforms, delta_timestamps=delta_timestamps, @@ -574,10 +566,6 @@ def __init__( _resolved_meta=True, ) - @property - def schema(self): - return _target_schema(self._frames_table) - def read_indices(self, indices, columns): return self._frame_rows.read_indices(indices, columns) @@ -585,7 +573,8 @@ def _validate_physical_metadata(self): return True def _open_frame_rows(self, projection): - rows = _PaimonTableFrameReader(self._frames_table, projection) + rows = _PaimonTableFrameReader( + self._frames_table, columns=projection) if rows.num_rows != self._total_frames: raise ValueError( "Paimon table has %d rows but metadata declares %d frames." diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py index a032011912d1..a63945fd869d 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/reader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -27,9 +27,9 @@ class _PaimonTableFrameReader: """Read logical frame rows from one indexed Paimon table.""" - def __init__(self, frames_table, projection): + def __init__(self, frames_table, *, columns): self._table, self.snapshot_id, splits = _indexed_read_table( - frames_table, projection) + frames_table, columns) snapshot = self._table.snapshot_manager().get_snapshot_by_id( self.snapshot_id) self.num_rows = snapshot.next_row_id diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 7bfe16093f13..d5f2b47ee31b 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -99,10 +99,6 @@ class _ManualDatasetReader(pmm.PaimonDatasetReader): - @property - def schema(self): - raise NotImplementedError - def read_indices(self, indices, columns): raise NotImplementedError @@ -384,11 +380,8 @@ class Reader(pmm.PaimonDatasetReader): def __init__(self, metadata, **kwargs): self.calls = [] self.closed = False - super().__init__(metadata, **kwargs) - - @property - def schema(self): - return _schema_from_info(info) + super().__init__( + metadata, schema=_schema_from_info(info), **kwargs) def read_indices(self, indices, columns): self.calls.append((indices, columns)) @@ -419,6 +412,8 @@ def close(self): metadata, delta_timestamps={"action": [-0.1, 0.0, 0.1]}, ) + with self.assertRaisesRegex(TypeError, "tag_name"): + Reader(metadata, tag_name="snapshot-b") dataset = pmm.PaimonLeRobotDataset(reader) self.assertIsInstance(dataset.reader, pmm.PaimonDatasetReader) @@ -428,7 +423,7 @@ def close(self): self.assertEqual([((0, 1, 2), tuple(info["features"]))], reader.calls) - self.assertIsNone(dataset.tag_name) + self.assertFalse(hasattr(dataset, "tag_name")) self.assertEqual("dataset-version-12", dataset.meta.revision) self.assertEqual((2,), dataset.features["observation.state"]["shape"]) self.assertEqual([1.0], dataset.meta.stats["action"]["mean"].tolist()) @@ -879,6 +874,8 @@ def test_dataset_return_uint8_requires_bool(self): "pypaimon.multimodal.lerobot.dataset." "_load_dataset", return_value=loaded), patch( + "pypaimon.multimodal.lerobot.dataset._target_schema", + return_value=pa.schema([])), patch( "pypaimon.multimodal.lerobot.dataset.sys.version_info", (3, 10)): for invalid in (0, 1, None, "true"): From 5e3929733166489849c98ec07569567b36a2eb88 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 06:34:24 -0700 Subject: [PATCH 14/16] [python] Make custom reader schema optional --- docs/docs/pypaimon/lerobot.md | 7 +++---- paimon-python/pypaimon/multimodal/lerobot/dataset.py | 7 ++++--- paimon-python/pypaimon/tests/multimodal_lerobot_test.py | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index cb6fa3aae688..99110cdb5d7a 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -237,7 +237,6 @@ class CustomDatasetReader(PaimonDatasetReader): self._source = source super().__init__( metadata, - schema=source.schema, file_io=getattr(source, "file_io", None), **kwargs, ) @@ -251,6 +250,6 @@ reader = CustomDatasetReader( dataset = PaimonLeRobotDataset(reader) ``` -`source` exposes `schema`, `read_indices`, and optional `file_io`. -`PaimonDatasetReader` reuses the standard Episode, delta-window, media, and -Torch handling. +`source` exposes `read_indices` and optional `file_io`. Pass +`schema=source.schema` for eager schema validation. `PaimonDatasetReader` +reuses the standard Episode, delta-window, media, and Torch handling. diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index a999c42704d3..251dcb6ed079 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -91,7 +91,7 @@ def __init__( self, meta, *, - schema, + schema=None, file_io=None, episodes=None, image_transforms=None, @@ -103,10 +103,9 @@ def __init__( _resolved_meta=False): _require_dataset_python() metadata = meta if _resolved_meta else _reader_metadata(meta) - if not isinstance(schema, pa.Schema): + if schema is not None and not isinstance(schema, pa.Schema): raise TypeError( "PaimonDatasetReader schema must be a pyarrow.Schema.") - self.schema = schema self.file_io = file_io info = self._init_dataset( metadata, @@ -118,6 +117,8 @@ def __init__( video_backend, return_uint8, ) + schema = schema if schema is not None else _schema_from_info(info) + self.schema = schema projection, validation_context, subtasks = \ self._init_frame_contract( schema, info, self._validate_physical_metadata()) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index d5f2b47ee31b..746da31fce95 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -380,8 +380,7 @@ class Reader(pmm.PaimonDatasetReader): def __init__(self, metadata, **kwargs): self.calls = [] self.closed = False - super().__init__( - metadata, schema=_schema_from_info(info), **kwargs) + super().__init__(metadata, **kwargs) def read_indices(self, indices, columns): self.calls.append((indices, columns)) @@ -417,6 +416,7 @@ def close(self): dataset = pmm.PaimonLeRobotDataset(reader) self.assertIsInstance(dataset.reader, pmm.PaimonDatasetReader) + self.assertEqual(_schema_from_info(info), dataset.reader.schema) self.assertIsNone(dataset.reader.absolute_to_relative_idx) self.assertTrue(repr(dataset).startswith("PaimonLeRobotDataset(")) sample, _ = dataset.__getitems__([1, 2]) From 87345950dd358f022d44d6e80fd21a0f89114be1 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 07:04:42 -0700 Subject: [PATCH 15/16] [python] Validate custom reader batches --- .../pypaimon/multimodal/lerobot/dataset.py | 32 ++++++++++++++++--- .../pypaimon/tests/multimodal_lerobot_test.py | 21 ++++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 251dcb6ed079..f8b727843879 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -130,9 +130,8 @@ def __init__( def read_indices(self, indices, columns): """Return one row per requested absolute index as a PyArrow Table. - ``indices`` are unique; result order is unrestricted. Returned rows - must contain every requested column without missing or duplicate - indices. + ``indices`` are unique; result order is unrestricted. Result columns + must match ``schema`` and contain every requested index exactly once. """ def _validate_physical_metadata(self): @@ -216,6 +215,11 @@ def _init_metadata(self): def _init_episodes(self, episodes): self._episode_ranges = _episode_ranges( self.meta, self._total_frames, self._total_episodes) + if (self._episode_ranges is None + and (self._total_frames or self._total_episodes)): + raise ValueError( + "LeRobot metadata must define episodes for a non-empty " + "dataset.") self._episode_ends = [end for _, end in self._episode_ranges] \ if self._episode_ranges is not None else None self.episodes = _selected_episodes(episodes, self._total_episodes) @@ -493,6 +497,7 @@ def _read_rows(self, indices, projection): rows, projection, indices, + self.schema, self._validation_context, self.tolerance_s, self._features, @@ -635,6 +640,16 @@ def __getitem__(self, index): def __getitems__(self, indices): return self.reader.get_items(indices) + @property + def return_uint8(self): + return self.reader.return_uint8 + + @return_uint8.setter + def return_uint8(self, value): + if not isinstance(value, bool): + raise TypeError("return_uint8 must be a boolean.") + self.reader.return_uint8 = value + def set_image_transforms(self, image_transforms): self.reader.set_image_transforms(image_transforms) @@ -1041,8 +1056,8 @@ def _validate_reader_schema(expected_schema, actual_schema, source): def _read_reader_rows( - reader, projection, indices, validation_context, tolerance_s, - features): + reader, projection, indices, expected_schema, validation_context, + tolerance_s, features): values = reader.read_indices(tuple(indices), tuple(projection)) if not isinstance(values, pa.Table): raise TypeError( @@ -1052,6 +1067,13 @@ def _read_reader_rows( raise ValueError( "PaimonDatasetReader result is missing fields: %s" % sorted(missing)) + for name in projection: + expected_type = expected_schema.field(name).type + actual_type = values.schema.field(name).type + if actual_type != expected_type: + raise ValueError( + "PaimonDatasetReader field %s expects %s, found %s." + % (name, expected_type, actual_type)) rows = _arrow_rows(values.select(projection), features) expected = set(indices) result = {} diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 746da31fce95..5bf5ecb199cc 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -407,6 +407,10 @@ def close(self): "tasks": ["pick"], "stats": {"action": {"mean": [1.0]}}, } + missing_episodes = dict(metadata) + missing_episodes.pop("episodes") + with self.assertRaisesRegex(ValueError, "must define episodes"): + Reader(missing_episodes) reader = Reader( metadata, delta_timestamps={"action": [-0.1, 0.0, 0.1]}, @@ -435,6 +439,23 @@ def close(self): sample["action"], torch.tensor([0.0, 1.0, 2.0])) self.assertEqual([False, False, False], sample["action_is_pad"].tolist()) + dataset.return_uint8 = True + self.assertTrue(reader.return_uint8) + with self.assertRaisesRegex(TypeError, "return_uint8"): + dataset.return_uint8 = 1 + + wrong_schema = reader.schema.set( + reader.schema.get_field_index("action"), + pa.field("action", pa.int64()), + ) + with patch.object(reader, "read_indices") as read: + read.return_value = pa.Table.from_pylist( + [{name: rows[0][name] for name in info["features"]}], + schema=wrong_schema, + ) + with self.assertRaisesRegex( + ValueError, "field action expects float, found int64"): + dataset[0] dataset.close() self.assertTrue(reader.closed) From 1cf87f1cad39d8e49e7afb79f4ee878153c4678a Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 07:18:39 -0700 Subject: [PATCH 16/16] [python] Proxy LeRobot image transforms --- paimon-python/pypaimon/multimodal/lerobot/dataset.py | 8 ++++++++ paimon-python/pypaimon/tests/multimodal_lerobot_test.py | 7 +++++++ 2 files changed, 15 insertions(+) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index f8b727843879..6abbd34babac 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -650,6 +650,14 @@ def return_uint8(self, value): raise TypeError("return_uint8 must be a boolean.") self.reader.return_uint8 = value + @property + def image_transforms(self): + return self.reader.image_transforms + + @image_transforms.setter + def image_transforms(self, value): + self.reader.set_image_transforms(value) + def set_image_transforms(self, image_transforms): self.reader.set_image_transforms(image_transforms) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 5bf5ecb199cc..a9790b4e1607 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -443,6 +443,13 @@ def close(self): self.assertTrue(reader.return_uint8) with self.assertRaisesRegex(TypeError, "return_uint8"): dataset.return_uint8 = 1 + image_transforms = Mock() + dataset.image_transforms = image_transforms + self.assertIs(image_transforms, reader.image_transforms) + dataset.image_transforms = None + self.assertIsNone(reader.image_transforms) + with self.assertRaisesRegex(TypeError, "image_transforms"): + dataset.image_transforms = 1 wrong_schema = reader.schema.set( reader.schema.get_field_index("action"),