JingsongLi commented on code in PR #9498: URL: https://github.com/apache/paimon/pull/9498#discussion_r3896862752
########## paimon-python/pypaimon/multimodal/lerobot/dataset.py: ########## @@ -0,0 +1,758 @@ +# 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( Review Comment: [P2] Please reject unsupported lazy-reader fallbacks before loading payload columns. `to_torch()` silently calls `_materialize()` when the row-ID path is unavailable; a reachable example is query authorization masking `_ROW_ID`. In that case constructing this wrapper reads and retains the full projected table even though the documented contract says payload columns remain lazy, and enabling deltas materializes a second projection as well. This can turn dataset construction into O(table size) payload I/O and memory use on training-scale tables. Please preflight the lazy eligibility and fail with an actionable error, or provide a lazy routing path for these cases. ########## paimon-python/pypaimon/multimodal/lerobot/dataset.py: ########## @@ -0,0 +1,758 @@ +# 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)): + root = Path(metadata) + if root.name == "meta": Review Comment: [P2] Please disambiguate the dataset root and `meta/` directory structurally instead of using the basename. A valid dataset rooted at `/datasets/meta` stores its metadata at `/datasets/meta/meta/info.json`, but this branch rewrites the root to `/datasets` and then rejects it. Check `<input>/meta/info.json` first as the dataset-root form; only if that is absent should `<input>/info.json` be treated as the metadata-directory form and resolved to its parent. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
