From 286235b13335bfac2e8abaec44903128b55f030e Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 02:43:05 -0700 Subject: [PATCH 01/32] [python] Persist LeRobot dataset metadata in Paimon --- docs/docs/pypaimon/multimodal-api.mdx | 17 +- .../pypaimon/multimodal/connection.py | 45 +- .../pypaimon/multimodal/lerobot/api.py | 153 ++++- .../pypaimon/multimodal/lerobot/loader.py | 149 ++++- .../pypaimon/multimodal/lerobot/metadata.py | 620 ++++++++++++++++++ .../pypaimon/tests/multimodal_lerobot_test.py | 573 +++++++++++++++- 6 files changed, 1503 insertions(+), 54 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/lerobot/metadata.py diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index e8732875b176..c28b0b5009a9 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -524,7 +524,9 @@ source drift. Calling it again with the same input appends the rows again. `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 self-contained LeRobot dataset backed by the frame table, +`__datasets`, `
__episodes`, and `
__tasks`. A manifest row is +marked `READY` only after all component snapshots are committed and tagged. ```shell pip install 'pypaimon[lerobot]' @@ -554,13 +556,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. +All four tables share a `dataset_id` (the target identifier by default) and a +unique `metadata_version` per import. A `READY` manifest is appended only after +all component snapshots are committed and tagged. Set `dataset_id=` to choose +a stable ID and select `metadata_version` to read a specific imported version. +`drop_table()` removes all four tables. -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/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 3c55032b846a..2d1c97936b64 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -26,6 +26,7 @@ TableAlreadyExistException, TableNotExistException, ) +from pypaimon.common.identifier import Identifier from pypaimon.multimodal.table import MultimodalTable, _to_arrow_table _DEFAULT_OPTIONS = { @@ -123,6 +124,7 @@ def load_from_lerobot( source, *, batch_size: int = 1024, + dataset_id: Optional[str] = None, options=None, source_options=None): """Import LeRobot Dataset v3 and return the committed snapshot ID.""" @@ -132,13 +134,54 @@ def load_from_lerobot( table_name, source, batch_size=batch_size, + dataset_id=dataset_id, options=options, source_options=source_options, ) def drop_table(self, name: str, ignore_if_not_exists: bool = False): + identifier = self._identifier(name) + owner_id = None + try: + from pypaimon.multimodal.lerobot.metadata import ( + _DEFAULT_DATASET_ID_OPTION, + _OWNER_ID_OPTION, + ) + raw_table = self.catalog.get_table(identifier) + table_options = raw_table.table_schema.options + if _DEFAULT_DATASET_ID_OPTION in table_options: + if Identifier.from_string( + identifier).get_branch_name() is not None: + raise ValueError( + "Dropping a managed LeRobot table branch is not " + "supported; drop the branch through the Catalog.") + owner_id = table_options.get(_OWNER_ID_OPTION) + except (DatabaseNotExistException, TableNotExistException): + pass + + companions = [] + if owner_id is not None: + from pypaimon.multimodal.lerobot.metadata import \ + _companion_table_identifiers + for companion in _companion_table_identifiers( + raw_table).values(): + try: + table = self.catalog.get_table(companion) + except (DatabaseNotExistException, TableNotExistException): + continue + actual = table.table_schema.options.get(_OWNER_ID_OPTION) + if actual != owner_id: + raise ValueError( + "Refusing to drop %s because it belongs to a " + "different table." % companion) + companions.append(companion) + for companion in companions: + self.catalog.drop_table( + companion, + ignore_if_not_exists=True, + ) self.catalog.drop_table( - self._identifier(name), + identifier, ignore_if_not_exists=ignore_if_not_exists, ) diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 0b7f76d78a46..963534df08cf 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -26,6 +26,17 @@ DatabaseNotExistException, TableNotExistException, ) +from pypaimon.multimodal.lerobot.metadata import ( + _DEFAULT_DATASET_ID_OPTION, + _OWNER_ID_OPTION, + _frame_schema, + _load_dataset_metadata, + _managed_table_options, + _new_id, + _prepare_metadata_tables, + _publish_dataset, + _reject_subtasks, +) from pypaimon.multimodal.lerobot.loader import ( _strict_lerobot_table, _write_dataset, @@ -56,14 +67,16 @@ def load_from_lerobot( source, *, batch_size: int = 1024, + dataset_id: Optional[str] = None, 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. + A missing target table is created from LeRobot metadata. Episode, task, and + dataset metadata are stored in companion Paimon tables. ``dataset_id`` + defaults to the target identifier. FileIO URI + credentials come only from ``source_options`` and are not inherited from + the target Catalog. """ if sys.version_info < (3, 10): raise RuntimeError( @@ -85,18 +98,39 @@ def load_from_lerobot( _schema_from_info(local_info, include_task=False) total_frames, _, total_tasks = \ _validated_counts(local_info, resolved_source.path) - if total_frames == 0: + if total_frames == 0 and ( + resolved_source.root is not None + or resolved_source.file_io is not None): source_schema = _schema_from_info( local_info, include_task=total_tasks > 0, ) - _validated_table( + _reject_subtasks(None, resolved_source) + table, owner_id = _validated_table( connection, table_name, source_schema, options, resolved_source, ) + resolved_dataset_id = _resolved_dataset_id( + dataset_id, table) + metadata = _load_dataset_metadata( + None, local_info, resolved_source) + tables = _prepare_metadata_tables( + connection, table.raw_table, owner_id) + metadata_version = _new_id() + _publish_dataset( + connection, + tables, + resolved_dataset_id, + metadata_version, + local_info, + resolved_source, + metadata, + table.identifier, + None, + ) return None LeRobotDataset = _import_lerobot_dataset() dataset = _open_resolved_dataset( @@ -107,26 +141,60 @@ def load_from_lerobot( row_count, _, _ = \ _validated_counts(info, resolved_source.path) - source_schema = _schema_from_info( + lerobot_schema = _schema_from_info( info, include_task=_has_tasks(dataset, info)) - table = _validated_table( + _reject_subtasks(dataset, resolved_source) + table, owner_id = _validated_table( connection, table_name, - source_schema, + lerobot_schema, options, resolved_source, ) + resolved_dataset_id = _resolved_dataset_id( + dataset_id, table) + metadata = _load_dataset_metadata( + dataset, info, resolved_source) + tables = _prepare_metadata_tables( + connection, table.raw_table, owner_id) + metadata_version = _new_id() if row_count == 0: + _publish_dataset( + connection, + tables, + resolved_dataset_id, + metadata_version, + info, + resolved_source, + metadata, + table.identifier, + None, + ) return None - return _write_dataset( + snapshot_id = _write_dataset( table, dataset, info, resolved_source, - source_schema, + lerobot_schema, batch_size, + resolved_dataset_id, + metadata_version, + metadata, ) + _publish_dataset( + connection, + tables, + resolved_dataset_id, + metadata_version, + info, + resolved_source, + metadata, + table.identifier, + snapshot_id, + ) + return snapshot_id finally: close = getattr(dataset, "close", None) if callable(close): @@ -159,10 +227,57 @@ def _required_count(info, name, source): return int(value) +def _resolved_dataset_id(value, table): + if value is not None and ( + not isinstance(value, str) or not value.strip()): + raise ValueError("dataset_id must be a non-empty string.") + if value is None: + default_id = table.raw_table.table_schema.options.get( + _DEFAULT_DATASET_ID_OPTION) + if not default_id: + raise ValueError( + "Self-contained LeRobot table %s has no default dataset_id." + % table.identifier) + return default_id + return value.strip() + + def _validated_table( connection, table_name, source_schema, options, source): - table = _get_or_create_table( - connection, table_name, source_schema, options) + try: + table = connection.get_table(table_name) + except (DatabaseNotExistException, TableNotExistException): + owner_id = _new_id() + create_options = dict(options or {}) + managed_options = _managed_table_options( + connection._identifier(table_name), owner_id) + 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) + table = connection.create_table( + table_name, + schema=_frame_schema(source_schema), + options=create_options, + ) + _validate_target_schema(table, _frame_schema(source_schema), source) + return table, owner_id + + owner_id = table.raw_table.table_schema.options.get(_OWNER_ID_OPTION) + if owner_id is None: + raise ValueError( + "Existing LeRobot target %s is not managed by " + "load_from_lerobot; use a new target table." % table.identifier) + if table.raw_table.identifier.get_branch_name() is not None: + raise ValueError( + "Self-contained LeRobot import does not support table branches.") + _validate_target_schema(table, _frame_schema(source_schema), source) + return table, owner_id + + +def _validate_target_schema(table, source_schema, source): target_schema = _target_schema(table.raw_table) _validate_lerobot_schema( source_schema, target_schema, source.path) @@ -172,15 +287,3 @@ def _validated_table( source, 0, ) - return table - - -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 index 1f46e047d2b4..d635fa191981 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -25,6 +25,7 @@ from pypaimon.multimodal.arrow_utils import strict_arrow_table from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.lerobot.metadata import _with_frame_identity from pypaimon.multimodal.lerobot.schema import _feature_shape from pypaimon.multimodal.table import _target_schema @@ -67,7 +68,10 @@ def _write_dataset( info, source, source_schema, - batch_size): + batch_size, + dataset_id, + metadata_version, + metadata): target_schema = _target_schema(table.raw_table) write_builder = table.raw_table.new_batch_write_builder() table_write = None @@ -75,15 +79,33 @@ def _write_dataset( commit_started = False batch_count = 0 row_count = 0 + episodes = metadata["episodes"] + task_names = { + row["task_index"]: row["task"] for row in metadata["tasks"] + } + observed_tasks = {} 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): batch = _read_batch( - dataset, info, begin, end, source_schema) + dataset, info, begin, end, source_schema, task_names) + seen_tasks = _validate_frame_controls( + batch, + int(info["fps"]), + episode_index, + episode_begin, + begin, + task_indices, + ) + observed_tasks.setdefault(episode_index, set()).update( + seen_tasks) + batch = _with_frame_identity( + batch, dataset_id, metadata_version) batch = _strict_lerobot_table( batch, target_schema, @@ -94,6 +116,8 @@ def _write_dataset( batch_count += 1 row_count += batch.num_rows + _validate_episode_tasks(episodes, observed_tasks) + expected_rows = int(info.get("total_frames", len(dataset))) if row_count != expected_rows: raise ValueError( @@ -119,26 +143,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,7 +178,96 @@ def _episode_batches(dataset, info, batch_size): % (expected_begin, total_frames)) -def _read_batch(dataset, info, begin, end, schema): +def _validate_frame_controls( + batch, + fps, + episode_index, + episode_begin, + begin, + task_indices): + required = [ + "index", "episode_index", "frame_index", "timestamp", "task_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] + if (isinstance(timestamp, bool) + or not isinstance(timestamp, numbers.Real) + or not math.isclose( + float(timestamp), frame_index / fps, + rel_tol=0.0, abs_tol=1e-4)): + raise ValueError( + "LeRobot frame %d has timestamp %r; expected %r." + % (index, timestamp, frame_index / fps)) + 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) + return seen_tasks + + +def _validate_episode_tasks(episodes, observed_tasks): + for episode in episodes: + episode_index = episode["episode_index"] + expected = set(episode["task_indices"]) + actual = observed_tasks.get(episode_index, set()) + 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, task_names=None): read_batch = getattr(dataset, "read_batch", None) if callable(read_batch): raw = read_batch(begin, end) @@ -183,8 +303,9 @@ def _read_batch(dataset, info, begin, end, schema): if "task" in schema.names: task_indices = raw.column("task_index").to_pylist() + tasks = dataset.meta.tasks if task_names is None else task_names arrays.append(pa.array( - [_task_name(dataset.meta.tasks, value) for value in task_indices], + [_task_name(tasks, value) for value in task_indices], type=pa.string(), )) fields.append(schema.field("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..4de4f2355b0b --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -0,0 +1,620 @@ +# 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. + +"""Self-contained metadata tables for imported LeRobot datasets.""" + +import hashlib +import json +import numbers +import uuid +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 + + +_DATASET_ID = "dataset_id" +_METADATA_VERSION = "metadata_version" +_OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" +_DEFAULT_DATASET_ID_OPTION = "pypaimon.lerobot.dataset-id" +_TABLE_SUFFIXES = { + "datasets": "__datasets", + "episodes": "__episodes", + "tasks": "__tasks", +} +_COMPANION_OPTION_KEYS = { + name: "pypaimon.lerobot.%s-table" % name + for name in _TABLE_SUFFIXES +} + +_FRAME_ID_FIELDS = [ + pa.field(_DATASET_ID, pa.string(), nullable=False), + pa.field(_METADATA_VERSION, pa.string(), nullable=False), +] +_DATASETS_SCHEMA = pa.schema([ + pa.field(_DATASET_ID, pa.string(), nullable=False), + pa.field(_METADATA_VERSION, pa.string(), nullable=False), + pa.field("status", pa.string(), nullable=False), + pa.field("format", pa.string(), nullable=False), + pa.field("format_version", pa.string(), nullable=False), + pa.field("fps", pa.int64(), nullable=False), + pa.field("features_json", pa.string(), nullable=False), + pa.field("info_json", pa.string(), nullable=False), + pa.field("global_stats_json", pa.string()), + pa.field("total_frames", pa.int64(), nullable=False), + pa.field("total_episodes", pa.int64(), nullable=False), + pa.field("total_tasks", pa.int64(), nullable=False), + pa.field("frames_snapshot_id", pa.int64()), + pa.field("episodes_snapshot_id", pa.int64()), + pa.field("tasks_snapshot_id", pa.int64()), + pa.field("source_uri", pa.string(), nullable=False), + pa.field("metadata_checksum", pa.string(), nullable=False), +]) +_EPISODES_SCHEMA = pa.schema([ + pa.field(_DATASET_ID, pa.string(), nullable=False), + pa.field(_METADATA_VERSION, pa.string(), nullable=False), + 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("length", pa.int64(), nullable=False), + pa.field("task_indices", pa.list_(pa.int64()), nullable=False), + pa.field("split", pa.string()), + pa.field("episode_stats_json", pa.string()), + pa.field("episode_metadata_json", pa.string(), nullable=False), +]) +_TASKS_SCHEMA = pa.schema([ + pa.field(_DATASET_ID, pa.string(), nullable=False), + pa.field(_METADATA_VERSION, pa.string(), nullable=False), + pa.field("task_index", pa.int64(), nullable=False), + pa.field("task", pa.string(), nullable=False), + pa.field("task_metadata_json", pa.string()), +]) + + +def _frame_schema(source_schema): + reserved = [ + field.name for field in _FRAME_ID_FIELDS + if field.name in source_schema.names + ] + if reserved: + raise ValueError( + "LeRobot features use reserved Paimon fields: %s" % reserved) + return pa.schema(list(source_schema) + _FRAME_ID_FIELDS) + + +def _with_frame_identity(table, dataset_id, metadata_version): + size = table.num_rows + return pa.Table.from_arrays( + list(table.columns) + [ + pa.array([dataset_id] * size, type=pa.string()), + pa.array([metadata_version] * size, type=pa.string()), + ], + schema=pa.schema(list(table.schema) + _FRAME_ID_FIELDS), + ) + + +def _load_dataset_metadata(dataset, info, source): + fps = _positive_integer(info.get("fps"), "fps") + stats = _source_stats(dataset, source) + task_records = _source_tasks(dataset, source, int(info["total_tasks"])) + tasks, task_indices = _task_rows(task_records, int(info["total_tasks"])) + total_episodes = int(info["total_episodes"]) + episode_records = [] if total_episodes == 0 \ + else _source_episodes(dataset, source) + episodes = _episode_rows( + episode_records, + task_indices, + info, + int(info["total_frames"]), + total_episodes, + ) + canonical = { + "info": info, + "stats": stats, + "episodes": episodes, + "tasks": tasks, + } + return { + "fps": fps, + "features_json": _canonical_json(info["features"]), + "info_json": _canonical_json(info), + "global_stats_json": ( + None if stats is None else _canonical_json(stats)), + "episodes": episodes, + "tasks": tasks, + "metadata_checksum": "sha256:" + hashlib.sha256( + _canonical_json(canonical).encode("utf-8") + ).hexdigest(), + } + + +def _new_id(): + return uuid.uuid4().hex + + +def _companion_identifier(frames_identifier, suffix): + identifier = Identifier.from_string(str(frames_identifier)) + if identifier.is_system_table(): + raise ValueError( + "LeRobot target cannot be a Paimon system table: %s" + % frames_identifier) + return Identifier( + identifier.get_database_name(), + identifier.get_table_name() + suffix, + branch=identifier.get_branch_name(), + ).get_full_name() + + +def _managed_table_options(frames_identifier, owner_id): + identifier = Identifier.from_string(str(frames_identifier)) + if identifier.get_branch_name() is not None: + raise ValueError( + "Self-contained LeRobot import does not support table branches.") + result = { + _OWNER_ID_OPTION: owner_id, + _DEFAULT_DATASET_ID_OPTION: str(frames_identifier), + } + 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, owner_id): + schemas = { + "datasets": _DATASETS_SCHEMA, + "episodes": _EPISODES_SCHEMA, + "tasks": _TASKS_SCHEMA, + } + identifiers = _companion_table_identifiers(frames_table) + tables = {} + for name, schema in schemas.items(): + identifier = identifiers[name] + try: + table = connection.catalog.get_table(identifier) + except (DatabaseNotExistException, TableNotExistException): + paimon_schema = PaimonSchema.from_pyarrow_schema( + schema, + options={ + "bucket": "-1", + _OWNER_ID_OPTION: owner_id, + }, + ) + try: + connection.catalog.create_table( + identifier, paimon_schema, False) + except TableAlreadyExistException: + pass + table = connection.catalog.get_table(identifier) + if table.table_schema.primary_keys: + raise ValueError( + "LeRobot metadata table %s must be append-only." % identifier) + actual = _target_schema(table) + if not actual.equals(schema, check_metadata=False): + raise ValueError( + "LeRobot metadata table %s has schema %s; expected %s." + % (identifier, actual, schema)) + actual_owner_id = table.table_schema.options.get(_OWNER_ID_OPTION) + if actual_owner_id != owner_id: + raise ValueError( + "LeRobot metadata table %s belongs to a different target " + "table. Drop the stale companion tables before importing." + % identifier) + tables[name] = table + return tables + + +def _publish_dataset( + connection, + tables, + dataset_id, + metadata_version, + info, + source, + metadata, + frames_identifier, + frames_snapshot_id): + episodes = _versioned_table( + metadata["episodes"], + _EPISODES_SCHEMA, + dataset_id, + metadata_version, + ) + tasks = _versioned_table( + metadata["tasks"], + _TASKS_SCHEMA, + dataset_id, + metadata_version, + ) + episodes_snapshot_id = _append_arrow(tables["episodes"], episodes) + tasks_snapshot_id = _append_arrow(tables["tasks"], tasks) + + tag = "pypaimon-lerobot-%s" % metadata_version + for identifier, snapshot_id in ( + (frames_identifier, frames_snapshot_id), + (tables["episodes"].identifier, episodes_snapshot_id), + (tables["tasks"].identifier, tasks_snapshot_id)): + if snapshot_id is not None: + _create_tag(connection.catalog, identifier, tag, snapshot_id) + + manifest = _manifest_row( + dataset_id, + metadata_version, + info, + source, + metadata, + frames_snapshot_id, + episodes_snapshot_id, + tasks_snapshot_id, + ) + _append_arrow(tables["datasets"], pa.Table.from_pylist( + [manifest], schema=_DATASETS_SCHEMA)) + + +def _manifest_row( + dataset_id, + metadata_version, + info, + source, + metadata, + frames_snapshot_id, + episodes_snapshot_id, + tasks_snapshot_id): + return { + _DATASET_ID: dataset_id, + _METADATA_VERSION: metadata_version, + "status": "READY", + "format": "lerobot", + "format_version": str(info["codebase_version"]), + "fps": metadata["fps"], + "features_json": metadata["features_json"], + "info_json": metadata["info_json"], + "global_stats_json": metadata["global_stats_json"], + "total_frames": int(info["total_frames"]), + "total_episodes": int(info["total_episodes"]), + "total_tasks": int(info["total_tasks"]), + "frames_snapshot_id": frames_snapshot_id, + "episodes_snapshot_id": episodes_snapshot_id, + "tasks_snapshot_id": tasks_snapshot_id, + "source_uri": str(source.path), + "metadata_checksum": metadata["metadata_checksum"], + } + + +def _versioned_table(rows, schema, dataset_id, metadata_version): + values = [] + for row in rows: + value = dict(row) + value[_DATASET_ID] = dataset_id + value[_METADATA_VERSION] = metadata_version + values.append(value) + return pa.Table.from_pylist(values, schema=schema) + + +def _append_arrow(table, data): + if data.num_rows == 0: + return None + builder = table.new_batch_write_builder() + table_write = builder.new_write() + table_commit = builder.new_commit() + commit_started = False + recorder = _SnapshotRecorder() + table_commit.add_commit_callback(recorder) + try: + table_write.write_arrow(data) + 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 not commit_started: + table_write.abort() + raise + finally: + try: + table_write.close() + finally: + table_commit.close() + + +def _create_tag(catalog, identifier, tag_name, snapshot_id): + 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) + + +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 [] + 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).to_pylist() + path = _metadata_root(dataset, source) / "meta" / "tasks.parquet" + try: + return pq.read_table(path).to_pylist() + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot task 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, + _remote_parquet_files, + _remote_path, + ) + directory = _remote_path(source.path, "meta/episodes") + paths = _remote_parquet_files(source.file_io, directory) + tables = [ + _read_remote_parquet(source.file_io, path) for path in paths + ] + else: + directory = _metadata_root(dataset, source) / "meta" / "episodes" + paths = sorted(directory.rglob("*.parquet")) + try: + tables = [pq.read_table(path) for path in paths] + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot Episode metadata %s: %s" + % (directory, error)) from error + rows = [] + for table in tables: + rows.extend(table.to_pylist()) + rows.sort(key=lambda row: _integer( + row.get("episode_index"), "episode_index")) + return rows + + +def _reject_subtasks(dataset, source): + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import _remote_path + path = _remote_path(source.path, "meta/subtasks.parquet") + try: + source.file_io.get_file_status(path) + except FileNotFoundError: + return + else: + path = _metadata_root(dataset, source) / "meta" / "subtasks.parquet" + if not path.is_file(): + return + raise ValueError( + "LeRobot subtask metadata is not supported yet: %s" % path) + + +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_rows(records, total_tasks): + rows = [None] * total_tasks + by_name = {} + for record in records: + index = _integer(record.get("task_index"), "task_index") + task = record.get("task", record.get("name")) + if task is None: + task = record.get("__index_level_0__") + if index < 0 or index >= total_tasks or task is None \ + or rows[index] is not None: + raise ValueError("LeRobot task metadata is invalid: %s" % record) + task = str(task) + if task in by_name: + raise ValueError("LeRobot task metadata repeats task %r." % task) + by_name[task] = index + extra = dict(record) + for key in ("task_index", "task", "name", "__index_level_0__"): + extra.pop(key, None) + rows[index] = { + "task_index": index, + "task": task, + "task_metadata_json": ( + _canonical_json(extra) if extra else None), + } + if any(row is None for row in rows): + raise ValueError( + "LeRobot task metadata does not cover [0, %d)." % total_tasks) + return rows, by_name + + +def _episode_rows(records, task_indices, info, total_frames, total_episodes): + if len(records) != total_episodes: + raise ValueError( + "LeRobot metadata reports %d Episodes but %d were found." + % (total_episodes, len(records))) + splits = _episode_splits(info.get("splits"), total_episodes) + rows = [] + expected_begin = 0 + for ordinal, record in enumerate(records): + index = _integer(record.get("episode_index"), "episode_index") + begin = _integer( + record.get("dataset_from_index"), "dataset_from_index") + end = _integer(record.get("dataset_to_index"), "dataset_to_index") + length = _integer(record.get("length"), "length") + if index != ordinal or begin != expected_begin or end <= begin \ + or length != end - begin: + raise ValueError( + "LeRobot Episode %d has inconsistent index, range, or length." + % ordinal) + names = record.get("tasks", []) + if isinstance(names, str): + names = [names] + if task_indices and not names: + raise ValueError( + "LeRobot Episode %d does not declare any task." % ordinal) + try: + episode_task_indices = [task_indices[str(name)] for name in names] + except (KeyError, TypeError) as error: + raise ValueError( + "LeRobot Episode %d refers to an unknown task." % ordinal + ) from error + if len(set(episode_task_indices)) != len(episode_task_indices): + raise ValueError( + "LeRobot Episode %d repeats a task." % ordinal) + + stats = { + key[len("stats/"):]: value + for key, value in record.items() if key.startswith("stats/") + } + extra = { + key: value for key, value in record.items() + if key not in { + "episode_index", "dataset_from_index", "dataset_to_index", + "length", "tasks" + } and not key.startswith("stats/") + } + rows.append({ + "episode_index": index, + "dataset_from_index": begin, + "dataset_to_index": end, + "length": length, + "task_indices": episode_task_indices, + "split": splits[index], + "episode_stats_json": ( + _canonical_json(stats) if stats else None), + "episode_metadata_json": _canonical_json(extra), + }) + expected_begin = end + if expected_begin != total_frames: + raise ValueError( + "LeRobot Episode ranges cover %d frames but metadata reports %d." + % (expected_begin, total_frames)) + return rows + + +def _episode_splits(value, total_episodes): + result = [None] * total_episodes + if not isinstance(value, dict): + return result + for name, bounds in value.items(): + if not isinstance(bounds, str) or bounds.count(":") != 1: + continue + begin_text, end_text = bounds.split(":") + try: + begin, end = int(begin_text), int(end_text) + except ValueError: + continue + if begin < 0 or end < begin or end > total_episodes: + continue + for index in range(begin, end): + if result[index] is None: + result[index] = str(name) + elif result[index] != str(name): + result[index] = None + return result + + +def _canonical_json(value): + return json.dumps( + _json_value(value), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + +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/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 7f1f8d10a8ce..2fd8b3aa4569 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -26,11 +26,20 @@ import numpy as np import pyarrow as pa import pyarrow.fs as pafs +import pyarrow.parquet as pq +from pypaimon import Schema as PaimonSchema +from pypaimon.catalog.catalog_exception import TableNotExistException 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.metadata import ( + _DATASETS_SCHEMA, + _OWNER_ID_OPTION, + _frame_schema, + _managed_table_options, +) from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, @@ -64,8 +73,19 @@ 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() + + 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", "owner") + def test_dataset_open_never_downloads_videos(self): calls = [] @@ -341,6 +361,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]}, }, @@ -357,6 +378,13 @@ def test_empty_local_dataset_returns_before_opening_lerobot(self): table = connection.get_table("empty_frames") self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) + manifests = _catalog_rows( + connection, "empty_frames__datasets") + self.assertEqual(["READY"], [ + row["status"] for row in manifests + ]) + self.assertEqual(0, manifests[0]["total_frames"]) + self.assertIsNone(manifests[0]["frames_snapshot_id"]) finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -371,6 +399,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 +432,104 @@ def test_empty_fast_path_validates_required_counts(self): finally: shutil.rmtree(temp_dir, ignore_errors=True) + def test_empty_dataset_preserves_tasks_and_stats(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"), + }) + + self.assertIsNone(connection.load_from_lerobot("frames", source)) + manifest = _catalog_rows(connection, "frames__datasets")[0] + self.assertIsNotNone(manifest["global_stats_json"]) + self.assertEqual(1, manifest["total_tasks"]) + self.assertEqual(1, manifest["tasks_snapshot_id"]) + self.assertEqual( + "pick", _catalog_rows(connection, "frames__tasks")[0]["task"]) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + def test_subtasks_are_rejected_before_any_snapshot(self): + 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]}, + }, + })) + pq.write_table(pa.table({ + "subtask_index": [0], + "subtask": ["reach"], + }), source / "meta" / "subtasks.parquet") + connection = pmm.connect(options={ + "warehouse": str(temp_dir / "warehouse"), + }) + + with self.assertRaisesRegex(ValueError, "subtask metadata"): + connection.load_from_lerobot("frames", source) + with self.assertRaises(TableNotExistException): + connection.get_table("frames") + 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) + table = connection.get_table("frames") + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + with self.assertRaises(TableNotExistException): + connection.catalog.get_table( + connection._identifier("frames__datasets")) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + def test_source_values_are_safely_converted(self): class Dataset: @@ -632,6 +759,8 @@ 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.assertEqual("STRING NOT NULL", types["dataset_id"]) + self.assertEqual("STRING NOT NULL", types["metadata_version"]) rows = table.scan().select([ "episode_index", @@ -644,6 +773,8 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): "observation.matrix", "action", "reward", + "dataset_id", + "metadata_version", ]).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]) @@ -655,6 +786,58 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): rows[4]["observation.matrix"]) self.assertAlmostEqual(0.2, rows[4]["timestamp"], places=6) self.assertEqual(1.0, rows[4]["reward"]) + self.assertEqual( + [str(table.identifier)] * 5, + [row["dataset_id"] for row in rows], + ) + manifests = _catalog_rows(self.connection, "robot_data__datasets") + self.assertEqual(["READY"], [ + row["status"] for row in manifests + ]) + manifest = manifests[0] + self.assertEqual(str(table.identifier), manifest["dataset_id"]) + self.assertEqual(32, len(manifest["metadata_version"])) + self.assertEqual(snapshot_id, manifest["frames_snapshot_id"]) + self.assertEqual(1, manifest["episodes_snapshot_id"]) + self.assertEqual(1, manifest["tasks_snapshot_id"]) + self.assertEqual("lerobot", manifest["format"]) + self.assertEqual("v3.0", manifest["format_version"]) + self.assertIsNotNone(manifest["global_stats_json"]) + self.assertTrue(manifest["metadata_checksum"].startswith("sha256:")) + self.assertEqual( + [manifest["metadata_version"]] * 5, + [row["metadata_version"] for row in rows], + ) + tag = "pypaimon-lerobot-%s" % manifest["metadata_version"] + self.assertEqual( + snapshot_id, + 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") + 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([[0], [1]], [row["task_indices"] for row in episodes]) + self.assertEqual(["train", "train"], [row["split"] for row in episodes]) + self.assertTrue(all( + row["episode_stats_json"] is not None for row in episodes)) + + tasks = _catalog_rows(self.connection, "robot_data__tasks") + self.assertEqual([(0, "pick"), (1, "place")], [ + (row["task_index"], row["task"]) for row in tasks + ]) self.assertEqual( snapshot_id, table.raw_table.snapshot_manager().get_latest_snapshot().id, @@ -679,6 +862,147 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): "robot_data", self.image_source, batch_size=4) self.assertEqual(2, appended_snapshot_id) self.assertEqual(10, table.scan().to_arrow().num_rows) + manifests = _catalog_rows(self.connection, "robot_data__datasets") + ready_versions = [ + row["metadata_version"] for row in manifests + if row["status"] == "READY" + ] + self.assertEqual(2, len(ready_versions)) + self.assertEqual(2, len(set(ready_versions))) + versions = table.scan().select([ + "metadata_version" + ]).to_arrow().column(0).to_pylist() + self.assertEqual([5, 5], sorted( + versions.count(version) for version in set(versions))) + + 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) + datasets = self.connection.catalog.get_table( + self.connection._identifier( + table_name + "__datasets")) + self.assertIsNone( + datasets.snapshot_manager().get_latest_snapshot()) + table = self.connection.get_table(table_name) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_explicit_dataset_id_is_shared_by_all_tables(self): + dataset_id = "aloha_pick_cube@3" + self.connection.load_from_lerobot( + "custom_id", self.image_source, dataset_id=dataset_id) + + for name in ( + "custom_id", "custom_id__datasets", + "custom_id__episodes", "custom_id__tasks"): + rows = _catalog_rows(self.connection, name) + self.assertEqual({dataset_id}, { + row["dataset_id"] for row in rows + }) + + def test_frame_task_uses_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) + frames = self.connection.get_table("reordered_tasks").scan().select([ + "index", "task_index", "task" + ]).to_arrow().sort_by("index").to_pylist() + published = { + row["task_index"]: row["task"] + for row in _catalog_rows( + self.connection, "reordered_tasks__tasks") + } + self.assertTrue(all( + row["task"] == published[row["task_index"]] + 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) + datasets = self.connection.catalog.get_table( + self.connection._identifier("extra_episode_task__datasets")) + self.assertIsNone( + datasets.snapshot_manager().get_latest_snapshot()) + table = self.connection.get_table("extra_episode_task") + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + 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) + datasets = self.connection.catalog.get_table( + self.connection._identifier("missing_tasks__datasets")) + self.assertIsNone( + datasets.snapshot_manager().get_latest_snapshot()) + table = self.connection.get_table("missing_tasks") + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" @@ -708,7 +1032,7 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): ["pick", "pick", "place", "place", "place"], [row["task"] for row in rows], ) - self.assertFalse(any( + self.assertTrue(any( path.endswith("meta/stats.json") for path in source_file_io.opened_paths )) @@ -717,6 +1041,236 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): if "/data/" in path and path.endswith(".parquet") ])) + 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._Hdf5SourceFileIO", + return_value=source_file_io): + self.assertIsNone(self.connection.load_from_lerobot( + "empty_oss", source)) + self.assertEqual( + 0, + _catalog_rows( + self.connection, "empty_oss__datasets")[0]["total_episodes"], + ) + + def test_tag_falls_back_for_catalogs_without_tag_api(self): + with patch.object( + self.connection.catalog, + "create_tag", + side_effect=NotImplementedError): + snapshot_id = self.connection.load_from_lerobot( + "tag_fallback", self.image_source) + + manifest = _catalog_rows( + self.connection, "tag_fallback__datasets")[0] + tag = "pypaimon-lerobot-%s" % manifest["metadata_version"] + table = self.connection.get_table("tag_fallback") + self.assertEqual( + snapshot_id, + table.raw_table.tag_manager().get(tag).id, + ) + + def test_failed_publication_leaves_no_manifest(self): + with patch( + "pypaimon.multimodal.lerobot.api._publish_dataset", + side_effect=RuntimeError("publish failed")): + with self.assertRaisesRegex(RuntimeError, "publish failed"): + self.connection.load_from_lerobot( + "failed_publish", self.image_source) + + datasets = self.connection.catalog.get_table( + self.connection._identifier("failed_publish__datasets")) + self.assertIsNone( + datasets.snapshot_manager().get_latest_snapshot()) + frames = self.connection.get_table("failed_publish").scan().select([ + "metadata_version" + ]).to_arrow() + self.assertEqual(5, frames.num_rows) + versions = set(frames.column(0).to_pylist()) + self.assertEqual(1, len(versions)) + self.assertEqual(32, len(next(iter(versions)))) + + def test_existing_unmanaged_table_is_rejected(self): + info = json.loads((self.image_source / "meta" / "info.json").read_text()) + schema = _schema_from_info(info, include_task=True) + table = self.connection.create_table("unmanaged", schema=schema) + + with self.assertRaisesRegex(ValueError, "is not managed"): + self.connection.load_from_lerobot( + "unmanaged", self.image_source) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_companion_tables_must_be_append_only(self): + info = json.loads((self.image_source / "meta" / "info.json").read_text()) + source_schema = _schema_from_info(info, include_task=True) + owner_id = "test-owner" + table = self.connection.create_table( + "invalid_group", + schema=_frame_schema(source_schema), + options=_managed_table_options( + self.connection._identifier("invalid_group"), owner_id), + ) + identifier = self.connection._identifier("invalid_group__datasets") + self.connection.catalog.create_table( + identifier, + PaimonSchema.from_pyarrow_schema( + _DATASETS_SCHEMA, + primary_keys=["metadata_version"], + options={ + "bucket": "1", + _OWNER_ID_OPTION: owner_id, + }, + ), + False, + ) + + with self.assertRaisesRegex(ValueError, "must be append-only"): + self.connection.load_from_lerobot( + "invalid_group", self.image_source) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + def test_stale_companion_tables_are_rejected(self): + info = json.loads((self.image_source / "meta" / "info.json").read_text()) + source_schema = _schema_from_info(info, include_task=True) + table = self.connection.create_table( + "stale_group", + schema=_frame_schema(source_schema), + options=_managed_table_options( + self.connection._identifier("stale_group"), "new-owner"), + ) + identifier = self.connection._identifier("stale_group__datasets") + self.connection.catalog.create_table( + identifier, + PaimonSchema.from_pyarrow_schema( + _DATASETS_SCHEMA, + options={ + "bucket": "-1", + _OWNER_ID_OPTION: "old-owner", + }, + ), + False, + ) + + with self.assertRaisesRegex(ValueError, "different target table"): + self.connection.load_from_lerobot( + "stale_group", self.image_source) + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + with self.assertRaisesRegex(ValueError, "Refusing to drop"): + self.connection.drop_table("stale_group") + self.connection.get_table("stale_group") + self.connection.catalog.get_table(identifier) + + def test_drop_table_removes_companion_tables(self): + self.connection.load_from_lerobot("drop_group", self.image_source) + self.connection.drop_table("drop_group") + + for name in ( + "drop_group", "drop_group__datasets", + "drop_group__episodes", "drop_group__tasks"): + with self.subTest(name=name): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(name)) + + def test_drop_table_can_retry_after_companion_failure(self): + self.connection.load_from_lerobot( + "retry_drop", self.image_source) + original_drop = self.connection.catalog.drop_table + failed = [False] + + def flaky_drop(identifier, ignore_if_not_exists=False): + if str(identifier).endswith("__episodes") and not failed[0]: + failed[0] = True + raise RuntimeError("injected drop failure") + return original_drop(identifier, ignore_if_not_exists) + + with patch.object( + self.connection.catalog, + "drop_table", + side_effect=flaky_drop): + with self.assertRaisesRegex(RuntimeError, "injected"): + self.connection.drop_table("retry_drop") + self.connection.get_table("retry_drop") + + self.connection.drop_table("retry_drop") + for name in ( + "retry_drop", "retry_drop__datasets", + "retry_drop__episodes", "retry_drop__tasks"): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(name)) + + def test_companion_table_can_be_dropped_directly(self): + self.connection.load_from_lerobot( + "direct_drop", self.image_source) + self.connection.drop_table("direct_drop__tasks") + + self.connection.get_table("direct_drop") + with self.assertRaises(TableNotExistException): + self.connection.get_table("direct_drop__tasks") + self.connection.drop_table("direct_drop") + + def test_drop_table_rejects_managed_branch(self): + self.connection.load_from_lerobot( + "branch_drop", self.image_source) + self.connection.catalog.create_branch( + self.connection._identifier("branch_drop"), "dev") + + with self.assertRaisesRegex(ValueError, "table branch"): + self.connection.drop_table("branch_drop$branch_dev") + for name in ( + "branch_drop", "branch_drop__datasets", + "branch_drop__episodes", "branch_drop__tasks"): + self.connection.catalog.get_table( + self.connection._identifier(name)) + + def test_table_group_survives_frame_table_rename(self): + self.connection.load_from_lerobot("before_rename", self.image_source) + self.connection.catalog.rename_table( + self.connection._identifier("before_rename"), + self.connection._identifier("after_rename"), + ) + + self.assertEqual(2, self.connection.load_from_lerobot( + "after_rename", self.image_source)) + manifests = _catalog_rows( + self.connection, "before_rename__datasets") + ready = [row for row in manifests if row["status"] == "READY"] + self.assertEqual(2, len(ready)) + self.assertEqual(2, len({ + row["metadata_version"] for row in ready + })) + self.assertEqual( + {self.connection._identifier("before_rename")}, + {row["dataset_id"] for row in ready}, + ) + + self.connection.drop_table("after_rename") + for name in ( + "after_rename", "before_rename__datasets", + "before_rename__episodes", "before_rename__tasks"): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(name)) + 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) @@ -767,17 +1321,22 @@ def test_existing_incompatible_schema_fails_without_snapshot(self): for name, replacement in incompatible_fields.items(): with self.subTest(name=name): table_name = "incompatible_%s" % name + options = { + "file.format": "parquet", + "vector.file.format": "parquet", + } + options.update(_managed_table_options( + self.connection._identifier(table_name), + "owner-%s" % name, + )) table = self.connection.create_table( table_name, - schema=pa.schema([ + schema=_frame_schema(pa.schema([ replacement if field.name == replacement.name else field for field in schema - ]), - options={ - "file.format": "parquet", - "vector.file.format": "parquet", - }, + ])), + options=options, ) with self.assertRaisesRegex( From ebf91408e530a538ba714a7bdd7eaeb7ba5ebc22 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 09:12:41 -0700 Subject: [PATCH 02/32] [python] Publish LeRobot imports by component snapshots --- docs/docs/pypaimon/multimodal-api.mdx | 25 +-- paimon-python/README.md | 4 +- paimon-python/pypaimon/multimodal/__init__.py | 2 + .../pypaimon/multimodal/connection.py | 2 +- .../pypaimon/multimodal/lerobot/__init__.py | 6 +- .../pypaimon/multimodal/lerobot/api.py | 83 +++++++-- .../pypaimon/multimodal/lerobot/loader.py | 4 +- .../pypaimon/multimodal/lerobot/metadata.py | 88 ++++++--- .../pypaimon/tests/multimodal_lerobot_test.py | 173 ++++++++++-------- 9 files changed, 254 insertions(+), 133 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index c28b0b5009a9..96e8ce11c3c7 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -524,28 +524,31 @@ source drift. Calling it again with the same input appends the rows again. `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 creates a self-contained LeRobot dataset backed by the frame table, +frame, and creates a LeRobot dataset backed by the frame table, `
__datasets`, `
__episodes`, and `
__tasks`. A manifest row is -marked `READY` only after all component snapshots are committed and tagged. +reserved as `PENDING` before the component writes and marked `READY` only after +all component snapshots are committed and tagged. Readers ignore `PENDING`. ```shell pip install 'pypaimon[lerobot]' ``` ```python -snapshot_id = conn.load_from_lerobot( +result = conn.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(snapshot_id) +print(result.version_id) +print(result.frames_snapshot_id) ``` -The return value is `None` when the source has no frames. +The result contains the `dataset_id`, `version_id`, and the frame, Episode, and +Task snapshot IDs. Empty components have no snapshot. For FileIO URIs, pass credentials through `source_options`: ```python -snapshot_id = conn.load_from_lerobot( +result = conn.load_from_lerobot( "robot_data", "oss://source-bucket/lerobot_dataset", source_options={ @@ -556,11 +559,11 @@ snapshot_id = conn.load_from_lerobot( ) ``` -All four tables share a `dataset_id` (the target identifier by default) and a -unique `metadata_version` per import. A `READY` manifest is appended only after -all component snapshots are committed and tagged. Set `dataset_id=` to choose -a stable ID and select `metadata_version` to read a specific imported version. -`drop_table()` removes all four tables. +Component rows carry `dataset_id` and stable LeRobot indices. `version_id` +exists only in the Dataset manifest, where it resolves the exact frame, +Episode, and Task snapshots. These snapshots define the published Dataset +state; tags retain them. Set `dataset_id=` to choose a stable ID. The one-time +importer accepts each `dataset_id` once. `drop_table()` removes all four tables. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. diff --git a/paimon-python/README.md b/paimon-python/README.md index a008dbe2608d..19b2b5009742 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -44,11 +44,11 @@ pip install 'pypaimon[lerobot]' import pypaimon.multimodal as pmm connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) -snapshot_id = connection.load_from_lerobot( +result = connection.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(snapshot_id) +print(result.version_id, result.frames_snapshot_id) ``` The schema comes from `meta/info.json`. Each frame becomes one row; media uses diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index b669267b41da..287b672a0231 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -29,6 +29,7 @@ Hdf5File, Hdf5LoadResult, ) +from pypaimon.multimodal.lerobot import LeRobotLoadResult from pypaimon.multimodal.table import ( MultimodalTable, TextRoute, @@ -51,6 +52,7 @@ "BlobStore", "Hdf5File", "Hdf5LoadResult", + "LeRobotLoadResult", "MultimodalConnection", "MultimodalTable", "NoSuchKey", diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 2d1c97936b64..b78bd1d823ed 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -127,7 +127,7 @@ def load_from_lerobot( dataset_id: Optional[str] = None, options=None, source_options=None): - """Import LeRobot Dataset v3 and return the committed snapshot ID.""" + """Import LeRobot Dataset v3 and return its published table state.""" from pypaimon.multimodal.lerobot import load_from_lerobot return load_from_lerobot( self, diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index a40f2a8ccef0..d00c91307057 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -16,9 +16,13 @@ """One-time LeRobot Dataset v3 import into a multimodal Paimon table.""" -from pypaimon.multimodal.lerobot.api import load_from_lerobot +from pypaimon.multimodal.lerobot.api import ( + LeRobotLoadResult, + load_from_lerobot, +) __all__ = [ + "LeRobotLoadResult", "load_from_lerobot", ] diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 963534df08cf..8d0c4713fd60 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -18,6 +18,7 @@ import numbers import sys +from dataclasses import dataclass from typing import Mapping, Optional import pyarrow as pa @@ -36,6 +37,7 @@ _prepare_metadata_tables, _publish_dataset, _reject_subtasks, + _reserve_dataset_version, ) from pypaimon.multimodal.lerobot.loader import ( _strict_lerobot_table, @@ -61,6 +63,17 @@ from pypaimon.multimodal.table import _target_schema +@dataclass(frozen=True) +class LeRobotLoadResult: + """Published Paimon state for one imported LeRobot Dataset.""" + + dataset_id: str + version_id: str + frames_snapshot_id: Optional[int] + episodes_snapshot_id: Optional[int] + tasks_snapshot_id: Optional[int] + + def load_from_lerobot( connection, table_name: str, @@ -69,8 +82,9 @@ def load_from_lerobot( batch_size: int = 1024, dataset_id: Optional[str] = None, options: Optional[Mapping[str, object]] = None, - source_options: Optional[Mapping[str, object]] = None): - """Import LeRobot Dataset v3 and return the committed snapshot ID. + source_options: Optional[Mapping[str, object]] = None, +) -> LeRobotLoadResult: + """Import LeRobot Dataset v3 and return its published Paimon state. A missing target table is created from LeRobot metadata. Episode, task, and dataset metadata are stored in companion Paimon tables. ``dataset_id`` @@ -119,19 +133,33 @@ def load_from_lerobot( None, local_info, resolved_source) tables = _prepare_metadata_tables( connection, table.raw_table, owner_id) - metadata_version = _new_id() - _publish_dataset( + version_id = _new_id() + _reserve_dataset_version( + tables["datasets"], + resolved_dataset_id, + version_id, + local_info, + resolved_source, + metadata, + ) + episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( connection, tables, resolved_dataset_id, - metadata_version, + version_id, local_info, resolved_source, metadata, table.identifier, None, ) - return None + return LeRobotLoadResult( + dataset_id=resolved_dataset_id, + version_id=version_id, + frames_snapshot_id=None, + episodes_snapshot_id=episodes_snapshot_id, + tasks_snapshot_id=tasks_snapshot_id, + ) LeRobotDataset = _import_lerobot_dataset() dataset = _open_resolved_dataset( LeRobotDataset, resolved_source, local_info) @@ -157,22 +185,36 @@ def load_from_lerobot( dataset, info, resolved_source) tables = _prepare_metadata_tables( connection, table.raw_table, owner_id) - metadata_version = _new_id() + version_id = _new_id() + _reserve_dataset_version( + tables["datasets"], + resolved_dataset_id, + version_id, + info, + resolved_source, + metadata, + ) if row_count == 0: - _publish_dataset( + episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( connection, tables, resolved_dataset_id, - metadata_version, + version_id, info, resolved_source, metadata, table.identifier, None, ) - return None - snapshot_id = _write_dataset( + return LeRobotLoadResult( + dataset_id=resolved_dataset_id, + version_id=version_id, + frames_snapshot_id=None, + episodes_snapshot_id=episodes_snapshot_id, + tasks_snapshot_id=tasks_snapshot_id, + ) + frames_snapshot_id = _write_dataset( table, dataset, info, @@ -180,21 +222,26 @@ def load_from_lerobot( lerobot_schema, batch_size, resolved_dataset_id, - metadata_version, metadata, ) - _publish_dataset( + episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( connection, tables, resolved_dataset_id, - metadata_version, + version_id, info, resolved_source, metadata, table.identifier, - snapshot_id, + frames_snapshot_id, + ) + return LeRobotLoadResult( + dataset_id=resolved_dataset_id, + version_id=version_id, + frames_snapshot_id=frames_snapshot_id, + episodes_snapshot_id=episodes_snapshot_id, + tasks_snapshot_id=tasks_snapshot_id, ) - return snapshot_id finally: close = getattr(dataset, "close", None) if callable(close): @@ -236,7 +283,7 @@ def _resolved_dataset_id(value, table): _DEFAULT_DATASET_ID_OPTION) if not default_id: raise ValueError( - "Self-contained LeRobot table %s has no default dataset_id." + "LeRobot table %s has no default dataset_id." % table.identifier) return default_id return value.strip() @@ -272,7 +319,7 @@ def _validated_table( "load_from_lerobot; use a new target table." % table.identifier) if table.raw_table.identifier.get_branch_name() is not None: raise ValueError( - "Self-contained LeRobot import does not support table branches.") + "LeRobot import does not support table branches.") _validate_target_schema(table, _frame_schema(source_schema), source) return table, owner_id diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index d635fa191981..ce1dc9345454 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -70,7 +70,6 @@ def _write_dataset( source_schema, batch_size, dataset_id, - metadata_version, metadata): target_schema = _target_schema(table.raw_table) write_builder = table.raw_table.new_batch_write_builder() @@ -104,8 +103,7 @@ def _write_dataset( ) observed_tasks.setdefault(episode_index, set()).update( seen_tasks) - batch = _with_frame_identity( - batch, dataset_id, metadata_version) + batch = _with_frame_identity(batch, dataset_id) batch = _strict_lerobot_table( batch, target_schema, diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 4de4f2355b0b..da12b20cc2e2 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -14,12 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Self-contained metadata tables for imported LeRobot datasets.""" +"""Metadata tables for imported LeRobot datasets.""" import hashlib import json import numbers import uuid +from datetime import datetime, timezone from pathlib import Path import pyarrow as pa @@ -37,7 +38,7 @@ _DATASET_ID = "dataset_id" -_METADATA_VERSION = "metadata_version" +_VERSION_ID = "version_id" _OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" _DEFAULT_DATASET_ID_OPTION = "pypaimon.lerobot.dataset-id" _TABLE_SUFFIXES = { @@ -52,12 +53,13 @@ _FRAME_ID_FIELDS = [ pa.field(_DATASET_ID, pa.string(), nullable=False), - pa.field(_METADATA_VERSION, pa.string(), nullable=False), ] _DATASETS_SCHEMA = pa.schema([ pa.field(_DATASET_ID, pa.string(), nullable=False), - pa.field(_METADATA_VERSION, pa.string(), nullable=False), + pa.field(_VERSION_ID, pa.string(), nullable=False), + pa.field("parent_version_id", pa.string()), pa.field("status", pa.string(), nullable=False), + pa.field("published_at", pa.timestamp("us", tz="UTC")), pa.field("format", pa.string(), nullable=False), pa.field("format_version", pa.string(), nullable=False), pa.field("fps", pa.int64(), nullable=False), @@ -75,7 +77,6 @@ ]) _EPISODES_SCHEMA = pa.schema([ pa.field(_DATASET_ID, pa.string(), nullable=False), - pa.field(_METADATA_VERSION, pa.string(), nullable=False), 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), @@ -87,7 +88,6 @@ ]) _TASKS_SCHEMA = pa.schema([ pa.field(_DATASET_ID, pa.string(), nullable=False), - pa.field(_METADATA_VERSION, pa.string(), nullable=False), pa.field("task_index", pa.int64(), nullable=False), pa.field("task", pa.string(), nullable=False), pa.field("task_metadata_json", pa.string()), @@ -105,12 +105,11 @@ def _frame_schema(source_schema): return pa.schema(list(source_schema) + _FRAME_ID_FIELDS) -def _with_frame_identity(table, dataset_id, metadata_version): +def _with_frame_identity(table, dataset_id): size = table.num_rows return pa.Table.from_arrays( list(table.columns) + [ pa.array([dataset_id] * size, type=pa.string()), - pa.array([metadata_version] * size, type=pa.string()), ], schema=pa.schema(list(table.schema) + _FRAME_ID_FIELDS), ) @@ -172,7 +171,7 @@ def _managed_table_options(frames_identifier, owner_id): identifier = Identifier.from_string(str(frames_identifier)) if identifier.get_branch_name() is not None: raise ValueError( - "Self-contained LeRobot import does not support table branches.") + "LeRobot import does not support table branches.") result = { _OWNER_ID_OPTION: owner_id, _DEFAULT_DATASET_ID_OPTION: str(frames_identifier), @@ -240,32 +239,59 @@ def _prepare_metadata_tables(connection, frames_table, owner_id): return tables +def _reserve_dataset_version( + datasets_table, + dataset_id, + version_id, + info, + source, + metadata): + _ensure_dataset_is_new(datasets_table, dataset_id) + pending = _manifest_row( + dataset_id, + version_id, + None, + "PENDING", + None, + info, + source, + metadata, + None, + None, + None, + ) + snapshot_id = _append_arrow( + datasets_table, + pa.Table.from_pylist([pending], schema=_DATASETS_SCHEMA), + ) + if snapshot_id is None: + raise RuntimeError("LeRobot version reservation created no snapshot.") + + def _publish_dataset( connection, tables, dataset_id, - metadata_version, + version_id, info, source, metadata, frames_identifier, frames_snapshot_id): - episodes = _versioned_table( + episodes = _dataset_table( metadata["episodes"], _EPISODES_SCHEMA, dataset_id, - metadata_version, ) - tasks = _versioned_table( + tasks = _dataset_table( metadata["tasks"], _TASKS_SCHEMA, dataset_id, - metadata_version, ) episodes_snapshot_id = _append_arrow(tables["episodes"], episodes) tasks_snapshot_id = _append_arrow(tables["tasks"], tasks) - tag = "pypaimon-lerobot-%s" % metadata_version + tag = "pypaimon-lerobot-%s" % version_id for identifier, snapshot_id in ( (frames_identifier, frames_snapshot_id), (tables["episodes"].identifier, episodes_snapshot_id), @@ -275,7 +301,10 @@ def _publish_dataset( manifest = _manifest_row( dataset_id, - metadata_version, + version_id, + None, + "READY", + datetime.now(timezone.utc), info, source, metadata, @@ -285,11 +314,15 @@ def _publish_dataset( ) _append_arrow(tables["datasets"], pa.Table.from_pylist( [manifest], schema=_DATASETS_SCHEMA)) + return episodes_snapshot_id, tasks_snapshot_id def _manifest_row( dataset_id, - metadata_version, + version_id, + parent_version_id, + status, + published_at, info, source, metadata, @@ -298,8 +331,10 @@ def _manifest_row( tasks_snapshot_id): return { _DATASET_ID: dataset_id, - _METADATA_VERSION: metadata_version, - "status": "READY", + _VERSION_ID: version_id, + "parent_version_id": parent_version_id, + "status": status, + "published_at": published_at, "format": "lerobot", "format_version": str(info["codebase_version"]), "fps": metadata["fps"], @@ -317,16 +352,27 @@ def _manifest_row( } -def _versioned_table(rows, schema, dataset_id, metadata_version): +def _dataset_table(rows, schema, dataset_id): values = [] for row in rows: value = dict(row) value[_DATASET_ID] = dataset_id - value[_METADATA_VERSION] = metadata_version values.append(value) return pa.Table.from_pylist(values, schema=schema) +def _ensure_dataset_is_new(datasets_table, dataset_id): + snapshot = datasets_table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + return + builder = datasets_table.new_read_builder() + plan = builder.new_scan().plan() + rows = builder.new_read().to_arrow(plan.splits()).to_pylist() + if any(row[_DATASET_ID] == dataset_id for row in rows): + raise ValueError( + "LeRobot dataset_id %s has already been imported." % dataset_id) + + def _append_arrow(table, data): if data.num_rows == 0: return None diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 2fd8b3aa4569..66e0cfa2ef78 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -372,19 +372,25 @@ 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)) + result = connection.load_from_lerobot( + "empty_frames", source) import_lerobot.assert_not_called() + self.assertEqual("default.empty_frames", result.dataset_id) + self.assertEqual(32, len(result.version_id)) + self.assertIsNone(result.frames_snapshot_id) + self.assertIsNone(result.episodes_snapshot_id) + self.assertIsNone(result.tasks_snapshot_id) table = connection.get_table("empty_frames") self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) manifests = _catalog_rows( connection, "empty_frames__datasets") - self.assertEqual(["READY"], [ + self.assertEqual(["PENDING", "READY"], [ row["status"] for row in manifests ]) - self.assertEqual(0, manifests[0]["total_frames"]) - self.assertIsNone(manifests[0]["frames_snapshot_id"]) + self.assertEqual(result.version_id, manifests[1]["version_id"]) + self.assertEqual(0, manifests[1]["total_frames"]) + self.assertIsNone(manifests[1]["frames_snapshot_id"]) finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -460,8 +466,10 @@ def test_empty_dataset_preserves_tasks_and_stats(self): "warehouse": str(temp_dir / "warehouse"), }) - self.assertIsNone(connection.load_from_lerobot("frames", source)) - manifest = _catalog_rows(connection, "frames__datasets")[0] + result = connection.load_from_lerobot("frames", source) + self.assertIsNone(result.frames_snapshot_id) + self.assertEqual(1, result.tasks_snapshot_id) + manifest = _catalog_rows(connection, "frames__datasets")[1] self.assertIsNotNone(manifest["global_stats_json"]) self.assertEqual(1, manifest["total_tasks"]) self.assertEqual(1, manifest["tasks_snapshot_id"]) @@ -741,11 +749,15 @@ 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): + result = self.connection.load_from_lerobot( "robot_data", self.image_source, batch_size=2) - self.assertEqual(1, snapshot_id) + self.assertIsInstance(result, pmm.LeRobotLoadResult) + self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, result.episodes_snapshot_id) + self.assertEqual(1, result.tasks_snapshot_id) + self.assertEqual(32, len(result.version_id)) table = self.connection.get_table("robot_data") schema = table.raw_table.fields @@ -760,7 +772,8 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): self.assertEqual("BIGINT NOT NULL", types["episode_index"]) self.assertEqual("BLOB NOT NULL", types["observation.image"]) self.assertEqual("STRING NOT NULL", types["dataset_id"]) - self.assertEqual("STRING NOT NULL", types["metadata_version"]) + self.assertNotIn("metadata_version", types) + self.assertNotIn("version_id", types) rows = table.scan().select([ "episode_index", @@ -774,7 +787,6 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): "action", "reward", "dataset_id", - "metadata_version", ]).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]) @@ -791,26 +803,27 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): [row["dataset_id"] for row in rows], ) manifests = _catalog_rows(self.connection, "robot_data__datasets") - self.assertEqual(["READY"], [ + self.assertEqual(["PENDING", "READY"], [ row["status"] for row in manifests ]) - manifest = manifests[0] + manifest = manifests[1] self.assertEqual(str(table.identifier), manifest["dataset_id"]) - self.assertEqual(32, len(manifest["metadata_version"])) - self.assertEqual(snapshot_id, manifest["frames_snapshot_id"]) - self.assertEqual(1, manifest["episodes_snapshot_id"]) - self.assertEqual(1, manifest["tasks_snapshot_id"]) + self.assertEqual(result.version_id, manifest["version_id"]) + self.assertIsNone(manifest["parent_version_id"]) + self.assertIsNotNone(manifest["published_at"]) + self.assertEqual( + result.frames_snapshot_id, manifest["frames_snapshot_id"]) + self.assertEqual( + result.episodes_snapshot_id, manifest["episodes_snapshot_id"]) + self.assertEqual( + result.tasks_snapshot_id, manifest["tasks_snapshot_id"]) self.assertEqual("lerobot", manifest["format"]) self.assertEqual("v3.0", manifest["format_version"]) self.assertIsNotNone(manifest["global_stats_json"]) self.assertTrue(manifest["metadata_checksum"].startswith("sha256:")) + tag = "pypaimon-lerobot-%s" % manifest["version_id"] self.assertEqual( - [manifest["metadata_version"]] * 5, - [row["metadata_version"] for row in rows], - ) - tag = "pypaimon-lerobot-%s" % manifest["metadata_version"] - self.assertEqual( - snapshot_id, + result.frames_snapshot_id, self.connection.catalog.get_tag( table.identifier, tag).snapshot.id, ) @@ -824,6 +837,12 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): ) 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) self.assertEqual([(0, 0, 2), (1, 2, 5)], [ (row["episode_index"], row["dataset_from_index"], row["dataset_to_index"]) @@ -835,11 +854,16 @@ def test_import_infers_schema_preserves_episodes_and_appends(self): row["episode_stats_json"] is not None for row in episodes)) 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.assertEqual([(0, "pick"), (1, "place")], [ (row["task_index"], row["task"]) for row in tasks ]) self.assertEqual( - snapshot_id, + result.frames_snapshot_id, table.raw_table.snapshot_manager().get_latest_snapshot().id, ) @@ -858,22 +882,10 @@ 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) - manifests = _catalog_rows(self.connection, "robot_data__datasets") - ready_versions = [ - row["metadata_version"] for row in manifests - if row["status"] == "READY" - ] - self.assertEqual(2, len(ready_versions)) - self.assertEqual(2, len(set(ready_versions))) - versions = table.scan().select([ - "metadata_version" - ]).to_arrow().column(0).to_pylist() - self.assertEqual([5, 5], sorted( - versions.count(version) for version in set(versions))) + with self.assertRaisesRegex(ValueError, "already been imported"): + self.connection.load_from_lerobot( + "robot_data", self.image_source, batch_size=4) + self.assertEqual(5, table.scan().to_arrow().num_rows) def test_frame_controls_must_match_published_episode_metadata(self): cases = [ @@ -903,11 +915,11 @@ def test_frame_controls_must_match_published_episode_metadata(self): with self.assertRaisesRegex( ValueError, "has %s" % column): self.connection.load_from_lerobot(table_name, source) - datasets = self.connection.catalog.get_table( - self.connection._identifier( - table_name + "__datasets")) - self.assertIsNone( - datasets.snapshot_manager().get_latest_snapshot()) + manifests = _catalog_rows( + self.connection, table_name + "__datasets") + self.assertEqual(["PENDING"], [ + row["status"] for row in manifests + ]) table = self.connection.get_table(table_name) self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) @@ -965,10 +977,11 @@ def test_episode_tasks_must_exactly_match_frame_tasks(self): ValueError, "declares task indices"): self.connection.load_from_lerobot( "extra_episode_task", source) - datasets = self.connection.catalog.get_table( - self.connection._identifier("extra_episode_task__datasets")) - self.assertIsNone( - datasets.snapshot_manager().get_latest_snapshot()) + self.assertEqual( + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, "extra_episode_task__datasets")], + ) table = self.connection.get_table("extra_episode_task") self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) @@ -996,10 +1009,11 @@ def test_nonempty_dataset_cannot_publish_without_tasks(self): with self.assertRaisesRegex(ValueError, "task_index"): self.connection.load_from_lerobot("missing_tasks", source) - datasets = self.connection.catalog.get_table( - self.connection._identifier("missing_tasks__datasets")) - self.assertIsNone( - datasets.snapshot_manager().get_latest_snapshot()) + self.assertEqual( + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, "missing_tasks__datasets")], + ) table = self.connection.get_table("missing_tasks") self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) @@ -1011,13 +1025,13 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): with patch( "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): - snapshot_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "oss_images", source, batch_size=2, ) - self.assertEqual(1, snapshot_id) + self.assertEqual(1, result.frames_snapshot_id) table = self.connection.get_table("oss_images") rows = table.scan().select([ "episode_index", "frame_index", "index", "task" @@ -1060,8 +1074,9 @@ def test_empty_oss_source_does_not_require_episode_directory(self): with patch( "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): - self.assertIsNone(self.connection.load_from_lerobot( - "empty_oss", source)) + result = self.connection.load_from_lerobot( + "empty_oss", source) + self.assertIsNone(result.frames_snapshot_id) self.assertEqual( 0, _catalog_rows( @@ -1073,19 +1088,19 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): self.connection.catalog, "create_tag", side_effect=NotImplementedError): - snapshot_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "tag_fallback", self.image_source) manifest = _catalog_rows( - self.connection, "tag_fallback__datasets")[0] - tag = "pypaimon-lerobot-%s" % manifest["metadata_version"] + self.connection, "tag_fallback__datasets")[1] + tag = "pypaimon-lerobot-%s" % manifest["version_id"] table = self.connection.get_table("tag_fallback") self.assertEqual( - snapshot_id, + result.frames_snapshot_id, table.raw_table.tag_manager().get(tag).id, ) - def test_failed_publication_leaves_no_manifest(self): + def test_failed_publication_leaves_pending_manifest(self): with patch( "pypaimon.multimodal.lerobot.api._publish_dataset", side_effect=RuntimeError("publish failed")): @@ -1093,17 +1108,16 @@ def test_failed_publication_leaves_no_manifest(self): self.connection.load_from_lerobot( "failed_publish", self.image_source) - datasets = self.connection.catalog.get_table( - self.connection._identifier("failed_publish__datasets")) - self.assertIsNone( - datasets.snapshot_manager().get_latest_snapshot()) + manifests = _catalog_rows( + self.connection, "failed_publish__datasets") + self.assertEqual(["PENDING"], [ + row["status"] for row in manifests + ]) frames = self.connection.get_table("failed_publish").scan().select([ - "metadata_version" + "dataset_id" ]).to_arrow() self.assertEqual(5, frames.num_rows) - versions = set(frames.column(0).to_pylist()) - self.assertEqual(1, len(versions)) - self.assertEqual(32, len(next(iter(versions)))) + self.assertEqual(1, len(set(frames.column(0).to_pylist()))) def test_existing_unmanaged_table_is_rejected(self): info = json.loads((self.image_source / "meta" / "info.json").read_text()) @@ -1131,7 +1145,7 @@ def test_companion_tables_must_be_append_only(self): identifier, PaimonSchema.from_pyarrow_schema( _DATASETS_SCHEMA, - primary_keys=["metadata_version"], + primary_keys=["version_id"], options={ "bucket": "1", _OWNER_ID_OPTION: owner_id, @@ -1249,17 +1263,24 @@ def test_table_group_survives_frame_table_rename(self): self.connection._identifier("after_rename"), ) - self.assertEqual(2, self.connection.load_from_lerobot( - "after_rename", self.image_source)) + result = self.connection.load_from_lerobot( + "after_rename", + self.image_source, + dataset_id="renamed-dataset", + ) + self.assertEqual(2, result.frames_snapshot_id) manifests = _catalog_rows( self.connection, "before_rename__datasets") ready = [row for row in manifests if row["status"] == "READY"] self.assertEqual(2, len(ready)) self.assertEqual(2, len({ - row["metadata_version"] for row in ready + row["version_id"] for row in ready })) self.assertEqual( - {self.connection._identifier("before_rename")}, + { + self.connection._identifier("before_rename"), + "renamed-dataset", + }, {row["dataset_id"] for row in ready}, ) From 7a23db9e0b4131c8568f88292ea6071bef5e4895 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 09:38:51 -0700 Subject: [PATCH 03/32] [python] Make LeRobot imports retry-safe --- docs/docs/pypaimon/multimodal-api.mdx | 4 +- .../pypaimon/multimodal/connection.py | 7 +- .../pypaimon/multimodal/lerobot/api.py | 244 +++++++-------- .../pypaimon/multimodal/lerobot/metadata.py | 22 +- .../pypaimon/tests/multimodal_lerobot_test.py | 280 +++++------------- 5 files changed, 190 insertions(+), 367 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 96e8ce11c3c7..5faeded0ab23 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -563,7 +563,9 @@ Component rows carry `dataset_id` and stable LeRobot indices. `version_id` exists only in the Dataset manifest, where it resolves the exact frame, Episode, and Task snapshots. These snapshots define the published Dataset state; tags retain them. Set `dataset_id=` to choose a stable ID. The one-time -importer accepts each `dataset_id` once. `drop_table()` removes all four tables. +importer requires a new target table. A failed import removes the table group +so the call can be retried. `drop_table()` removes all four tables; companion +tables cannot be dropped separately through `MultimodalConnection`. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index b78bd1d823ed..76868a88a303 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -127,7 +127,7 @@ def load_from_lerobot( dataset_id: Optional[str] = None, options=None, source_options=None): - """Import LeRobot Dataset v3 and return its published table state.""" + """Import LeRobot Dataset v3 into a new Paimon table group.""" from pypaimon.multimodal.lerobot import load_from_lerobot return load_from_lerobot( self, @@ -149,6 +149,11 @@ def drop_table(self, name: str, ignore_if_not_exists: bool = False): ) raw_table = self.catalog.get_table(identifier) table_options = raw_table.table_schema.options + if (_OWNER_ID_OPTION in table_options + and _DEFAULT_DATASET_ID_OPTION not in table_options): + raise ValueError( + "%s is a managed LeRobot companion table; drop its " + "frame table instead." % identifier) if _DEFAULT_DATASET_ID_OPTION in table_options: if Identifier.from_string( identifier).get_branch_name() is not None: diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 8d0c4713fd60..8db75c27f555 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -21,15 +21,10 @@ from dataclasses import dataclass from typing import Mapping, Optional -import pyarrow as pa - -from pypaimon.catalog.catalog_exception import ( - DatabaseNotExistException, - TableNotExistException, -) +from pypaimon.catalog.catalog_exception import TableAlreadyExistException from pypaimon.multimodal.lerobot.metadata import ( _DEFAULT_DATASET_ID_OPTION, - _OWNER_ID_OPTION, + _drop_import_tables, _frame_schema, _load_dataset_metadata, _managed_table_options, @@ -39,14 +34,10 @@ _reject_subtasks, _reserve_dataset_version, ) -from pypaimon.multimodal.lerobot.loader import ( - _strict_lerobot_table, - _write_dataset, -) +from pypaimon.multimodal.lerobot.loader import _write_dataset from pypaimon.multimodal.lerobot.schema import ( _require_v3, _schema_from_info, - _validate_lerobot_schema, ) from pypaimon.multimodal.lerobot.source import ( _has_tasks, @@ -60,7 +51,6 @@ _validated_source_options, _validate_source_kerberos, ) -from pypaimon.multimodal.table import _target_schema @dataclass(frozen=True) @@ -86,7 +76,7 @@ def load_from_lerobot( ) -> LeRobotLoadResult: """Import LeRobot Dataset v3 and return its published Paimon state. - A missing target table is created from LeRobot metadata. Episode, task, and + A new target table is created from LeRobot metadata. Episode, task, and dataset metadata are stored in companion Paimon tables. ``dataset_id`` defaults to the target identifier. FileIO URI credentials come only from ``source_options`` and are not inherited from @@ -120,45 +110,19 @@ def load_from_lerobot( include_task=total_tasks > 0, ) _reject_subtasks(None, resolved_source) - table, owner_id = _validated_table( - connection, - table_name, - source_schema, - options, - resolved_source, - ) - resolved_dataset_id = _resolved_dataset_id( - dataset_id, table) metadata = _load_dataset_metadata( None, local_info, resolved_source) - tables = _prepare_metadata_tables( - connection, table.raw_table, owner_id) - version_id = _new_id() - _reserve_dataset_version( - tables["datasets"], - resolved_dataset_id, - version_id, - local_info, - resolved_source, - metadata, - ) - episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( + return _import_dataset( connection, - tables, - resolved_dataset_id, - version_id, + table_name, + None, local_info, resolved_source, + source_schema, + batch_size, + dataset_id, + options, metadata, - table.identifier, - None, - ) - return LeRobotLoadResult( - dataset_id=resolved_dataset_id, - version_id=version_id, - frames_snapshot_id=None, - episodes_snapshot_id=episodes_snapshot_id, - tasks_snapshot_id=tasks_snapshot_id, ) LeRobotDataset = _import_lerobot_dataset() dataset = _open_resolved_dataset( @@ -166,86 +130,97 @@ def load_from_lerobot( 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) lerobot_schema = _schema_from_info( info, include_task=_has_tasks(dataset, info)) _reject_subtasks(dataset, resolved_source) - table, owner_id = _validated_table( - connection, - table_name, - lerobot_schema, - options, - resolved_source, - ) - resolved_dataset_id = _resolved_dataset_id( - dataset_id, table) metadata = _load_dataset_metadata( dataset, info, resolved_source) - tables = _prepare_metadata_tables( - connection, table.raw_table, owner_id) - version_id = _new_id() - _reserve_dataset_version( - tables["datasets"], - resolved_dataset_id, - version_id, + return _import_dataset( + connection, + table_name, + dataset, info, resolved_source, + lerobot_schema, + batch_size, + dataset_id, + options, metadata, ) + finally: + close = getattr(dataset, "close", None) + if callable(close): + close() - if row_count == 0: - episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( - connection, - tables, - resolved_dataset_id, - version_id, - info, - resolved_source, - metadata, - table.identifier, - None, - ) - return LeRobotLoadResult( - dataset_id=resolved_dataset_id, - version_id=version_id, - frames_snapshot_id=None, - episodes_snapshot_id=episodes_snapshot_id, - tasks_snapshot_id=tasks_snapshot_id, - ) + +def _import_dataset( + connection, + table_name, + dataset, + info, + source, + source_schema, + batch_size, + dataset_id, + options, + metadata): + table, owner_id = _create_target_table( + connection, table_name, source_schema, options) + try: + resolved_dataset_id = _resolved_dataset_id(dataset_id, table) + tables = _prepare_metadata_tables( + connection, table.raw_table, owner_id) + version_id = _new_id() + _reserve_dataset_version( + tables["datasets"], + resolved_dataset_id, + version_id, + info, + source, + metadata, + ) + frames_snapshot_id = None + if int(info["total_frames"]) > 0: frames_snapshot_id = _write_dataset( table, dataset, info, - resolved_source, - lerobot_schema, + source, + source_schema, batch_size, resolved_dataset_id, metadata, ) - episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( - connection, - tables, - resolved_dataset_id, - version_id, - info, - resolved_source, - metadata, - table.identifier, - frames_snapshot_id, - ) - return LeRobotLoadResult( - dataset_id=resolved_dataset_id, - version_id=version_id, - frames_snapshot_id=frames_snapshot_id, - episodes_snapshot_id=episodes_snapshot_id, - tasks_snapshot_id=tasks_snapshot_id, - ) - finally: - close = getattr(dataset, "close", None) - if callable(close): - close() + episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( + connection, + tables, + resolved_dataset_id, + version_id, + info, + source, + metadata, + table.identifier, + frames_snapshot_id, + ) + return LeRobotLoadResult( + dataset_id=resolved_dataset_id, + version_id=version_id, + frames_snapshot_id=frames_snapshot_id, + episodes_snapshot_id=episodes_snapshot_id, + tasks_snapshot_id=tasks_snapshot_id, + ) + except BaseException as error: + try: + _drop_import_tables( + connection.catalog, table.raw_table, owner_id) + except BaseException as cleanup_error: + raise RuntimeError( + "LeRobot import failed and cleanup also failed: %s" + % cleanup_error + ) from error + raise def _validated_counts(info, source): @@ -289,48 +264,27 @@ def _resolved_dataset_id(value, table): return value.strip() -def _validated_table( - connection, table_name, source_schema, options, source): +def _create_target_table( + connection, table_name, source_schema, options): + owner_id = _new_id() + create_options = dict(options or {}) + managed_options = _managed_table_options( + connection._identifier(table_name), owner_id) + 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: - table = connection.get_table(table_name) - except (DatabaseNotExistException, TableNotExistException): - owner_id = _new_id() - create_options = dict(options or {}) - managed_options = _managed_table_options( - connection._identifier(table_name), owner_id) - 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) table = connection.create_table( table_name, schema=_frame_schema(source_schema), options=create_options, ) - _validate_target_schema(table, _frame_schema(source_schema), source) - return table, owner_id - - owner_id = table.raw_table.table_schema.options.get(_OWNER_ID_OPTION) - if owner_id is None: + except TableAlreadyExistException as error: raise ValueError( - "Existing LeRobot target %s is not managed by " - "load_from_lerobot; use a new target table." % table.identifier) - if table.raw_table.identifier.get_branch_name() is not None: - raise ValueError( - "LeRobot import does not support table branches.") - _validate_target_schema(table, _frame_schema(source_schema), source) + "LeRobot target %s already exists; use a new target table." + % connection._identifier(table_name) + ) from error return table, owner_id - - -def _validate_target_schema(table, source_schema, source): - 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, - ) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index da12b20cc2e2..0c4993df5304 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -246,7 +246,6 @@ def _reserve_dataset_version( info, source, metadata): - _ensure_dataset_is_new(datasets_table, dataset_id) pending = _manifest_row( dataset_id, version_id, @@ -361,16 +360,17 @@ def _dataset_table(rows, schema, dataset_id): return pa.Table.from_pylist(values, schema=schema) -def _ensure_dataset_is_new(datasets_table, dataset_id): - snapshot = datasets_table.snapshot_manager().get_latest_snapshot() - if snapshot is None: - return - builder = datasets_table.new_read_builder() - plan = builder.new_scan().plan() - rows = builder.new_read().to_arrow(plan.splits()).to_pylist() - if any(row[_DATASET_ID] == dataset_id for row in rows): - raise ValueError( - "LeRobot dataset_id %s has already been imported." % dataset_id) +def _drop_import_tables(catalog, frames_table, owner_id): + identifiers = list( + _companion_table_identifiers(frames_table).values()) + identifiers.append(frames_table.identifier.get_full_name()) + for identifier in identifiers: + try: + table = catalog.get_table(identifier) + except (DatabaseNotExistException, TableNotExistException): + continue + if table.table_schema.options.get(_OWNER_ID_OPTION) == owner_id: + catalog.drop_table(identifier, ignore_if_not_exists=True) def _append_arrow(table, data): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 66e0cfa2ef78..219e8ba37afa 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -19,7 +19,9 @@ 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 @@ -28,18 +30,12 @@ import pyarrow.fs as pafs import pyarrow.parquet as pq -from pypaimon import Schema as PaimonSchema from pypaimon.catalog.catalog_exception import TableNotExistException 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.metadata import ( - _DATASETS_SCHEMA, - _OWNER_ID_OPTION, - _frame_schema, - _managed_table_options, -) +from pypaimon.multimodal.lerobot.metadata import _managed_table_options from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, @@ -529,12 +525,8 @@ def test_invalid_fps_creates_no_snapshot_or_manifest(self): with self.assertRaisesRegex(ValueError, "fps must be positive"): connection.load_from_lerobot("frames", source) - table = connection.get_table("frames") - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) with self.assertRaises(TableNotExistException): - connection.catalog.get_table( - connection._identifier("frames__datasets")) + connection.get_table("frames") finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -882,7 +874,7 @@ def test_import_infers_schema_and_preserves_episodes(self): [imported[index] for index in range(5)], ) - with self.assertRaisesRegex(ValueError, "already been imported"): + 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) @@ -915,14 +907,10 @@ def test_frame_controls_must_match_published_episode_metadata(self): with self.assertRaisesRegex( ValueError, "has %s" % column): self.connection.load_from_lerobot(table_name, source) - manifests = _catalog_rows( - self.connection, table_name + "__datasets") - self.assertEqual(["PENDING"], [ - row["status"] for row in manifests - ]) - table = self.connection.get_table(table_name) - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) + for suffix in ("", "__datasets", "__episodes", "__tasks"): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(table_name + suffix)) def test_explicit_dataset_id_is_shared_by_all_tables(self): dataset_id = "aloha_pick_cube@3" @@ -977,14 +965,8 @@ def test_episode_tasks_must_exactly_match_frame_tasks(self): ValueError, "declares task indices"): self.connection.load_from_lerobot( "extra_episode_task", source) - self.assertEqual( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, "extra_episode_task__datasets")], - ) - table = self.connection.get_table("extra_episode_task") - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) + with self.assertRaises(TableNotExistException): + self.connection.get_table("extra_episode_task") def test_nonempty_dataset_cannot_publish_without_tasks(self): source = self.temp_dir / "missing_tasks" @@ -1009,14 +991,8 @@ def test_nonempty_dataset_cannot_publish_without_tasks(self): with self.assertRaisesRegex(ValueError, "task_index"): self.connection.load_from_lerobot("missing_tasks", source) - self.assertEqual( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, "missing_tasks__datasets")], - ) - table = self.connection.get_table("missing_tasks") - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) + with self.assertRaises(TableNotExistException): + self.connection.get_table("missing_tasks") def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" @@ -1100,7 +1076,7 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): table.raw_table.tag_manager().get(tag).id, ) - def test_failed_publication_leaves_pending_manifest(self): + def test_failed_publication_is_cleaned_and_can_retry(self): with patch( "pypaimon.multimodal.lerobot.api._publish_dataset", side_effect=RuntimeError("publish failed")): @@ -1108,89 +1084,65 @@ def test_failed_publication_leaves_pending_manifest(self): self.connection.load_from_lerobot( "failed_publish", self.image_source) - manifests = _catalog_rows( - self.connection, "failed_publish__datasets") - self.assertEqual(["PENDING"], [ - row["status"] for row in manifests - ]) - frames = self.connection.get_table("failed_publish").scan().select([ - "dataset_id" - ]).to_arrow() - self.assertEqual(5, frames.num_rows) - self.assertEqual(1, len(set(frames.column(0).to_pylist()))) + for suffix in ("", "__datasets", "__episodes", "__tasks"): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier("failed_publish" + suffix)) + + result = self.connection.load_from_lerobot( + "failed_publish", self.image_source) + self.assertEqual(1, result.frames_snapshot_id) - def test_existing_unmanaged_table_is_rejected(self): + 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) - table = self.connection.create_table("unmanaged", schema=schema) + table = self.connection.create_table("existing", schema=schema) - with self.assertRaisesRegex(ValueError, "is not managed"): + with self.assertRaisesRegex(ValueError, "already exists"): self.connection.load_from_lerobot( - "unmanaged", self.image_source) + "existing", self.image_source) self.assertIsNone( table.raw_table.snapshot_manager().get_latest_snapshot()) - def test_companion_tables_must_be_append_only(self): - info = json.loads((self.image_source / "meta" / "info.json").read_text()) - source_schema = _schema_from_info(info, include_task=True) - owner_id = "test-owner" - table = self.connection.create_table( - "invalid_group", - schema=_frame_schema(source_schema), - options=_managed_table_options( - self.connection._identifier("invalid_group"), owner_id), - ) - identifier = self.connection._identifier("invalid_group__datasets") - self.connection.catalog.create_table( - identifier, - PaimonSchema.from_pyarrow_schema( - _DATASETS_SCHEMA, - primary_keys=["version_id"], - options={ - "bucket": "1", - _OWNER_ID_OPTION: owner_id, - }, - ), - False, - ) + def test_concurrent_import_cannot_claim_the_same_target(self): + from pypaimon.multimodal.lerobot import api - with self.assertRaisesRegex(ValueError, "must be append-only"): - self.connection.load_from_lerobot( - "invalid_group", self.image_source) - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) + original_reserve = api._reserve_dataset_version + reserved = threading.Event() + release = threading.Event() - def test_stale_companion_tables_are_rejected(self): - info = json.loads((self.image_source / "meta" / "info.json").read_text()) - source_schema = _schema_from_info(info, include_task=True) - table = self.connection.create_table( - "stale_group", - schema=_frame_schema(source_schema), - options=_managed_table_options( - self.connection._identifier("stale_group"), "new-owner"), - ) - identifier = self.connection._identifier("stale_group__datasets") - self.connection.catalog.create_table( - identifier, - PaimonSchema.from_pyarrow_schema( - _DATASETS_SCHEMA, - options={ - "bucket": "-1", - _OWNER_ID_OPTION: "old-owner", - }, - ), - False, - ) + def reserve_then_wait(*args, **kwargs): + result = original_reserve(*args, **kwargs) + reserved.set() + release.wait(10) + return result - with self.assertRaisesRegex(ValueError, "different target table"): - self.connection.load_from_lerobot( - "stale_group", self.image_source) - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) - with self.assertRaisesRegex(ValueError, "Refusing to drop"): - self.connection.drop_table("stale_group") - self.connection.get_table("stale_group") - self.connection.catalog.get_table(identifier) + with patch.object( + api, + "_reserve_dataset_version", + side_effect=reserve_then_wait): + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + self.connection.load_from_lerobot, + "concurrent", + self.image_source, + ) + try: + self.assertTrue(reserved.wait(10)) + with self.assertRaisesRegex( + ValueError, "already exists"): + self.connection.load_from_lerobot( + "concurrent", self.image_source) + finally: + release.set() + result = future.result(timeout=30) + + self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual( + 5, + self.connection.get_table( + "concurrent").scan().to_arrow().num_rows, + ) def test_drop_table_removes_companion_tables(self): self.connection.load_from_lerobot("drop_group", self.image_source) @@ -1232,14 +1184,16 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.catalog.get_table( self.connection._identifier(name)) - def test_companion_table_can_be_dropped_directly(self): + def test_companion_table_cannot_be_dropped_directly(self): self.connection.load_from_lerobot( "direct_drop", self.image_source) - self.connection.drop_table("direct_drop__tasks") + + with self.assertRaisesRegex(ValueError, "companion table"): + self.connection.drop_table("direct_drop__tasks") self.connection.get_table("direct_drop") - with self.assertRaises(TableNotExistException): - self.connection.get_table("direct_drop__tasks") + self.connection.catalog.get_table( + self.connection._identifier("direct_drop__tasks")) self.connection.drop_table("direct_drop") def test_drop_table_rejects_managed_branch(self): @@ -1263,26 +1217,9 @@ def test_table_group_survives_frame_table_rename(self): self.connection._identifier("after_rename"), ) - result = self.connection.load_from_lerobot( - "after_rename", - self.image_source, - dataset_id="renamed-dataset", - ) - self.assertEqual(2, result.frames_snapshot_id) - manifests = _catalog_rows( - self.connection, "before_rename__datasets") - ready = [row for row in manifests if row["status"] == "READY"] - self.assertEqual(2, len(ready)) - self.assertEqual(2, len({ - row["version_id"] for row in ready - })) - self.assertEqual( - { - self.connection._identifier("before_rename"), - "renamed-dataset", - }, - {row["dataset_id"] for row in ready}, - ) + with self.assertRaisesRegex(ValueError, "already exists"): + self.connection.load_from_lerobot( + "after_rename", self.image_source) self.connection.drop_table("after_rename") for name in ( @@ -1292,80 +1229,5 @@ def test_table_group_survives_frame_table_rename(self): self.connection.catalog.get_table( self.connection._identifier(name)) - 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=_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 - options = { - "file.format": "parquet", - "vector.file.format": "parquet", - } - options.update(_managed_table_options( - self.connection._identifier(table_name), - "owner-%s" % name, - )) - table = self.connection.create_table( - table_name, - schema=_frame_schema(pa.schema([ - replacement if field.name == replacement.name - else field - for field in schema - ])), - options=options, - ) - - 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() From 580233a84b7faf9ad60e4194d5e8a264825107f5 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 10:34:43 -0700 Subject: [PATCH 04/32] [python] Preserve LeRobot import result on cleanup --- .../pypaimon/multimodal/connection.py | 6 +- .../pypaimon/multimodal/lerobot/api.py | 3 +- .../pypaimon/multimodal/lerobot/source.py | 13 +++- .../pypaimon/tests/multimodal_lerobot_test.py | 74 +++++++++++++++++++ 4 files changed, 93 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 76868a88a303..a32dedd36a42 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -73,6 +73,7 @@ def create_table( identifier = self._identifier(name) already_exists = _table_exists(self.catalog, identifier) paimon_schema = _to_paimon_schema(schema, data, options, partitioned) + _validate_multimodal_schema(paimon_schema, identifier) self._create_database_for(identifier) try: @@ -217,7 +218,10 @@ def _table_exists(catalog, identifier: str) -> bool: def _validate_multimodal_table(table, identifier: str): - table_schema = table.table_schema + _validate_multimodal_schema(table.table_schema, identifier) + + +def _validate_multimodal_schema(table_schema, identifier: str): options = table_schema.options if str(options.get("data-evolution.enabled", "false")).lower() != "true": raise ValueError( diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 8db75c27f555..d0a0841a6787 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -40,6 +40,7 @@ _schema_from_info, ) from pypaimon.multimodal.lerobot.source import ( + _close_quietly, _has_tasks, _import_lerobot_dataset, _load_hub_info, @@ -152,7 +153,7 @@ def load_from_lerobot( finally: close = getattr(dataset, "close", None) if callable(close): - close() + _close_quietly(dataset, "dataset") def _import_dataset( diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index eaf995d69eda..af03227bc8c8 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -17,6 +17,7 @@ """LeRobot source resolution for local, Hub, and FileIO datasets.""" import json +import logging import posixpath from bisect import bisect_right from contextlib import closing, contextmanager @@ -39,6 +40,9 @@ from pypaimon.multimodal.lerobot.loader import _encode_media_frame +_LOGGER = logging.getLogger(__name__) + + @dataclass(frozen=True) class _LeRobotSource: path: str @@ -103,7 +107,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): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 219e8ba37afa..2a77a3bb7f3d 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -656,6 +656,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", @@ -1093,6 +1108,65 @@ def test_failed_publication_is_cleaned_and_can_retry(self): "failed_publish", self.image_source) self.assertEqual(1, result.frames_snapshot_id) + 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")) + + result = self.connection.load_from_lerobot( + "invalid_options", self.image_source) + self.assertEqual(1, result.frames_snapshot_id) + + 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): + result = self.connection.load_from_lerobot( + "close_failure", self.image_source) + + self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual( + ["PENDING", "READY"], + [row["status"] for row in _catalog_rows( + self.connection, "close_failure__datasets")], + ) + + 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._Hdf5SourceFileIO", + return_value=source_file_io): + result = self.connection.load_from_lerobot( + "source_close_failure", source) + + self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual( + ["PENDING", "READY"], + [row["status"] for row in _catalog_rows( + self.connection, "source_close_failure__datasets")], + ) + 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) From f127d2aaf110b60d5a396aff027e8af436ef56ca Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 11:00:01 -0700 Subject: [PATCH 05/32] [python] Clean up failed LeRobot table creation --- .../pypaimon/multimodal/connection.py | 12 ++++- .../pypaimon/multimodal/lerobot/api.py | 24 ++++++++- .../pypaimon/tests/multimodal_lerobot_test.py | 50 ++++++++++++++++++- .../pypaimon/tests/multimodal_table_test.py | 29 +++++++++++ 4 files changed, 111 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index a32dedd36a42..ec772a59fda5 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -72,8 +72,16 @@ def create_table( """Create a multimodal table and optionally add initial data.""" identifier = self._identifier(name) already_exists = _table_exists(self.catalog, identifier) - paimon_schema = _to_paimon_schema(schema, data, options, partitioned) - _validate_multimodal_schema(paimon_schema, identifier) + if already_exists and ignore_if_exists: + return self.get_table(name) + try: + paimon_schema = _to_paimon_schema( + schema, data, options, partitioned) + _validate_multimodal_schema(paimon_schema, identifier) + except ValueError: + if ignore_if_exists and _table_exists(self.catalog, identifier): + return self.get_table(name) + raise self._create_database_for(identifier) try: diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index d0a0841a6787..2bfe936ea17a 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -21,9 +21,14 @@ from dataclasses import dataclass from typing import Mapping, Optional -from pypaimon.catalog.catalog_exception import TableAlreadyExistException +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableAlreadyExistException, + TableNotExistException, +) from pypaimon.multimodal.lerobot.metadata import ( _DEFAULT_DATASET_ID_OPTION, + _OWNER_ID_OPTION, _drop_import_tables, _frame_schema, _load_dataset_metadata, @@ -288,4 +293,21 @@ def _create_target_table( "LeRobot target %s already exists; use a new target table." % connection._identifier(table_name) ) from error + except BaseException as error: + try: + frames_table = connection.catalog.get_table( + connection._identifier(table_name)) + except (DatabaseNotExistException, TableNotExistException): + frames_table = None + if frames_table is not None and frames_table.table_schema.options.get( + _OWNER_ID_OPTION) == owner_id: + try: + _drop_import_tables( + connection.catalog, frames_table, owner_id) + except BaseException as cleanup_error: + raise RuntimeError( + "LeRobot target creation failed and cleanup also failed: " + "%s" % cleanup_error + ) from error + raise return table, owner_id diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 2a77a3bb7f3d..2135b524c0fd 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -35,7 +35,10 @@ from pypaimon.common.options import Options from pypaimon.multimodal.hdf5 import _Hdf5SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot -from pypaimon.multimodal.lerobot.metadata import _managed_table_options +from pypaimon.multimodal.lerobot.metadata import ( + _managed_table_options, + _OWNER_ID_OPTION, +) from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, @@ -1123,6 +1126,51 @@ def test_invalid_target_options_do_not_leave_table(self): "invalid_options", self.image_source) self.assertEqual(1, result.frames_snapshot_id) + def test_target_open_failure_is_cleaned_and_can_retry(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) + result = self.connection.load_from_lerobot( + "failed_open", self.image_source) + + self.assertEqual(1, result.frames_snapshot_id) + + def test_target_open_failure_does_not_drop_another_owner(self): + original_create = self.connection.create_table + + def create_other_owner(*args, **kwargs): + options = dict(kwargs["options"]) + options[_OWNER_ID_OPTION] = "other-owner" + kwargs["options"] = options + original_create(*args, **kwargs) + raise RuntimeError("get failed") + + with patch.object( + self.connection, + "create_table", + side_effect=create_other_owner): + with self.assertRaisesRegex(RuntimeError, "get failed"): + self.connection.load_from_lerobot( + "other_owner", self.image_source) + + table = self.connection.catalog.get_table( + self.connection._identifier("other_owner")) + self.assertEqual( + "other-owner", + table.table_schema.options[_OWNER_ID_OPTION], + ) + def test_dataset_close_failure_does_not_override_success(self): from pypaimon.multimodal.lerobot import api diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index 85794d232c9b..bec750bfb2d6 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -760,6 +760,35 @@ def test_get_table_rejects_primary_key_table(self): with self.assertRaisesRegex(ValueError, "primary keys"): self.conn.get_table("pk") + def test_create_table_ignores_invalid_options_when_table_exists(self): + schema = _schema({"id": pa.int32()}) + expected = self.conn.create_table("existing", schema=schema) + + actual = self.conn.create_table( + "existing", + schema=_schema({ + "id": pa.int32(), + "embedding": _vector(3), + }), + options={"data-evolution.enabled": "false"}, + ignore_if_exists=True, + ) + + self.assertEqual(expected.identifier, actual.identifier) + with patch( + "pypaimon.multimodal.connection._table_exists", + side_effect=[False, True]): + raced = self.conn.create_table( + "existing", + schema=_schema({ + "id": pa.int32(), + "embedding": _vector(3), + }), + options={"data-evolution.enabled": "false"}, + ignore_if_exists=True, + ) + self.assertEqual(expected.identifier, raced.identifier) + def test_create_table_can_add_initial_data_and_get_by_short_name(self): self.conn.create_table( "users", From 81540dcab24794a83e04d4615b40806f03ce75f6 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 11:11:42 -0700 Subject: [PATCH 06/32] [python] Handle concurrent multimodal table deletion --- .../pypaimon/multimodal/connection.py | 12 +++-- .../pypaimon/tests/multimodal_table_test.py | 46 ++++++++++++++++++- 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index ec772a59fda5..dad5baeb0c26 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -73,14 +73,20 @@ def create_table( identifier = self._identifier(name) already_exists = _table_exists(self.catalog, identifier) if already_exists and ignore_if_exists: - return self.get_table(name) + try: + return self.get_table(name) + except (DatabaseNotExistException, TableNotExistException): + pass try: paimon_schema = _to_paimon_schema( schema, data, options, partitioned) _validate_multimodal_schema(paimon_schema, identifier) except ValueError: - if ignore_if_exists and _table_exists(self.catalog, identifier): - return self.get_table(name) + if ignore_if_exists: + try: + return self.get_table(name) + except (DatabaseNotExistException, TableNotExistException): + pass raise self._create_database_for(identifier) diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index bec750bfb2d6..378f7705264d 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -777,7 +777,7 @@ def test_create_table_ignores_invalid_options_when_table_exists(self): self.assertEqual(expected.identifier, actual.identifier) with patch( "pypaimon.multimodal.connection._table_exists", - side_effect=[False, True]): + return_value=False): raced = self.conn.create_table( "existing", schema=_schema({ @@ -789,6 +789,50 @@ def test_create_table_ignores_invalid_options_when_table_exists(self): ) self.assertEqual(expected.identifier, raced.identifier) + def test_create_table_handles_concurrent_delete_when_ignoring(self): + schema = _schema({"id": pa.int32()}) + self.conn.create_table("deleted", schema=schema) + original_get = self.conn.get_table + deleted = [False] + + def delete_once(name): + if not deleted[0]: + deleted[0] = True + self.conn.catalog.drop_table("default.deleted", False) + return original_get(name) + + with patch.object( + self.conn, "get_table", side_effect=delete_once): + table = self.conn.create_table( + "deleted", schema=schema, ignore_if_exists=True) + self.assertEqual("default.deleted", table.identifier) + + self.conn.create_table("fallback_deleted", schema=schema) + deleted[0] = False + + def delete_fallback_once(name): + if not deleted[0]: + deleted[0] = True + self.conn.catalog.drop_table( + "default.fallback_deleted", False) + return original_get(name) + + with patch( + "pypaimon.multimodal.connection._table_exists", + return_value=False): + with patch.object( + self.conn, + "get_table", + side_effect=delete_fallback_once): + with self.assertRaisesRegex( + ValueError, "data-evolution.enabled"): + self.conn.create_table( + "fallback_deleted", + schema=schema, + options={"data-evolution.enabled": "false"}, + ignore_if_exists=True, + ) + def test_create_table_can_add_initial_data_and_get_by_short_name(self): self.conn.create_table( "users", From 0d9d05385edb47d3f53b02d740d846986d2ceae7 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Tue, 1 Sep 2026 23:34:55 -0700 Subject: [PATCH 07/32] [python] Preserve initial data across create races --- .../pypaimon/multimodal/connection.py | 6 +++-- .../pypaimon/tests/multimodal_table_test.py | 25 ++++++++++++++++++- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index dad5baeb0c26..72910abe993f 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -90,15 +90,17 @@ def create_table( raise self._create_database_for(identifier) + created = False try: self.catalog.create_table( - identifier, paimon_schema, ignore_if_exists) + identifier, paimon_schema, False) + created = True except TableAlreadyExistException: if not ignore_if_exists: raise table = self.get_table(name) - if data is not None and not already_exists: + if data is not None and created: table.add(data) return table diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index 378f7705264d..8ea88ad18276 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -804,8 +804,14 @@ def delete_once(name): with patch.object( self.conn, "get_table", side_effect=delete_once): table = self.conn.create_table( - "deleted", schema=schema, ignore_if_exists=True) + "deleted", + data=pa.table({"id": [1, 2, 3]}), + schema=schema, + ignore_if_exists=True, + ) self.assertEqual("default.deleted", table.identifier) + self.assertEqual( + [1, 2, 3], table.scan().to_arrow()["id"].to_pylist()) self.conn.create_table("fallback_deleted", schema=schema) deleted[0] = False @@ -833,6 +839,23 @@ def delete_fallback_once(name): ignore_if_exists=True, ) + def test_create_table_does_not_add_data_to_concurrent_winner(self): + schema = _schema({"id": pa.int32()}) + winner = self.conn.create_table("winner", schema=schema) + + with patch( + "pypaimon.multimodal.connection._table_exists", + return_value=False): + actual = self.conn.create_table( + "winner", + data=pa.table({"id": [1, 2, 3]}), + schema=schema, + ignore_if_exists=True, + ) + + self.assertEqual(winner.identifier, actual.identifier) + self.assertEqual([], actual.scan().to_arrow().to_pylist()) + def test_create_table_can_add_initial_data_and_get_by_short_name(self): self.conn.create_table( "users", From aa9932887ff5c68399c59dcb492491de035ee751 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 01:33:20 -0700 Subject: [PATCH 08/32] [python] Align LeRobot versions with table tags --- docs/docs/pypaimon/multimodal-api.mdx | 22 +- paimon-python/README.md | 3 +- .../pypaimon/multimodal/connection.py | 10 +- .../pypaimon/multimodal/lerobot/api.py | 52 +-- .../pypaimon/multimodal/lerobot/loader.py | 3 - .../pypaimon/multimodal/lerobot/metadata.py | 320 ++++++++---------- .../pypaimon/multimodal/lerobot/source.py | 11 + .../pypaimon/tests/multimodal_lerobot_test.py | 145 +++----- 8 files changed, 228 insertions(+), 338 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 5faeded0ab23..9fb38f202387 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -525,9 +525,10 @@ source drift. Calling it again with the same input appends the rows again. `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 creates a LeRobot dataset backed by the frame table, -`
__datasets`, `
__episodes`, and `
__tasks`. A manifest row is +`
__versions`, `
__episodes`, and `
__tasks`. A manifest row is reserved as `PENDING` before the component writes and marked `READY` only after -all component snapshots are committed and tagged. Readers ignore `PENDING`. +all components are committed and tagged with the same numeric `version_id`. +Readers ignore `PENDING`. ```shell pip install 'pypaimon[lerobot]' @@ -542,8 +543,8 @@ print(result.version_id) print(result.frames_snapshot_id) ``` -The result contains the `dataset_id`, `version_id`, and the frame, Episode, and -Task snapshot IDs. Empty components have no snapshot. +The result contains the `version_id` and the frame, Episode, and Task snapshot +IDs. For FileIO URIs, pass credentials through `source_options`: @@ -559,13 +560,12 @@ result = conn.load_from_lerobot( ) ``` -Component rows carry `dataset_id` and stable LeRobot indices. `version_id` -exists only in the Dataset manifest, where it resolves the exact frame, -Episode, and Task snapshots. These snapshots define the published Dataset -state; tags retain them. Set `dataset_id=` to choose a stable ID. The one-time -importer requires a new target table. A failed import removes the table group -so the call can be retried. `drop_table()` removes all four tables; companion -tables cannot be dropped separately through `MultimodalConnection`. +The component tables retain the native LeRobot V3 schemas and contain no +version columns. A READY row in `
__versions` identifies a release; +reading the same tag from the three component tables reconstructs that release. +The one-time importer requires a new target table. A failed import removes the +table group so the call can be retried. `drop_table()` removes all four tables; +companion tables cannot be dropped separately through `MultimodalConnection`. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. diff --git a/paimon-python/README.md b/paimon-python/README.md index 19b2b5009742..ecebb0504127 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -52,7 +52,8 @@ print(result.version_id, result.frames_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. +BLOB columns. The import creates frame, Episode, task, and version tables and +tags the three component tables with `result.version_id`. # HDF5 to multimodal tables diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 72910abe993f..f80ac6487056 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -141,7 +141,6 @@ def load_from_lerobot( source, *, batch_size: int = 1024, - dataset_id: Optional[str] = None, options=None, source_options=None): """Import LeRobot Dataset v3 into a new Paimon table group.""" @@ -151,7 +150,6 @@ def load_from_lerobot( table_name, source, batch_size=batch_size, - dataset_id=dataset_id, options=options, source_options=source_options, ) @@ -161,17 +159,17 @@ def drop_table(self, name: str, ignore_if_not_exists: bool = False): owner_id = None try: from pypaimon.multimodal.lerobot.metadata import ( - _DEFAULT_DATASET_ID_OPTION, _OWNER_ID_OPTION, + _is_managed_root, ) raw_table = self.catalog.get_table(identifier) table_options = raw_table.table_schema.options - if (_OWNER_ID_OPTION in table_options - and _DEFAULT_DATASET_ID_OPTION not in table_options): + managed_root = _is_managed_root(table_options) + if _OWNER_ID_OPTION in table_options and not managed_root: raise ValueError( "%s is a managed LeRobot companion table; drop its " "frame table instead." % identifier) - if _DEFAULT_DATASET_ID_OPTION in table_options: + if managed_root: if Identifier.from_string( identifier).get_branch_name() is not None: raise ValueError( diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 2bfe936ea17a..9fa37e8a2783 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -27,13 +27,11 @@ TableNotExistException, ) from pypaimon.multimodal.lerobot.metadata import ( - _DEFAULT_DATASET_ID_OPTION, _OWNER_ID_OPTION, _drop_import_tables, - _frame_schema, _load_dataset_metadata, _managed_table_options, - _new_id, + _new_owner_id, _prepare_metadata_tables, _publish_dataset, _reject_subtasks, @@ -63,8 +61,7 @@ class LeRobotLoadResult: """Published Paimon state for one imported LeRobot Dataset.""" - dataset_id: str - version_id: str + version_id: int frames_snapshot_id: Optional[int] episodes_snapshot_id: Optional[int] tasks_snapshot_id: Optional[int] @@ -76,17 +73,15 @@ def load_from_lerobot( source, *, batch_size: int = 1024, - dataset_id: Optional[str] = None, options: Optional[Mapping[str, object]] = None, source_options: Optional[Mapping[str, object]] = None, ) -> LeRobotLoadResult: """Import LeRobot Dataset v3 and return its published Paimon state. A new target table is created from LeRobot metadata. Episode, task, and - dataset metadata are stored in companion Paimon tables. ``dataset_id`` - defaults to the target identifier. FileIO URI - credentials come only from ``source_options`` and are not inherited from - the target Catalog. + 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( @@ -126,7 +121,6 @@ def load_from_lerobot( resolved_source, source_schema, batch_size, - dataset_id, options, metadata, ) @@ -151,7 +145,6 @@ def load_from_lerobot( resolved_source, lerobot_schema, batch_size, - dataset_id, options, metadata, ) @@ -169,22 +162,17 @@ def _import_dataset( source, source_schema, batch_size, - dataset_id, options, metadata): table, owner_id = _create_target_table( connection, table_name, source_schema, options) try: - resolved_dataset_id = _resolved_dataset_id(dataset_id, table) tables = _prepare_metadata_tables( - connection, table.raw_table, owner_id) - version_id = _new_id() + connection, table.raw_table, owner_id, metadata) + version_id = 1 _reserve_dataset_version( - tables["datasets"], - resolved_dataset_id, + tables["versions"], version_id, - info, - source, metadata, ) frames_snapshot_id = None @@ -196,22 +184,17 @@ def _import_dataset( source, source_schema, batch_size, - resolved_dataset_id, metadata, ) episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( connection, tables, - resolved_dataset_id, version_id, - info, - source, metadata, table.identifier, frames_snapshot_id, ) return LeRobotLoadResult( - dataset_id=resolved_dataset_id, version_id=version_id, frames_snapshot_id=frames_snapshot_id, episodes_snapshot_id=episodes_snapshot_id, @@ -255,24 +238,9 @@ def _required_count(info, name, source): return int(value) -def _resolved_dataset_id(value, table): - if value is not None and ( - not isinstance(value, str) or not value.strip()): - raise ValueError("dataset_id must be a non-empty string.") - if value is None: - default_id = table.raw_table.table_schema.options.get( - _DEFAULT_DATASET_ID_OPTION) - if not default_id: - raise ValueError( - "LeRobot table %s has no default dataset_id." - % table.identifier) - return default_id - return value.strip() - - def _create_target_table( connection, table_name, source_schema, options): - owner_id = _new_id() + owner_id = _new_owner_id() create_options = dict(options or {}) managed_options = _managed_table_options( connection._identifier(table_name), owner_id) @@ -285,7 +253,7 @@ def _create_target_table( try: table = connection.create_table( table_name, - schema=_frame_schema(source_schema), + schema=source_schema, options=create_options, ) except TableAlreadyExistException as error: diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index ce1dc9345454..5ea9fba715d8 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -25,7 +25,6 @@ from pypaimon.multimodal.arrow_utils import strict_arrow_table from pypaimon.multimodal.hdf5 import _SnapshotRecorder -from pypaimon.multimodal.lerobot.metadata import _with_frame_identity from pypaimon.multimodal.lerobot.schema import _feature_shape from pypaimon.multimodal.table import _target_schema @@ -69,7 +68,6 @@ def _write_dataset( source, source_schema, batch_size, - dataset_id, metadata): target_schema = _target_schema(table.raw_table) write_builder = table.raw_table.new_batch_write_builder() @@ -103,7 +101,6 @@ def _write_dataset( ) observed_tasks.setdefault(episode_index, set()).update( seen_tasks) - batch = _with_frame_identity(batch, dataset_id) batch = _strict_lerobot_table( batch, target_schema, diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 0c4993df5304..1ff0ff26340b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -14,13 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Metadata tables for imported LeRobot datasets.""" +"""LeRobot component tables and version publication.""" -import hashlib import json import numbers import uuid -from datetime import datetime, timezone from pathlib import Path import pyarrow as pa @@ -37,12 +35,10 @@ from pypaimon.multimodal.table import _target_schema -_DATASET_ID = "dataset_id" _VERSION_ID = "version_id" _OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" -_DEFAULT_DATASET_ID_OPTION = "pypaimon.lerobot.dataset-id" _TABLE_SUFFIXES = { - "datasets": "__datasets", + "versions": "__versions", "episodes": "__episodes", "tasks": "__tasks", } @@ -51,106 +47,67 @@ for name in _TABLE_SUFFIXES } -_FRAME_ID_FIELDS = [ - pa.field(_DATASET_ID, pa.string(), nullable=False), -] -_DATASETS_SCHEMA = pa.schema([ - pa.field(_DATASET_ID, pa.string(), nullable=False), - pa.field(_VERSION_ID, pa.string(), nullable=False), - pa.field("parent_version_id", pa.string()), +_VERSIONS_SCHEMA = pa.schema([ + pa.field(_VERSION_ID, pa.int64(), nullable=False), pa.field("status", pa.string(), nullable=False), - pa.field("published_at", pa.timestamp("us", tz="UTC")), - pa.field("format", pa.string(), nullable=False), - pa.field("format_version", pa.string(), nullable=False), - pa.field("fps", pa.int64(), nullable=False), - pa.field("features_json", pa.string(), nullable=False), pa.field("info_json", pa.string(), nullable=False), - pa.field("global_stats_json", pa.string()), - pa.field("total_frames", pa.int64(), nullable=False), - pa.field("total_episodes", pa.int64(), nullable=False), - pa.field("total_tasks", pa.int64(), nullable=False), - pa.field("frames_snapshot_id", pa.int64()), - pa.field("episodes_snapshot_id", pa.int64()), - pa.field("tasks_snapshot_id", pa.int64()), - pa.field("source_uri", pa.string(), nullable=False), - pa.field("metadata_checksum", pa.string(), nullable=False), + pa.field("stats_json", pa.string()), +]) +_EMPTY_TASKS_SCHEMA = pa.schema([ + pa.field("task_index", pa.int64(), nullable=False), + pa.field("task", pa.string(), nullable=False), ]) -_EPISODES_SCHEMA = pa.schema([ - pa.field(_DATASET_ID, 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), - pa.field("task_indices", pa.list_(pa.int64()), nullable=False), - pa.field("split", pa.string()), - pa.field("episode_stats_json", pa.string()), - pa.field("episode_metadata_json", pa.string(), nullable=False), ]) -_TASKS_SCHEMA = pa.schema([ - pa.field(_DATASET_ID, pa.string(), nullable=False), - pa.field("task_index", pa.int64(), nullable=False), - pa.field("task", pa.string(), nullable=False), - pa.field("task_metadata_json", pa.string()), -]) - - -def _frame_schema(source_schema): - reserved = [ - field.name for field in _FRAME_ID_FIELDS - if field.name in source_schema.names - ] - if reserved: - raise ValueError( - "LeRobot features use reserved Paimon fields: %s" % reserved) - return pa.schema(list(source_schema) + _FRAME_ID_FIELDS) - - -def _with_frame_identity(table, dataset_id): - size = table.num_rows - return pa.Table.from_arrays( - list(table.columns) + [ - pa.array([dataset_id] * size, type=pa.string()), - ], - schema=pa.schema(list(table.schema) + _FRAME_ID_FIELDS), - ) +_EPISODE_CONTROL_COLUMNS = [ + "episode_index", + "dataset_from_index", + "dataset_to_index", + "tasks", + "length", +] def _load_dataset_metadata(dataset, info, source): fps = _positive_integer(info.get("fps"), "fps") stats = _source_stats(dataset, source) - task_records = _source_tasks(dataset, source, int(info["total_tasks"])) - tasks, task_indices = _task_rows(task_records, int(info["total_tasks"])) + tasks_table = _source_tasks( + dataset, source, int(info["total_tasks"])) + tasks, task_indices = _task_rows( + tasks_table.to_pylist(), int(info["total_tasks"])) total_episodes = int(info["total_episodes"]) - episode_records = [] if total_episodes == 0 \ - else _source_episodes(dataset, source) + episode_source = ( + _source_episodes(dataset, source) + if total_episodes > 0 + else {"paths": [], "rows": [], "schema": _EMPTY_EPISODES_SCHEMA} + ) episodes = _episode_rows( - episode_records, + episode_source["rows"], task_indices, info, int(info["total_frames"]), total_episodes, ) - canonical = { - "info": info, - "stats": stats, - "episodes": episodes, - "tasks": tasks, - } return { "fps": fps, - "features_json": _canonical_json(info["features"]), "info_json": _canonical_json(info), - "global_stats_json": ( + "stats_json": ( None if stats is None else _canonical_json(stats)), "episodes": episodes, "tasks": tasks, - "metadata_checksum": "sha256:" + hashlib.sha256( - _canonical_json(canonical).encode("utf-8") - ).hexdigest(), + "episodes_schema": episode_source["schema"], + "episode_paths": episode_source["paths"], + "tasks_table": tasks_table, + "source": source, } -def _new_id(): +def _new_owner_id(): return uuid.uuid4().hex @@ -172,16 +129,18 @@ def _managed_table_options(frames_identifier, owner_id): if identifier.get_branch_name() is not None: raise ValueError( "LeRobot import does not support table branches.") - result = { - _OWNER_ID_OPTION: owner_id, - _DEFAULT_DATASET_ID_OPTION: str(frames_identifier), - } + result = {_OWNER_ID_OPTION: owner_id} for name, suffix in _TABLE_SUFFIXES.items(): result[_COMPANION_OPTION_KEYS[name]] = _companion_identifier( frames_identifier, suffix) return result +def _is_managed_root(options): + return _OWNER_ID_OPTION in options and all( + key in options for key in _COMPANION_OPTION_KEYS.values()) + + def _companion_table_identifiers(frames_table): options = frames_table.table_schema.options identifiers = {} @@ -195,11 +154,11 @@ def _companion_table_identifiers(frames_table): return identifiers -def _prepare_metadata_tables(connection, frames_table, owner_id): +def _prepare_metadata_tables(connection, frames_table, owner_id, metadata): schemas = { - "datasets": _DATASETS_SCHEMA, - "episodes": _EPISODES_SCHEMA, - "tasks": _TASKS_SCHEMA, + "versions": _VERSIONS_SCHEMA, + "episodes": metadata["episodes_schema"], + "tasks": metadata["tasks_table"].schema, } identifiers = _companion_table_identifiers(frames_table) tables = {} @@ -240,28 +199,13 @@ def _prepare_metadata_tables(connection, frames_table, owner_id): def _reserve_dataset_version( - datasets_table, - dataset_id, + versions_table, version_id, - info, - source, metadata): - pending = _manifest_row( - dataset_id, - version_id, - None, - "PENDING", - None, - info, - source, - metadata, - None, - None, - None, - ) + pending = _manifest_row(version_id, "PENDING", metadata) snapshot_id = _append_arrow( - datasets_table, - pa.Table.from_pylist([pending], schema=_DATASETS_SCHEMA), + versions_table, + pa.Table.from_pylist([pending], schema=_VERSIONS_SCHEMA), ) if snapshot_id is None: raise RuntimeError("LeRobot version reservation created no snapshot.") @@ -270,96 +214,47 @@ def _reserve_dataset_version( def _publish_dataset( connection, tables, - dataset_id, version_id, - info, - source, metadata, frames_identifier, frames_snapshot_id): - episodes = _dataset_table( - metadata["episodes"], - _EPISODES_SCHEMA, - dataset_id, + episodes_snapshot_id = _append_arrow_tables( + tables["episodes"], + _source_episode_tables(metadata), ) - tasks = _dataset_table( - metadata["tasks"], - _TASKS_SCHEMA, - dataset_id, - ) - episodes_snapshot_id = _append_arrow(tables["episodes"], episodes) - tasks_snapshot_id = _append_arrow(tables["tasks"], tasks) + tasks_snapshot_id = _append_arrow( + tables["tasks"], metadata["tasks_table"]) - tag = "pypaimon-lerobot-%s" % version_id + if None in ( + frames_snapshot_id, episodes_snapshot_id, tasks_snapshot_id): + raise ValueError( + "LeRobot tag-backed import requires non-empty frame, Episode, " + "and task components.") + tag = str(version_id) for identifier, snapshot_id in ( (frames_identifier, frames_snapshot_id), (tables["episodes"].identifier, episodes_snapshot_id), (tables["tasks"].identifier, tasks_snapshot_id)): - if snapshot_id is not None: - _create_tag(connection.catalog, identifier, tag, snapshot_id) + _create_tag(connection.catalog, identifier, tag, snapshot_id) - manifest = _manifest_row( - dataset_id, - version_id, - None, - "READY", - datetime.now(timezone.utc), - info, - source, - metadata, - frames_snapshot_id, - episodes_snapshot_id, - tasks_snapshot_id, - ) - _append_arrow(tables["datasets"], pa.Table.from_pylist( - [manifest], schema=_DATASETS_SCHEMA)) + manifest = _manifest_row(version_id, "READY", metadata) + _append_arrow(tables["versions"], pa.Table.from_pylist( + [manifest], schema=_VERSIONS_SCHEMA)) return episodes_snapshot_id, tasks_snapshot_id def _manifest_row( - dataset_id, version_id, - parent_version_id, status, - published_at, - info, - source, - metadata, - frames_snapshot_id, - episodes_snapshot_id, - tasks_snapshot_id): + metadata): return { - _DATASET_ID: dataset_id, _VERSION_ID: version_id, - "parent_version_id": parent_version_id, "status": status, - "published_at": published_at, - "format": "lerobot", - "format_version": str(info["codebase_version"]), - "fps": metadata["fps"], - "features_json": metadata["features_json"], "info_json": metadata["info_json"], - "global_stats_json": metadata["global_stats_json"], - "total_frames": int(info["total_frames"]), - "total_episodes": int(info["total_episodes"]), - "total_tasks": int(info["total_tasks"]), - "frames_snapshot_id": frames_snapshot_id, - "episodes_snapshot_id": episodes_snapshot_id, - "tasks_snapshot_id": tasks_snapshot_id, - "source_uri": str(source.path), - "metadata_checksum": metadata["metadata_checksum"], + "stats_json": metadata["stats_json"], } -def _dataset_table(rows, schema, dataset_id): - values = [] - for row in rows: - value = dict(row) - value[_DATASET_ID] = dataset_id - values.append(value) - return pa.Table.from_pylist(values, schema=schema) - - def _drop_import_tables(catalog, frames_table, owner_id): identifiers = list( _companion_table_identifiers(frames_table).values()) @@ -374,8 +269,10 @@ def _drop_import_tables(catalog, frames_table, owner_id): def _append_arrow(table, data): - if data.num_rows == 0: - return None + return _append_arrow_tables(table, [data]) + + +def _append_arrow_tables(table, tables): builder = table.new_batch_write_builder() table_write = builder.new_write() table_commit = builder.new_commit() @@ -383,7 +280,20 @@ def _append_arrow(table, data): recorder = _SnapshotRecorder() table_commit.add_commit_callback(recorder) try: - table_write.write_arrow(data) + 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 + if row_count == 0: + table_write.abort() + return None messages = table_write.prepare_commit() commit_started = True table_commit.commit(messages) @@ -432,17 +342,17 @@ def _source_stats(dataset, source): def _source_tasks(dataset, source, total_tasks): if total_tasks == 0: - return [] + 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).to_pylist() + return _read_remote_parquet(source.file_io, path) path = _metadata_root(dataset, source) / "meta" / "tasks.parquet" try: - return pq.read_table(path).to_pylist() + return pq.read_table(path) except (OSError, ValueError, pa.ArrowException) as error: raise ValueError( "Cannot read LeRobot task metadata %s: %s" % (path, error) @@ -453,29 +363,67 @@ def _source_episodes(dataset, source): if source.file_io is not None: from pypaimon.multimodal.lerobot.source import ( _read_remote_parquet, + _read_remote_parquet_schema, _remote_parquet_files, _remote_path, ) directory = _remote_path(source.path, "meta/episodes") paths = _remote_parquet_files(source.file_io, directory) - tables = [ - _read_remote_parquet(source.file_io, path) for path in paths - ] + + def read(path, columns=None): + return _read_remote_parquet( + source.file_io, path, columns=columns) + + 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")) - try: - tables = [pq.read_table(path) for path in paths] - except (OSError, ValueError, pa.ArrowException) as error: - raise ValueError( - "Cannot read LeRobot Episode metadata %s: %s" - % (directory, error)) from error + + def read(path, columns=None): + return pq.read_table(path, columns=columns) + + read_schema = pq.read_schema + if not paths: + return { + "paths": [], + "rows": [], + "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.") + tables = [read(path, columns=_EPISODE_CONTROL_COLUMNS) + for path in paths] + except (OSError, ValueError, pa.ArrowException) as error: + raise ValueError( + "Cannot read LeRobot Episode metadata %s: %s" + % (directory, error)) from error rows = [] for table in tables: rows.extend(table.to_pylist()) rows.sort(key=lambda row: _integer( row.get("episode_index"), "episode_index")) - return rows + return {"paths": paths, "rows": rows, "schema": schema} + + +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 _reject_subtasks(dataset, source): diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index af03227bc8c8..b7fa703dd903 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -466,6 +466,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) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 2135b524c0fd..9dcc8a09faa9 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -56,6 +56,7 @@ _remote_source_path, _validate_info_paths, ) +from pypaimon.multimodal.table import _target_schema try: from lerobot.datasets.lerobot_dataset import LeRobotDataset @@ -350,7 +351,7 @@ def test_remote_episode_metadata_projects_stats_columns(self): columns=_RemoteLeRobotDataset._EPISODE_COLUMNS, ) - 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" @@ -371,25 +372,11 @@ def test_empty_local_dataset_returns_before_opening_lerobot(self): with patch( "pypaimon.multimodal.lerobot.api._import_lerobot_dataset" ) as import_lerobot: - result = 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() - self.assertEqual("default.empty_frames", result.dataset_id) - self.assertEqual(32, len(result.version_id)) - self.assertIsNone(result.frames_snapshot_id) - self.assertIsNone(result.episodes_snapshot_id) - self.assertIsNone(result.tasks_snapshot_id) - table = connection.get_table("empty_frames") - self.assertIsNone( - table.raw_table.snapshot_manager().get_latest_snapshot()) - manifests = _catalog_rows( - connection, "empty_frames__datasets") - self.assertEqual(["PENDING", "READY"], [ - row["status"] for row in manifests - ]) - self.assertEqual(result.version_id, manifests[1]["version_id"]) - self.assertEqual(0, manifests[1]["total_frames"]) - self.assertIsNone(manifests[1]["frames_snapshot_id"]) + with self.assertRaises(TableNotExistException): + connection.get_table("empty_frames") finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -437,7 +424,7 @@ def test_empty_fast_path_validates_required_counts(self): finally: shutil.rmtree(temp_dir, ignore_errors=True) - def test_empty_dataset_preserves_tasks_and_stats(self): + def test_empty_dataset_with_tasks_is_rejected(self): temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_empty_meta_")) try: source = temp_dir / "source" @@ -465,15 +452,10 @@ def test_empty_dataset_preserves_tasks_and_stats(self): "warehouse": str(temp_dir / "warehouse"), }) - result = connection.load_from_lerobot("frames", source) - self.assertIsNone(result.frames_snapshot_id) - self.assertEqual(1, result.tasks_snapshot_id) - manifest = _catalog_rows(connection, "frames__datasets")[1] - self.assertIsNotNone(manifest["global_stats_json"]) - self.assertEqual(1, manifest["total_tasks"]) - self.assertEqual(1, manifest["tasks_snapshot_id"]) - self.assertEqual( - "pick", _catalog_rows(connection, "frames__tasks")[0]["task"]) + 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) @@ -767,7 +749,7 @@ def test_import_infers_schema_and_preserves_episodes(self): self.assertEqual(1, result.frames_snapshot_id) self.assertEqual(1, result.episodes_snapshot_id) self.assertEqual(1, result.tasks_snapshot_id) - self.assertEqual(32, len(result.version_id)) + self.assertEqual(1, result.version_id) table = self.connection.get_table("robot_data") schema = table.raw_table.fields @@ -781,7 +763,7 @@ def test_import_infers_schema_and_preserves_episodes(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.assertEqual("STRING NOT NULL", types["dataset_id"]) + self.assertNotIn("dataset_id", types) self.assertNotIn("metadata_version", types) self.assertNotIn("version_id", types) @@ -796,7 +778,6 @@ def test_import_infers_schema_and_preserves_episodes(self): "observation.matrix", "action", "reward", - "dataset_id", ]).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]) @@ -808,30 +789,19 @@ def test_import_infers_schema_and_preserves_episodes(self): rows[4]["observation.matrix"]) self.assertAlmostEqual(0.2, rows[4]["timestamp"], places=6) self.assertEqual(1.0, rows[4]["reward"]) - self.assertEqual( - [str(table.identifier)] * 5, - [row["dataset_id"] for row in rows], - ) - manifests = _catalog_rows(self.connection, "robot_data__datasets") + manifests = _catalog_rows(self.connection, "robot_data__versions") self.assertEqual(["PENDING", "READY"], [ row["status"] for row in manifests ]) manifest = manifests[1] - self.assertEqual(str(table.identifier), manifest["dataset_id"]) self.assertEqual(result.version_id, manifest["version_id"]) - self.assertIsNone(manifest["parent_version_id"]) - self.assertIsNotNone(manifest["published_at"]) - self.assertEqual( - result.frames_snapshot_id, manifest["frames_snapshot_id"]) + self.assertEqual("v3.0", json.loads( + manifest["info_json"])["codebase_version"]) + self.assertIsNotNone(manifest["stats_json"]) self.assertEqual( - result.episodes_snapshot_id, manifest["episodes_snapshot_id"]) - self.assertEqual( - result.tasks_snapshot_id, manifest["tasks_snapshot_id"]) - self.assertEqual("lerobot", manifest["format"]) - self.assertEqual("v3.0", manifest["format_version"]) - self.assertIsNotNone(manifest["global_stats_json"]) - self.assertTrue(manifest["metadata_checksum"].startswith("sha256:")) - tag = "pypaimon-lerobot-%s" % manifest["version_id"] + {"version_id", "status", "info_json", "stats_json"}, + set(manifest)) + tag = str(manifest["version_id"]) self.assertEqual( result.frames_snapshot_id, self.connection.catalog.get_tag( @@ -853,15 +823,21 @@ def test_import_infers_schema_and_preserves_episodes(self): "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([[0], [1]], [row["task_indices"] for row in episodes]) - self.assertEqual(["train", "train"], [row["split"] for row in episodes]) - self.assertTrue(all( - row["episode_stats_json"] is not None 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 = { @@ -869,8 +845,16 @@ def test_import_infers_schema_and_preserves_episodes(self): 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"]) for row in tasks + (row["task_index"], row[task_name]) for row in tasks ]) self.assertEqual( result.frames_snapshot_id, @@ -925,24 +909,11 @@ def test_frame_controls_must_match_published_episode_metadata(self): with self.assertRaisesRegex( ValueError, "has %s" % column): self.connection.load_from_lerobot(table_name, source) - for suffix in ("", "__datasets", "__episodes", "__tasks"): + for suffix in ("", "__versions", "__episodes", "__tasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier(table_name + suffix)) - def test_explicit_dataset_id_is_shared_by_all_tables(self): - dataset_id = "aloha_pick_cube@3" - self.connection.load_from_lerobot( - "custom_id", self.image_source, dataset_id=dataset_id) - - for name in ( - "custom_id", "custom_id__datasets", - "custom_id__episodes", "custom_id__tasks"): - rows = _catalog_rows(self.connection, name) - self.assertEqual({dataset_id}, { - row["dataset_id"] for row in rows - }) - def test_frame_task_uses_published_task_mapping(self): source = self.temp_dir / "reordered_tasks" shutil.copytree(self.image_source, source) @@ -954,10 +925,12 @@ def test_frame_task_uses_published_task_mapping(self): frames = self.connection.get_table("reordered_tasks").scan().select([ "index", "task_index", "task" ]).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"] - for row in _catalog_rows( - self.connection, "reordered_tasks__tasks") + row["task_index"]: row[task_name] for row in task_rows } self.assertTrue(all( row["task"] == published[row["task_index"]] @@ -1068,14 +1041,8 @@ def test_empty_oss_source_does_not_require_episode_directory(self): with patch( "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): - result = self.connection.load_from_lerobot( - "empty_oss", source) - self.assertIsNone(result.frames_snapshot_id) - self.assertEqual( - 0, - _catalog_rows( - self.connection, "empty_oss__datasets")[0]["total_episodes"], - ) + 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( @@ -1086,8 +1053,8 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): "tag_fallback", self.image_source) manifest = _catalog_rows( - self.connection, "tag_fallback__datasets")[1] - tag = "pypaimon-lerobot-%s" % manifest["version_id"] + self.connection, "tag_fallback__versions")[1] + tag = str(manifest["version_id"]) table = self.connection.get_table("tag_fallback") self.assertEqual( result.frames_snapshot_id, @@ -1102,7 +1069,7 @@ def test_failed_publication_is_cleaned_and_can_retry(self): self.connection.load_from_lerobot( "failed_publish", self.image_source) - for suffix in ("", "__datasets", "__episodes", "__tasks"): + for suffix in ("", "__versions", "__episodes", "__tasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier("failed_publish" + suffix)) @@ -1192,7 +1159,7 @@ def open_with_failing_close(*args, **kwargs): self.assertEqual( ["PENDING", "READY"], [row["status"] for row in _catalog_rows( - self.connection, "close_failure__datasets")], + self.connection, "close_failure__versions")], ) def test_source_close_failure_does_not_override_success(self): @@ -1212,7 +1179,7 @@ def test_source_close_failure_does_not_override_success(self): self.assertEqual( ["PENDING", "READY"], [row["status"] for row in _catalog_rows( - self.connection, "source_close_failure__datasets")], + self.connection, "source_close_failure__versions")], ) def test_existing_target_is_rejected(self): @@ -1271,7 +1238,7 @@ def test_drop_table_removes_companion_tables(self): self.connection.drop_table("drop_group") for name in ( - "drop_group", "drop_group__datasets", + "drop_group", "drop_group__versions", "drop_group__episodes", "drop_group__tasks"): with self.subTest(name=name): with self.assertRaises(TableNotExistException): @@ -1300,7 +1267,7 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.drop_table("retry_drop") for name in ( - "retry_drop", "retry_drop__datasets", + "retry_drop", "retry_drop__versions", "retry_drop__episodes", "retry_drop__tasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( @@ -1327,7 +1294,7 @@ def test_drop_table_rejects_managed_branch(self): with self.assertRaisesRegex(ValueError, "table branch"): self.connection.drop_table("branch_drop$branch_dev") for name in ( - "branch_drop", "branch_drop__datasets", + "branch_drop", "branch_drop__versions", "branch_drop__episodes", "branch_drop__tasks"): self.connection.catalog.get_table( self.connection._identifier(name)) @@ -1345,7 +1312,7 @@ def test_table_group_survives_frame_table_rename(self): self.connection.drop_table("after_rename") for name in ( - "after_rename", "before_rename__datasets", + "after_rename", "before_rename__versions", "before_rename__episodes", "before_rename__tasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( From 222aa3eb2cf6fe5a7ab57ae67a96c1eb0f343cc7 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 02:33:46 -0700 Subject: [PATCH 09/32] [python] Return LeRobot version ID --- docs/docs/pypaimon/multimodal-api.mdx | 11 ++--- paimon-python/README.md | 6 +-- paimon-python/pypaimon/multimodal/__init__.py | 2 - .../pypaimon/multimodal/lerobot/__init__.py | 6 +-- .../pypaimon/multimodal/lerobot/api.py | 24 ++-------- .../pypaimon/multimodal/lerobot/metadata.py | 1 - .../pypaimon/tests/multimodal_lerobot_test.py | 47 +++++++++---------- 7 files changed, 35 insertions(+), 62 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 9fb38f202387..6867ac211d47 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -535,21 +535,20 @@ pip install 'pypaimon[lerobot]' ``` ```python -result = conn.load_from_lerobot( +version_id = conn.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(result.version_id) -print(result.frames_snapshot_id) +print(version_id) ``` -The result contains the `version_id` and the frame, Episode, and Task snapshot -IDs. +The returned `version_id` is the common tag name for the frame, Episode, and +Task tables. For FileIO URIs, pass credentials through `source_options`: ```python -result = conn.load_from_lerobot( +version_id = conn.load_from_lerobot( "robot_data", "oss://source-bucket/lerobot_dataset", source_options={ diff --git a/paimon-python/README.md b/paimon-python/README.md index ecebb0504127..cb51263a67c7 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -44,16 +44,16 @@ pip install 'pypaimon[lerobot]' import pypaimon.multimodal as pmm connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) -result = connection.load_from_lerobot( +version_id = connection.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(result.version_id, result.frames_snapshot_id) +print(version_id) ``` The 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 `result.version_id`. +tags the three component tables with the returned `version_id`. # HDF5 to multimodal tables diff --git a/paimon-python/pypaimon/multimodal/__init__.py b/paimon-python/pypaimon/multimodal/__init__.py index 287b672a0231..b669267b41da 100644 --- a/paimon-python/pypaimon/multimodal/__init__.py +++ b/paimon-python/pypaimon/multimodal/__init__.py @@ -29,7 +29,6 @@ Hdf5File, Hdf5LoadResult, ) -from pypaimon.multimodal.lerobot import LeRobotLoadResult from pypaimon.multimodal.table import ( MultimodalTable, TextRoute, @@ -52,7 +51,6 @@ "BlobStore", "Hdf5File", "Hdf5LoadResult", - "LeRobotLoadResult", "MultimodalConnection", "MultimodalTable", "NoSuchKey", diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index d00c91307057..a40f2a8ccef0 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -16,13 +16,9 @@ """One-time LeRobot Dataset v3 import into a multimodal Paimon table.""" -from pypaimon.multimodal.lerobot.api import ( - LeRobotLoadResult, - load_from_lerobot, -) +from pypaimon.multimodal.lerobot.api import load_from_lerobot __all__ = [ - "LeRobotLoadResult", "load_from_lerobot", ] diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 9fa37e8a2783..687874f757ab 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -18,7 +18,6 @@ import numbers import sys -from dataclasses import dataclass from typing import Mapping, Optional from pypaimon.catalog.catalog_exception import ( @@ -57,16 +56,6 @@ ) -@dataclass(frozen=True) -class LeRobotLoadResult: - """Published Paimon state for one imported LeRobot Dataset.""" - - version_id: int - frames_snapshot_id: Optional[int] - episodes_snapshot_id: Optional[int] - tasks_snapshot_id: Optional[int] - - def load_from_lerobot( connection, table_name: str, @@ -75,8 +64,8 @@ def load_from_lerobot( batch_size: int = 1024, options: Optional[Mapping[str, object]] = None, source_options: Optional[Mapping[str, object]] = None, -) -> LeRobotLoadResult: - """Import LeRobot Dataset v3 and return its published Paimon state. +) -> 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. @@ -186,7 +175,7 @@ def _import_dataset( batch_size, metadata, ) - episodes_snapshot_id, tasks_snapshot_id = _publish_dataset( + _publish_dataset( connection, tables, version_id, @@ -194,12 +183,7 @@ def _import_dataset( table.identifier, frames_snapshot_id, ) - return LeRobotLoadResult( - version_id=version_id, - frames_snapshot_id=frames_snapshot_id, - episodes_snapshot_id=episodes_snapshot_id, - tasks_snapshot_id=tasks_snapshot_id, - ) + return version_id except BaseException as error: try: _drop_import_tables( diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 1ff0ff26340b..21cd383b624b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -240,7 +240,6 @@ def _publish_dataset( manifest = _manifest_row(version_id, "READY", metadata) _append_arrow(tables["versions"], pa.Table.from_pylist( [manifest], schema=_VERSIONS_SCHEMA)) - return episodes_snapshot_id, tasks_snapshot_id def _manifest_row( diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 9dcc8a09faa9..3496036e2409 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -742,14 +742,10 @@ def _create_image_dataset(root): dataset.finalize() def test_import_infers_schema_and_preserves_episodes(self): - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "robot_data", self.image_source, batch_size=2) - self.assertIsInstance(result, pmm.LeRobotLoadResult) - self.assertEqual(1, result.frames_snapshot_id) - self.assertEqual(1, result.episodes_snapshot_id) - self.assertEqual(1, result.tasks_snapshot_id) - self.assertEqual(1, result.version_id) + self.assertEqual(1, version_id) table = self.connection.get_table("robot_data") schema = table.raw_table.fields @@ -794,7 +790,7 @@ def test_import_infers_schema_and_preserves_episodes(self): row["status"] for row in manifests ]) manifest = manifests[1] - self.assertEqual(result.version_id, manifest["version_id"]) + self.assertEqual(version_id, manifest["version_id"]) self.assertEqual("v3.0", json.loads( manifest["info_json"])["codebase_version"]) self.assertIsNotNone(manifest["stats_json"]) @@ -803,7 +799,7 @@ def test_import_infers_schema_and_preserves_episodes(self): set(manifest)) tag = str(manifest["version_id"]) self.assertEqual( - result.frames_snapshot_id, + 1, self.connection.catalog.get_tag( table.identifier, tag).snapshot.id, ) @@ -857,7 +853,7 @@ def test_import_infers_schema_and_preserves_episodes(self): (row["task_index"], row[task_name]) for row in tasks ]) self.assertEqual( - result.frames_snapshot_id, + 1, table.raw_table.snapshot_manager().get_latest_snapshot().id, ) @@ -992,13 +988,13 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): with patch( "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "oss_images", source, batch_size=2, ) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) table = self.connection.get_table("oss_images") rows = table.scan().select([ "episode_index", "frame_index", "index", "task" @@ -1049,15 +1045,16 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): self.connection.catalog, "create_tag", side_effect=NotImplementedError): - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "tag_fallback", self.image_source) manifest = _catalog_rows( self.connection, "tag_fallback__versions")[1] tag = str(manifest["version_id"]) table = self.connection.get_table("tag_fallback") + self.assertEqual(1, version_id) self.assertEqual( - result.frames_snapshot_id, + table.raw_table.snapshot_manager().get_latest_snapshot().id, table.raw_table.tag_manager().get(tag).id, ) @@ -1074,9 +1071,9 @@ def test_failed_publication_is_cleaned_and_can_retry(self): self.connection.catalog.get_table( self.connection._identifier("failed_publish" + suffix)) - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "failed_publish", self.image_source) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) def test_invalid_target_options_do_not_leave_table(self): with self.assertRaisesRegex(ValueError, "data-evolution.enabled"): @@ -1089,9 +1086,9 @@ def test_invalid_target_options_do_not_leave_table(self): self.connection.catalog.get_table( self.connection._identifier("invalid_options")) - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "invalid_options", self.image_source) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) def test_target_open_failure_is_cleaned_and_can_retry(self): original_get = self.connection.get_table @@ -1108,10 +1105,10 @@ def fail_once(name): with self.assertRaisesRegex(RuntimeError, "get failed"): self.connection.load_from_lerobot( "failed_open", self.image_source) - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "failed_open", self.image_source) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) def test_target_open_failure_does_not_drop_another_owner(self): original_create = self.connection.create_table @@ -1152,10 +1149,10 @@ def open_with_failing_close(*args, **kwargs): api, "_open_resolved_dataset", side_effect=open_with_failing_close): - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "close_failure", self.image_source) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) self.assertEqual( ["PENDING", "READY"], [row["status"] for row in _catalog_rows( @@ -1172,10 +1169,10 @@ def test_source_close_failure_does_not_override_success(self): with patch( "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): - result = self.connection.load_from_lerobot( + version_id = self.connection.load_from_lerobot( "source_close_failure", source) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) self.assertEqual( ["PENDING", "READY"], [row["status"] for row in _catalog_rows( @@ -1224,9 +1221,9 @@ def reserve_then_wait(*args, **kwargs): "concurrent", self.image_source) finally: release.set() - result = future.result(timeout=30) + version_id = future.result(timeout=30) - self.assertEqual(1, result.frames_snapshot_id) + self.assertEqual(1, version_id) self.assertEqual( 5, self.connection.get_table( From dbc09220f09deed78187755ecd78cfeb92c74c80 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 03:21:55 -0700 Subject: [PATCH 10/32] [python] Guard initial LeRobot publication --- .../pypaimon/multimodal/lerobot/metadata.py | 23 +++++--- .../pypaimon/tests/multimodal_lerobot_test.py | 55 ++++++++++++++++++- 2 files changed, 68 insertions(+), 10 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 21cd383b624b..99e9127e30fc 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -207,8 +207,7 @@ def _reserve_dataset_version( versions_table, pa.Table.from_pylist([pending], schema=_VERSIONS_SCHEMA), ) - if snapshot_id is None: - raise RuntimeError("LeRobot version reservation created no snapshot.") + _require_initial_snapshot("versions", snapshot_id) def _publish_dataset( @@ -218,18 +217,15 @@ def _publish_dataset( metadata, frames_identifier, frames_snapshot_id): + _require_initial_snapshot("frames", frames_snapshot_id) episodes_snapshot_id = _append_arrow_tables( tables["episodes"], _source_episode_tables(metadata), ) + _require_initial_snapshot("episodes", episodes_snapshot_id) tasks_snapshot_id = _append_arrow( tables["tasks"], metadata["tasks_table"]) - - if None in ( - frames_snapshot_id, episodes_snapshot_id, tasks_snapshot_id): - raise ValueError( - "LeRobot tag-backed import requires non-empty frame, Episode, " - "and task components.") + _require_initial_snapshot("tasks", tasks_snapshot_id) tag = str(version_id) for identifier, snapshot_id in ( (frames_identifier, frames_snapshot_id), @@ -242,6 +238,17 @@ def _publish_dataset( [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, status, diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 3496036e2409..c5f144ca4e1d 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1035,7 +1035,7 @@ def test_empty_oss_source_does_not_require_episode_directory(self): source_file_io = _RemoteLeRobotFileIO(local_source, source) with patch( - "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", + "pypaimon.multimodal.lerobot.source._SourceFileIO", return_value=source_file_io): with self.assertRaisesRegex(ValueError, "non-empty"): self.connection.load_from_lerobot("empty_oss", source) @@ -1167,7 +1167,7 @@ def test_source_close_failure_does_not_override_success(self): with self.assertLogs( "pypaimon.multimodal.lerobot.source", level="WARNING"): with patch( - "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", + "pypaimon.multimodal.lerobot.source._SourceFileIO", return_value=source_file_io): version_id = self.connection.load_from_lerobot( "source_close_failure", source) @@ -1230,6 +1230,57 @@ def reserve_then_wait(*args, **kwargs): "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): + task_names = { + row["task_index"]: row["task"] + for row in metadata["tasks"] + } + table.add(_read_batch( + dataset, + info, + 0, + 1, + source_schema, + task_names, + )) + 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) + + for name in ( + "concurrent_append", + "concurrent_append__versions", + "concurrent_append__episodes", + "concurrent_append__tasks"): + with self.subTest(name=name): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(name)) + def test_drop_table_removes_companion_tables(self): self.connection.load_from_lerobot("drop_group", self.image_source) self.connection.drop_table("drop_group") From a8df2a541dbdbd7b90f1ea6fa79b27b016e8c577 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 07:07:55 -0700 Subject: [PATCH 11/32] [python] Align LeRobot task and subtask metadata --- docs/docs/pypaimon/multimodal-api.mdx | 12 +- .../pypaimon/multimodal/lerobot/api.py | 16 +- .../pypaimon/multimodal/lerobot/loader.py | 37 +--- .../pypaimon/multimodal/lerobot/metadata.py | 74 ++++--- .../pypaimon/multimodal/lerobot/schema.py | 17 +- .../pypaimon/multimodal/lerobot/source.py | 5 - .../pypaimon/tests/multimodal_lerobot_test.py | 184 +++++++++++++----- 7 files changed, 210 insertions(+), 135 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 6867ac211d47..c1a84396b0d6 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -525,8 +525,9 @@ source drift. Calling it again with the same input appends the rows again. `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 creates a LeRobot dataset backed by the frame table, -`
__versions`, `
__episodes`, and `
__tasks`. A manifest row is -reserved as `PENDING` before the component writes and marked `READY` only after +`
__versions`, `
__episodes`, `
__tasks`, and an optional +`
__subtasks`. Task text remains in the metadata table; frames retain +`task_index`. A manifest row is reserved as `PENDING` and marked `READY` after all components are committed and tagged with the same numeric `version_id`. Readers ignore `PENDING`. @@ -542,8 +543,7 @@ version_id = conn.load_from_lerobot( print(version_id) ``` -The returned `version_id` is the common tag name for the frame, Episode, and -Task tables. +The returned `version_id` is the common tag name for all dataset components. For FileIO URIs, pass credentials through `source_options`: @@ -561,9 +561,9 @@ version_id = conn.load_from_lerobot( The component tables retain the native LeRobot V3 schemas and contain no version columns. A READY row in `
__versions` identifies a release; -reading the same tag from the three component tables reconstructs that release. +reading the same tag from its component tables reconstructs that release. The one-time importer requires a new target table. A failed import removes the -table group so the call can be retried. `drop_table()` removes all four tables; +table group so the call can be retried. `drop_table()` removes the table group; companion tables cannot be dropped separately through `MultimodalConnection`. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 687874f757ab..2f6205ce36b7 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -33,7 +33,6 @@ _new_owner_id, _prepare_metadata_tables, _publish_dataset, - _reject_subtasks, _reserve_dataset_version, ) from pypaimon.multimodal.lerobot.loader import _write_dataset @@ -43,7 +42,6 @@ ) from pypaimon.multimodal.lerobot.source import ( _close_quietly, - _has_tasks, _import_lerobot_dataset, _load_hub_info, _open_resolved_dataset, @@ -89,17 +87,13 @@ 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 = \ + _schema_from_info(local_info) + total_frames, _, _ = \ _validated_counts(local_info, resolved_source.path) if total_frames == 0 and ( resolved_source.root is not None or resolved_source.file_io is not None): - source_schema = _schema_from_info( - local_info, - include_task=total_tasks > 0, - ) - _reject_subtasks(None, resolved_source) + source_schema = _schema_from_info(local_info) metadata = _load_dataset_metadata( None, local_info, resolved_source) return _import_dataset( @@ -121,9 +115,7 @@ def load_from_lerobot( _require_v3(info, resolved_source.path) _validated_counts(info, resolved_source.path) - lerobot_schema = _schema_from_info( - info, include_task=_has_tasks(dataset, info)) - _reject_subtasks(dataset, resolved_source) + lerobot_schema = _schema_from_info(info) metadata = _load_dataset_metadata( dataset, info, resolved_source) return _import_dataset( diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index 5ea9fba715d8..dd95056c58e8 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -77,9 +77,6 @@ def _write_dataset( batch_count = 0 row_count = 0 episodes = metadata["episodes"] - task_names = { - row["task_index"]: row["task"] for row in metadata["tasks"] - } observed_tasks = {} snapshot_recorder = _SnapshotRecorder() @@ -90,7 +87,7 @@ def _write_dataset( for episode_index, episode_begin, task_indices, begin, end in \ _episode_batches(dataset, info, batch_size, episodes): batch = _read_batch( - dataset, info, begin, end, source_schema, task_names) + dataset, info, begin, end, source_schema) seen_tasks = _validate_frame_controls( batch, int(info["fps"]), @@ -213,14 +210,18 @@ def _validate_frame_controls( 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), frame_index / fps, + 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, frame_index / fps)) + % (index, timestamp, expected_timestamp)) task_index = _control_integer( values["task_index"][offset], "task_index", index) if task_index not in allowed_tasks: @@ -262,7 +263,7 @@ def _control_integer(value, name, frame_index): return int(value) -def _read_batch(dataset, info, begin, end, schema, task_names=None): +def _read_batch(dataset, info, begin, end, schema): read_batch = getattr(dataset, "read_batch", None) if callable(read_batch): raw = read_batch(begin, end) @@ -296,14 +297,6 @@ def _read_batch(dataset, info, begin, end, schema, task_names=None): arrays.append(_safe_array(values, field, name, dtype)) fields.append(field) - if "task" in schema.names: - task_indices = raw.column("task_index").to_pylist() - tasks = dataset.meta.tasks if task_names is None else task_names - arrays.append(pa.array( - [_task_name(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)) @@ -471,17 +464,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 index 99e9127e30fc..3dd527f41cee 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -41,6 +41,7 @@ "versions": "__versions", "episodes": "__episodes", "tasks": "__tasks", + "subtasks": "__subtasks", } _COMPANION_OPTION_KEYS = { name: "pypaimon.lerobot.%s-table" % name @@ -52,6 +53,7 @@ pa.field("status", pa.string(), 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), @@ -78,8 +80,9 @@ def _load_dataset_metadata(dataset, info, source): stats = _source_stats(dataset, source) tasks_table = _source_tasks( dataset, source, int(info["total_tasks"])) - tasks, task_indices = _task_rows( + _, task_indices = _task_rows( tasks_table.to_pylist(), int(info["total_tasks"])) + subtasks_table = _source_subtasks(dataset, source) total_episodes = int(info["total_episodes"]) episode_source = ( _source_episodes(dataset, source) @@ -99,10 +102,10 @@ def _load_dataset_metadata(dataset, info, source): "stats_json": ( None if stats is None else _canonical_json(stats)), "episodes": episodes, - "tasks": tasks, "episodes_schema": episode_source["schema"], "episode_paths": episode_source["paths"], "tasks_table": tasks_table, + "subtasks_table": subtasks_table, "source": source, } @@ -160,7 +163,18 @@ def _prepare_metadata_tables(connection, frames_table, owner_id, metadata): "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] @@ -226,11 +240,19 @@ def _publish_dataset( 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 ( - (frames_identifier, frames_snapshot_id), - (tables["episodes"].identifier, episodes_snapshot_id), - (tables["tasks"].identifier, tasks_snapshot_id)): + for identifier, snapshot_id in component_snapshots: _create_tag(connection.catalog, identifier, tag, snapshot_id) manifest = _manifest_row(version_id, "READY", metadata) @@ -258,6 +280,7 @@ def _manifest_row( "status": status, "info_json": metadata["info_json"], "stats_json": metadata["stats_json"], + "has_subtasks": metadata["subtasks_table"] is not None, } @@ -365,6 +388,29 @@ def _source_tasks(dataset, source, total_tasks): ) 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 ( @@ -432,22 +478,6 @@ def _source_episode_tables(metadata): % (path, error)) from error -def _reject_subtasks(dataset, source): - if source.file_io is not None: - from pypaimon.multimodal.lerobot.source import _remote_path - path = _remote_path(source.path, "meta/subtasks.parquet") - try: - source.file_io.get_file_status(path) - except FileNotFoundError: - return - else: - path = _metadata_root(dataset, source) / "meta" / "subtasks.parquet" - if not path.is_file(): - return - raise ValueError( - "LeRobot subtask metadata is not supported yet: %s" % path) - - def _metadata_root(dataset, source): if source.root is not None: return Path(source.root) diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py b/paimon-python/pypaimon/multimodal/lerobot/schema.py index 6b8ccfb8b537..8df069be4aa4 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/schema.py +++ b/paimon-python/pypaimon/multimodal/lerobot/schema.py @@ -47,22 +47,15 @@ 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_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 b7fa703dd903..ed1d0b1ba832 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -492,8 +492,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 c5f144ca4e1d..7e72a91b4ced 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -36,13 +36,14 @@ from pypaimon.multimodal.hdf5 import _Hdf5SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot from pypaimon.multimodal.lerobot.metadata import ( + _load_dataset_metadata, _managed_table_options, _OWNER_ID_OPTION, ) from pypaimon.multimodal.lerobot.loader import ( _image_bytes, _read_batch, - _task_name, + _validate_frame_controls, ) from pypaimon.multimodal.lerobot.schema import ( _schema_from_info, @@ -197,10 +198,21 @@ 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_optional_dependency_error_is_actionable(self): original_import = builtins.__import__ @@ -224,10 +236,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) @@ -240,7 +252,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": { @@ -249,7 +261,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({ @@ -263,7 +275,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( @@ -459,7 +471,7 @@ def test_empty_dataset_with_tasks_is_rejected(self): finally: shutil.rmtree(temp_dir, ignore_errors=True) - def test_subtasks_are_rejected_before_any_snapshot(self): + def test_optional_subtasks_keep_their_native_schema(self): temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_subtasks_")) try: source = temp_dir / "source" @@ -478,14 +490,20 @@ def test_subtasks_are_rejected_before_any_snapshot(self): "subtask_index": [0], "subtask": ["reach"], }), source / "meta" / "subtasks.parquet") - connection = pmm.connect(options={ - "warehouse": str(temp_dir / "warehouse"), - }) + 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", + ), + ) - with self.assertRaisesRegex(ValueError, "subtask metadata"): - connection.load_from_lerobot("frames", source) - with self.assertRaises(TableNotExistException): - connection.get_table("frames") + expected = pq.read_table(source / "meta" / "subtasks.parquet") + self.assertTrue(metadata["subtasks_table"].equals(expected)) finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -540,7 +558,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) @@ -555,7 +573,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()) @@ -762,6 +780,7 @@ def test_import_infers_schema_and_preserves_episodes(self): 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", @@ -769,7 +788,6 @@ def test_import_infers_schema_and_preserves_episodes(self): "timestamp", "index", "task_index", - "task", "observation.state", "observation.matrix", "action", @@ -778,8 +796,8 @@ def test_import_infers_schema_and_preserves_episodes(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"]) @@ -795,8 +813,10 @@ def test_import_infers_schema_and_preserves_episodes(self): manifest["info_json"])["codebase_version"]) self.assertIsNotNone(manifest["stats_json"]) self.assertEqual( - {"version_id", "status", "info_json", "stats_json"}, + {"version_id", "status", "info_json", "stats_json", + "has_subtasks"}, set(manifest)) + self.assertFalse(manifest["has_subtasks"]) tag = str(manifest["version_id"]) self.assertEqual( 1, @@ -877,6 +897,66 @@ def test_import_infers_schema_and_preserves_episodes(self): "robot_data", self.image_source, batch_size=4) self.assertEqual(5, table.scan().to_arrow().num_rows) + 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="subtask"), + )) + 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"), + ) + self.assertTrue(_catalog_rows( + self.connection, "with_subtasks__versions")[1]["has_subtasks"]) + self.assertEqual( + 1, + self.connection.catalog.get_tag( + self.connection._identifier("with_subtasks__subtasks"), + str(version_id), + ).snapshot.id, + ) + def test_frame_controls_must_match_published_episode_metadata(self): cases = [ ("index", 99), @@ -905,12 +985,14 @@ def test_frame_controls_must_match_published_episode_metadata(self): with self.assertRaisesRegex( ValueError, "has %s" % column): self.connection.load_from_lerobot(table_name, source) - for suffix in ("", "__versions", "__episodes", "__tasks"): + for suffix in ( + "", "__versions", "__episodes", "__tasks", + "__subtasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier(table_name + suffix)) - def test_frame_task_uses_published_task_mapping(self): + 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" @@ -918,8 +1000,12 @@ def test_frame_task_uses_published_task_mapping(self): pq.write_table(tasks.take(pa.array([1, 0])), path) self.connection.load_from_lerobot("reordered_tasks", source) - frames = self.connection.get_table("reordered_tasks").scan().select([ - "index", "task_index", "task" + 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") @@ -928,10 +1014,9 @@ def test_frame_task_uses_published_task_mapping(self): published = { row["task_index"]: row[task_name] for row in task_rows } + self.assertEqual({0: "pick", 1: "place"}, published) self.assertTrue(all( - row["task"] == published[row["task_index"]] - for row in frames - )) + 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" @@ -997,7 +1082,7 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): 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 @@ -1005,10 +1090,8 @@ 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.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 @@ -1035,7 +1118,7 @@ def test_empty_oss_source_does_not_require_episode_directory(self): source_file_io = _RemoteLeRobotFileIO(local_source, source) with patch( - "pypaimon.multimodal.lerobot.source._SourceFileIO", + "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): with self.assertRaisesRegex(ValueError, "non-empty"): self.connection.load_from_lerobot("empty_oss", source) @@ -1066,7 +1149,8 @@ def test_failed_publication_is_cleaned_and_can_retry(self): self.connection.load_from_lerobot( "failed_publish", self.image_source) - for suffix in ("", "__versions", "__episodes", "__tasks"): + for suffix in ( + "", "__versions", "__episodes", "__tasks", "__subtasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier("failed_publish" + suffix)) @@ -1167,7 +1251,7 @@ def test_source_close_failure_does_not_override_success(self): with self.assertLogs( "pypaimon.multimodal.lerobot.source", level="WARNING"): with patch( - "pypaimon.multimodal.lerobot.source._SourceFileIO", + "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", return_value=source_file_io): version_id = self.connection.load_from_lerobot( "source_close_failure", source) @@ -1181,7 +1265,7 @@ def test_source_close_failure_does_not_override_success(self): 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) + schema = _schema_from_info(info) table = self.connection.create_table("existing", schema=schema) with self.assertRaisesRegex(ValueError, "already exists"): @@ -1243,17 +1327,12 @@ def append_then_write( source_schema, batch_size, metadata): - task_names = { - row["task_index"]: row["task"] - for row in metadata["tasks"] - } table.add(_read_batch( dataset, info, 0, 1, source_schema, - task_names, )) return original_write( table, @@ -1275,7 +1354,8 @@ def append_then_write( "concurrent_append", "concurrent_append__versions", "concurrent_append__episodes", - "concurrent_append__tasks"): + "concurrent_append__tasks", + "concurrent_append__subtasks"): with self.subTest(name=name): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( @@ -1287,7 +1367,8 @@ def test_drop_table_removes_companion_tables(self): for name in ( "drop_group", "drop_group__versions", - "drop_group__episodes", "drop_group__tasks"): + "drop_group__episodes", "drop_group__tasks", + "drop_group__subtasks"): with self.subTest(name=name): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( @@ -1316,7 +1397,8 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.drop_table("retry_drop") for name in ( "retry_drop", "retry_drop__versions", - "retry_drop__episodes", "retry_drop__tasks"): + "retry_drop__episodes", "retry_drop__tasks", + "retry_drop__subtasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier(name)) @@ -1346,6 +1428,9 @@ def test_drop_table_rejects_managed_branch(self): "branch_drop__episodes", "branch_drop__tasks"): self.connection.catalog.get_table( self.connection._identifier(name)) + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier("branch_drop__subtasks")) def test_table_group_survives_frame_table_rename(self): self.connection.load_from_lerobot("before_rename", self.image_source) @@ -1361,7 +1446,8 @@ def test_table_group_survives_frame_table_rename(self): self.connection.drop_table("after_rename") for name in ( "after_rename", "before_rename__versions", - "before_rename__episodes", "before_rename__tasks"): + "before_rename__episodes", "before_rename__tasks", + "before_rename__subtasks"): with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier(name)) From b9fdcefa5e1302e9ff5604c6942c14e048cf0f52 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 07:10:13 -0700 Subject: [PATCH 12/32] [python] Stabilize LeRobot FileIO tests --- .../pypaimon/tests/multimodal_lerobot_test.py | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 7e72a91b4ced..147e98babc2d 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -35,6 +35,7 @@ from pypaimon.common.options import Options from pypaimon.multimodal.hdf5 import _Hdf5SourceFileIO from pypaimon.multimodal.lerobot import load_from_lerobot +import pypaimon.multimodal.lerobot.source as lerobot_source from pypaimon.multimodal.lerobot.metadata import ( _load_dataset_metadata, _managed_table_options, @@ -65,6 +66,13 @@ LeRobotDataset = None +_SOURCE_FILE_IO = ( + "_SourceFileIO" + if hasattr(lerobot_source, "_SourceFileIO") + else "_Hdf5SourceFileIO" +) + + def _replaced_contract(field, old, new): description = field.metadata[b"description"].decode("utf-8") if old not in description: @@ -183,9 +191,8 @@ def test_double_encoded_file_uri_cannot_escape_source(self): 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 patch.object( + lerobot_source, _SOURCE_FILE_IO) as source_file_io: with self.assertRaisesRegex(ValueError, "process-isolated"): load_from_lerobot( Mock(), @@ -1070,8 +1077,9 @@ 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", + with patch.object( + lerobot_source, + _SOURCE_FILE_IO, return_value=source_file_io): version_id = self.connection.load_from_lerobot( "oss_images", @@ -1117,8 +1125,9 @@ def test_empty_oss_source_does_not_require_episode_directory(self): source = "oss://source-bucket/empty-robot" source_file_io = _RemoteLeRobotFileIO(local_source, source) - with patch( - "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", + with patch.object( + lerobot_source, + _SOURCE_FILE_IO, return_value=source_file_io): with self.assertRaisesRegex(ValueError, "non-empty"): self.connection.load_from_lerobot("empty_oss", source) @@ -1250,8 +1259,9 @@ def test_source_close_failure_does_not_override_success(self): with self.assertLogs( "pypaimon.multimodal.lerobot.source", level="WARNING"): - with patch( - "pypaimon.multimodal.lerobot.source._Hdf5SourceFileIO", + with patch.object( + lerobot_source, + _SOURCE_FILE_IO, return_value=source_file_io): version_id = self.connection.load_from_lerobot( "source_close_failure", source) From 82430d49435c41f0e2b1183079485e06381ab300 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 08:02:24 -0700 Subject: [PATCH 13/32] [python] Stabilize LeRobot metadata publication --- .../pypaimon/multimodal/lerobot/loader.py | 13 +- .../pypaimon/multimodal/lerobot/metadata.py | 164 ++++++++---------- .../pypaimon/multimodal/lerobot/source.py | 11 -- .../pypaimon/tests/multimodal_lerobot_test.py | 150 ++++++++++++++++ 4 files changed, 230 insertions(+), 108 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index dd95056c58e8..847a9291c67e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -95,6 +95,7 @@ def _write_dataset( episode_begin, begin, task_indices, + metadata["subtask_indices"], ) observed_tasks.setdefault(episode_index, set()).update( seen_tasks) @@ -176,10 +177,13 @@ def _validate_frame_controls( episode_index, episode_begin, begin, - task_indices): + 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( @@ -231,6 +235,13 @@ def _validate_frame_controls( % (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 diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 3dd527f41cee..be4865161d5d 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -80,14 +80,15 @@ def _load_dataset_metadata(dataset, info, source): stats = _source_stats(dataset, source) tasks_table = _source_tasks( dataset, source, int(info["total_tasks"])) - _, task_indices = _task_rows( + task_indices = _task_indices( tasks_table.to_pylist(), 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": [], "rows": [], "schema": _EMPTY_EPISODES_SCHEMA} + else {"tables": [], "rows": [], "schema": _EMPTY_EPISODES_SCHEMA} ) episodes = _episode_rows( episode_source["rows"], @@ -103,10 +104,10 @@ def _load_dataset_metadata(dataset, info, source): None if stats is None else _canonical_json(stats)), "episodes": episodes, "episodes_schema": episode_source["schema"], - "episode_paths": episode_source["paths"], + "episode_tables": episode_source["tables"], "tasks_table": tasks_table, "subtasks_table": subtasks_table, - "source": source, + "subtask_indices": subtask_indices, } @@ -115,16 +116,28 @@ def _new_owner_id(): def _companion_identifier(frames_identifier, suffix): - identifier = Identifier.from_string(str(frames_identifier)) + 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) - return Identifier( + companion = Identifier( identifier.get_database_name(), identifier.get_table_name() + suffix, branch=identifier.get_branch_name(), - ).get_full_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, owner_id): @@ -234,7 +247,7 @@ def _publish_dataset( _require_initial_snapshot("frames", frames_snapshot_id) episodes_snapshot_id = _append_arrow_tables( tables["episodes"], - _source_episode_tables(metadata), + metadata["episode_tables"], ) _require_initial_snapshot("episodes", episodes_snapshot_id) tasks_snapshot_id = _append_arrow( @@ -287,7 +300,7 @@ def _manifest_row( def _drop_import_tables(catalog, frames_table, owner_id): identifiers = list( _companion_table_identifiers(frames_table).values()) - identifiers.append(frames_table.identifier.get_full_name()) + identifiers.append(frames_table.identifier) for identifier in identifiers: try: table = catalog.get_table(identifier) @@ -303,12 +316,14 @@ def _append_arrow(table, data): def _append_arrow_tables(table, tables): builder = table.new_batch_write_builder() - table_write = builder.new_write() - table_commit = builder.new_commit() + table_write = None + table_commit = None commit_started = False recorder = _SnapshotRecorder() - table_commit.add_commit_callback(recorder) 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: @@ -330,14 +345,16 @@ def _append_arrow_tables(table, tables): raise RuntimeError("LeRobot metadata commit has no snapshot id.") return recorder.snapshot_id except BaseException: - if not commit_started: + if table_write is not None and not commit_started: table_write.abort() raise finally: try: - table_write.close() + if table_write is not None: + table_write.close() finally: - table_commit.close() + if table_commit is not None: + table_commit.close() def _create_tag(catalog, identifier, tag_name, snapshot_id): @@ -415,7 +432,6 @@ def _source_episodes(dataset, source): if source.file_io is not None: from pypaimon.multimodal.lerobot.source import ( _read_remote_parquet, - _read_remote_parquet_schema, _remote_parquet_files, _remote_path, ) @@ -426,8 +442,6 @@ def read(path, columns=None): return _read_remote_parquet( source.file_io, path, columns=columns) - 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")) @@ -435,47 +449,28 @@ def read_schema(path): def read(path, columns=None): return pq.read_table(path, columns=columns) - read_schema = pq.read_schema if not paths: return { - "paths": [], + "tables": [], "rows": [], "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:]): + tables = [read(path) for path in paths] + schema = tables[0].schema + if any(not table.schema.equals(schema, check_metadata=False) + for table in tables[1:]): raise ValueError("Episode Parquet schemas are inconsistent.") - tables = [read(path, columns=_EPISODE_CONTROL_COLUMNS) - for path in paths] except (OSError, ValueError, pa.ArrowException) as error: raise ValueError( "Cannot read LeRobot Episode metadata %s: %s" % (directory, error)) from error rows = [] for table in tables: - rows.extend(table.to_pylist()) + rows.extend(table.select(_EPISODE_CONTROL_COLUMNS).to_pylist()) rows.sort(key=lambda row: _integer( row.get("episode_index"), "episode_index")) - return {"paths": paths, "rows": rows, "schema": schema} - - -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 + return {"tables": tables, "rows": rows, "schema": schema} def _metadata_root(dataset, source): @@ -488,8 +483,8 @@ def _metadata_root(dataset, source): return Path(root) -def _task_rows(records, total_tasks): - rows = [None] * total_tasks +def _task_indices(records, total_tasks): + seen = [False] * total_tasks by_name = {} for record in records: index = _integer(record.get("task_index"), "task_index") @@ -497,25 +492,40 @@ def _task_rows(records, total_tasks): if task is None: task = record.get("__index_level_0__") if index < 0 or index >= total_tasks or task is None \ - or rows[index] is not None: + or seen[index]: raise ValueError("LeRobot task metadata is invalid: %s" % record) task = str(task) if task in by_name: raise ValueError("LeRobot task metadata repeats task %r." % task) by_name[task] = index - extra = dict(record) - for key in ("task_index", "task", "name", "__index_level_0__"): - extra.pop(key, None) - rows[index] = { - "task_index": index, - "task": task, - "task_metadata_json": ( - _canonical_json(extra) if extra else None), - } - if any(row is None for row in rows): + seen[index] = True + if not all(seen): raise ValueError( "LeRobot task metadata does not cover [0, %d)." % total_tasks) - return rows, by_name + 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.") + for expected, value in enumerate( + subtasks_table.column("subtask_index").to_pylist()): + if _integer(value, "subtask_index") != expected: + raise ValueError( + "LeRobot subtask metadata must cover [0, %d) in order." + % subtasks_table.num_rows) + return range(subtasks_table.num_rows) def _episode_rows(records, task_indices, info, total_frames, total_episodes): @@ -523,7 +533,6 @@ def _episode_rows(records, task_indices, info, total_frames, total_episodes): raise ValueError( "LeRobot metadata reports %d Episodes but %d were found." % (total_episodes, len(records))) - splits = _episode_splits(info.get("splits"), total_episodes) rows = [] expected_begin = 0 for ordinal, record in enumerate(records): @@ -553,27 +562,12 @@ def _episode_rows(records, task_indices, info, total_frames, total_episodes): raise ValueError( "LeRobot Episode %d repeats a task." % ordinal) - stats = { - key[len("stats/"):]: value - for key, value in record.items() if key.startswith("stats/") - } - extra = { - key: value for key, value in record.items() - if key not in { - "episode_index", "dataset_from_index", "dataset_to_index", - "length", "tasks" - } and not key.startswith("stats/") - } rows.append({ "episode_index": index, "dataset_from_index": begin, "dataset_to_index": end, "length": length, "task_indices": episode_task_indices, - "split": splits[index], - "episode_stats_json": ( - _canonical_json(stats) if stats else None), - "episode_metadata_json": _canonical_json(extra), }) expected_begin = end if expected_begin != total_frames: @@ -583,28 +577,6 @@ def _episode_rows(records, task_indices, info, total_frames, total_episodes): return rows -def _episode_splits(value, total_episodes): - result = [None] * total_episodes - if not isinstance(value, dict): - return result - for name, bounds in value.items(): - if not isinstance(bounds, str) or bounds.count(":") != 1: - continue - begin_text, end_text = bounds.split(":") - try: - begin, end = int(begin_text), int(end_text) - except ValueError: - continue - if begin < 0 or end < begin or end > total_episodes: - continue - for index in range(begin, end): - if result[index] is None: - result[index] = str(name) - elif result[index] != str(name): - result[index] = None - return result - - def _canonical_json(value): return json.dumps( _json_value(value), diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index 68ddf4d4aa2d..7bdb03678570 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -466,17 +466,6 @@ 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) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 6c1523033669..2e87f5bf3161 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -32,12 +32,16 @@ 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, + _subtask_indices, _OWNER_ID_OPTION, ) from pypaimon.multimodal.lerobot.loader import ( @@ -87,6 +91,14 @@ def test_self_contained_import_rejects_table_branches(self): with self.assertRaisesRegex(ValueError, "does not support"): _managed_table_options("db.robot$branch_dev", "owner") + 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 = [] @@ -214,6 +226,32 @@ def test_timestamp_validation_quantizes_float32(self): 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_optional_dependency_error_is_actionable(self): original_import = builtins.__import__ @@ -484,6 +522,7 @@ def test_optional_subtasks_keep_their_native_schema(self): "fps": 30, "features": { "index": {"dtype": "int64", "shape": [1]}, + "subtask_index": {"dtype": "int64", "shape": [1]}, }, })) pq.write_table(pa.table({ @@ -504,6 +543,74 @@ def test_optional_subtasks_keep_their_native_schema(self): 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): + info = {"features": {"subtask_index": {}}} + with self.assertRaisesRegex(ValueError, "subtasks.parquet is missing"): + _subtask_indices(None, info) + with self.assertRaisesRegex(ValueError, "must cover"): + _subtask_indices(pa.table({ + "subtask_index": [1, 0], + "subtask": ["reach", "grasp"], + }), info) + + def test_native_metadata_does_not_require_json_values(self): + 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({ + "task_index": [0], + "task": ["pick"], + "native_bytes": [b"\xff"], + }), 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", + ) + metadata = _load_dataset_metadata( + None, + info, + _LeRobotSource( + path=str(source), + root=source, + repo_id="local/native-metadata", + ), + ) + + stored_episode = metadata["episode_tables"][0] + 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(), + ) finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -957,6 +1064,49 @@ def test_import_publishes_optional_subtasks(self): ).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), From 1b8c80793879687a0080326991ef4da056a3cef0 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 08:53:42 -0700 Subject: [PATCH 14/32] [python] Stream LeRobot episode metadata --- paimon-python/README.md | 7 +- .../pypaimon/multimodal/lerobot/api.py | 33 +++-- .../pypaimon/multimodal/lerobot/metadata.py | 118 +++++++++++------- .../pypaimon/multimodal/lerobot/source.py | 11 ++ .../pypaimon/tests/multimodal_lerobot_test.py | 44 ++++++- 5 files changed, 146 insertions(+), 67 deletions(-) diff --git a/paimon-python/README.md b/paimon-python/README.md index fa0788e9bca4..ac766d5c866e 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -51,9 +51,10 @@ version_id = connection.load_from_lerobot( print(version_id) ``` -The 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`. +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/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 2f6205ce36b7..5cd7be88b552 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -27,13 +27,16 @@ ) from pypaimon.multimodal.lerobot.metadata import ( _OWNER_ID_OPTION, + _append_arrow_tables, _drop_import_tables, _load_dataset_metadata, _managed_table_options, _new_owner_id, _prepare_metadata_tables, + _positive_integer, _publish_dataset, _reserve_dataset_version, + _validated_episode_tables, ) from pypaimon.multimodal.lerobot.loader import _write_dataset from pypaimon.multimodal.lerobot.schema import ( @@ -88,25 +91,8 @@ def load_from_lerobot( _require_v3(local_info, resolved_source.path) _validate_info_paths(local_info) _schema_from_info(local_info) - total_frames, _, _ = \ - _validated_counts(local_info, resolved_source.path) - if total_frames == 0 and ( - resolved_source.root is not None - or resolved_source.file_io is not None): - source_schema = _schema_from_info(local_info) - metadata = _load_dataset_metadata( - None, local_info, resolved_source) - return _import_dataset( - connection, - table_name, - None, - local_info, - resolved_source, - source_schema, - batch_size, - options, - metadata, - ) + _positive_integer(local_info.get("fps"), "fps") + _validated_counts(local_info, resolved_source.path) LeRobotDataset = _import_lerobot_dataset() dataset = _open_resolved_dataset( LeRobotDataset, resolved_source, local_info) @@ -156,6 +142,11 @@ def _import_dataset( version_id, metadata, ) + episodes_snapshot_id = _append_arrow_tables( + tables["episodes"], + _validated_episode_tables(metadata), + flush_each=True, + ) frames_snapshot_id = None if int(info["total_frames"]) > 0: frames_snapshot_id = _write_dataset( @@ -174,6 +165,7 @@ def _import_dataset( metadata, table.identifier, frames_snapshot_id, + episodes_snapshot_id, ) return version_id except BaseException as error: @@ -197,6 +189,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 diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index be4865161d5d..a35b8e7065ba 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -88,25 +88,23 @@ def _load_dataset_metadata(dataset, info, source): episode_source = ( _source_episodes(dataset, source) if total_episodes > 0 - else {"tables": [], "rows": [], "schema": _EMPTY_EPISODES_SCHEMA} - ) - episodes = _episode_rows( - episode_source["rows"], - task_indices, - info, - int(info["total_frames"]), - total_episodes, + 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)), - "episodes": episodes, + None if stats is None else _canonical_json( + stats, allow_nan=True)), + "episodes": None, "episodes_schema": episode_source["schema"], - "episode_tables": episode_source["tables"], + "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, } @@ -243,12 +241,9 @@ def _publish_dataset( version_id, metadata, frames_identifier, - frames_snapshot_id): + frames_snapshot_id, + episodes_snapshot_id): _require_initial_snapshot("frames", frames_snapshot_id) - episodes_snapshot_id = _append_arrow_tables( - tables["episodes"], - metadata["episode_tables"], - ) _require_initial_snapshot("episodes", episodes_snapshot_id) tasks_snapshot_id = _append_arrow( tables["tasks"], metadata["tasks_table"]) @@ -314,8 +309,11 @@ def _append_arrow(table, data): return _append_arrow_tables(table, [data]) -def _append_arrow_tables(table, tables): - builder = table.new_batch_write_builder() +def _append_arrow_tables(table, tables, flush_each=False): + builder = ( + table.new_stream_write_builder() + if flush_each else table.new_batch_write_builder() + ) table_write = None table_commit = None commit_started = False @@ -325,6 +323,7 @@ def _append_arrow_tables(table, tables): table_commit = builder.new_commit() table_commit.add_commit_callback(recorder) row_count = 0 + messages = [] target_schema = _target_schema(table) for data in tables: if data.num_rows == 0: @@ -335,12 +334,19 @@ def _append_arrow_tables(table, tables): % (data.schema, target_schema)) table_write.write_arrow(data) row_count += data.num_rows + if flush_each: + messages = table_write.prepare_commit(0) + del data if row_count == 0: table_write.abort() return None - messages = table_write.prepare_commit() + if not flush_each: + messages = table_write.prepare_commit() commit_started = True - table_commit.commit(messages) + if flush_each: + table_commit.commit(messages, 0) + else: + table_commit.commit(messages) if recorder.snapshot_id is None: raise RuntimeError("LeRobot metadata commit has no snapshot id.") return recorder.snapshot_id @@ -431,46 +437,69 @@ def _source_subtasks(dataset, source): def _source_episodes(dataset, source): if source.file_io is not None: from pypaimon.multimodal.lerobot.source import ( - _read_remote_parquet, + _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(path, columns=None): - return _read_remote_parquet( - source.file_io, path, columns=columns) + 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")) - - def read(path, columns=None): - return pq.read_table(path, columns=columns) + read_schema = pq.read_schema if not paths: return { - "tables": [], - "rows": [], + "paths": [], "schema": _EMPTY_EPISODES_SCHEMA, } try: - tables = [read(path) for path in paths] - schema = tables[0].schema - if any(not table.schema.equals(schema, check_metadata=False) - for table in tables[1:]): + 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): rows = [] - for table in tables: + for table in _source_episode_tables(metadata): rows.extend(table.select(_EPISODE_CONTROL_COLUMNS).to_pylist()) + yield table + del table rows.sort(key=lambda row: _integer( row.get("episode_index"), "episode_index")) - return {"tables": tables, "rows": rows, "schema": schema} + metadata["episodes"] = _episode_rows( + rows, + metadata["task_indices"], + metadata["total_frames"], + metadata["total_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): @@ -519,16 +548,21 @@ def _subtask_indices(subtasks_table, info): if "subtask_index" not in subtasks_table.column_names: raise ValueError( "LeRobot subtask metadata is missing subtask_index.") - for expected, value in enumerate( - subtasks_table.column("subtask_index").to_pylist()): - if _integer(value, "subtask_index") != expected: + records = subtasks_table.to_pylist() + for expected, record in enumerate(records): + label = record.get("subtask", record.get("name")) + if label is None: + label = record.get("__index_level_0__") + if _integer(record.get("subtask_index"), "subtask_index") \ + != expected or not isinstance(label, str) or not label: raise ValueError( - "LeRobot subtask metadata must cover [0, %d) in order." + "LeRobot subtask metadata must provide ordered numeric and " + "text mappings for [0, %d)." % subtasks_table.num_rows) return range(subtasks_table.num_rows) -def _episode_rows(records, task_indices, info, total_frames, total_episodes): +def _episode_rows(records, task_indices, total_frames, total_episodes): if len(records) != total_episodes: raise ValueError( "LeRobot metadata reports %d Episodes but %d were found." @@ -577,13 +611,13 @@ def _episode_rows(records, task_indices, info, total_frames, total_episodes): return rows -def _canonical_json(value): +def _canonical_json(value, allow_nan=False): return json.dumps( _json_value(value), ensure_ascii=False, sort_keys=True, separators=(",", ":"), - allow_nan=False, + allow_nan=allow_nan, ) diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index 7bdb03678570..68ddf4d4aa2d 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -466,6 +466,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) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 2e87f5bf3161..e553399032f3 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -23,7 +23,7 @@ import unittest from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import call, Mock, patch import numpy as np import pyarrow as pa @@ -42,6 +42,7 @@ _load_dataset_metadata, _managed_table_options, _subtask_indices, + _validated_episode_tables, _OWNER_ID_OPTION, ) from pypaimon.multimodal.lerobot.loader import ( @@ -252,6 +253,29 @@ def test_metadata_writer_closes_after_commit_creation_failure(self): writer.abort.assert_called_once_with() writer.close.assert_called_once_with() + def test_episode_shards_are_flushed_incrementally(self): + data = pa.table({"episode_index": [0]}) + table = Mock() + builder = table.new_stream_write_builder.return_value + writer = builder.new_write.return_value + writer.prepare_commit.return_value = [Mock()] + + 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(), flush_each=True) + + self.assertEqual([call(0), call(0)], + writer.prepare_commit.call_args_list) + writer.abort.assert_called_once_with() + writer.close.assert_called_once_with() + def test_optional_dependency_error_is_actionable(self): original_import = builtins.__import__ @@ -551,11 +575,15 @@ def test_subtask_metadata_must_match_frame_feature(self): info = {"features": {"subtask_index": {}}} with self.assertRaisesRegex(ValueError, "subtasks.parquet is missing"): _subtask_indices(None, info) - with self.assertRaisesRegex(ValueError, "must cover"): + with self.assertRaisesRegex(ValueError, "numeric and text mappings"): _subtask_indices(pa.table({ "subtask_index": [1, 0], "subtask": ["reach", "grasp"], }), info) + with self.assertRaisesRegex(ValueError, "numeric and text mappings"): + _subtask_indices(pa.table({ + "subtask_index": [0], + }), info) def test_native_metadata_does_not_require_json_values(self): temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_native_")) @@ -590,6 +618,10 @@ def test_native_metadata_does_not_require_json_values(self): 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, @@ -600,7 +632,10 @@ def test_native_metadata_does_not_require_json_values(self): ), ) - stored_episode = metadata["episode_tables"][0] + 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(), @@ -611,6 +646,9 @@ def test_native_metadata_does_not_require_json_values(self): 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) From 5fd90844678da9ddedecb756cf0c78793eede4bf Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 20:17:36 -0700 Subject: [PATCH 15/32] [python] Make LeRobot cleanup generation safe --- .../pypaimon/multimodal/connection.py | 22 +---- .../pypaimon/multimodal/lerobot/api.py | 1 - .../pypaimon/multimodal/lerobot/metadata.py | 56 ++++++++---- .../pypaimon/tests/multimodal_lerobot_test.py | 88 ++++++++++++++++--- 4 files changed, 116 insertions(+), 51 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index a3371cb0eee9..99ca8e8360d4 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -203,27 +203,11 @@ def drop_table(self, name: str, ignore_if_not_exists: bool = False): except (DatabaseNotExistException, TableNotExistException): pass - companions = [] if owner_id is not None: from pypaimon.multimodal.lerobot.metadata import \ - _companion_table_identifiers - for companion in _companion_table_identifiers( - raw_table).values(): - try: - table = self.catalog.get_table(companion) - except (DatabaseNotExistException, TableNotExistException): - continue - actual = table.table_schema.options.get(_OWNER_ID_OPTION) - if actual != owner_id: - raise ValueError( - "Refusing to drop %s because it belongs to a " - "different table." % companion) - companions.append(companion) - for companion in companions: - self.catalog.drop_table( - companion, - ignore_if_not_exists=True, - ) + _drop_import_tables + _drop_import_tables(self.catalog, raw_table, owner_id) + return self.catalog.drop_table( identifier, ignore_if_not_exists=ignore_if_not_exists, diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 5cd7be88b552..0aa5fa7893ce 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -145,7 +145,6 @@ def _import_dataset( episodes_snapshot_id = _append_arrow_tables( tables["episodes"], _validated_episode_tables(metadata), - flush_each=True, ) frames_snapshot_id = None if int(info["total_frames"]) > 0: diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index a35b8e7065ba..ccd27568fd2e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -297,23 +297,48 @@ def _drop_import_tables(catalog, frames_table, owner_id): _companion_table_identifiers(frames_table).values()) identifiers.append(frames_table.identifier) for identifier in identifiers: + _drop_owned_table(catalog, identifier, owner_id) + + +def _drop_owned_table(catalog, identifier, owner_id): + """Move a table aside before checking ownership and deleting it.""" + source = ( + identifier + if isinstance(identifier, Identifier) + else Identifier.from_string(str(identifier)) + ) + quarantine = Identifier( + source.get_database_name(), + "__pypaimon_drop_%s" % uuid.uuid4().hex, + ) + try: + catalog.rename_table(source, quarantine) + except (DatabaseNotExistException, TableNotExistException): + return + + try: + table = catalog.get_table(quarantine) + actual_owner = table.table_schema.options.get(_OWNER_ID_OPTION) + if actual_owner != owner_id: + raise ValueError( + "Refusing to drop %s because it belongs to a different " + "table." % source) + catalog.drop_table(quarantine) + except BaseException: try: - table = catalog.get_table(identifier) - except (DatabaseNotExistException, TableNotExistException): - continue - if table.table_schema.options.get(_OWNER_ID_OPTION) == owner_id: - catalog.drop_table(identifier, ignore_if_not_exists=True) + catalog.rename_table(quarantine, source) + except (DatabaseNotExistException, TableAlreadyExistException, + TableNotExistException): + pass + raise def _append_arrow(table, data): return _append_arrow_tables(table, [data]) -def _append_arrow_tables(table, tables, flush_each=False): - builder = ( - table.new_stream_write_builder() - if flush_each else table.new_batch_write_builder() - ) +def _append_arrow_tables(table, tables): + builder = table.new_batch_write_builder() table_write = None table_commit = None commit_started = False @@ -323,7 +348,6 @@ def _append_arrow_tables(table, tables, flush_each=False): table_commit = builder.new_commit() table_commit.add_commit_callback(recorder) row_count = 0 - messages = [] target_schema = _target_schema(table) for data in tables: if data.num_rows == 0: @@ -334,19 +358,13 @@ def _append_arrow_tables(table, tables, flush_each=False): % (data.schema, target_schema)) table_write.write_arrow(data) row_count += data.num_rows - if flush_each: - messages = table_write.prepare_commit(0) del data if row_count == 0: table_write.abort() return None - if not flush_each: - messages = table_write.prepare_commit() + messages = table_write.prepare_commit() commit_started = True - if flush_each: - table_commit.commit(messages, 0) - else: - table_commit.commit(messages) + table_commit.commit(messages) if recorder.snapshot_id is None: raise RuntimeError("LeRobot metadata commit has no snapshot id.") return recorder.snapshot_id diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index e553399032f3..ea1ff2fd18d8 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -23,7 +23,7 @@ import unittest from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from unittest.mock import call, Mock, patch +from unittest.mock import Mock, patch import numpy as np import pyarrow as pa @@ -253,12 +253,11 @@ def test_metadata_writer_closes_after_commit_creation_failure(self): writer.abort.assert_called_once_with() writer.close.assert_called_once_with() - def test_episode_shards_are_flushed_incrementally(self): + def test_episode_shards_use_normal_batch_rolling(self): data = pa.table({"episode_index": [0]}) table = Mock() - builder = table.new_stream_write_builder.return_value + builder = table.new_batch_write_builder.return_value writer = builder.new_write.return_value - writer.prepare_commit.return_value = [Mock()] def shards(): yield data @@ -269,10 +268,10 @@ def shards(): "pypaimon.multimodal.lerobot.metadata._target_schema", return_value=data.schema): with self.assertRaisesRegex(RuntimeError, "two shards"): - _append_arrow_tables(table, shards(), flush_each=True) + _append_arrow_tables(table, shards()) - self.assertEqual([call(0), call(0)], - writer.prepare_commit.call_args_list) + 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() @@ -1042,6 +1041,29 @@ def test_import_infers_schema_and_preserves_episodes(self): "robot_data", self.image_source, batch_size=4) self.assertEqual(5, table.scan().to_arrow().num_rows) + 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 @@ -1566,18 +1588,27 @@ def test_drop_table_can_retry_after_companion_failure(self): self.connection.load_from_lerobot( "retry_drop", self.image_source) original_drop = self.connection.catalog.drop_table + original_rename = self.connection.catalog.rename_table failed = [False] + episode_quarantine = [None] + + def track_rename(source, target): + if source.get_table_name().endswith("__episodes"): + episode_quarantine[0] = target.get_full_name() + return original_rename(source, target) def flaky_drop(identifier, ignore_if_not_exists=False): - if str(identifier).endswith("__episodes") and not failed[0]: + if identifier.get_full_name() == episode_quarantine[0] \ + and not failed[0]: failed[0] = True raise RuntimeError("injected drop failure") return original_drop(identifier, ignore_if_not_exists) - with patch.object( - self.connection.catalog, - "drop_table", - side_effect=flaky_drop): + with patch.object(self.connection.catalog, "rename_table", + side_effect=track_rename), patch.object( + self.connection.catalog, + "drop_table", + side_effect=flaky_drop): with self.assertRaisesRegex(RuntimeError, "injected"): self.connection.drop_table("retry_drop") self.connection.get_table("retry_drop") @@ -1591,6 +1622,39 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.catalog.get_table( self.connection._identifier(name)) + def test_drop_table_does_not_delete_recreated_companion(self): + self.connection.load_from_lerobot( + "drop_race", self.image_source) + identifier = self.connection._identifier("drop_race__versions") + old_table = self.connection.catalog.get_table(identifier) + replacement_schema = old_table.table_schema.to_schema() + replacement_schema.options = dict(replacement_schema.options) + replacement_schema.options[_OWNER_ID_OPTION] = "other-owner" + original_rename = self.connection.catalog.rename_table + replaced = [False] + + def replace_before_rename(source, target): + if source.get_full_name() == identifier and not replaced[0]: + replaced[0] = True + self.connection.catalog.drop_table(source) + self.connection.catalog.create_table( + source, replacement_schema, False) + return original_rename(source, target) + + with patch.object( + self.connection.catalog, + "rename_table", + side_effect=replace_before_rename): + with self.assertRaisesRegex(ValueError, "different table"): + self.connection.drop_table("drop_race") + + replacement = self.connection.catalog.get_table(identifier) + self.assertEqual( + "other-owner", + replacement.table_schema.options[_OWNER_ID_OPTION], + ) + self.connection.get_table("drop_race") + def test_companion_table_cannot_be_dropped_directly(self): self.connection.load_from_lerobot( "direct_drop", self.image_source) From e98d1f210fff9e56f26eb0eb9105dee42c5dc336 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 22:42:04 -0700 Subject: [PATCH 16/32] [python] Validate native LeRobot control metadata --- .../pypaimon/multimodal/lerobot/api.py | 3 + .../pypaimon/multimodal/lerobot/metadata.py | 148 +++++++++++------- .../pypaimon/multimodal/lerobot/schema.py | 25 +++ .../pypaimon/tests/multimodal_lerobot_test.py | 27 ++++ 4 files changed, 144 insertions(+), 59 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 0aa5fa7893ce..b0c47bd03d0b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -42,6 +42,7 @@ from pypaimon.multimodal.lerobot.schema import ( _require_v3, _schema_from_info, + _validate_v3_control_features, ) from pypaimon.multimodal.lerobot.source import ( _close_quietly, @@ -93,6 +94,7 @@ def load_from_lerobot( _schema_from_info(local_info) _positive_integer(local_info.get("fps"), "fps") _validated_counts(local_info, resolved_source.path) + _validate_v3_control_features(local_info) LeRobotDataset = _import_lerobot_dataset() dataset = _open_resolved_dataset( LeRobotDataset, resolved_source, local_info) @@ -100,6 +102,7 @@ def load_from_lerobot( info = dict(dataset.meta.info) _require_v3(info, resolved_source.path) _validated_counts(info, resolved_source.path) + _validate_v3_control_features(info) lerobot_schema = _schema_from_info(info) metadata = _load_dataset_metadata( diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index ccd27568fd2e..557ee4ed5e1b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -16,6 +16,7 @@ """LeRobot component tables and version publication.""" +from array import array import json import numbers import uuid @@ -75,6 +76,39 @@ ] +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) @@ -489,19 +523,64 @@ def read_schema(path): def _validated_episode_tables(metadata): - rows = [] + episodes = _EpisodeIndex() + expected_begin = 0 for table in _source_episode_tables(metadata): - rows.extend(table.select(_EPISODE_CONTROL_COLUMNS).to_pylist()) + 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 - rows.sort(key=lambda row: _integer( - row.get("episode_index"), "episode_index")) - metadata["episodes"] = _episode_rows( - rows, - metadata["task_indices"], - metadata["total_frames"], - metadata["total_episodes"], - ) + 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): @@ -580,55 +659,6 @@ def _subtask_indices(subtasks_table, info): return range(subtasks_table.num_rows) -def _episode_rows(records, task_indices, total_frames, total_episodes): - if len(records) != total_episodes: - raise ValueError( - "LeRobot metadata reports %d Episodes but %d were found." - % (total_episodes, len(records))) - rows = [] - expected_begin = 0 - for ordinal, record in enumerate(records): - index = _integer(record.get("episode_index"), "episode_index") - begin = _integer( - record.get("dataset_from_index"), "dataset_from_index") - end = _integer(record.get("dataset_to_index"), "dataset_to_index") - length = _integer(record.get("length"), "length") - if index != ordinal or begin != expected_begin or end <= begin \ - or length != end - begin: - raise ValueError( - "LeRobot Episode %d has inconsistent index, range, or length." - % ordinal) - names = record.get("tasks", []) - if isinstance(names, str): - names = [names] - if task_indices and not names: - raise ValueError( - "LeRobot Episode %d does not declare any task." % ordinal) - try: - episode_task_indices = [task_indices[str(name)] for name in names] - except (KeyError, TypeError) as error: - raise ValueError( - "LeRobot Episode %d refers to an unknown task." % ordinal - ) from error - if len(set(episode_task_indices)) != len(episode_task_indices): - raise ValueError( - "LeRobot Episode %d repeats a task." % ordinal) - - rows.append({ - "episode_index": index, - "dataset_from_index": begin, - "dataset_to_index": end, - "length": length, - "task_indices": episode_task_indices, - }) - expected_begin = end - if expected_begin != total_frames: - raise ValueError( - "LeRobot Episode ranges cover %d frames but metadata reports %d." - % (expected_begin, total_frames)) - return rows - - def _canonical_json(value, allow_nan=False): return json.dumps( _json_value(value), diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py b/paimon-python/pypaimon/multimodal/lerobot/schema.py index 8df069be4aa4..6ee158978fe1 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_CONTROL_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", "")) @@ -58,6 +66,23 @@ def _schema_from_info(info): ]) +def _validate_v3_control_features(info): + features = info.get("features") + if not isinstance(features, dict): + raise ValueError("LeRobot metadata features must be an object.") + expected = dict(_V3_CONTROL_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 control feature %s must have dtype=%s and " + "shape=[1]." % (name, dtype)) + + 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: diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index ea1ff2fd18d8..cb3c18cccb3d 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -53,6 +53,7 @@ from pypaimon.multimodal.lerobot.schema import ( _schema_from_info, _validate_lerobot_schema, + _validate_v3_control_features, ) from pypaimon.multimodal.lerobot.source import ( _LeRobotSource, @@ -389,6 +390,32 @@ def test_existing_schema_preserves_lerobot_feature_contract(self): ValueError, "cannot be converted"): _validate_lerobot_schema(source, target, "dataset") + def test_v3_control_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_control_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, "control feature %s" % name): + _validate_v3_control_features({"features": invalid}) + + missing = dict(features) + del missing["task_index"] + with self.assertRaisesRegex(ValueError, "control feature task_index"): + _validate_v3_control_features({"features": missing}) + def test_remote_episode_metadata_projects_stats_columns(self): source = _LeRobotSource( path="oss://bucket/robot", From 7d06ab56dc6b9da44b3e0b2f16825335dffef70c Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 22:47:46 -0700 Subject: [PATCH 17/32] [python] Preserve LeRobot task index metadata --- .../pypaimon/multimodal/lerobot/metadata.py | 24 +++++++++++++++--- .../pypaimon/tests/multimodal_lerobot_test.py | 25 +++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 557ee4ed5e1b..a335fe0547bb 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -38,6 +38,7 @@ _VERSION_ID = "version_id" _OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" +_PANDAS_METADATA_OPTION = "pypaimon.lerobot.pandas-metadata" _TABLE_SUFFIXES = { "versions": "__versions", "episodes": "__episodes", @@ -226,12 +227,17 @@ def _prepare_metadata_tables(connection, frames_table, owner_id, metadata): try: table = connection.catalog.get_table(identifier) except (DatabaseNotExistException, TableNotExistException): + options = { + "bucket": "-1", + _OWNER_ID_OPTION: owner_id, + } + 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={ - "bucket": "-1", - _OWNER_ID_OPTION: owner_id, - }, + options=options, ) try: connection.catalog.create_table( @@ -257,6 +263,16 @@ def _prepare_metadata_tables(connection, frames_table, owner_id, metadata): 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 _reserve_dataset_version( versions_table, version_id, diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index cb3c18cccb3d..c8894c7b97e8 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -41,6 +41,7 @@ _companion_identifier, _load_dataset_metadata, _managed_table_options, + _restore_pandas_metadata, _subtask_indices, _validated_episode_tables, _OWNER_ID_OPTION, @@ -87,6 +88,13 @@ def _catalog_rows(connection, name): 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): @@ -931,6 +939,8 @@ def _create_image_dataset(root): dataset.finalize() 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) @@ -1043,6 +1053,14 @@ def test_import_infers_schema_and_preserves_episodes(self): 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, @@ -1141,6 +1159,13 @@ def test_import_publishes_optional_subtasks(self): 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")[1]["has_subtasks"]) self.assertEqual( From 47e573479b42e3a570870f9ae1d29d0816c898d5 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 23:22:06 -0700 Subject: [PATCH 18/32] [python] Bound LeRobot import cleanup and episode state --- .../pypaimon/multimodal/lerobot/loader.py | 34 +++++---- .../pypaimon/multimodal/lerobot/metadata.py | 63 ++++++++--------- .../pypaimon/multimodal/lerobot/source.py | 69 +++++++++++++++---- .../pypaimon/tests/multimodal_lerobot_test.py | 39 ++++++++++- 4 files changed, 147 insertions(+), 58 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index 847a9291c67e..295eee99c363 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -77,7 +77,9 @@ def _write_dataset( batch_count = 0 row_count = 0 episodes = metadata["episodes"] - observed_tasks = {} + current_episode = None + expected_tasks = set() + observed_tasks = set() snapshot_recorder = _SnapshotRecorder() try: @@ -86,6 +88,13 @@ def _write_dataset( table_commit.add_commit_callback(snapshot_recorder) 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( @@ -97,8 +106,7 @@ def _write_dataset( task_indices, metadata["subtask_indices"], ) - observed_tasks.setdefault(episode_index, set()).update( - seen_tasks) + observed_tasks.update(seen_tasks) batch = _strict_lerobot_table( batch, target_schema, @@ -109,7 +117,9 @@ def _write_dataset( batch_count += 1 row_count += batch.num_rows - _validate_episode_tasks(episodes, observed_tasks) + 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: @@ -245,16 +255,12 @@ def _validate_frame_controls( return seen_tasks -def _validate_episode_tasks(episodes, observed_tasks): - for episode in episodes: - episode_index = episode["episode_index"] - expected = set(episode["task_indices"]) - actual = observed_tasks.get(episode_index, set()) - if actual != expected: - raise ValueError( - "LeRobot Episode %d declares task indices %s but its " - "frames use %s." - % (episode_index, sorted(expected), sorted(actual))) +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): diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index a335fe0547bb..58dd8e6ddd52 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -346,40 +346,41 @@ def _drop_import_tables(catalog, frames_table, owner_id): identifiers = list( _companion_table_identifiers(frames_table).values()) identifiers.append(frames_table.identifier) - for identifier in identifiers: - _drop_owned_table(catalog, identifier, owner_id) - - -def _drop_owned_table(catalog, identifier, owner_id): - """Move a table aside before checking ownership and deleting it.""" - source = ( - identifier - if isinstance(identifier, Identifier) - else Identifier.from_string(str(identifier)) - ) - quarantine = Identifier( - source.get_database_name(), - "__pypaimon_drop_%s" % uuid.uuid4().hex, - ) + quarantined = [] try: - catalog.rename_table(source, quarantine) - except (DatabaseNotExistException, TableNotExistException): - return + for identifier in identifiers: + source = ( + identifier + if isinstance(identifier, Identifier) + else Identifier.from_string(str(identifier)) + ) + quarantine = Identifier( + source.get_database_name(), + "__pypaimon_drop_%s" % uuid.uuid4().hex, + ) + try: + catalog.rename_table(source, quarantine) + except (DatabaseNotExistException, TableNotExistException): + continue + quarantined.append((source, quarantine)) - try: - table = catalog.get_table(quarantine) - actual_owner = table.table_schema.options.get(_OWNER_ID_OPTION) - if actual_owner != owner_id: - raise ValueError( - "Refusing to drop %s because it belongs to a different " - "table." % source) - catalog.drop_table(quarantine) + for source, quarantine in quarantined: + table = catalog.get_table(quarantine) + actual_owner = table.table_schema.options.get(_OWNER_ID_OPTION) + if actual_owner != owner_id: + raise ValueError( + "Refusing to drop %s because it belongs to a different " + "table." % source) + + for _, quarantine in quarantined: + catalog.drop_table(quarantine) except BaseException: - try: - catalog.rename_table(quarantine, source) - except (DatabaseNotExistException, TableAlreadyExistException, - TableNotExistException): - pass + for source, quarantine in reversed(quarantined): + try: + catalog.rename_table(quarantine, source) + except (DatabaseNotExistException, TableAlreadyExistException, + TableNotExistException): + pass raise diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index 68ddf4d4aa2d..f2d29348171b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -16,6 +16,7 @@ """LeRobot source resolution for local, Hub, and FileIO datasets.""" +from array import array import json import logging import posixpath @@ -54,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): @@ -206,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 @@ -268,19 +297,35 @@ 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)) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index c8894c7b97e8..8d5baf95b94e 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -15,6 +15,7 @@ # limitations under the License. import builtins +from array import array import json import shutil import sys @@ -451,13 +452,15 @@ 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_is_rejected_before_opening_lerobot(self): temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_empty_")) @@ -1086,6 +1089,20 @@ def test_import_infers_schema_and_preserves_episodes(self): "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) @@ -1707,6 +1724,26 @@ def replace_before_rename(source, target): ) self.connection.get_table("drop_race") + def test_drop_table_validates_group_before_deleting(self): + self.connection.load_from_lerobot( + "mixed_owner", self.image_source) + identifier = self.connection._identifier("mixed_owner__tasks") + table = self.connection.catalog.get_table(identifier) + schema = table.table_schema.to_schema() + schema.options = dict(schema.options) + schema.options[_OWNER_ID_OPTION] = "other-owner" + self.connection.catalog.drop_table(identifier) + self.connection.catalog.create_table(identifier, schema, False) + + with self.assertRaisesRegex(ValueError, "different table"): + self.connection.drop_table("mixed_owner") + + for name in ( + "mixed_owner", "mixed_owner__versions", + "mixed_owner__episodes", "mixed_owner__tasks"): + self.connection.catalog.get_table( + self.connection._identifier(name)) + def test_companion_table_cannot_be_dropped_directly(self): self.connection.load_from_lerobot( "direct_drop", self.image_source) From 7a58c04859f2e430b2805644866b9086c5f3d0fb Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Wed, 2 Sep 2026 23:43:12 -0700 Subject: [PATCH 19/32] [python] Clean only owned LeRobot import tables --- .../pypaimon/multimodal/lerobot/api.py | 12 +++- .../pypaimon/multimodal/lerobot/metadata.py | 69 +++++++++++++++---- .../pypaimon/tests/multimodal_lerobot_test.py | 27 ++++++-- 3 files changed, 86 insertions(+), 22 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index b0c47bd03d0b..2caeb4a715ad 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -173,7 +173,11 @@ def _import_dataset( except BaseException as error: try: _drop_import_tables( - connection.catalog, table.raw_table, owner_id) + connection.catalog, + table.raw_table, + owner_id, + owned_only=True, + ) except BaseException as cleanup_error: raise RuntimeError( "LeRobot import failed and cleanup also failed: %s" @@ -244,7 +248,11 @@ def _create_target_table( _OWNER_ID_OPTION) == owner_id: try: _drop_import_tables( - connection.catalog, frames_table, owner_id) + connection.catalog, + frames_table, + owner_id, + owned_only=True, + ) except BaseException as cleanup_error: raise RuntimeError( "LeRobot target creation failed and cleanup also failed: " diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 58dd8e6ddd52..da7ae69e0772 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -342,7 +342,8 @@ def _manifest_row( } -def _drop_import_tables(catalog, frames_table, owner_id): +def _drop_import_tables( + catalog, frames_table, owner_id, owned_only=False): identifiers = list( _companion_table_identifiers(frames_table).values()) identifiers.append(frames_table.identifier) @@ -364,24 +365,64 @@ def _drop_import_tables(catalog, frames_table, owner_id): continue quarantined.append((source, quarantine)) + owned = [] + foreign = [] for source, quarantine in quarantined: table = catalog.get_table(quarantine) actual_owner = table.table_schema.options.get(_OWNER_ID_OPTION) - if actual_owner != owner_id: - raise ValueError( - "Refusing to drop %s because it belongs to a different " - "table." % source) + target = owned if actual_owner == owner_id else foreign + target.append((source, quarantine)) + except BaseException as error: + failures = _restore_quarantined(catalog, quarantined) + if failures: + raise RuntimeError( + "Failed to restore quarantined LeRobot tables: %s" + % ", ".join(failures)) from error + raise + + if foreign and not owned_only: + error = ValueError( + "Refusing to drop %s because it belongs to a different table." + % foreign[0][0]) + failures = _restore_quarantined(catalog, quarantined) + if failures: + raise RuntimeError( + "Failed to restore quarantined LeRobot tables: %s" + % ", ".join(failures)) from error + raise error + + restore_failures = _restore_quarantined(catalog, foreign) + drop_failures = _drop_quarantined(catalog, owned) + if drop_failures: + drop_failures = _drop_quarantined(catalog, drop_failures) + if restore_failures or drop_failures: + raise RuntimeError( + "LeRobot cleanup left quarantined tables: %s" + % ", ".join(restore_failures + [ + str(quarantine) for _, quarantine in drop_failures + ])) + + +def _restore_quarantined(catalog, tables): + failures = [] + for source, quarantine in reversed(tables): + try: + catalog.rename_table(quarantine, source) + except BaseException: + failures.append(str(quarantine)) + return failures + - for _, quarantine in quarantined: +def _drop_quarantined(catalog, tables): + failures = [] + for source, quarantine in tables: + try: catalog.drop_table(quarantine) - except BaseException: - for source, quarantine in reversed(quarantined): - try: - catalog.rename_table(quarantine, source) - except (DatabaseNotExistException, TableAlreadyExistException, - TableNotExistException): - pass - raise + except (DatabaseNotExistException, TableNotExistException): + pass + except BaseException: + failures.append((source, quarantine)) + return failures def _append_arrow(table, data): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 8d5baf95b94e..470225c4f400 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1438,6 +1438,25 @@ def test_failed_publication_is_cleaned_and_can_retry(self): "failed_publish", self.image_source) self.assertEqual(1, version_id) + def test_stale_companion_does_not_block_retry(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, "different target"): + self.connection.load_from_lerobot("stale", self.image_source) + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier("stale")) + + self.connection.catalog.drop_table(stale) + self.assertEqual( + 1, + self.connection.load_from_lerobot("stale", self.image_source), + ) + def test_invalid_target_options_do_not_leave_table(self): with self.assertRaisesRegex(ValueError, "data-evolution.enabled"): self.connection.load_from_lerobot( @@ -1653,7 +1672,7 @@ def test_drop_table_removes_companion_tables(self): self.connection.catalog.get_table( self.connection._identifier(name)) - def test_drop_table_can_retry_after_companion_failure(self): + def test_drop_table_retries_companion_failure(self): self.connection.load_from_lerobot( "retry_drop", self.image_source) original_drop = self.connection.catalog.drop_table @@ -1678,11 +1697,7 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.catalog, "drop_table", side_effect=flaky_drop): - with self.assertRaisesRegex(RuntimeError, "injected"): - self.connection.drop_table("retry_drop") - self.connection.get_table("retry_drop") - - self.connection.drop_table("retry_drop") + self.connection.drop_table("retry_drop") for name in ( "retry_drop", "retry_drop__versions", "retry_drop__episodes", "retry_drop__tasks", From f90fd76f450ef2273e1a4a3dc27300d50f2ea6fb Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 00:08:50 -0700 Subject: [PATCH 20/32] [python] Reconcile LeRobot cleanup generations --- .../pypaimon/multimodal/lerobot/metadata.py | 43 ++++++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 83 ++++++++++++++++++- 2 files changed, 115 insertions(+), 11 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index da7ae69e0772..88618ea38a14 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -355,6 +355,14 @@ def _drop_import_tables( if isinstance(identifier, Identifier) else Identifier.from_string(str(identifier)) ) + if owned_only: + try: + table = catalog.get_table(source) + except (DatabaseNotExistException, TableNotExistException): + continue + if table.table_schema.options.get( + _OWNER_ID_OPTION) != owner_id: + continue quarantine = Identifier( source.get_database_name(), "__pypaimon_drop_%s" % uuid.uuid4().hex, @@ -392,9 +400,10 @@ def _drop_import_tables( raise error restore_failures = _restore_quarantined(catalog, foreign) - drop_failures = _drop_quarantined(catalog, owned) + drop_failures = _drop_quarantined(catalog, owned, owner_id) if drop_failures: - drop_failures = _drop_quarantined(catalog, drop_failures) + drop_failures = _drop_quarantined( + catalog, drop_failures, owner_id) if restore_failures or drop_failures: raise RuntimeError( "LeRobot cleanup left quarantined tables: %s" @@ -406,17 +415,29 @@ def _drop_import_tables( def _restore_quarantined(catalog, tables): failures = [] for source, quarantine in reversed(tables): - try: - catalog.rename_table(quarantine, source) - except BaseException: - failures.append(str(quarantine)) + for attempt in range(2): + try: + if not _table_exists(catalog, quarantine): + break + if _table_exists(catalog, source): + failures.append(str(quarantine)) + break + catalog.rename_table(quarantine, source) + break + except BaseException: + if attempt == 1: + failures.append(str(quarantine)) return failures -def _drop_quarantined(catalog, tables): +def _drop_quarantined(catalog, tables, owner_id): failures = [] for source, quarantine in tables: try: + table = catalog.get_table(quarantine) + if table.table_schema.options.get(_OWNER_ID_OPTION) != owner_id: + failures.append((source, quarantine)) + continue catalog.drop_table(quarantine) except (DatabaseNotExistException, TableNotExistException): pass @@ -425,6 +446,14 @@ def _drop_quarantined(catalog, tables): return failures +def _table_exists(catalog, identifier): + try: + catalog.get_table(identifier) + return True + except (DatabaseNotExistException, TableNotExistException): + return False + + def _append_arrow(table, data): return _append_arrow_tables(table, [data]) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 470225c4f400..8c4d2689a2f0 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1442,11 +1442,25 @@ def test_stale_companion_does_not_block_retry(self): self.connection.load_from_lerobot( "other_group", self.image_source) stale = self.connection._identifier("stale__tasks") + stale_name = Identifier.from_string(stale).get_full_name() self.connection.catalog.rename_table( self.connection._identifier("other_group__tasks"), stale) - with self.assertRaisesRegex(ValueError, "different target"): - self.connection.load_from_lerobot("stale", self.image_source) + original_rename = self.connection.catalog.rename_table + renamed_sources = [] + + def track_rename(source, target): + renamed_sources.append(source.get_full_name()) + return original_rename(source, target) + + with patch.object( + self.connection.catalog, + "rename_table", + side_effect=track_rename): + with self.assertRaisesRegex(ValueError, "different target"): + self.connection.load_from_lerobot( + "stale", self.image_source) + self.assertNotIn(stale_name, renamed_sources) with self.assertRaises(TableNotExistException): self.connection.catalog.get_table( self.connection._identifier("stale")) @@ -1706,6 +1720,52 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.catalog.get_table( self.connection._identifier(name)) + def test_drop_retry_does_not_delete_reused_quarantine(self): + self.connection.load_from_lerobot( + "reused_quarantine", self.image_source) + source = self.connection._identifier( + "reused_quarantine__episodes") + source_name = Identifier.from_string(source).get_full_name() + replacement_schema = self.connection.catalog.get_table( + source).table_schema.to_schema() + replacement_schema.options = dict(replacement_schema.options) + replacement_schema.options[_OWNER_ID_OPTION] = "other-owner" + original_drop = self.connection.catalog.drop_table + original_rename = self.connection.catalog.rename_table + quarantine = [None] + injected = [False] + + def track_rename(rename_source, target): + if rename_source.get_full_name() == source_name: + quarantine[0] = target + return original_rename(rename_source, target) + + def drop_then_lose_response(identifier, ignore_if_not_exists=False): + if identifier == quarantine[0] and not injected[0]: + injected[0] = True + original_drop(identifier, ignore_if_not_exists) + self.connection.catalog.create_table( + identifier, replacement_schema, False) + raise RuntimeError("response lost") + return original_drop(identifier, ignore_if_not_exists) + + with patch.object( + self.connection.catalog, + "rename_table", + side_effect=track_rename), patch.object( + self.connection.catalog, + "drop_table", + side_effect=drop_then_lose_response): + with self.assertRaisesRegex(RuntimeError, "quarantined"): + self.connection.drop_table("reused_quarantine") + + replacement = self.connection.catalog.get_table(quarantine[0]) + self.assertEqual( + "other-owner", + replacement.table_schema.options[_OWNER_ID_OPTION], + ) + original_drop(quarantine[0]) + def test_drop_table_does_not_delete_recreated_companion(self): self.connection.load_from_lerobot( "drop_race", self.image_source) @@ -1749,10 +1809,25 @@ def test_drop_table_validates_group_before_deleting(self): schema.options[_OWNER_ID_OPTION] = "other-owner" self.connection.catalog.drop_table(identifier) self.connection.catalog.create_table(identifier, schema, False) + original_rename = self.connection.catalog.rename_table + failed = [False] - with self.assertRaisesRegex(ValueError, "different table"): - self.connection.drop_table("mixed_owner") + def fail_first_restore(source, target): + if source.get_table_name().startswith("__pypaimon_drop_") \ + and target.get_table_name() == "mixed_owner__versions" \ + and not failed[0]: + failed[0] = True + raise RuntimeError("transient restore failure") + return original_rename(source, target) + + with patch.object( + self.connection.catalog, + "rename_table", + side_effect=fail_first_restore): + with self.assertRaisesRegex(ValueError, "different table"): + self.connection.drop_table("mixed_owner") + self.assertTrue(failed[0]) for name in ( "mixed_owner", "mixed_owner__versions", "mixed_owner__episodes", "mixed_owner__tasks"): From 847c507d118a2941bf49d5e9cb214853f0f2bd18 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 00:25:46 -0700 Subject: [PATCH 21/32] [python] Reconcile uncertain LeRobot table renames --- .../pypaimon/multimodal/lerobot/metadata.py | 147 ++++++++++++------ .../pypaimon/tests/multimodal_lerobot_test.py | 79 ++++++++++ 2 files changed, 181 insertions(+), 45 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 88618ea38a14..11a87bb0df2b 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -39,6 +39,7 @@ _VERSION_ID = "version_id" _OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" _PANDAS_METADATA_OPTION = "pypaimon.lerobot.pandas-metadata" +_MISSING = object() _TABLE_SUFFIXES = { "versions": "__versions", "episodes": "__episodes", @@ -355,31 +356,19 @@ def _drop_import_tables( if isinstance(identifier, Identifier) else Identifier.from_string(str(identifier)) ) - if owned_only: - try: - table = catalog.get_table(source) - except (DatabaseNotExistException, TableNotExistException): - continue - if table.table_schema.options.get( - _OWNER_ID_OPTION) != owner_id: - continue + expected_owner = _table_owner(catalog, source) + if expected_owner is _MISSING: + continue + if owned_only and expected_owner != owner_id: + continue quarantine = Identifier( source.get_database_name(), "__pypaimon_drop_%s" % uuid.uuid4().hex, ) - try: - catalog.rename_table(source, quarantine) - except (DatabaseNotExistException, TableNotExistException): - continue - quarantined.append((source, quarantine)) - - owned = [] - foreign = [] - for source, quarantine in quarantined: - table = catalog.get_table(quarantine) - actual_owner = table.table_schema.options.get(_OWNER_ID_OPTION) - target = owned if actual_owner == owner_id else foreign - target.append((source, quarantine)) + entry = [source, quarantine, expected_owner] + quarantined.append(entry) + if not _move_to_quarantine(catalog, entry): + quarantined.pop() except BaseException as error: failures = _restore_quarantined(catalog, quarantined) if failures: @@ -388,6 +377,8 @@ def _drop_import_tables( % ", ".join(failures)) from error raise + owned = [entry for entry in quarantined if entry[2] == owner_id] + foreign = [entry for entry in quarantined if entry[2] != owner_id] if foreign and not owned_only: error = ValueError( "Refusing to drop %s because it belongs to a different table." @@ -408,50 +399,116 @@ def _drop_import_tables( raise RuntimeError( "LeRobot cleanup left quarantined tables: %s" % ", ".join(restore_failures + [ - str(quarantine) for _, quarantine in drop_failures + str(quarantine) for _, quarantine, _ in drop_failures ])) +def _move_to_quarantine(catalog, entry): + source, quarantine, expected_owner = entry + error = None + rename_attempted = False + for _ in range(2): + try: + source_owner = _table_owner(catalog, source) + quarantine_owner = _table_owner(catalog, quarantine) + if quarantine_owner == expected_owner \ + and source_owner != expected_owner: + return True + if rename_attempted and quarantine_owner is not _MISSING \ + and source_owner != expected_owner: + entry[2] = quarantine_owner + return True + if source_owner is _MISSING and quarantine_owner is _MISSING: + return False + if source_owner != expected_owner \ + or quarantine_owner is not _MISSING: + raise RuntimeError( + "LeRobot table generation changed while quarantining %s." + % source) + rename_attempted = True + catalog.rename_table(source, quarantine) + except BaseException as current_error: + error = current_error + source_owner = _table_owner(catalog, source) + quarantine_owner = _table_owner(catalog, quarantine) + if quarantine_owner == expected_owner \ + and source_owner != expected_owner: + return True + if rename_attempted and quarantine_owner is not _MISSING \ + and source_owner != expected_owner: + entry[2] = quarantine_owner + return True + if source_owner is _MISSING and quarantine_owner is _MISSING: + return False + raise RuntimeError( + "Cannot determine the result of quarantining LeRobot table %s." + % source) from error + + def _restore_quarantined(catalog, tables): failures = [] - for source, quarantine in reversed(tables): - for attempt in range(2): + for source, quarantine, expected_owner in reversed(tables): + restored = False + for _ in range(2): try: - if not _table_exists(catalog, quarantine): + source_owner = _table_owner(catalog, source) + quarantine_owner = _table_owner(catalog, quarantine) + if source_owner == expected_owner \ + and quarantine_owner is _MISSING: + restored = True break - if _table_exists(catalog, source): - failures.append(str(quarantine)) + if source_owner is not _MISSING \ + or quarantine_owner != expected_owner: break catalog.rename_table(quarantine, source) - break except BaseException: - if attempt == 1: - failures.append(str(quarantine)) + pass + if not restored: + try: + restored = _table_owner(catalog, source) == expected_owner \ + and _table_owner(catalog, quarantine) is _MISSING + except BaseException: + pass + if not restored: + failures.append(str(quarantine)) return failures def _drop_quarantined(catalog, tables, owner_id): failures = [] - for source, quarantine in tables: - try: - table = catalog.get_table(quarantine) - if table.table_schema.options.get(_OWNER_ID_OPTION) != owner_id: - failures.append((source, quarantine)) - continue - catalog.drop_table(quarantine) - except (DatabaseNotExistException, TableNotExistException): - pass - except BaseException: - failures.append((source, quarantine)) + for entry in tables: + source, quarantine, expected_owner = entry + dropped = False + for _ in range(2): + try: + source_owner = _table_owner(catalog, source) + quarantine_owner = _table_owner(catalog, quarantine) + if quarantine_owner is _MISSING: + dropped = source_owner != expected_owner + break + if quarantine_owner != expected_owner \ + or expected_owner != owner_id: + break + catalog.drop_table(quarantine) + except BaseException: + pass + if not dropped: + try: + dropped = _table_owner(catalog, quarantine) is _MISSING \ + and _table_owner(catalog, source) != expected_owner + except BaseException: + pass + if not dropped: + failures.append(entry) return failures -def _table_exists(catalog, identifier): +def _table_owner(catalog, identifier): try: - catalog.get_table(identifier) - return True + table = catalog.get_table(identifier) + return table.table_schema.options.get(_OWNER_ID_OPTION) except (DatabaseNotExistException, TableNotExistException): - return False + return _MISSING def _append_arrow(table, data): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 8c4d2689a2f0..24b0a33be503 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1720,6 +1720,40 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.catalog.get_table( self.connection._identifier(name)) + def test_quarantine_rename_unknown_result_is_reconciled(self): + for suffix, error_type in ( + ("runtime", RuntimeError), + ("missing", TableNotExistException)): + name = "rename_lost_%s" % suffix + self.connection.load_from_lerobot(name, self.image_source) + source = Identifier.from_string(self.connection._identifier( + "%s__episodes" % name)) + original_rename = self.connection.catalog.rename_table + injected = [False] + + def rename_then_lose_response(rename_source, target): + result = original_rename(rename_source, target) + if rename_source == source and not injected[0]: + injected[0] = True + if error_type is TableNotExistException: + raise error_type(rename_source) + raise error_type("response lost") + return result + + with patch.object( + self.connection.catalog, + "rename_table", + side_effect=rename_then_lose_response): + self.connection.drop_table(name) + + self.assertTrue(injected[0]) + for table_name in ( + name, "%s__versions" % name, + "%s__episodes" % name, "%s__tasks" % name): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(table_name)) + def test_drop_retry_does_not_delete_reused_quarantine(self): self.connection.load_from_lerobot( "reused_quarantine", self.image_source) @@ -1834,6 +1868,51 @@ def fail_first_restore(source, target): self.connection.catalog.get_table( self.connection._identifier(name)) + def test_restore_rejects_replaced_canonical_generation(self): + self.connection.load_from_lerobot( + "restore_replaced", self.image_source) + tasks = self.connection._identifier("restore_replaced__tasks") + tasks_schema = self.connection.catalog.get_table( + tasks).table_schema.to_schema() + tasks_schema.options = dict(tasks_schema.options) + tasks_schema.options[_OWNER_ID_OPTION] = "foreign-owner" + self.connection.catalog.drop_table(tasks) + self.connection.catalog.create_table(tasks, tasks_schema, False) + + versions = Identifier.from_string( + self.connection._identifier("restore_replaced__versions")) + versions_schema = self.connection.catalog.get_table( + versions).table_schema.to_schema() + versions_schema.options = dict(versions_schema.options) + versions_schema.options[_OWNER_ID_OPTION] = "replacement-owner" + original_drop = self.connection.catalog.drop_table + original_rename = self.connection.catalog.rename_table + injected = [False] + + def replace_before_restore(source, target): + if source.get_table_name().startswith("__pypaimon_drop_") \ + and target == versions and not injected[0]: + injected[0] = True + original_drop(source) + self.connection.catalog.create_table( + target, versions_schema, False) + raise TableNotExistException(source) + return original_rename(source, target) + + with patch.object( + self.connection.catalog, + "rename_table", + side_effect=replace_before_restore): + with self.assertRaisesRegex( + RuntimeError, "Failed to restore quarantined"): + self.connection.drop_table("restore_replaced") + + replacement = self.connection.catalog.get_table(versions) + self.assertEqual( + "replacement-owner", + replacement.table_schema.options[_OWNER_ID_OPTION], + ) + def test_companion_table_cannot_be_dropped_directly(self): self.connection.load_from_lerobot( "direct_drop", self.image_source) From 05f423af3b40db94c059eaa9f37d846e7dab487e Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 00:41:51 -0700 Subject: [PATCH 22/32] [python] Avoid retrying ambiguous LeRobot table drops --- .../pypaimon/multimodal/lerobot/metadata.py | 25 +++++------ .../pypaimon/tests/multimodal_lerobot_test.py | 42 ++++++++++++------- 2 files changed, 38 insertions(+), 29 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 11a87bb0df2b..98ca043fb2fd 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -392,9 +392,6 @@ def _drop_import_tables( restore_failures = _restore_quarantined(catalog, foreign) drop_failures = _drop_quarantined(catalog, owned, owner_id) - if drop_failures: - drop_failures = _drop_quarantined( - catalog, drop_failures, owner_id) if restore_failures or drop_failures: raise RuntimeError( "LeRobot cleanup left quarantined tables: %s" @@ -479,19 +476,17 @@ def _drop_quarantined(catalog, tables, owner_id): for entry in tables: source, quarantine, expected_owner = entry dropped = False - for _ in range(2): - try: - source_owner = _table_owner(catalog, source) - quarantine_owner = _table_owner(catalog, quarantine) - if quarantine_owner is _MISSING: - dropped = source_owner != expected_owner - break - if quarantine_owner != expected_owner \ - or expected_owner != owner_id: - break + try: + source_owner = _table_owner(catalog, source) + quarantine_owner = _table_owner(catalog, quarantine) + if quarantine_owner is _MISSING: + dropped = source_owner != expected_owner + elif quarantine_owner == expected_owner \ + and expected_owner == owner_id: catalog.drop_table(quarantine) - except BaseException: - pass + dropped = True + except BaseException: + pass if not dropped: try: dropped = _table_owner(catalog, quarantine) is _MISSING \ diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 24b0a33be503..97d02ffcb728 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1686,12 +1686,20 @@ def test_drop_table_removes_companion_tables(self): self.connection.catalog.get_table( self.connection._identifier(name)) - def test_drop_table_retries_companion_failure(self): + def test_drop_failure_is_not_retried(self): self.connection.load_from_lerobot( "retry_drop", self.image_source) + episode_name = self.connection._identifier( + "retry_drop__episodes") + episode_table = self.connection.catalog.get_table(episode_name) + expected_owner = episode_table.table_schema.options[ + _OWNER_ID_OPTION] + replacement_schema = episode_table.table_schema.to_schema() + replacement_schema.options = dict(replacement_schema.options) + replacement_schema.options[_OWNER_ID_OPTION] = "other-owner" original_drop = self.connection.catalog.drop_table original_rename = self.connection.catalog.rename_table - failed = [False] + attempts = [0] episode_quarantine = [None] def track_rename(source, target): @@ -1700,10 +1708,13 @@ def track_rename(source, target): return original_rename(source, target) def flaky_drop(identifier, ignore_if_not_exists=False): - if identifier.get_full_name() == episode_quarantine[0] \ - and not failed[0]: - failed[0] = True - raise RuntimeError("injected drop failure") + if identifier.get_full_name() == episode_quarantine[0]: + attempts[0] += 1 + if attempts[0] == 1: + raise RuntimeError("injected drop failure") + original_drop(identifier, ignore_if_not_exists) + self.connection.catalog.create_table( + identifier, replacement_schema, False) return original_drop(identifier, ignore_if_not_exists) with patch.object(self.connection.catalog, "rename_table", @@ -1711,14 +1722,17 @@ def flaky_drop(identifier, ignore_if_not_exists=False): self.connection.catalog, "drop_table", side_effect=flaky_drop): - self.connection.drop_table("retry_drop") - for name in ( - "retry_drop", "retry_drop__versions", - "retry_drop__episodes", "retry_drop__tasks", - "retry_drop__subtasks"): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(name)) + with self.assertRaisesRegex(RuntimeError, "quarantined"): + self.connection.drop_table("retry_drop") + + self.assertEqual(1, attempts[0]) + remaining = self.connection.catalog.get_table( + episode_quarantine[0]) + self.assertEqual( + expected_owner, + remaining.table_schema.options.get(_OWNER_ID_OPTION), + ) + original_drop(episode_quarantine[0]) def test_quarantine_rename_unknown_result_is_reconciled(self): for suffix, error_type in ( From 60c858c4655e6e94634f8baf7655b2e4608e4c18 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 01:16:21 -0700 Subject: [PATCH 23/32] [python] Avoid table renames during LeRobot cleanup --- .../pypaimon/multimodal/lerobot/metadata.py | 168 +++----------- .../pypaimon/tests/multimodal_lerobot_test.py | 215 +++--------------- 2 files changed, 60 insertions(+), 323 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 98ca043fb2fd..26639b1e29ba 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -348,154 +348,42 @@ def _drop_import_tables( identifiers = list( _companion_table_identifiers(frames_table).values()) identifiers.append(frames_table.identifier) - quarantined = [] - try: - for identifier in identifiers: - source = ( - identifier - if isinstance(identifier, Identifier) - else Identifier.from_string(str(identifier)) - ) - expected_owner = _table_owner(catalog, source) - if expected_owner is _MISSING: - continue - if owned_only and expected_owner != owner_id: + owned = [] + for identifier in identifiers: + identifier = ( + identifier + if isinstance(identifier, Identifier) + else Identifier.from_string(str(identifier)) + ) + actual_owner = _table_owner(catalog, identifier) + if actual_owner is _MISSING: + continue + if actual_owner != owner_id: + if owned_only: continue - quarantine = Identifier( - source.get_database_name(), - "__pypaimon_drop_%s" % uuid.uuid4().hex, - ) - entry = [source, quarantine, expected_owner] - quarantined.append(entry) - if not _move_to_quarantine(catalog, entry): - quarantined.pop() - except BaseException as error: - failures = _restore_quarantined(catalog, quarantined) - if failures: - raise RuntimeError( - "Failed to restore quarantined LeRobot tables: %s" - % ", ".join(failures)) from error - raise - - owned = [entry for entry in quarantined if entry[2] == owner_id] - foreign = [entry for entry in quarantined if entry[2] != owner_id] - if foreign and not owned_only: - error = ValueError( - "Refusing to drop %s because it belongs to a different table." - % foreign[0][0]) - failures = _restore_quarantined(catalog, quarantined) - if failures: - raise RuntimeError( - "Failed to restore quarantined LeRobot tables: %s" - % ", ".join(failures)) from error - raise error - - restore_failures = _restore_quarantined(catalog, foreign) - drop_failures = _drop_quarantined(catalog, owned, owner_id) - if restore_failures or drop_failures: - raise RuntimeError( - "LeRobot cleanup left quarantined tables: %s" - % ", ".join(restore_failures + [ - str(quarantine) for _, quarantine, _ in drop_failures - ])) - - -def _move_to_quarantine(catalog, entry): - source, quarantine, expected_owner = entry - error = None - rename_attempted = False - for _ in range(2): - try: - source_owner = _table_owner(catalog, source) - quarantine_owner = _table_owner(catalog, quarantine) - if quarantine_owner == expected_owner \ - and source_owner != expected_owner: - return True - if rename_attempted and quarantine_owner is not _MISSING \ - and source_owner != expected_owner: - entry[2] = quarantine_owner - return True - if source_owner is _MISSING and quarantine_owner is _MISSING: - return False - if source_owner != expected_owner \ - or quarantine_owner is not _MISSING: - raise RuntimeError( - "LeRobot table generation changed while quarantining %s." - % source) - rename_attempted = True - catalog.rename_table(source, quarantine) - except BaseException as current_error: - error = current_error - source_owner = _table_owner(catalog, source) - quarantine_owner = _table_owner(catalog, quarantine) - if quarantine_owner == expected_owner \ - and source_owner != expected_owner: - return True - if rename_attempted and quarantine_owner is not _MISSING \ - and source_owner != expected_owner: - entry[2] = quarantine_owner - return True - if source_owner is _MISSING and quarantine_owner is _MISSING: - return False - raise RuntimeError( - "Cannot determine the result of quarantining LeRobot table %s." - % source) from error - - -def _restore_quarantined(catalog, tables): - failures = [] - for source, quarantine, expected_owner in reversed(tables): - restored = False - for _ in range(2): - try: - source_owner = _table_owner(catalog, source) - quarantine_owner = _table_owner(catalog, quarantine) - if source_owner == expected_owner \ - and quarantine_owner is _MISSING: - restored = True - break - if source_owner is not _MISSING \ - or quarantine_owner != expected_owner: - break - catalog.rename_table(quarantine, source) - except BaseException: - pass - if not restored: - try: - restored = _table_owner(catalog, source) == expected_owner \ - and _table_owner(catalog, quarantine) is _MISSING - except BaseException: - pass - if not restored: - failures.append(str(quarantine)) - return failures - + raise ValueError( + "Refusing to drop %s because it belongs to a different " + "table." % identifier) + owned.append(identifier) -def _drop_quarantined(catalog, tables, owner_id): failures = [] - for entry in tables: - source, quarantine, expected_owner = entry - dropped = False + for identifier in owned: try: - source_owner = _table_owner(catalog, source) - quarantine_owner = _table_owner(catalog, quarantine) - if quarantine_owner is _MISSING: - dropped = source_owner != expected_owner - elif quarantine_owner == expected_owner \ - and expected_owner == owner_id: - catalog.drop_table(quarantine) - dropped = True + if _table_owner(catalog, identifier) != owner_id: + failures.append(str(identifier)) + continue + catalog.drop_table(identifier) except BaseException: - pass - if not dropped: try: - dropped = _table_owner(catalog, quarantine) is _MISSING \ - and _table_owner(catalog, source) != expected_owner + if _table_owner(catalog, identifier) is _MISSING: + continue except BaseException: pass - if not dropped: - failures.append(entry) - return failures + failures.append(str(identifier)) + if failures: + raise RuntimeError( + "LeRobot cleanup could not drop tables: %s" + % ", ".join(failures)) def _table_owner(catalog, identifier): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 97d02ffcb728..5e6bb0ad1d4b 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1689,163 +1689,67 @@ def test_drop_table_removes_companion_tables(self): def test_drop_failure_is_not_retried(self): self.connection.load_from_lerobot( "retry_drop", self.image_source) - episode_name = self.connection._identifier( - "retry_drop__episodes") + episode_name = Identifier.from_string(self.connection._identifier( + "retry_drop__episodes")) episode_table = self.connection.catalog.get_table(episode_name) expected_owner = episode_table.table_schema.options[ _OWNER_ID_OPTION] - replacement_schema = episode_table.table_schema.to_schema() - replacement_schema.options = dict(replacement_schema.options) - replacement_schema.options[_OWNER_ID_OPTION] = "other-owner" original_drop = self.connection.catalog.drop_table - original_rename = self.connection.catalog.rename_table attempts = [0] - episode_quarantine = [None] - - def track_rename(source, target): - if source.get_table_name().endswith("__episodes"): - episode_quarantine[0] = target.get_full_name() - return original_rename(source, target) def flaky_drop(identifier, ignore_if_not_exists=False): - if identifier.get_full_name() == episode_quarantine[0]: + if identifier == episode_name: attempts[0] += 1 - if attempts[0] == 1: - raise RuntimeError("injected drop failure") - original_drop(identifier, ignore_if_not_exists) - self.connection.catalog.create_table( - identifier, replacement_schema, False) + raise RuntimeError("injected drop failure") return original_drop(identifier, ignore_if_not_exists) - with patch.object(self.connection.catalog, "rename_table", - side_effect=track_rename), patch.object( - self.connection.catalog, - "drop_table", - side_effect=flaky_drop): - with self.assertRaisesRegex(RuntimeError, "quarantined"): + with patch.object( + self.connection.catalog, + "drop_table", + side_effect=flaky_drop): + with self.assertRaisesRegex(RuntimeError, "could not drop"): self.connection.drop_table("retry_drop") self.assertEqual(1, attempts[0]) - remaining = self.connection.catalog.get_table( - episode_quarantine[0]) + remaining = self.connection.catalog.get_table(episode_name) self.assertEqual( expected_owner, remaining.table_schema.options.get(_OWNER_ID_OPTION), ) - original_drop(episode_quarantine[0]) - - def test_quarantine_rename_unknown_result_is_reconciled(self): - for suffix, error_type in ( - ("runtime", RuntimeError), - ("missing", TableNotExistException)): - name = "rename_lost_%s" % suffix - self.connection.load_from_lerobot(name, self.image_source) - source = Identifier.from_string(self.connection._identifier( - "%s__episodes" % name)) - original_rename = self.connection.catalog.rename_table - injected = [False] - - def rename_then_lose_response(rename_source, target): - result = original_rename(rename_source, target) - if rename_source == source and not injected[0]: - injected[0] = True - if error_type is TableNotExistException: - raise error_type(rename_source) - raise error_type("response lost") - return result + original_drop(episode_name) - with patch.object( - self.connection.catalog, - "rename_table", - side_effect=rename_then_lose_response): - self.connection.drop_table(name) - - self.assertTrue(injected[0]) - for table_name in ( - name, "%s__versions" % name, - "%s__episodes" % name, "%s__tasks" % name): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(table_name)) - - def test_drop_retry_does_not_delete_reused_quarantine(self): + def test_drop_response_loss_is_reconciled_without_retry(self): self.connection.load_from_lerobot( - "reused_quarantine", self.image_source) - source = self.connection._identifier( - "reused_quarantine__episodes") - source_name = Identifier.from_string(source).get_full_name() - replacement_schema = self.connection.catalog.get_table( - source).table_schema.to_schema() - replacement_schema.options = dict(replacement_schema.options) - replacement_schema.options[_OWNER_ID_OPTION] = "other-owner" + "lost_drop_response", self.image_source) + episode_name = Identifier.from_string(self.connection._identifier( + "lost_drop_response__episodes")) original_drop = self.connection.catalog.drop_table - original_rename = self.connection.catalog.rename_table - quarantine = [None] - injected = [False] - - def track_rename(rename_source, target): - if rename_source.get_full_name() == source_name: - quarantine[0] = target - return original_rename(rename_source, target) + attempts = [0] def drop_then_lose_response(identifier, ignore_if_not_exists=False): - if identifier == quarantine[0] and not injected[0]: - injected[0] = True - original_drop(identifier, ignore_if_not_exists) - self.connection.catalog.create_table( - identifier, replacement_schema, False) + result = original_drop(identifier, ignore_if_not_exists) + if identifier == episode_name: + attempts[0] += 1 raise RuntimeError("response lost") - return original_drop(identifier, ignore_if_not_exists) + return result with patch.object( self.connection.catalog, - "rename_table", - side_effect=track_rename), patch.object( - self.connection.catalog, - "drop_table", - side_effect=drop_then_lose_response): - with self.assertRaisesRegex(RuntimeError, "quarantined"): - self.connection.drop_table("reused_quarantine") - - replacement = self.connection.catalog.get_table(quarantine[0]) - self.assertEqual( - "other-owner", - replacement.table_schema.options[_OWNER_ID_OPTION], - ) - original_drop(quarantine[0]) + "drop_table", + side_effect=drop_then_lose_response): + self.connection.drop_table("lost_drop_response") + + self.assertEqual(1, attempts[0]) - def test_drop_table_does_not_delete_recreated_companion(self): + def test_drop_table_does_not_rename_table_directories(self): self.connection.load_from_lerobot( - "drop_race", self.image_source) - identifier = self.connection._identifier("drop_race__versions") - old_table = self.connection.catalog.get_table(identifier) - replacement_schema = old_table.table_schema.to_schema() - replacement_schema.options = dict(replacement_schema.options) - replacement_schema.options[_OWNER_ID_OPTION] = "other-owner" - original_rename = self.connection.catalog.rename_table - replaced = [False] - - def replace_before_rename(source, target): - if source.get_full_name() == identifier and not replaced[0]: - replaced[0] = True - self.connection.catalog.drop_table(source) - self.connection.catalog.create_table( - source, replacement_schema, False) - return original_rename(source, target) + "direct_group_drop", self.image_source) with patch.object( self.connection.catalog, "rename_table", - side_effect=replace_before_rename): - with self.assertRaisesRegex(ValueError, "different table"): - self.connection.drop_table("drop_race") - - replacement = self.connection.catalog.get_table(identifier) - self.assertEqual( - "other-owner", - replacement.table_schema.options[_OWNER_ID_OPTION], - ) - self.connection.get_table("drop_race") + side_effect=AssertionError("rename must not be used")): + self.connection.drop_table("direct_group_drop") def test_drop_table_validates_group_before_deleting(self): self.connection.load_from_lerobot( @@ -1857,76 +1761,21 @@ def test_drop_table_validates_group_before_deleting(self): schema.options[_OWNER_ID_OPTION] = "other-owner" self.connection.catalog.drop_table(identifier) self.connection.catalog.create_table(identifier, schema, False) - original_rename = self.connection.catalog.rename_table - failed = [False] - - def fail_first_restore(source, target): - if source.get_table_name().startswith("__pypaimon_drop_") \ - and target.get_table_name() == "mixed_owner__versions" \ - and not failed[0]: - failed[0] = True - raise RuntimeError("transient restore failure") - return original_rename(source, target) with patch.object( self.connection.catalog, - "rename_table", - side_effect=fail_first_restore): + "drop_table", + wraps=self.connection.catalog.drop_table) as drop: with self.assertRaisesRegex(ValueError, "different table"): self.connection.drop_table("mixed_owner") + drop.assert_not_called() - self.assertTrue(failed[0]) for name in ( "mixed_owner", "mixed_owner__versions", "mixed_owner__episodes", "mixed_owner__tasks"): self.connection.catalog.get_table( self.connection._identifier(name)) - def test_restore_rejects_replaced_canonical_generation(self): - self.connection.load_from_lerobot( - "restore_replaced", self.image_source) - tasks = self.connection._identifier("restore_replaced__tasks") - tasks_schema = self.connection.catalog.get_table( - tasks).table_schema.to_schema() - tasks_schema.options = dict(tasks_schema.options) - tasks_schema.options[_OWNER_ID_OPTION] = "foreign-owner" - self.connection.catalog.drop_table(tasks) - self.connection.catalog.create_table(tasks, tasks_schema, False) - - versions = Identifier.from_string( - self.connection._identifier("restore_replaced__versions")) - versions_schema = self.connection.catalog.get_table( - versions).table_schema.to_schema() - versions_schema.options = dict(versions_schema.options) - versions_schema.options[_OWNER_ID_OPTION] = "replacement-owner" - original_drop = self.connection.catalog.drop_table - original_rename = self.connection.catalog.rename_table - injected = [False] - - def replace_before_restore(source, target): - if source.get_table_name().startswith("__pypaimon_drop_") \ - and target == versions and not injected[0]: - injected[0] = True - original_drop(source) - self.connection.catalog.create_table( - target, versions_schema, False) - raise TableNotExistException(source) - return original_rename(source, target) - - with patch.object( - self.connection.catalog, - "rename_table", - side_effect=replace_before_restore): - with self.assertRaisesRegex( - RuntimeError, "Failed to restore quarantined"): - self.connection.drop_table("restore_replaced") - - replacement = self.connection.catalog.get_table(versions) - self.assertEqual( - "replacement-owner", - replacement.table_schema.options[_OWNER_ID_OPTION], - ) - def test_companion_table_cannot_be_dropped_directly(self): self.connection.load_from_lerobot( "direct_drop", self.image_source) From a33837ae4ac1c8ca6affd1b8e8197d5eaa556b1e Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 01:35:48 -0700 Subject: [PATCH 24/32] [python] Keep LeRobot root until companions drop --- paimon-python/pypaimon/multimodal/lerobot/metadata.py | 9 ++++++++- .../pypaimon/tests/multimodal_lerobot_test.py | 10 +++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 26639b1e29ba..e91601a3cb4c 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -345,9 +345,14 @@ def _manifest_row( def _drop_import_tables( catalog, frames_table, owner_id, owned_only=False): + root_identifier = ( + frames_table.identifier + if isinstance(frames_table.identifier, Identifier) + else Identifier.from_string(str(frames_table.identifier)) + ) identifiers = list( _companion_table_identifiers(frames_table).values()) - identifiers.append(frames_table.identifier) + identifiers.append(root_identifier) owned = [] for identifier in identifiers: identifier = ( @@ -368,6 +373,8 @@ def _drop_import_tables( failures = [] for identifier in owned: + if identifier == root_identifier and failures: + break try: if _table_owner(catalog, identifier) != owner_id: failures.append(str(identifier)) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 5e6bb0ad1d4b..1d7fbcff1e72 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1716,7 +1716,15 @@ def flaky_drop(identifier, ignore_if_not_exists=False): expected_owner, remaining.table_schema.options.get(_OWNER_ID_OPTION), ) - original_drop(episode_name) + self.connection.get_table("retry_drop") + + self.connection.drop_table("retry_drop") + for name in ( + "retry_drop", "retry_drop__versions", + "retry_drop__episodes", "retry_drop__tasks"): + with self.assertRaises(TableNotExistException): + self.connection.catalog.get_table( + self.connection._identifier(name)) def test_drop_response_loss_is_reconciled_without_retry(self): self.connection.load_from_lerobot( From a5f744f3062a6f7abdbd91c8cc0750fb0ebfd327 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 02:03:39 -0700 Subject: [PATCH 25/32] [docs] Simplify LeRobot version description --- docs/docs/pypaimon/multimodal-api.mdx | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index a79b3766e8ba..45be0beb30c5 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -652,9 +652,8 @@ version_id = conn.load_from_lerobot( ) ``` -The component tables retain the native LeRobot V3 schemas and contain no -version columns. A READY row in `
__versions` identifies a release; -reading the same tag from its component tables reconstructs that release. +A READY row in `
__versions` identifies a release; reading the same tag +from its component tables reconstructs that release. The one-time importer requires a new target table. A failed import removes the table group so the call can be retried. `drop_table()` removes the table group; companion tables cannot be dropped separately through `MultimodalConnection`. From 27f7d393651cfacdbdee08e5f389fb588c500f93 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 02:09:38 -0700 Subject: [PATCH 26/32] [docs] Avoid promising transactional import cleanup --- docs/docs/pypaimon/multimodal-api.mdx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 45be0beb30c5..282871f52e47 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -654,9 +654,9 @@ version_id = conn.load_from_lerobot( A READY row in `
__versions` identifies a release; reading the same tag from its component tables reconstructs that release. -The one-time importer requires a new target table. A failed import removes the -table group so the call can be retried. `drop_table()` removes the table group; -companion tables cannot be dropped separately through `MultimodalConnection`. +The one-time importer requires a new target table. `drop_table()` removes the +table group; companion tables cannot be dropped separately through +`MultimodalConnection`. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. From f99e287f5d3bfff9afc254e289c81371712419b8 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 02:22:59 -0700 Subject: [PATCH 27/32] [python] Notify callbacks for duplicate commits --- .../pypaimon/tests/multimodal_lerobot_test.py | 32 +++++++++++++ .../tests/write/commit_callback_test.py | 45 +++++++++++++++++++ .../pypaimon/write/file_store_commit.py | 39 +++++++++++----- 3 files changed, 106 insertions(+), 10 deletions(-) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 1d7fbcff1e72..a92417098ab7 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1438,6 +1438,38 @@ def test_failed_publication_is_cleaned_and_can_retry(self): "failed_publish", self.image_source) self.assertEqual(1, version_id) + def test_lost_commit_response_does_not_fail_import(self): + from pypaimon.snapshot.renaming_snapshot_commit import ( + RenamingSnapshotCommit, + ) + + real_commit = RenamingSnapshotCommit.commit + lost = [False] + + def commit_then_lose_response( + snapshot_commit, base_snapshot_uuid, snapshot, statistics): + result = real_commit( + snapshot_commit, base_snapshot_uuid, snapshot, statistics) + if result and not lost[0]: + lost[0] = True + raise TimeoutError("lost snapshot commit response") + return result + + with patch.object( + RenamingSnapshotCommit, + "commit", + new=commit_then_lose_response): + version_id = self.connection.load_from_lerobot( + "lost_commit_response", self.image_source) + + self.assertTrue(lost[0]) + self.assertEqual(1, version_id) + self.assertEqual( + ["PENDING", "READY"], + [row["status"] for row in _catalog_rows( + self.connection, "lost_commit_response__versions")], + ) + def test_stale_companion_does_not_block_retry(self): self.connection.load_from_lerobot( "other_group", self.image_source) diff --git a/paimon-python/pypaimon/tests/write/commit_callback_test.py b/paimon-python/pypaimon/tests/write/commit_callback_test.py index 036031edefec..9addc1e812db 100644 --- a/paimon-python/pypaimon/tests/write/commit_callback_test.py +++ b/paimon-python/pypaimon/tests/write/commit_callback_test.py @@ -115,6 +115,51 @@ def test_callback_receives_correct_snapshot_data(self): table_write.close() table_commit.close() + def test_callback_invoked_after_lost_commit_response(self): + table = self._create_table( + 'test_callback_response_loss', + options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }, + ) + write_builder = table.new_batch_write_builder() + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + + callback = RecordingCallback() + table_commit.add_commit_callback(callback) + real_commit = table_commit.file_store_commit.snapshot_commit.commit + attempts = [] + + def commit_then_lose_response(base_snapshot_uuid, snapshot, statistics): + attempts.append(snapshot.id) + self.assertTrue(real_commit( + base_snapshot_uuid, snapshot, statistics)) + raise TimeoutError('lost snapshot commit response') + + table_commit.file_store_commit.snapshot_commit.commit = ( + commit_then_lose_response) + table_commit.file_store_commit._commit_retry_wait = lambda _: None + + data = pa.Table.from_pydict({ + 'id': [1, 2], + 'name': ['a', 'b'], + 'dt': ['p1', 'p1'], + }, schema=self.pa_schema) + table_write.write_arrow(data) + table_commit.commit(table_write.prepare_commit()) + + self.assertEqual([1], attempts) + self.assertEqual(1, len(callback.contexts)) + self.assertEqual(1, callback.contexts[0].snapshot.id) + self.assertGreater(len(callback.contexts[0].commit_entries), 0) + for entry in callback.contexts[0].commit_entries: + self.assertIsNotNone(entry.file.first_row_id) + + table_write.close() + table_commit.close() + def test_multiple_callbacks(self): table = self._create_table('test_multi_callbacks') write_builder = table.new_batch_write_builder() diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 5f8f054aa812..f3fb278e15a5 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -599,7 +599,11 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str commit_result_may_be_uncertain: bool = False) -> CommitResult: start_millis = int(time.time() * 1000) if self._is_duplicate_commit( - retry_result, latest_snapshot, commit_identifier, commit_kind): + retry_result, + latest_snapshot, + commit_identifier, + commit_kind, + notify_callbacks=True): return SuccessResult() latest_snapshot_id = latest_snapshot.id if latest_snapshot else 0 @@ -862,14 +866,8 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str commit_kind, ) - if self.commit_callbacks: - context = CommitCallbackContext( - snapshot=snapshot_data, - commit_entries=commit_entries, - identifier=commit_identifier, - ) - for callback in self.commit_callbacks: - callback.call(context) + self._notify_commit_callbacks( + snapshot_data, commit_entries, commit_identifier) return SuccessResult() @@ -927,7 +925,8 @@ def _is_duplicate_commit( retry_result, latest_snapshot, commit_identifier, - commit_kind) -> bool: + commit_kind, + notify_callbacks=False) -> bool: if (isinstance(retry_result, CommitFailRetryResult) and latest_snapshot is not None): start_check_snapshot_id = 1 # Snapshot.FIRST_SNAPSHOT_ID @@ -953,9 +952,29 @@ def _is_duplicate_commit( f"Commit already completed (snapshot {snapshot_id}), " f"user: {self.commit_user}, identifier: {commit_identifier}" ) + if notify_callbacks and self.commit_callbacks: + entries = [] + for manifest in self.manifest_list_manager.read_delta( + snapshot): + entries.extend(self.manifest_file_manager.read( + manifest.file_name, drop_stats=False)) + self._notify_commit_callbacks( + snapshot, entries, commit_identifier) return True return False + def _notify_commit_callbacks( + self, snapshot, commit_entries, commit_identifier): + if not self.commit_callbacks: + return + context = CommitCallbackContext( + snapshot=snapshot, + commit_entries=commit_entries, + identifier=commit_identifier, + ) + for callback in self.commit_callbacks: + callback.call(context) + def _create_dynamic_partition_filter(self, commit_messages: List[CommitMessage]): """Build a partition filter from the unique partitions present in commit_messages.""" predicate_builder = PredicateBuilder(self.table.partition_keys_fields) From e1ebdcbcee87381b7e385489d47857957a3bb238 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 02:36:33 -0700 Subject: [PATCH 28/32] [python] Keep failed LeRobot imports for inspection --- docs/docs/pypaimon/multimodal-api.mdx | 4 +- .../pypaimon/multimodal/connection.py | 32 +- .../pypaimon/multimodal/lerobot/api.py | 114 +++---- .../pypaimon/multimodal/lerobot/metadata.py | 124 ++------ .../pypaimon/tests/multimodal_lerobot_test.py | 281 +++--------------- 5 files changed, 93 insertions(+), 462 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 282871f52e47..30a769c15858 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -654,9 +654,7 @@ version_id = conn.load_from_lerobot( A READY row in `
__versions` identifies a release; reading the same tag from its component tables reconstructs that release. -The one-time importer requires a new target table. `drop_table()` removes the -table group; companion tables cannot be dropped separately through -`MultimodalConnection`. +The one-time importer requires a new target table. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 99ca8e8360d4..31a1bdc88490 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -26,7 +26,6 @@ TableAlreadyExistException, TableNotExistException, ) -from pypaimon.common.identifier import Identifier from pypaimon.multimodal.table import MultimodalTable, _to_arrow_table _DEFAULT_OPTIONS = { @@ -179,37 +178,8 @@ def load_from_rosbag( ) def drop_table(self, name: str, ignore_if_not_exists: bool = False): - identifier = self._identifier(name) - owner_id = None - try: - from pypaimon.multimodal.lerobot.metadata import ( - _OWNER_ID_OPTION, - _is_managed_root, - ) - raw_table = self.catalog.get_table(identifier) - table_options = raw_table.table_schema.options - managed_root = _is_managed_root(table_options) - if _OWNER_ID_OPTION in table_options and not managed_root: - raise ValueError( - "%s is a managed LeRobot companion table; drop its " - "frame table instead." % identifier) - if managed_root: - if Identifier.from_string( - identifier).get_branch_name() is not None: - raise ValueError( - "Dropping a managed LeRobot table branch is not " - "supported; drop the branch through the Catalog.") - owner_id = table_options.get(_OWNER_ID_OPTION) - except (DatabaseNotExistException, TableNotExistException): - pass - - if owner_id is not None: - from pypaimon.multimodal.lerobot.metadata import \ - _drop_import_tables - _drop_import_tables(self.catalog, raw_table, owner_id) - return self.catalog.drop_table( - identifier, + self._identifier(name), ignore_if_not_exists=ignore_if_not_exists, ) diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 2caeb4a715ad..4a47acb9d176 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -20,18 +20,11 @@ import sys from typing import Mapping, Optional -from pypaimon.catalog.catalog_exception import ( - DatabaseNotExistException, - TableAlreadyExistException, - TableNotExistException, -) +from pypaimon.catalog.catalog_exception import TableAlreadyExistException from pypaimon.multimodal.lerobot.metadata import ( - _OWNER_ID_OPTION, _append_arrow_tables, - _drop_import_tables, _load_dataset_metadata, _managed_table_options, - _new_owner_id, _prepare_metadata_tables, _positive_integer, _publish_dataset, @@ -134,56 +127,41 @@ def _import_dataset( batch_size, options, metadata): - table, owner_id = _create_target_table( + table = _create_target_table( connection, table_name, source_schema, options) - try: - tables = _prepare_metadata_tables( - connection, table.raw_table, owner_id, metadata) - version_id = 1 - _reserve_dataset_version( - tables["versions"], - version_id, - metadata, - ) - 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, + tables = _prepare_metadata_tables( + connection, table.raw_table, metadata) + version_id = 1 + _reserve_dataset_version( + tables["versions"], + version_id, + metadata, + ) + 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, - table.identifier, - frames_snapshot_id, - episodes_snapshot_id, ) - return version_id - except BaseException as error: - try: - _drop_import_tables( - connection.catalog, - table.raw_table, - owner_id, - owned_only=True, - ) - except BaseException as cleanup_error: - raise RuntimeError( - "LeRobot import failed and cleanup also failed: %s" - % cleanup_error - ) from error - raise + _publish_dataset( + connection, + tables, + version_id, + metadata, + table.identifier, + frames_snapshot_id, + episodes_snapshot_id, + ) + return version_id def _validated_counts(info, source): @@ -217,10 +195,9 @@ def _required_count(info, name, source): def _create_target_table( connection, table_name, source_schema, options): - owner_id = _new_owner_id() create_options = dict(options or {}) managed_options = _managed_table_options( - connection._identifier(table_name), owner_id) + connection._identifier(table_name)) reserved_options = set(managed_options).intersection(create_options) if reserved_options: raise ValueError( @@ -238,25 +215,4 @@ def _create_target_table( "LeRobot target %s already exists; use a new target table." % connection._identifier(table_name) ) from error - except BaseException as error: - try: - frames_table = connection.catalog.get_table( - connection._identifier(table_name)) - except (DatabaseNotExistException, TableNotExistException): - frames_table = None - if frames_table is not None and frames_table.table_schema.options.get( - _OWNER_ID_OPTION) == owner_id: - try: - _drop_import_tables( - connection.catalog, - frames_table, - owner_id, - owned_only=True, - ) - except BaseException as cleanup_error: - raise RuntimeError( - "LeRobot target creation failed and cleanup also failed: " - "%s" % cleanup_error - ) from error - raise - return table, owner_id + return table diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index e91601a3cb4c..3e2c500ab6d1 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -19,7 +19,6 @@ from array import array import json import numbers -import uuid from pathlib import Path import pyarrow as pa @@ -37,9 +36,7 @@ _VERSION_ID = "version_id" -_OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" _PANDAS_METADATA_OPTION = "pypaimon.lerobot.pandas-metadata" -_MISSING = object() _TABLE_SUFFIXES = { "versions": "__versions", "episodes": "__episodes", @@ -145,10 +142,6 @@ def _load_dataset_metadata(dataset, info, source): } -def _new_owner_id(): - return uuid.uuid4().hex - - def _companion_identifier(frames_identifier, suffix): identifier = ( frames_identifier @@ -174,23 +167,18 @@ def _quote_identifier_part(value): return "`%s`" % value if "." in value else value -def _managed_table_options(frames_identifier, owner_id): +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 = {_OWNER_ID_OPTION: owner_id} + result = {} for name, suffix in _TABLE_SUFFIXES.items(): result[_COMPANION_OPTION_KEYS[name]] = _companion_identifier( frames_identifier, suffix) return result -def _is_managed_root(options): - return _OWNER_ID_OPTION in options and all( - key in options for key in _COMPANION_OPTION_KEYS.values()) - - def _companion_table_identifiers(frames_table): options = frames_table.table_schema.options identifiers = {} @@ -204,7 +192,7 @@ def _companion_table_identifiers(frames_table): return identifiers -def _prepare_metadata_tables(connection, frames_table, owner_id, metadata): +def _prepare_metadata_tables(connection, frames_table, metadata): schemas = { "versions": _VERSIONS_SCHEMA, "episodes": metadata["episodes_schema"], @@ -225,41 +213,23 @@ def _prepare_metadata_tables(connection, frames_table, owner_id, metadata): 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: - table = connection.catalog.get_table(identifier) - except (DatabaseNotExistException, TableNotExistException): - options = { - "bucket": "-1", - _OWNER_ID_OPTION: owner_id, - } - 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: - pass - table = connection.catalog.get_table(identifier) - if table.table_schema.primary_keys: - raise ValueError( - "LeRobot metadata table %s must be append-only." % identifier) - actual = _target_schema(table) - if not actual.equals(schema, check_metadata=False): + connection.catalog.create_table( + identifier, paimon_schema, False) + except TableAlreadyExistException as error: raise ValueError( - "LeRobot metadata table %s has schema %s; expected %s." - % (identifier, actual, schema)) - actual_owner_id = table.table_schema.options.get(_OWNER_ID_OPTION) - if actual_owner_id != owner_id: - raise ValueError( - "LeRobot metadata table %s belongs to a different target " - "table. Drop the stale companion tables before importing." - % identifier) + "LeRobot metadata table %s already exists." % identifier + ) from error + table = connection.catalog.get_table(identifier) tables[name] = table return tables @@ -343,64 +313,6 @@ def _manifest_row( } -def _drop_import_tables( - catalog, frames_table, owner_id, owned_only=False): - root_identifier = ( - frames_table.identifier - if isinstance(frames_table.identifier, Identifier) - else Identifier.from_string(str(frames_table.identifier)) - ) - identifiers = list( - _companion_table_identifiers(frames_table).values()) - identifiers.append(root_identifier) - owned = [] - for identifier in identifiers: - identifier = ( - identifier - if isinstance(identifier, Identifier) - else Identifier.from_string(str(identifier)) - ) - actual_owner = _table_owner(catalog, identifier) - if actual_owner is _MISSING: - continue - if actual_owner != owner_id: - if owned_only: - continue - raise ValueError( - "Refusing to drop %s because it belongs to a different " - "table." % identifier) - owned.append(identifier) - - failures = [] - for identifier in owned: - if identifier == root_identifier and failures: - break - try: - if _table_owner(catalog, identifier) != owner_id: - failures.append(str(identifier)) - continue - catalog.drop_table(identifier) - except BaseException: - try: - if _table_owner(catalog, identifier) is _MISSING: - continue - except BaseException: - pass - failures.append(str(identifier)) - if failures: - raise RuntimeError( - "LeRobot cleanup could not drop tables: %s" - % ", ".join(failures)) - - -def _table_owner(catalog, identifier): - try: - table = catalog.get_table(identifier) - return table.table_schema.options.get(_OWNER_ID_OPTION) - except (DatabaseNotExistException, TableNotExistException): - return _MISSING - - def _append_arrow(table, data): return _append_arrow_tables(table, [data]) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index a92417098ab7..901f19b50e41 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -45,7 +45,6 @@ _restore_pandas_metadata, _subtask_indices, _validated_episode_tables, - _OWNER_ID_OPTION, ) from pypaimon.multimodal.lerobot.loader import ( _image_bytes, @@ -100,7 +99,7 @@ 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", "owner") + _managed_table_options("db.robot$branch_dev") def test_companion_identifier_preserves_quoted_components(self): name = _companion_identifier( @@ -1264,12 +1263,12 @@ def test_frame_controls_must_match_published_episode_metadata(self): with self.assertRaisesRegex( ValueError, "has %s" % column): self.connection.load_from_lerobot(table_name, source) - for suffix in ( - "", "__versions", "__episodes", "__tasks", - "__subtasks"): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(table_name + suffix)) + self.connection.get_table(table_name) + self.assertEqual( + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, table_name + "__versions")], + ) def test_task_text_remains_in_published_task_mapping(self): source = self.temp_dir / "reordered_tasks" @@ -1316,8 +1315,12 @@ def test_episode_tasks_must_exactly_match_frame_tasks(self): ValueError, "declares task indices"): self.connection.load_from_lerobot( "extra_episode_task", source) - with self.assertRaises(TableNotExistException): - self.connection.get_table("extra_episode_task") + self.connection.get_table("extra_episode_task") + self.assertEqual( + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, "extra_episode_task__versions")], + ) def test_nonempty_dataset_cannot_publish_without_tasks(self): source = self.temp_dir / "missing_tasks" @@ -1342,8 +1345,12 @@ def test_nonempty_dataset_cannot_publish_without_tasks(self): with self.assertRaisesRegex(ValueError, "task_index"): self.connection.load_from_lerobot("missing_tasks", source) - with self.assertRaises(TableNotExistException): - self.connection.get_table("missing_tasks") + self.connection.get_table("missing_tasks") + self.assertEqual( + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, "missing_tasks__versions")], + ) def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" @@ -1420,7 +1427,7 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): table.raw_table.tag_manager().get(tag).id, ) - def test_failed_publication_is_cleaned_and_can_retry(self): + def test_failed_publication_remains_pending(self): with patch( "pypaimon.multimodal.lerobot.api._publish_dataset", side_effect=RuntimeError("publish failed")): @@ -1428,15 +1435,12 @@ def test_failed_publication_is_cleaned_and_can_retry(self): self.connection.load_from_lerobot( "failed_publish", self.image_source) - for suffix in ( - "", "__versions", "__episodes", "__tasks", "__subtasks"): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier("failed_publish" + suffix)) - - version_id = self.connection.load_from_lerobot( - "failed_publish", self.image_source) - self.assertEqual(1, version_id) + self.connection.get_table("failed_publish") + self.assertEqual( + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, "failed_publish__versions")], + ) def test_lost_commit_response_does_not_fail_import(self): from pypaimon.snapshot.renaming_snapshot_commit import ( @@ -1470,38 +1474,19 @@ def commit_then_lose_response( self.connection, "lost_commit_response__versions")], ) - def test_stale_companion_does_not_block_retry(self): + def test_existing_companion_is_rejected(self): self.connection.load_from_lerobot( "other_group", self.image_source) stale = self.connection._identifier("stale__tasks") - stale_name = Identifier.from_string(stale).get_full_name() self.connection.catalog.rename_table( self.connection._identifier("other_group__tasks"), stale) - original_rename = self.connection.catalog.rename_table - renamed_sources = [] - - def track_rename(source, target): - renamed_sources.append(source.get_full_name()) - return original_rename(source, target) - - with patch.object( - self.connection.catalog, - "rename_table", - side_effect=track_rename): - with self.assertRaisesRegex(ValueError, "different target"): - self.connection.load_from_lerobot( - "stale", self.image_source) - self.assertNotIn(stale_name, renamed_sources) - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier("stale")) + with self.assertRaisesRegex(ValueError, "already exists"): + self.connection.load_from_lerobot( + "stale", self.image_source) - self.connection.catalog.drop_table(stale) - self.assertEqual( - 1, - 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"): @@ -1518,7 +1503,7 @@ def test_invalid_target_options_do_not_leave_table(self): "invalid_options", self.image_source) self.assertEqual(1, version_id) - def test_target_open_failure_is_cleaned_and_can_retry(self): + def test_target_open_failure_leaves_created_table(self): original_get = self.connection.get_table failed = [False] @@ -1533,35 +1518,8 @@ def fail_once(name): with self.assertRaisesRegex(RuntimeError, "get failed"): self.connection.load_from_lerobot( "failed_open", self.image_source) - version_id = self.connection.load_from_lerobot( - "failed_open", self.image_source) - - self.assertEqual(1, version_id) - def test_target_open_failure_does_not_drop_another_owner(self): - original_create = self.connection.create_table - - def create_other_owner(*args, **kwargs): - options = dict(kwargs["options"]) - options[_OWNER_ID_OPTION] = "other-owner" - kwargs["options"] = options - original_create(*args, **kwargs) - raise RuntimeError("get failed") - - with patch.object( - self.connection, - "create_table", - side_effect=create_other_owner): - with self.assertRaisesRegex(RuntimeError, "get failed"): - self.connection.load_from_lerobot( - "other_owner", self.image_source) - - table = self.connection.catalog.get_table( - self.connection._identifier("other_owner")) - self.assertEqual( - "other-owner", - table.table_schema.options[_OWNER_ID_OPTION], - ) + self.connection.get_table("failed_open") def test_dataset_close_failure_does_not_override_success(self): from pypaimon.multimodal.lerobot import api @@ -1694,176 +1652,13 @@ def append_then_write( self.connection.load_from_lerobot( "concurrent_append", self.image_source) - for name in ( - "concurrent_append", - "concurrent_append__versions", - "concurrent_append__episodes", - "concurrent_append__tasks", - "concurrent_append__subtasks"): - with self.subTest(name=name): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(name)) - - def test_drop_table_removes_companion_tables(self): - self.connection.load_from_lerobot("drop_group", self.image_source) - self.connection.drop_table("drop_group") - - for name in ( - "drop_group", "drop_group__versions", - "drop_group__episodes", "drop_group__tasks", - "drop_group__subtasks"): - with self.subTest(name=name): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(name)) - - def test_drop_failure_is_not_retried(self): - self.connection.load_from_lerobot( - "retry_drop", self.image_source) - episode_name = Identifier.from_string(self.connection._identifier( - "retry_drop__episodes")) - episode_table = self.connection.catalog.get_table(episode_name) - expected_owner = episode_table.table_schema.options[ - _OWNER_ID_OPTION] - original_drop = self.connection.catalog.drop_table - attempts = [0] - - def flaky_drop(identifier, ignore_if_not_exists=False): - if identifier == episode_name: - attempts[0] += 1 - raise RuntimeError("injected drop failure") - return original_drop(identifier, ignore_if_not_exists) - - with patch.object( - self.connection.catalog, - "drop_table", - side_effect=flaky_drop): - with self.assertRaisesRegex(RuntimeError, "could not drop"): - self.connection.drop_table("retry_drop") - - self.assertEqual(1, attempts[0]) - remaining = self.connection.catalog.get_table(episode_name) + self.connection.get_table("concurrent_append") self.assertEqual( - expected_owner, - remaining.table_schema.options.get(_OWNER_ID_OPTION), - ) - self.connection.get_table("retry_drop") - - self.connection.drop_table("retry_drop") - for name in ( - "retry_drop", "retry_drop__versions", - "retry_drop__episodes", "retry_drop__tasks"): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(name)) - - def test_drop_response_loss_is_reconciled_without_retry(self): - self.connection.load_from_lerobot( - "lost_drop_response", self.image_source) - episode_name = Identifier.from_string(self.connection._identifier( - "lost_drop_response__episodes")) - original_drop = self.connection.catalog.drop_table - attempts = [0] - - def drop_then_lose_response(identifier, ignore_if_not_exists=False): - result = original_drop(identifier, ignore_if_not_exists) - if identifier == episode_name: - attempts[0] += 1 - raise RuntimeError("response lost") - return result - - with patch.object( - self.connection.catalog, - "drop_table", - side_effect=drop_then_lose_response): - self.connection.drop_table("lost_drop_response") - - self.assertEqual(1, attempts[0]) - - def test_drop_table_does_not_rename_table_directories(self): - self.connection.load_from_lerobot( - "direct_group_drop", self.image_source) - - with patch.object( - self.connection.catalog, - "rename_table", - side_effect=AssertionError("rename must not be used")): - self.connection.drop_table("direct_group_drop") - - def test_drop_table_validates_group_before_deleting(self): - self.connection.load_from_lerobot( - "mixed_owner", self.image_source) - identifier = self.connection._identifier("mixed_owner__tasks") - table = self.connection.catalog.get_table(identifier) - schema = table.table_schema.to_schema() - schema.options = dict(schema.options) - schema.options[_OWNER_ID_OPTION] = "other-owner" - self.connection.catalog.drop_table(identifier) - self.connection.catalog.create_table(identifier, schema, False) - - with patch.object( - self.connection.catalog, - "drop_table", - wraps=self.connection.catalog.drop_table) as drop: - with self.assertRaisesRegex(ValueError, "different table"): - self.connection.drop_table("mixed_owner") - drop.assert_not_called() - - for name in ( - "mixed_owner", "mixed_owner__versions", - "mixed_owner__episodes", "mixed_owner__tasks"): - self.connection.catalog.get_table( - self.connection._identifier(name)) - - def test_companion_table_cannot_be_dropped_directly(self): - self.connection.load_from_lerobot( - "direct_drop", self.image_source) - - with self.assertRaisesRegex(ValueError, "companion table"): - self.connection.drop_table("direct_drop__tasks") - - self.connection.get_table("direct_drop") - self.connection.catalog.get_table( - self.connection._identifier("direct_drop__tasks")) - self.connection.drop_table("direct_drop") - - def test_drop_table_rejects_managed_branch(self): - self.connection.load_from_lerobot( - "branch_drop", self.image_source) - self.connection.catalog.create_branch( - self.connection._identifier("branch_drop"), "dev") - - with self.assertRaisesRegex(ValueError, "table branch"): - self.connection.drop_table("branch_drop$branch_dev") - for name in ( - "branch_drop", "branch_drop__versions", - "branch_drop__episodes", "branch_drop__tasks"): - self.connection.catalog.get_table( - self.connection._identifier(name)) - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier("branch_drop__subtasks")) - - def test_table_group_survives_frame_table_rename(self): - self.connection.load_from_lerobot("before_rename", self.image_source) - self.connection.catalog.rename_table( - self.connection._identifier("before_rename"), - self.connection._identifier("after_rename"), + ["PENDING"], + [row["status"] for row in _catalog_rows( + self.connection, "concurrent_append__versions")], ) - with self.assertRaisesRegex(ValueError, "already exists"): - self.connection.load_from_lerobot( - "after_rename", self.image_source) - - self.connection.drop_table("after_rename") - for name in ( - "after_rename", "before_rename__versions", - "before_rename__episodes", "before_rename__tasks", - "before_rename__subtasks"): - with self.assertRaises(TableNotExistException): - self.connection.catalog.get_table( - self.connection._identifier(name)) if __name__ == "__main__": unittest.main() From d0bd8c0259e9817c5d64b60636d5136a4491136b Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 02:40:47 -0700 Subject: [PATCH 29/32] [python] Clarify LeRobot required feature validation --- paimon-python/pypaimon/multimodal/lerobot/api.py | 6 +++--- .../pypaimon/multimodal/lerobot/schema.py | 8 ++++---- .../pypaimon/tests/multimodal_lerobot_test.py | 14 +++++++------- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 4a47acb9d176..18ad5bc35c36 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -35,7 +35,7 @@ from pypaimon.multimodal.lerobot.schema import ( _require_v3, _schema_from_info, - _validate_v3_control_features, + _validate_v3_required_features, ) from pypaimon.multimodal.lerobot.source import ( _close_quietly, @@ -87,7 +87,7 @@ def load_from_lerobot( _schema_from_info(local_info) _positive_integer(local_info.get("fps"), "fps") _validated_counts(local_info, resolved_source.path) - _validate_v3_control_features(local_info) + _validate_v3_required_features(local_info) LeRobotDataset = _import_lerobot_dataset() dataset = _open_resolved_dataset( LeRobotDataset, resolved_source, local_info) @@ -95,7 +95,7 @@ def load_from_lerobot( info = dict(dataset.meta.info) _require_v3(info, resolved_source.path) _validated_counts(info, resolved_source.path) - _validate_v3_control_features(info) + _validate_v3_required_features(info) lerobot_schema = _schema_from_info(info) metadata = _load_dataset_metadata( diff --git a/paimon-python/pypaimon/multimodal/lerobot/schema.py b/paimon-python/pypaimon/multimodal/lerobot/schema.py index 6ee158978fe1..de16c95a6478 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/schema.py +++ b/paimon-python/pypaimon/multimodal/lerobot/schema.py @@ -37,7 +37,7 @@ "string": pa.string(), } -_V3_CONTROL_DTYPES = { +_V3_REQUIRED_FEATURE_DTYPES = { "timestamp": "float32", "frame_index": "int64", "episode_index": "int64", @@ -66,11 +66,11 @@ def _schema_from_info(info): ]) -def _validate_v3_control_features(info): +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_CONTROL_DTYPES) + expected = dict(_V3_REQUIRED_FEATURE_DTYPES) if "subtask_index" in features: expected["subtask_index"] = "int64" for name, dtype in expected.items(): @@ -79,7 +79,7 @@ def _validate_v3_control_features(info): or str(feature.get("dtype", "")) != dtype \ or _feature_shape(feature, name) != (1,): raise ValueError( - "LeRobot V3 control feature %s must have dtype=%s and " + "LeRobot V3 required feature %s must have dtype=%s and " "shape=[1]." % (name, dtype)) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 901f19b50e41..ce44ea5465bd 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -54,7 +54,7 @@ from pypaimon.multimodal.lerobot.schema import ( _schema_from_info, _validate_lerobot_schema, - _validate_v3_control_features, + _validate_v3_required_features, ) from pypaimon.multimodal.lerobot.source import ( _LeRobotSource, @@ -398,7 +398,7 @@ def test_existing_schema_preserves_lerobot_feature_contract(self): ValueError, "cannot be converted"): _validate_lerobot_schema(source, target, "dataset") - def test_v3_control_features_have_native_types(self): + def test_v3_required_features_have_native_types(self): features = { "timestamp": {"dtype": "float32", "shape": [1]}, "frame_index": {"dtype": "int64", "shape": [1]}, @@ -406,7 +406,7 @@ def test_v3_control_features_have_native_types(self): "index": {"dtype": "int64", "shape": [1]}, "task_index": {"dtype": "int64", "shape": [1]}, } - _validate_v3_control_features({"features": features}) + _validate_v3_required_features({"features": features}) for name, replacement in ( ("timestamp", {"dtype": "float64", "shape": [1]}), @@ -416,13 +416,13 @@ def test_v3_control_features_have_native_types(self): invalid = dict(features) invalid[name] = replacement with self.assertRaisesRegex( - ValueError, "control feature %s" % name): - _validate_v3_control_features({"features": invalid}) + ValueError, "required feature %s" % name): + _validate_v3_required_features({"features": invalid}) missing = dict(features) del missing["task_index"] - with self.assertRaisesRegex(ValueError, "control feature task_index"): - _validate_v3_control_features({"features": missing}) + 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( From e2806b4d7da18001cae478c131c9831232c10924 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 03:14:09 -0700 Subject: [PATCH 30/32] [python] Reconcile LeRobot publication retries --- .../pypaimon/multimodal/lerobot/metadata.py | 61 +++++++++---- .../pypaimon/multimodal/lerobot/source.py | 35 ++++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 87 +++++++++++++++---- .../tests/write/commit_callback_test.py | 42 +++++++++ .../pypaimon/write/file_store_commit.py | 7 ++ 5 files changed, 196 insertions(+), 36 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 3e2c500ab6d1..1c13a0e94c81 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -114,7 +114,7 @@ def _load_dataset_metadata(dataset, info, source): tasks_table = _source_tasks( dataset, source, int(info["total_tasks"])) task_indices = _task_indices( - tasks_table.to_pylist(), int(info["total_tasks"])) + 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"]) @@ -363,11 +363,35 @@ def _append_arrow_tables(table, tables): def _create_tag(catalog, identifier, tag_name, snapshot_id): try: - catalog.create_tag( - identifier, tag_name, snapshot_id=snapshot_id) + 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: - catalog.get_table(identifier).create_tag( - tag_name, snapshot_id=snapshot_id) + snapshot = catalog.get_table(identifier).tag_manager().get(tag_name) + return None if snapshot is None else snapshot.id def _source_stats(dataset, source): @@ -555,18 +579,23 @@ def _metadata_root(dataset, source): return Path(root) -def _task_indices(records, total_tasks): +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.get("task", record.get("name")) - if task is None: - task = record.get("__index_level_0__") - if index < 0 or index >= total_tasks or task is None \ + 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) - task = str(task) if task in by_name: raise ValueError("LeRobot task metadata repeats task %r." % task) by_name[task] = index @@ -591,11 +620,13 @@ def _subtask_indices(subtasks_table, info): if "subtask_index" not in subtasks_table.column_names: raise ValueError( "LeRobot subtask metadata is missing subtask_index.") - records = subtasks_table.to_pylist() + 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.get("subtask", record.get("name")) - if label is None: - label = record.get("__index_level_0__") + label = record[label_column] if _integer(record.get("subtask_index"), "subtask_index") \ != expected or not isinstance(label, str) or not label: raise ValueError( diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index f2d29348171b..04547db80403 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -332,16 +332,17 @@ def _load_tasks(self, info): 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." @@ -388,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) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index ce44ea5465bd..a2d28f1d8114 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -570,6 +570,8 @@ def test_empty_dataset_with_tasks_is_rejected(self): 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" @@ -585,10 +587,10 @@ def test_optional_subtasks_keep_their_native_schema(self): "subtask_index": {"dtype": "int64", "shape": [1]}, }, })) - pq.write_table(pa.table({ - "subtask_index": [0], - "subtask": ["reach"], - }), source / "meta" / "subtasks.parquet") + 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( @@ -608,20 +610,25 @@ def test_optional_subtasks_keep_their_native_schema(self): 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(pa.table({ - "subtask_index": [1, 0], - "subtask": ["reach", "grasp"], - }), info) - 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" @@ -636,11 +643,13 @@ def test_native_metadata_does_not_require_json_values(self): "index": {"dtype": "int64", "shape": [1]}, }, } - pq.write_table(pa.table({ - "task_index": [0], - "task": ["pick"], - "native_bytes": [b"\xff"], - }), source / "meta" / "tasks.parquet") + 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], @@ -1153,7 +1162,7 @@ def test_import_publishes_optional_subtasks(self): ), path) subtasks = pa.Table.from_pandas(pd.DataFrame( {"subtask_index": [0, 1]}, - index=pd.Index(["reach", "grasp"], name="subtask"), + index=pd.Index(["reach", "grasp"], name="instruction"), )) pq.write_table(subtasks, source / "meta" / "subtasks.parquet") @@ -1387,6 +1396,28 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): if "/data/" in path and path.endswith(".parquet") ])) + 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) @@ -1427,6 +1458,32 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): 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( + ["PENDING", "READY"], + [row["status"] for row in _catalog_rows( + self.connection, "tag_response_loss__versions")], + ) + def test_failed_publication_remains_pending(self): with patch( "pypaimon.multimodal.lerobot.api._publish_dataset", diff --git a/paimon-python/pypaimon/tests/write/commit_callback_test.py b/paimon-python/pypaimon/tests/write/commit_callback_test.py index 9addc1e812db..9ea0d4e0394e 100644 --- a/paimon-python/pypaimon/tests/write/commit_callback_test.py +++ b/paimon-python/pypaimon/tests/write/commit_callback_test.py @@ -23,6 +23,7 @@ import pyarrow as pa from pypaimon import CatalogFactory, Schema +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.write.commit_callback import CommitCallback, CommitCallbackContext @@ -160,6 +161,47 @@ def commit_then_lose_response(base_snapshot_uuid, snapshot, statistics): table_write.close() table_commit.close() + def test_empty_overwrite_callback_after_lost_commit_response(self): + table = self._create_table('test_empty_overwrite_response_loss') + builder = table.new_batch_write_builder() + table_write = builder.new_write() + initial_commit = builder.new_commit() + table_write.write_arrow(pa.Table.from_pydict({ + 'id': [1], + 'name': ['a'], + 'dt': ['p1'], + }, schema=self.pa_schema)) + initial_commit.commit(table_write.prepare_commit()) + table_write.close() + initial_commit.close() + + table_commit = table.new_batch_write_builder().new_commit() + callback = RecordingCallback() + table_commit.add_commit_callback(callback) + real_commit = table_commit.file_store_commit.snapshot_commit.commit + attempts = [] + + def commit_then_lose_response( + base_snapshot_uuid, snapshot, statistics): + attempts.append(snapshot.id) + self.assertTrue(real_commit( + base_snapshot_uuid, snapshot, statistics)) + raise TimeoutError('lost snapshot commit response') + + table_commit.file_store_commit.snapshot_commit.commit = ( + commit_then_lose_response) + table_commit.file_store_commit._commit_retry_wait = lambda _: None + table_commit.file_store_commit.truncate_table( + BATCH_COMMIT_IDENTIFIER) + + self.assertEqual([2], attempts) + self.assertEqual(1, len(callback.contexts)) + self.assertEqual(2, callback.contexts[0].snapshot.id) + read_builder = table.new_read_builder() + splits = read_builder.new_scan().plan().splits() + self.assertEqual(0, read_builder.new_read().to_arrow(splits).num_rows) + table_commit.close() + def test_multiple_callbacks(self): table = self._create_table('test_multi_callbacks') write_builder = table.new_batch_write_builder() diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index f3fb278e15a5..cf71636a0e0d 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -494,6 +494,13 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, start_time_ms = int(time.time() * 1000) while True: latest_snapshot = self.snapshot_manager.get_latest_snapshot() + if retry_result is not None and self._is_duplicate_commit( + retry_result, + latest_snapshot, + commit_identifier, + commit_kind, + notify_callbacks=True): + break commit_entries = ( rewritten_commit_entries if rewritten_commit_entries is not None From 178e3560e45190434c9ba3f0f233bb4cccd47675 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 04:17:32 -0700 Subject: [PATCH 31/32] [python] Separate generic fixes from LeRobot import --- .../pypaimon/multimodal/connection.py | 28 +----- .../pypaimon/tests/multimodal_lerobot_test.py | 32 ------- .../pypaimon/tests/multimodal_table_test.py | 96 ------------------- .../tests/write/commit_callback_test.py | 87 ----------------- .../pypaimon/write/file_store_commit.py | 46 ++------- 5 files changed, 14 insertions(+), 275 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 31a1bdc88490..938622536763 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -71,35 +71,18 @@ def create_table( """Create a multimodal table and optionally add initial data.""" identifier = self._identifier(name) already_exists = _table_exists(self.catalog, identifier) - if already_exists and ignore_if_exists: - try: - return self.get_table(name) - except (DatabaseNotExistException, TableNotExistException): - pass - try: - paimon_schema = _to_paimon_schema( - schema, data, options, partitioned) - _validate_multimodal_schema(paimon_schema, identifier) - except ValueError: - if ignore_if_exists: - try: - return self.get_table(name) - except (DatabaseNotExistException, TableNotExistException): - pass - raise + paimon_schema = _to_paimon_schema(schema, data, options, partitioned) self._create_database_for(identifier) - created = False try: self.catalog.create_table( - identifier, paimon_schema, False) - created = True + identifier, paimon_schema, ignore_if_exists) except TableAlreadyExistException: if not ignore_if_exists: raise table = self.get_table(name) - if data is not None and created: + if data is not None and not already_exists: table.add(data) return table @@ -210,10 +193,7 @@ def _table_exists(catalog, identifier: str) -> bool: def _validate_multimodal_table(table, identifier: str): - _validate_multimodal_schema(table.table_schema, identifier) - - -def _validate_multimodal_schema(table_schema, identifier: str): + table_schema = table.table_schema options = table_schema.options if str(options.get("data-evolution.enabled", "false")).lower() != "true": raise ValueError( diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index a2d28f1d8114..74d78ce5c33b 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1499,38 +1499,6 @@ def test_failed_publication_remains_pending(self): self.connection, "failed_publish__versions")], ) - def test_lost_commit_response_does_not_fail_import(self): - from pypaimon.snapshot.renaming_snapshot_commit import ( - RenamingSnapshotCommit, - ) - - real_commit = RenamingSnapshotCommit.commit - lost = [False] - - def commit_then_lose_response( - snapshot_commit, base_snapshot_uuid, snapshot, statistics): - result = real_commit( - snapshot_commit, base_snapshot_uuid, snapshot, statistics) - if result and not lost[0]: - lost[0] = True - raise TimeoutError("lost snapshot commit response") - return result - - with patch.object( - RenamingSnapshotCommit, - "commit", - new=commit_then_lose_response): - version_id = self.connection.load_from_lerobot( - "lost_commit_response", self.image_source) - - self.assertTrue(lost[0]) - self.assertEqual(1, version_id) - self.assertEqual( - ["PENDING", "READY"], - [row["status"] for row in _catalog_rows( - self.connection, "lost_commit_response__versions")], - ) - def test_existing_companion_is_rejected(self): self.connection.load_from_lerobot( "other_group", self.image_source) diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index f2d81b2a44a8..ddf84f5696e1 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -760,102 +760,6 @@ def test_get_table_rejects_primary_key_table(self): with self.assertRaisesRegex(ValueError, "primary keys"): self.conn.get_table("pk") - def test_create_table_ignores_invalid_options_when_table_exists(self): - schema = _schema({"id": pa.int32()}) - expected = self.conn.create_table("existing", schema=schema) - - actual = self.conn.create_table( - "existing", - schema=_schema({ - "id": pa.int32(), - "embedding": _vector(3), - }), - options={"data-evolution.enabled": "false"}, - ignore_if_exists=True, - ) - - self.assertEqual(expected.identifier, actual.identifier) - with patch( - "pypaimon.multimodal.connection._table_exists", - return_value=False): - raced = self.conn.create_table( - "existing", - schema=_schema({ - "id": pa.int32(), - "embedding": _vector(3), - }), - options={"data-evolution.enabled": "false"}, - ignore_if_exists=True, - ) - self.assertEqual(expected.identifier, raced.identifier) - - def test_create_table_handles_concurrent_delete_when_ignoring(self): - schema = _schema({"id": pa.int32()}) - self.conn.create_table("deleted", schema=schema) - original_get = self.conn.get_table - deleted = [False] - - def delete_once(name): - if not deleted[0]: - deleted[0] = True - self.conn.catalog.drop_table("default.deleted", False) - return original_get(name) - - with patch.object( - self.conn, "get_table", side_effect=delete_once): - table = self.conn.create_table( - "deleted", - data=pa.table({"id": [1, 2, 3]}), - schema=schema, - ignore_if_exists=True, - ) - self.assertEqual("default.deleted", table.identifier) - self.assertEqual( - [1, 2, 3], table.scan().to_arrow()["id"].to_pylist()) - - self.conn.create_table("fallback_deleted", schema=schema) - deleted[0] = False - - def delete_fallback_once(name): - if not deleted[0]: - deleted[0] = True - self.conn.catalog.drop_table( - "default.fallback_deleted", False) - return original_get(name) - - with patch( - "pypaimon.multimodal.connection._table_exists", - return_value=False): - with patch.object( - self.conn, - "get_table", - side_effect=delete_fallback_once): - with self.assertRaisesRegex( - ValueError, "data-evolution.enabled"): - self.conn.create_table( - "fallback_deleted", - schema=schema, - options={"data-evolution.enabled": "false"}, - ignore_if_exists=True, - ) - - def test_create_table_does_not_add_data_to_concurrent_winner(self): - schema = _schema({"id": pa.int32()}) - winner = self.conn.create_table("winner", schema=schema) - - with patch( - "pypaimon.multimodal.connection._table_exists", - return_value=False): - actual = self.conn.create_table( - "winner", - data=pa.table({"id": [1, 2, 3]}), - schema=schema, - ignore_if_exists=True, - ) - - self.assertEqual(winner.identifier, actual.identifier) - self.assertEqual([], actual.scan().to_arrow().to_pylist()) - def test_create_table_can_add_initial_data_and_get_by_short_name(self): self.conn.create_table( "users", diff --git a/paimon-python/pypaimon/tests/write/commit_callback_test.py b/paimon-python/pypaimon/tests/write/commit_callback_test.py index 9ea0d4e0394e..036031edefec 100644 --- a/paimon-python/pypaimon/tests/write/commit_callback_test.py +++ b/paimon-python/pypaimon/tests/write/commit_callback_test.py @@ -23,7 +23,6 @@ import pyarrow as pa from pypaimon import CatalogFactory, Schema -from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.write.commit_callback import CommitCallback, CommitCallbackContext @@ -116,92 +115,6 @@ def test_callback_receives_correct_snapshot_data(self): table_write.close() table_commit.close() - def test_callback_invoked_after_lost_commit_response(self): - table = self._create_table( - 'test_callback_response_loss', - options={ - 'row-tracking.enabled': 'true', - 'data-evolution.enabled': 'true', - }, - ) - write_builder = table.new_batch_write_builder() - table_write = write_builder.new_write() - table_commit = write_builder.new_commit() - - callback = RecordingCallback() - table_commit.add_commit_callback(callback) - real_commit = table_commit.file_store_commit.snapshot_commit.commit - attempts = [] - - def commit_then_lose_response(base_snapshot_uuid, snapshot, statistics): - attempts.append(snapshot.id) - self.assertTrue(real_commit( - base_snapshot_uuid, snapshot, statistics)) - raise TimeoutError('lost snapshot commit response') - - table_commit.file_store_commit.snapshot_commit.commit = ( - commit_then_lose_response) - table_commit.file_store_commit._commit_retry_wait = lambda _: None - - data = pa.Table.from_pydict({ - 'id': [1, 2], - 'name': ['a', 'b'], - 'dt': ['p1', 'p1'], - }, schema=self.pa_schema) - table_write.write_arrow(data) - table_commit.commit(table_write.prepare_commit()) - - self.assertEqual([1], attempts) - self.assertEqual(1, len(callback.contexts)) - self.assertEqual(1, callback.contexts[0].snapshot.id) - self.assertGreater(len(callback.contexts[0].commit_entries), 0) - for entry in callback.contexts[0].commit_entries: - self.assertIsNotNone(entry.file.first_row_id) - - table_write.close() - table_commit.close() - - def test_empty_overwrite_callback_after_lost_commit_response(self): - table = self._create_table('test_empty_overwrite_response_loss') - builder = table.new_batch_write_builder() - table_write = builder.new_write() - initial_commit = builder.new_commit() - table_write.write_arrow(pa.Table.from_pydict({ - 'id': [1], - 'name': ['a'], - 'dt': ['p1'], - }, schema=self.pa_schema)) - initial_commit.commit(table_write.prepare_commit()) - table_write.close() - initial_commit.close() - - table_commit = table.new_batch_write_builder().new_commit() - callback = RecordingCallback() - table_commit.add_commit_callback(callback) - real_commit = table_commit.file_store_commit.snapshot_commit.commit - attempts = [] - - def commit_then_lose_response( - base_snapshot_uuid, snapshot, statistics): - attempts.append(snapshot.id) - self.assertTrue(real_commit( - base_snapshot_uuid, snapshot, statistics)) - raise TimeoutError('lost snapshot commit response') - - table_commit.file_store_commit.snapshot_commit.commit = ( - commit_then_lose_response) - table_commit.file_store_commit._commit_retry_wait = lambda _: None - table_commit.file_store_commit.truncate_table( - BATCH_COMMIT_IDENTIFIER) - - self.assertEqual([2], attempts) - self.assertEqual(1, len(callback.contexts)) - self.assertEqual(2, callback.contexts[0].snapshot.id) - read_builder = table.new_read_builder() - splits = read_builder.new_scan().plan().splits() - self.assertEqual(0, read_builder.new_read().to_arrow(splits).num_rows) - table_commit.close() - def test_multiple_callbacks(self): table = self._create_table('test_multi_callbacks') write_builder = table.new_batch_write_builder() diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index cf71636a0e0d..5f8f054aa812 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -494,13 +494,6 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, start_time_ms = int(time.time() * 1000) while True: latest_snapshot = self.snapshot_manager.get_latest_snapshot() - if retry_result is not None and self._is_duplicate_commit( - retry_result, - latest_snapshot, - commit_identifier, - commit_kind, - notify_callbacks=True): - break commit_entries = ( rewritten_commit_entries if rewritten_commit_entries is not None @@ -606,11 +599,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str commit_result_may_be_uncertain: bool = False) -> CommitResult: start_millis = int(time.time() * 1000) if self._is_duplicate_commit( - retry_result, - latest_snapshot, - commit_identifier, - commit_kind, - notify_callbacks=True): + retry_result, latest_snapshot, commit_identifier, commit_kind): return SuccessResult() latest_snapshot_id = latest_snapshot.id if latest_snapshot else 0 @@ -873,8 +862,14 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str commit_kind, ) - self._notify_commit_callbacks( - snapshot_data, commit_entries, commit_identifier) + if self.commit_callbacks: + context = CommitCallbackContext( + snapshot=snapshot_data, + commit_entries=commit_entries, + identifier=commit_identifier, + ) + for callback in self.commit_callbacks: + callback.call(context) return SuccessResult() @@ -932,8 +927,7 @@ def _is_duplicate_commit( retry_result, latest_snapshot, commit_identifier, - commit_kind, - notify_callbacks=False) -> bool: + commit_kind) -> bool: if (isinstance(retry_result, CommitFailRetryResult) and latest_snapshot is not None): start_check_snapshot_id = 1 # Snapshot.FIRST_SNAPSHOT_ID @@ -959,29 +953,9 @@ def _is_duplicate_commit( f"Commit already completed (snapshot {snapshot_id}), " f"user: {self.commit_user}, identifier: {commit_identifier}" ) - if notify_callbacks and self.commit_callbacks: - entries = [] - for manifest in self.manifest_list_manager.read_delta( - snapshot): - entries.extend(self.manifest_file_manager.read( - manifest.file_name, drop_stats=False)) - self._notify_commit_callbacks( - snapshot, entries, commit_identifier) return True return False - def _notify_commit_callbacks( - self, snapshot, commit_entries, commit_identifier): - if not self.commit_callbacks: - return - context = CommitCallbackContext( - snapshot=snapshot, - commit_entries=commit_entries, - identifier=commit_identifier, - ) - for callback in self.commit_callbacks: - callback.call(context) - def _create_dynamic_partition_filter(self, commit_messages: List[CommitMessage]): """Build a partition filter from the unique partitions present in commit_messages.""" predicate_builder = PredicateBuilder(self.table.partition_keys_fields) From 0704cf6a881076d4baf01d19d29f7e5cebccacd6 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 3 Sep 2026 06:41:21 -0700 Subject: [PATCH 32/32] [python] Publish LeRobot versions after component tags --- docs/docs/pypaimon/multimodal-api.mdx | 12 ++- .../pypaimon/multimodal/lerobot/api.py | 6 -- .../pypaimon/multimodal/lerobot/metadata.py | 17 +--- .../pypaimon/tests/multimodal_lerobot_test.py | 94 +++++++------------ 4 files changed, 44 insertions(+), 85 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 30a769c15858..c50acc7eb1d9 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -620,9 +620,9 @@ repository. It derives the schema from `meta/info.json`, writes one row per 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`. A manifest row is reserved as `PENDING` and marked `READY` after -all components are committed and tagged with the same numeric `version_id`. -Readers ignore `PENDING`. +`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]' @@ -652,8 +652,10 @@ version_id = conn.load_from_lerobot( ) ``` -A READY row in `
__versions` identifies a release; reading the same tag -from its component tables reconstructs that release. +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. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 18ad5bc35c36..6078ac7b245e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -28,7 +28,6 @@ _prepare_metadata_tables, _positive_integer, _publish_dataset, - _reserve_dataset_version, _validated_episode_tables, ) from pypaimon.multimodal.lerobot.loader import _write_dataset @@ -132,11 +131,6 @@ def _import_dataset( tables = _prepare_metadata_tables( connection, table.raw_table, metadata) version_id = 1 - _reserve_dataset_version( - tables["versions"], - version_id, - metadata, - ) episodes_snapshot_id = _append_arrow_tables( tables["episodes"], _validated_episode_tables(metadata), diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 1c13a0e94c81..ee3879444ade 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -50,7 +50,6 @@ _VERSIONS_SCHEMA = pa.schema([ pa.field(_VERSION_ID, pa.int64(), nullable=False), - pa.field("status", pa.string(), nullable=False), pa.field("info_json", pa.string(), nullable=False), pa.field("stats_json", pa.string()), pa.field("has_subtasks", pa.bool_(), nullable=False), @@ -244,18 +243,6 @@ def _restore_pandas_metadata(table, data): return data.replace_schema_metadata(metadata) -def _reserve_dataset_version( - versions_table, - version_id, - metadata): - pending = _manifest_row(version_id, "PENDING", metadata) - snapshot_id = _append_arrow( - versions_table, - pa.Table.from_pylist([pending], schema=_VERSIONS_SCHEMA), - ) - _require_initial_snapshot("versions", snapshot_id) - - def _publish_dataset( connection, tables, @@ -284,7 +271,7 @@ def _publish_dataset( for identifier, snapshot_id in component_snapshots: _create_tag(connection.catalog, identifier, tag, snapshot_id) - manifest = _manifest_row(version_id, "READY", metadata) + manifest = _manifest_row(version_id, metadata) _append_arrow(tables["versions"], pa.Table.from_pylist( [manifest], schema=_VERSIONS_SCHEMA)) @@ -302,11 +289,9 @@ def _require_initial_snapshot(component, snapshot_id): def _manifest_row( version_id, - status, metadata): return { _VERSION_ID: version_id, - "status": status, "info_json": metadata["info_json"], "stats_json": metadata["stats_json"], "has_subtasks": metadata["subtasks_table"] is not None, diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 74d78ce5c33b..97c82b43d40a 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -996,17 +996,14 @@ def test_import_infers_schema_and_preserves_episodes(self): 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(["PENDING", "READY"], [ - row["status"] for row in manifests - ]) - manifest = manifests[1] + 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( - {"version_id", "status", "info_json", "stats_json", - "has_subtasks"}, + {"version_id", "info_json", "stats_json", "has_subtasks"}, set(manifest)) self.assertFalse(manifest["has_subtasks"]) tag = str(manifest["version_id"]) @@ -1192,7 +1189,7 @@ def test_import_publishes_optional_subtasks(self): subtasks_table, subtasks_arrow).to_pandas(), ) self.assertTrue(_catalog_rows( - self.connection, "with_subtasks__versions")[1]["has_subtasks"]) + self.connection, "with_subtasks__versions")[0]["has_subtasks"]) self.assertEqual( 1, self.connection.catalog.get_tag( @@ -1273,11 +1270,8 @@ def test_frame_controls_must_match_published_episode_metadata(self): ValueError, "has %s" % column): self.connection.load_from_lerobot(table_name, source) self.connection.get_table(table_name) - self.assertEqual( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, table_name + "__versions")], - ) + 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" @@ -1325,11 +1319,8 @@ def test_episode_tasks_must_exactly_match_frame_tasks(self): self.connection.load_from_lerobot( "extra_episode_task", source) self.connection.get_table("extra_episode_task") - self.assertEqual( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, "extra_episode_task__versions")], - ) + 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" @@ -1355,11 +1346,8 @@ def test_nonempty_dataset_cannot_publish_without_tasks(self): with self.assertRaisesRegex(ValueError, "task_index"): self.connection.load_from_lerobot("missing_tasks", source) self.connection.get_table("missing_tasks") - self.assertEqual( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, "missing_tasks__versions")], - ) + 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" @@ -1449,7 +1437,7 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): "tag_fallback", self.image_source) manifest = _catalog_rows( - self.connection, "tag_fallback__versions")[1] + self.connection, "tag_fallback__versions")[0] tag = str(manifest["version_id"]) table = self.connection.get_table("tag_fallback") self.assertEqual(1, version_id) @@ -1479,25 +1467,21 @@ def create_then_lose_response(*args, **kwargs): self.assertTrue(lost[0]) self.assertEqual(1, version_id) self.assertEqual( - ["PENDING", "READY"], - [row["status"] for row in _catalog_rows( - self.connection, "tag_response_loss__versions")], - ) + [1], + [row["version_id"] for row in _catalog_rows( + self.connection, "tag_response_loss__versions")]) - def test_failed_publication_remains_pending(self): + def test_tag_failure_remains_unpublished(self): with patch( - "pypaimon.multimodal.lerobot.api._publish_dataset", - side_effect=RuntimeError("publish failed")): - with self.assertRaisesRegex(RuntimeError, "publish failed"): + "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( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, "failed_publish__versions")], - ) + self.assertEqual([], _catalog_rows( + self.connection, "failed_publish__versions")) def test_existing_companion_is_rejected(self): self.connection.load_from_lerobot( @@ -1565,10 +1549,9 @@ def open_with_failing_close(*args, **kwargs): self.assertEqual(1, version_id) self.assertEqual( - ["PENDING", "READY"], - [row["status"] for row in _catalog_rows( - self.connection, "close_failure__versions")], - ) + [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" @@ -1585,10 +1568,9 @@ def test_source_close_failure_does_not_override_success(self): self.assertEqual(1, version_id) self.assertEqual( - ["PENDING", "READY"], - [row["status"] for row in _catalog_rows( - self.connection, "source_close_failure__versions")], - ) + [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()) @@ -1604,20 +1586,19 @@ def test_existing_target_is_rejected(self): def test_concurrent_import_cannot_claim_the_same_target(self): from pypaimon.multimodal.lerobot import api - original_reserve = api._reserve_dataset_version - reserved = threading.Event() + original_prepare = api._prepare_metadata_tables + root_created = threading.Event() release = threading.Event() - def reserve_then_wait(*args, **kwargs): - result = original_reserve(*args, **kwargs) - reserved.set() + def prepare_then_wait(*args, **kwargs): + root_created.set() release.wait(10) - return result + return original_prepare(*args, **kwargs) with patch.object( api, - "_reserve_dataset_version", - side_effect=reserve_then_wait): + "_prepare_metadata_tables", + side_effect=prepare_then_wait): with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit( self.connection.load_from_lerobot, @@ -1625,7 +1606,7 @@ def reserve_then_wait(*args, **kwargs): self.image_source, ) try: - self.assertTrue(reserved.wait(10)) + self.assertTrue(root_created.wait(10)) with self.assertRaisesRegex( ValueError, "already exists"): self.connection.load_from_lerobot( @@ -1678,11 +1659,8 @@ def append_then_write( "concurrent_append", self.image_source) self.connection.get_table("concurrent_append") - self.assertEqual( - ["PENDING"], - [row["status"] for row in _catalog_rows( - self.connection, "concurrent_append__versions")], - ) + self.assertEqual([], _catalog_rows( + self.connection, "concurrent_append__versions")) if __name__ == "__main__":