diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index b8ecc00913b6..216570ec5e30 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -204,7 +204,7 @@ jobs: run: | python -m pip install --upgrade pip pip install torch --index-url https://download.pytorch.org/whl/cpu - python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 datasketches + python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 datasketches 'datasets>=4,<4.1' python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/ - name: Run lint-python.sh shell: bash diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index c50acc7eb1d9..e7793ca9bedd 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -661,6 +661,24 @@ The one-time importer requires a new target table. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. +For map-style training, read a published version directly from Paimon: + +```python +from torch.utils.data import DataLoader +from pypaimon.multimodal import PaimonLeRobotDataset + +dataset = PaimonLeRobotDataset( + conn.get_table("robot_data"), + version_id=version_id, +) +loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) +``` + +If `version_id` is omitted, the latest published version is used. Metadata is +available through `dataset.meta`. Payload columns remain lazy. Pass +`index_mapping=dataset.index_mapping` to readers of the same table snapshot to +avoid rescanning the control columns. + ## Overwrite `overwrite` accepts the same input formats as `add` and replaces existing data diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index 53717d0155ea..584edc8431cf 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -29,6 +29,7 @@ Hdf5File, Hdf5LoadResult, ) +from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset from pypaimon.multimodal.rosbag import ( RosbagLoadResult, RosbagSource, @@ -60,6 +61,7 @@ "MultimodalTable", "NoSuchKey", "ObjectInfo", + "PaimonLeRobotDataset", "PutObjectResult", "RosbagLoadResult", "RosbagSource", diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index a40f2a8ccef0..40e6b2dc3ce6 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -14,11 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""One-time LeRobot Dataset v3 import into a multimodal Paimon table.""" +"""LeRobot Dataset v3 integration for multimodal Paimon tables.""" from pypaimon.multimodal.lerobot.api import load_from_lerobot +from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset __all__ = [ + "PaimonLeRobotDataset", "load_from_lerobot", ] diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py new file mode 100644 index 000000000000..7383ff3f64f4 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -0,0 +1,1049 @@ +# 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. + +"""LeRobot-compatible map-style reads from a multimodal Paimon table.""" + +import bisect +import io +import json +import math +import operator +from array import array + +import pyarrow as pa + +from pypaimon.catalog.catalog_exception import TagNotExistException +from pypaimon.multimodal.lerobot.metadata import ( + _companion_table_identifiers, + _restore_pandas_metadata, + _tag_snapshot_id, +) +from pypaimon.multimodal.lerobot.schema import ( + _feature_shape, + _require_v3, + _schema_from_info, + _validate_lerobot_schema, +) +from pypaimon.multimodal.table import _target_schema, _time_travel_table + + +_TORCH_DTYPE_NAMES = { + "bool": "bool", + "boolean": "bool", + "int8": "int8", + "int16": "int16", + "int32": "int32", + "int64": "int64", + "uint8": "uint8", + "uint16": "uint16", + "uint32": "uint32", + "float16": "float16", + "float32": "float32", + "float64": "float64", +} + + +class _LeRobotIndexMapping: + """Reusable semantic-index mapping bound to one table snapshot.""" + + def __init__( + self, + table_identifier, + snapshot_id, + metadata_signature, + tolerance_s, + positions): + self._table_identifier = table_identifier + self._snapshot_id = snapshot_id + self._metadata_signature = metadata_signature + self._tolerance_s = tolerance_s + self._positions = positions + + +class PaimonLeRobotDataset: + """Map-style LeRobot reader backed by Paimon's lazy Torch dataset. + + The published version and its LeRobot metadata are resolved from the + Paimon table group and remain available through :attr:`meta`. + + Set ``return_uint8=True`` to keep 8-bit images in their decoded + ``torch.uint8`` representation instead of normalizing them to float32. + Higher-bit-depth images retain the existing float32 behavior. + """ + + def __init__( + self, + table, + *, + version_id=None, + episodes=None, + image_transforms=None, + delta_timestamps=None, + tolerance_s=1e-4, + index_mapping=None, + blob_parallelism=16, + return_uint8=False): + raw_table, self.meta, self.version_id = _load_published_version( + table, version_id) + self.repo_id = self.meta.repo_id + self.image_transforms = image_transforms + self.delta_timestamps = delta_timestamps + self.tolerance_s = float(tolerance_s) + if not math.isfinite(self.tolerance_s) or self.tolerance_s < 0: + raise ValueError("tolerance_s must be finite and non-negative.") + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + if not isinstance(return_uint8, bool): + raise TypeError("return_uint8 must be a boolean.") + self.return_uint8 = return_uint8 + if image_transforms is not None and not callable(image_transforms): + raise TypeError("image_transforms must be callable or None.") + + info = dict(_metadata_member(self.meta, "info", {})) + _require_v3(info, self.repo_id) + self._features = dict( + _metadata_member(self.meta, "features", info.get("features"))) + if not self._features: + raise ValueError("LeRobot metadata must define features.") + self._image_keys = [ + name for name, feature in self._features.items() + if feature.get("dtype") == "image" + ] + video_keys = [ + name for name, feature in self._features.items() + if feature.get("dtype") == "video" + ] + if video_keys: + raise NotImplementedError( + "PaimonLeRobotDataset currently supports image-backed " + "features only; video features are not yet supported: %s" + % video_keys) + + self._total_frames = int( + _metadata_member( + self.meta, "total_frames", info.get("total_frames", -1))) + self._total_episodes = int( + _metadata_member( + self.meta, "total_episodes", info.get("total_episodes", -1))) + self._total_tasks = int( + _metadata_member( + self.meta, "total_tasks", info.get("total_tasks", -1))) + if self._total_frames < 0 or self._total_episodes < 0: + raise ValueError( + "LeRobot metadata must define total_frames and " + "total_episodes.") + if self._total_tasks < 0: + raise ValueError("LeRobot metadata must define total_tasks.") + + self._episode_ranges = _episode_ranges( + self.meta, self._total_frames, self._total_episodes) + 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) + if self.episodes is not None and self._episode_ranges is None: + raise ValueError("Episode selection requires episode metadata.") + self._selected_ranges = None + if self.episodes is not None: + # LeRobot exposes the caller's episode order but its Parquet filter + # returns frames in their stored dataset order. + range_episodes = sorted(self.episodes) + self._selected_ranges = [ + self._episode_ranges[index] for index in range_episodes + ] + self._selected_ends = [] + size = 0 + for begin, end in self._selected_ranges: + size += end - begin + self._selected_ends.append(size) + + self._fps = int( + _metadata_member(self.meta, "fps", info.get("fps", 0))) + if self._fps <= 0: + raise ValueError("LeRobot metadata fps must be positive.") + self._delta_indices = _delta_indices( + delta_timestamps, + self._fps, + self.tolerance_s, + self._features, + ) + if self._delta_indices and self._episode_ranges is None: + raise ValueError("delta_timestamps requires episode metadata.") + + 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) + control_contract = _control_contract( + 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)) + + self._dataset, splits, read_table, snapshot_id = _lazy_torch_dataset( + raw_table, projection) + if len(self._dataset) != self._total_frames: + raise ValueError( + "Paimon table has %d rows but metadata declares %d frames." + % (len(self._dataset), self._total_frames)) + table_identifier = str(table.identifier) + if index_mapping is None: + self._index_mapping = _semantic_index_mapping( + read_table, + splits, + self._total_frames, + table_identifier, + snapshot_id, + control_contract, + self.tolerance_s, + ) + else: + self._index_mapping = _reuse_index_mapping( + index_mapping, + table_identifier, + snapshot_id, + control_contract["signature"], + self._total_frames, + self.tolerance_s, + ) + self._index_positions = self._index_mapping._positions + self._file_io = read_table.file_io + self._task_names = control_contract["task_names"] + self._subtask_names = control_contract["subtask_names"] + self._delta_dataset = None + if self._delta_indices: + delta_projection = ["index"] + [ + key for key in self._delta_indices if key != "index" + ] + self._delta_dataset = _lazy_torch_dataset_for_splits( + read_table, delta_projection, splits) + + @property + def features(self): + return self._features + + @property + def fps(self): + return self._fps + + @property + def index_mapping(self): + """Mapping reusable by another reader of the same table snapshot.""" + return self._index_mapping + + @property + def num_frames(self): + if self.episodes is None: + return self._total_frames + return self._selected_ends[-1] if self._selected_ends else 0 + + @property + def num_episodes(self): + return self._total_episodes if self.episodes is None \ + else len(self.episodes) + + def __len__(self): + return self.num_frames + + 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): + relative = [_normalize_index(index, len(self)) for index in indices] + if not relative: + return [] + absolute = [self._absolute_index(index) for index in relative] + plans = [self._plan(index) for index in absolute] + + base_indices = sorted(set(absolute)) + base_rows = _read_rows( + self._dataset, base_indices, self._index_positions) + delta_indices = sorted({ + position + for plan in plans + for positions in plan["windows"].values() + for position in positions + if position not in base_rows + }) + delta_rows = _read_rows( + self._delta_dataset, delta_indices, self._index_positions) \ + if delta_indices else {} + + _resolve_image_blobs( + self._file_io, + [base_rows, delta_rows], + self._image_keys, + self.blob_parallelism, + ) + _materialize_labels( + base_rows, self._task_names, self._subtask_names) + converted = { + position: _torch_row( + row, self._features, self.return_uint8) + for position, row in base_rows.items() + } + converted.update({ + position: _torch_row( + row, self._features, self.return_uint8) + for position, row in delta_rows.items() + }) + + import torch + duplicates = _duplicate_indices(plans) + result = [] + for plan in plans: + item = dict(converted[plan["index"]]) + if plan["index"] in duplicates: + item = { + key: value.clone() if torch.is_tensor(value) else value + for key, value in item.items() + } + for key, positions in plan["windows"].items(): + item[key] = torch.stack([ + converted[position][key] for position in positions + ]) + item.update(plan["padding"]) + if self.image_transforms is not None: + for key in self._image_keys: + item[key] = self.image_transforms(item[key]) + result.append(item) + return result + + def set_image_transforms(self, image_transforms): + if image_transforms is not None and not callable(image_transforms): + raise TypeError("image_transforms must be callable or None.") + self.image_transforms = image_transforms + + def clear_image_transforms(self): + self.image_transforms = None + + def _absolute_index(self, index): + if self._selected_ranges is None: + return index + range_index = bisect.bisect_right(self._selected_ends, index) + previous_end = self._selected_ends[range_index - 1] \ + if range_index else 0 + return self._selected_ranges[range_index][0] + index - previous_end + + def _plan(self, index): + windows = {} + padding = {} + if self._delta_indices: + episode = bisect.bisect_right(self._episode_ends, index) + begin, end = self._episode_ranges[episode] + import torch + for key, deltas in self._delta_indices.items(): + windows[key] = [ + min(max(index + delta, begin), end - 1) + for delta in deltas + ] + padding["%s_is_pad" % key] = torch.BoolTensor([ + not begin <= index + delta < end for delta in deltas + ]) + return {"index": index, "windows": windows, "padding": padding} + + def __repr__(self): + return ( + "%s(repo_id=%r, episodes=%d, frames=%d, features=%r)" + % (self.__class__.__name__, self.repo_id, self.num_episodes, + self.num_frames, list(self.features))) + + +class _PaimonLeRobotMetadata: + + def __init__( + self, repo_id, version_id, info, stats, episodes, tasks, + subtasks): + self.repo_id = repo_id + self.revision = str(version_id) + self.info = info + self.stats = stats + self.episodes = episodes + self.tasks = tasks + self.subtasks = subtasks + + def __getattr__(self, name): + info = self.__dict__.get("info", {}) + try: + return info[name] + except KeyError as error: + raise AttributeError(name) from error + + @property + def image_keys(self): + return [ + name for name, feature in self.features.items() + if feature["dtype"] == "image" + ] + + @property + def video_keys(self): + return [ + name for name, feature in self.features.items() + if feature["dtype"] == "video" + ] + + @property + def camera_keys(self): + return [ + name for name, feature in self.features.items() + if feature["dtype"] in ("image", "video") + ] + + @property + def names(self): + return { + name: feature.get("names") + for name, feature in self.features.items() + } + + @property + def shapes(self): + return { + name: tuple(feature["shape"]) + for name, feature in self.features.items() + } + + def get_task_index(self, task): + if task not in self.tasks.index: + return None + return int(self.tasks.loc[task].task_index) + + +def _load_published_version(table, version_id): + raw_table = getattr(table, "raw_table", None) + if raw_table is None: + raise TypeError("table must be a MultimodalTable.") + identifiers = _companion_table_identifiers(raw_table) + catalog = table.catalog + manifests = _read_arrow(catalog.get_table( + identifiers["versions"])).to_pylist() + version_id, manifest = _select_manifest(manifests, version_id) + tag = str(version_id) + + frames = _tagged_table(catalog, raw_table, tag) + episodes_table = _tagged_table( + catalog, catalog.get_table(identifiers["episodes"]), tag) + episodes = _episode_dataset(episodes_table) + tasks_table = _tagged_table( + catalog, catalog.get_table(identifiers["tasks"]), tag) + tasks = _component_dataframe(tasks_table, "task_index") + subtasks = None + if manifest["has_subtasks"]: + subtasks_table = _tagged_table( + catalog, catalog.get_table(identifiers["subtasks"]), tag) + subtasks = _component_dataframe(subtasks_table, "subtask_index") + + info = _json_object(manifest["info_json"], "info_json") + for feature in info.get("features", {}).values(): + feature["shape"] = tuple(feature["shape"]) + stats = None if manifest["stats_json"] is None else _numpy_stats( + _json_object(manifest["stats_json"], "stats_json")) + metadata = _PaimonLeRobotMetadata( + str(table.identifier), version_id, info, stats, episodes, tasks, + subtasks) + return frames, metadata, version_id + + +def _select_manifest(manifests, version_id): + if not manifests: + raise ValueError("Paimon LeRobot table has no published versions.") + if version_id is None: + version_id = max(row["version_id"] for row in manifests) + else: + try: + version_id = operator.index(version_id) + except TypeError as error: + raise ValueError( + "version_id must be an integer or None.") from error + if isinstance(version_id, bool): + raise ValueError("version_id must be an integer or None.") + matches = [ + row for row in manifests if row["version_id"] == version_id + ] + if len(matches) != 1: + raise ValueError( + "Paimon LeRobot version %d has %d manifest rows." + % (version_id, len(matches))) + return version_id, matches[0] + + +def _tagged_table(catalog, table, tag): + try: + snapshot_id = _tag_snapshot_id(catalog, table.identifier, tag) + except TagNotExistException: + snapshot_id = None + if snapshot_id is None: + raise ValueError( + "Paimon LeRobot component %s is missing tag %s." + % (table.identifier, tag)) + return _time_travel_table(table, tag_name=tag) + + +def _read_arrow(table, projection=None): + builder = table.new_read_builder() + if projection is not None: + builder = builder.with_projection(projection) + plan = builder.new_scan().plan() + return builder.new_read().to_arrow(plan.splits()) + + +def _episode_dataset(table): + try: + from datasets import Dataset + except ImportError as error: + raise ImportError( + "PaimonLeRobotDataset requires datasets from " + "'pypaimon[lerobot]'.") from error + + projection = [ + name for name in _target_schema(table).names + if not name.startswith("stats/") + ] + data = _read_arrow(table, projection).sort_by("episode_index") + return Dataset(data) + + +def _component_dataframe(table, index_field): + data = _read_arrow(table).sort_by(index_field) + return _restore_pandas_metadata(table, data).to_pandas() + + +def _json_object(value, field): + try: + result = json.loads(value) + except (TypeError, ValueError) as error: + raise ValueError( + "Paimon LeRobot manifest %s is invalid JSON." % field + ) from error + if not isinstance(result, dict): + raise ValueError( + "Paimon LeRobot manifest %s must contain an object." % field) + return result + + +def _numpy_stats(value): + if isinstance(value, dict): + return {name: _numpy_stats(item) for name, item in value.items()} + import numpy as np + return np.array(value) + + +def _metadata_member(metadata, name, default=None): + value = getattr(metadata, name, None) + return default if value is None else value + + +def _episode_row(episodes, ordinal): + return episodes.iloc[ordinal] if hasattr(episodes, "iloc") \ + else episodes[ordinal] + + +def _episode_ranges(metadata, total_frames, total_episodes): + episodes = _metadata_member(metadata, "episodes") + if episodes is None: + return None + if len(episodes) != total_episodes: + raise ValueError( + "LeRobot episode metadata contains %d rows, expected %d." + % (len(episodes), total_episodes)) + ranges = [] + expected = 0 + for ordinal in range(total_episodes): + row = _episode_row(episodes, ordinal) + try: + index = operator.index(row["episode_index"]) + begin = operator.index(row["dataset_from_index"]) + end = operator.index(row["dataset_to_index"]) + length = operator.index(row["length"]) + except (KeyError, TypeError) as error: + raise ValueError( + "LeRobot episode %d metadata must contain integer controls." + % ordinal) from error + if index != ordinal: + raise ValueError( + "LeRobot episode row %d has episode_index=%d." + % (ordinal, index)) + if begin != expected or end <= begin: + raise ValueError( + "LeRobot episode %d has invalid frame range [%d, %d)." + % (ordinal, begin, end)) + if length != end - begin: + raise ValueError( + "LeRobot episode %d has length %d, expected %d." + % (ordinal, length, end - begin)) + ranges.append((begin, end)) + expected = end + if expected != total_frames: + raise ValueError( + "LeRobot episode ranges cover %d frames, expected %d." + % (expected, total_frames)) + return ranges + + +def _validate_component_metadata(features, total_tasks, tasks, subtasks): + task_count = 0 if tasks is None else len(tasks) + if task_count != total_tasks: + raise ValueError( + "LeRobot task metadata contains %d rows, expected %d." + % (task_count, total_tasks)) + has_subtasks = subtasks is not None + has_subtask_feature = "subtask_index" in features + if has_subtasks != has_subtask_feature: + raise ValueError( + "Paimon LeRobot manifest has_subtasks does not match the " + "subtask_index feature.") + + +def _control_contract( + metadata, episode_ranges, fps, tasks, subtasks, timestamp_type): + task_names = _index_names(tasks, "task_index") + subtask_names = _index_names(subtasks, "subtask_index") + episode_tasks = _episode_tasks(metadata, len(episode_ranges)) \ + if episode_ranges is not None else None + signature = ( + tuple(episode_ranges) if episode_ranges is not None else None, + fps, + tuple(sorted(task_names.items())) if task_names is not None else None, + tuple(sorted(subtask_names.items())) if subtask_names is not None + else None, + episode_tasks, + str(timestamp_type), + ) + return { + "episode_ranges": episode_ranges, + "episode_ends": ( + [end for _, end in episode_ranges] + if episode_ranges is not None else None), + "fps": fps, + "task_names": task_names, + "subtask_names": subtask_names, + "episode_tasks": episode_tasks, + "timestamp_type": timestamp_type, + "signature": signature, + } + + +def _index_names(values, index_field): + if values is None or len(values) == 0: + return None + if not hasattr(values, "iterrows"): + return { + index: str(value) for index, value in enumerate(values) + } + result = {} + for name, row in values.iterrows(): + try: + index = operator.index(row[index_field]) + except (KeyError, TypeError) as error: + raise ValueError( + "LeRobot %s metadata must contain integer indices." + % index_field) from error + if index in result: + raise ValueError( + "LeRobot %s metadata contains duplicate index %d." + % (index_field, index)) + result[index] = str(name) + if sorted(result) != list(range(len(result))): + raise ValueError( + "LeRobot %s metadata indices must be contiguous." + % index_field) + return result + + +def _episode_tasks(metadata, total_episodes): + episodes = _metadata_member(metadata, "episodes") + if episodes is None: + return None + result = [] + for ordinal in range(total_episodes): + row = _episode_row(episodes, ordinal) + tasks = row.get("tasks") if hasattr(row, "get") else None + if tasks is None: + result.append(None) + elif isinstance(tasks, str): + result.append((tasks,)) + else: + result.append(tuple(sorted(str(task) for task in tasks))) + return tuple(result) + + +def _selected_episodes(episodes, total_episodes): + if episodes is None: + return None + selected = [] + seen = set() + for value in episodes: + try: + index = operator.index(value) + except TypeError as error: + raise ValueError( + "episodes must contain integer indices.") from error + if index < 0 or index >= total_episodes: + raise ValueError( + "episodes must contain indices in [0, %d)." % total_episodes) + if index in seen: + raise ValueError("episodes must not contain duplicate indices.") + seen.add(index) + selected.append(index) + return selected + + +def _delta_indices(delta_timestamps, fps, tolerance_s, features): + if delta_timestamps is None: + return None + if fps <= 0: + raise ValueError("LeRobot metadata fps must be positive.") + result = {} + for key, timestamps in delta_timestamps.items(): + if key not in features: + raise ValueError("Unknown LeRobot delta feature: %s" % key) + deltas = [] + for timestamp in timestamps: + index = round(float(timestamp) * fps) + if abs(float(timestamp) - index / fps) > tolerance_s: + raise ValueError( + "delta_timestamps for %s must be multiples of 1/%d." + % (key, fps)) + deltas.append(index) + result[key] = deltas + return result + + +def _lazy_torch_dataset(raw_table, projection): + from pypaimon.common.options.core_options import CoreOptions + read_table = raw_table.copy({ + CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true" + }) + builder = read_table.new_read_builder().with_projection(projection) + plan = builder.new_scan().plan() + splits = plan.splits() + return ( + _required_lazy_torch_dataset(builder.new_read(), splits), + splits, + read_table, + plan.snapshot_id, + ) + + +def _lazy_torch_dataset_for_splits(read_table, projection, splits): + builder = read_table.new_read_builder().with_projection(projection) + return _required_lazy_torch_dataset(builder.new_read(), splits) + + +def _required_lazy_torch_dataset(table_read, splits): + from pypaimon.read.datasource.torch_dataset import TorchDataset + return TorchDataset.lazy(table_read, splits) + + +def _semantic_index_mapping( + read_table, + splits, + size, + table_identifier, + snapshot_id, + control_contract, + tolerance_s): + projection = [ + "index", "episode_index", "frame_index", "timestamp", "task_index" + ] + if control_contract["subtask_names"] is not None: + projection.append("subtask_index") + index_dataset = _lazy_torch_dataset_for_splits( + read_table, projection, splits) + if len(index_dataset) != size: + raise ValueError( + "Paimon index contains %d rows, expected %d." + % (len(index_dataset), size)) + + positions = None + batch_size = 65536 + for begin in range(0, size, batch_size): + end = min(begin + batch_size, size) + rows = index_dataset.__getitems__(range(begin, end)) + if len(rows) != end - begin: + raise ValueError( + "Paimon index read returned %d rows for range [%d, %d)." + % (len(rows), begin, end)) + for offset, row in enumerate(rows): + physical = begin + offset + try: + index = operator.index(row["index"]) + except (KeyError, TypeError) as error: + raise ValueError( + "Paimon LeRobot index must contain integers.") from error + if index < 0 or index >= size: + raise ValueError( + "Paimon LeRobot index %d is outside [0, %d)." + % (index, size)) + _validate_control_row( + row, index, control_contract, tolerance_s) + if positions is None and index == physical: + continue + if positions is None: + positions = array("q", [-1]) * size + for previous in range(physical): + positions[previous] = previous + if positions[index] >= 0: + raise ValueError( + "Paimon LeRobot index contains duplicate value %d." + % index) + positions[index] = physical + + if positions is not None and any(position < 0 for position in positions): + raise ValueError("Paimon LeRobot index is not contiguous.") + return _LeRobotIndexMapping( + table_identifier, + snapshot_id, + control_contract["signature"], + tolerance_s, + range(size) if positions is None else positions, + ) + + +def _reuse_index_mapping( + mapping, + table_identifier, + snapshot_id, + metadata_signature, + size, + tolerance_s): + if not isinstance(mapping, _LeRobotIndexMapping): + raise TypeError( + "index_mapping must come from PaimonLeRobotDataset.index_mapping.") + if mapping._table_identifier != table_identifier: + raise ValueError("index_mapping belongs to a different Paimon table.") + if mapping._snapshot_id != snapshot_id: + raise ValueError("index_mapping belongs to a different Paimon snapshot.") + if mapping._metadata_signature != metadata_signature: + raise ValueError("index_mapping belongs to different LeRobot metadata.") + if len(mapping._positions) != size: + raise ValueError("index_mapping has an incompatible frame count.") + if mapping._tolerance_s > tolerance_s: + raise ValueError( + "index_mapping was validated with a looser tolerance_s.") + return mapping + + +def _validate_control_row(row, index, contract, tolerance_s): + ranges = contract["episode_ranges"] + if ranges is None: + return + episode = bisect.bisect_right(contract["episode_ends"], index) + begin, _ = ranges[episode] + frame = index - begin + _validate_int_control(row, "episode_index", index, episode) + _validate_int_control(row, "frame_index", index, frame) + + expected_timestamp = pa.scalar( + frame / contract["fps"], type=contract["timestamp_type"]).as_py() + try: + timestamp = float(row["timestamp"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + "Paimon timestamp at LeRobot index %d must be numeric." % index + ) from error + if not math.isfinite(timestamp) \ + or abs(timestamp - expected_timestamp) > tolerance_s: + raise ValueError( + "Paimon timestamp at LeRobot index %d is %r, metadata expects %r." + % (index, timestamp, expected_timestamp)) + + task_names = contract["task_names"] + task_index = _control_int(row, "task_index", index) + task = None if task_names is None else task_names.get(task_index) + if task is None: + raise ValueError( + "Paimon task_index at LeRobot index %d is absent from metadata: " + "%r." % (index, task_index)) + episode_tasks = contract["episode_tasks"] + allowed = episode_tasks[episode] if episode_tasks is not None else None + if allowed is not None and task not in allowed: + raise ValueError( + "Paimon task at LeRobot index %d is not assigned to episode %d: " + "%r." % (index, episode, task)) + + subtask_names = contract["subtask_names"] + if subtask_names is not None: + subtask_index = _control_int(row, "subtask_index", index) + if subtask_index not in subtask_names: + raise ValueError( + "Paimon subtask_index at LeRobot index %d is absent from " + "metadata: %r." % (index, subtask_index)) + + +def _validate_int_control(row, field, index, expected): + actual = _control_int(row, field, index) + if actual != expected: + raise ValueError( + "Paimon %s at LeRobot index %d is %r, metadata expects %r." + % (field, index, actual, expected)) + + +def _control_int(row, field, index): + try: + return operator.index(row[field]) + except (KeyError, TypeError) as error: + raise ValueError( + "Paimon %s at LeRobot index %d must be an integer." + % (field, index)) from error + + +def _read_rows(dataset, indices, index_positions): + if not indices: + return {} + positions = [ + index_positions[index] for index in indices + ] + rows = dataset.__getitems__(positions) + result = {} + for index, row in zip(indices, rows): + if int(row["index"]) != index: + raise ValueError( + "Paimon row mapped to LeRobot index %d contains index=%r." + % (index, row["index"])) + result[index] = row + return result + + +def _duplicate_indices(plans): + seen = set() + duplicates = set() + for plan in plans: + index = plan["index"] + if index in seen: + duplicates.add(index) + seen.add(index) + return duplicates + + +def _resolve_image_blobs( + file_io, row_groups, image_keys, parallelism): + from pypaimon.multimodal.blob_read import fetch_blob_bodies + + values = {key: [] for key in image_keys} + targets = {key: [] for key in image_keys} + for rows in row_groups: + for row in rows.values(): + for key in image_keys: + if key in row: + targets[key].append(row) + values[key].append(row[key]) + used = [key for key in image_keys if values[key]] + if not used: + return + bodies = fetch_blob_bodies( + file_io, values, used, parallelism) + for key in used: + for row, body in zip(targets[key], bodies[key]): + row[key] = body + + +def _materialize_labels(rows, task_names, subtask_names): + for row in rows.values(): + task_index = operator.index(row["task_index"]) + row["task"] = task_names[task_index] + if subtask_names is not None: + subtask_index = operator.index(row["subtask_index"]) + row["subtask"] = subtask_names[subtask_index] + + +def _torch_row(row, features, return_uint8=False): + import torch + + result = dict(row) + for key, feature in features.items(): + if key not in result: + continue + value = result[key] + if feature.get("dtype") == "image": + result[key] = _image_tensor( + value, feature, return_uint8=return_uint8) + elif feature.get("dtype") != "string": + dtype = getattr(torch, _TORCH_DTYPE_NAMES[feature.get("dtype")]) + result[key] = torch.tensor(value, dtype=dtype) + return result + + +def _image_tensor(payload, feature, return_uint8=False): + if payload is None: + raise ValueError("LeRobot image feature contains a null frame.") + import numpy as np + import torch + try: + from PIL import Image + except ImportError as error: + raise ImportError( + "PaimonLeRobotDataset requires Pillow from " + "'pypaimon[lerobot]'.") from error + + expected_shape = _feature_shape(feature, "image") + if len(expected_shape) != 3: + raise ValueError( + "LeRobot image feature must have three dimensions.") + names = feature.get("names") or [] + payload_shape = expected_shape[1:] + expected_shape[:1] \ + if names and names[0] in ("channel", "channels") \ + else expected_shape + with Image.open(io.BytesIO(payload)) as image: + array = np.array(image, copy=True) + if array.ndim == 2: + array = array[:, :, None] + if array.shape != payload_shape: + raise ValueError( + "LeRobot image payload has shape %s, expected %s." + % (array.shape, payload_shape)) + normalize = array.dtype == np.uint8 + tensor = torch.from_numpy(array).permute(2, 0, 1) + if normalize and return_uint8: + return tensor + # Preserve high-bit-depth and floating-point images in native units. + tensor = tensor.float() + return tensor.div_(255) if normalize else tensor + + +def _normalize_index(index, size): + index = operator.index(index) + if index < 0: + index += size + if index < 0 or index >= size: + raise IndexError("PaimonLeRobotDataset index out of range") + return index + + +def _positive_int(value, name): + try: + value = operator.index(value) + except TypeError as error: + raise ValueError("%s must be a positive integer." % name) from error + if isinstance(value, bool) or value <= 0: + raise ValueError("%s must be a positive integer." % name) + return value diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index de4bb2cacef3..48085976ad9e 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -189,7 +189,11 @@ class TorchDataset(Dataset): rows into Python objects. """ - def __init__(self, table_read: TableRead, splits: List[Split]): + def __init__( + self, + table_read: TableRead, + splits: List[Split], + ): """ Initialize TorchDataset. @@ -197,6 +201,16 @@ def __init__(self, table_read: TableRead, splits: List[Split]): table_read: TableRead instance for reading data splits: List of splits to read """ + self._initialize(table_read, splits, require_lazy=False) + + @classmethod + def lazy(cls, table_read: TableRead, splits: List[Split]): + """Create a dataset which rejects materialization fallbacks.""" + dataset = cls.__new__(cls) + dataset._initialize(table_read, splits, require_lazy=True) + return dataset + + def _initialize(self, table_read, splits, require_lazy): self.table_read = table_read self.splits = splits self._data = None @@ -221,8 +235,16 @@ def __init__(self, table_read: TableRead, splits: List[Split]): SpecialFields.ROW_ID.name).combine_chunks() if pc.count_distinct(self._row_ids).as_py() != len( self._row_ids): + if require_lazy: + raise ValueError( + "Lazy TorchDataset requires visible and unique " + "_ROW_ID values.") self._materialize() else: + if require_lazy: + raise ValueError( + "Lazy TorchDataset requires row tracking, data " + "evolution, a visible _ROW_ID, and supported splits.") self._materialize() def _supports_lazy_row_id_read(self) -> bool: @@ -232,7 +254,9 @@ def _supports_lazy_row_id_read(self) -> bool: return False if self.table_read.include_row_kind: return False - if self.table_read.nested_name_paths: + if self.table_read.nested_name_paths and any( + len(path) > 1 + for path in self.table_read.nested_name_paths): return False if any(self._row_id_is_masked(split) for split in self.splits): return False diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 97c82b43d40a..d372ee9f94c6 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -16,6 +16,7 @@ import builtins from array import array +import io import json import shutil import sys @@ -37,6 +38,11 @@ from pypaimon.common.options import Options from pypaimon.multimodal.source_utils import _SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot +from pypaimon.multimodal.lerobot.dataset import ( + _image_tensor, + _selected_episodes, + _select_manifest, +) from pypaimon.multimodal.lerobot.metadata import ( _append_arrow_tables, _companion_identifier, @@ -97,6 +103,17 @@ def _catalog_arrow(connection, name): class LeRobotValidationTest(unittest.TestCase): + def test_dataset_selects_latest_or_requested_published_version(self): + manifests = [ + {"version_id": 1}, + {"version_id": 3}, + {"version_id": 2}, + ] + self.assertEqual((3, manifests[1]), _select_manifest(manifests, None)) + self.assertEqual((2, manifests[2]), _select_manifest(manifests, 2)) + with self.assertRaisesRegex(ValueError, "version 4 has 0"): + _select_manifest(manifests, 4) + def test_self_contained_import_rejects_table_branches(self): with self.assertRaisesRegex(ValueError, "does not support"): _managed_table_options("db.robot$branch_dev") @@ -109,6 +126,195 @@ def test_companion_identifier_preserves_quoted_components(self): self.assertEqual("db.name", identifier.get_database_name()) self.assertEqual("robot.data__tasks", identifier.get_table_name()) + def test_image_tensor_preserves_declared_channels(self): + try: + from PIL import Image + import torch + except ImportError as error: + self.skipTest(str(error)) + + cases = [ + ("L", np.full((4, 5), 64, dtype=np.uint8), + [4, 5, 1], [64]), + ("RGB", np.tile( + np.array([32, 64, 96], dtype=np.uint8), (4, 5, 1)), + [4, 5, 3], [32, 64, 96]), + ("RGBA", np.tile( + np.array([32, 64, 96, 128], dtype=np.uint8), (4, 5, 1)), + [4, 5, 4], [32, 64, 96, 128]), + ] + for mode, values, shape, expected in cases: + with self.subTest(mode=mode): + output = io.BytesIO() + Image.fromarray(values, mode=mode).save(output, format="PNG") + feature = {"dtype": "image", "shape": shape} + tensor = _image_tensor(output.getvalue(), feature) + uint8_tensor = _image_tensor( + output.getvalue(), feature, return_uint8=True) + + self.assertEqual(torch.float32, tensor.dtype) + self.assertEqual(torch.uint8, uint8_tensor.dtype) + self.assertEqual( + [shape[2], shape[0], shape[1]], list(tensor.shape)) + self.assertEqual(list(tensor.shape), list(uint8_tensor.shape)) + self.assertGreaterEqual(float(tensor.min()), 0.0) + self.assertLessEqual(float(tensor.max()), 1.0) + self.assertEqual(expected, uint8_tensor[:, 0, 0].tolist()) + torch.testing.assert_close( + tensor, uint8_tensor.float().div(255)) + + output = io.BytesIO() + Image.fromarray(cases[1][1], mode="RGB").save(output, format="PNG") + tensor = _image_tensor(output.getvalue(), { + "dtype": "image", + "shape": [3, 4, 5], + "names": ["channels", "height", "width"], + }) + self.assertEqual([3, 4, 5], list(tensor.shape)) + + uint8_tensor = _image_tensor( + output.getvalue(), + { + "dtype": "image", + "shape": [3, 4, 5], + "names": ["channels", "height", "width"], + }, + return_uint8=True, + ) + self.assertEqual(torch.uint8, uint8_tensor.dtype) + self.assertEqual([32, 64, 96], uint8_tensor[:, 0, 0].tolist()) + original = uint8_tensor.clone() + uint8_tensor.zero_() + reread = _image_tensor( + output.getvalue(), + { + "dtype": "image", + "shape": [3, 4, 5], + "names": ["channels", "height", "width"], + }, + return_uint8=True, + ) + self.assertTrue(torch.equal(original, reread)) + self.assertNotEqual(uint8_tensor.data_ptr(), reread.data_ptr()) + + depth = np.array([ + [0, 1000, 4095], + [8192, 32768, 65535], + ], dtype=np.uint16) + output = io.BytesIO() + Image.fromarray(depth).save(output, format="PNG") + feature = { + "dtype": "image", + "shape": [2, 3, 1], + "info": { + "is_depth_map": True, + "depth_unit": "mm", + }, + } + tensor = _image_tensor(output.getvalue(), feature) + uint8_requested = _image_tensor( + output.getvalue(), + feature, + return_uint8=True, + ) + self.assertEqual([1, 2, 3], list(tensor.shape)) + self.assertEqual(torch.float32, tensor.dtype) + self.assertEqual(depth.astype(np.float32).tolist(), tensor[0].tolist()) + self.assertEqual(torch.float32, uint8_requested.dtype) + self.assertTrue(torch.equal(tensor, uint8_requested)) + + def test_dataset_uint8_getitem_matches_getitems(self): + try: + from PIL import Image + import torch + except ImportError as error: + self.skipTest(str(error)) + + features = { + "index": {"dtype": "int64", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + "observation.left": { + "dtype": "image", "shape": [4, 5, 3]}, + "observation.wrist": { + "dtype": "image", "shape": [4, 5, 1]}, + } + + def jpeg(mode, values): + output = io.BytesIO() + Image.fromarray(values, mode=mode).save( + output, format="JPEG", quality=100) + return output.getvalue() + + rows = [] + for index in range(2): + rows.append({ + "index": index, + "task_index": 0, + "observation.left": jpeg( + "RGB", np.full((4, 5, 3), 40 + index, np.uint8)), + "observation.wrist": jpeg( + "L", np.full((4, 5), 80 + index, np.uint8)), + }) + + class Rows: + + def __getitems__(self, positions): + return [dict(rows[position]) for position in positions] + + dataset = object.__new__(pmm.PaimonLeRobotDataset) + dataset._total_frames = 2 + dataset.episodes = None + dataset._selected_ranges = None + dataset._delta_indices = {} + dataset._dataset = Rows() + dataset._index_positions = {0: 0, 1: 1} + dataset._delta_dataset = None + dataset._file_io = Mock() + dataset._image_keys = [ + "observation.left", "observation.wrist"] + dataset.blob_parallelism = 1 + dataset._task_names = ["task"] + dataset._subtask_names = None + dataset._features = features + dataset.return_uint8 = True + dataset.image_transforms = None + + with patch( + "pypaimon.multimodal.lerobot.dataset._resolve_image_blobs"): + single = dataset[1] + batched = dataset.__getitems__([1, 0]) + dataset.return_uint8 = False + normalized = dataset[1] + for key in dataset._image_keys: + self.assertEqual(torch.uint8, single[key].dtype) + self.assertTrue(torch.equal(single[key], batched[0][key])) + torch.testing.assert_close( + normalized[key], single[key].float().div(255)) + + def test_dataset_return_uint8_requires_bool(self): + published = ( + Mock(), + Mock(repo_id="pypaimon/invalid-return-uint8"), + 1, + ) + with patch( + "pypaimon.multimodal.lerobot.dataset." + "_load_published_version", + return_value=published): + for invalid in (0, 1, None, "true"): + with self.subTest(return_uint8=invalid): + with self.assertRaisesRegex( + TypeError, "return_uint8 must be a boolean"): + pmm.PaimonLeRobotDataset( + Mock(), return_uint8=invalid) + + def test_selected_episodes_preserves_caller_order(self): + self.assertEqual([1, 0], _selected_episodes([1, 0], 2)) + with self.assertRaisesRegex(ValueError, "duplicate"): + _selected_episodes([1, 1], 2) + with self.assertRaisesRegex(ValueError, "indices in"): + _selected_episodes([2], 2) + def test_dataset_open_never_downloads_videos(self): calls = [] @@ -1190,6 +1396,10 @@ def test_import_publishes_optional_subtasks(self): ) self.assertTrue(_catalog_rows( self.connection, "with_subtasks__versions")[0]["has_subtasks"]) + dataset = pmm.PaimonLeRobotDataset(frames, version_id=version_id) + self.assertEqual(["reach", "grasp"], list( + dataset.meta.subtasks.index)) + self.assertEqual("reach", dataset[0]["subtask"]) self.assertEqual( 1, self.connection.catalog.get_tag( @@ -1349,6 +1559,109 @@ def test_nonempty_dataset_cannot_publish_without_tasks(self): self.assertEqual([], _catalog_rows( self.connection, "missing_tasks__versions")) + def test_paimon_dataset_reads_lazy_batches_with_lerobot_metadata(self): + import torch + + version_id = self.connection.load_from_lerobot( + "training_data", self.image_source, batch_size=2) + table = self.connection.get_table("training_data") + dataset = pmm.PaimonLeRobotDataset( + table, + version_id=version_id, + delta_timestamps={"action": [-0.1, 0.0, 0.1]}, + blob_parallelism=3, + ) + + self.assertEqual(1, dataset.version_id) + self.assertEqual(5, len(dataset)) + self.assertEqual(2, dataset.num_episodes) + self.assertIsNotNone(dataset.meta.stats) + self.assertTrue(all( + isinstance(feature["shape"], tuple) + for feature in dataset.features.values() + )) + self.assertEqual(["pick", "place"], list(dataset.meta.tasks.index)) + episodes = dataset.meta.episodes + self.assertEqual(2, len(episodes)) + self.assertEqual([0, 2], list(episodes["dataset_from_index"])) + self.assertEqual([2, 5], list(episodes["dataset_to_index"])) + self.assertFalse(any( + name.startswith("stats/") + for name in episodes.column_names)) + self.assertIsNone(dataset._dataset._data) + + from pypaimon.multimodal.blob_read import fetch_blob_bodies + with patch( + "pypaimon.multimodal.blob_read.fetch_blob_bodies", + wraps=fetch_blob_bodies) as fetch: + last, first = dataset.__getitems__([4, 0]) + self.assertEqual(3, fetch.call_args.args[3]) + self.assertEqual("place", last["task"]) + self.assertEqual([3, 8, 10], list(last["observation.image"].shape)) + self.assertAlmostEqual( + 100.0 / 255.0, + float(last["observation.image"].mean()), + places=5, + ) + self.assertEqual( + [[1.0, -1.0], [2.0, -2.0], [2.0, -2.0]], + last["action"].tolist(), + ) + self.assertEqual([False, False, True], + last["action_is_pad"].tolist()) + self.assertEqual([True, False, False], + first["action_is_pad"].tolist()) + + uint8_dataset = pmm.PaimonLeRobotDataset( + table, + version_id=version_id, + index_mapping=dataset.index_mapping, + return_uint8=True, + ) + uint8_sample = uint8_dataset[4] + uint8_batch = uint8_dataset.__getitems__([4, 0]) + uint8_image = uint8_sample["observation.image"] + self.assertEqual("torch.uint8", str(uint8_image.dtype)) + self.assertEqual([3, 8, 10], list(uint8_image.shape)) + self.assertEqual(100.0, float(uint8_image.float().mean())) + self.assertTrue(torch.equal( + uint8_image, uint8_batch[0]["observation.image"])) + torch.testing.assert_close( + last["observation.image"], uint8_image.float().div(255)) + + reordered = pmm.PaimonLeRobotDataset( + table, + episodes=[1, 0], + index_mapping=dataset.index_mapping, + ) + self.assertEqual([1, 0], reordered.episodes) + self.assertEqual(5, len(reordered)) + self.assertEqual(0, int(reordered[0]["episode_index"])) + self.assertEqual(1, int(reordered[-1]["episode_index"])) + + table.add(pa.Table.from_pylist([{ + "index": 999, + "episode_index": 99, + "frame_index": 0, + "timestamp": 0.0, + "task_index": 0, + "observation.state": [0.0, 0.0, 0.0], + "observation.matrix": [[0.0, 0.0], [0.0, 0.0]], + "action": [0.0, 0.0], + "reward": 0.0, + "observation.image": _image_bytes( + np.zeros((8, 10, 3), dtype=np.uint8), self.temp_dir), + }], schema=_target_schema(table.raw_table))) + self.assertEqual(6, table.scan().to_arrow().num_rows) + + episode = pmm.PaimonLeRobotDataset( + table, + episodes=[1], + ) + self.assertEqual(3, len(episode)) + self.assertEqual(1, episode.num_episodes) + self.assertEqual(2, int(episode[0]["index"])) + def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" source_file_io = _RemoteLeRobotFileIO(self.image_source, source) diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 0dc269110d5e..01ba96656c75 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -34,6 +34,19 @@ from pypaimon import CatalogFactory, Schema from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.multimodal.lerobot.dataset import ( + PaimonLeRobotDataset, + _episode_ranges, + _validate_component_metadata, + _validate_control_row, +) +from pypaimon.multimodal.lerobot.metadata import ( + _VERSIONS_SCHEMA, + _append_arrow, + _managed_table_options, + _prepare_metadata_tables, +) +from pypaimon.multimodal.lerobot.schema import _schema_from_info from pypaimon.multimodal.table import MultimodalTable from pypaimon.read.datasource.torch_dataset import ( @@ -571,6 +584,298 @@ def test_non_streaming_row_tracking_reads_batches_lazily(self): sorted(actual_ids), ) + def test_paimon_lerobot_dataset_reuses_lazy_map_reader(self): + import pandas as pd + + 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]}, + 'observation.half': {'dtype': 'float16', 'shape': [2]}, + 'observation.double': {'dtype': 'float64', 'shape': [1]}, + 'observation.small': {'dtype': 'int8', 'shape': [1]}, + 'action': {'dtype': 'float32', 'shape': [2]}, + } + info = { + 'codebase_version': 'v3.0', + 'total_frames': 5, + 'total_episodes': 2, + 'total_tasks': 1, + 'fps': 10, + 'features': features, + } + arrow_schema = _schema_from_info(info) + identifier = 'default.test_paimon_lerobot_dataset' + table_options = { + 'data-evolution.enabled': 'true', + 'row-tracking.enabled': 'true', + 'blob-as-descriptor': 'true', + 'vector.file.format': 'parquet', + } + table_options.update(_managed_table_options(identifier)) + schema = Schema.from_pyarrow_schema( + arrow_schema, + partition_keys=['frame_index'], + options=table_options, + ) + self.catalog.create_table(identifier, schema, False) + raw_table = self.catalog.get_table(identifier) + table = MultimodalTable(self.catalog, identifier, raw_table) + table.add(pa.Table.from_pylist([ + { + 'index': index, + 'episode_index': 0 if index < 2 else 1, + 'frame_index': index if index < 2 else index - 2, + 'timestamp': (index if index < 2 else index - 2) / 10, + 'task_index': 0, + 'observation.state': [float(index), float(index + 1)], + 'observation.half': [float(index), float(index + 1)], + 'observation.double': 1e40, + 'observation.small': index, + 'action': [float(index), float(-index)], + } + for index in range(5) + ], schema=arrow_schema)) + episodes = pa.Table.from_pylist([ + { + 'episode_index': 0, + 'dataset_from_index': 0, + 'dataset_to_index': 2, + 'tasks': ['pick'], + 'length': 2, + }, + { + 'episode_index': 1, + 'dataset_from_index': 2, + 'dataset_to_index': 5, + 'tasks': ['pick'], + 'length': 3, + }, + ]) + tasks = pa.Table.from_pandas(pd.DataFrame( + {'task_index': [0]}, + index=pd.Index(['pick'], name='task'), + )) + component_tables = _prepare_metadata_tables( + SimpleNamespace(catalog=self.catalog), + raw_table, + { + 'episodes_schema': episodes.schema, + 'tasks_table': tasks, + 'subtasks_table': None, + }, + ) + raw_table.create_tag('1', snapshot_id=1) + for name, data in (('episodes', episodes), ('tasks', tasks)): + self.assertEqual(1, _append_arrow(component_tables[name], data)) + component_tables[name].create_tag('1', snapshot_id=1) + _append_arrow(component_tables['versions'], pa.Table.from_pylist([{ + 'version_id': 1, + 'info_json': json.dumps(info), + 'stats_json': json.dumps({'action': {}}), + 'has_subtasks': False, + }], schema=_VERSIONS_SCHEMA)) + + dataset = PaimonLeRobotDataset( + table, + episodes=[1, 0], + delta_timestamps={'action': [-0.1, 0.0, 0.1]}, + ) + + self.assertEqual(1, dataset.version_id) + self.assertEqual(['pick'], list(dataset.meta.tasks.index)) + self.assertIsNone(dataset._dataset._data) + self.assertEqual(5, len(dataset)) + first, last = dataset.__getitems__([0, 2]) + self.assertEqual(0, int(first['index'])) + self.assertEqual(torch.float16, first['observation.half'].dtype) + self.assertEqual(torch.float64, first['observation.double'].dtype) + self.assertEqual(1e40, first['observation.double'].item()) + self.assertEqual(torch.int8, first['observation.small'].dtype) + self.assertEqual( + [[0.0, 0.0], [0.0, 0.0], [1.0, -1.0]], + first['action'].tolist(), + ) + self.assertEqual([True, False, False], + first['action_is_pad'].tolist()) + self.assertEqual(2, int(last['index'])) + self.assertEqual([True, False, False], + last['action_is_pad'].tolist()) + self.assertEqual(3, int(dataset[3]['index'])) + self.assertEqual('pick', dataset[3]['task']) + restored = pickle.loads(pickle.dumps(dataset)) + self.assertEqual('pick', restored[0]['task']) + + duplicate_a, duplicate_b = dataset.__getitems__([0, 0]) + self.assertIsNot(duplicate_a['action'], duplicate_b['action']) + duplicate_a['action'].add_(1) + self.assertNotEqual( + duplicate_a['action'].tolist(), duplicate_b['action'].tolist()) + + batch = next(iter(DataLoader( + dataset, batch_size=2, num_workers=2, shuffle=False))) + self.assertEqual([0, 1], batch['index'].tolist()) + + with patch( + 'pypaimon.multimodal.lerobot.dataset.' + '_semantic_index_mapping') as rebuild: + reused = PaimonLeRobotDataset( + table, + episodes=[1], + index_mapping=pickle.loads( + pickle.dumps(dataset.index_mapping)), + blob_parallelism=3, + ) + rebuild.assert_not_called() + self.assertEqual(3, reused.blob_parallelism) + self.assertEqual(2, int(reused[0]['index'])) + with self.assertRaisesRegex(ValueError, 'looser tolerance_s'): + PaimonLeRobotDataset( + table, + index_mapping=dataset.index_mapping, + tolerance_s=1e-5, + ) + with self.assertRaisesRegex( + ValueError, 'blob_parallelism must be a positive integer'): + PaimonLeRobotDataset(table, blob_parallelism=0) + with self.assertRaisesRegex( + ValueError, "version 2 has 0 manifest rows"): + PaimonLeRobotDataset(table, version_id=2) + + component_tables['tasks'].delete_tag('1') + with self.assertRaisesRegex(ValueError, "missing tag 1"): + PaimonLeRobotDataset(table) + component_tables['tasks'].create_tag('1', snapshot_id=1) + + auth = TableQueryAuthResult( + filter=None, + column_masking={ + '_ROW_ID': json.dumps({'name': 'NULL'}), + }, + ) + raw_table.catalog_environment.table_query_auth = ( + lambda options, table_identifier: lambda select: auth + ) + with self.assertRaisesRegex( + ValueError, 'requires .*visible _ROW_ID'): + PaimonLeRobotDataset(table) + + def test_paimon_lerobot_timestamp_uses_float32_quantization(self): + frame_index = 61441 + contract = { + 'episode_ranges': [(0, frame_index + 1)], + 'episode_ends': [frame_index + 1], + 'fps': 30, + 'task_names': {0: 'pick'}, + 'subtask_names': None, + 'episode_tasks': (('pick',),), + 'timestamp_type': pa.float32(), + } + _validate_control_row( + { + 'episode_index': 0, + 'frame_index': frame_index, + 'timestamp': pa.scalar( + frame_index / 30, type=pa.float32()).as_py(), + 'task_index': 0, + }, + frame_index, + contract, + 1e-4, + ) + with self.assertRaisesRegex(ValueError, 'metadata expects'): + _validate_control_row( + { + 'episode_index': 0, + 'frame_index': frame_index, + 'timestamp': float('nan'), + 'task_index': 0, + }, + frame_index, + contract, + 1e-4, + ) + + def test_paimon_lerobot_tolerance_must_be_finite(self): + metadata = SimpleNamespace(repo_id='test/repo') + with patch( + 'pypaimon.multimodal.lerobot.dataset.' + '_load_published_version', + return_value=(None, metadata, 1)): + for tolerance in (float('nan'), float('inf'), float('-inf')): + with self.subTest(tolerance=tolerance): + with self.assertRaisesRegex( + ValueError, 'finite and non-negative'): + PaimonLeRobotDataset(None, tolerance_s=tolerance) + + def test_paimon_lerobot_episode_metadata_must_be_exact(self): + valid = { + 'episode_index': 0, + 'dataset_from_index': 0, + 'dataset_to_index': 1, + 'length': 1, + } + cases = ( + ( + 'extra row', + [valid, dict(valid, episode_index=1)], + 'contains 2 rows, expected 1', + ), + ( + 'wrong episode index', + [dict(valid, episode_index=1)], + 'row 0 has episode_index=1', + ), + ( + 'wrong length', + [dict(valid, length=2)], + 'episode 0 has length 2, expected 1', + ), + ) + for name, episodes, message in cases: + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, message): + _episode_ranges( + SimpleNamespace(episodes=episodes), 1, 1) + + def test_paimon_lerobot_component_metadata_must_match_manifest(self): + tasks = [{'task_index': 0}] + cases = ( + ( + 'task count', + {}, + 2, + tasks, + None, + 'task metadata contains 1 rows, expected 2', + ), + ( + 'missing subtasks', + {'subtask_index': {}}, + 1, + tasks, + None, + 'has_subtasks does not match', + ), + ( + 'unexpected subtasks', + {}, + 1, + tasks, + [{'subtask_index': 0}], + 'has_subtasks does not match', + ), + ) + for name, features, total_tasks, task_rows, subtasks, message \ + in cases: + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, message): + _validate_component_metadata( + features, total_tasks, task_rows, subtasks) + def test_non_streaming_row_tracking_without_data_evolution_materializes(self): schema = Schema.from_pyarrow_schema( self.pa_schema, diff --git a/paimon-python/setup.py b/paimon-python/setup.py index 12e24632bc3e..80e1aabeab6f 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -247,6 +247,7 @@ def read_requirements(): # by LeRobot's media dependencies. 'datasets>=4,<4.1; python_version>="3.10"', 'pandas>=2.2.2,<3; python_version>="3.10"', + 'torch>=2.3; python_version>="3.10"', 'lerobot>=0.4.4,<0.5; python_version>="3.10"', ], 'ray': [