diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index 5b35fff4f83a..53e05082dfd2 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -139,6 +139,10 @@ jobs: fi python -m pip install 'h5py>=3,<4' python -c "import h5py; print('h5py', h5py.__version__)" + if [[ "${{ matrix.python-version }}" == "3.10" ]]; then + python -m pip install './paimon-python[lerobot]' + python -c "import datasets, lerobot; print('datasets', datasets.__version__, 'lerobot', lerobot.__version__)" + fi if [[ "${{ matrix.python-version }}" == "3.11" ]]; then # Exercise the 0.4 API in one lane until its wheel is published. diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 6ca8d85634b3..38d3056d6d5f 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -318,6 +318,59 @@ HDF5 core itself has no Ray dependency. provenance columns, maintain a source ledger, skip prior inputs, or detect source drift. Calling it again with the same input appends the rows again. +## Load LeRobot Dataset v3 + +`load_from_lerobot` converts a LeRobot Dataset v3 from a local directory, +FileIO URI, or Hugging Face repository to a normal Paimon table once. It +derives the schema from `meta/info.json`, keeps one row per frame, preserves +Episode/frame indices and task text, and commits all batches in one Snapshot. + +```shell +pip install 'pypaimon[lerobot]' +``` + +```python +result = conn.load_from_lerobot( + "robot_data", + "/data/lerobot_dataset", +) +print(result.episode_count, result.row_count, result.snapshot_id) +``` + +An existing local directory is always used directly. FileIO-supported directory +URIs such as `oss://bucket/lerobot_dataset` use credentials from the explicit +`source_options` argument, never from the target Catalog. Another non-path +string is passed to the official LeRobot API as a Hugging Face `repo_id`. + +```python +result = conn.load_from_lerobot( + "robot_data", + "oss://source-bucket/lerobot_dataset", + source_options={ + "fs.oss.endpoint": "oss-cn-hangzhou.aliyuncs.com", + "fs.oss.accessKeyId": "SOURCE_ACCESS_KEY_ID", + "fs.oss.accessKeySecret": "SOURCE_ACCESS_KEY_SECRET", + }, +) +``` + +URI metadata and Parquet files are read directly through Paimon FileIO. The +LeRobot 0.4 video decoder requires a local path, so only the active MP4 chunk +for each video feature is cached temporarily and removed after import. Local +space is bounded by the active video chunks, not the complete dataset. + +If the target is absent, it is created from metadata. Existing targets use the +same strict schema validation and append semantics as `load_from_hdf5`. +One-dimensional numbers map to `VECTOR`; higher-rank values map to nested +`ARRAY`; `image` and `video` map to `BLOB`. Images retain their compressed +bytes. Videos are decoded once and stored as per-frame PNG BLOBs, so the shared +MP4 is not repeated in every row. + +Use `feature_mapping={"observation.state": "state"}` for explicit renames. An +optional `transform` receives each Arrow table and must preserve its row count +and Episode/frame/index order. Version 2.x, `uint64`, language event +structures, and depth-video semantics are not supported in this first version. + ## Overwrite `overwrite` accepts the same input formats as `add` and replaces existing data diff --git a/paimon-python/README.md b/paimon-python/README.md index 0e213985af8b..407a5d414f2c 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -31,6 +31,32 @@ pip3 install dist/*.tar.gz The command will install the package and core dependencies to your local Python environment. +# LeRobot Dataset v3 to multimodal tables + +Install the optional dependency, then import a local, FileIO URI, or Hugging +Face v3 dataset once. The target table is created from `meta/info.json` when +absent; later calls append. + +```commandline +pip install 'pypaimon[lerobot]' +``` + +```python +import pypaimon.multimodal as pmm + +connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) +result = connection.load_from_lerobot( + "robot_data", + "/data/lerobot_dataset", +) +print(result.row_count, result.snapshot_id) +``` + +Each LeRobot frame becomes one Paimon row. Numeric vectors remain vectors, +images use their existing compressed bytes, and video frames are decoded once +and stored as per-frame PNG BLOBs; the source MP4 is never copied into every +row. + # HDF5 to multimodal tables HDF5 loading requires Python 3.8 or newer. Install the optional dependency and diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index 96afcbbea392..ac04c4e13009 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 import LeRobotLoadResult from pypaimon.multimodal.table import ( MultimodalTable, TextRoute, @@ -47,6 +48,7 @@ "BlobStore", "Hdf5File", "Hdf5LoadResult", + "LeRobotLoadResult", "MultimodalConnection", "MultimodalTable", "NoSuchKey", diff --git a/paimon-python/pypaimon/multimodal/arrow_utils.py b/paimon-python/pypaimon/multimodal/arrow_utils.py new file mode 100644 index 000000000000..76d056ed9211 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/arrow_utils.py @@ -0,0 +1,139 @@ +# 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. + +"""Shared Arrow schema validation for multimodal format importers.""" + +import pyarrow as pa +import pyarrow.compute as pc + + +def strict_arrow_table( + data, + target_schema, + source_path, + batch_index, + format_name): + if isinstance(data, pa.RecordBatch): + table = pa.Table.from_batches([data]) + elif isinstance(data, pa.Table): + table = data + else: + raise ValueError( + "%s transform must return Arrow data or an iterable of Arrow data." + % format_name) + + missing = [ + name for name in target_schema.names if name not in table.column_names + ] + if missing: + raise ValueError( + "%s batch %d from %s is missing columns: %s" + % (format_name, batch_index, source_path, missing)) + extra = [ + name for name in table.column_names if name not in target_schema.names + ] + if extra: + raise ValueError( + "%s batch %d from %s has unexpected columns: %s" + % (format_name, batch_index, source_path, extra)) + if table.column_names != target_schema.names: + raise ValueError( + "%s batch %d from %s has columns in the wrong order: %s; " + "expected %s." + % (format_name, batch_index, source_path, table.column_names, + target_schema.names)) + try: + _validate_nested_nullability(table, target_schema) + if table.schema.equals(target_schema, check_metadata=False): + return table + casted = table.cast(target_schema, safe=True) + _validate_nested_nullability(casted, target_schema) + return casted + except (ValueError, TypeError, NotImplementedError) as error: + raise ValueError( + "%s batch %d from %s cannot be converted to the table schema: %s" + % (format_name, batch_index, source_path, error)) from error + + +def _validate_nested_nullability(table, schema): + for field, column in zip(schema, table.columns): + for chunk in column.chunks: + _validate_array_nullability(chunk, field, field.name) + + +def _validate_array_nullability(array, field, path): + if not field.nullable and array.null_count: + raise ValueError( + "non-nullable field %s contains %d null value(s)" + % (path, array.null_count)) + + target_type = field.type + source_type = array.type + if (pa.types.is_list(target_type) + or pa.types.is_large_list(target_type) + or pa.types.is_fixed_size_list(target_type)): + if not (pa.types.is_list(source_type) + or pa.types.is_large_list(source_type) + or pa.types.is_fixed_size_list(source_type)): + return + _validate_array_nullability( + pc.list_flatten(array), + target_type.value_field, + "%s.%s" % (path, target_type.value_field.name), + ) + return + + if pa.types.is_map(target_type): + if not pa.types.is_map(source_type): + return + start = array.offsets[0].as_py() + stop = array.offsets[-1].as_py() + length = stop - start + offsets = pc.subtract( + array.offsets, + pa.scalar(start, type=array.offsets.type), + ) + entries = pa.StructArray.from_arrays( + [array.keys.slice(start, length), + array.items.slice(start, length)], + fields=[source_type.key_field, source_type.item_field], + ) + logical_entries = pc.list_flatten(pa.ListArray.from_arrays( + offsets, + entries, + mask=pc.is_null(array), + )) + _validate_array_nullability( + logical_entries.field(0), target_type.key_field, + "%s.%s" % (path, target_type.key_field.name)) + _validate_array_nullability( + logical_entries.field(1), target_type.item_field, + "%s.%s" % (path, target_type.item_field.name)) + return + + if pa.types.is_struct(target_type): + if not pa.types.is_struct(source_type): + return + parent_valid = pc.is_valid(array) if array.null_count else None + for index, child_field in enumerate(target_type): + child = array.field(index) + if parent_valid is not None: + child = pc.filter(child, parent_valid) + _validate_array_nullability( + child, + child_field, + "%s.%s" % (path, child_field.name), + ) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 51085910a44c..a7bd2bcb5d28 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -117,6 +117,29 @@ def load_from_hdf5( source_options=source_options, ) + def load_from_lerobot( + self, + table_name: str, + source, + *, + transform=None, + feature_mapping=None, + batch_size: int = 1024, + options=None, + source_options=None): + """Import one LeRobot Dataset v3 in a single append commit.""" + from pypaimon.multimodal.lerobot import load_from_lerobot + return load_from_lerobot( + self, + table_name, + source, + transform=transform, + feature_mapping=feature_mapping, + batch_size=batch_size, + options=options, + source_options=source_options, + ) + def drop_table(self, name: str, ignore_if_not_exists: bool = False): self.catalog.drop_table( self._identifier(name), diff --git a/paimon-python/pypaimon/multimodal/hdf5.py b/paimon-python/pypaimon/multimodal/hdf5.py index 68389cf4e28c..e55b5752a623 100644 --- a/paimon-python/pypaimon/multimodal/hdf5.py +++ b/paimon-python/pypaimon/multimodal/hdf5.py @@ -26,13 +26,13 @@ from urllib.parse import quote, unquote, urlparse, urlunparse import pyarrow as pa -import pyarrow.compute as pc import pyarrow.fs as pafs from pypaimon.common.options import Options from pypaimon.filesystem.local_file_io import _file_uri_path from pypaimon.filesystem.pyarrow_file_io import LegacyOssDirectoryListingError from pypaimon.filesystem.resolving_file_io import ResolvingFileIO +from pypaimon.multimodal.arrow_utils import strict_arrow_table from pypaimon.multimodal.table import _target_schema from pypaimon.write.commit_callback import CommitCallback @@ -494,116 +494,13 @@ def _require_seekable(stream, source): def _strict_arrow_table(data, target_schema, source, batch_index): - if isinstance(data, pa.RecordBatch): - table = pa.Table.from_batches([data]) - elif isinstance(data, pa.Table): - table = data - else: - raise ValueError( - "HDF5 transform must return Arrow data or an iterable of Arrow data.") - - missing = [ - name for name in target_schema.names if name not in table.column_names - ] - if missing: - raise ValueError( - "HDF5 batch %d from %s is missing columns: %s" - % (batch_index, source.path, missing)) - extra = [ - name for name in table.column_names if name not in target_schema.names - ] - if extra: - raise ValueError( - "HDF5 batch %d from %s has unexpected columns: %s" - % (batch_index, source.path, extra)) - if table.column_names != target_schema.names: - raise ValueError( - "HDF5 batch %d from %s has columns in the wrong order: %s; " - "expected %s." - % (batch_index, source.path, table.column_names, - target_schema.names)) - try: - _validate_nested_nullability(table, target_schema) - if table.schema.equals(target_schema, check_metadata=False): - return table - casted = table.cast(target_schema, safe=True) - _validate_nested_nullability(casted, target_schema) - return casted - except (ValueError, TypeError, NotImplementedError) as error: - raise ValueError( - "HDF5 batch %d from %s cannot be converted to the table schema: %s" - % (batch_index, source.path, error)) from error - - -def _validate_nested_nullability(table, schema): - for field, column in zip(schema, table.columns): - for chunk in column.chunks: - _validate_array_nullability(chunk, field, field.name) - - -def _validate_array_nullability(array, field, path): - if not field.nullable and array.null_count: - raise ValueError( - "non-nullable field %s contains %d null value(s)" - % (path, array.null_count)) - - target_type = field.type - source_type = array.type - if (pa.types.is_list(target_type) - or pa.types.is_large_list(target_type) - or pa.types.is_fixed_size_list(target_type)): - if not (pa.types.is_list(source_type) - or pa.types.is_large_list(source_type) - or pa.types.is_fixed_size_list(source_type)): - return - _validate_array_nullability( - pc.list_flatten(array), - target_type.value_field, - "%s.%s" % (path, target_type.value_field.name), - ) - return - - if pa.types.is_map(target_type): - if not pa.types.is_map(source_type): - return - start = array.offsets[0].as_py() - stop = array.offsets[-1].as_py() - length = stop - start - offsets = pc.subtract( - array.offsets, - pa.scalar(start, type=array.offsets.type), - ) - entries = pa.StructArray.from_arrays( - [array.keys.slice(start, length), - array.items.slice(start, length)], - fields=[source_type.key_field, source_type.item_field], - ) - logical_entries = pc.list_flatten(pa.ListArray.from_arrays( - offsets, - entries, - mask=pc.is_null(array), - )) - _validate_array_nullability( - logical_entries.field(0), target_type.key_field, - "%s.%s" % (path, target_type.key_field.name)) - _validate_array_nullability( - logical_entries.field(1), target_type.item_field, - "%s.%s" % (path, target_type.item_field.name)) - return - - if pa.types.is_struct(target_type): - if not pa.types.is_struct(source_type): - return - parent_valid = pc.is_valid(array) if array.null_count else None - for index, child_field in enumerate(target_type): - child = array.field(index) - if parent_valid is not None: - child = pc.filter(child, parent_valid) - _validate_array_nullability( - child, - child_field, - "%s.%s" % (path, child_field.name), - ) + return strict_arrow_table( + data, + target_schema, + source.path, + batch_index, + "HDF5", + ) def _arrow_batches(transformed): diff --git a/paimon-python/pypaimon/multimodal/lerobot.py b/paimon-python/pypaimon/multimodal/lerobot.py new file mode 100644 index 000000000000..6b4ac2d16487 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot.py @@ -0,0 +1,973 @@ +# 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. + +"""One-time LeRobot Dataset v3 import into a multimodal Paimon table.""" + +import io +import json +import shutil +import sys +import tempfile +from bisect import bisect_right +from contextlib import closing, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Mapping, Optional + +import pyarrow as pa +import pyarrow.fs as pafs +import pyarrow.parquet as pq + +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableNotExistException, +) +from pypaimon.common.options import Options +from pypaimon.filesystem.pyarrow_file_io import LegacyOssDirectoryListingError +from pypaimon.multimodal.arrow_utils import strict_arrow_table +from pypaimon.multimodal.hdf5 import ( + _Hdf5SourceFileIO, + _SnapshotRecorder, + _normalize_source_path, + _qualified_status_path, + _validated_source_options, +) +from pypaimon.multimodal.table import _target_schema + + +_DEFAULT_TABLE_OPTIONS = { + "file.format": "parquet", + "vector.file.format": "parquet", +} +_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(), +} +_MEDIA_DTYPES = ("image", "video") + + +@dataclass(frozen=True) +class LeRobotLoadResult: + """Counts and snapshot for one ``load_from_lerobot`` call.""" + + episode_count: int + batch_count: int + row_count: int + snapshot_id: Optional[int] + + +@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 + + +def load_from_lerobot( + connection, + table_name: str, + source, + *, + transform: Optional[Callable] = None, + feature_mapping: Optional[Mapping[str, str]] = None, + batch_size: int = 1024, + options: Optional[Mapping[str, object]] = None, + source_options: Optional[Mapping[str, object]] = None): + """Import one LeRobot Dataset v3 and commit all frames once. + + 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 transform is not None and not callable(transform): + raise ValueError("transform must be callable or None.") + 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) + with _resolved_source(source, validated_source_options) as ( + resolved_source, local_info): + if local_info is not None: + _require_v3(local_info, resolved_source.path) + LeRobotDataset = _import_lerobot_dataset() + dataset = _open_resolved_dataset( + LeRobotDataset, resolved_source, local_info) + try: + info = dict(dataset.meta.info) + _require_v3(info, resolved_source.path) + + source_schema, mapped_names = _schema_from_info( + info, feature_mapping, + include_task=_has_tasks(dataset, info)) + table = _get_or_create_table( + connection, table_name, source_schema, options) + target_schema = _target_schema(table.raw_table) + _strict_lerobot_table( + pa.Table.from_batches([], schema=source_schema), + target_schema, + resolved_source, + 0, + ) + + episode_count = int(info.get("total_episodes", 0)) + row_count = int(info.get("total_frames", len(dataset))) + if row_count == 0: + return LeRobotLoadResult( + episode_count=episode_count, + batch_count=0, + row_count=0, + snapshot_id=None, + ) + return _write_dataset( + table, + dataset, + info, + resolved_source, + source_schema, + mapped_names, + transform, + batch_size, + ) + finally: + close = getattr(dataset, "close", None) + if callable(close): + close() + + +@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 _strict_lerobot_table(data, target_schema, source, batch_index): + return strict_arrow_table( + data, + target_schema, + source.path, + batch_index, + "LeRobot", + ) + + +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 _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 _open_dataset(LeRobotDataset, source): + try: + if source.root is not None: + return LeRobotDataset( + repo_id=source.repo_id, + root=source.root, + video_backend="pyav", + ) + return LeRobotDataset(repo_id=source.repo_id, video_backend="pyav") + 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._episodes_by_index = { + int(episode["episode_index"]): episode + for episode in self._episodes + } + self._data_ranges = self._build_data_ranges(info) + self._cached_data_path = None + self._cached_data_table = None + self._video_cache = {} + self._video_temp_dir = None + + def __len__(self): + return int(self.meta.info.get("total_frames", 0)) + + def close(self): + self._cached_data_table = None + self._video_cache.clear() + if self._video_temp_dir is not None: + self._video_temp_dir.cleanup() + self._video_temp_dir = 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_path(self.source.path, relative_path) + 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 = image_path if "://" in image_path else \ + _remote_path(self.source.path, image_path) + return _read_remote_bytes(self._file_io, source_path) + return _encode_media_frame(value) + + def read_video_values(self, name, raw): + episode_indices = raw.column("episode_index").to_pylist() + if not episode_indices or len(set(episode_indices)) != 1: + raise ValueError( + "LeRobot video batches must contain exactly one Episode.") + episode_index = int(_python_scalar(episode_indices[0])) + episode = self._episodes_by_index.get(episode_index) + if episode is None: + raise ValueError( + "LeRobot metadata is missing Episode %d." % episode_index) + info = self.meta.info + relative_path = info["video_path"].format( + video_key=name, + chunk_index=int(episode[ + "videos/%s/chunk_index" % name]), + file_index=int(episode[ + "videos/%s/file_index" % name]), + ) + source_path = _remote_path(self.source.path, relative_path) + local_path = self._cached_video_path(name, source_path) + start = float(episode["videos/%s/from_timestamp" % name]) + timestamps = [ + start + float(_python_scalar(value)) + for value in raw.column("timestamp").to_pylist() + ] + try: + from lerobot.datasets.video_utils import decode_video_frames + except ImportError as error: + raise ImportError( + "LeRobot video import requires the video dependencies from " + "'pypaimon[lerobot]'.") from error + frames = decode_video_frames( + local_path, + timestamps, + 1e-4, + "pyav", + ) + if len(frames) != len(timestamps): + raise ValueError( + "LeRobot video %s returned %d frames; expected %d." + % (source_path, len(frames), len(timestamps))) + return [_encode_media_frame(frame) for frame in frames] + + 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()) + rows.sort(key=lambda row: int(row["episode_index"])) + if len(rows) != episode_count: + raise ValueError( + "LeRobot metadata reports %d Episodes but %d were found." + % (episode_count, len(rows))) + return rows + + def _load_tasks(self, info): + task_count = int(info.get("total_tasks", 0)) + if task_count == 0: + return [] + path = _remote_path(self.source.path, "meta/tasks.parquet") + rows = _read_remote_parquet(self._file_io, path).to_pylist() + tasks = [None] * task_count + for row in rows: + index = int(row["task_index"]) + name = row.get("__index_level_0__") + if name is None: + name = row.get("task", row.get("name")) + if index < 0 or index >= task_count or name is None: + raise ValueError("LeRobot task metadata is invalid: %s" % row) + tasks[index] = str(name) + if any(task is None for task in tasks): + raise ValueError( + "LeRobot metadata reports %d tasks but %d were found." + % (task_count, len(rows))) + return tasks + + def _build_data_ranges(self, info): + ranges = {} + for episode in self._episodes: + path = self._data_path(episode, info) + begin = int(episode["dataset_from_index"]) + end = int(episode["dataset_to_index"]) + if path in ranges: + previous_begin, previous_end = ranges[path] + if begin != previous_end: + raise ValueError( + "LeRobot data file %s has non-contiguous Episode " + "ranges." % path) + ranges[path] = previous_begin, end + else: + ranges[path] = begin, end + return ranges + + def _episode_for_range(self, begin, end): + index = bisect_right(self._episode_starts, begin) - 1 + if index < 0: + raise ValueError("LeRobot frame range starts before Episode 0.") + episode = self._episodes[index] + episode_end = int(episode["dataset_to_index"]) + if end > episode_end: + raise ValueError("LeRobot frame batch crosses an Episode boundary.") + return episode + + @staticmethod + def _data_path(episode, info): + return info["data_path"].format( + chunk_index=int(episode["data/chunk_index"]), + file_index=int(episode["data/file_index"]), + ) + + def _cached_video_path(self, name, source_path): + cached = self._video_cache.get(name) + if cached is not None and cached[0] == source_path: + return cached[1] + if cached is not None: + try: + cached[1].unlink() + except FileNotFoundError: + pass + if self._video_temp_dir is None: + self._video_temp_dir = tempfile.TemporaryDirectory( + prefix="pypaimon_lerobot_video_") + output = tempfile.NamedTemporaryFile( + dir=self._video_temp_dir.name, + suffix=".mp4", + delete=False, + ) + try: + stream = self._file_io.new_input_stream(source_path) + with closing(stream) as source_stream: + shutil.copyfileobj(source_stream, output) + finally: + output.close() + local_path = Path(output.name) + self._video_cache[name] = source_path, local_path + return local_path + + +def _remote_path(root, relative_path): + return "%s/%s" % (root.rstrip("/"), relative_path.lstrip("/")) + + +def _read_remote_bytes(source_file_io, path): + stream = source_file_io.new_input_stream(path) + with closing(stream) as source_stream: + return source_stream.read() + + +def _read_remote_json(source_file_io, path): + try: + return json.loads(_read_remote_bytes( + source_file_io, path).decode("utf-8")) + except (OSError, UnicodeError, ValueError) as error: + raise ValueError( + "Cannot read LeRobot metadata %s: %s" % (path, error)) from error + + +def _read_remote_parquet(source_file_io, path): + stream = source_file_io.new_input_stream(path) + with closing(stream) as source_stream: + try: + return pq.read_table(source_stream) + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot Parquet file %s: %s" + % (path, error)) from error + + +def _remote_parquet_files(source_file_io, directory): + try: + statuses = source_file_io.list_status(directory) + except LegacyOssDirectoryListingError as error: + raise ValueError( + "LeRobot URI directory listing is unavailable at %s; use " + "Jindo or upgrade PyArrow." % directory) from error + paths = [] + for status in statuses: + path = _qualified_status_path(directory, status) + if status.type == pafs.FileType.Directory: + paths.extend(_remote_parquet_files(source_file_io, path)) + elif status.type == pafs.FileType.File and path.endswith(".parquet"): + paths.append(path) + return sorted(paths) + + +def _has_tasks(dataset, info): + return int(info.get("total_tasks", 0)) > 0 \ + and getattr(dataset.meta, "tasks", None) is not None + + +def _schema_from_info(info, feature_mapping, 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.") + source_names = list(features) + if include_task: + source_names.append("task") + mapping = _validated_feature_mapping(feature_mapping, source_names) + + fields = [] + mapped_names = {} + for source_name, feature in features.items(): + target_name = mapping.get(source_name, source_name) + mapped_names[source_name] = target_name + fields.append(_feature_field(target_name, source_name, feature)) + if include_task: + target_name = mapping.get("task", "task") + mapped_names["task"] = target_name + fields.append(pa.field( + target_name, + pa.string(), + nullable=False, + metadata={b"description": b"LeRobot task"}, + )) + return pa.schema(fields), mapped_names + + +def _validated_feature_mapping(feature_mapping, source_names): + if feature_mapping is None: + return {} + if not isinstance(feature_mapping, Mapping): + raise ValueError("feature_mapping must be a mapping or None.") + mapping = dict(feature_mapping) + unknown = sorted(set(mapping).difference(source_names)) + if unknown: + raise ValueError("feature_mapping contains unknown features: %s" % unknown) + targets = [] + for source_name in source_names: + target = mapping.get(source_name, source_name) + if not isinstance(target, str) or not target: + raise ValueError( + "feature_mapping target for %s must be a non-empty string." + % source_name) + targets.append(target) + duplicates = sorted({name for name in targets if targets.count(name) > 1}) + if duplicates: + raise ValueError( + "feature_mapping produces duplicate columns: %s" % duplicates) + return mapping + + +def _feature_field(target_name, source_name, feature): + if not isinstance(feature, dict): + raise ValueError( + "LeRobot feature %s metadata must be an object." % source_name) + dtype = str(feature.get("dtype", "")) + shape = _feature_shape(feature, source_name) + if dtype in _MEDIA_DTYPES: + if dtype == "video" and _is_depth_video(feature): + raise ValueError( + "LeRobot depth-video feature %s is not supported; its " + "decoded values cannot be losslessly stored as PNG frames." + % source_name) + 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, source_name, suffix)) + if pa.types.is_string(scalar_type) and shape not in ((), (1,)): + raise ValueError( + "LeRobot string feature %s must be scalar." % source_name) + arrow_type = _tensor_type(scalar_type, shape) + description = "LeRobot dtype=%s, shape=%s" % (dtype, list(shape)) + return pa.field( + target_name, + arrow_type, + nullable=False, + metadata={b"description": description.encode("utf-8")}, + ) + + +def _is_depth_video(feature): + info = feature.get("info") or feature.get("video_info") or {} + return bool(info.get("is_depth_map") + or info.get("video.is_depth_map")) + + +def _feature_shape(feature, name): + shape = feature.get("shape", ()) + if shape is None: + shape = () + if not isinstance(shape, (list, tuple)): + raise ValueError("LeRobot feature %s has an invalid shape: %r" % (name, shape)) + try: + result = tuple(int(size) for size in shape) + except (TypeError, ValueError) as error: + raise ValueError( + "LeRobot feature %s has an invalid shape: %r" % (name, shape)) from error + if any(size <= 0 for size in result): + raise ValueError("LeRobot feature %s has an invalid shape: %r" % (name, shape)) + return result + + +def _tensor_type(scalar_type, shape): + if shape in ((), (1,)): + return scalar_type + if len(shape) == 1: + return pa.list_(scalar_type, shape[0]) + result = pa.list_(scalar_type, shape[-1]) + for unused_size in reversed(shape[1:-1]): + result = pa.list_(result) + return pa.list_(result) + + +def _get_or_create_table(connection, table_name, schema, options): + try: + return connection.get_table(table_name) + except (DatabaseNotExistException, TableNotExistException): + table_options = dict(_DEFAULT_TABLE_OPTIONS) + if options: + table_options.update({str(key): str(value) + for key, value in options.items()}) + return connection.create_table( + table_name, + schema=schema, + options=table_options, + ) + + +def _write_dataset( + table, + dataset, + info, + source, + source_schema, + mapped_names, + transform, + 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, mapped_names) + if transform is not None: + transformed = transform(batch) + transformed = _as_arrow_table(transformed) + _validate_transform_order(batch, transformed, mapped_names) + batch = transformed + 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 LeRobotLoadResult( + episode_count=int(info.get("total_episodes", 0)), + batch_count=batch_count, + row_count=row_count, + snapshot_id=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, mapped_names): + 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"] + video_names = [name for name, feature in features.items() + if feature["dtype"] == "video"] + remote_video_reader = getattr(dataset, "read_video_values", None) + if callable(remote_video_reader): + video_values = { + name: remote_video_reader(name, raw) + for name in video_names + } + else: + video_values = {name: [] for name in video_names} + if video_names: + for index in range(begin, end): + item = dataset[index] + for name in video_names: + video_values[name].append( + _encode_media_frame(item[name])) + + arrays = [] + fields = [] + for source_name, feature in features.items(): + target_name = mapped_names[source_name] + field = schema.field(target_name) + dtype = feature["dtype"] + if dtype == "video": + values = video_values[source_name] + else: + if source_name not in raw.column_names: + raise ValueError( + "LeRobot data is missing metadata feature %s." % source_name) + values = raw.column(source_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, source_name) + for value in values] + arrays.append(pa.array(values, type=field.type)) + fields.append(field) + + if "task" in mapped_names: + task_indices = raw.column("task_index").to_pylist() + arrays.append(pa.array( + [_task_name(dataset.meta.tasks, value) for value in task_indices], + type=pa.string(), + )) + fields.append(schema.field(mapped_names["task"])) + return pa.Table.from_arrays(arrays, schema=pa.schema(fields)) + + +def _normalize_value(value, feature, name): + shape = _feature_shape(feature, name) + if shape in ((), (1,)): + if isinstance(value, (list, tuple)): + if len(value) != 1: + raise ValueError( + "LeRobot feature %s expected shape %s, got %s." + % (name, shape, _value_shape(value))) + return value[0] + return _python_scalar(value) + actual_shape = _value_shape(value) + if actual_shape != shape: + raise ValueError( + "LeRobot feature %s expected shape %s, got %s." + % (name, shape, actual_shape)) + return value + + +def _value_shape(value): + if hasattr(value, "shape"): + return tuple(int(size) for size in value.shape) + if isinstance(value, (list, tuple)): + if not value: + return (0,) + child = _value_shape(value[0]) + if any(_value_shape(item) != child for item in value[1:]): + return (len(value), -1) + return (len(value),) + child + return () + + +def _python_scalar(value): + item = getattr(value, "item", None) + if callable(item): + return item() + return value + + +def _image_bytes(value, root): + 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: + path = Path(image_path) + if not path.is_absolute(): + path = Path(root) / path + return path.read_bytes() + return _encode_media_frame(value) + + +def _encode_media_frame(value): + try: + import numpy as np + from PIL import Image + except ImportError as error: + raise ImportError( + "LeRobot media import requires numpy and Pillow from the " + "'pypaimon[lerobot]' extra.") from error + + if isinstance(value, Image.Image): + image = value + else: + detach = getattr(value, "detach", None) + if callable(detach): + value = detach().cpu().numpy() + array = np.asarray(value) + if array.ndim == 3 and array.shape[0] in (1, 3, 4): + array = np.transpose(array, (1, 2, 0)) + if np.issubdtype(array.dtype, np.floating): + array = np.rint(np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8) + if array.ndim == 3 and array.shape[2] == 1: + array = array[:, :, 0] + try: + image = Image.fromarray(array) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + "Unsupported LeRobot media frame shape or dtype: %s, %s." + % (array.shape, array.dtype)) from error + output = io.BytesIO() + image.save(output, format="PNG") + return output.getvalue() + + +def _task_name(tasks, task_index): + index = int(_python_scalar(task_index)) + if hasattr(tasks, "iloc"): + return str(tasks.iloc[index].name) + task = tasks[index] + if isinstance(task, dict): + return str(task.get("task", task.get("name"))) + return str(task) + + +def _as_arrow_table(value): + if isinstance(value, pa.RecordBatch): + return pa.Table.from_batches([value]) + if isinstance(value, pa.Table): + return value + raise ValueError("LeRobot transform must return one Arrow table or batch.") + + +def _validate_transform_order(before, after, mapped_names): + if before.num_rows != after.num_rows: + raise ValueError( + "LeRobot transform must preserve the number of frames in each batch.") + for source_name in ("episode_index", "frame_index", "index"): + target_name = mapped_names.get(source_name) + if target_name is None or target_name not in after.column_names: + continue + if not before.column(target_name).equals(after.column(target_name)): + raise ValueError( + "LeRobot transform must preserve %s order." % source_name) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py new file mode 100644 index 000000000000..8f6d47528a94 --- /dev/null +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -0,0 +1,510 @@ +# 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. + +import builtins +import json +import shutil +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.fs as pafs + +import pypaimon.multimodal as pmm +from pypaimon.multimodal.lerobot import ( + _import_lerobot_dataset, + _schema_from_info, +) + +try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset +except ImportError: + LeRobotDataset = None + + +class LeRobotValidationTest(unittest.TestCase): + + def test_optional_dependency_error_is_actionable(self): + original_import = builtins.__import__ + + def reject_lerobot(name, *args, **kwargs): + if name.startswith("lerobot"): + raise ImportError("missing for test") + return original_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=reject_lerobot): + with self.assertRaisesRegex( + ImportError, r"install 'pypaimon\[lerobot\]'"): + _import_lerobot_dataset() + + def test_schema_comes_from_metadata_and_rejects_unsupported_types(self): + info = { + "features": { + "scalar": {"dtype": "uint16", "shape": [1]}, + "vector": {"dtype": "float32", "shape": [3]}, + "tensor": {"dtype": "float64", "shape": [2, 3]}, + "image": {"dtype": "image", "shape": [8, 10, 3]}, + } + } + schema, names = _schema_from_info( + info, {"scalar": "renamed"}, include_task=True) + + self.assertEqual( + ["renamed", "vector", "tensor", "image", "task"], + schema.names, + ) + self.assertEqual(pa.int32(), schema.field("renamed").type) + self.assertEqual(pa.list_(pa.float32(), 3), schema.field("vector").type) + self.assertEqual( + pa.list_(pa.list_(pa.float64(), 3)), + schema.field("tensor").type, + ) + self.assertEqual(pa.large_binary(), schema.field("image").type) + self.assertEqual("renamed", names["scalar"]) + + info["features"]["scalar"]["dtype"] = "uint64" + with self.assertRaisesRegex(ValueError, "no lossless Paimon integer"): + _schema_from_info(info, None, include_task=False) + + info["features"] = { + "depth": { + "dtype": "video", + "shape": [8, 10, 1], + "info": {"is_depth_map": True}, + } + } + with self.assertRaisesRegex(ValueError, "depth-video"): + _schema_from_info(info, None, include_task=False) + + def test_local_v2_is_rejected_before_opening(self): + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_v2_")) + try: + info_dir = temp_dir / "meta" + info_dir.mkdir() + (info_dir / "info.json").write_text(json.dumps({ + "codebase_version": "v2.1", + "features": {"index": {"dtype": "int64", "shape": [1]}}, + })) + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + with self.assertRaisesRegex(ValueError, "supports LeRobot Dataset v3 only"): + connection.load_from_lerobot("frames", temp_dir) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + +class _RemoteLeRobotFileIO: + + def __init__(self, local_root, remote_root): + self.local_root = Path(local_root) + self.remote_root = remote_root.rstrip("/") + self.opened_paths = [] + self.close_count = 0 + + def _local_path(self, remote_path): + prefix = self.remote_root + "/" + if remote_path == self.remote_root: + return self.local_root + if not remote_path.startswith(prefix): + raise FileNotFoundError(remote_path) + return self.local_root / remote_path[len(prefix):] + + def _status(self, local_path): + relative = local_path.relative_to(self.local_root).as_posix() + remote_path = self.remote_root + if relative != ".": + remote_path += "/" + relative + native_path = remote_path.split("://", 1)[1] + file_type = pafs.FileType.Directory if local_path.is_dir() \ + else pafs.FileType.File + return pafs.FileInfo(native_path, file_type) + + def get_file_status(self, remote_path): + local_path = self._local_path(remote_path) + if not local_path.exists(): + raise FileNotFoundError(remote_path) + return self._status(local_path) + + def list_status(self, remote_path): + return [self._status(path) for path in sorted( + self._local_path(remote_path).iterdir())] + + def new_input_stream(self, remote_path): + self.opened_paths.append(remote_path) + return self._local_path(remote_path).open("rb") + + def close(self): + self.close_count += 1 + + +@unittest.skipUnless( + sys.version_info >= (3, 10) and LeRobotDataset is not None, + "LeRobot 0.4.x requires Python 3.10+ and the lerobot extra", +) +class LeRobotImportTest(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.source_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_source_")) + cls.image_source = cls.source_dir / "images" + cls.video_source = cls.source_dir / "videos" + cls._create_image_dataset(cls.image_source) + cls._create_video_dataset(cls.video_source) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.source_dir, ignore_errors=True) + + def setUp(self): + self.temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_test_")) + self.connection = pmm.connect(options={ + "warehouse": str(self.temp_dir / "warehouse"), + }) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + @staticmethod + def _create_image_dataset(root): + dataset = LeRobotDataset.create( + repo_id="pypaimon/local-image-test", + root=root, + fps=10, + use_videos=False, + image_writer_processes=0, + image_writer_threads=0, + features={ + "observation.state": { + "dtype": "float32", + "shape": (3,), + "names": ["x", "y", "z"], + }, + "observation.matrix": { + "dtype": "float32", + "shape": (2, 2), + "names": None, + }, + "action": { + "dtype": "float32", + "shape": (2,), + "names": ["x", "y"], + }, + "reward": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + "observation.image": { + "dtype": "image", + "shape": (8, 10, 3), + "names": ["height", "width", "channels"], + }, + }, + ) + for episode_index, length in enumerate((2, 3)): + for frame_index in range(length): + value = episode_index * 80 + frame_index * 10 + dataset.add_frame({ + "observation.state": np.array( + [episode_index, frame_index, episode_index + frame_index], + dtype=np.float32, + ), + "observation.matrix": np.array( + [[episode_index, frame_index], [frame_index, episode_index]], + dtype=np.float32, + ), + "action": np.array( + [frame_index, -frame_index], dtype=np.float32), + "reward": np.array([float(frame_index == length - 1)], + dtype=np.float32), + "observation.image": np.full( + (8, 10, 3), value, dtype=np.uint8), + "task": "pick" if episode_index == 0 else "place", + }) + dataset.save_episode() + dataset.finalize() + + @staticmethod + def _create_video_dataset(root): + dataset = LeRobotDataset.create( + repo_id="pypaimon/local-video-test", + root=root, + fps=5, + use_videos=True, + video_backend="pyav", + vcodec="h264", + image_writer_processes=0, + image_writer_threads=0, + features={ + "observation.video": { + "dtype": "video", + "shape": (8, 10, 3), + "names": ["height", "width", "channels"], + }, + "action": { + "dtype": "float32", + "shape": (2,), + "names": None, + }, + }, + ) + for frame_index in range(3): + dataset.add_frame({ + "observation.video": np.full( + (8, 10, 3), frame_index * 80, dtype=np.uint8), + "action": np.array( + [frame_index, -frame_index], dtype=np.float32), + "task": "move", + }) + dataset.save_episode() + dataset.finalize() + + def test_import_infers_schema_preserves_episodes_and_appends(self): + result = self.connection.load_from_lerobot( + "robot_data", self.image_source, batch_size=2) + + self.assertEqual(2, result.episode_count) + self.assertEqual(3, result.batch_count) + self.assertEqual(5, result.row_count) + self.assertEqual(1, result.snapshot_id) + + table = self.connection.get_table("robot_data") + schema = table.raw_table.fields + types = {field.name: str(field.type) for field in schema} + self.assertEqual("VECTOR NOT NULL", types["observation.state"]) + self.assertEqual( + "ARRAY> NOT NULL", + types["observation.matrix"], + ) + self.assertEqual("VECTOR NOT NULL", types["action"]) + self.assertEqual("FLOAT NOT NULL", types["timestamp"]) + self.assertEqual("BIGINT NOT NULL", types["episode_index"]) + self.assertEqual("BLOB NOT NULL", types["observation.image"]) + + rows = table.scan().select([ + "episode_index", + "frame_index", + "timestamp", + "index", + "task_index", + "task", + "observation.state", + "observation.matrix", + "action", + "reward", + ]).to_arrow().sort_by("index").to_pylist() + self.assertEqual([0, 0, 1, 1, 1], [row["episode_index"] for row in rows]) + self.assertEqual([0, 1, 0, 1, 2], [row["frame_index"] for row in rows]) + self.assertEqual([0, 1, 2, 3, 4], [row["index"] for row in rows]) + self.assertEqual(["pick", "pick", "place", "place", "place"], + [row["task"] for row in rows]) + self.assertEqual([1.0, -1.0], rows[1]["action"]) + self.assertEqual([[1.0, 2.0], [2.0, 1.0]], + rows[4]["observation.matrix"]) + self.assertAlmostEqual(0.2, rows[4]["timestamp"], places=6) + self.assertEqual(1.0, rows[4]["reward"]) + self.assertEqual( + result.snapshot_id, + table.raw_table.snapshot_manager().get_latest_snapshot().id, + ) + + scalar, blobs = table.scan().select([ + "index", "observation.image"] + ).read_blobs() + imported = dict(zip( + scalar.column("index").to_pylist(), blobs["observation.image"])) + source = LeRobotDataset( + repo_id="pypaimon/local-image-test", + root=self.image_source, + video_backend="pyav", + ).hf_dataset.with_format("arrow")[:] + expected = source.column("observation.image").to_pylist() + self.assertEqual( + [value["bytes"] for value in expected], + [imported[index] for index in range(5)], + ) + + appended = self.connection.load_from_lerobot( + "robot_data", self.image_source, batch_size=4) + self.assertEqual(2, appended.snapshot_id) + self.assertEqual(10, table.scan().to_arrow().num_rows) + + def test_video_frames_are_independent_blob_payloads(self): + result = self.connection.load_from_lerobot( + "video_data", self.video_source, batch_size=2) + self.assertEqual(3, result.row_count) + self.assertEqual(2, result.batch_count) + + table = self.connection.get_table("video_data") + scalar, blobs = table.scan().select([ + "index", "observation.video"] + ).read_blobs() + self.assertEqual(3, scalar.num_rows) + bodies = blobs["observation.video"] + self.assertTrue(all(body.startswith(b"\x89PNG\r\n\x1a\n") + for body in bodies)) + mp4 = next(self.video_source.rglob("*.mp4")).read_bytes() + self.assertTrue(all(body != mp4 for body in bodies)) + + def test_oss_source_streams_parquet_and_preserves_episodes(self): + source = "oss://source-bucket/robot-images" + source_file_io = _RemoteLeRobotFileIO(self.image_source, source) + + with patch( + "pypaimon.multimodal.lerobot._Hdf5SourceFileIO", + return_value=source_file_io): + result = self.connection.load_from_lerobot( + "oss_images", + source, + batch_size=2, + ) + + self.assertEqual(2, result.episode_count) + self.assertEqual(5, result.row_count) + table = self.connection.get_table("oss_images") + rows = table.scan().select([ + "episode_index", "frame_index", "index", "task" + ]).to_arrow().sort_by("index").to_pylist() + self.assertEqual([0, 0, 1, 1, 1], [ + row["episode_index"] for row in rows + ]) + self.assertEqual([0, 1, 0, 1, 2], [ + row["frame_index"] for row in rows + ]) + self.assertEqual( + ["pick", "pick", "place", "place", "place"], + [row["task"] for row in rows], + ) + self.assertFalse(any( + path.endswith("meta/stats.json") + for path in source_file_io.opened_paths + )) + self.assertEqual(1, len([ + path for path in source_file_io.opened_paths + if "/data/" in path and path.endswith(".parquet") + ])) + + def test_oss_video_uses_explicit_options_and_one_file_cache(self): + source = "oss://source-bucket/robot-video" + source_file_io = _RemoteLeRobotFileIO(self.video_source, source) + source_options = { + "fs.oss.endpoint": "oss-cn-test.example.com", + "fs.oss.accessKeyId": "source-key", + "fs.oss.accessKeySecret": "source-secret", + } + + with patch( + "pypaimon.multimodal.lerobot._Hdf5SourceFileIO", + return_value=source_file_io) as source_file_io_class: + result = self.connection.load_from_lerobot( + "oss_video", + source, + source_options=source_options, + ) + + self.assertEqual(3, result.row_count) + self.assertEqual(1, result.snapshot_id) + self.assertEqual(1, source_file_io.close_count) + self.assertEqual( + source_options, + source_file_io_class.call_args.args[0].to_map(), + ) + self.assertEqual( + 1, + len([path for path in source_file_io.opened_paths + if path.endswith(".mp4")]), + ) + self.assertEqual( + len(source_file_io.opened_paths), + len(set(source_file_io.opened_paths)), + ) + self.assertFalse(any( + path.endswith("meta/stats.json") + for path in source_file_io.opened_paths + )) + table = self.connection.get_table("oss_video") + unused_scalar, blobs = table.scan().select([ + "index", "observation.video" + ]).read_blobs() + self.assertTrue(all( + body.startswith(b"\x89PNG\r\n\x1a\n") + for body in blobs["observation.video"] + )) + + def test_existing_incompatible_schema_fails_without_snapshot(self): + info = json.loads((self.image_source / "meta" / "info.json").read_text()) + schema, unused_names = _schema_from_info( + info, None, include_task=True) + fields = [ + pa.field(field.name, pa.string(), nullable=False) + if field.name == "action" else field + for field in schema + ] + table = self.connection.create_table( + "incompatible", + schema=pa.schema(fields), + options={ + "file.format": "parquet", + "vector.file.format": "parquet", + }, + ) + + with self.assertRaisesRegex(ValueError, "cannot be converted"): + self.connection.load_from_lerobot( + "incompatible", self.image_source) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_feature_mapping_and_transform_preserve_order(self): + boundaries = [] + + def transform(batch): + episodes = batch.column("episode_index").to_pylist() + frames = batch.column("frame_index").to_pylist() + self.assertEqual(1, len(set(episodes))) + boundaries.append((episodes[0], frames)) + columns = [ + pc.add(batch.column(name), 1) + if name == "reward" else batch.column(name) + for name in batch.column_names + ] + return pa.Table.from_arrays(columns, schema=batch.schema) + + result = self.connection.load_from_lerobot( + "mapped", + self.image_source, + feature_mapping={"observation.state": "state"}, + transform=transform, + batch_size=3, + ) + self.assertEqual(5, result.row_count) + self.assertEqual([(0, [0, 1]), (1, [0, 1, 2])], boundaries) + table = self.connection.get_table("mapped") + self.assertIn("state", [field.name for field in table.raw_table.fields]) + rewards = table.scan().select([ + "index", "reward" + ]).to_arrow().sort_by("index").column("reward").to_pylist() + self.assertEqual([1.0, 2.0, 1.0, 1.0, 2.0], rewards) + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/setup.py b/paimon-python/setup.py index bf83fe4386fd..a2bfa57f47fa 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -237,6 +237,12 @@ def read_requirements(): # HDF5 loading is explicitly guarded and documented as Python 3.8+. 'h5py>=3,<4; python_version>="3.8"', ], + 'lerobot': [ + # datasets 4.1+ may select PyArrow 21+, while PyPaimon currently + # supports PyArrow <20. LeRobot 0.4.x uses the v3 dataset format. + 'datasets>=4,<4.1; python_version>="3.10"', + 'lerobot>=0.4.4,<0.5; python_version>="3.10"', + ], 'ray': [ 'ray>=2.10,<3; python_version>="3.8"', ],