diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index e8732875b176..89ccc304fbaa 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -562,6 +562,28 @@ keep their compressed bytes. Only v3 is supported. Video features, `uint64`, and language event structures are rejected. +For map-style training, reuse the source dataset's `meta/` contract while +reading frames lazily from the imported table: + +```python +from torch.utils.data import DataLoader +from pypaimon.multimodal import PaimonLeRobotDataset + +dataset = PaimonLeRobotDataset( + conn.get_table("robot_data"), + "/data/lerobot_dataset/meta", + blob_parallelism=8, +) +loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) +``` + +The metadata argument supplies the standard info, episodes, tasks, and stats +through `dataset.meta`; the imported table does not persist the complete +`meta/` directory. Payload columns remain lazy. Pass +`index_mapping=dataset.index_mapping` to readers of the same table snapshot to +avoid rescanning the control columns. The first version supports image-backed datasets; +video-backed reads remain a follow-up. + ## 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 b669267b41da..bc658a71c8da 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.table import ( MultimodalTable, TextRoute, @@ -55,6 +56,7 @@ "MultimodalTable", "NoSuchKey", "ObjectInfo", + "PaimonLeRobotDataset", "PutObjectResult", "TextRoute", "VectorRoute", 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..614ad1383f8f --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -0,0 +1,765 @@ +# 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 operator +from array import array +from pathlib import Path + +from pypaimon.multimodal.lerobot.schema import ( + _feature_shape, + _require_v3, + _schema_from_info, + _validate_lerobot_schema, +) +from pypaimon.multimodal.table import _target_schema + + +class _LeRobotIndexMapping: + """Reusable semantic-index mapping bound to one table snapshot.""" + + def __init__( + self, + table_identifier, + snapshot_id, + metadata_signature, + positions): + self._table_identifier = table_identifier + self._snapshot_id = snapshot_id + self._metadata_signature = metadata_signature + self._positions = positions + + +class PaimonLeRobotDataset: + """Map-style LeRobot reader backed by Paimon's lazy Torch dataset. + + ``metadata`` is a ``LeRobotDatasetMetadata`` object or a local LeRobot v3 + dataset/``meta`` directory. Paimon supplies frame data; metadata remains + available through :attr:`meta` for LeRobot training code. + """ + + def __init__( + self, + table, + metadata, + *, + episodes=None, + image_transforms=None, + delta_timestamps=None, + tolerance_s=1e-4, + index_mapping=None, + blob_parallelism=16): + self.meta = _resolve_metadata(metadata) + self.repo_id = getattr( + self.meta, "repo_id", getattr(table, "identifier", "paimon")) + self.image_transforms = image_transforms + self.delta_timestamps = delta_timestamps + self.tolerance_s = float(tolerance_s) + if self.tolerance_s < 0: + raise ValueError("tolerance_s must be non-negative.") + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + 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))) + if self._total_frames < 0 or self._total_episodes < 0: + raise ValueError( + "LeRobot metadata must define total_frames and " + "total_episodes.") + + 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: + self._selected_ranges = [ + self._episode_ranges[index] for index in self.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.") + + raw_table = getattr(table, "raw_table", None) + if raw_table is None: + raise TypeError("table must be a MultimodalTable.") + target_schema = _target_schema(raw_table) + table_fields = set(target_schema.names) + tasks = _metadata_member(self.meta, "tasks") + include_task = int(info.get("total_tasks", 0)) > 0 \ + or (tasks is not None and len(tasks) > 0) + source_schema = _schema_from_info( + info, include_task=include_task) + _validate_lerobot_schema(source_schema, target_schema, self.repo_id) + control_contract = _control_contract( + self.meta, self._episode_ranges, self._fps, tasks) + projection = list(self._features) + if include_task and "task" not in projection: + projection.append("task") + 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._index_positions = self._index_mapping._positions + self._file_io = read_table.file_io + 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 {} + + _materialize_images( + self._file_io, + [base_rows, delta_rows], + self._image_keys, + self.blob_parallelism, + ) + converted = { + position: _torch_row(row, self._features) + for position, row in base_rows.items() + } + converted.update({ + position: _torch_row(row, self._features) + 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))) + + +def _resolve_metadata(metadata): + if isinstance(metadata, (str, Path)): + path = Path(metadata) + if (path / "meta" / "info.json").is_file(): + root = path + elif (path / "info.json").is_file(): + root = path.parent + else: + raise ValueError( + "metadata must point to a LeRobot v3 dataset or meta " + "directory.") + try: + from lerobot.datasets.lerobot_dataset import \ + LeRobotDatasetMetadata + except ImportError as error: + raise ImportError( + "PaimonLeRobotDataset requires LeRobot; install " + "'pypaimon[lerobot]'.") from error + return LeRobotDatasetMetadata( + repo_id="local/pypaimon", root=root) + if not hasattr(metadata, "info"): + raise TypeError( + "metadata must be LeRobotDatasetMetadata or a local meta path.") + return metadata + + +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 + ranges = [] + expected = 0 + for ordinal in range(total_episodes): + row = _episode_row(episodes, ordinal) + begin = int(row["dataset_from_index"]) + end = int(row["dataset_to_index"]) + if begin != expected or end <= begin: + raise ValueError( + "LeRobot episode %d has invalid frame range [%d, %d)." + % (ordinal, begin, end)) + 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 _control_contract(metadata, episode_ranges, fps, tasks): + task_indices = _task_indices(tasks) + 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_indices.items())) if task_indices is not None + else None, + episode_tasks, + ) + return { + "episode_ranges": episode_ranges, + "episode_ends": ( + [end for _, end in episode_ranges] + if episode_ranges is not None else None), + "fps": fps, + "task_indices": task_indices, + "episode_tasks": episode_tasks, + "signature": signature, + } + + +def _task_indices(tasks): + if tasks is None or len(tasks) == 0: + return None + if hasattr(tasks, "iterrows"): + return { + str(task): operator.index(row["task_index"]) + for task, row in tasks.iterrows() + } + if isinstance(tasks, dict): + return { + str(task): operator.index(index) + for task, index in tasks.items() + } + return {str(task): index for index, task in enumerate(tasks)} + + +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 sorted(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(table_read, splits, require_lazy=True) + + +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["task_indices"] is not None: + projection.append("task") + 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"], + range(size) if positions is None else positions, + ) + + +def _reuse_index_mapping( + mapping, + table_identifier, + snapshot_id, + metadata_signature, + size): + 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.") + 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 = frame / contract["fps"] + 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 abs(timestamp - expected_timestamp) > tolerance_s: + raise ValueError( + "Paimon timestamp at LeRobot index %d is %r, metadata expects %r." + % (index, timestamp, expected_timestamp)) + + task_indices = contract["task_indices"] + if task_indices is None: + return + task = str(row.get("task")) + expected_task_index = task_indices.get(task) + if expected_task_index is None: + raise ValueError( + "Paimon task at LeRobot index %d is absent from metadata: %r." + % (index, task)) + _validate_int_control( + row, "task_index", index, expected_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)) + + +def _validate_int_control(row, field, index, expected): + try: + actual = 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 + if actual != expected: + raise ValueError( + "Paimon %s at LeRobot index %d is %r, metadata expects %r." + % (field, index, actual, expected)) + + +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 _materialize_images( + 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 _torch_row(row, features): + 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) + elif feature.get("dtype") != "string": + dtype = torch.float16 \ + if feature.get("dtype") == "float16" else None + result[key] = torch.tensor(value, dtype=dtype) + return result + + +def _image_tensor(payload, feature): + 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)) + return torch.from_numpy(array).permute(2, 0, 1).float().div_(255) + + +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..b18bdabe5ac7 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -189,13 +189,20 @@ 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], + *, + require_lazy: bool = False, + ): """ Initialize TorchDataset. Args: table_read: TableRead instance for reading data splits: List of splits to read + require_lazy: Fail instead of materializing unsupported reads """ self.table_read = table_read self.splits = splits @@ -221,8 +228,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 +247,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 7f1f8d10a8ce..e449c7c79967 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -15,6 +15,7 @@ # limitations under the License. import builtins +import io import json import shutil import sys @@ -31,6 +32,10 @@ from pypaimon.common.options import Options from pypaimon.multimodal.hdf5 import _Hdf5SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot +from pypaimon.multimodal.lerobot.dataset import ( + _image_tensor, + _resolve_metadata, +) from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, @@ -66,6 +71,43 @@ def _replaced_contract(field, old, new): class LeRobotValidationTest(unittest.TestCase): + def test_image_tensor_preserves_declared_channels(self): + try: + from PIL import Image + 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") + tensor = _image_tensor( + output.getvalue(), {"dtype": "image", "shape": shape}) + self.assertEqual( + [shape[2], shape[0], shape[1]], list(tensor.shape)) + for actual, value in zip(tensor[:, 0, 0], expected): + self.assertAlmostEqual( + value / 255, float(actual), places=6) + + 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)) + def test_dataset_open_never_downloads_videos(self): calls = [] @@ -680,6 +722,56 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): self.assertEqual(2, appended_snapshot_id) self.assertEqual(10, table.scan().to_arrow().num_rows) + def test_paimon_dataset_reads_lazy_batches_with_lerobot_metadata(self): + self.connection.load_from_lerobot( + "training_data", self.image_source, batch_size=2) + dataset = pmm.PaimonLeRobotDataset( + self.connection.get_table("training_data"), + self.image_source / "meta", + delta_timestamps={"action": [-0.1, 0.0, 0.1]}, + blob_parallelism=3, + ) + + self.assertEqual(5, len(dataset)) + self.assertEqual(2, dataset.num_episodes) + self.assertIsNotNone(dataset.meta.stats) + 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()) + + episode = pmm.PaimonLeRobotDataset( + self.connection.get_table("training_data"), + self.image_source, + episodes=[1], + ) + self.assertEqual(3, len(episode)) + self.assertEqual(1, episode.num_episodes) + self.assertEqual(2, int(episode[0]["index"])) + + root_named_meta = self.temp_dir / "meta" + shutil.copytree(self.image_source / "meta", root_named_meta / "meta") + self.assertEqual(5, _resolve_metadata(root_named_meta).total_frames) + def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" source_file_io = _RemoteLeRobotFileIO(self.image_source, source) @@ -787,5 +879,6 @@ def test_existing_incompatible_schema_fails_without_snapshot(self): self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) + if __name__ == "__main__": unittest.main() diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 0dc269110d5e..e79c229adcb4 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -34,6 +34,8 @@ from pypaimon import CatalogFactory, Schema from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.multimodal.lerobot.dataset import PaimonLeRobotDataset +from pypaimon.multimodal.lerobot.schema import _schema_from_info from pypaimon.multimodal.table import MultimodalTable from pypaimon.read.datasource.torch_dataset import ( @@ -571,6 +573,190 @@ def test_non_streaming_row_tracking_reads_batches_lazily(self): sorted(actual_ids), ) + def test_paimon_lerobot_dataset_reuses_lazy_map_reader(self): + 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]}, + '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, include_task=True) + table_options = { + 'data-evolution.enabled': 'true', + 'row-tracking.enabled': 'true', + 'blob-as-descriptor': 'true', + 'vector.file.format': 'parquet', + } + schema = Schema.from_pyarrow_schema( + arrow_schema, + partition_keys=['frame_index'], + options=table_options, + ) + identifier = 'default.test_paimon_lerobot_dataset' + 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)], + 'action': [float(index), float(-index)], + 'task': 'pick', + } + for index in range(5) + ], schema=arrow_schema)) + metadata = SimpleNamespace( + info=info, + features=features, + fps=10, + total_frames=5, + total_episodes=2, + episodes=[ + {'dataset_from_index': 0, 'dataset_to_index': 2}, + {'dataset_from_index': 2, 'dataset_to_index': 5}, + ], + tasks=['pick'], + stats={'action': {}}, + repo_id='pypaimon/test', + ) + dataset = PaimonLeRobotDataset( + table, + metadata, + episodes=[1, 0], + delta_timestamps={'action': [-0.1, 0.0, 0.1]}, + ) + + 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( + [[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'])) + + 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, + metadata, + 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, 'blob_parallelism must be a positive integer'): + PaimonLeRobotDataset(table, metadata, blob_parallelism=0) + + mismatched_episodes = SimpleNamespace( + **dict(metadata.__dict__, episodes=[ + {'dataset_from_index': 0, 'dataset_to_index': 3}, + {'dataset_from_index': 3, 'dataset_to_index': 5}, + ]) + ) + with self.assertRaisesRegex( + ValueError, 'episode_index at LeRobot index 2'): + PaimonLeRobotDataset(table, mismatched_episodes) + with self.assertRaisesRegex( + ValueError, 'different LeRobot metadata'): + PaimonLeRobotDataset( + table, + mismatched_episodes, + index_mapping=dataset.index_mapping, + ) + + mismatched_tasks = SimpleNamespace( + **dict(metadata.__dict__, tasks=['place']) + ) + with self.assertRaisesRegex( + ValueError, 'task at LeRobot index .* absent from metadata'): + PaimonLeRobotDataset(table, mismatched_tasks) + + incompatible_features = dict(features) + incompatible_features['action'] = { + 'dtype': 'float32', 'shape': [3] + } + incompatible_info = dict(info, features=incompatible_features) + incompatible_metadata = SimpleNamespace( + **dict(metadata.__dict__, info=incompatible_info, + features=incompatible_features) + ) + with self.assertRaisesRegex(ValueError, "cannot be converted"): + PaimonLeRobotDataset(table, incompatible_metadata) + + no_task_identifier = '%s_no_task' % identifier + self.catalog.create_table( + no_task_identifier, + Schema.from_pyarrow_schema( + _schema_from_info(info, include_task=False), + options=table_options, + ), + False, + ) + no_task_table = MultimodalTable( + self.catalog, + no_task_identifier, + self.catalog.get_table(no_task_identifier), + ) + with self.assertRaisesRegex( + ValueError, "missing.*task"): + PaimonLeRobotDataset(no_task_table, metadata) + + 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 patch.object( + TableRead, 'to_arrow', side_effect=AssertionError( + 'lazy preflight must not materialize payload columns')): + with self.assertRaisesRegex( + ValueError, 'requires .*visible _ROW_ID'): + PaimonLeRobotDataset(table, metadata) + def test_non_streaming_row_tracking_without_data_evolution_materializes(self): schema = Schema.from_pyarrow_schema( self.pa_schema,