diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index c50acc7eb1d9..bdf780859f09 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -616,32 +616,39 @@ orphan files for normal Paimon cleanup. ## Load LeRobot Dataset v3 `load_from_lerobot` imports a local directory, FileIO URI, or Hugging Face -repository. It derives the schema from `meta/info.json`, writes one row per -frame, and creates a LeRobot dataset backed by the frame table, -`__versions`, `
__episodes`, `
__tasks`, and an optional -`
__subtasks`. Task text remains in the metadata table; frames retain -`task_index`. After all components are committed and tagged with the same -numeric `version_id`, one row is appended to `
__versions` to publish the -version. +repository. It derives the schema from `meta/info.json` and writes one row per +frame. The import creates `
__episodes`, `
__tasks`, `
__info`, +and optional `
__stats` and `
__subtasks` companion tables. Task text +remains in the task table; frames retain `task_index`. + +Info and stats use `key STRING, value STRING`: each top-level info property or +stats feature becomes one row. Each value is JSON-encoded, preserving nested +objects, arrays, nulls, and scalar types. Decode it with `json.loads`; for example, +`fps` stores `30`, while `codebase_version` stores `"v3.0"` (including quotes). +Statistics may contain `NaN` and `Infinity`, supported by Python's JSON decoder. +Missing or empty stats create no stats table. The frame table's +`pypaimon.lerobot.-table` options identify the components created. ```shell pip install 'pypaimon[lerobot]' ``` ```python -version_id = conn.load_from_lerobot( +conn.load_from_lerobot( "robot_data", "/data/lerobot_dataset", + tag_name="initial-import", # Optional: tag all imported component snapshots. ) -print(version_id) ``` -The returned `version_id` is the common tag name for all dataset components. +The call returns `None` on success. Omitting `tag_name` imports the tables +without creating tags. The one-time importer requires a new target table and +a non-empty source dataset; subsequent table edits use the normal Paimon APIs. For FileIO URIs, pass credentials through `source_options`: ```python -version_id = conn.load_from_lerobot( +conn.load_from_lerobot( "robot_data", "oss://source-bucket/lerobot_dataset", source_options={ @@ -652,11 +659,37 @@ version_id = conn.load_from_lerobot( ) ``` -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. +Before training, finish any related data/metadata updates and pause writes to +this table group. Create a common named tag over the current component snapshots: + +```python +tag = "train-2026-09-07" +snapshots = conn.create_lerobot_tag("robot_data", tag) +frames = conn.get_table("robot_data").scan(tag_name=tag).to_arrow() + +# Companion tables are ordinary Paimon tables, read with the same tag. +info_table = conn.catalog.get_table("default.robot_data__info").copy( + {"scan.tag-name": tag}) +builder = info_table.new_read_builder() +info_rows = builder.new_read().to_arrow(builder.new_scan().plan().splits()) + +import json +info = { + row["key"]: json.loads(row["value"]) + for row in info_rows.to_pylist() +} +``` + +`create_lerobot_tag` returns component names mapped to snapshot IDs; these IDs +may differ across tables. Later appends do not change tagged reads. Read every +required component (including training statistics) through the same tag, and +never fall back to latest if a tag is missing. + +Cross-table tagging is not atomic. Use a tag only after the creation call +succeeds. A failure may leave partial tags; retry with writes still paused and +unchanged snapshots, or choose a new name after repairing the group. Existing +tags are never moved to different snapshots. Retain or delete component tags +together, and keep writers paused until the call returns. 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 ac766d5c866e..f864ff265bf0 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -44,17 +44,29 @@ pip install 'pypaimon[lerobot]' import pypaimon.multimodal as pmm connection = pmm.connect(options={"warehouse": "/tmp/warehouse"}) -version_id = connection.load_from_lerobot( +connection.load_from_lerobot( "robot_data", "/data/lerobot_dataset", ) -print(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`. +Episode, task, info, and optional stats/subtask tables. Info and stats use +`key STRING, value STRING` rows, with each value JSON-encoded to preserve +nested metadata. Decode values with `json.loads`. + +Before training, pause writes and create a shared tag: + +```python +connection.create_lerobot_tag("robot_data", "train-2026-09-07") +frames = connection.get_table("robot_data").scan( + tag_name="train-2026-09-07").to_arrow() +``` + +Read every metadata component with the same tag. Use the tag only after creation +succeeds; cross-table tagging is not atomic. Alternatively, pass `tag_name` to +`load_from_lerobot` to tag the imported snapshots immediately. # HDF5 to multimodal tables diff --git a/paimon-python/pypaimon/multimodal/connection.py b/paimon-python/pypaimon/multimodal/connection.py index 31a1bdc88490..4cea5786db6e 100644 --- a/paimon-python/pypaimon/multimodal/connection.py +++ b/paimon-python/pypaimon/multimodal/connection.py @@ -141,18 +141,29 @@ def load_from_lerobot( *, batch_size: int = 1024, options=None, - source_options=None): + source_options=None, + tag_name=None) -> None: """Import LeRobot Dataset v3 into a new Paimon table group.""" from pypaimon.multimodal.lerobot import load_from_lerobot - return load_from_lerobot( + load_from_lerobot( self, table_name, source, batch_size=batch_size, options=options, source_options=source_options, + tag_name=tag_name, ) + def create_lerobot_tag(self, table_name: str, tag_name: str): + """Pin all LeRobot components; pause group writes until this returns. + + Returns component snapshot IDs. Use the tag only after success and + retain it on every component for the lifetime of a training run. + """ + from pypaimon.multimodal.lerobot.metadata import create_lerobot_tag + return create_lerobot_tag(self, table_name, tag_name) + def load_from_rosbag( self, table_name: str, diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 6078ac7b245e..84e3d2eed5c1 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -22,12 +22,14 @@ from pypaimon.catalog.catalog_exception import TableAlreadyExistException from pypaimon.multimodal.lerobot.metadata import ( + _COMPANION_OPTION_KEYS, _append_arrow_tables, _load_dataset_metadata, _managed_table_options, _prepare_metadata_tables, _positive_integer, - _publish_dataset, + _commit_metadata, + _validate_tag_name, _validated_episode_tables, ) from pypaimon.multimodal.lerobot.loader import _write_dataset @@ -58,11 +60,13 @@ def load_from_lerobot( batch_size: int = 1024, options: Optional[Mapping[str, object]] = None, source_options: Optional[Mapping[str, object]] = None, -) -> int: - """Import LeRobot Dataset v3 and return its version ID. + tag_name: Optional[str] = None, +) -> None: + """Import LeRobot Dataset v3 into a new Paimon table group. A new target table is created from LeRobot metadata. Episode, task, and - version metadata are stored in companion Paimon tables. + info/stats metadata are stored in companion Paimon tables. If provided, + ``tag_name`` pins all components to their imported snapshots. FileIO URI credentials come only from ``source_options`` and are not inherited from the target Catalog. """ @@ -74,6 +78,9 @@ def load_from_lerobot( or batch_size <= 0: raise ValueError("batch_size must be a positive integer.") + if tag_name is not None: + _validate_tag_name(tag_name) + validated_source_options = _validated_source_options(source_options) _validate_source_kerberos( [source], validated_source_options, "LeRobot") @@ -99,7 +106,7 @@ def load_from_lerobot( lerobot_schema = _schema_from_info(info) metadata = _load_dataset_metadata( dataset, info, resolved_source) - return _import_dataset( + _import_dataset( connection, table_name, dataset, @@ -109,6 +116,7 @@ def load_from_lerobot( batch_size, options, metadata, + tag_name, ) finally: close = getattr(dataset, "close", None) @@ -125,12 +133,12 @@ def _import_dataset( source_schema, batch_size, options, - metadata): + metadata, + tag_name): table = _create_target_table( - connection, table_name, source_schema, options) + connection, table_name, source_schema, options, metadata) tables = _prepare_metadata_tables( connection, table.raw_table, metadata) - version_id = 1 episodes_snapshot_id = _append_arrow_tables( tables["episodes"], _validated_episode_tables(metadata), @@ -146,16 +154,15 @@ def _import_dataset( batch_size, metadata, ) - _publish_dataset( + _commit_metadata( connection, tables, - version_id, + tag_name, metadata, table.identifier, frames_snapshot_id, episodes_snapshot_id, ) - return version_id def _validated_counts(info, source): @@ -188,11 +195,12 @@ def _required_count(info, name, source): def _create_target_table( - connection, table_name, source_schema, options): + connection, table_name, source_schema, options, metadata): create_options = dict(options or {}) managed_options = _managed_table_options( - connection._identifier(table_name)) - reserved_options = set(managed_options).intersection(create_options) + connection._identifier(table_name), metadata) + reserved_options = set(_COMPANION_OPTION_KEYS.values()).intersection( + create_options) if reserved_options: raise ValueError( "%s are managed by load_from_lerobot." diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index ee3879444ade..0f568f47693a 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""LeRobot component tables and version publication.""" +"""LeRobot component tables and training tags.""" from array import array import json @@ -26,19 +26,18 @@ from pypaimon import Schema as PaimonSchema from pypaimon.catalog.catalog_exception import ( - DatabaseNotExistException, TableAlreadyExistException, - TableNotExistException, + TagNotExistException, ) from pypaimon.common.identifier import Identifier from pypaimon.multimodal.hdf5 import _SnapshotRecorder from pypaimon.multimodal.table import _target_schema -_VERSION_ID = "version_id" _PANDAS_METADATA_OPTION = "pypaimon.lerobot.pandas-metadata" _TABLE_SUFFIXES = { - "versions": "__versions", + "info": "__info", + "stats": "__stats", "episodes": "__episodes", "tasks": "__tasks", "subtasks": "__subtasks", @@ -48,12 +47,6 @@ for name in _TABLE_SUFFIXES } -_VERSIONS_SCHEMA = pa.schema([ - pa.field(_VERSION_ID, pa.int64(), nullable=False), - pa.field("info_json", pa.string(), nullable=False), - pa.field("stats_json", pa.string()), - pa.field("has_subtasks", pa.bool_(), nullable=False), -]) _EMPTY_TASKS_SCHEMA = pa.schema([ pa.field("task_index", pa.int64(), nullable=False), pa.field("task", pa.string(), nullable=False), @@ -110,6 +103,7 @@ def __getitem__(self, index): def _load_dataset_metadata(dataset, info, source): fps = _positive_integer(info.get("fps"), "fps") stats = _source_stats(dataset, source) + stats_table = None if stats is None else _metadata_table(stats) tasks_table = _source_tasks( dataset, source, int(info["total_tasks"])) task_indices = _task_indices( @@ -124,10 +118,9 @@ def _load_dataset_metadata(dataset, info, source): ) return { "fps": fps, - "info_json": _canonical_json(info), - "stats_json": ( - None if stats is None else _canonical_json( - stats, allow_nan=True)), + "info_table": _metadata_table(info), + "stats_table": (stats_table if stats_table is not None + and stats_table.num_rows > 0 else None), "episodes": None, "episodes_schema": episode_source["schema"], "episode_paths": episode_source["paths"], @@ -166,13 +159,16 @@ def _quote_identifier_part(value): return "`%s`" % value if "." in value else value -def _managed_table_options(frames_identifier): +def _managed_table_options(frames_identifier, metadata=None): identifier = Identifier.from_string(str(frames_identifier)) if identifier.get_branch_name() is not None: raise ValueError( "LeRobot import does not support table branches.") result = {} for name, suffix in _TABLE_SUFFIXES.items(): + if metadata is not None and name in ("stats", "subtasks") \ + and metadata[name + "_table"] is None: + continue result[_COMPANION_OPTION_KEYS[name]] = _companion_identifier( frames_identifier, suffix) return result @@ -184,6 +180,8 @@ def _companion_table_identifiers(frames_table): for name, key in _COMPANION_OPTION_KEYS.items(): value = options.get(key) if not value: + if name in ("stats", "subtasks"): + continue raise ValueError( "LeRobot table %s is missing managed option %s." % (frames_table.identifier, key)) @@ -193,22 +191,14 @@ def _companion_table_identifiers(frames_table): def _prepare_metadata_tables(connection, frames_table, metadata): schemas = { - "versions": _VERSIONS_SCHEMA, + "info": metadata["info_table"].schema, "episodes": metadata["episodes_schema"], "tasks": metadata["tasks_table"].schema, } - if metadata["subtasks_table"] is not None: - schemas["subtasks"] = metadata["subtasks_table"].schema + for name in ("stats", "subtasks"): + if metadata[name + "_table"] is not None: + schemas[name] = metadata[name + "_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] @@ -243,43 +233,72 @@ def _restore_pandas_metadata(table, data): return data.replace_schema_metadata(metadata) -def _publish_dataset( +def _commit_metadata( connection, tables, - version_id, + tag_name, metadata, frames_identifier, frames_snapshot_id, episodes_snapshot_id): _require_initial_snapshot("frames", frames_snapshot_id) _require_initial_snapshot("episodes", episodes_snapshot_id) - tasks_snapshot_id = _append_arrow( - tables["tasks"], metadata["tasks_table"]) - _require_initial_snapshot("tasks", tasks_snapshot_id) component_snapshots = [ - (frames_identifier, frames_snapshot_id), (tables["episodes"].identifier, episodes_snapshot_id), - (tables["tasks"].identifier, tasks_snapshot_id), ] - if metadata["subtasks_table"] is not None: - subtasks_snapshot_id = _append_arrow( - tables["subtasks"], metadata["subtasks_table"]) - _require_initial_snapshot("subtasks", subtasks_snapshot_id) - component_snapshots.append( - (tables["subtasks"].identifier, subtasks_snapshot_id)) - tag = str(version_id) - for identifier, snapshot_id in component_snapshots: - _create_tag(connection.catalog, identifier, tag, snapshot_id) + for name in ("tasks", "subtasks", "stats", "info"): + if name not in tables: + continue + snapshot_id = _append_arrow(tables[name], metadata[name + "_table"]) + _require_initial_snapshot(name, snapshot_id) + component_snapshots.append((tables[name].identifier, snapshot_id)) + # Tag the root last so a failed component tag does not expose a root tag. + component_snapshots.append((frames_identifier, frames_snapshot_id)) + if tag_name is not None: + for identifier, snapshot_id in component_snapshots: + _create_tag(connection.catalog, identifier, tag_name, snapshot_id) + + +def create_lerobot_tag(connection, table_name, tag_name): + """Tag the current snapshots of a LeRobot table group for training. + + Pause group writes until this call returns. Tags across tables are not an + atomic transaction: use the name only after success, and read every + component with that tag (never fall back to latest). Failed calls may leave + partial tags. Retrying is safe while the component snapshots are unchanged. + Returns a mapping from component name to tagged snapshot ID. + """ + _validate_tag_name(tag_name) + frames = connection.catalog.get_table(connection._identifier(table_name)) + identifiers = _companion_table_identifiers(frames) + identifiers["frames"] = frames.identifier + snapshots = {} + for name, identifier in identifiers.items(): + table = connection.catalog.get_table(identifier) + snapshot = table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + raise ValueError("LeRobot component %s has no snapshot." % name) + snapshots[name] = snapshot.id + existing = _tag_snapshot_id(connection.catalog, identifier, tag_name) + if existing is not None and existing != snapshot.id: + raise ValueError( + "LeRobot tag %s on %s already points to snapshot %s; " + "use a new tag name." % (tag_name, identifier, existing)) + for name, identifier in identifiers.items(): + _create_tag(connection.catalog, identifier, tag_name, snapshots[name]) + return snapshots + - manifest = _manifest_row(version_id, metadata) - _append_arrow(tables["versions"], pa.Table.from_pylist( - [manifest], schema=_VERSIONS_SCHEMA)) +def _validate_tag_name(tag_name): + if not isinstance(tag_name, str) or not tag_name.strip() \ + or any(character in tag_name for character in ("/", "\\", "\x00")): + raise ValueError("tag_name must be a non-blank name without path separators.") def _require_initial_snapshot(component, snapshot_id): if snapshot_id is None: raise ValueError( - "LeRobot tag-backed import requires a non-empty %s component." + "LeRobot import requires a non-empty %s component." % component) if snapshot_id != 1: raise RuntimeError( @@ -287,17 +306,6 @@ def _require_initial_snapshot(component, snapshot_id): "expected snapshot 1, found %d." % (component, snapshot_id)) -def _manifest_row( - version_id, - metadata): - return { - _VERSION_ID: version_id, - "info_json": metadata["info_json"], - "stats_json": metadata["stats_json"], - "has_subtasks": metadata["subtasks_table"] is not None, - } - - def _append_arrow(table, data): return _append_arrow_tables(table, [data]) @@ -374,6 +382,8 @@ def _tag_snapshot_id(catalog, identifier, tag_name): try: response = catalog.get_tag(identifier, tag_name) snapshot = response.snapshot + except TagNotExistException: + return None except NotImplementedError: snapshot = catalog.get_table(identifier).tag_manager().get(tag_name) return None if snapshot is None else snapshot.id @@ -621,14 +631,16 @@ def _subtask_indices(subtasks_table, info): return range(subtasks_table.num_rows) -def _canonical_json(value, allow_nan=False): - return json.dumps( - _json_value(value), - ensure_ascii=False, - sort_keys=True, - separators=(",", ":"), - allow_nan=allow_nan, - ) +def _metadata_table(value): + if not isinstance(value, dict): + raise ValueError("LeRobot info and stats metadata must be objects.") + return pa.table({ + "key": pa.array(list(value), type=pa.string()), + "value": pa.array([ + json.dumps(_json_value(item), ensure_ascii=False, separators=(",", ":")) + for item in value.values() + ], type=pa.string()), + }) def _json_value(value): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 97c82b43d40a..196a795be341 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -36,12 +36,14 @@ from pypaimon.common.identifier import Identifier from pypaimon.common.options import Options from pypaimon.multimodal.source_utils import _SourceFileIO +from pypaimon.multimodal.connection import MultimodalConnection 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, + _metadata_table, _restore_pandas_metadata, _subtask_indices, _validated_episode_tables, @@ -95,8 +97,159 @@ def _catalog_arrow(connection, name): return table, builder.new_read().to_arrow(plan.splits()) +def _catalog_metadata(connection, name): + return { + row["key"]: json.loads(row["value"]) + for row in _catalog_rows(connection, name) + } + + class LeRobotValidationTest(unittest.TestCase): + def test_metadata_json_preserves_nested_values(self): + values = { + "name": "机器人", + "count": 2 ** 64, + "custom": {"labels": ["pick", None], "enabled": True}, + } + table = _metadata_table(values) + self.assertEqual(pa.schema([("key", pa.string()), ("value", pa.string())]), + table.schema) + self.assertEqual(values, { + row["key"]: json.loads(row["value"]) + for row in table.to_pylist() + }) + + def test_invalid_training_tag_fails_before_catalog_access(self): + for tag_name in (None, 1, "", " ", "a/b", "a\\b", "a\x00b"): + with self.subTest(tag_name=tag_name): + connection = Mock() + with self.assertRaisesRegex(ValueError, "tag_name"): + MultimodalConnection.create_lerobot_tag( + connection, "robot", tag_name) + self.assertEqual([], connection.mock_calls) + + def test_import_validates_tag_before_source_access(self): + connection = Mock() + with patch("pypaimon.multimodal.lerobot.api._resolved_source", + side_effect=RuntimeError("source accessed")) as resolve: + for tag_name in (1, "", " ", "a/b", "a\\b", "a\x00b"): + with self.subTest(tag_name=tag_name): + with self.assertRaisesRegex(ValueError, "tag_name"): + load_from_lerobot(connection, "robot", "source", + tag_name=tag_name) + resolve.assert_not_called() + self.assertEqual([], connection.mock_calls) + with self.assertRaisesRegex(RuntimeError, "source accessed"): + load_from_lerobot(connection, "robot", "source", tag_name=None) + resolve.assert_called_once() + + @patch("pypaimon.multimodal.lerobot.api._import_lerobot_dataset", + return_value=Mock()) + def test_training_tag_uses_current_component_snapshots(self, _): + import pandas as pd + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source" + (source / "meta" / "episodes").mkdir(parents=True) + (source / "data").mkdir() + info = { + "codebase_version": "v3.0", + "total_frames": 1, + "total_episodes": 1, + "total_tasks": 1, + "fps": 30, + "data_path": "data/file.parquet", + "features": { + name: {"dtype": dtype, "shape": [1]} + for name, dtype in ( + ("index", "int64"), ("episode_index", "int64"), + ("frame_index", "int64"), ("task_index", "int64"), + ("timestamp", "float32"), + ) + }, + "custom": {"labels": ["pick", None], "enabled": True}, + } + (source / "meta" / "info.json").write_text(json.dumps(info)) + stats = {"timestamp": {"min": [0.0], "max": [0.0]}} + (source / "meta" / "stats.json").write_text(json.dumps(stats)) + pq.write_table(pa.Table.from_pandas(pd.DataFrame( + {"task_index": [0]}, index=pd.Index(["pick"], name="task"), + )), source / "meta" / "tasks.parquet") + pq.write_table(pa.table({ + "episode_index": [0], "dataset_from_index": [0], + "dataset_to_index": [1], "tasks": [["pick"]], "length": [1], + "data/chunk_index": [0], "data/file_index": [0], + }), source / "meta" / "episodes" / "file.parquet") + frames = pa.table({ + "index": [0], "episode_index": [0], "frame_index": [0], + "task_index": [0], "timestamp": pa.array([0], pa.float32()), + }) + pq.write_table(frames, source / "data" / "file.parquet") + connection = pmm.connect(options={"warehouse": str(root / "wh")}) + remote = "oss://source-bucket/robot" + with patch( + "pypaimon.multimodal.lerobot.source._SourceFileIO", + return_value=_RemoteLeRobotFileIO(source, remote)): + self.assertIsNone(connection.load_from_lerobot("robot", remote)) + + self.assertEqual(info, _catalog_metadata(connection, "robot__info")) + self.assertEqual(stats, _catalog_metadata(connection, "robot__stats")) + table = connection.get_table("robot") + self.assertEqual([], table.raw_table.tag_manager().list_tags()) + table.add(frames) + snapshots = connection.create_lerobot_tag("robot", "training") + self.assertEqual({ + "frames": 2, "info": 1, "stats": 1, "episodes": 1, "tasks": 1, + }, snapshots) + table.add(frames) + self.assertEqual(2, table.scan(tag_name="training").to_arrow().num_rows) + self.assertEqual(3, table.scan().to_arrow().num_rows) + for component, snapshot_id in snapshots.items(): + name = "robot" if component == "frames" else "robot__" + component + self.assertEqual(snapshot_id, connection.catalog.get_tag( + connection._identifier(name), "training").snapshot.id) + + with patch.object(connection.catalog, "create_tag") as create_tag: + with self.assertRaisesRegex(ValueError, "already points"): + connection.create_lerobot_tag("robot", "training") + create_tag.assert_not_called() + + create_tag = connection.catalog.create_tag + attempts = [] + + def fail_second_component(*args, **kwargs): + attempts.append(args[0]) + if len(attempts) == 2: + raise RuntimeError("tag failed") + return create_tag(*args, **kwargs) + + with patch.object(connection.catalog, "create_tag", + side_effect=fail_second_component): + with self.assertRaisesRegex(RuntimeError, "tag failed"): + connection.create_lerobot_tag("robot", "retry") + self.assertFalse(table.raw_table.tag_manager().tag_exists("retry")) + self.assertEqual(3, connection.create_lerobot_tag( + "robot", "retry")["frames"]) + self.assertEqual(3, table.scan(tag_name="retry").to_arrow().num_rows) + + connection.catalog.drop_table(connection._identifier("robot__tasks")) + with patch.object(connection.catalog, "create_tag") as create_tag: + with self.assertRaises(TableNotExistException): + connection.create_lerobot_tag("robot", "incomplete") + create_tag.assert_not_called() + + (source / "meta" / "stats.json").unlink() + with patch( + "pypaimon.multimodal.lerobot.source._SourceFileIO", + return_value=_RemoteLeRobotFileIO(source, remote)): + connection.load_from_lerobot("no_stats", remote, tag_name="ready") + with self.assertRaises(TableNotExistException): + connection.get_table("no_stats__stats") + self.assertEqual({"frames": 1, "info": 1, "episodes": 1, "tasks": 1}, + connection.create_lerobot_tag("no_stats", "training")) + def test_self_contained_import_rejects_table_branches(self): with self.assertRaisesRegex(ValueError, "does not support"): _managed_table_options("db.robot$branch_dev") @@ -691,13 +844,16 @@ 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"]) + stored_stats = { + row["key"]: json.loads(row["value"]) + for row in metadata["stats_table"].to_pylist() + } self.assertTrue(np.isnan(stored_stats["mean"])) self.assertTrue(np.isinf(stored_stats["max"])) finally: shutil.rmtree(temp_dir, ignore_errors=True) - def test_invalid_fps_creates_no_snapshot_or_manifest(self): + def test_invalid_fps_creates_no_table(self): temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_fps_")) try: source = temp_dir / "source" @@ -952,10 +1108,10 @@ def _create_image_dataset(root): def test_import_infers_schema_and_preserves_episodes(self): import pandas as pd - version_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "robot_data", self.image_source, batch_size=2) - self.assertEqual(1, version_id) + self.assertIsNone(result) table = self.connection.get_table("robot_data") schema = table.raw_table.fields @@ -970,8 +1126,6 @@ def test_import_infers_schema_and_preserves_episodes(self): self.assertEqual("BIGINT NOT NULL", types["episode_index"]) self.assertEqual("BLOB NOT NULL", types["observation.image"]) self.assertNotIn("dataset_id", types) - self.assertNotIn("metadata_version", types) - self.assertNotIn("version_id", types) self.assertNotIn("task", types) rows = table.scan().select([ @@ -995,31 +1149,21 @@ 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"]) - manifests = _catalog_rows(self.connection, "robot_data__versions") - self.assertEqual(1, len(manifests)) - manifest = manifests[0] - self.assertEqual(version_id, manifest["version_id"]) - self.assertEqual("v3.0", json.loads( - manifest["info_json"])["codebase_version"]) - self.assertIsNotNone(manifest["stats_json"]) self.assertEqual( - {"version_id", "info_json", "stats_json", "has_subtasks"}, - set(manifest)) - self.assertFalse(manifest["has_subtasks"]) - tag = str(manifest["version_id"]) + json.loads((self.image_source / "meta" / "info.json").read_text()), + _catalog_metadata(self.connection, "robot_data__info"), + ) self.assertEqual( - 1, - self.connection.catalog.get_tag( - table.identifier, tag).snapshot.id, + json.loads((self.image_source / "meta" / "stats.json").read_text()), + _catalog_metadata(self.connection, "robot_data__stats"), ) - 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, - ) + for name in ("robot_data__info", "robot_data__stats"): + fields = self.connection.catalog.get_table( + self.connection._identifier(name)).fields + self.assertEqual({"key": "STRING", "value": "STRING"}, { + field.name: str(field.type) for field in fields + }) + self.assertEqual([], table.raw_table.tag_manager().list_tags()) episodes = _catalog_rows(self.connection, "robot_data__episodes") episode_fields = { @@ -1027,7 +1171,6 @@ def test_import_infers_schema_and_preserves_episodes(self): self.connection._identifier( "robot_data__episodes")).fields } - self.assertNotIn("version_id", episode_fields) source_episode_schema = pq.read_schema(next( (self.image_source / "meta" / "episodes").rglob("*.parquet"))) self.assertTrue(_target_schema( @@ -1049,7 +1192,6 @@ def test_import_infers_schema_and_preserves_episodes(self): field.name for field in self.connection.catalog.get_table( self.connection._identifier("robot_data__tasks")).fields } - self.assertNotIn("version_id", task_fields) self.assertTrue(_target_schema( self.connection.catalog.get_table(self.connection._identifier( "robot_data__tasks")) @@ -1163,8 +1305,8 @@ def test_import_publishes_optional_subtasks(self): )) pq.write_table(subtasks, source / "meta" / "subtasks.parquet") - version_id = self.connection.load_from_lerobot( - "with_subtasks", source) + result = self.connection.load_from_lerobot( + "with_subtasks", source, tag_name="training") frames = self.connection.get_table("with_subtasks") self.assertNotIn("subtask", [ @@ -1188,13 +1330,12 @@ def test_import_publishes_optional_subtasks(self): _restore_pandas_metadata( subtasks_table, subtasks_arrow).to_pandas(), ) - self.assertTrue(_catalog_rows( - self.connection, "with_subtasks__versions")[0]["has_subtasks"]) + self.assertIsNone(result) self.assertEqual( 1, self.connection.catalog.get_tag( self.connection._identifier("with_subtasks__subtasks"), - str(version_id), + "training", ).snapshot.id, ) @@ -1209,8 +1350,9 @@ def test_import_preserves_quoted_database_name(self): self.assertEqual([ "robot", "robot__episodes", + "robot__info", + "robot__stats", "robot__tasks", - "robot__versions", ], sorted(table_names)) def test_import_reuses_validated_episode_metadata(self): @@ -1271,7 +1413,7 @@ def test_frame_controls_must_match_published_episode_metadata(self): self.connection.load_from_lerobot(table_name, source) self.connection.get_table(table_name) self.assertEqual([], _catalog_rows( - self.connection, table_name + "__versions")) + self.connection, table_name + "__info")) def test_task_text_remains_in_published_task_mapping(self): source = self.temp_dir / "reordered_tasks" @@ -1320,7 +1462,7 @@ def test_episode_tasks_must_exactly_match_frame_tasks(self): "extra_episode_task", source) self.connection.get_table("extra_episode_task") self.assertEqual([], _catalog_rows( - self.connection, "extra_episode_task__versions")) + self.connection, "extra_episode_task__info")) def test_nonempty_dataset_cannot_publish_without_tasks(self): source = self.temp_dir / "missing_tasks" @@ -1347,7 +1489,7 @@ def test_nonempty_dataset_cannot_publish_without_tasks(self): self.connection.load_from_lerobot("missing_tasks", source) self.connection.get_table("missing_tasks") self.assertEqual([], _catalog_rows( - self.connection, "missing_tasks__versions")) + self.connection, "missing_tasks__info")) def test_oss_source_streams_parquet_and_preserves_episodes(self): source = "oss://source-bucket/robot-images" @@ -1356,13 +1498,13 @@ def test_oss_source_streams_parquet_and_preserves_episodes(self): with patch( "pypaimon.multimodal.lerobot.source._SourceFileIO", return_value=source_file_io): - version_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "oss_images", source, batch_size=2, ) - self.assertEqual(1, version_id) + self.assertIsNone(result) table = self.connection.get_table("oss_images") rows = table.scan().select([ "episode_index", "frame_index", "index", "task_index" @@ -1433,14 +1575,12 @@ def test_tag_falls_back_for_catalogs_without_tag_api(self): self.connection.catalog, "create_tag", side_effect=NotImplementedError): - version_id = self.connection.load_from_lerobot( - "tag_fallback", self.image_source) + result = self.connection.load_from_lerobot( + "tag_fallback", self.image_source, tag_name="training") - manifest = _catalog_rows( - self.connection, "tag_fallback__versions")[0] - tag = str(manifest["version_id"]) + tag = "training" table = self.connection.get_table("tag_fallback") - self.assertEqual(1, version_id) + self.assertIsNone(result) self.assertEqual( table.raw_table.snapshot_manager().get_latest_snapshot().id, table.raw_table.tag_manager().get(tag).id, @@ -1461,27 +1601,27 @@ def create_then_lose_response(*args, **kwargs): self.connection.catalog, "create_tag", side_effect=create_then_lose_response): - version_id = self.connection.load_from_lerobot( - "tag_response_loss", self.image_source) + result = self.connection.load_from_lerobot( + "tag_response_loss", self.image_source, tag_name="training") self.assertTrue(lost[0]) - self.assertEqual(1, version_id) - self.assertEqual( - [1], - [row["version_id"] for row in _catalog_rows( - self.connection, "tag_response_loss__versions")]) + self.assertIsNone(result) + self.assertEqual(1, self.connection.catalog.get_tag( + self.connection._identifier("tag_response_loss"), + "training").snapshot.id) - def test_tag_failure_remains_unpublished(self): + def test_tag_failure_leaves_imported_data(self): with patch( "pypaimon.multimodal.lerobot.metadata._create_tag", side_effect=RuntimeError("tag failed")): with self.assertRaisesRegex(RuntimeError, "tag failed"): self.connection.load_from_lerobot( - "failed_publish", self.image_source) + "failed_publish", self.image_source, tag_name="training") - self.connection.get_table("failed_publish") - self.assertEqual([], _catalog_rows( - self.connection, "failed_publish__versions")) + self.assertEqual(5, self.connection.get_table( + "failed_publish").scan().to_arrow().num_rows) + self.assertEqual("v3.0", _catalog_metadata( + self.connection, "failed_publish__info")["codebase_version"]) def test_existing_companion_is_rejected(self): self.connection.load_from_lerobot( @@ -1508,9 +1648,9 @@ def test_invalid_target_options_do_not_leave_table(self): self.connection.catalog.get_table( self.connection._identifier("invalid_options")) - version_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "invalid_options", self.image_source) - self.assertEqual(1, version_id) + self.assertIsNone(result) def test_target_open_failure_leaves_created_table(self): original_get = self.connection.get_table @@ -1544,14 +1684,12 @@ def open_with_failing_close(*args, **kwargs): api, "_open_resolved_dataset", side_effect=open_with_failing_close): - version_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "close_failure", self.image_source) - self.assertEqual(1, version_id) - self.assertEqual( - [1], - [row["version_id"] for row in _catalog_rows( - self.connection, "close_failure__versions")]) + self.assertIsNone(result) + self.assertEqual("v3.0", _catalog_metadata( + self.connection, "close_failure__info")["codebase_version"]) def test_source_close_failure_does_not_override_success(self): source = "oss://source-bucket/robot-images" @@ -1563,14 +1701,12 @@ def test_source_close_failure_does_not_override_success(self): with patch( "pypaimon.multimodal.lerobot.source._SourceFileIO", return_value=source_file_io): - version_id = self.connection.load_from_lerobot( + result = self.connection.load_from_lerobot( "source_close_failure", source) - self.assertEqual(1, version_id) - self.assertEqual( - [1], - [row["version_id"] for row in _catalog_rows( - self.connection, "source_close_failure__versions")]) + self.assertIsNone(result) + self.assertEqual("v3.0", _catalog_metadata( + self.connection, "source_close_failure__info")["codebase_version"]) def test_existing_target_is_rejected(self): info = json.loads((self.image_source / "meta" / "info.json").read_text()) @@ -1613,16 +1749,16 @@ def prepare_then_wait(*args, **kwargs): "concurrent", self.image_source) finally: release.set() - version_id = future.result(timeout=30) + result = future.result(timeout=30) - self.assertEqual(1, version_id) + self.assertIsNone(result) self.assertEqual( 5, self.connection.get_table( "concurrent").scan().to_arrow().num_rows, ) - def test_concurrent_append_cannot_enter_published_version(self): + def test_concurrent_append_rejects_initial_import(self): from pypaimon.multimodal.lerobot import api original_write = api._write_dataset @@ -1660,7 +1796,7 @@ def append_then_write( self.connection.get_table("concurrent_append") self.assertEqual([], _catalog_rows( - self.connection, "concurrent_append__versions")) + self.connection, "concurrent_append__info")) if __name__ == "__main__":