This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new c947e5ceaf [python] Simplify LeRobot metadata tables and training tags
(#9676)
c947e5ceaf is described below
commit c947e5ceafb4e144ea2605bc505ec6206700f9ae
Author: Yann Byron <[email protected]>
AuthorDate: Tue Sep 8 14:18:01 2026 +0800
[python] Simplify LeRobot metadata tables and training tags (#9676)
---
docs/docs/pypaimon/multimodal-api.mdx | 65 +++--
paimon-python/README.md | 20 +-
paimon-python/pypaimon/multimodal/connection.py | 15 +-
paimon-python/pypaimon/multimodal/lerobot/api.py | 36 ++-
.../pypaimon/multimodal/lerobot/metadata.py | 146 ++++++-----
.../pypaimon/tests/multimodal_lerobot_test.py | 290 +++++++++++++++------
6 files changed, 392 insertions(+), 180 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-api.mdx
b/docs/docs/pypaimon/multimodal-api.mdx
index c50acc7eb1..bdf780859f 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,
-`<table>__versions`, `<table>__episodes`, `<table>__tasks`, and an optional
-`<table>__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 `<table>__versions` to publish the
-version.
+repository. It derives the schema from `meta/info.json` and writes one row per
+frame. The import creates `<table>__episodes`, `<table>__tasks`,
`<table>__info`,
+and optional `<table>__stats` and `<table>__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.<component>-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 `<table>__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 ac766d5c86..f864ff265b 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 31a1bdc884..4cea5786db 100644
--- a/paimon-python/pypaimon/multimodal/connection.py
+++ b/paimon-python/pypaimon/multimodal/connection.py
@@ -141,18 +141,29 @@ class MultimodalConnection:
*,
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 6078ac7b24..84e3d2eed5 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/api.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/api.py
@@ -22,12 +22,14 @@ from typing import Mapping, Optional
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 ee3879444a..0f568f4769 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 @@ import pyarrow.parquet as pq
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 @@ _COMPANION_OPTION_KEYS = {
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 @@ class _EpisodeIndex:
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 97c82b43d4..196a795be3 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -36,12 +36,14 @@ 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.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 @@ class LeRobotValidationTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
))
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 @@ class LeRobotImportTest(unittest.TestCase):
_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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
"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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
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 @@ class LeRobotImportTest(unittest.TestCase):
"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 @@ class LeRobotImportTest(unittest.TestCase):
self.connection.get_table("concurrent_append")
self.assertEqual([], _catalog_rows(
- self.connection, "concurrent_append__versions"))
+ self.connection, "concurrent_append__info"))
if __name__ == "__main__":
