From 2c5b4f9cfb8962d9eba4949962841e0c57259685 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Fri, 28 Aug 2026 00:55:11 -0700 Subject: [PATCH 1/5] [python] Add LeRobot v3 dataset import --- .github/workflows/paimon-python-checks.yml | 4 + docs/docs/pypaimon/multimodal-api.mdx | 35 + paimon-python/README.md | 26 + paimon-python/pypaimon/multimodal/__init__.py | 2 + .../pypaimon/multimodal/connection.py | 21 + paimon-python/pypaimon/multimodal/hdf5.py | 22 +- paimon-python/pypaimon/multimodal/lerobot.py | 621 ++++++++++++++++++ .../pypaimon/tests/multimodal_lerobot_test.py | 380 +++++++++++ paimon-python/setup.py | 6 + 9 files changed, 1107 insertions(+), 10 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/lerobot.py create mode 100644 paimon-python/pypaimon/tests/multimodal_lerobot_test.py 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..09a3580d71f0 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -318,6 +318,41 @@ 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 local LeRobot Dataset v3 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", + batch_size=1024, +) +print(result.episode_count, result.row_count, result.snapshot_id) +``` + +An existing local directory is always used directly; another non-path string +is passed to the official LeRobot API as a Hugging Face `repo_id`. + +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..1552f4257e7b 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 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", + batch_size=1024, +) +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/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 51085910a44c..32a8f9e0a8ac 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -117,6 +117,27 @@ 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): + """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, + ) + 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..a2cbadb60a2e 100644 --- a/paimon-python/pypaimon/multimodal/hdf5.py +++ b/paimon-python/pypaimon/multimodal/hdf5.py @@ -493,34 +493,36 @@ def _require_seekable(stream, source): ) from error -def _strict_arrow_table(data, target_schema, source, batch_index): +def _strict_arrow_table( + data, target_schema, source, batch_index, format_name="HDF5"): 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.") + "%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( - "HDF5 batch %d from %s is missing columns: %s" - % (batch_index, source.path, missing)) + "%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( - "HDF5 batch %d from %s has unexpected columns: %s" - % (batch_index, source.path, extra)) + "%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( - "HDF5 batch %d from %s has columns in the wrong order: %s; " + "%s batch %d from %s has columns in the wrong order: %s; " "expected %s." - % (batch_index, source.path, table.column_names, + % (format_name, batch_index, source.path, table.column_names, target_schema.names)) try: _validate_nested_nullability(table, target_schema) @@ -531,8 +533,8 @@ def _strict_arrow_table(data, target_schema, source, batch_index): 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 + "%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): diff --git a/paimon-python/pypaimon/multimodal/lerobot.py b/paimon-python/pypaimon/multimodal/lerobot.py new file mode 100644 index 000000000000..e1e4f5b32b3e --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot.py @@ -0,0 +1,621 @@ +# 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 sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Mapping, Optional + +import pyarrow as pa + +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableNotExistException, +) +from pypaimon.multimodal.hdf5 import ( + _SnapshotRecorder, + _strict_arrow_table, +) +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 + + +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): + """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`. + """ + 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.") + + resolved_source, local_info = _resolve_source(source) + if local_info is not None: + _require_v3(local_info, resolved_source.path) + LeRobotDataset = _import_lerobot_dataset() + dataset = _open_dataset(LeRobotDataset, resolved_source) + 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_arrow_table( + pa.Table.from_batches([], schema=source_schema), + target_schema, + resolved_source, + 0, + format_name="LeRobot", + ) + + 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, + ) + + +def _resolve_source(source): + if isinstance(source, Path): + root = source.expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError("LeRobot source directory does not exist: %s" % root) + return _local_source(root) + 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(): + return _local_source(candidate.resolve()) + if candidate.is_absolute() or value.startswith((".", "~")): + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % candidate) + return _LeRobotSource(path=value, root=None, repo_id=value), None + + +def _local_source(root): + info_path = root / "meta" / "info.json" + if not info_path.is_file(): + raise ValueError( + "LeRobot source is missing meta/info.json: %s" % 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(root), + root=root, + repo_id="local/pypaimon-import", + ), + info, + ) + + +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 _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_arrow_table( + batch, + target_schema, + source, + batch_count, + format_name="LeRobot", + ) + 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): + 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"] + 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": + 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..6b62b49464a6 --- /dev/null +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -0,0 +1,380 @@ +# 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 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) + + +@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_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"', ], From 09eddd584ca92a7dc422c45a2df725dfe2ef1706 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Fri, 28 Aug 2026 01:35:57 -0700 Subject: [PATCH 2/5] [python] Support OSS sources in LeRobot import --- docs/docs/pypaimon/multimodal-api.mdx | 31 ++- paimon-python/README.md | 6 +- .../pypaimon/multimodal/connection.py | 4 +- paimon-python/pypaimon/multimodal/lerobot.py | 198 +++++++++++++----- .../pypaimon/tests/multimodal_lerobot_test.py | 88 ++++++++ 5 files changed, 265 insertions(+), 62 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 09a3580d71f0..44a475d1df9f 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -320,10 +320,10 @@ source drift. Calling it again with the same input appends the rows again. ## Load LeRobot Dataset v3 -`load_from_lerobot` converts a local LeRobot Dataset v3 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. +`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]' @@ -333,13 +333,30 @@ pip install 'pypaimon[lerobot]' result = conn.load_from_lerobot( "robot_data", "/data/lerobot_dataset", - batch_size=1024, ) print(result.episode_count, result.row_count, result.snapshot_id) ``` -An existing local directory is always used directly; another non-path string -is passed to the official LeRobot API as a Hugging Face `repo_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", + }, +) +``` + +The official LeRobot reader requires a local root. URI sources are therefore +copied once to a temporary directory, imported, and removed. Each MP4 is copied +once, not once per frame; ensure local temporary storage can hold the source. 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`. diff --git a/paimon-python/README.md b/paimon-python/README.md index 1552f4257e7b..407a5d414f2c 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -33,8 +33,9 @@ The command will install the package and core dependencies to your local Python # LeRobot Dataset v3 to multimodal tables -Install the optional dependency, then import a local v3 dataset once. The -target table is created from `meta/info.json` when absent; later calls append. +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]' @@ -47,7 +48,6 @@ connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) result = connection.load_from_lerobot( "robot_data", "/data/lerobot_dataset", - batch_size=1024, ) print(result.row_count, result.snapshot_id) ``` diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 32a8f9e0a8ac..a7bd2bcb5d28 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -125,7 +125,8 @@ def load_from_lerobot( transform=None, feature_mapping=None, batch_size: int = 1024, - options=None): + 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( @@ -136,6 +137,7 @@ def load_from_lerobot( 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): diff --git a/paimon-python/pypaimon/multimodal/lerobot.py b/paimon-python/pypaimon/multimodal/lerobot.py index e1e4f5b32b3e..19836a28ccbc 100644 --- a/paimon-python/pypaimon/multimodal/lerobot.py +++ b/paimon-python/pypaimon/multimodal/lerobot.py @@ -18,20 +18,32 @@ import io import json +import posixpath +import shutil import sys +import tempfile +from contextlib import closing, contextmanager from dataclasses import dataclass -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Callable, Mapping, Optional +from urllib.parse import unquote, urlparse import pyarrow as pa +import pyarrow.fs as pafs 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.hdf5 import ( + _Hdf5SourceFileIO, _SnapshotRecorder, + _normalize_source_path, + _qualified_status_path, _strict_arrow_table, + _validated_source_options, ) from pypaimon.multimodal.table import _target_schema @@ -83,12 +95,14 @@ def load_from_lerobot( transform: Optional[Callable] = None, feature_mapping: Optional[Mapping[str, str]] = None, batch_size: int = 1024, - options: Optional[Mapping[str, object]] = None): + 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`. + :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( @@ -100,81 +114,100 @@ def load_from_lerobot( or batch_size <= 0: raise ValueError("batch_size must be a positive integer.") - resolved_source, local_info = _resolve_source(source) - if local_info is not None: - _require_v3(local_info, resolved_source.path) - LeRobotDataset = _import_lerobot_dataset() - dataset = _open_dataset(LeRobotDataset, resolved_source) - 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_arrow_table( - pa.Table.from_batches([], schema=source_schema), - target_schema, - resolved_source, - 0, - format_name="LeRobot", - ) + 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_dataset(LeRobotDataset, resolved_source) + 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_arrow_table( + pa.Table.from_batches([], schema=source_schema), + target_schema, + resolved_source, + 0, + format_name="LeRobot", + ) - 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, + 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, ) - return _write_dataset( - table, - dataset, - info, - resolved_source, - source_schema, - mapped_names, - transform, - batch_size, - ) -def _resolve_source(source): +@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) - return _local_source(root) + 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.") + raise ValueError( + "source must be a local directory or Hugging Face repo_id.") value = source.strip() candidate = Path(value).expanduser() if candidate.is_dir(): - return _local_source(candidate.resolve()) + yield _local_source(candidate.resolve()) + return if candidate.is_absolute() or value.startswith((".", "~")): raise FileNotFoundError( "LeRobot source directory does not exist: %s" % candidate) - return _LeRobotSource(path=value, root=None, repo_id=value), None + if "://" not in value: + yield _LeRobotSource(path=value, root=None, repo_id=value), None + return + source_uri = _normalize_source_path(value).rstrip("/") + with tempfile.TemporaryDirectory(prefix="pypaimon_lerobot_source_") \ + as temp_dir: + root = Path(temp_dir) + _materialize_remote_source( + source_uri, root, Options(source_options)) + yield _local_source(root, source_uri) -def _local_source(root): + +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" % root) + "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 + "Cannot read LeRobot metadata %s: %s" + % (info_path, error)) from error return ( _LeRobotSource( - path=str(root), + path=str(display_path or root), root=root, repo_id="local/pypaimon-import", ), @@ -182,6 +215,69 @@ def _local_source(root): ) +def _materialize_remote_source(source_uri, root, options): + source_file_io = _Hdf5SourceFileIO(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) + _copy_remote_directory( + source_file_io, source_uri, source_uri, root) + finally: + source_file_io.close() + + +def _copy_remote_directory(source_file_io, source_root, directory, local_root): + try: + children = 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 + for status in children: + source_path = _qualified_status_path(directory, status) + relative_path = _relative_source_path(source_root, source_path) + local_path = local_root / PurePosixPath(relative_path) + if status.type == pafs.FileType.Directory: + local_path.mkdir(parents=True, exist_ok=True) + _copy_remote_directory( + source_file_io, source_root, source_path, local_root) + elif status.type == pafs.FileType.File: + local_path.parent.mkdir(parents=True, exist_ok=True) + stream = source_file_io.new_input_stream(source_path) + with closing(stream) as source_stream: + with local_path.open("wb") as output: + shutil.copyfileobj(source_stream, output) + else: + raise ValueError( + "Unsupported LeRobot source status for path: %s" + % source_path) + + +def _relative_source_path(source_root, source_path): + root_uri = urlparse(source_root) + path_uri = urlparse(source_path) + if (root_uri.scheme.lower(), root_uri.netloc) != ( + path_uri.scheme.lower(), path_uri.netloc): + raise ValueError( + "LeRobot source entry is outside %s: %s" + % (source_root, source_path)) + relative = posixpath.relpath( + unquote(path_uri.path), unquote(root_uri.path) or "/") + if relative in ("", ".") or relative == ".." \ + or relative.startswith("../"): + raise ValueError( + "LeRobot source entry is outside %s: %s" + % (source_root, source_path)) + return relative + + def _require_v3(info, source): version = str(info.get("codebase_version", "")) if not (version == "v3" or version.startswith("v3.")): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 6b62b49464a6..8197500a08cd 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -26,6 +26,7 @@ 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 ( @@ -111,6 +112,50 @@ def test_local_v2_is_rejected_before_opening(self): 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", @@ -320,6 +365,49 @@ def test_video_frames_are_independent_blob_payloads(self): mp4 = next(self.video_source.rglob("*.mp4")).read_bytes() self.assertTrue(all(body != mp4 for body in bodies)) + def test_oss_source_uses_explicit_options_and_copies_each_file_once(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)), + ) + 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( From cae41debd1d014a3755d5ad3c5e6755af90a68a7 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Fri, 28 Aug 2026 01:58:50 -0700 Subject: [PATCH 3/5] [python] Keep LeRobot schema validation isolated --- paimon-python/pypaimon/multimodal/hdf5.py | 22 ++++----- paimon-python/pypaimon/multimodal/lerobot.py | 50 ++++++++++++++++++-- 2 files changed, 55 insertions(+), 17 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/hdf5.py b/paimon-python/pypaimon/multimodal/hdf5.py index a2cbadb60a2e..68389cf4e28c 100644 --- a/paimon-python/pypaimon/multimodal/hdf5.py +++ b/paimon-python/pypaimon/multimodal/hdf5.py @@ -493,36 +493,34 @@ def _require_seekable(stream, source): ) from error -def _strict_arrow_table( - data, target_schema, source, batch_index, format_name="HDF5"): +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( - "%s transform must return Arrow data or an iterable of Arrow data." - % format_name) + "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( - "%s batch %d from %s is missing columns: %s" - % (format_name, batch_index, source.path, missing)) + "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( - "%s batch %d from %s has unexpected columns: %s" - % (format_name, batch_index, source.path, extra)) + "HDF5 batch %d from %s has unexpected columns: %s" + % (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; " + "HDF5 batch %d from %s has columns in the wrong order: %s; " "expected %s." - % (format_name, batch_index, source.path, table.column_names, + % (batch_index, source.path, table.column_names, target_schema.names)) try: _validate_nested_nullability(table, target_schema) @@ -533,8 +531,8 @@ def _strict_arrow_table( 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 + "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): diff --git a/paimon-python/pypaimon/multimodal/lerobot.py b/paimon-python/pypaimon/multimodal/lerobot.py index 19836a28ccbc..c6dc97917c98 100644 --- a/paimon-python/pypaimon/multimodal/lerobot.py +++ b/paimon-python/pypaimon/multimodal/lerobot.py @@ -42,7 +42,7 @@ _SnapshotRecorder, _normalize_source_path, _qualified_status_path, - _strict_arrow_table, + _validate_nested_nullability, _validated_source_options, ) from pypaimon.multimodal.table import _target_schema @@ -129,12 +129,11 @@ def load_from_lerobot( table = _get_or_create_table( connection, table_name, source_schema, options) target_schema = _target_schema(table.raw_table) - _strict_arrow_table( + _strict_lerobot_table( pa.Table.from_batches([], schema=source_schema), target_schema, resolved_source, 0, - format_name="LeRobot", ) episode_count = int(info.get("total_episodes", 0)) @@ -278,6 +277,48 @@ def _relative_source_path(source_root, source_path): return relative +def _strict_lerobot_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( + "LeRobot transform must return an Arrow table or record batch.") + + missing = [ + name for name in target_schema.names if name not in table.column_names + ] + if missing: + raise ValueError( + "LeRobot 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( + "LeRobot batch %d from %s has unexpected columns: %s" + % (batch_index, source.path, extra)) + if table.column_names != target_schema.names: + raise ValueError( + "LeRobot 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( + "LeRobot batch %d from %s cannot be converted to the table " + "schema: %s" % (batch_index, source.path, error)) from error + + def _require_v3(info, source): version = str(info.get("codebase_version", "")) if not (version == "v3" or version.startswith("v3.")): @@ -480,12 +521,11 @@ def _write_dataset( transformed = _as_arrow_table(transformed) _validate_transform_order(batch, transformed, mapped_names) batch = transformed - batch = _strict_arrow_table( + batch = _strict_lerobot_table( batch, target_schema, source, batch_count, - format_name="LeRobot", ) table_write.write_arrow(batch) batch_count += 1 From deebb989ab81da12df72167035e5dbf9fdbea713 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Fri, 28 Aug 2026 02:03:31 -0700 Subject: [PATCH 4/5] [python] Share strict Arrow import validation --- .../pypaimon/multimodal/arrow_utils.py | 139 ++++++++++++++++++ paimon-python/pypaimon/multimodal/hdf5.py | 119 +-------------- paimon-python/pypaimon/multimodal/lerobot.py | 48 +----- 3 files changed, 155 insertions(+), 151 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/arrow_utils.py 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/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 index c6dc97917c98..4c7431a3924e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot.py +++ b/paimon-python/pypaimon/multimodal/lerobot.py @@ -37,12 +37,12 @@ ) 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, - _validate_nested_nullability, _validated_source_options, ) from pypaimon.multimodal.table import _target_schema @@ -278,45 +278,13 @@ def _relative_source_path(source_root, source_path): def _strict_lerobot_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( - "LeRobot transform must return an Arrow table or record batch.") - - missing = [ - name for name in target_schema.names if name not in table.column_names - ] - if missing: - raise ValueError( - "LeRobot 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( - "LeRobot batch %d from %s has unexpected columns: %s" - % (batch_index, source.path, extra)) - if table.column_names != target_schema.names: - raise ValueError( - "LeRobot 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( - "LeRobot batch %d from %s cannot be converted to the table " - "schema: %s" % (batch_index, source.path, error)) from error + return strict_arrow_table( + data, + target_schema, + source.path, + batch_index, + "LeRobot", + ) def _require_v3(info, source): From e5e2b2085de3e0b463440fc5f71aaf4e0abf4362 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Fri, 28 Aug 2026 02:18:37 -0700 Subject: [PATCH 5/5] [python] Stream remote LeRobot imports through FileIO --- docs/docs/pypaimon/multimodal-api.mdx | 7 +- paimon-python/pypaimon/multimodal/lerobot.py | 476 +++++++++++++----- .../pypaimon/tests/multimodal_lerobot_test.py | 44 +- 3 files changed, 409 insertions(+), 118 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 44a475d1df9f..38d3056d6d5f 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -354,9 +354,10 @@ result = conn.load_from_lerobot( ) ``` -The official LeRobot reader requires a local root. URI sources are therefore -copied once to a temporary directory, imported, and removed. Each MP4 is copied -once, not once per frame; ensure local temporary storage can hold the source. +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`. diff --git a/paimon-python/pypaimon/multimodal/lerobot.py b/paimon-python/pypaimon/multimodal/lerobot.py index 4c7431a3924e..6b4ac2d16487 100644 --- a/paimon-python/pypaimon/multimodal/lerobot.py +++ b/paimon-python/pypaimon/multimodal/lerobot.py @@ -18,18 +18,18 @@ import io import json -import posixpath import shutil import sys import tempfile +from bisect import bisect_right from contextlib import closing, contextmanager from dataclasses import dataclass -from pathlib import Path, PurePosixPath +from pathlib import Path from typing import Callable, Mapping, Optional -from urllib.parse import unquote, urlparse import pyarrow as pa import pyarrow.fs as pafs +import pyarrow.parquet as pq from pypaimon.catalog.catalog_exception import ( DatabaseNotExistException, @@ -85,6 +85,14 @@ 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( @@ -120,41 +128,48 @@ def load_from_lerobot( if local_info is not None: _require_v3(local_info, resolved_source.path) LeRobotDataset = _import_lerobot_dataset() - dataset = _open_dataset(LeRobotDataset, resolved_source) - 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, - ) + 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, + 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, ) - 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 @@ -183,12 +198,30 @@ def _resolved_source(source, source_options): return source_uri = _normalize_source_path(value).rstrip("/") - with tempfile.TemporaryDirectory(prefix="pypaimon_lerobot_source_") \ - as temp_dir: - root = Path(temp_dir) - _materialize_remote_source( - source_uri, root, Options(source_options)) - yield _local_source(root, source_uri) + 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): @@ -214,69 +247,6 @@ def _local_source(root, display_path=None): ) -def _materialize_remote_source(source_uri, root, options): - source_file_io = _Hdf5SourceFileIO(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) - _copy_remote_directory( - source_file_io, source_uri, source_uri, root) - finally: - source_file_io.close() - - -def _copy_remote_directory(source_file_io, source_root, directory, local_root): - try: - children = 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 - for status in children: - source_path = _qualified_status_path(directory, status) - relative_path = _relative_source_path(source_root, source_path) - local_path = local_root / PurePosixPath(relative_path) - if status.type == pafs.FileType.Directory: - local_path.mkdir(parents=True, exist_ok=True) - _copy_remote_directory( - source_file_io, source_root, source_path, local_root) - elif status.type == pafs.FileType.File: - local_path.parent.mkdir(parents=True, exist_ok=True) - stream = source_file_io.new_input_stream(source_path) - with closing(stream) as source_stream: - with local_path.open("wb") as output: - shutil.copyfileobj(source_stream, output) - else: - raise ValueError( - "Unsupported LeRobot source status for path: %s" - % source_path) - - -def _relative_source_path(source_root, source_path): - root_uri = urlparse(source_root) - path_uri = urlparse(source_path) - if (root_uri.scheme.lower(), root_uri.netloc) != ( - path_uri.scheme.lower(), path_uri.netloc): - raise ValueError( - "LeRobot source entry is outside %s: %s" - % (source_root, source_path)) - relative = posixpath.relpath( - unquote(path_uri.path), unquote(root_uri.path) or "/") - if relative in ("", ".") or relative == ".." \ - or relative.startswith("../"): - raise ValueError( - "LeRobot source entry is outside %s: %s" - % (source_root, source_path)) - return relative - - def _strict_lerobot_table(data, target_schema, source, batch_index): return strict_arrow_table( data, @@ -321,6 +291,268 @@ def _open_dataset(LeRobotDataset, source): % (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 @@ -558,7 +790,11 @@ def _episode_batches(dataset, info, batch_size): def _read_batch(dataset, info, begin, end, schema, mapped_names): - raw = dataset.hf_dataset.with_format("arrow")[begin:end] + 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): @@ -566,12 +802,20 @@ def _read_batch(dataset, info, begin, end, schema, mapped_names): features = info["features"] video_names = [name for name, feature in features.items() if feature["dtype"] == "video"] - 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])) + 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 = [] @@ -587,8 +831,12 @@ def _read_batch(dataset, info, begin, end, schema, mapped_names): "LeRobot data is missing metadata feature %s." % source_name) values = raw.column(source_name).to_pylist() if dtype == "image": - values = [_image_bytes(value, dataset.root) - for value in values] + 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] diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 8197500a08cd..8f6d47528a94 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -365,7 +365,45 @@ def test_video_frames_are_independent_blob_payloads(self): mp4 = next(self.video_source.rglob("*.mp4")).read_bytes() self.assertTrue(all(body != mp4 for body in bodies)) - def test_oss_source_uses_explicit_options_and_copies_each_file_once(self): + 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 = { @@ -399,6 +437,10 @@ def test_oss_source_uses_explicit_options_and_copies_each_file_once(self): 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"