JingsongLi commented on code in PR #9498: URL: https://github.com/apache/paimon/pull/9498#discussion_r3969635019
########## paimon-python/pypaimon/multimodal/lerobot/dataset.py: ########## @@ -0,0 +1,1177 @@ +# 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 +import os +import sys + +import pyarrow as pa + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.lerobot.metadata import ( + _companion_table_identifiers, + _restore_pandas_metadata, + _tag_snapshot_id, + _validate_tag_name, +) +from pypaimon.multimodal.lerobot.loader import _DECLARED_NUMERIC_RANGES +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 +from pypaimon.read.query_auth_split import QueryAuthSplit + + +_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", +} + +_IMAGE_READ_ATTEMPTS = 3 + +_CONTROL_FEATURES = frozenset({ + "index", + "episode_index", + "frame_index", + "timestamp", + "task_index", + "subtask_index", +}) + + +class PaimonLeRobotDataset: + """Map-style LeRobot reader backed by indexed Paimon reads. + + LeRobot metadata is resolved from the Paimon table group and remains + available through :attr:`meta`. + + 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, + *, + tag_name=None, + episodes=None, + image_transforms=None, + delta_timestamps=None, + tolerance_s=1e-4, + blob_parallelism=16, + return_uint8=False): + if sys.version_info < (3, 10): + raise RuntimeError( + "PaimonLeRobotDataset requires Python 3.10 or newer; " + "install and run 'pypaimon[lerobot]' on a supported Python " + "version.") + raw_table, self.meta = _load_dataset(table, tag_name) + self.tag_name = tag_name + 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 = self._init_metadata() + self._init_episodes(episodes) + self._init_reader(raw_table, info) + + def _init_metadata(self): + 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._fps = int( + _metadata_member(self.meta, "fps", info.get("fps", 0))) + if self._fps <= 0: + raise ValueError("LeRobot metadata fps must be positive.") + return info + + def _init_episodes(self, 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: + # 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._delta_indices = _delta_indices( + self.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.") + + def _init_reader(self, raw_table, info): + target_schema = _target_schema(raw_table) + table_fields = set(target_schema.names) + tasks = _metadata_member(self.meta, "tasks") + subtasks = _metadata_member(self.meta, "subtasks") + _validate_component_metadata( + self._features, self._total_tasks, tasks, subtasks) + source_schema = _schema_from_info(info) + _validate_lerobot_schema(source_schema, target_schema, self.repo_id) + validation_context = _build_frame_validation_context( + self.meta, + self._episode_ranges, + self._fps, + tasks, + subtasks, + source_schema.field("timestamp").type, + ) + projection = list(self._features) + missing = set(projection) - table_fields + if missing: + raise ValueError( + "Paimon table is missing LeRobot fields: %s" + % sorted(missing)) + + self._read_table, self._snapshot_id, splits = _indexed_read_table( + raw_table, projection) + snapshot = self._read_table.snapshot_manager().get_snapshot_by_id( + self._snapshot_id) + if snapshot.next_row_id != self._total_frames: + raise ValueError( + "Paimon table has %d rows but metadata declares %d frames." + % (snapshot.next_row_id, self._total_frames)) + self._projection = projection + self._frame_locator = _FrameLocator( + self._read_table, snapshot, splits) + self._validation_context = validation_context + self._file_io = self._read_table.file_io + self._task_names = validation_context["task_names"] + self._subtask_names = validation_context["subtask_names"] + self._delta_projection = None + if self._delta_indices: + self._delta_projection = list(dict.fromkeys( + [ + "index", "episode_index", "frame_index", "timestamp", + "task_index", + ] + + (["subtask_index"] if subtasks is not None else []) + + list(self._delta_indices) + )) + + @property + def features(self): + return self._features + + @property + def fps(self): + return self._fps + + @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): + dataset_indices = [ + _normalize_index(index, len(self)) for index in indices + ] + if not dataset_indices: + return [] + frame_indices = [ + self._global_index(index) for index in dataset_indices + ] + plans = [self._plan(index) for index in frame_indices] + + unique_frame_indices = sorted(set(frame_indices)) + unique_frame_index_set = set(unique_frame_indices) + delta_indices = sorted({ + position + for plan in plans + for positions in plan["windows"].values() + for position in positions + if position not in unique_frame_index_set + }) + lookup_indices = sorted(unique_frame_index_set.union(delta_indices)) + splits, needs_filter = self._frame_locator.locate(lookup_indices) + rows = self._read_rows( + lookup_indices, self._projection, splits, needs_filter) + base_rows = { + index: rows[index] for index in unique_frame_indices + } + delta_rows = { + index: { + name: rows[index][name] for name in self._delta_projection + } + for index in delta_indices + } if delta_indices else {} + + _attach_task_labels( + base_rows, self._task_names, self._subtask_names) + row_groups = [base_rows, delta_rows] + image_sources = _image_blob_sources( + row_groups, self._image_keys) + for attempt in range(_IMAGE_READ_ATTEMPTS): + if attempt: + _restore_image_blob_sources(image_sources) + try: + _resolve_image_blobs( + self._file_io, + row_groups, + self._image_keys, + self.blob_parallelism, + ) + 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() + }) + break + except OSError: + if attempt + 1 == _IMAGE_READ_ATTEMPTS: + raise + + 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 _read_rows( + self, indices, projection, splits=None, needs_filter=True): + if not indices: + return {} + return _read_rows_by_index( + self._read_table, + projection, + indices, + self._validation_context, + self.tolerance_s, + self._features, + splits, + needs_filter, + ) + + 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 _global_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 _FrameLocator: + """Locate LeRobot frame rows in one fixed Paimon snapshot.""" + + def __init__(self, table, snapshot, splits): + self._table = table + self._snapshot = snapshot + self._scanner = None + self._scanner_initialized = False + self._process_id = os.getpid() + self._set_splits(splits) + + def _set_splits(self, splits): + from pypaimon.read.datasource.torch_dataset import ( + SplitRangeIndex, + row_ranges_for_split, + ) + + self._splits = splits + self._split_ranges = [ + row_ranges_for_split(split) for split in splits + ] + self._split_range_index = SplitRangeIndex(self._split_ranges) + + def locate(self, indices): + """Return narrowed splits and whether rows still need filtering.""" + self._ensure_process() + predicate = _index_predicate(self._table, indices) + try: + scanner = self._index_scanner(predicate) + except Exception as error: + raise RuntimeError( + "Failed to open the Paimon global index for LeRobot frame " + "lookups.") from error + if scanner is None: + raise RuntimeError( + "PaimonLeRobotDataset requires a readable global index on " + "the frame 'index' column.") + try: + evaluation = scanner.scan_with_coverage(predicate) + if evaluation is None: + raise RuntimeError( + "The Paimon global index could not evaluate the LeRobot " + "frame index predicate.") + unindexed = scanner.unindexed_ranges( + predicate, + search_mode=self._table.options.scalar_index_search_mode(), + contributing_field_ids=evaluation.contributing_field_ids, + ) + ranges = evaluation.result.results().to_range_list() + unindexed + from pypaimon.read.datasource.torch_dataset import ( + select_indexed_splits, + ) + from pypaimon.utils.range import Range + return select_indexed_splits( + self._splits, + self._split_ranges, + self._split_range_index, + Range.sort_and_merge_overlap(ranges, True), + ), bool(unindexed) + except RuntimeError: + raise + except Exception as error: + raise RuntimeError( + "Failed to query the Paimon global index for LeRobot " + "frames.") from error + + def _ensure_process(self): + process_id = os.getpid() + if process_id == self._process_id: + return + self._scanner = None + self._scanner_initialized = False + self._set_splits(self._splits) + self._process_id = process_id + + def _index_scanner(self, predicate): + if not self._scanner_initialized: + from pypaimon.globalindex import DataEvolutionGlobalIndexScanner + self._scanner = DataEvolutionGlobalIndexScanner.create( + self._table, + predicate=predicate, + snapshot=self._snapshot, + ) + self._scanner_initialized = True + return self._scanner + + def close(self): + scanner = self._scanner + self._scanner = None + self._scanner_initialized = False + if scanner is not None and self._process_id == os.getpid(): + scanner.close() + + def __getstate__(self): + state = self.__dict__.copy() + state["_scanner"] = None + state["_scanner_initialized"] = False + state["_process_id"] = None + state["_split_ranges"] = None + state["_split_range_index"] = None + return state + + def __del__(self): + try: + self.close() + except Exception: + pass + + +class _PaimonLeRobotMetadata: + + def __init__( + self, repo_id, tag_name, info, stats, episodes, tasks, + subtasks): + self.repo_id = repo_id + self.revision = tag_name + 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_dataset(table, tag_name): + raw_table = getattr(table, "raw_table", None) + if raw_table is None: + raise TypeError("table must be a MultimodalTable.") + if tag_name is not None: + _validate_tag_name(tag_name) + identifiers = _companion_table_identifiers(raw_table) + catalog = table.catalog + frames = _component_table(catalog, raw_table, tag_name) + episodes_table = _component_table( + catalog, catalog.get_table(identifiers["episodes"]), tag_name) + episodes = _episode_dataset(episodes_table) + tasks_table = _component_table( + catalog, catalog.get_table(identifiers["tasks"]), tag_name) + tasks = _component_dataframe(tasks_table, "task_index") + subtasks = None + if "subtasks" in identifiers: + subtasks_table = _component_table( + catalog, catalog.get_table(identifiers["subtasks"]), tag_name) + subtasks = _component_dataframe(subtasks_table, "subtask_index") + + info = _metadata_object(_component_table( + catalog, catalog.get_table(identifiers["info"]), tag_name), "info") + for feature in info.get("features", {}).values(): + feature["shape"] = tuple(feature["shape"]) + stats = None + if "stats" in identifiers: + stats = _numpy_stats(_metadata_object(_component_table( + catalog, catalog.get_table(identifiers["stats"]), tag_name), + "stats")) + metadata = _PaimonLeRobotMetadata( + str(table.identifier), tag_name, info, stats, episodes, tasks, + subtasks) + return frames, metadata + + +def _component_table(catalog, table, tag_name): + if tag_name is None: + return table + snapshot_id = _tag_snapshot_id(catalog, table.identifier, tag_name) + if snapshot_id is None: + raise ValueError( + "Paimon LeRobot component %s is missing tag %s." + % (table.identifier, tag_name)) + return _time_travel_table(table, tag_name=tag_name) + + +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 _metadata_object(table, name): + result = {} + for row in _read_arrow(table).to_pylist(): + key = row.get("key") + if not isinstance(key, str) or key in result: + raise ValueError( + "Paimon LeRobot %s metadata contains an invalid key." + % name) + try: + result[key] = json.loads(row.get("value")) + except (TypeError, ValueError) as error: + raise ValueError( + "Paimon LeRobot %s metadata value for %r is invalid JSON." + % (name, key)) from error + 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 subtask metadata does not match the " + "subtask_index feature.") + + +def _build_frame_validation_context( + 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 + 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, + } + + +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]) Review Comment: [P2] Preserve integer task IDs when iterating metadata `iterrows()` does not preserve column dtypes. For example, a task DataFrame with integer `task_index=[0, 1]`, floating-point `quality=[0.9, 0.8]`, and text index `['pick', 'place']` yields float64 rows, so `operator.index(row[index_field])` raises even though the ID column is still int64. The importer accepts and preserves these extra native metadata columns, so a successfully imported dataset fails to initialize here. I reproduced the same failure for both tasks and subtasks. Please iterate the index column directly, e.g. `zip(values.index, values[index_field])`, and add a regression case with a floating-point metadata column. -- 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]
