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..fb794630ee27 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -318,6 +318,48 @@ 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` imports a local directory, FileIO URI, or Hugging Face +repository. It derives the schema from `meta/info.json`, writes one row per +frame, and commits once. + +```shell +pip install 'pypaimon[lerobot]' +``` + +```python +snapshot_id = conn.load_from_lerobot( + "robot_data", + "/data/lerobot_dataset", +) +print(snapshot_id) +``` + +The return value is `None` when the source has no frames. + +For FileIO URIs, pass credentials through `source_options`: + +```python +snapshot_id = 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", + }, +) +``` + +Missing tables are created from metadata; existing tables use strict schema +validation and append semantics. Scalars map to scalar types, vectors to +`VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images +keep their compressed bytes. + +Only v3 is supported. Video features, `uint64`, and language event structures +are rejected. + ## 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..a008dbe2608d 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -31,6 +31,29 @@ pip3 install dist/*.tar.gz The command will install the package and core dependencies to your local Python environment. +# Load LeRobot Dataset v3 + +Install the optional dependency, then import a local directory, FileIO URI, or +Hugging Face repository: + +```commandline +pip install 'pypaimon[lerobot]' +``` + +```python +import pypaimon.multimodal as pmm + +connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) +snapshot_id = connection.load_from_lerobot( + "robot_data", + "/data/lerobot_dataset", +) +print(snapshot_id) +``` + +The schema comes from `meta/info.json`. Each frame becomes one row; media uses +BLOB columns. Missing tables are created and later calls append. + # HDF5 to multimodal tables HDF5 loading requires Python 3.8 or newer. Install the optional dependency and diff --git a/paimon-python/pypaimon/multimodal/arrow_utils.py b/paimon-python/pypaimon/multimodal/arrow_utils.py new file mode 100644 index 000000000000..76d056ed9211 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/arrow_utils.py @@ -0,0 +1,139 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared Arrow schema validation for multimodal format importers.""" + +import pyarrow as pa +import pyarrow.compute as pc + + +def strict_arrow_table( + data, + target_schema, + source_path, + batch_index, + format_name): + if isinstance(data, pa.RecordBatch): + table = pa.Table.from_batches([data]) + elif isinstance(data, pa.Table): + table = data + else: + raise ValueError( + "%s transform must return Arrow data or an iterable of Arrow data." + % format_name) + + missing = [ + name for name in target_schema.names if name not in table.column_names + ] + if missing: + raise ValueError( + "%s batch %d from %s is missing columns: %s" + % (format_name, batch_index, source_path, missing)) + extra = [ + name for name in table.column_names if name not in target_schema.names + ] + if extra: + raise ValueError( + "%s batch %d from %s has unexpected columns: %s" + % (format_name, batch_index, source_path, extra)) + if table.column_names != target_schema.names: + raise ValueError( + "%s batch %d from %s has columns in the wrong order: %s; " + "expected %s." + % (format_name, batch_index, source_path, table.column_names, + target_schema.names)) + try: + _validate_nested_nullability(table, target_schema) + if table.schema.equals(target_schema, check_metadata=False): + return table + casted = table.cast(target_schema, safe=True) + _validate_nested_nullability(casted, target_schema) + return casted + except (ValueError, TypeError, NotImplementedError) as error: + raise ValueError( + "%s batch %d from %s cannot be converted to the table schema: %s" + % (format_name, batch_index, source_path, error)) from error + + +def _validate_nested_nullability(table, schema): + for field, column in zip(schema, table.columns): + for chunk in column.chunks: + _validate_array_nullability(chunk, field, field.name) + + +def _validate_array_nullability(array, field, path): + if not field.nullable and array.null_count: + raise ValueError( + "non-nullable field %s contains %d null value(s)" + % (path, array.null_count)) + + target_type = field.type + source_type = array.type + if (pa.types.is_list(target_type) + or pa.types.is_large_list(target_type) + or pa.types.is_fixed_size_list(target_type)): + if not (pa.types.is_list(source_type) + or pa.types.is_large_list(source_type) + or pa.types.is_fixed_size_list(source_type)): + return + _validate_array_nullability( + pc.list_flatten(array), + target_type.value_field, + "%s.%s" % (path, target_type.value_field.name), + ) + return + + if pa.types.is_map(target_type): + if not pa.types.is_map(source_type): + return + start = array.offsets[0].as_py() + stop = array.offsets[-1].as_py() + length = stop - start + offsets = pc.subtract( + array.offsets, + pa.scalar(start, type=array.offsets.type), + ) + entries = pa.StructArray.from_arrays( + [array.keys.slice(start, length), + array.items.slice(start, length)], + fields=[source_type.key_field, source_type.item_field], + ) + logical_entries = pc.list_flatten(pa.ListArray.from_arrays( + offsets, + entries, + mask=pc.is_null(array), + )) + _validate_array_nullability( + logical_entries.field(0), target_type.key_field, + "%s.%s" % (path, target_type.key_field.name)) + _validate_array_nullability( + logical_entries.field(1), target_type.item_field, + "%s.%s" % (path, target_type.item_field.name)) + return + + if pa.types.is_struct(target_type): + if not pa.types.is_struct(source_type): + return + parent_valid = pc.is_valid(array) if array.null_count else None + for index, child_field in enumerate(target_type): + child = array.field(index) + if parent_valid is not None: + child = pc.filter(child, parent_valid) + _validate_array_nullability( + child, + child_field, + "%s.%s" % (path, child_field.name), + ) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 51085910a44c..3c55032b846a 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -117,6 +117,25 @@ def load_from_hdf5( source_options=source_options, ) + def load_from_lerobot( + self, + table_name: str, + source, + *, + batch_size: int = 1024, + options=None, + source_options=None): + """Import LeRobot Dataset v3 and return the committed snapshot ID.""" + from pypaimon.multimodal.lerobot import load_from_lerobot + return load_from_lerobot( + self, + table_name, + source, + batch_size=batch_size, + options=options, + source_options=source_options, + ) + def drop_table(self, name: str, ignore_if_not_exists: bool = False): self.catalog.drop_table( self._identifier(name), diff --git a/paimon-python/pypaimon/multimodal/hdf5.py b/paimon-python/pypaimon/multimodal/hdf5.py index 68389cf4e28c..f50e9e135f7c 100644 --- a/paimon-python/pypaimon/multimodal/hdf5.py +++ b/paimon-python/pypaimon/multimodal/hdf5.py @@ -26,13 +26,18 @@ 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.source_utils import ( + _source_path_text, + _validated_source_options, + _validate_source_kerberos, +) from pypaimon.multimodal.table import _target_schema from pypaimon.write.commit_callback import CommitCallback @@ -114,6 +119,9 @@ def new_input_stream(self, path): file_io, native_path = self._resolve(path) return file_io.new_input_stream(native_path) + def to_filesystem_path(self, path): + return self._resolve(path)[1] + def close(self): self._resolver.close() @@ -359,17 +367,6 @@ def _raise_legacy_directory_listing_error(path, error): % path) from error -def _source_path_text(value): - try: - path = os.fspath(value) - except TypeError as error: - raise ValueError( - "paths must contain only filesystem paths or URIs.") from error - if isinstance(path, bytes): - raise ValueError("paths must contain only filesystem paths or URIs.") - return path - - def _normalize_source_path(value): path = _source_path_text(value) parsed = urlparse(path) @@ -442,40 +439,6 @@ def _path_values(paths): "paths must be a path or an iterable of paths.") from error -def _validated_source_options(source_options): - if source_options is None: - return {} - if not isinstance(source_options, Mapping): - raise ValueError("source_options must be a mapping.") - return dict(source_options) - - -def _validate_source_kerberos(paths, source_options): - source_principal = ( - source_options.get("security.kerberos.login.principal") - or source_options.get("security.principal") - ) - source_keytab = ( - source_options.get("security.kerberos.login.keytab") - or source_options.get("security.keytab") - ) - if not source_principal and not source_keytab: - return - if bool(source_principal) != bool(source_keytab): - raise ValueError( - "Source Kerberos principal and keytab must be both set or both " - "unset.") - if not any( - urlparse(_source_path_text(path)).scheme.lower() - in ("hdfs", "viewfs") for path in paths): - return - raise ValueError( - "HDF5 sources cannot use an explicit Kerberos keytab in a shared " - "process because kinit overwrites process-global credentials. " - "Run the load in a process-isolated worker with a pre-acquired " - "ticket cache and omit the source principal and keytab options.") - - def _require_seekable(stream, source): required = ("read", "seek", "tell") if any(not callable(getattr(stream, method, None)) for method in required): @@ -494,116 +457,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/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py new file mode 100644 index 000000000000..a40f2a8ccef0 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -0,0 +1,24 @@ +# 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.""" + +from pypaimon.multimodal.lerobot.api import load_from_lerobot + + +__all__ = [ + "load_from_lerobot", +] diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py new file mode 100644 index 000000000000..3d52c0c39644 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -0,0 +1,131 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Public LeRobot import API.""" + +import sys +from typing import Mapping, Optional + +import pyarrow as pa + +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableNotExistException, +) +from pypaimon.multimodal.lerobot.loader import ( + _strict_lerobot_table, + _write_dataset, +) +from pypaimon.multimodal.lerobot.schema import ( + _require_v3, + _schema_from_info, + _validate_lerobot_schema, +) +from pypaimon.multimodal.lerobot.source import ( + _has_tasks, + _import_lerobot_dataset, + _load_hub_info, + _open_resolved_dataset, + _resolved_source, + _validate_info_paths, +) +from pypaimon.multimodal.source_utils import ( + _validated_source_options, + _validate_source_kerberos, +) +from pypaimon.multimodal.table import _target_schema + + +def load_from_lerobot( + connection, + table_name: str, + source, + *, + batch_size: int = 1024, + options: Optional[Mapping[str, object]] = None, + source_options: Optional[Mapping[str, object]] = None): + """Import LeRobot Dataset v3 and return the committed snapshot ID. + + A missing target table is created from LeRobot metadata. An existing table + receives the same strict schema validation and append semantics as + :meth:`MultimodalConnection.load_from_hdf5`. FileIO URI credentials come + only from ``source_options`` and are not inherited from the target Catalog. + """ + if sys.version_info < (3, 10): + raise RuntimeError( + "load_from_lerobot requires Python 3.10 or newer; install and " + "run 'pypaimon[lerobot]' on a supported Python version.") + if isinstance(batch_size, bool) or not isinstance(batch_size, int) \ + or batch_size <= 0: + raise ValueError("batch_size must be a positive integer.") + + validated_source_options = _validated_source_options(source_options) + _validate_source_kerberos( + [source], validated_source_options, "LeRobot") + with _resolved_source(source, validated_source_options) as ( + resolved_source, local_info): + if local_info is None: + local_info = _load_hub_info(resolved_source) + _require_v3(local_info, resolved_source.path) + _validate_info_paths(local_info) + _schema_from_info(local_info, include_task=False) + LeRobotDataset = _import_lerobot_dataset() + dataset = _open_resolved_dataset( + LeRobotDataset, resolved_source, local_info) + try: + info = dict(dataset.meta.info) + _require_v3(info, resolved_source.path) + + source_schema = _schema_from_info( + info, include_task=_has_tasks(dataset, info)) + table = _get_or_create_table( + connection, table_name, source_schema, options) + target_schema = _target_schema(table.raw_table) + _validate_lerobot_schema( + source_schema, target_schema, resolved_source.path) + _strict_lerobot_table( + pa.Table.from_batches([], schema=source_schema), + target_schema, + resolved_source, + 0, + ) + + row_count = int(info.get("total_frames", len(dataset))) + if row_count == 0: + return None + return _write_dataset( + table, + dataset, + info, + resolved_source, + source_schema, + batch_size, + ) + finally: + close = getattr(dataset, "close", None) + if callable(close): + close() + + +def _get_or_create_table(connection, table_name, schema, options): + try: + return connection.get_table(table_name) + except (DatabaseNotExistException, TableNotExistException): + return connection.create_table( + table_name, + schema=schema, + options=options, + ) diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py new file mode 100644 index 000000000000..ecd9565ab4f3 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -0,0 +1,278 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LeRobot frame conversion and batch writing.""" + +import io +from pathlib import Path + +import pyarrow as pa + +from pypaimon.multimodal.arrow_utils import strict_arrow_table +from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.lerobot.schema import _feature_shape +from pypaimon.multimodal.table import _target_schema + + +def _strict_lerobot_table(data, target_schema, source, batch_index): + return strict_arrow_table( + data, + target_schema, + source.path, + batch_index, + "LeRobot", + ) + + +def _write_dataset( + table, + dataset, + info, + source, + source_schema, + batch_size): + target_schema = _target_schema(table.raw_table) + write_builder = table.raw_table.new_batch_write_builder() + table_write = None + table_commit = None + commit_started = False + batch_count = 0 + row_count = 0 + snapshot_recorder = _SnapshotRecorder() + + try: + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + table_commit.add_commit_callback(snapshot_recorder) + for begin, end in _episode_batches(dataset, info, batch_size): + batch = _read_batch( + dataset, info, begin, end, source_schema) + batch = _strict_lerobot_table( + batch, + target_schema, + source, + batch_count, + ) + table_write.write_arrow(batch) + batch_count += 1 + row_count += batch.num_rows + + expected_rows = int(info.get("total_frames", len(dataset))) + if row_count != expected_rows: + raise ValueError( + "LeRobot metadata reports %d frames but import produced %d." + % (expected_rows, row_count)) + messages = table_write.prepare_commit() + commit_started = True + table_commit.commit(messages) + if snapshot_recorder.snapshot_id is None: + raise RuntimeError( + "LeRobot append committed without reporting a snapshot id.") + return snapshot_recorder.snapshot_id + except BaseException: + if table_write is not None and not commit_started: + table_write.abort() + raise + finally: + try: + if table_write is not None: + table_write.close() + finally: + if table_commit is not None: + table_commit.close() + + +def _episode_batches(dataset, info, batch_size): + episodes = getattr(dataset.meta, "episodes", None) + episode_count = int(info.get("total_episodes", 0)) + total_frames = int(info.get("total_frames", len(dataset))) + if episodes is None: + raise ValueError("LeRobot v3 metadata is missing episode boundaries.") + expected_begin = 0 + for ordinal in range(episode_count): + episode = episodes.iloc[ordinal] if hasattr(episodes, "iloc") \ + else episodes[ordinal] + begin = int(_python_scalar(episode["dataset_from_index"])) + end = int(_python_scalar(episode["dataset_to_index"])) + if begin != expected_begin or end <= begin: + raise ValueError( + "LeRobot episode %d has invalid frame range [%d, %d); " + "expected it to start at %d." + % (ordinal, begin, end, expected_begin)) + while begin < end: + batch_end = min(begin + batch_size, end) + yield begin, batch_end + begin = batch_end + expected_begin = end + if expected_begin != total_frames: + raise ValueError( + "LeRobot episode ranges cover %d frames but metadata reports %d." + % (expected_begin, total_frames)) + + +def _read_batch(dataset, info, begin, end, schema): + read_batch = getattr(dataset, "read_batch", None) + if callable(read_batch): + raw = read_batch(begin, end) + else: + raw = dataset.hf_dataset.with_format("arrow")[begin:end] + if isinstance(raw, pa.RecordBatch): + raw = pa.Table.from_batches([raw]) + elif not isinstance(raw, pa.Table): + raw = pa.Table.from_pydict(raw) + features = info["features"] + + arrays = [] + fields = [] + for name, feature in features.items(): + field = schema.field(name) + dtype = feature["dtype"] + if name not in raw.column_names: + raise ValueError( + "LeRobot data is missing metadata feature %s." % name) + values = raw.column(name).to_pylist() + if dtype == "image": + image_reader = getattr(dataset, "image_bytes", None) + if callable(image_reader): + values = [image_reader(value) for value in values] + else: + values = [_image_bytes(value, dataset.root) + for value in values] + else: + values = [_normalize_value(value, feature, name) + for value in values] + arrays.append(pa.array(values, type=field.type)) + fields.append(field) + + if "task" in schema.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("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: + root = Path(root).resolve() + path = Path(image_path) + if not path.is_absolute(): + path = root / path + path = path.resolve() + try: + path.relative_to(root) + except ValueError as error: + raise ValueError( + "LeRobot image path must stay within the source " + "directory: %s" % image_path) from error + 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 index < 0 or index >= len(tasks): + raise ValueError( + "LeRobot task_index %d is outside [0, %d)." + % (index, len(tasks))) + 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) diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py b/paimon-python/pypaimon/multimodal/lerobot/schema.py new file mode 100644 index 000000000000..ebb54ecbfd62 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/schema.py @@ -0,0 +1,159 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LeRobot metadata validation and Arrow schema conversion.""" + +import pyarrow as pa + + +_SCALAR_DTYPES = { + "bool": pa.bool_(), + "boolean": pa.bool_(), + "int8": pa.int8(), + "int16": pa.int16(), + "int32": pa.int32(), + "int64": pa.int64(), + "uint8": pa.int16(), + "uint16": pa.int32(), + "uint32": pa.int64(), + "float16": pa.float32(), + "float32": pa.float32(), + "float64": pa.float64(), + "string": pa.string(), +} + + +def _require_v3(info, source): + version = str(info.get("codebase_version", "")) + if not (version == "v3" or version.startswith("v3.")): + raise ValueError( + "load_from_lerobot supports LeRobot Dataset v3 only; %s reports " + "codebase_version=%r. Upgrade the dataset to v3 first." + % (source, version or None)) + + +def _schema_from_info(info, include_task): + features = info.get("features") + if not isinstance(features, dict) or not features: + raise ValueError("LeRobot metadata features must be a non-empty object.") + + fields = [] + for name, feature in features.items(): + fields.append(_feature_field(name, feature)) + if include_task: + fields.append(pa.field( + "task", + pa.string(), + nullable=False, + metadata={b"description": b"LeRobot task"}, + )) + return pa.schema(fields) + + +def _validate_lerobot_schema(source_schema, target_schema, source): + """Require an existing table to preserve the LeRobot feature contract.""" + for source_field in source_schema: + target_index = target_schema.get_field_index(source_field.name) + if target_index < 0: + # The shared schema validator reports missing columns consistently. + continue + target_field = target_schema.field(target_index) + if source_field.type != target_field.type: + raise ValueError( + "LeRobot feature %s from %s cannot be converted to the " + "table schema: expected %s, found %s." + % (source_field.name, source, source_field.type, + target_field.type)) + + source_description = _description(source_field) + if not source_description.startswith("LeRobot dtype="): + continue + target_description = _description(target_field) + if target_description != source_description: + raise ValueError( + "LeRobot feature %s from %s cannot be converted to the " + "table schema: expected %s, found %s." + % (source_field.name, source, source_description, + target_description or "no LeRobot feature metadata")) + + +def _description(field): + if not field.metadata: + return "" + return field.metadata.get(b"description", b"").decode("utf-8") + + +def _feature_field(name, feature): + if not isinstance(feature, dict): + raise ValueError( + "LeRobot feature %s metadata must be an object." % name) + dtype = str(feature.get("dtype", "")) + shape = _feature_shape(feature, name) + if dtype == "video": + raise ValueError( + "LeRobot video feature %s is not supported yet; use an " + "image-based dataset." % name) + if dtype == "image": + arrow_type = pa.large_binary() + else: + scalar_type = _SCALAR_DTYPES.get(dtype) + if scalar_type is None: + suffix = " (uint64 has no lossless Paimon integer mapping)" \ + if dtype == "uint64" else "" + raise ValueError( + "Unsupported LeRobot dtype %r for feature %s%s." + % (dtype, name, suffix)) + if pa.types.is_string(scalar_type) and shape not in ((), (1,)): + raise ValueError( + "LeRobot string feature %s must be scalar." % name) + arrow_type = _tensor_type(scalar_type, shape) + description = "LeRobot dtype=%s, shape=%s" % (dtype, list(shape)) + return pa.field( + name, + arrow_type, + nullable=False, + metadata={b"description": description.encode("utf-8")}, + ) + + +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) diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py new file mode 100644 index 000000000000..985f7330461d --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -0,0 +1,466 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LeRobot source resolution for local, Hub, and FileIO datasets.""" + +import json +import posixpath +from bisect import bisect_right +from contextlib import closing, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Optional +from urllib.parse import quote, unquote, urlparse, urlunparse + +import pyarrow as pa +import pyarrow.fs as pafs +import pyarrow.parquet as pq + +from pypaimon.common.options import Options +from pypaimon.filesystem.pyarrow_file_io import LegacyOssDirectoryListingError +from pypaimon.multimodal.hdf5 import ( + _Hdf5SourceFileIO, + _normalize_source_path, + _qualified_status_path, +) +from pypaimon.multimodal.lerobot.loader import _encode_media_frame + + +@dataclass(frozen=True) +class _LeRobotSource: + path: str + root: Optional[Path] + repo_id: str + file_io: object = None + + +@dataclass(frozen=True) +class _RemoteLeRobotMeta: + info: dict + episodes: list + tasks: list + + +@contextmanager +def _resolved_source(source, source_options): + if isinstance(source, Path): + root = source.expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % root) + yield _local_source(root) + return + if not isinstance(source, str) or not source.strip(): + raise ValueError( + "source must be a local directory or Hugging Face repo_id.") + + value = source.strip() + candidate = Path(value).expanduser() + if candidate.is_dir(): + yield _local_source(candidate.resolve()) + return + if candidate.is_absolute() or value.startswith((".", "~")): + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % candidate) + if "://" not in value: + yield _LeRobotSource(path=value, root=None, repo_id=value), None + return + + source_uri = _normalize_source_path(value).rstrip("/") + source_file_io = _Hdf5SourceFileIO(Options(source_options)) + try: + try: + status = source_file_io.get_file_status(source_uri) + except FileNotFoundError as error: + raise FileNotFoundError( + "LeRobot source directory does not exist: %s" % source_uri + ) from error + if status.type != pafs.FileType.Directory: + raise ValueError( + "LeRobot URI source must be a directory: %s" % source_uri) + info = _read_remote_json( + source_file_io, _remote_path(source_uri, "meta/info.json")) + yield ( + _LeRobotSource( + path=source_uri, + root=None, + repo_id="", + file_io=source_file_io, + ), + info, + ) + finally: + source_file_io.close() + + +def _local_source(root, display_path=None): + info_path = root / "meta" / "info.json" + if not info_path.is_file(): + raise ValueError( + "LeRobot source is missing meta/info.json: %s" + % (display_path or root)) + try: + with info_path.open("r", encoding="utf-8") as file: + info = json.load(file) + except (OSError, ValueError) as error: + raise ValueError( + "Cannot read LeRobot metadata %s: %s" + % (info_path, error)) from error + return ( + _LeRobotSource( + path=str(display_path or root), + root=root, + repo_id="local/pypaimon-import", + ), + info, + ) + + +def _import_lerobot_dataset(): + try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + except ImportError as error: + raise ImportError( + "load_from_lerobot requires LeRobot; install " + "'pypaimon[lerobot]'.") from error + return LeRobotDataset + + +def _load_hub_info(source): + try: + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + except ImportError as error: + raise ImportError( + "load_from_lerobot requires LeRobot; install " + "'pypaimon[lerobot]'.") from error + try: + return dict(LeRobotDatasetMetadata(repo_id=source.repo_id).info) + except Exception as error: + raise ValueError( + "Cannot open LeRobot Dataset v3 metadata %s: %s" + % (source.path, error)) from error + + +def _open_dataset(LeRobotDataset, source): + try: + if source.root is not None: + return LeRobotDataset( + repo_id=source.repo_id, + root=source.root, + download_videos=False, + ) + return LeRobotDataset( + repo_id=source.repo_id, + download_videos=False, + ) + except Exception as error: + raise ValueError( + "Cannot open LeRobot Dataset v3 source %s: %s" + % (source.path, error)) from error + + +def _open_resolved_dataset(LeRobotDataset, source, info): + if source.file_io is not None: + return _RemoteLeRobotDataset(source, info) + return _open_dataset(LeRobotDataset, source) + + +class _RemoteLeRobotDataset: + + def __init__(self, source, info): + self.source = source + self.root = source.path + self._file_io = source.file_io + self._episodes = self._load_episodes(info) + self._tasks = self._load_tasks(info) + self.meta = _RemoteLeRobotMeta(info, self._episodes, self._tasks) + self._episode_starts = [ + int(episode["dataset_from_index"]) + for episode in self._episodes + ] + self._data_ranges = self._build_data_ranges(info) + self._cached_data_path = None + self._cached_data_table = None + + def __len__(self): + return int(self.meta.info.get("total_frames", 0)) + + def close(self): + self._cached_data_table = None + + def read_batch(self, begin, end): + episode = self._episode_for_range(begin, end) + relative_path = self._data_path(episode, self.meta.info) + source_path = _remote_source_path( + self.source.path, + relative_path, + "info.data_path", + self._file_io, + ) + if source_path != self._cached_data_path: + table = _read_remote_parquet(self._file_io, source_path) + expected_begin, expected_end = self._data_ranges[relative_path] + expected_rows = expected_end - expected_begin + if table.num_rows != expected_rows: + raise ValueError( + "LeRobot data file %s has %d rows; metadata expects %d." + % (source_path, table.num_rows, expected_rows)) + self._cached_data_path = source_path + self._cached_data_table = table + file_begin = self._data_ranges[relative_path][0] + return self._cached_data_table.slice(begin - file_begin, end - begin) + + def image_bytes(self, value): + if value is None: + raise ValueError("LeRobot image feature contains a null frame.") + if isinstance(value, (bytes, bytearray, memoryview)): + return bytes(value) + if isinstance(value, dict): + body = value.get("bytes") + if body is not None: + return bytes(body) + image_path = value.get("path") + if image_path: + source_path = _remote_source_path( + self.source.path, + image_path, + "image path", + self._file_io, + ) + return _read_remote_bytes(self._file_io, source_path) + return _encode_media_frame(value) + + def _load_episodes(self, info): + episode_count = int(info.get("total_episodes", 0)) + if episode_count == 0: + return [] + directory = _remote_path(self.source.path, "meta/episodes") + paths = _remote_parquet_files(self._file_io, directory) + rows = [] + for path in paths: + rows.extend(_read_remote_parquet( + self._file_io, path).to_pylist()) + 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): + path = info["data_path"].format( + chunk_index=int(episode["data/chunk_index"]), + file_index=int(episode["data/file_index"]), + ) + return _relative_dataset_path(path, "info.data_path") + + +def _remote_path(root, relative_path): + return "%s/%s" % (root.rstrip("/"), relative_path.lstrip("/")) + + +def _relative_dataset_path(path, name): + if not isinstance(path, str) or not path: + raise ValueError("LeRobot %s must be a relative path." % name) + path = _fully_unquote(path) + if "\\" in path or urlparse(path).scheme or path.startswith(("/", "~")): + raise ValueError( + "LeRobot %s must stay within the source directory: %s" + % (name, path)) + normalized = posixpath.normpath(path) + if normalized in (".", "..") or normalized.startswith("../"): + raise ValueError( + "LeRobot %s must stay within the source directory: %s" + % (name, path)) + return normalized + + +def _remote_source_path(root, path, name, source_file_io=None): + if not isinstance(path, str) or not path: + _relative_dataset_path(path, name) + root_uri = urlparse(root) + path_uri = urlparse(path) + root_path = posixpath.normpath(_fully_unquote(root_uri.path) or "/") + if path_uri.scheme: + if path_uri.query or path_uri.fragment \ + or path_uri.scheme.lower() != root_uri.scheme.lower() \ + or path_uri.netloc != root_uri.netloc: + raise ValueError( + "LeRobot %s must stay within the source directory: %s" + % (name, path)) + source_path = posixpath.normpath( + _fully_unquote(path_uri.path) or "/") + else: + relative_path = _relative_dataset_path(path, name) + source_path = posixpath.normpath(posixpath.join( + root_path, + relative_path, + )) + if posixpath.commonpath([root_path, source_path]) != root_path: + raise ValueError( + "LeRobot %s must stay within the source directory: %s" + % (name, path)) + source_uri = urlunparse(( + root_uri.scheme, + root_uri.netloc, + quote(source_path, safe="/:"), + "", + "", + "", + )) + _validate_filesystem_containment( + source_file_io, root, source_uri, name, path) + return source_uri + + +def _fully_unquote(path): + decoded = unquote(path) + while decoded != path: + path = decoded + decoded = unquote(path) + return decoded + + +def _validate_filesystem_containment( + source_file_io, root, source, name, original_path): + to_filesystem_path = getattr( + source_file_io, "to_filesystem_path", None) + if not callable(to_filesystem_path): + return + root_path = _fully_unquote(to_filesystem_path(root)) + source_path = _fully_unquote(to_filesystem_path(source)) + if urlparse(root).scheme.lower() == "file": + root_path = Path(root_path).resolve() + source_path = Path(source_path).resolve() + try: + source_path.relative_to(root_path) + except ValueError as error: + raise ValueError( + "LeRobot %s must stay within the source directory: %s" + % (name, original_path)) from error + return + else: + root_path = posixpath.normpath(root_path) + source_path = posixpath.normpath(source_path) + relative_path = posixpath.relpath(source_path, root_path) + if relative_path == ".." or relative_path.startswith("../"): + raise ValueError( + "LeRobot %s must stay within the source directory: %s" + % (name, original_path)) + + +def _validate_info_paths(info): + for name in ("data_path", "video_path"): + path = info.get(name) + if path is not None: + _relative_dataset_path(path, "info.%s" % name) + + +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 diff --git a/paimon-python/pypaimon/multimodal/source_utils.py b/paimon-python/pypaimon/multimodal/source_utils.py new file mode 100644 index 000000000000..1e3c58014e39 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/source_utils.py @@ -0,0 +1,67 @@ +# 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 validation for external multimodal sources.""" + +import os +from typing import Mapping +from urllib.parse import urlparse + + +def _source_path_text(value): + try: + path = os.fspath(value) + except TypeError as error: + raise ValueError( + "paths must contain only filesystem paths or URIs.") from error + if isinstance(path, bytes): + raise ValueError("paths must contain only filesystem paths or URIs.") + return path + + +def _validated_source_options(source_options): + if source_options is None: + return {} + if not isinstance(source_options, Mapping): + raise ValueError("source_options must be a mapping.") + return dict(source_options) + + +def _validate_source_kerberos(paths, source_options, source_name="HDF5"): + source_principal = ( + source_options.get("security.kerberos.login.principal") + or source_options.get("security.principal") + ) + source_keytab = ( + source_options.get("security.kerberos.login.keytab") + or source_options.get("security.keytab") + ) + if not source_principal and not source_keytab: + return + if bool(source_principal) != bool(source_keytab): + raise ValueError( + "Source Kerberos principal and keytab must be both set or both " + "unset.") + if not any( + urlparse(_source_path_text(path)).scheme.lower() + in ("hdfs", "viewfs") for path in paths): + return + raise ValueError( + "%s sources cannot use an explicit Kerberos keytab in a shared " + "process because kinit overwrites process-global credentials. " + "Run the load in a process-isolated worker with a pre-acquired " + "ticket cache and omit the source principal and keytab options." + % 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..26a5630e7448 --- /dev/null +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -0,0 +1,600 @@ +# 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 Mock, patch + +import numpy as np +import pyarrow as pa +import pyarrow.fs as pafs + +import pypaimon.multimodal as pmm +from pypaimon.common.options import Options +from pypaimon.multimodal.hdf5 import _Hdf5SourceFileIO +from pypaimon.multimodal.lerobot import load_from_lerobot +from pypaimon.multimodal.lerobot.loader import ( + _image_bytes, + _task_name, +) +from pypaimon.multimodal.lerobot.schema import ( + _schema_from_info, + _validate_lerobot_schema, +) +from pypaimon.multimodal.lerobot.source import ( + _LeRobotSource, + _import_lerobot_dataset, + _open_dataset, + _remote_source_path, + _validate_info_paths, +) + +try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset +except ImportError: + LeRobotDataset = None + + +class LeRobotValidationTest(unittest.TestCase): + + def test_dataset_open_never_downloads_videos(self): + calls = [] + + class Dataset: + + def __init__(self, **kwargs): + calls.append(kwargs) + + _open_dataset(Dataset, _LeRobotSource( + path="lerobot/example", + root=None, + repo_id="lerobot/example", + )) + self.assertFalse(calls[0]["download_videos"]) + + def test_dataset_paths_cannot_escape_source(self): + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_paths_")) + try: + root = temp_dir / "source" + root.mkdir() + inside = root / "frame.png" + inside.write_bytes(b"frame") + outside = temp_dir / "secret" + outside.write_bytes(b"secret") + + self.assertEqual( + b"frame", + _image_bytes({"path": str(inside)}, root), + ) + for path in ("../secret", str(outside)): + with self.assertRaisesRegex(ValueError, "within the source"): + _image_bytes({"path": path}, root) + + _validate_info_paths({ + "data_path": "data/chunk-{chunk_index:03d}/file.parquet", + }) + for path in ( + "../secret.parquet", + "%2e%2e/secret.parquet", + "%252e%252e/secret.parquet"): + with self.assertRaisesRegex(ValueError, "info.data_path"): + _validate_info_paths({"data_path": path}) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + def test_remote_image_paths_cannot_escape_source(self): + root = "oss://bucket/datasets/robot" + self.assertEqual( + root + "/images/frame.png", + _remote_source_path(root, "images/frame.png", "image path"), + ) + self.assertEqual( + root + "/images/frame.png", + _remote_source_path( + root, + root + "/images/frame.png", + "image path", + ), + ) + for path in ( + "../private", + "%2e%2e/private", + "%252e%252e/private", + "oss://other/private", + "oss://bucket/datasets/robot/../../private"): + with self.assertRaisesRegex(ValueError, "within the source"): + _remote_source_path(root, path, "image path") + + def test_double_encoded_file_uri_cannot_escape_source(self): + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_uri_")) + source_file_io = _Hdf5SourceFileIO(Options({})) + try: + root = temp_dir / "source" + root.mkdir() + (root / "frame one").write_bytes(b"frame") + (temp_dir / "secret").write_bytes(b"secret") + source_path = _remote_source_path( + root.as_uri(), + "frame%20one", + "image path", + source_file_io, + ) + with source_file_io.new_input_stream(source_path) as stream: + self.assertEqual(b"frame", stream.read()) + with self.assertRaisesRegex(ValueError, "within the source"): + _remote_source_path( + root.as_uri(), + "%252e%252e/secret", + "image path", + source_file_io, + ) + finally: + source_file_io.close() + shutil.rmtree(temp_dir, ignore_errors=True) + + def test_hdfs_source_rejects_explicit_keytab_before_resolution(self): + with patch( + "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO" + ) as source_file_io: + with self.assertRaisesRegex(ValueError, "process-isolated"): + load_from_lerobot( + Mock(), + "frames", + "hdfs://source-ns/robot", + source_options={ + "security.kerberos.login.principal": "source@REALM", + "security.kerberos.login.keytab": "/source.keytab", + }, + ) + source_file_io.assert_not_called() + + def test_negative_task_index_is_rejected(self): + with self.assertRaisesRegex(ValueError, "task_index -1"): + _task_name(["pick", "place"], -1) + self.assertEqual("place", _task_name(["pick", "place"], 1)) + + 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 = _schema_from_info(info, include_task=True) + + self.assertEqual( + ["scalar", "vector", "tensor", "image", "task"], + schema.names, + ) + self.assertEqual(pa.int32(), schema.field("scalar").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) + + info["features"]["scalar"]["dtype"] = "uint64" + with self.assertRaisesRegex(ValueError, "no lossless Paimon integer"): + _schema_from_info(info, include_task=False) + + info["features"] = { + "camera": { + "dtype": "video", + "shape": [8, 10, 3], + } + } + with self.assertRaisesRegex(ValueError, "video feature camera.*not supported"): + _schema_from_info(info, include_task=False) + + def test_existing_schema_preserves_lerobot_feature_contract(self): + source = _schema_from_info({ + "features": { + "scalar": {"dtype": "float32", "shape": [1]}, + "vector": {"dtype": "float32", "shape": [3]}, + "tensor": {"dtype": "float32", "shape": [2, 3]}, + "image": {"dtype": "image", "shape": [8, 10, 3]}, + } + }, include_task=False) + + replacements = { + "shape": pa.field( + "tensor", + source.field("tensor").type, + nullable=False, + metadata={ + b"description": b"LeRobot dtype=float32, shape=[5, 3]", + }, + ), + "dtype": pa.field( + "scalar", + pa.float64(), + nullable=False, + metadata={ + b"description": b"LeRobot dtype=float64, shape=[1]", + }, + ), + "array": pa.field( + "vector", + pa.list_(pa.float32()), + nullable=False, + metadata=source.field("vector").metadata, + ), + "bytes": pa.field( + "image", + pa.binary(), + nullable=False, + metadata=source.field("image").metadata, + ), + } + for name, replacement in replacements.items(): + with self.subTest(name=name): + target = pa.schema([ + replacement if field.name == replacement.name else field + for field in source + ]) + with self.assertRaisesRegex( + ValueError, "cannot be converted"): + _validate_lerobot_schema(source, target, "dataset") + + 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) + + def test_local_video_is_rejected_before_opening(self): + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_video_")) + try: + info_dir = temp_dir / "meta" + info_dir.mkdir() + (info_dir / "info.json").write_text(json.dumps({ + "codebase_version": "v3.0", + "features": { + "camera": {"dtype": "video", "shape": [8, 10, 3]}, + }, + })) + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + with self.assertRaisesRegex( + ValueError, "video feature camera.*not supported"): + connection.load_from_lerobot("frames", temp_dir) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + +class _RemoteLeRobotFileIO: + + def __init__(self, local_root, remote_root): + self.local_root = Path(local_root) + self.remote_root = remote_root.rstrip("/") + self.opened_paths = [] + self.close_count = 0 + + def _local_path(self, remote_path): + prefix = self.remote_root + "/" + if remote_path == self.remote_root: + return self.local_root + if not remote_path.startswith(prefix): + raise FileNotFoundError(remote_path) + return self.local_root / remote_path[len(prefix):] + + def _status(self, local_path): + relative = local_path.relative_to(self.local_root).as_posix() + remote_path = self.remote_root + if relative != ".": + remote_path += "/" + relative + native_path = remote_path.split("://", 1)[1] + file_type = pafs.FileType.Directory if local_path.is_dir() \ + else pafs.FileType.File + return pafs.FileInfo(native_path, file_type) + + def get_file_status(self, remote_path): + local_path = self._local_path(remote_path) + if not local_path.exists(): + raise FileNotFoundError(remote_path) + return self._status(local_path) + + def list_status(self, remote_path): + return [self._status(path) for path in sorted( + self._local_path(remote_path).iterdir())] + + def new_input_stream(self, remote_path): + self.opened_paths.append(remote_path) + return self._local_path(remote_path).open("rb") + + def close(self): + self.close_count += 1 + + +@unittest.skipUnless( + sys.version_info >= (3, 10) and LeRobotDataset is not None, + "LeRobot 0.4.x requires Python 3.10+ and the lerobot extra", +) +class LeRobotImportTest(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.source_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_source_")) + cls.image_source = cls.source_dir / "images" + cls._create_image_dataset(cls.image_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() + + def test_import_infers_schema_preserves_episodes_and_appends(self): + snapshot_id = self.connection.load_from_lerobot( + "robot_data", self.image_source, batch_size=2) + + self.assertEqual(1, 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( + 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, + ).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_snapshot_id = 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_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.source._Hdf5SourceFileIO", + return_value=source_file_io): + snapshot_id = self.connection.load_from_lerobot( + "oss_images", + source, + batch_size=2, + ) + + self.assertEqual(1, snapshot_id) + 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_existing_incompatible_schema_fails_without_snapshot(self): + info = json.loads((self.image_source / "meta" / "info.json").read_text()) + schema = _schema_from_info(info, include_task=True) + incompatible_fields = { + "shape": pa.field( + "observation.matrix", + schema.field("observation.matrix").type, + nullable=False, + metadata={ + b"description": b"LeRobot dtype=float32, shape=[5, 2]", + }, + ), + "dtype": pa.field( + "action", + pa.list_(pa.float64(), 2), + nullable=False, + metadata={ + b"description": b"LeRobot dtype=float64, shape=[2]", + }, + ), + "array": pa.field( + "action", + pa.list_(pa.float32()), + nullable=False, + metadata=schema.field("action").metadata, + ), + "bytes": pa.field( + "observation.image", + pa.binary(), + nullable=False, + metadata=schema.field("observation.image").metadata, + ), + } + for name, replacement in incompatible_fields.items(): + with self.subTest(name=name): + table_name = "incompatible_%s" % name + table = self.connection.create_table( + table_name, + schema=pa.schema([ + replacement if field.name == replacement.name + else field + for field in schema + ]), + options={ + "file.format": "parquet", + "vector.file.format": "parquet", + }, + ) + + with self.assertRaisesRegex( + ValueError, "cannot be converted"): + self.connection.load_from_lerobot( + table_name, self.image_source) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/setup.py b/paimon-python/setup.py index bf83fe4386fd..0bbe0b41a99d 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -237,6 +237,14 @@ 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. Pandas 2.2.2+ supports NumPy 2.x selected + # by LeRobot's media dependencies. + 'datasets>=4,<4.1; python_version>="3.10"', + 'pandas>=2.2.2,<3; python_version>="3.10"', + 'lerobot>=0.4.4,<0.5; python_version>="3.10"', + ], 'ray': [ 'ray>=2.10,<3; python_version>="3.8"', ],