diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index b3c59043a8ac..c50acc7eb1d9 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -617,26 +617,31 @@ orphan files for normal Paimon cleanup. `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. +frame, and creates a LeRobot dataset backed by the frame table, +`__versions`, `
__episodes`, `
__tasks`, and an optional +`
__subtasks`. Task text remains in the metadata table; frames retain +`task_index`. After all components are committed and tagged with the same +numeric `version_id`, one row is appended to `
__versions` to publish the +version. ```shell pip install 'pypaimon[lerobot]' ``` ```python -snapshot_id = conn.load_from_lerobot( +version_id = conn.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(snapshot_id) +print(version_id) ``` -The return value is `None` when the source has no frames. +The returned `version_id` is the common tag name for all dataset components. For FileIO URIs, pass credentials through `source_options`: ```python -snapshot_id = conn.load_from_lerobot( +version_id = conn.load_from_lerobot( "robot_data", "oss://source-bucket/lerobot_dataset", source_options={ @@ -647,13 +652,14 @@ snapshot_id = conn.load_from_lerobot( ) ``` -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. +A row in `
__versions` identifies a published release. Readers must first +resolve that row, then read every required component through its matching tag; +a missing tag is an error and must not fall back to the latest snapshot. +The component tags are immutable and must be retained or deleted together. +The one-time importer requires a new target table. -Only v3 is supported. Video features, `uint64`, and language event structures -are rejected. +Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested +`ARRAY`, and images to `BLOB`. Images keep their compressed bytes. ## Overwrite diff --git a/paimon-python/README.md b/paimon-python/README.md index 75ee07090d23..ac766d5c866e 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -44,15 +44,17 @@ pip install 'pypaimon[lerobot]' import pypaimon.multimodal as pmm connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) -snapshot_id = connection.load_from_lerobot( +version_id = connection.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(snapshot_id) +print(version_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. +The source dataset must be non-empty. Its schema comes from `meta/info.json`. +Each frame becomes one row; media uses BLOB columns. The import creates frame, +Episode, task, and version tables and tags the three component tables with the +returned `version_id`. # HDF5 to multimodal tables diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 8f51cdaf9725..938622536763 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -125,7 +125,7 @@ def load_from_lerobot( batch_size: int = 1024, options=None, source_options=None): - """Import LeRobot Dataset v3 and return the committed snapshot ID.""" + """Import LeRobot Dataset v3 into a new Paimon table group.""" from pypaimon.multimodal.lerobot import load_from_lerobot return load_from_lerobot( self, diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 0b7f76d78a46..6078ac7b245e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -20,23 +20,24 @@ 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.catalog.catalog_exception import TableAlreadyExistException +from pypaimon.multimodal.lerobot.metadata import ( + _append_arrow_tables, + _load_dataset_metadata, + _managed_table_options, + _prepare_metadata_tables, + _positive_integer, + _publish_dataset, + _validated_episode_tables, ) +from pypaimon.multimodal.lerobot.loader import _write_dataset from pypaimon.multimodal.lerobot.schema import ( _require_v3, _schema_from_info, - _validate_lerobot_schema, + _validate_v3_required_features, ) from pypaimon.multimodal.lerobot.source import ( - _has_tasks, + _close_quietly, _import_lerobot_dataset, _load_hub_info, _open_resolved_dataset, @@ -47,7 +48,6 @@ _validated_source_options, _validate_source_kerberos, ) -from pypaimon.multimodal.table import _target_schema def load_from_lerobot( @@ -57,13 +57,14 @@ def load_from_lerobot( *, 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. + source_options: Optional[Mapping[str, object]] = None, +) -> int: + """Import LeRobot Dataset v3 and return its version ID. + + A new target table is created from LeRobot metadata. Episode, task, and + version metadata are stored in companion Paimon tables. + FileIO URI credentials come only from ``source_options`` and are not + inherited from the target Catalog. """ if sys.version_info < (3, 10): raise RuntimeError( @@ -82,55 +83,79 @@ def load_from_lerobot( 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) - total_frames, _, total_tasks = \ - _validated_counts(local_info, resolved_source.path) - if total_frames == 0: - source_schema = _schema_from_info( - local_info, - include_task=total_tasks > 0, - ) - _validated_table( - connection, - table_name, - source_schema, - options, - resolved_source, - ) - return None + _schema_from_info(local_info) + _positive_integer(local_info.get("fps"), "fps") + _validated_counts(local_info, resolved_source.path) + _validate_v3_required_features(local_info) 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) - row_count, _, _ = \ - _validated_counts(info, resolved_source.path) + _validated_counts(info, resolved_source.path) + _validate_v3_required_features(info) - source_schema = _schema_from_info( - info, include_task=_has_tasks(dataset, info)) - table = _validated_table( + lerobot_schema = _schema_from_info(info) + metadata = _load_dataset_metadata( + dataset, info, resolved_source) + return _import_dataset( connection, table_name, - source_schema, - options, - resolved_source, - ) - - if row_count == 0: - return None - return _write_dataset( - table, dataset, info, resolved_source, - source_schema, + lerobot_schema, batch_size, + options, + metadata, ) finally: close = getattr(dataset, "close", None) if callable(close): - close() + _close_quietly(dataset, "dataset") + + +def _import_dataset( + connection, + table_name, + dataset, + info, + source, + source_schema, + batch_size, + options, + metadata): + table = _create_target_table( + connection, table_name, source_schema, options) + tables = _prepare_metadata_tables( + connection, table.raw_table, metadata) + version_id = 1 + episodes_snapshot_id = _append_arrow_tables( + tables["episodes"], + _validated_episode_tables(metadata), + ) + frames_snapshot_id = None + if int(info["total_frames"]) > 0: + frames_snapshot_id = _write_dataset( + table, + dataset, + info, + source, + source_schema, + batch_size, + metadata, + ) + _publish_dataset( + connection, + tables, + version_id, + metadata, + table.identifier, + frames_snapshot_id, + episodes_snapshot_id, + ) + return version_id def _validated_counts(info, source): @@ -142,6 +167,9 @@ def _validated_counts(info, source): "LeRobot metadata %s has inconsistent counts: total_frames=%d " "and total_episodes=%d must both be zero or both be positive." % (source, total_frames, total_episodes)) + if total_frames == 0: + raise ValueError( + "load_from_lerobot requires a non-empty LeRobot Dataset v3.") return total_frames, total_episodes, total_tasks @@ -159,28 +187,26 @@ def _required_count(info, name, source): return int(value) -def _validated_table( - connection, table_name, source_schema, options, source): - 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, source.path) - _strict_lerobot_table( - pa.Table.from_batches([], schema=source_schema), - target_schema, - source, - 0, - ) - return table - - -def _get_or_create_table(connection, table_name, schema, options): +def _create_target_table( + connection, table_name, source_schema, options): + create_options = dict(options or {}) + managed_options = _managed_table_options( + connection._identifier(table_name)) + reserved_options = set(managed_options).intersection(create_options) + if reserved_options: + raise ValueError( + "%s are managed by load_from_lerobot." + % sorted(reserved_options)) + create_options.update(managed_options) try: - return connection.get_table(table_name) - except (DatabaseNotExistException, TableNotExistException): - return connection.create_table( + table = connection.create_table( table_name, - schema=schema, - options=options, + schema=source_schema, + options=create_options, ) + except TableAlreadyExistException as error: + raise ValueError( + "LeRobot target %s already exists; use a new target table." + % connection._identifier(table_name) + ) from error + return table diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index 1f46e047d2b4..295eee99c363 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -67,7 +67,8 @@ def _write_dataset( info, source, source_schema, - batch_size): + batch_size, + metadata): target_schema = _target_schema(table.raw_table) write_builder = table.raw_table.new_batch_write_builder() table_write = None @@ -75,15 +76,37 @@ def _write_dataset( commit_started = False batch_count = 0 row_count = 0 + episodes = metadata["episodes"] + current_episode = None + expected_tasks = set() + observed_tasks = set() 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): + for episode_index, episode_begin, task_indices, begin, end in \ + _episode_batches(dataset, info, batch_size, episodes): + if episode_index != current_episode: + if current_episode is not None: + _validate_episode_tasks( + current_episode, expected_tasks, observed_tasks) + current_episode = episode_index + expected_tasks = set(task_indices) + observed_tasks = set() batch = _read_batch( dataset, info, begin, end, source_schema) + seen_tasks = _validate_frame_controls( + batch, + int(info["fps"]), + episode_index, + episode_begin, + begin, + task_indices, + metadata["subtask_indices"], + ) + observed_tasks.update(seen_tasks) batch = _strict_lerobot_table( batch, target_schema, @@ -94,6 +117,10 @@ def _write_dataset( batch_count += 1 row_count += batch.num_rows + if current_episode is not None: + _validate_episode_tasks( + current_episode, expected_tasks, observed_tasks) + expected_rows = int(info.get("total_frames", len(dataset))) if row_count != expected_rows: raise ValueError( @@ -119,26 +146,33 @@ def _write_dataset( table_commit.close() -def _episode_batches(dataset, info, batch_size): - episodes = getattr(dataset.meta, "episodes", None) +def _episode_batches(dataset, info, batch_size, episodes): 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] + episode_index = int(_python_scalar(episode["episode_index"])) begin = int(_python_scalar(episode["dataset_from_index"])) end = int(_python_scalar(episode["dataset_to_index"])) - if begin != expected_begin or end <= begin: + length = int(_python_scalar(episode["length"])) + if episode_index != ordinal or begin != expected_begin \ + or end <= begin or length != end - begin: raise ValueError( - "LeRobot episode %d has invalid frame range [%d, %d); " - "expected it to start at %d." - % (ordinal, begin, end, expected_begin)) + "LeRobot episode %d has an invalid index, range, or length." + % ordinal) + episode_begin = begin + task_indices = episode.get("task_indices", ()) while begin < end: batch_end = min(begin + batch_size, end) - yield begin, batch_end + yield ( + episode_index, + episode_begin, + task_indices, + begin, + batch_end, + ) begin = batch_end expected_begin = end if expected_begin != total_frames: @@ -147,6 +181,105 @@ def _episode_batches(dataset, info, batch_size): % (expected_begin, total_frames)) +def _validate_frame_controls( + batch, + fps, + episode_index, + episode_begin, + begin, + task_indices, + subtask_indices=None): + required = [ + "index", "episode_index", "frame_index", "timestamp", "task_index" + ] + if subtask_indices is not None: + required.append("subtask_index") + missing = [name for name in required if name not in batch.column_names] + if missing: + raise ValueError( + "LeRobot frame data is missing control columns: %s." + % ", ".join(missing)) + + allowed_tasks = set(task_indices) + values = { + name: batch.column(name).to_pylist() + for name in required + } + seen_tasks = set() + for offset in range(batch.num_rows): + index = begin + offset + frame_index = index - episode_begin + _require_control_integer( + values["index"][offset], "index", index, index) + _require_control_integer( + values["episode_index"][offset], + "episode_index", + episode_index, + index, + ) + _require_control_integer( + values["frame_index"][offset], + "frame_index", + frame_index, + index, + ) + timestamp = values["timestamp"][offset] + expected_timestamp = pa.scalar( + frame_index / fps, + type=batch.schema.field("timestamp").type, + ).as_py() + if (isinstance(timestamp, bool) + or not isinstance(timestamp, numbers.Real) + or not math.isclose( + float(timestamp), float(expected_timestamp), + rel_tol=0.0, abs_tol=1e-4)): + raise ValueError( + "LeRobot frame %d has timestamp %r; expected %r." + % (index, timestamp, expected_timestamp)) + task_index = _control_integer( + values["task_index"][offset], "task_index", index) + if task_index not in allowed_tasks: + raise ValueError( + "LeRobot frame %d has task_index %d outside Episode %d " + "tasks %s." + % (index, task_index, episode_index, + sorted(allowed_tasks))) + seen_tasks.add(task_index) + if subtask_indices is not None: + subtask_index = _control_integer( + values["subtask_index"][offset], "subtask_index", index) + if subtask_index not in subtask_indices: + raise ValueError( + "LeRobot frame %d has subtask_index %d outside [0, %d)." + % (index, subtask_index, len(subtask_indices))) + return seen_tasks + + +def _validate_episode_tasks(episode_index, expected, actual): + if actual != expected: + raise ValueError( + "LeRobot Episode %d declares task indices %s but its " + "frames use %s." + % (episode_index, sorted(expected), sorted(actual))) + + +def _require_control_integer(value, name, expected, frame_index): + actual = _control_integer(value, name, frame_index) + if actual != expected: + raise ValueError( + "LeRobot frame %d has %s %d; expected %d." + % (frame_index, name, actual, expected)) + + +def _control_integer(value, name, frame_index): + value = _python_scalar(value) + if isinstance(value, bool) or not isinstance(value, numbers.Integral): + raise ValueError( + "LeRobot frame %d has non-integer %s %r." + % (frame_index, name, value)) + return int(value) + + def _read_batch(dataset, info, begin, end, schema): read_batch = getattr(dataset, "read_batch", None) if callable(read_batch): @@ -181,13 +314,6 @@ def _read_batch(dataset, info, begin, end, schema): arrays.append(_safe_array(values, field, name, dtype)) 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)) @@ -355,17 +481,3 @@ def _encode_media_frame(value): 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/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py new file mode 100644 index 000000000000..ee3879444ade --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -0,0 +1,666 @@ +# 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 component tables and version publication.""" + +from array import array +import json +import numbers +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq + +from pypaimon import Schema as PaimonSchema +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableAlreadyExistException, + TableNotExistException, +) +from pypaimon.common.identifier import Identifier +from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.table import _target_schema + + +_VERSION_ID = "version_id" +_PANDAS_METADATA_OPTION = "pypaimon.lerobot.pandas-metadata" +_TABLE_SUFFIXES = { + "versions": "__versions", + "episodes": "__episodes", + "tasks": "__tasks", + "subtasks": "__subtasks", +} +_COMPANION_OPTION_KEYS = { + name: "pypaimon.lerobot.%s-table" % name + for name in _TABLE_SUFFIXES +} + +_VERSIONS_SCHEMA = pa.schema([ + pa.field(_VERSION_ID, pa.int64(), nullable=False), + pa.field("info_json", pa.string(), nullable=False), + pa.field("stats_json", pa.string()), + pa.field("has_subtasks", pa.bool_(), nullable=False), +]) +_EMPTY_TASKS_SCHEMA = pa.schema([ + pa.field("task_index", pa.int64(), nullable=False), + pa.field("task", pa.string(), nullable=False), +]) +_EMPTY_EPISODES_SCHEMA = pa.schema([ + pa.field("episode_index", pa.int64(), nullable=False), + pa.field("dataset_from_index", pa.int64(), nullable=False), + pa.field("dataset_to_index", pa.int64(), nullable=False), + pa.field("tasks", pa.list_(pa.string()), nullable=False), + pa.field("length", pa.int64(), nullable=False), +]) +_EPISODE_CONTROL_COLUMNS = [ + "episode_index", + "dataset_from_index", + "dataset_to_index", + "tasks", + "length", +] + + +class _EpisodeIndex: + + def __init__(self): + self._ranges = array("q") + self._task_offsets = array("q", [0]) + self._task_indices = array("q") + + def append(self, begin, end, task_indices): + self._ranges.extend((begin, end)) + self._task_indices.extend(task_indices) + self._task_offsets.append(len(self._task_indices)) + + def __len__(self): + return len(self._ranges) // 2 + + def __getitem__(self, index): + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + task_begin = self._task_offsets[index] + task_end = self._task_offsets[index + 1] + begin = self._ranges[index * 2] + end = self._ranges[index * 2 + 1] + return { + "episode_index": index, + "dataset_from_index": begin, + "dataset_to_index": end, + "length": end - begin, + "task_indices": self._task_indices[task_begin:task_end], + } + + +def _load_dataset_metadata(dataset, info, source): + fps = _positive_integer(info.get("fps"), "fps") + stats = _source_stats(dataset, source) + tasks_table = _source_tasks( + dataset, source, int(info["total_tasks"])) + task_indices = _task_indices( + tasks_table, int(info["total_tasks"])) + subtasks_table = _source_subtasks(dataset, source) + subtask_indices = _subtask_indices(subtasks_table, info) + total_episodes = int(info["total_episodes"]) + episode_source = ( + _source_episodes(dataset, source) + if total_episodes > 0 + else {"paths": [], "schema": _EMPTY_EPISODES_SCHEMA} + ) + return { + "fps": fps, + "info_json": _canonical_json(info), + "stats_json": ( + None if stats is None else _canonical_json( + stats, allow_nan=True)), + "episodes": None, + "episodes_schema": episode_source["schema"], + "episode_paths": episode_source["paths"], + "tasks_table": tasks_table, + "subtasks_table": subtasks_table, + "source": source, + "task_indices": task_indices, + "total_frames": int(info["total_frames"]), + "total_episodes": total_episodes, + "subtask_indices": subtask_indices, + } + + +def _companion_identifier(frames_identifier, suffix): + identifier = ( + frames_identifier + if isinstance(frames_identifier, Identifier) + else Identifier.from_string(str(frames_identifier)) + ) + if identifier.is_system_table(): + raise ValueError( + "LeRobot target cannot be a Paimon system table: %s" + % frames_identifier) + companion = Identifier( + identifier.get_database_name(), + identifier.get_table_name() + suffix, + branch=identifier.get_branch_name(), + ) + return "%s.%s" % ( + _quote_identifier_part(companion.get_database_name()), + _quote_identifier_part(companion.get_object_name()), + ) + + +def _quote_identifier_part(value): + return "`%s`" % value if "." in value else value + + +def _managed_table_options(frames_identifier): + identifier = Identifier.from_string(str(frames_identifier)) + if identifier.get_branch_name() is not None: + raise ValueError( + "LeRobot import does not support table branches.") + result = {} + for name, suffix in _TABLE_SUFFIXES.items(): + result[_COMPANION_OPTION_KEYS[name]] = _companion_identifier( + frames_identifier, suffix) + return result + + +def _companion_table_identifiers(frames_table): + options = frames_table.table_schema.options + identifiers = {} + for name, key in _COMPANION_OPTION_KEYS.items(): + value = options.get(key) + if not value: + raise ValueError( + "LeRobot table %s is missing managed option %s." + % (frames_table.identifier, key)) + identifiers[name] = value + return identifiers + + +def _prepare_metadata_tables(connection, frames_table, metadata): + schemas = { + "versions": _VERSIONS_SCHEMA, + "episodes": metadata["episodes_schema"], + "tasks": metadata["tasks_table"].schema, + } + if metadata["subtasks_table"] is not None: + schemas["subtasks"] = metadata["subtasks_table"].schema + identifiers = _companion_table_identifiers(frames_table) + if metadata["subtasks_table"] is None: + try: + connection.catalog.get_table(identifiers["subtasks"]) + except (DatabaseNotExistException, TableNotExistException): + pass + else: + raise ValueError( + "LeRobot metadata table %s already exists." + % identifiers["subtasks"]) + tables = {} + for name, schema in schemas.items(): + identifier = identifiers[name] + options = {"bucket": "-1"} + pandas_metadata = (schema.metadata or {}).get(b"pandas") + if pandas_metadata is not None: + options[_PANDAS_METADATA_OPTION] = pandas_metadata.decode( + "utf-8") + paimon_schema = PaimonSchema.from_pyarrow_schema( + schema, + options=options, + ) + try: + connection.catalog.create_table( + identifier, paimon_schema, False) + except TableAlreadyExistException as error: + raise ValueError( + "LeRobot metadata table %s already exists." % identifier + ) from error + table = connection.catalog.get_table(identifier) + tables[name] = table + return tables + + +def _restore_pandas_metadata(table, data): + pandas_metadata = table.table_schema.options.get( + _PANDAS_METADATA_OPTION) + if pandas_metadata is None: + return data + metadata = dict(data.schema.metadata or {}) + metadata[b"pandas"] = pandas_metadata.encode("utf-8") + return data.replace_schema_metadata(metadata) + + +def _publish_dataset( + connection, + tables, + version_id, + metadata, + frames_identifier, + frames_snapshot_id, + episodes_snapshot_id): + _require_initial_snapshot("frames", frames_snapshot_id) + _require_initial_snapshot("episodes", episodes_snapshot_id) + tasks_snapshot_id = _append_arrow( + tables["tasks"], metadata["tasks_table"]) + _require_initial_snapshot("tasks", tasks_snapshot_id) + component_snapshots = [ + (frames_identifier, frames_snapshot_id), + (tables["episodes"].identifier, episodes_snapshot_id), + (tables["tasks"].identifier, tasks_snapshot_id), + ] + if metadata["subtasks_table"] is not None: + subtasks_snapshot_id = _append_arrow( + tables["subtasks"], metadata["subtasks_table"]) + _require_initial_snapshot("subtasks", subtasks_snapshot_id) + component_snapshots.append( + (tables["subtasks"].identifier, subtasks_snapshot_id)) + tag = str(version_id) + for identifier, snapshot_id in component_snapshots: + _create_tag(connection.catalog, identifier, tag, snapshot_id) + + manifest = _manifest_row(version_id, metadata) + _append_arrow(tables["versions"], pa.Table.from_pylist( + [manifest], schema=_VERSIONS_SCHEMA)) + + +def _require_initial_snapshot(component, snapshot_id): + if snapshot_id is None: + raise ValueError( + "LeRobot tag-backed import requires a non-empty %s component." + % component) + if snapshot_id != 1: + raise RuntimeError( + "LeRobot initial import detected concurrent writes to %s; " + "expected snapshot 1, found %d." % (component, snapshot_id)) + + +def _manifest_row( + version_id, + metadata): + return { + _VERSION_ID: version_id, + "info_json": metadata["info_json"], + "stats_json": metadata["stats_json"], + "has_subtasks": metadata["subtasks_table"] is not None, + } + + +def _append_arrow(table, data): + return _append_arrow_tables(table, [data]) + + +def _append_arrow_tables(table, tables): + builder = table.new_batch_write_builder() + table_write = None + table_commit = None + commit_started = False + recorder = _SnapshotRecorder() + try: + table_write = builder.new_write() + table_commit = builder.new_commit() + table_commit.add_commit_callback(recorder) + row_count = 0 + target_schema = _target_schema(table) + for data in tables: + if data.num_rows == 0: + continue + if not data.schema.equals(target_schema, check_metadata=False): + raise ValueError( + "LeRobot component schema %s does not match target %s." + % (data.schema, target_schema)) + table_write.write_arrow(data) + row_count += data.num_rows + del data + if row_count == 0: + table_write.abort() + return None + messages = table_write.prepare_commit() + commit_started = True + table_commit.commit(messages) + if recorder.snapshot_id is None: + raise RuntimeError("LeRobot metadata commit has no snapshot id.") + return 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 _create_tag(catalog, identifier, tag_name, snapshot_id): + try: + try: + catalog.create_tag( + identifier, tag_name, snapshot_id=snapshot_id) + except NotImplementedError: + catalog.get_table(identifier).create_tag( + tag_name, snapshot_id=snapshot_id) + except Exception as error: + try: + actual_snapshot_id = _tag_snapshot_id( + catalog, identifier, tag_name) + except Exception: + raise error + if actual_snapshot_id == snapshot_id: + return + if actual_snapshot_id is not None: + raise RuntimeError( + "LeRobot tag %s on %s points to snapshot %s; expected %s." + % (tag_name, identifier, actual_snapshot_id, snapshot_id) + ) from error + raise error + + +def _tag_snapshot_id(catalog, identifier, tag_name): + try: + response = catalog.get_tag(identifier, tag_name) + snapshot = response.snapshot + except NotImplementedError: + snapshot = catalog.get_table(identifier).tag_manager().get(tag_name) + return None if snapshot is None else snapshot.id + + +def _source_stats(dataset, source): + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import ( + _read_remote_json, + _remote_path, + ) + path = _remote_path(source.path, "meta/stats.json") + try: + source.file_io.get_file_status(path) + except FileNotFoundError: + return None + return _read_remote_json(source.file_io, path) + root = _metadata_root(dataset, source) + path = root / "meta" / "stats.json" + if not path.is_file(): + return None + with path.open("r", encoding="utf-8") as file: + return json.load(file) + + +def _source_tasks(dataset, source, total_tasks): + if total_tasks == 0: + return pa.Table.from_pylist([], schema=_EMPTY_TASKS_SCHEMA) + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import ( + _read_remote_parquet, + _remote_path, + ) + path = _remote_path(source.path, "meta/tasks.parquet") + return _read_remote_parquet(source.file_io, path) + path = _metadata_root(dataset, source) / "meta" / "tasks.parquet" + try: + return pq.read_table(path) + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot task metadata %s: %s" % (path, error) + ) from error + + +def _source_subtasks(dataset, source): + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import ( + _read_remote_parquet, + _remote_path, + ) + path = _remote_path(source.path, "meta/subtasks.parquet") + try: + source.file_io.get_file_status(path) + except FileNotFoundError: + return None + return _read_remote_parquet(source.file_io, path) + path = _metadata_root(dataset, source) / "meta" / "subtasks.parquet" + if not path.is_file(): + return None + try: + return pq.read_table(path) + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot subtask metadata %s: %s" % (path, error) + ) from error + + +def _source_episodes(dataset, source): + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import ( + _read_remote_parquet_schema, + _remote_parquet_files, + _remote_path, + ) + directory = _remote_path(source.path, "meta/episodes") + paths = _remote_parquet_files(source.file_io, directory) + + def read_schema(path): + return _read_remote_parquet_schema(source.file_io, path) + + else: + directory = _metadata_root(dataset, source) / "meta" / "episodes" + paths = sorted(directory.rglob("*.parquet")) + read_schema = pq.read_schema + + if not paths: + return { + "paths": [], + "schema": _EMPTY_EPISODES_SCHEMA, + } + try: + schemas = [read_schema(path) for path in paths] + schema = schemas[0] + if any(not item.equals(schema, check_metadata=False) + for item in schemas[1:]): + raise ValueError("Episode Parquet schemas are inconsistent.") + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot Episode metadata %s: %s" + % (directory, error)) from error + return {"paths": paths, "schema": schema} + + +def _validated_episode_tables(metadata): + episodes = _EpisodeIndex() + expected_begin = 0 + for table in _source_episode_tables(metadata): + controls = table.select(_EPISODE_CONTROL_COLUMNS) + columns = { + name: controls.column(name) + for name in _EPISODE_CONTROL_COLUMNS + } + for offset in range(controls.num_rows): + index = _integer( + columns["episode_index"][offset].as_py(), + "episode_index", + ) + begin = _integer( + columns["dataset_from_index"][offset].as_py(), + "dataset_from_index", + ) + end = _integer( + columns["dataset_to_index"][offset].as_py(), + "dataset_to_index", + ) + length = _integer( + columns["length"][offset].as_py(), "length") + if index != len(episodes) or begin != expected_begin \ + or end <= begin or length != end - begin: + raise ValueError( + "LeRobot Episode %d has inconsistent index, range, " + "or length." % len(episodes)) + names = columns["tasks"][offset].as_py() or [] + if isinstance(names, str): + names = [names] + if metadata["task_indices"] and not names: + raise ValueError( + "LeRobot Episode %d does not declare any task." % index) + try: + task_indices = [ + metadata["task_indices"][str(name)] for name in names + ] + except (KeyError, TypeError) as error: + raise ValueError( + "LeRobot Episode %d refers to an unknown task." % index + ) from error + if len(set(task_indices)) != len(task_indices): + raise ValueError( + "LeRobot Episode %d repeats a task." % index) + episodes.append(begin, end, task_indices) + expected_begin = end + yield table + del table + if len(episodes) != metadata["total_episodes"]: + raise ValueError( + "LeRobot metadata reports %d Episodes but %d were found." + % (metadata["total_episodes"], len(episodes))) + if expected_begin != metadata["total_frames"]: + raise ValueError( + "LeRobot Episode ranges cover %d frames but metadata reports %d." + % (expected_begin, metadata["total_frames"])) + metadata["episodes"] = episodes + + +def _source_episode_tables(metadata): + source = metadata["source"] + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import _read_remote_parquet + for path in metadata["episode_paths"]: + yield _read_remote_parquet(source.file_io, path) + else: + for path in metadata["episode_paths"]: + try: + yield pq.read_table(path) + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot Episode metadata %s: %s" + % (path, error)) from error + + +def _metadata_root(dataset, source): + if source.root is not None: + return Path(source.root) + root = getattr(dataset, "root", None) + if root is None: + raise ValueError( + "Cannot resolve cached LeRobot metadata for %s." % source.path) + return Path(root) + + +def _task_indices(tasks_table, total_tasks): + if total_tasks == 0: + return {} + from pypaimon.multimodal.lerobot.source import _pandas_index_column + label_column = _pandas_index_column(tasks_table.schema, "task") + records = tasks_table.select([ + "task_index", label_column + ]).to_pylist() + seen = [False] * total_tasks + by_name = {} + for record in records: + index = _integer(record.get("task_index"), "task_index") + task = record[label_column] + if index < 0 or index >= total_tasks \ + or not isinstance(task, str) or not task \ + or seen[index]: + raise ValueError("LeRobot task metadata is invalid: %s" % record) + if task in by_name: + raise ValueError("LeRobot task metadata repeats task %r." % task) + by_name[task] = index + seen[index] = True + if not all(seen): + raise ValueError( + "LeRobot task metadata does not cover [0, %d)." % total_tasks) + return by_name + + +def _subtask_indices(subtasks_table, info): + has_feature = "subtask_index" in info["features"] + if subtasks_table is None: + if has_feature: + raise ValueError( + "LeRobot frames declare subtask_index but " + "meta/subtasks.parquet is missing.") + return None + if not has_feature: + raise ValueError( + "LeRobot meta/subtasks.parquet requires a subtask_index feature.") + if "subtask_index" not in subtasks_table.column_names: + raise ValueError( + "LeRobot subtask metadata is missing subtask_index.") + from pypaimon.multimodal.lerobot.source import _pandas_index_column + label_column = _pandas_index_column(subtasks_table.schema, "subtask") + records = subtasks_table.select([ + "subtask_index", label_column + ]).to_pylist() + for expected, record in enumerate(records): + label = record[label_column] + if _integer(record.get("subtask_index"), "subtask_index") \ + != expected or not isinstance(label, str) or not label: + raise ValueError( + "LeRobot subtask metadata must provide ordered numeric and " + "text mappings for [0, %d)." + % subtasks_table.num_rows) + return range(subtasks_table.num_rows) + + +def _canonical_json(value, allow_nan=False): + return json.dumps( + _json_value(value), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=allow_nan, + ) + + +def _json_value(value): + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, dict): + return {str(key): _json_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_value(item) for item in value] + as_py = getattr(value, "as_py", None) + if callable(as_py): + return _json_value(as_py()) + tolist = getattr(value, "tolist", None) + if callable(tolist): + return _json_value(tolist()) + item = getattr(value, "item", None) + if callable(item): + return _json_value(item()) + raise TypeError("LeRobot metadata contains a non-JSON value: %r" % value) + + +def _positive_integer(value, name): + result = _integer(value, name) + if result <= 0: + raise ValueError("LeRobot metadata %s must be positive." % name) + return result + + +def _integer(value, name): + item = getattr(value, "item", None) + if callable(item): + value = item() + if isinstance(value, bool) or not isinstance(value, numbers.Integral): + raise ValueError("LeRobot metadata %s must be an integer." % name) + return int(value) diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py b/paimon-python/pypaimon/multimodal/lerobot/schema.py index 6b8ccfb8b537..de16c95a6478 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/schema.py +++ b/paimon-python/pypaimon/multimodal/lerobot/schema.py @@ -37,6 +37,14 @@ "string": pa.string(), } +_V3_REQUIRED_FEATURE_DTYPES = { + "timestamp": "float32", + "frame_index": "int64", + "episode_index": "int64", + "index": "int64", + "task_index": "int64", +} + def _require_v3(info, source): version = str(info.get("codebase_version", "")) @@ -47,22 +55,32 @@ def _require_v3(info, source): % (source, version or None)) -def _schema_from_info(info, include_task): +def _schema_from_info(info): 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) + return pa.schema([ + _feature_field(name, feature) + for name, feature in features.items() + ]) + + +def _validate_v3_required_features(info): + features = info.get("features") + if not isinstance(features, dict): + raise ValueError("LeRobot metadata features must be an object.") + expected = dict(_V3_REQUIRED_FEATURE_DTYPES) + if "subtask_index" in features: + expected["subtask_index"] = "int64" + for name, dtype in expected.items(): + feature = features.get(name) + if not isinstance(feature, dict) \ + or str(feature.get("dtype", "")) != dtype \ + or _feature_shape(feature, name) != (1,): + raise ValueError( + "LeRobot V3 required feature %s must have dtype=%s and " + "shape=[1]." % (name, dtype)) def _validate_lerobot_schema(source_schema, target_schema, source): diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index 8d71e3eecb34..04547db80403 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -16,7 +16,9 @@ """LeRobot source resolution for local, Hub, and FileIO datasets.""" +from array import array import json +import logging import posixpath from bisect import bisect_right from contextlib import closing, contextmanager @@ -39,6 +41,9 @@ from pypaimon.multimodal.lerobot.loader import _encode_media_frame +_LOGGER = logging.getLogger(__name__) + + @dataclass(frozen=True) class _LeRobotSource: path: str @@ -50,10 +55,41 @@ class _LeRobotSource: @dataclass(frozen=True) class _RemoteLeRobotMeta: info: dict - episodes: list + episodes: object tasks: list +class _RemoteEpisodeIndex: + + def __init__(self): + self.starts = array("q") + self._ends = array("q") + self._chunk_indices = array("q") + self._file_indices = array("q") + + def append(self, begin, end, chunk_index, file_index): + self.starts.append(begin) + self._ends.append(end) + self._chunk_indices.append(chunk_index) + self._file_indices.append(file_index) + + def __len__(self): + return len(self.starts) + + def __getitem__(self, index): + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + return { + "episode_index": index, + "dataset_from_index": self.starts[index], + "dataset_to_index": self._ends[index], + "data/chunk_index": self._chunk_indices[index], + "data/file_index": self._file_indices[index], + } + + @contextmanager def _resolved_source(source, source_options): if isinstance(source, Path): @@ -103,7 +139,14 @@ def _resolved_source(source, source_options): info, ) finally: - source_file_io.close() + _close_quietly(source_file_io, "source FileIO") + + +def _close_quietly(resource, name): + try: + resource.close() + except Exception: + _LOGGER.warning("Failed to close LeRobot %s.", name, exc_info=True) def _local_source(root, display_path=None): @@ -195,10 +238,7 @@ def __init__(self, source, info): 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._episode_starts = self._episodes.starts self._data_ranges = self._build_data_ranges(info) self._cached_data_path = None self._cached_data_table = None @@ -257,35 +297,52 @@ def _load_episodes(self, info): return [] directory = _remote_path(self.source.path, "meta/episodes") paths = _remote_parquet_files(self._file_io, directory) - rows = [] + episodes = _RemoteEpisodeIndex() for path in paths: - rows.extend(_read_remote_parquet( + table = _read_remote_parquet( self._file_io, path, columns=self._EPISODE_COLUMNS, - ).to_pylist()) - rows.sort(key=lambda row: int(row["episode_index"])) - if len(rows) != episode_count: + ) + columns = { + name: table.column(name) + for name in self._EPISODE_COLUMNS + } + for offset in range(table.num_rows): + episode_index = int( + columns["episode_index"][offset].as_py()) + if episode_index != len(episodes): + raise ValueError( + "LeRobot Episode metadata must be ordered by " + "episode_index.") + episodes.append( + int(columns["dataset_from_index"][offset].as_py()), + int(columns["dataset_to_index"][offset].as_py()), + int(columns["data/chunk_index"][offset].as_py()), + int(columns["data/file_index"][offset].as_py()), + ) + if len(episodes) != episode_count: raise ValueError( "LeRobot metadata reports %d Episodes but %d were found." - % (episode_count, len(rows))) - return rows + % (episode_count, len(episodes))) + return episodes 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() + table = _read_remote_parquet(self._file_io, path) + name_column = _pandas_index_column(table.schema, "task") + rows = table.select(["task_index", name_column]).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: + name = row[name_column] + if index < 0 or index >= task_count \ + or not isinstance(name, str) or not name: raise ValueError("LeRobot task metadata is invalid: %s" % row) - tasks[index] = str(name) + tasks[index] = name if any(task is None for task in tasks): raise ValueError( "LeRobot metadata reports %d tasks but %d were found." @@ -332,6 +389,28 @@ def _remote_path(root, relative_path): return "%s/%s" % (root.rstrip("/"), relative_path.lstrip("/")) +def _pandas_index_column(schema, component): + encoded = (schema.metadata or {}).get(b"pandas") + try: + pandas_metadata = json.loads(encoded.decode("utf-8")) + index_columns = pandas_metadata["index_columns"] + except (AttributeError, KeyError, TypeError, ValueError, + UnicodeDecodeError) as error: + raise ValueError( + "LeRobot %s metadata must contain a Pandas index." + % component + ) from error + if not isinstance(index_columns, list) \ + or len(index_columns) != 1 \ + or not isinstance(index_columns[0], str) \ + or index_columns[0] not in schema.names \ + or index_columns[0] == component + "_index": + raise ValueError( + "LeRobot %s metadata must contain one text Pandas index." + % component) + return index_columns[0] + + def _relative_dataset_path(path, name): if not isinstance(path, str) or not path: raise ValueError("LeRobot %s must be a relative path." % name) @@ -455,6 +534,17 @@ def _read_remote_parquet(source_file_io, path, columns=None): % (path, error)) from error +def _read_remote_parquet_schema(source_file_io, path): + stream = source_file_io.new_input_stream(path) + with closing(stream) as source_stream: + try: + return pq.read_schema(source_stream) + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot Parquet schema %s: %s" + % (path, error)) from error + + def _remote_parquet_files(source_file_io, directory): try: statuses = source_file_io.list_status(directory) @@ -470,8 +560,3 @@ def _remote_parquet_files(source_file_io, directory): 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/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 7362c91d2337..97c82b43d40a 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -15,30 +15,46 @@ # limitations under the License. import builtins +from array import array import json import shutil import sys import tempfile +import threading import unittest +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from unittest.mock import Mock, patch import numpy as np import pyarrow as pa import pyarrow.fs as pafs +import pyarrow.parquet as pq +from pypaimon.catalog.catalog_exception import TableNotExistException import pypaimon.multimodal as pmm +from pypaimon.common.identifier import Identifier from pypaimon.common.options import Options from pypaimon.multimodal.source_utils import _SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot +from pypaimon.multimodal.lerobot.metadata import ( + _append_arrow_tables, + _companion_identifier, + _load_dataset_metadata, + _managed_table_options, + _restore_pandas_metadata, + _subtask_indices, + _validated_episode_tables, +) from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, - _task_name, + _validate_frame_controls, ) from pypaimon.multimodal.lerobot.schema import ( _schema_from_info, _validate_lerobot_schema, + _validate_v3_required_features, ) from pypaimon.multimodal.lerobot.source import ( _LeRobotSource, @@ -48,6 +64,7 @@ _remote_source_path, _validate_info_paths, ) +from pypaimon.multimodal.table import _target_schema try: from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -64,8 +81,34 @@ def _replaced_contract(field, old, new): } +def _catalog_rows(connection, name): + table = connection.catalog.get_table(connection._identifier(name)) + builder = table.new_read_builder() + plan = builder.new_scan().plan() + return builder.new_read().to_arrow(plan.splits()).to_pylist() + + +def _catalog_arrow(connection, name): + table = connection.catalog.get_table(connection._identifier(name)) + builder = table.new_read_builder() + plan = builder.new_scan().plan() + return table, builder.new_read().to_arrow(plan.splits()) + + class LeRobotValidationTest(unittest.TestCase): + def test_self_contained_import_rejects_table_branches(self): + with self.assertRaisesRegex(ValueError, "does not support"): + _managed_table_options("db.robot$branch_dev") + + def test_companion_identifier_preserves_quoted_components(self): + name = _companion_identifier( + "`db.name`.`robot.data`", "__tasks") + identifier = Identifier.from_string(name) + + self.assertEqual("db.name", identifier.get_database_name()) + self.assertEqual("robot.data__tasks", identifier.get_table_name()) + def test_dataset_open_never_downloads_videos(self): calls = [] @@ -177,10 +220,69 @@ def test_hdfs_source_rejects_explicit_keytab_before_resolution(self): ) 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_timestamp_validation_quantizes_float32(self): + frame_index = 61441 + batch = pa.table({ + "index": pa.array([frame_index], type=pa.int64()), + "episode_index": pa.array([0], type=pa.int64()), + "frame_index": pa.array([frame_index], type=pa.int64()), + "timestamp": pa.array([frame_index / 30], type=pa.float32()), + "task_index": pa.array([0], type=pa.int64()), + }) + + self.assertEqual( + {0}, + _validate_frame_controls( + batch, 30, 0, 0, frame_index, [0]), + ) + + def test_subtask_index_must_reference_metadata(self): + batch = pa.table({ + "index": pa.array([0], type=pa.int64()), + "episode_index": pa.array([0], type=pa.int64()), + "frame_index": pa.array([0], type=pa.int64()), + "timestamp": pa.array([0], type=pa.float32()), + "task_index": pa.array([0], type=pa.int64()), + "subtask_index": pa.array([2], type=pa.int64()), + }) + + with self.assertRaisesRegex(ValueError, "subtask_index 2 outside"): + _validate_frame_controls( + batch, 30, 0, 0, 0, [0], range(2)) + + def test_metadata_writer_closes_after_commit_creation_failure(self): + table = Mock() + builder = table.new_batch_write_builder.return_value + writer = builder.new_write.return_value + builder.new_commit.side_effect = RuntimeError("commit init failed") + + with self.assertRaisesRegex(RuntimeError, "commit init failed"): + _append_arrow_tables(table, []) + + writer.abort.assert_called_once_with() + writer.close.assert_called_once_with() + + def test_episode_shards_use_normal_batch_rolling(self): + data = pa.table({"episode_index": [0]}) + table = Mock() + builder = table.new_batch_write_builder.return_value + writer = builder.new_write.return_value + + def shards(): + yield data + yield data + raise RuntimeError("stop after two shards") + + with patch( + "pypaimon.multimodal.lerobot.metadata._target_schema", + return_value=data.schema): + with self.assertRaisesRegex(RuntimeError, "two shards"): + _append_arrow_tables(table, shards()) + + self.assertEqual(2, writer.write_arrow.call_count) + writer.prepare_commit.assert_not_called() + writer.abort.assert_called_once_with() + writer.close.assert_called_once_with() def test_optional_dependency_error_is_actionable(self): original_import = builtins.__import__ @@ -204,10 +306,10 @@ def test_schema_comes_from_metadata_and_rejects_unsupported_types(self): "image": {"dtype": "image", "shape": [8, 10, 3]}, } } - schema = _schema_from_info(info, include_task=True) + schema = _schema_from_info(info) self.assertEqual( - ["scalar", "vector", "tensor", "image", "task"], + ["scalar", "vector", "tensor", "image"], schema.names, ) self.assertEqual(pa.int32(), schema.field("scalar").type) @@ -220,7 +322,7 @@ def test_schema_comes_from_metadata_and_rejects_unsupported_types(self): info["features"]["scalar"]["dtype"] = "uint64" with self.assertRaisesRegex(ValueError, "no lossless Paimon integer"): - _schema_from_info(info, include_task=False) + _schema_from_info(info) info["features"] = { "camera": { @@ -229,7 +331,7 @@ def test_schema_comes_from_metadata_and_rejects_unsupported_types(self): } } with self.assertRaisesRegex(ValueError, "video feature camera.*not supported"): - _schema_from_info(info, include_task=False) + _schema_from_info(info) def test_existing_schema_preserves_lerobot_feature_contract(self): source = _schema_from_info({ @@ -243,7 +345,7 @@ def test_existing_schema_preserves_lerobot_feature_contract(self): "tensor": {"dtype": "float32", "shape": [2, 3]}, "image": {"dtype": "image", "shape": [8, 10, 3]}, } - }, include_task=False) + }) replacements = { "shape": pa.field( @@ -296,6 +398,32 @@ def test_existing_schema_preserves_lerobot_feature_contract(self): ValueError, "cannot be converted"): _validate_lerobot_schema(source, target, "dataset") + def test_v3_required_features_have_native_types(self): + features = { + "timestamp": {"dtype": "float32", "shape": [1]}, + "frame_index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "index": {"dtype": "int64", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + } + _validate_v3_required_features({"features": features}) + + for name, replacement in ( + ("timestamp", {"dtype": "float64", "shape": [1]}), + ("frame_index", {"dtype": "int32", "shape": [1]}), + ("episode_index", {"dtype": "int64", "shape": [2]})): + with self.subTest(name=name): + invalid = dict(features) + invalid[name] = replacement + with self.assertRaisesRegex( + ValueError, "required feature %s" % name): + _validate_v3_required_features({"features": invalid}) + + missing = dict(features) + del missing["task_index"] + with self.assertRaisesRegex(ValueError, "required feature task_index"): + _validate_v3_required_features({"features": missing}) + def test_remote_episode_metadata_projects_stats_columns(self): source = _LeRobotSource( path="oss://bucket/robot", @@ -323,15 +451,17 @@ def test_remote_episode_metadata_projects_stats_columns(self): "pypaimon.multimodal.lerobot.source._read_remote_parquet", return_value=episode_table, ) as read_parquet: - _RemoteLeRobotDataset(source, info) + dataset = _RemoteLeRobotDataset(source, info) read_parquet.assert_called_once_with( source.file_io, "oss://bucket/robot/meta/episodes/file.parquet", columns=_RemoteLeRobotDataset._EPISODE_COLUMNS, ) + self.assertIsInstance(dataset._episode_starts, array) + self.assertNotIsInstance(dataset.meta.episodes, list) - def test_empty_local_dataset_returns_before_opening_lerobot(self): + def test_empty_local_dataset_is_rejected_before_opening_lerobot(self): temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_empty_")) try: source = temp_dir / "source" @@ -341,6 +471,7 @@ def test_empty_local_dataset_returns_before_opening_lerobot(self): "total_frames": 0, "total_episodes": 0, "total_tasks": 0, + "fps": 30, "features": { "index": {"dtype": "int64", "shape": [1]}, }, @@ -351,12 +482,11 @@ def test_empty_local_dataset_returns_before_opening_lerobot(self): with patch( "pypaimon.multimodal.lerobot.api._import_lerobot_dataset" ) as import_lerobot: - self.assertIsNone(connection.load_from_lerobot( - "empty_frames", source)) + with self.assertRaisesRegex(ValueError, "non-empty"): + connection.load_from_lerobot("empty_frames", source) import_lerobot.assert_not_called() - table = connection.get_table("empty_frames") - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) + with self.assertRaises(TableNotExistException): + connection.get_table("empty_frames") finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -371,6 +501,7 @@ def test_empty_fast_path_validates_required_counts(self): "total_frames": 0, "total_episodes": 0, "total_tasks": 0, + "fps": 30, "features": { "index": {"dtype": "int64", "shape": [1]}, }, @@ -403,6 +534,195 @@ def test_empty_fast_path_validates_required_counts(self): finally: shutil.rmtree(temp_dir, ignore_errors=True) + def test_empty_dataset_with_tasks_is_rejected(self): + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_empty_meta_")) + try: + source = temp_dir / "source" + (source / "meta").mkdir(parents=True) + info = { + "codebase_version": "v3.0", + "total_frames": 0, + "total_episodes": 0, + "total_tasks": 1, + "fps": 30, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + }, + } + (source / "meta" / "info.json").write_text(json.dumps(info)) + (source / "meta" / "stats.json").write_text(json.dumps({ + "index": {"min": [0], "max": [0]}, + })) + pq.write_table(pa.table({ + "task_index": [0], + "task": ["pick"], + }), source / "meta" / "tasks.parquet") + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + + with self.assertRaisesRegex(ValueError, "non-empty"): + connection.load_from_lerobot("frames", source) + with self.assertRaises(TableNotExistException): + connection.get_table("frames") + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + def test_optional_subtasks_keep_their_native_schema(self): + import pandas as pd + + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_subtasks_")) + try: + source = temp_dir / "source" + (source / "meta").mkdir(parents=True) + (source / "meta" / "info.json").write_text(json.dumps({ + "codebase_version": "v3.0", + "total_frames": 0, + "total_episodes": 0, + "total_tasks": 0, + "fps": 30, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + "subtask_index": {"dtype": "int64", "shape": [1]}, + }, + })) + pq.write_table(pa.Table.from_pandas(pd.DataFrame( + {"subtask_index": [0]}, + index=pd.Index(["reach"], name="instruction"), + )), source / "meta" / "subtasks.parquet") + source_info = json.loads( + (source / "meta" / "info.json").read_text()) + metadata = _load_dataset_metadata( + None, + source_info, + _LeRobotSource( + path=str(source), + root=source, + repo_id="local/subtasks", + ), + ) + + expected = pq.read_table(source / "meta" / "subtasks.parquet") + self.assertTrue(metadata["subtasks_table"].equals(expected)) + self.assertEqual([0], list(metadata["subtask_indices"])) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + def test_subtask_metadata_must_match_frame_feature(self): + import pandas as pd + + info = {"features": {"subtask_index": {}}} + with self.assertRaisesRegex(ValueError, "subtasks.parquet is missing"): + _subtask_indices(None, info) + reordered = pa.Table.from_pandas(pd.DataFrame( + {"subtask_index": [1, 0]}, + index=pd.Index(["reach", "grasp"], name="instruction"), + )) + with self.assertRaisesRegex(ValueError, "numeric and text mappings"): + _subtask_indices(reordered, info) + with self.assertRaisesRegex(ValueError, "Pandas index"): + _subtask_indices(pa.table({ + "subtask_index": [0], + }), info) + + def test_native_metadata_does_not_require_json_values(self): + import pandas as pd + + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_native_")) + try: + source = temp_dir / "source" + (source / "meta" / "episodes").mkdir(parents=True) + info = { + "codebase_version": "v3.0", + "total_frames": 1, + "total_episodes": 1, + "total_tasks": 1, + "fps": 30, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + }, + } + pq.write_table(pa.Table.from_pandas(pd.DataFrame( + { + "task_index": [0], + "native_bytes": [b"\xff"], + }, + index=pd.Index(["pick"], name="instruction"), + )), source / "meta" / "tasks.parquet") + episode_table = pa.table({ + "episode_index": [0], + "dataset_from_index": [0], + "dataset_to_index": [1], + "tasks": [["pick"]], + "length": [1], + "native_bytes": [b"\xff"], + "stats/value/mean": [float("nan")], + }) + pq.write_table( + episode_table, + source / "meta" / "episodes" / "part.parquet", + ) + (source / "meta" / "stats.json").write_text(json.dumps({ + "mean": float("nan"), + "max": float("inf"), + })) + metadata = _load_dataset_metadata( + None, + info, + _LeRobotSource( + path=str(source), + root=source, + repo_id="local/native-metadata", + ), + ) + + self.assertIsNone(metadata["episodes"]) + self.assertEqual(1, len(metadata["episode_paths"])) + stored_episode = list(_validated_episode_tables(metadata))[0] + self.assertEqual(1, len(metadata["episodes"])) + self.assertEqual( + b"\xff", + stored_episode.column("native_bytes")[0].as_py(), + ) + self.assertTrue(np.isnan( + stored_episode.column("stats/value/mean")[0].as_py())) + self.assertEqual( + b"\xff", + metadata["tasks_table"].column("native_bytes")[0].as_py(), + ) + stored_stats = json.loads(metadata["stats_json"]) + self.assertTrue(np.isnan(stored_stats["mean"])) + self.assertTrue(np.isinf(stored_stats["max"])) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + def test_invalid_fps_creates_no_snapshot_or_manifest(self): + temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_fps_")) + try: + source = temp_dir / "source" + (source / "meta").mkdir(parents=True) + (source / "meta" / "info.json").write_text(json.dumps({ + "codebase_version": "v3.0", + "total_frames": 0, + "total_episodes": 0, + "total_tasks": 0, + "fps": 0, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + }, + })) + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + + with self.assertRaisesRegex(ValueError, "fps must be positive"): + connection.load_from_lerobot("frames", source) + with self.assertRaises(TableNotExistException): + connection.get_table("frames") + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + def test_source_values_are_safely_converted(self): class Dataset: @@ -428,7 +748,7 @@ def read_batch(self, unused_begin, unused_end): for feature, value, message in cases: with self.subTest(feature=feature, value=value): info = {"features": {"value": feature}} - schema = _schema_from_info(info, include_task=False) + schema = _schema_from_info(info) with self.assertRaisesRegex(ValueError, message): _read_batch(Dataset(value), info, 0, 1, schema) @@ -443,7 +763,7 @@ def read_batch(self, unused_begin, unused_end): for feature, value in boundary_cases: with self.subTest(feature=feature, value=value): info = {"features": {"value": feature}} - schema = _schema_from_info(info, include_task=False) + schema = _schema_from_info(info) result = _read_batch(Dataset(value), info, 0, 1, schema) self.assertEqual(value, result.column("value")[0].as_py()) @@ -529,6 +849,21 @@ def close(self): self.close_count += 1 +class _FailingCloseDataset: + + def __init__(self, dataset): + self._dataset = dataset + + def __getattr__(self, name): + return getattr(self._dataset, name) + + def __len__(self): + return len(self._dataset) + + def close(self): + raise RuntimeError("close failed") + + @unittest.skipUnless( sys.version_info >= (3, 10) and LeRobotDataset is not None, "LeRobot 0.4.x requires Python 3.10+ and the lerobot extra", @@ -614,11 +949,13 @@ def _create_image_dataset(root): dataset.save_episode() dataset.finalize() - def test_import_infers_schema_preserves_episodes_and_appends(self): - snapshot_id = self.connection.load_from_lerobot( + def test_import_infers_schema_and_preserves_episodes(self): + import pandas as pd + + version_id = self.connection.load_from_lerobot( "robot_data", self.image_source, batch_size=2) - self.assertEqual(1, snapshot_id) + self.assertEqual(1, version_id) table = self.connection.get_table("robot_data") schema = table.raw_table.fields @@ -632,6 +969,10 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): self.assertEqual("FLOAT NOT NULL", types["timestamp"]) self.assertEqual("BIGINT NOT NULL", types["episode_index"]) self.assertEqual("BLOB NOT NULL", types["observation.image"]) + self.assertNotIn("dataset_id", types) + self.assertNotIn("metadata_version", types) + self.assertNotIn("version_id", types) + self.assertNotIn("task", types) rows = table.scan().select([ "episode_index", @@ -639,7 +980,6 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): "timestamp", "index", "task_index", - "task", "observation.state", "observation.matrix", "action", @@ -648,15 +988,89 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): 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([0, 0, 1, 1, 1], + [row["task_index"] 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"]) + manifests = _catalog_rows(self.connection, "robot_data__versions") + self.assertEqual(1, len(manifests)) + manifest = manifests[0] + self.assertEqual(version_id, manifest["version_id"]) + self.assertEqual("v3.0", json.loads( + manifest["info_json"])["codebase_version"]) + self.assertIsNotNone(manifest["stats_json"]) self.assertEqual( - snapshot_id, + {"version_id", "info_json", "stats_json", "has_subtasks"}, + set(manifest)) + self.assertFalse(manifest["has_subtasks"]) + tag = str(manifest["version_id"]) + self.assertEqual( + 1, + self.connection.catalog.get_tag( + table.identifier, tag).snapshot.id, + ) + for name, expected_snapshot in ( + ("robot_data__episodes", 1), + ("robot_data__tasks", 1)): + self.assertEqual( + expected_snapshot, + self.connection.catalog.get_tag( + self.connection._identifier(name), tag).snapshot.id, + ) + + episodes = _catalog_rows(self.connection, "robot_data__episodes") + episode_fields = { + field.name for field in self.connection.catalog.get_table( + self.connection._identifier( + "robot_data__episodes")).fields + } + self.assertNotIn("version_id", episode_fields) + source_episode_schema = pq.read_schema(next( + (self.image_source / "meta" / "episodes").rglob("*.parquet"))) + self.assertTrue(_target_schema( + self.connection.catalog.get_table(self.connection._identifier( + "robot_data__episodes")) + ).equals(source_episode_schema, check_metadata=False)) + self.assertEqual([(0, 0, 2), (1, 2, 5)], [ + (row["episode_index"], row["dataset_from_index"], + row["dataset_to_index"]) + for row in episodes + ]) + self.assertEqual([["pick"], ["place"]], [ + row["tasks"] for row in episodes]) + self.assertTrue(any( + name.startswith("stats/") for name in episode_fields)) + + tasks = _catalog_rows(self.connection, "robot_data__tasks") + task_fields = { + field.name for field in self.connection.catalog.get_table( + self.connection._identifier("robot_data__tasks")).fields + } + self.assertNotIn("version_id", task_fields) + self.assertTrue(_target_schema( + self.connection.catalog.get_table(self.connection._identifier( + "robot_data__tasks")) + ).equals( + pq.read_schema(self.image_source / "meta" / "tasks.parquet"), + check_metadata=False, + )) + task_name = "task" if "task" in task_fields else "__index_level_0__" + self.assertEqual([(0, "pick"), (1, "place")], [ + (row["task_index"], row[task_name]) for row in tasks + ]) + tasks_table, tasks_arrow = _catalog_arrow( + self.connection, "robot_data__tasks") + pd.testing.assert_frame_equal( + pq.read_table( + self.image_source / "meta" / "tasks.parquet").to_pandas(), + _restore_pandas_metadata( + tasks_table, tasks_arrow).to_pandas(), + ) + self.assertEqual( + 1, table.raw_table.snapshot_manager().get_latest_snapshot().id, ) @@ -675,10 +1089,265 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): [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) + with self.assertRaisesRegex(ValueError, "already exists"): + self.connection.load_from_lerobot( + "robot_data", self.image_source, batch_size=4) + self.assertEqual(5, table.scan().to_arrow().num_rows) + + def test_episode_tasks_are_validated_incrementally(self): + from pypaimon.multimodal.lerobot import loader + + with patch.object( + loader, + "_validate_episode_tasks", + wraps=loader._validate_episode_tasks) as validate: + self.connection.load_from_lerobot( + "incremental_tasks", self.image_source, batch_size=1) + + self.assertEqual([0, 1], [ + call.args[0] for call in validate.call_args_list + ]) + + def test_episode_source_shards_share_paimon_files(self): + source = self.temp_dir / "episode_shards" + shutil.copytree(self.image_source, source) + episode_path = next( + (source / "meta" / "episodes").rglob("*.parquet")) + episodes = pq.read_table(episode_path) + episode_path.unlink() + for index in range(episodes.num_rows): + pq.write_table( + episodes.slice(index, 1), + episode_path.parent / ("part-%d.parquet" % index), + ) + + self.connection.load_from_lerobot("sharded_episodes", source) + table = self.connection.catalog.get_table( + self.connection._identifier("sharded_episodes__episodes")) + files = { + file.file_name + for split in table.new_read_builder().new_scan().plan().splits() + for file in split.files + } + self.assertEqual(1, len(files)) + + def test_import_publishes_optional_subtasks(self): + import pandas as pd + + source = self.temp_dir / "with_subtasks" + shutil.copytree(self.image_source, source) + info_path = source / "meta" / "info.json" + info = json.loads(info_path.read_text()) + info["features"]["subtask_index"] = { + "dtype": "int64", + "shape": [1], + "names": None, + } + info_path.write_text(json.dumps(info)) + + next_subtask = 0 + for path in sorted((source / "data").rglob("*.parquet")): + data = pq.read_table(path) + values = [ + (next_subtask + index) % 2 + for index in range(data.num_rows) + ] + next_subtask += data.num_rows + pq.write_table(data.append_column( + "subtask_index", + pa.array(values, type=pa.int64()), + ), path) + subtasks = pa.Table.from_pandas(pd.DataFrame( + {"subtask_index": [0, 1]}, + index=pd.Index(["reach", "grasp"], name="instruction"), + )) + pq.write_table(subtasks, source / "meta" / "subtasks.parquet") + + version_id = self.connection.load_from_lerobot( + "with_subtasks", source) + + frames = self.connection.get_table("with_subtasks") + self.assertNotIn("subtask", [ + field.name for field in frames.raw_table.fields + ]) + self.assertEqual([0, 1, 0, 1, 0], frames.scan().select([ + "index", "subtask_index" + ]).to_arrow().sort_by("index").column("subtask_index").to_pylist()) + subtasks_table = self.connection.catalog.get_table( + self.connection._identifier("with_subtasks__subtasks")) + self.assertTrue(_target_schema(subtasks_table).equals( + subtasks.schema, check_metadata=False)) + self.assertEqual( + subtasks.to_pylist(), + _catalog_rows(self.connection, "with_subtasks__subtasks"), + ) + _, subtasks_arrow = _catalog_arrow( + self.connection, "with_subtasks__subtasks") + pd.testing.assert_frame_equal( + subtasks.to_pandas(), + _restore_pandas_metadata( + subtasks_table, subtasks_arrow).to_pandas(), + ) + self.assertTrue(_catalog_rows( + self.connection, "with_subtasks__versions")[0]["has_subtasks"]) + self.assertEqual( + 1, + self.connection.catalog.get_tag( + self.connection._identifier("with_subtasks__subtasks"), + str(version_id), + ).snapshot.id, + ) + + def test_import_preserves_quoted_database_name(self): + self.connection.catalog.create_database( + "db.name", ignore_if_exists=False) + + self.connection.load_from_lerobot( + "`db.name`.robot", self.image_source) + + table_names = self.connection.catalog.list_tables("db.name") + self.assertEqual([ + "robot", + "robot__episodes", + "robot__tasks", + "robot__versions", + ], sorted(table_names)) + + def test_import_reuses_validated_episode_metadata(self): + from pypaimon.multimodal.lerobot import api + + source = self.temp_dir / "stable_episodes" + shutil.copytree(self.image_source, source) + episode_path = next((source / "meta" / "episodes").rglob("*.parquet")) + original_write = api._write_dataset + + def write_then_replace(*args, **kwargs): + snapshot_id = original_write(*args, **kwargs) + episodes = pq.read_table(episode_path) + tasks = episodes.column("tasks").to_pylist() + tasks[0] = ["place"] + pq.write_table(episodes.set_column( + episodes.schema.get_field_index("tasks"), + "tasks", + pa.array(tasks, type=episodes.schema.field("tasks").type), + ), episode_path) + return snapshot_id + + with patch.object( + api, "_write_dataset", side_effect=write_then_replace): + self.connection.load_from_lerobot("stable_episodes", source) + + published = _catalog_rows( + self.connection, "stable_episodes__episodes") + self.assertEqual(["pick"], published[0]["tasks"]) + + def test_frame_controls_must_match_published_episode_metadata(self): + cases = [ + ("index", 99), + ("episode_index", 0), + ("frame_index", 1), + ("timestamp", 0.2), + ("task_index", 0), + ] + for column, value in cases: + with self.subTest(column=column): + source = self.temp_dir / ("corrupt_" + column) + shutil.copytree(self.image_source, source) + path = next((source / "data").rglob("*.parquet")) + data = pq.read_table(path) + values = data.column(column).to_pylist() + values[2] = value + index = data.schema.get_field_index(column) + data = data.set_column( + index, + column, + pa.array(values, type=data.schema.field(index).type), + ) + pq.write_table(data, path) + + table_name = "corrupt_" + column + with self.assertRaisesRegex( + ValueError, "has %s" % column): + self.connection.load_from_lerobot(table_name, source) + self.connection.get_table(table_name) + self.assertEqual([], _catalog_rows( + self.connection, table_name + "__versions")) + + def test_task_text_remains_in_published_task_mapping(self): + source = self.temp_dir / "reordered_tasks" + shutil.copytree(self.image_source, source) + path = source / "meta" / "tasks.parquet" + tasks = pq.read_table(path) + pq.write_table(tasks.take(pa.array([1, 0])), path) + + self.connection.load_from_lerobot("reordered_tasks", source) + table = self.connection.get_table("reordered_tasks") + self.assertNotIn("task", [ + field.name for field in table.raw_table.fields + ]) + frames = table.scan().select([ + "index", "task_index" + ]).to_arrow().sort_by("index").to_pylist() + task_rows = _catalog_rows( + self.connection, "reordered_tasks__tasks") + task_name = ( + "task" if "task" in task_rows[0] else "__index_level_0__") + published = { + row["task_index"]: row[task_name] for row in task_rows + } + self.assertEqual({0: "pick", 1: "place"}, published) + self.assertTrue(all( + row["task_index"] in published for row in frames)) + + def test_episode_tasks_must_exactly_match_frame_tasks(self): + source = self.temp_dir / "extra_episode_task" + shutil.copytree(self.image_source, source) + path = next((source / "meta" / "episodes").rglob("*.parquet")) + episodes = pq.read_table(path) + tasks = episodes.column("tasks").to_pylist() + tasks[0] = ["pick", "place"] + index = episodes.schema.get_field_index("tasks") + episodes = episodes.set_column( + index, + "tasks", + pa.array(tasks, type=episodes.schema.field(index).type), + ) + pq.write_table(episodes, path) + + with self.assertRaisesRegex( + ValueError, "declares task indices"): + self.connection.load_from_lerobot( + "extra_episode_task", source) + self.connection.get_table("extra_episode_task") + self.assertEqual([], _catalog_rows( + self.connection, "extra_episode_task__versions")) + + def test_nonempty_dataset_cannot_publish_without_tasks(self): + source = self.temp_dir / "missing_tasks" + shutil.copytree(self.image_source, source) + info_path = source / "meta" / "info.json" + info = json.loads(info_path.read_text()) + info["total_tasks"] = 0 + info_path.write_text(json.dumps(info)) + episode_path = next( + (source / "meta" / "episodes").rglob("*.parquet")) + episodes = pq.read_table(episode_path) + index = episodes.schema.get_field_index("tasks") + episodes = episodes.set_column( + index, + "tasks", + pa.array( + [[] for _ in range(episodes.num_rows)], + type=episodes.schema.field(index).type, + ), + ) + pq.write_table(episodes, episode_path) + + with self.assertRaisesRegex(ValueError, "task_index"): + self.connection.load_from_lerobot("missing_tasks", source) + self.connection.get_table("missing_tasks") + self.assertEqual([], _catalog_rows( + self.connection, "missing_tasks__versions")) def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" @@ -687,16 +1356,16 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): with patch( "pypaimon.multimodal.lerobot.source._SourceFileIO", return_value=source_file_io): - snapshot_id = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "oss_images", source, batch_size=2, ) - self.assertEqual(1, snapshot_id) + self.assertEqual(1, version_id) table = self.connection.get_table("oss_images") rows = table.scan().select([ - "episode_index", "frame_index", "index", "task" + "episode_index", "frame_index", "index", "task_index" ]).to_arrow().sort_by("index").to_pylist() self.assertEqual([0, 0, 1, 1, 1], [ row["episode_index"] for row in rows @@ -704,11 +1373,9 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): 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( + self.assertEqual([0, 0, 1, 1, 1], [ + row["task_index"] for row in rows]) + self.assertTrue(any( path.endswith("meta/stats.json") for path in source_file_io.opened_paths )) @@ -717,75 +1384,284 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): if "/data/" in path and path.endswith(".parquet") ])) - def test_existing_incompatible_schema_fails_without_snapshot(self): + def test_remote_tasks_use_the_pandas_index_column(self): + local_source = self.temp_dir / "custom_task_index" + shutil.copytree(self.image_source, local_source) + path = local_source / "meta" / "tasks.parquet" + tasks = pq.read_table(path).to_pandas().rename_axis("instruction") + pq.write_table(pa.Table.from_pandas(tasks), path) + source = "oss://source-bucket/custom-task-index" + source_file_io = _RemoteLeRobotFileIO(local_source, source) + + with patch( + "pypaimon.multimodal.lerobot.source._SourceFileIO", + return_value=source_file_io): + self.connection.load_from_lerobot( + "custom_task_index", source) + + rows = _catalog_rows( + self.connection, "custom_task_index__tasks") + self.assertEqual( + [(0, "pick"), (1, "place")], + [(row["task_index"], row["instruction"]) for row in rows], + ) + + def test_empty_oss_source_does_not_require_episode_directory(self): + local_source = self.temp_dir / "empty_remote" + (local_source / "meta").mkdir(parents=True) + (local_source / "meta" / "info.json").write_text(json.dumps({ + "codebase_version": "v3.0", + "total_frames": 0, + "total_episodes": 0, + "total_tasks": 0, + "fps": 30, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + }, + })) + source = "oss://source-bucket/empty-robot" + source_file_io = _RemoteLeRobotFileIO(local_source, source) + + with patch( + "pypaimon.multimodal.lerobot.source._SourceFileIO", + return_value=source_file_io): + with self.assertRaisesRegex(ValueError, "non-empty"): + self.connection.load_from_lerobot("empty_oss", source) + + def test_tag_falls_back_for_catalogs_without_tag_api(self): + with patch.object( + self.connection.catalog, + "create_tag", + side_effect=NotImplementedError): + version_id = self.connection.load_from_lerobot( + "tag_fallback", self.image_source) + + manifest = _catalog_rows( + self.connection, "tag_fallback__versions")[0] + tag = str(manifest["version_id"]) + table = self.connection.get_table("tag_fallback") + self.assertEqual(1, version_id) + self.assertEqual( + table.raw_table.snapshot_manager().get_latest_snapshot().id, + table.raw_table.tag_manager().get(tag).id, + ) + + def test_tag_response_loss_is_reconciled(self): + create_tag = self.connection.catalog.create_tag + lost = [False] + + def create_then_lose_response(*args, **kwargs): + result = create_tag(*args, **kwargs) + if not lost[0]: + lost[0] = True + raise TimeoutError("lost tag response") + return result + + with patch.object( + self.connection.catalog, + "create_tag", + side_effect=create_then_lose_response): + version_id = self.connection.load_from_lerobot( + "tag_response_loss", self.image_source) + + self.assertTrue(lost[0]) + self.assertEqual(1, version_id) + self.assertEqual( + [1], + [row["version_id"] for row in _catalog_rows( + self.connection, "tag_response_loss__versions")]) + + def test_tag_failure_remains_unpublished(self): + with patch( + "pypaimon.multimodal.lerobot.metadata._create_tag", + side_effect=RuntimeError("tag failed")): + with self.assertRaisesRegex(RuntimeError, "tag failed"): + self.connection.load_from_lerobot( + "failed_publish", self.image_source) + + self.connection.get_table("failed_publish") + self.assertEqual([], _catalog_rows( + self.connection, "failed_publish__versions")) + + def test_existing_companion_is_rejected(self): + self.connection.load_from_lerobot( + "other_group", self.image_source) + stale = self.connection._identifier("stale__tasks") + self.connection.catalog.rename_table( + self.connection._identifier("other_group__tasks"), stale) + + with self.assertRaisesRegex(ValueError, "already exists"): + self.connection.load_from_lerobot( + "stale", self.image_source) + + self.connection.get_table("stale") + self.connection.catalog.get_table(stale) + + def test_invalid_target_options_do_not_leave_table(self): + with self.assertRaisesRegex(ValueError, "data-evolution.enabled"): + self.connection.load_from_lerobot( + "invalid_options", + self.image_source, + options={"data-evolution.enabled": "false"}, + ) + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier("invalid_options")) + + version_id = self.connection.load_from_lerobot( + "invalid_options", self.image_source) + self.assertEqual(1, version_id) + + def test_target_open_failure_leaves_created_table(self): + original_get = self.connection.get_table + failed = [False] + + def fail_once(name): + if not failed[0]: + failed[0] = True + raise RuntimeError("get failed") + return original_get(name) + + with patch.object( + self.connection, "get_table", side_effect=fail_once): + with self.assertRaisesRegex(RuntimeError, "get failed"): + self.connection.load_from_lerobot( + "failed_open", self.image_source) + + self.connection.get_table("failed_open") + + def test_dataset_close_failure_does_not_override_success(self): + from pypaimon.multimodal.lerobot import api + + original_open = api._open_resolved_dataset + + def open_with_failing_close(*args, **kwargs): + return _FailingCloseDataset(original_open(*args, **kwargs)) + + with self.assertLogs( + "pypaimon.multimodal.lerobot.source", level="WARNING"): + with patch.object( + api, + "_open_resolved_dataset", + side_effect=open_with_failing_close): + version_id = self.connection.load_from_lerobot( + "close_failure", self.image_source) + + self.assertEqual(1, version_id) + self.assertEqual( + [1], + [row["version_id"] for row in _catalog_rows( + self.connection, "close_failure__versions")]) + + def test_source_close_failure_does_not_override_success(self): + source = "oss://source-bucket/robot-images" + source_file_io = _RemoteLeRobotFileIO(self.image_source, source) + source_file_io.close = Mock(side_effect=RuntimeError("close failed")) + + with self.assertLogs( + "pypaimon.multimodal.lerobot.source", level="WARNING"): + with patch( + "pypaimon.multimodal.lerobot.source._SourceFileIO", + return_value=source_file_io): + version_id = self.connection.load_from_lerobot( + "source_close_failure", source) + + self.assertEqual(1, version_id) + self.assertEqual( + [1], + [row["version_id"] for row in _catalog_rows( + self.connection, "source_close_failure__versions")]) + + def test_existing_target_is_rejected(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=_replaced_contract( - schema.field("observation.matrix"), - "shape=[2,2]", - "shape=[5,2]", - ), - ), - "dtype": pa.field( - "action", - pa.list_(pa.float64(), 2), - nullable=False, - metadata=_replaced_contract( - schema.field("action"), - "dtype=float32", - "dtype=float64", - ), - ), - "names": pa.field( - "action", - schema.field("action").type, - nullable=False, - metadata=_replaced_contract( - schema.field("action"), - 'names=["x","y"]', - 'names=["y","x"]', - ), - ), - "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", - }, + schema = _schema_from_info(info) + table = self.connection.create_table("existing", schema=schema) + + with self.assertRaisesRegex(ValueError, "already exists"): + self.connection.load_from_lerobot( + "existing", self.image_source) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_concurrent_import_cannot_claim_the_same_target(self): + from pypaimon.multimodal.lerobot import api + + original_prepare = api._prepare_metadata_tables + root_created = threading.Event() + release = threading.Event() + + def prepare_then_wait(*args, **kwargs): + root_created.set() + release.wait(10) + return original_prepare(*args, **kwargs) + + with patch.object( + api, + "_prepare_metadata_tables", + side_effect=prepare_then_wait): + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + self.connection.load_from_lerobot, + "concurrent", + self.image_source, ) + try: + self.assertTrue(root_created.wait(10)) + with self.assertRaisesRegex( + ValueError, "already exists"): + self.connection.load_from_lerobot( + "concurrent", self.image_source) + finally: + release.set() + version_id = future.result(timeout=30) + + self.assertEqual(1, version_id) + self.assertEqual( + 5, + self.connection.get_table( + "concurrent").scan().to_arrow().num_rows, + ) + + def test_concurrent_append_cannot_enter_published_version(self): + from pypaimon.multimodal.lerobot import api + + original_write = api._write_dataset + + def append_then_write( + table, + dataset, + info, + source, + source_schema, + batch_size, + metadata): + table.add(_read_batch( + dataset, + info, + 0, + 1, + source_schema, + )) + return original_write( + table, + dataset, + info, + source, + source_schema, + batch_size, + metadata, + ) + + with patch.object( + api, "_write_dataset", side_effect=append_then_write): + with self.assertRaisesRegex(RuntimeError, "concurrent writes"): + self.connection.load_from_lerobot( + "concurrent_append", self.image_source) + + self.connection.get_table("concurrent_append") + self.assertEqual([], _catalog_rows( + self.connection, "concurrent_append__versions")) - 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()