JingsongLi commented on code in PR #9446: URL: https://github.com/apache/paimon/pull/9446#discussion_r3891142893
########## paimon-python/pypaimon/multimodal/lerobot/schema.py: ########## @@ -0,0 +1,159 @@ +# 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 metadata validation and Arrow schema conversion.""" + +import pyarrow as pa + + +_SCALAR_DTYPES = { + "bool": pa.bool_(), + "boolean": pa.bool_(), + "int8": pa.int8(), + "int16": pa.int16(), + "int32": pa.int32(), + "int64": pa.int64(), + "uint8": pa.int16(), + "uint16": pa.int32(), + "uint32": pa.int64(), + "float16": pa.float32(), + "float32": pa.float32(), + "float64": pa.float64(), + "string": pa.string(), +} + + +def _require_v3(info, source): + version = str(info.get("codebase_version", "")) + if not (version == "v3" or version.startswith("v3.")): + raise ValueError( + "load_from_lerobot supports LeRobot Dataset v3 only; %s reports " + "codebase_version=%r. Upgrade the dataset to v3 first." + % (source, version or None)) + + +def _schema_from_info(info, include_task): + features = info.get("features") + if not isinstance(features, dict) or not features: + raise ValueError("LeRobot metadata features must be a non-empty object.") + + fields = [] + for name, feature in features.items(): + fields.append(_feature_field(name, feature)) + if include_task: + fields.append(pa.field( + "task", + pa.string(), + nullable=False, + metadata={b"description": b"LeRobot task"}, + )) + return pa.schema(fields) + + +def _validate_lerobot_schema(source_schema, target_schema, source): + """Require an existing table to preserve the LeRobot feature contract.""" + for source_field in source_schema: + target_index = target_schema.get_field_index(source_field.name) + if target_index < 0: + # The shared schema validator reports missing columns consistently. + continue + target_field = target_schema.field(target_index) + if source_field.type != target_field.type: + raise ValueError( + "LeRobot feature %s from %s cannot be converted to the " + "table schema: expected %s, found %s." + % (source_field.name, source, source_field.type, + target_field.type)) + + source_description = _description(source_field) + if not source_description.startswith("LeRobot dtype="): + continue + target_description = _description(target_field) + if target_description != source_description: + raise ValueError( + "LeRobot feature %s from %s cannot be converted to the " + "table schema: expected %s, found %s." + % (source_field.name, source, source_description, + target_description or "no LeRobot feature metadata")) + + +def _description(field): + if not field.metadata: + return "" + return field.metadata.get(b"description", b"").decode("utf-8") + + +def _feature_field(name, feature): + if not isinstance(feature, dict): + raise ValueError( + "LeRobot feature %s metadata must be an object." % name) + dtype = str(feature.get("dtype", "")) + shape = _feature_shape(feature, name) + if dtype == "video": + raise ValueError( + "LeRobot video feature %s is not supported yet; use an " + "image-based dataset." % name) + if dtype == "image": + arrow_type = pa.large_binary() + else: + scalar_type = _SCALAR_DTYPES.get(dtype) + if scalar_type is None: + suffix = " (uint64 has no lossless Paimon integer mapping)" \ + if dtype == "uint64" else "" + raise ValueError( + "Unsupported LeRobot dtype %r for feature %s%s." + % (dtype, name, suffix)) + if pa.types.is_string(scalar_type) and shape not in ((), (1,)): + raise ValueError( + "LeRobot string feature %s must be scalar." % name) + arrow_type = _tensor_type(scalar_type, shape) + description = "LeRobot dtype=%s, shape=%s" % (dtype, list(shape)) Review Comment: **[P1] Preserve LeRobot component names in the append contract.** This description records only `dtype` and `shape`, so two valid features such as `names=["x", "y"]` and `names=["y", "x"]` produce identical schemas and `_validate_lerobot_schema` accepts the append. The loader then writes vector positions unchanged, silently mixing different component semantics in one training table. Please persist a canonical representation of `feature["names"]` in the field metadata and compare it for existing tables; add a regression test with reordered names. ########## paimon-python/pypaimon/multimodal/lerobot/source.py: ########## @@ -0,0 +1,466 @@ +# 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 source resolution for local, Hub, and FileIO datasets.""" + +import json +import posixpath +from bisect import bisect_right +from contextlib import closing, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Optional +from urllib.parse import quote, unquote, urlparse, urlunparse + +import pyarrow as pa +import pyarrow.fs as pafs +import pyarrow.parquet as pq + +from pypaimon.common.options import Options +from pypaimon.filesystem.pyarrow_file_io import LegacyOssDirectoryListingError +from pypaimon.multimodal.hdf5 import ( + _Hdf5SourceFileIO, + _normalize_source_path, + _qualified_status_path, +) +from pypaimon.multimodal.lerobot.loader import _encode_media_frame + + +@dataclass(frozen=True) +class _LeRobotSource: + path: str + root: Optional[Path] + repo_id: str + file_io: object = None + + +@dataclass(frozen=True) +class _RemoteLeRobotMeta: + info: dict + episodes: list + tasks: list + + +@contextmanager +def _resolved_source(source, source_options): + if isinstance(source, Path): + root = source.expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % root) + yield _local_source(root) + return + if not isinstance(source, str) or not source.strip(): + raise ValueError( + "source must be a local directory or Hugging Face repo_id.") + + value = source.strip() + candidate = Path(value).expanduser() + if candidate.is_dir(): + yield _local_source(candidate.resolve()) + return + if candidate.is_absolute() or value.startswith((".", "~")): + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % candidate) + if "://" not in value: + yield _LeRobotSource(path=value, root=None, repo_id=value), None + return + + source_uri = _normalize_source_path(value).rstrip("/") + source_file_io = _Hdf5SourceFileIO(Options(source_options)) + try: + try: + status = source_file_io.get_file_status(source_uri) + except FileNotFoundError as error: + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % source_uri + ) from error + if status.type != pafs.FileType.Directory: + raise ValueError( + "LeRobot URI source must be a directory: %s" % source_uri) + info = _read_remote_json( + source_file_io, _remote_path(source_uri, "meta/info.json")) + yield ( + _LeRobotSource( + path=source_uri, + root=None, + repo_id="", + file_io=source_file_io, + ), + info, + ) + finally: + source_file_io.close() + + +def _local_source(root, display_path=None): + info_path = root / "meta" / "info.json" + if not info_path.is_file(): + raise ValueError( + "LeRobot source is missing meta/info.json: %s" + % (display_path or root)) + try: + with info_path.open("r", encoding="utf-8") as file: + info = json.load(file) + except (OSError, ValueError) as error: + raise ValueError( + "Cannot read LeRobot metadata %s: %s" + % (info_path, error)) from error + return ( + _LeRobotSource( + path=str(display_path or root), + root=root, + repo_id="local/pypaimon-import", + ), + info, + ) + + +def _import_lerobot_dataset(): + try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + except ImportError as error: + raise ImportError( + "load_from_lerobot requires LeRobot; install " + "'pypaimon[lerobot]'.") from error + return LeRobotDataset + + +def _load_hub_info(source): + try: + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + except ImportError as error: + raise ImportError( + "load_from_lerobot requires LeRobot; install " + "'pypaimon[lerobot]'.") from error + try: + return dict(LeRobotDatasetMetadata(repo_id=source.repo_id).info) + except Exception as error: + raise ValueError( + "Cannot open LeRobot Dataset v3 metadata %s: %s" + % (source.path, error)) from error + + +def _open_dataset(LeRobotDataset, source): + try: + if source.root is not None: + return LeRobotDataset( + repo_id=source.repo_id, + root=source.root, + download_videos=False, + ) + return LeRobotDataset( + repo_id=source.repo_id, + download_videos=False, + ) + except Exception as error: + raise ValueError( + "Cannot open LeRobot Dataset v3 source %s: %s" + % (source.path, error)) from error + + +def _open_resolved_dataset(LeRobotDataset, source, info): + if source.file_io is not None: + return _RemoteLeRobotDataset(source, info) + return _open_dataset(LeRobotDataset, source) + + +class _RemoteLeRobotDataset: + + def __init__(self, source, info): + self.source = source + self.root = source.path + self._file_io = source.file_io + self._episodes = self._load_episodes(info) + self._tasks = self._load_tasks(info) + self.meta = _RemoteLeRobotMeta(info, self._episodes, self._tasks) + self._episode_starts = [ + int(episode["dataset_from_index"]) + for episode in self._episodes + ] + self._data_ranges = self._build_data_ranges(info) + self._cached_data_path = None + self._cached_data_table = None + + def __len__(self): + return int(self.meta.info.get("total_frames", 0)) + + def close(self): + self._cached_data_table = None + + def read_batch(self, begin, end): + episode = self._episode_for_range(begin, end) + relative_path = self._data_path(episode, self.meta.info) + source_path = _remote_source_path( + self.source.path, + relative_path, + "info.data_path", + self._file_io, + ) + if source_path != self._cached_data_path: + table = _read_remote_parquet(self._file_io, source_path) + expected_begin, expected_end = self._data_ranges[relative_path] + expected_rows = expected_end - expected_begin + if table.num_rows != expected_rows: + raise ValueError( + "LeRobot data file %s has %d rows; metadata expects %d." + % (source_path, table.num_rows, expected_rows)) + self._cached_data_path = source_path + self._cached_data_table = table + file_begin = self._data_ranges[relative_path][0] + return self._cached_data_table.slice(begin - file_begin, end - begin) + + def image_bytes(self, value): + if value is None: + raise ValueError("LeRobot image feature contains a null frame.") + if isinstance(value, (bytes, bytearray, memoryview)): + return bytes(value) + if isinstance(value, dict): + body = value.get("bytes") + if body is not None: + return bytes(body) + image_path = value.get("path") + if image_path: + source_path = _remote_source_path( + self.source.path, + image_path, + "image path", + self._file_io, + ) + return _read_remote_bytes(self._file_io, source_path) + return _encode_media_frame(value) + + def _load_episodes(self, info): + episode_count = int(info.get("total_episodes", 0)) + if episode_count == 0: + return [] + directory = _remote_path(self.source.path, "meta/episodes") + paths = _remote_parquet_files(self._file_io, directory) + rows = [] + for path in paths: + rows.extend(_read_remote_parquet( + self._file_io, path).to_pylist()) Review Comment: **[P1] Project the episode metadata columns before materializing rows.** `_read_remote_parquet` currently reads every column and `to_pylist()` retains all of them for the full import. Normal LeRobot v3 episode rows include per-feature `stats/*` arrays, while this class only uses five scalar lookup fields. Large FileIO datasets can therefore consume GBs of Python heap and OOM before the first data batch. Please read only `episode_index`, `dataset_from_index`, `dataset_to_index`, `data/chunk_index`, and `data/file_index`. ########## paimon-python/pypaimon/multimodal/lerobot/api.py: ########## @@ -0,0 +1,131 @@ +# 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. + +"""Public LeRobot import API.""" + +import sys +from typing import Mapping, Optional + +import pyarrow as pa + +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableNotExistException, +) +from pypaimon.multimodal.lerobot.loader import ( + _strict_lerobot_table, + _write_dataset, +) +from pypaimon.multimodal.lerobot.schema import ( + _require_v3, + _schema_from_info, + _validate_lerobot_schema, +) +from pypaimon.multimodal.lerobot.source import ( + _has_tasks, + _import_lerobot_dataset, + _load_hub_info, + _open_resolved_dataset, + _resolved_source, + _validate_info_paths, +) +from pypaimon.multimodal.source_utils import ( + _validated_source_options, + _validate_source_kerberos, +) +from pypaimon.multimodal.table import _target_schema + + +def load_from_lerobot( + connection, + table_name: str, + source, + *, + batch_size: int = 1024, + options: Optional[Mapping[str, object]] = None, + source_options: Optional[Mapping[str, object]] = None): + """Import LeRobot Dataset v3 and return the committed snapshot ID. + + A missing target table is created from LeRobot metadata. An existing table + receives the same strict schema validation and append semantics as + :meth:`MultimodalConnection.load_from_hdf5`. FileIO URI credentials come + only from ``source_options`` and are not inherited from the target Catalog. + """ + if sys.version_info < (3, 10): + raise RuntimeError( + "load_from_lerobot requires Python 3.10 or newer; install and " + "run 'pypaimon[lerobot]' on a supported Python version.") + if isinstance(batch_size, bool) or not isinstance(batch_size, int) \ + or batch_size <= 0: + raise ValueError("batch_size must be a positive integer.") + + validated_source_options = _validated_source_options(source_options) + _validate_source_kerberos( + [source], validated_source_options, "LeRobot") + with _resolved_source(source, validated_source_options) as ( + resolved_source, local_info): + if local_info is None: + local_info = _load_hub_info(resolved_source) + _require_v3(local_info, resolved_source.path) + _validate_info_paths(local_info) + _schema_from_info(local_info, include_task=False) + LeRobotDataset = _import_lerobot_dataset() + dataset = _open_resolved_dataset( Review Comment: **[P2] Handle an empty local dataset before opening LeRobot.** A valid zero-frame directory created by `LeRobotDataset.create()` has `meta/info.json` but no data/episode/task Parquet files. Opening it here makes LeRobot treat the missing files as a cache miss and attempt a Hub download using the synthetic `local/pypaimon-import` repo ID, so the documented `None` no-op path at line 108 is never reached. Please create or validate the table from `local_info`, then return `None` before constructing `LeRobotDataset` when `total_frames == 0`. ########## paimon-python/pypaimon/multimodal/lerobot/loader.py: ########## @@ -0,0 +1,278 @@ +# 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 frame conversion and batch writing.""" + +import io +from pathlib import Path + +import pyarrow as pa + +from pypaimon.multimodal.arrow_utils import strict_arrow_table +from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.lerobot.schema import _feature_shape +from pypaimon.multimodal.table import _target_schema + + +def _strict_lerobot_table(data, target_schema, source, batch_index): + return strict_arrow_table( + data, + target_schema, + source.path, + batch_index, + "LeRobot", + ) + + +def _write_dataset( + table, + dataset, + info, + source, + source_schema, + batch_size): + target_schema = _target_schema(table.raw_table) + write_builder = table.raw_table.new_batch_write_builder() + table_write = None + table_commit = None + commit_started = False + batch_count = 0 + row_count = 0 + snapshot_recorder = _SnapshotRecorder() + + try: + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + table_commit.add_commit_callback(snapshot_recorder) + for begin, end in _episode_batches(dataset, info, batch_size): + batch = _read_batch( + dataset, info, begin, end, source_schema) + batch = _strict_lerobot_table( + batch, + target_schema, + source, + batch_count, + ) + table_write.write_arrow(batch) + batch_count += 1 + row_count += batch.num_rows + + expected_rows = int(info.get("total_frames", len(dataset))) + if row_count != expected_rows: + raise ValueError( + "LeRobot metadata reports %d frames but import produced %d." + % (expected_rows, row_count)) + messages = table_write.prepare_commit() + commit_started = True + table_commit.commit(messages) + if snapshot_recorder.snapshot_id is None: + raise RuntimeError( + "LeRobot append committed without reporting a snapshot id.") + return snapshot_recorder.snapshot_id + except BaseException: + if table_write is not None and not commit_started: + table_write.abort() + raise + finally: + try: + if table_write is not None: + table_write.close() + finally: + if table_commit is not None: + table_commit.close() + + +def _episode_batches(dataset, info, batch_size): + episodes = getattr(dataset.meta, "episodes", None) + episode_count = int(info.get("total_episodes", 0)) + total_frames = int(info.get("total_frames", len(dataset))) + if episodes is None: + raise ValueError("LeRobot v3 metadata is missing episode boundaries.") + expected_begin = 0 + for ordinal in range(episode_count): + episode = episodes.iloc[ordinal] if hasattr(episodes, "iloc") \ + else episodes[ordinal] + begin = int(_python_scalar(episode["dataset_from_index"])) + end = int(_python_scalar(episode["dataset_to_index"])) + if begin != expected_begin or end <= begin: + raise ValueError( + "LeRobot episode %d has invalid frame range [%d, %d); " + "expected it to start at %d." + % (ordinal, begin, end, expected_begin)) + while begin < end: + batch_end = min(begin + batch_size, end) + yield begin, batch_end + begin = batch_end + expected_begin = end + if expected_begin != total_frames: + raise ValueError( + "LeRobot episode ranges cover %d frames but metadata reports %d." + % (expected_begin, total_frames)) + + +def _read_batch(dataset, info, begin, end, schema): + read_batch = getattr(dataset, "read_batch", None) + if callable(read_batch): + raw = read_batch(begin, end) + else: + raw = dataset.hf_dataset.with_format("arrow")[begin:end] + if isinstance(raw, pa.RecordBatch): + raw = pa.Table.from_batches([raw]) + elif not isinstance(raw, pa.Table): + raw = pa.Table.from_pydict(raw) + features = info["features"] + + arrays = [] + fields = [] + for name, feature in features.items(): + field = schema.field(name) + dtype = feature["dtype"] + if name not in raw.column_names: + raise ValueError( + "LeRobot data is missing metadata feature %s." % name) + values = raw.column(name).to_pylist() + if dtype == "image": + image_reader = getattr(dataset, "image_bytes", None) + if callable(image_reader): + values = [image_reader(value) for value in values] + else: + values = [_image_bytes(value, dataset.root) + for value in values] + else: + values = [_normalize_value(value, feature, name) + for value in values] + arrays.append(pa.array(values, type=field.type)) Review Comment: **[P2] Validate source values before coercing them to the target Arrow type.** `pa.array(values, type=field.type)` is not a safe cast: for example, a Parquet float value `1.5` under `int32` metadata becomes `1`, and an oversized float becomes `inf` under `float32`. Because `strict_arrow_table` runs afterwards, it sees only the already-coerced target type and cannot reject the mismatch. Please validate the raw Arrow element type/range or use an explicitly safe cast before rebuilding the 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]
