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