diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 63f39d5aea54..7367541c7e74 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -46,22 +46,18 @@ with row tracking, no primary key, and `bucket=-1`. The indices below identify LeRobot records; they are independent of Paimon's internal row IDs. These relationships are used by the reader and do not declare database foreign keys. -The diagram and field tables below describe **list imports**, including a -one-element list such as `source=["/data/source_a"]`. This form adds -`dataset_index` to frames and Episodes and also supports later appends. -[Single-source shorthand](#single-source-shorthand) has a smaller schema. +The same three-table schema is used for one source and for a list of sources. -[![LeRobot metadata files map to the Dataset table, Episode Parquet maps to the Episode table, and frame Parquet plus media map to the Frame table. Dataset and Episode indices relate the three tables.](/img/lerobot-data-model.svg)](/img/lerobot-data-model.svg) +[![LeRobot metadata files map to the Dataset table, Episode Parquet maps to the Episode table, and frame Parquet plus media map to the Frame table. Dataset IDs and Episode indices relate the three tables.](/img/lerobot-data-model.svg)](/img/lerobot-data-model.svg) ### Dataset fields `robot_data_datasets` has one row for each source, even when multiple sources -have the same label. Source ordinals follow the input order; appends allocate -new ordinals after the existing sources. +have the same label. UUIDv7 values retain the order of serial imports and appends. | Field | Paimon type | Meaning | | --- | --- | --- | -| `index` | `BIGINT` | Source ordinal, starting at 0. Referenced by frame and Episode `dataset_index`. | +| `id` | `STRING` | UUIDv7 assigned to this source; referenced by frame and Episode `dataset_id`. | | `name` | `STRING`, nullable | Optional label supplied through `dataset_name`; omitted names are null. | | `info` | `STRING` | JSON containing the source's original `meta/info.json`, including `features`, `fps`, robot type, and counts. FPS is not a separate column. | | `tasks` | `ARRAY>` | Source-local task IDs and their text from `meta/tasks.parquet`. | @@ -81,7 +77,7 @@ forced to `BIGINT`. | Field | Representation | Meaning | | --- | --- | --- | -| `dataset_index` | `BIGINT` | Source ordinal referencing `_datasets.index`. | +| `dataset_id` | `STRING` | Source UUID referencing `_datasets.id`. | | `episode_index` | Source integer type | Episode number within that source, starting at 0. | | `dataset_from_index` | Source integer type | Inclusive start of the Episode's source-local frame `index` range. | | `dataset_to_index` | Source integer type | Exclusive end of that frame range. | @@ -90,7 +86,7 @@ forced to `BIGINT`. | `meta`, `data`, `videos`, `stats` | `VARIANT` columns | Grouped source fields describing metadata locations, data shards, per-camera video locations/timestamps, and per-Episode feature statistics. Empty groups are `{}`. | | Other source fields | Source type, or a `VARIANT` for a `/` prefix group | Custom Episode metadata, such as quality labels or camera calibration. | -For list imports, a field such as `data/chunk_index` becomes the +A field such as `data/chunk_index` becomes the `"chunk_index"` key inside `_episodes.data`. A field such as `videos/observation.images.front/from_timestamp` becomes the `"observation.images.front/from_timestamp"` key inside `_episodes.videos`. @@ -110,7 +106,7 @@ names, including dots such as `observation.state`, remain column names. | Field | Paimon type | Meaning | | --- | --- | --- | -| `dataset_index` | `BIGINT` | Source ordinal referencing `_datasets.index`. | +| `dataset_id` | `STRING` | Source UUID referencing `_datasets.id`. | | `index` | `BIGINT` | Frame position within the original source, across all its Episodes. | | `episode_index` | `BIGINT` | Episode containing this frame, within the same source. | | `frame_index` | `BIGINT` | Frame position within the Episode; restarts at 0 for each Episode. | @@ -119,11 +115,11 @@ names, including dots such as `observation.state`, remain column names. | `subtask_index` | `BIGINT`, nullable, when declared | ID resolved through that source's `_datasets.subtasks`; null means no subtask label. | | Feature columns | Scalars, `ARRAY`, or `BLOB` | Observations, actions, and other features, mapped as described below. | -Within a table group, `(dataset_index, index)` identifies a frame and -`(dataset_index, episode_index)` identifies an Episode. Task/subtask IDs are +Within a table group, `(dataset_id, index)` identifies a frame and +`(dataset_id, episode_index)` identifies an Episode. Task/subtask IDs are also scoped to a source. For example, two sources can both contain frame `index=0`, Episode `episode_index=0`, and task `task_index=0` with different -task texts; `dataset_index` keeps them distinct. The training reader derives +task texts; `dataset_id` keeps them distinct. The training reader derives continuous sample positions at read time without rewriting the stored IDs. ### Mapping the native LeRobot v3 format @@ -152,19 +148,6 @@ videos become `BLOB`, so each camera remains a feature column rather than a separate Paimon table. See [video import options](#load-lerobot-dataset-v3) for optional Episode remuxing. -### Single-source shorthand - -For a new import with a single source string and no `dataset_name` or non-empty -`extra_dataset_fields`, `load_from_lerobot("robot_data", "/data/source_a")` -still creates the same three tables. Its Dataset row has `index=0` and -`name=null`, but frames and Episodes do not gain a source `dataset_index`. -Episode fields such as `data/chunk_index` and `stats/action/mean` retain their -original names and types instead of being grouped into VARIANT columns. - -To use the source-indexed schema above, or to append more sources later, pass -`source=["/data/source_a"]`. Supplying `dataset_name` or non-empty -`extra_dataset_fields` also selects this schema, even for one source. - ## Load LeRobot Dataset v3 `load_from_lerobot` imports a local directory, FileIO URI, or Hugging Face @@ -224,7 +207,8 @@ source is rejected before creating tables to avoid losing original frame IDs. Sources must share robot type and compatible definitions for common features. Each source retains its own FPS from `info`; different source frame rates are handled automatically. Use `PaimonLeRobotDataset` to train from mixed-FPS groups. Pass `dataset_index=N` to -that reader to select one source. Frame-rate handling is automatic and does not +that reader to select the source at position `N` in UUID order; this reader +argument is not the stored UUID. Frame-rate handling is automatic and does not require an import option. ### Schema evolution @@ -308,8 +292,7 @@ removing suffixes, and names that overwrite standard `_datasets` columns are rejected before creating target tables. Training readers project only required source metadata and skip extra text columns. -Append another Dataset to a source-local group with `append_to_lerobot`. -Create the group by importing a list, including a one-element list for one source: +Append another Dataset to a source-local group with `append_to_lerobot`: ```python conn.load_from_lerobot("robot_merged", ["/data/source_a", "/data/source_b"]) diff --git a/docs/static/img/lerobot-data-model.svg b/docs/static/img/lerobot-data-model.svg index 9559f9a4d8a5..12790b49f0cd 100644 --- a/docs/static/img/lerobot-data-model.svg +++ b/docs/static/img/lerobot-data-model.svg @@ -18,7 +18,7 @@ specific language governing permissions and limitations under the License. --> LeRobot v3 files and the three Paimon tables - For list imports, source metadata becomes one Dataset row, Episode Parquet becomes Episode rows, and frame Parquet plus image or video payloads become Frame rows. Dataset indices identify sources, and dataset plus Episode indices relate frames to Episodes. These are logical relationships, not primary or foreign key constraints. + Source metadata becomes one Dataset row, Episode Parquet becomes Episode rows, and frame Parquet plus image or video payloads become Frame rows. A UUID links each source Dataset to its Episodes and frames. These are logical relationships, not primary or foreign key constraints. @@ -27,7 +27,7 @@ under the License. LeRobot files โ†’ Paimon table group - List imports: preserve source-local IDs and add dataset_index. + Preserve source-local IDs; link the three tables with a Dataset UUID. LeRobot v3 source files Paimon tables ยท prefix robot_data @@ -42,12 +42,12 @@ under the License. robot_data_datasets One row per source - index, name, info + id (UUIDv7), name, info tasks, subtasks, stats 1 source โ†’ many Episodes - dataset_index references source index + dataset_id references the source UUID meta/episodes/**/*.parquet Episode boundaries and task labels @@ -56,12 +56,12 @@ under the License. robot_data_episodes One row per Episode - dataset_index, episode_index, frame range + dataset_id, episode_index, frame range length, tasks, metadata / statistics 1 Episode โ†’ many frames - match dataset_index + episode_index + match dataset_id + episode_index data/**/*.parquet Image payloads or videos/ MP4 shards @@ -70,7 +70,7 @@ under the License. robot_data_frames One row per frame - dataset_index, indices, timestamp, task IDs + dataset_id, indices, timestamp, task IDs Observations, actions, image / video BLOBs Logical relationships only. LeRobot indices are separate from Paimon's internal row IDs. diff --git a/paimon-python/pypaimon/multimodal/lerobot/aggregate.py b/paimon-python/pypaimon/multimodal/lerobot/aggregate.py index d1e9d5626848..d873484345ed 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/aggregate.py +++ b/paimon-python/pypaimon/multimodal/lerobot/aggregate.py @@ -18,6 +18,9 @@ import copy import json +import secrets +import time +import uuid from decimal import Decimal from contextlib import ExitStack, contextmanager @@ -34,8 +37,8 @@ from pypaimon.multimodal.lerobot.metadata import ( _EPISODE_CONTROL_COLUMNS, _append_arrow, _commit_metadata, _companion_identifier, _companion_table_identifiers, - _json_value, _load_dataset_metadata, _tag_snapshot_id, - _metadata_table, _prepare_metadata_tables, + _dataset_fps, _json_value, _load_dataset_metadata, _tag_snapshot_id, + _episode_task_indices, _metadata_table, _prepare_metadata_tables, _source_episode_tables, _validated_episode_tables, ) from pypaimon.multimodal.lerobot.schema import ( @@ -66,7 +69,7 @@ "total_chunks", "splits", "data_path", "video_path", "chunks_size", "data_files_size_in_mb", "video_files_size_in_mb", "features", } -_DATASET_INDEX_COLUMN = "dataset_index" +_DATASET_ID_COLUMN = "dataset_id" _EPISODE_VARIANT_GROUPS = ("meta", "data", "videos", "stats") _EPISODE_EXTRA_VARIANT_GROUP_DESCRIPTION = "pypaimon.lerobot.episode-variant-group" _VARIANT_TYPE = pa.struct([ @@ -82,7 +85,7 @@ pa.field("subtask", pa.string(), nullable=False), ])) _DATASETS_SCHEMA = pa.schema([ - pa.field("index", pa.int64(), nullable=False), + pa.field("id", pa.string(), nullable=False), pa.field("name", pa.string()), pa.field("info", pa.string(), nullable=False), pa.field("tasks", _SOURCE_TASKS_TYPE, nullable=False), @@ -91,6 +94,16 @@ ]) +def _next_dataset_uuid(previous=None): + """Generate a UUIDv7 after the last source, preserving serial append order.""" + timestamp_ms = time.time_ns() // 1_000_000 + if previous is not None: + timestamp_ms = max(timestamp_ms, (uuid.UUID(previous).int >> 80) + 1) + value = ((timestamp_ms << 80) | (7 << 76) | + (secrets.randbits(12) << 64) | (2 << 62) | secrets.randbits(62)) + return str(uuid.UUID(int=value)) + + def _json_table(table): return {row["key"]: json.loads(row["value"]) for row in table.to_pylist()} @@ -107,7 +120,7 @@ def _recover_incomplete_append(connection, table_name): identifiers = _companion_table_identifiers(frames) datasets = connection.catalog.get_table(identifiers["datasets"]) episodes = connection.catalog.get_table(identifiers["episodes"]) - source_rows = _read_arrow(datasets, ["index", "info"]).sort_by("index").to_pylist() + source_rows = _read_arrow(datasets, ["id", "info"]).sort_by("id").to_pylist() if not source_rows: raise ValueError("LeRobot append recovery requires Dataset metadata.") ranges = _source_dataset_ranges(source_rows) @@ -341,10 +354,10 @@ def _plan_merge(plans, extra_dataset_fields=None, schema_evolution=True, existing_sources=None): first = plans[0] mixed_fps = len({plan["metadata"]["fps"] for plan in plans}) > 1 - reserved_frame_fields = {_DATASET_INDEX_COLUMN} + reserved_frame_fields = {_DATASET_ID_COLUMN} for plan in plans: if reserved_frame_fields.intersection(plan["schema"].names): - raise ValueError("LeRobot source frame features use reserved source index columns.") + raise ValueError("LeRobot source frame features use reserved dataset_id column.") frame_schema = _merge_schemas([plan["schema"] for plan in plans], schema_evolution, _V3_REQUIRED_FEATURE_DTYPES) features, contracts = {}, {} @@ -366,9 +379,9 @@ def _plan_merge(plans, extra_dataset_fields=None, schema_evolution=True, groups = tuple(dict.fromkeys(group for plan in plans for group in _episode_variant_groups( plan["metadata"]["episodes_schema"]))) - if _DATASET_INDEX_COLUMN in groups: - raise ValueError("LeRobot Episode VARIANT group uses reserved column dataset_index.") - reserved_episode_fields = {_DATASET_INDEX_COLUMN, *groups} + if _DATASET_ID_COLUMN in groups: + raise ValueError("LeRobot Episode VARIANT group uses reserved column dataset_id.") + reserved_episode_fields = {_DATASET_ID_COLUMN, *groups} for ordinal, plan in enumerate(plans): conflicts = reserved_episode_fields.intersection( plan["metadata"]["episodes_schema"].names) @@ -414,7 +427,7 @@ def _plan_merge(plans, extra_dataset_fields=None, schema_evolution=True, metadata["tasks_table"] = _vocabulary(plans, "tasks", existing_sources) metadata["subtasks_table"] = _vocabulary(plans, "subtasks", existing_sources) for field in ( - pa.field(_DATASET_INDEX_COLUMN, pa.int64(), nullable=False), + pa.field(_DATASET_ID_COLUMN, pa.string(), nullable=False), pa.field("meta", _VARIANT_TYPE, nullable=False), pa.field("data", _VARIANT_TYPE, nullable=False), pa.field("videos", _VARIANT_TYPE, nullable=False), @@ -429,7 +442,7 @@ def _plan_merge(plans, extra_dataset_fields=None, schema_evolution=True, PaimonSchema.from_pyarrow_schema(metadata["episodes_schema"]) metadata["datasets_table"] = pa.Table.from_pylist([ { - "index": ordinal, + "id": plan[_DATASET_ID_COLUMN], "name": plan["dataset_name"], "info": json.dumps(_json_value(plan["info"]), ensure_ascii=False, separators=(",", ":")), @@ -474,8 +487,10 @@ def _plan_merge(plans, extra_dataset_fields=None, schema_evolution=True, if stats[0] is not None: decoded = [_json_table(s) for s in stats] names = dict.fromkeys(name for item in decoded for name in item) - merged_stats = {name: _merge_stats([item[name] for item in decoded if name in item]) - for name in names if name not in _CONTROLS} + merged_stats = { + name: (decoded[0][name] if len(decoded) == 1 else + _merge_stats([item[name] for item in decoded if name in item])) + for name in names if name not in _CONTROLS} metadata["info_table"] = _metadata_table(info) metadata["stats_table"] = (None if merged_stats is None else _metadata_table(merged_stats)) @@ -484,8 +499,8 @@ def _plan_merge(plans, extra_dataset_fields=None, schema_evolution=True, def _transform_frames(plan, batch, episode): batch = batch.append_column( - pa.field(_DATASET_INDEX_COLUMN, pa.int64(), nullable=False), - pa.array([plan[_DATASET_INDEX_COLUMN]] * batch.num_rows, type=pa.int64())) + pa.field(_DATASET_ID_COLUMN, pa.string(), nullable=False), + pa.array([plan[_DATASET_ID_COLUMN]] * batch.num_rows, type=pa.string())) for name in _CONTROLS: if name not in batch.column_names: continue @@ -526,10 +541,26 @@ def _merged_episodes(plans, schema): for group in _episode_variant_groups(plan["metadata"]["episodes_schema"]) if group not in _EPISODE_VARIANT_GROUPS)) for plan in plans: + validated = plan["metadata"]["episodes"] + task_indices = plan["metadata"]["task_indices"] + episode_offset = 0 for table in _source_episode_tables(plan["metadata"]): source_groups = set(_episode_variant_groups(table.schema)) records = [] for original in table.to_pylist(): + try: + expected = validated[episode_offset] + source_tasks = _episode_task_indices( + original["tasks"], task_indices, episode_offset) + except (IndexError, KeyError, TypeError, ValueError): + raise ValueError( + "LeRobot Episode metadata changed after validation.") from None + if (original["episode_index"] != episode_offset + or original["dataset_from_index"] != expected["dataset_from_index"] + or original["dataset_to_index"] != expected["dataset_to_index"] + or original["length"] != expected["length"] + or source_tasks != list(expected["task_indices"])): + raise ValueError("LeRobot Episode metadata changed after validation.") row = {name: original[name] for name in schema.names if name in original} for group in ("meta", "data", "videos"): @@ -546,7 +577,7 @@ def _merged_episodes(plans, schema): # original quantiles cannot describe the merged IDs. if feature not in _CONTROLS or not statistic.startswith("q"): stats.setdefault(feature, {})[statistic] = value - row[_DATASET_INDEX_COLUMN] = plan[_DATASET_INDEX_COLUMN] + row[_DATASET_ID_COLUMN] = plan[_DATASET_ID_COLUMN] episode = row["episode_index"] if "subtask_index" not in plan["control_stats"][episode]: if "subtask_index" in stats: @@ -556,7 +587,10 @@ def _merged_episodes(plans, schema): stats[name] = values row["stats"] = _variant_value(stats) records.append(row) + episode_offset += 1 yield pa.Table.from_pylist(records, schema=schema) + if episode_offset != len(validated): + raise ValueError("LeRobot Episode metadata changed after validation.") def _import_sources(connection, table_name, sources, batch_size, options, @@ -579,6 +613,7 @@ def _import_sources(connection, table_name, sources, batch_size, options, info = _load_hub_info(resolved) _require_v3(info, "source %d" % ordinal) _validate_info_paths(info) + _dataset_fps(info.get("fps")) _validated_counts(info, "source %d" % ordinal) _validate_v3_required_features(info) schema = _schema_from_info(info) @@ -591,7 +626,7 @@ def _import_sources(connection, table_name, sources, batch_size, options, dataset_resources.callback(_close_quietly, dataset, "dataset") plans.append({"dataset": dataset, "dataset_resources": dataset_resources, "source": resolved, "info": info, - _DATASET_INDEX_COLUMN: ordinal, + _DATASET_ID_COLUMN: None, "dataset_name": (dataset_name[ordinal] if dataset_name is not None else None), "schema": schema, "videos": videos, @@ -616,11 +651,11 @@ def _import_sources(connection, table_name, sources, batch_size, options, tag_name) is not None: raise ValueError("LeRobot tag %s already exists on %s." % (tag_name, component.identifier)) - if _DATASET_INDEX_COLUMN not in _target_schema(tables["episodes"]).names: + if _DATASET_ID_COLUMN not in _target_schema(tables["episodes"]).names: raise ValueError("LeRobot append requires a source-local table group.") source_rows = _read_arrow(tables["datasets"], [ - "index", "info", "tasks", "subtasks", "stats", - ]).sort_by("index").to_pylist() + "id", "info", "tasks", "subtasks", "stats", + ]).sort_by("id").to_pylist() if not source_rows: raise ValueError("LeRobot append requires an initialized table group.") ranges = _source_dataset_ranges(source_rows) @@ -636,12 +671,16 @@ def _import_sources(connection, table_name, sources, batch_size, options, "LeRobot group has incomplete %s data; call " "recover_lerobot_append before appending." % name) existing = (table, tables, source_rows, ranges) + previous = source_rows[-1]["id"] if existing is not None else None + for plan in plans: + previous = _next_dataset_uuid(previous) + plan[_DATASET_ID_COLUMN] = previous metadata, stats = _plan_merge( plans, extra_dataset_fields, schema_evolution if existing is None else True, source_rows if existing is not None else None) target_schema = metadata["frame_schema"] target_schema = target_schema.append( - pa.field(_DATASET_INDEX_COLUMN, pa.int64(), nullable=False)) + pa.field(_DATASET_ID_COLUMN, pa.string(), nullable=False)) if existing is None: table = _create_target_table(connection, table_name, target_schema, options, metadata, @@ -687,9 +726,9 @@ def _import_sources(connection, table_name, sources, batch_size, options, _validate_lerobot_schema(target_schema, _target_schema(table.raw_table), "append") # Validate every component before making any catalog changes. requests = [ - (table.raw_table, target_schema, (*_V3_REQUIRED_FEATURE_DTYPES, _DATASET_INDEX_COLUMN)), + (table.raw_table, target_schema, (*_V3_REQUIRED_FEATURE_DTYPES, _DATASET_ID_COLUMN)), (tables["episodes"], metadata["episodes_schema"], - (*_EPISODE_CONTROL_COLUMNS, _DATASET_INDEX_COLUMN)), + (*_EPISODE_CONTROL_COLUMNS, _DATASET_ID_COLUMN)), (tables["datasets"], metadata["datasets_table"].schema, tuple(field.name for field in _DATASETS_SCHEMA if not field.nullable)), ] @@ -706,15 +745,11 @@ def _import_sources(connection, table_name, sources, batch_size, options, tables = {name: connection.catalog.get_table(component.identifier) for name, component in tables.items()} frame_offset = ranges[-1][4] - dataset_offset = len(source_rows) for plan in plans: plan["frame_offset"] += frame_offset - plan[_DATASET_INDEX_COLUMN] += dataset_offset - datasets = metadata["datasets_table"].to_pylist() - for row in datasets: - row["index"] += dataset_offset metadata["datasets_table"] = pa.Table.from_pylist( - datasets, schema=_target_schema(tables["datasets"])) + metadata["datasets_table"].to_pylist(), + schema=_target_schema(tables["datasets"])) if existing is not None: stack.enter_context(_recover_failed_append(connection, table_name)) episode_schema = _target_schema(tables["episodes"]) @@ -771,7 +806,7 @@ def _import_sources(connection, table_name, sources, batch_size, options, _commit_metadata(connection, tables, tag_name, metadata, table.identifier, frames_snapshot, episodes_snapshot) else: - for column in ("dataset_index", "index"): + for column in ("dataset_id", "index"): if table.raw_table.create_global_index(column, index_type="btree") <= 0: raise RuntimeError("LeRobot append did not refresh %s index." % column) diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index 58e2cc0862a3..51fe28e005bb 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -16,47 +16,19 @@ """Public LeRobot import API.""" -import json import numbers import sys -from functools import partial from typing import Mapping, Optional -import pyarrow as pa - from pypaimon.catalog.catalog_exception import TableAlreadyExistException from pypaimon.multimodal.lerobot.metadata import ( _COMPANION_OPTION_KEYS, - _json_value, _companion_identifier, - _append_arrow_tables, - _load_dataset_metadata, _managed_table_options, - _prepare_metadata_tables, - _dataset_fps, - _commit_metadata, _validate_tag_name, - _validated_episode_tables, -) -from pypaimon.multimodal.lerobot.loader import _write_dataset -from pypaimon.multimodal.lerobot.schema import ( - _require_v3, - _schema_from_info, - _validate_v3_required_features, - _video_feature_names, -) -from pypaimon.multimodal.lerobot.source import ( - _close_quietly, - _import_lerobot_dataset, - _info_dict, - _load_hub_info, - _open_resolved_dataset, - _resolved_source, - _validate_info_paths, ) from pypaimon.multimodal.source_utils import ( _validated_source_options, - _validate_source_kerberos, ) from pypaimon.multimodal.lerobot.text_fields import _validated_text_fields @@ -119,149 +91,17 @@ def load_from_lerobot( if extra_dataset_fields is not None: from pypaimon.multimodal.lerobot.aggregate import _DATASETS_SCHEMA extra_dataset_fields = _validated_text_fields(extra_dataset_fields, _DATASETS_SCHEMA.names) - if append or isinstance(source, (list, tuple)) or dataset_name is not None \ - or extra_dataset_fields: - if not isinstance(source, (list, tuple)): - source = [source] - if not source: - raise ValueError("LeRobot sources must not be empty.") - if tag_name is not None: - _validate_tag_name(tag_name) - from pypaimon.multimodal.lerobot.aggregate import _import_sources - return _import_sources( - connection, table_name, source, batch_size, options, - _validated_source_options(source_options), tag_name, parallelism, - dataset_name, extra_dataset_fields, schema_evolution, append) - if parallelism > 1: - from pypaimon.multimodal.lerobot.ray_loader import _require_ray, _ray_source_uri - _require_ray() - source = _ray_source_uri(source) - + if not isinstance(source, (list, tuple)): + source = [source] + if not source: + raise ValueError("LeRobot sources must not be empty.") 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") - with _resolved_source(source, validated_source_options) as ( - resolved_source, local_info): - if local_info is None: - local_info = _load_hub_info(resolved_source) - _require_v3(local_info, resolved_source.path) - _validate_info_paths(local_info) - _schema_from_info(local_info) - _dataset_fps(local_info.get("fps")) - _validated_counts(local_info, resolved_source.path) - _validate_v3_required_features(local_info) - video_fields = _video_feature_names(local_info) - LeRobotDataset = (_import_lerobot_dataset() - if parallelism == 1 and resolved_source.file_io is None - else None) - dataset = _open_resolved_dataset( - LeRobotDataset, - resolved_source, - local_info, - download_videos=bool(video_fields), - ) - try: - info = _info_dict(dataset.meta.info) - _require_v3(info, resolved_source.path) - _validated_counts(info, resolved_source.path) - _validate_v3_required_features(info) - video_fields = _video_feature_names(info) - - lerobot_schema = _schema_from_info(info) - # Preserve source extensions omitted by LeRobot DatasetInfo.to_dict(). - metadata = _load_dataset_metadata( - dataset, local_info, resolved_source) - from pypaimon.multimodal.lerobot.aggregate import ( - _DATASETS_SCHEMA, _stats_json, _vocabulary_rows, - ) - metadata["datasets_table"] = pa.Table.from_pylist([{ - "index": 0, - "name": None, - "info": json.dumps(_json_value(local_info), ensure_ascii=False), - "tasks": _vocabulary_rows(metadata["tasks_table"], "task"), - "subtasks": _vocabulary_rows(metadata["subtasks_table"], "subtask"), - "stats": _stats_json(metadata["stats_table"]), - }], schema=_DATASETS_SCHEMA) - _import_dataset( - connection, - table_name, - dataset, - info, - resolved_source, - lerobot_schema, - batch_size, - options, - metadata, - tag_name, - video_fields, - parallelism, - validated_source_options, - ) - finally: - close = getattr(dataset, "close", None) - if callable(close): - _close_quietly(dataset, "dataset") - - -def _import_dataset( - connection, - table_name, - dataset, - info, - source, - source_schema, - batch_size, - options, - metadata, - tag_name, - video_fields, - parallelism=1, - source_options=None): - table = _create_target_table( - connection, - table_name, - source_schema, - options, - metadata, - video_fields, - ) - tables = _prepare_metadata_tables( - connection, table.raw_table, metadata) - episodes_snapshot_id = _append_arrow_tables( - tables["episodes"], - _validated_episode_tables(metadata, video_fields), - ) - frames_snapshot_id = None - if int(info["total_frames"]) > 0: - if parallelism > 1: - from pypaimon.multimodal.lerobot.ray_loader import _write_dataset_ray - write_dataset = partial( - _write_dataset_ray, parallelism=parallelism, - source_options=source_options) - else: - write_dataset = _write_dataset - frames_snapshot_id = write_dataset( - table, - dataset, - info, - source, - source_schema, - batch_size, - metadata, - video_fields, - ) - _commit_metadata( - connection, - tables, - tag_name, - metadata, - table.identifier, - frames_snapshot_id, - episodes_snapshot_id, - ) + from pypaimon.multimodal.lerobot.aggregate import _import_sources + return _import_sources( + connection, table_name, source, batch_size, options, + _validated_source_options(source_options), tag_name, parallelism, + dataset_name, extra_dataset_fields, schema_evolution, append) def _validated_counts(info, source): diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 4dd2880fbde7..314fb6455bc2 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -26,6 +26,7 @@ import pickle import sys import struct +import uuid import zlib from abc import ABC, abstractmethod from collections import OrderedDict @@ -689,9 +690,6 @@ def __init__( if sources is None or dataset_index >= len(sources): raise ValueError("LeRobot source dataset_index %d does not exist." % dataset_index) - source = sources[dataset_index] - if source.get("index", source.get("dataset_index")) != dataset_index: - raise ValueError("LeRobot source dataset indices are not contiguous.") begin, end = _source_dataset_ranges(sources)[dataset_index][1:3] episodes = list(range(begin, end)) self.tag_name = tag_name @@ -734,6 +732,9 @@ class _SourceEpisodeFrameReader: def __init__(self, reader, datasets, vocabulary_maps=None): self._reader = reader ranges = _source_dataset_ranges(datasets) + self._physical_ids = [source["id"] for source in datasets] + self._id_to_ordinal = {value: ordinal + for ordinal, value in enumerate(self._physical_ids)} self._frame_starts = [source[3] for source in ranges] self._frame_ends = [source[4] for source in ranges] self._episode_offsets = pa.array( @@ -761,8 +762,8 @@ def __init__(self, reader, datasets, vocabulary_maps=None): def read_indices(self, indices, columns): mapped = [name for name in self._vocabulary if name in columns] projection = list(columns) - if "dataset_index" not in projection: - projection.append("dataset_index") + if "dataset_id" not in projection: + projection.append("dataset_id") if "index" not in projection: projection.append("index") keys = [] @@ -770,20 +771,23 @@ def read_indices(self, indices, columns): dataset_index = bisect.bisect_right(self._frame_ends, index) if dataset_index >= len(self._frame_ends) or index < 0: raise IndexError("LeRobot frame index %d is out of range." % index) - keys.append((dataset_index, index - self._frame_starts[dataset_index])) + keys.append((self._physical_ids[dataset_index], + index - self._frame_starts[dataset_index])) data = self._reader.read_keys(keys, projection) if data.num_rows == 0: return data.select(columns) expected = set(keys) - actual = list(zip(data["dataset_index"].to_pylist(), + actual = list(zip(data["dataset_id"].to_pylist(), data["index"].to_pylist())) if len(actual) != len(expected) or set(actual) != expected: raise ValueError("LeRobot frame keys are missing or duplicated.") data = data.set_column( data.schema.get_field_index("index"), data.schema.field("index"), - pa.array([self._frame_starts[dataset_index] + local - for dataset_index, local in actual], type=pa.int64())) - source_ids = data["dataset_index"] + pa.array([self._frame_starts[self._id_to_ordinal[dataset_id]] + local + for dataset_id, local in actual], type=pa.int64())) + source_ids = pa.array([self._id_to_ordinal[value] + for value in data["dataset_id"].to_pylist()], + type=pa.int64()) if "episode_index" in columns: try: local_ids = data["episode_index"] @@ -1089,12 +1093,12 @@ def _load_dataset(table, tag_name): frames = _component_table(catalog, raw_table, tag_name) episodes_table = _component_table( catalog, catalog.get_table(identifiers["episodes"]), tag_name) - source_local_layout = "dataset_index" in _target_schema(episodes_table).names + source_local_layout = "dataset_id" in _target_schema(episodes_table).names datasets_table = _component_table( catalog, catalog.get_table(identifiers["datasets"]), tag_name) source_rows = _read_arrow(datasets_table, [ - "index", "name", "info", "tasks", "subtasks", "stats", - ]).sort_by("index").to_pylist() + "id", "name", "info", "tasks", "subtasks", "stats", + ]).sort_by("id").to_pylist() datasets = ([{key: value for key, value in row.items() if key not in ("tasks", "subtasks", "stats")} for row in source_rows] if source_local_layout else None) @@ -1225,8 +1229,10 @@ def _merged_source_stats(sources, episodes_table): raise ValueError("LeRobot sources disagree on stats metadata.") decoded = [json.loads(value) for value in values] names = dict.fromkeys(name for value in decoded for name in value) - merged = {name: _merge_stats([value[name] for value in decoded if name in value]) - for name in names if name not in _CONTROLS} + merged = { + name: (decoded[0][name] if len(decoded) == 1 else + _merge_stats([value[name] for value in decoded if name in value])) + for name in names if name not in _CONTROLS} columns = _target_schema(episodes_table).names if "stats" in columns: stats_rows = _read_arrow(episodes_table, ["stats"])["stats"].to_pylist() @@ -1288,12 +1294,21 @@ def _episode_dataset(table, datasets=None): if datasets is not None: # Merged Episodes may contain Arrow extensions unsupported by HF Dataset. - projection = _EPISODE_CONTROL_COLUMNS + ["dataset_index"] + projection = _EPISODE_CONTROL_COLUMNS + ["dataset_id"] else: projection = [name for name in _target_schema(table).names if not name.startswith("stats/") and name != "stats"] data = _read_arrow(table, projection) if datasets is not None: + physical_ids = [dataset["id"] for dataset in datasets] + ordinals = pc.index_in(data["dataset_id"], + value_set=pa.array(physical_ids, type=pa.string())) + if ordinals.null_count: + raise ValueError("LeRobot Episode references an unknown Dataset UUID.") + ordinals = pc.cast(ordinals, pa.int64()) + data = data.set_column( + data.schema.get_field_index("dataset_id"), + pa.field("dataset_index", pa.int64()), ordinals) data = data.sort_by([ ("dataset_index", "ascending"), ("episode_index", "ascending"), @@ -1417,10 +1432,16 @@ def _source_dataset_ranges(datasets): """Derive source ranges from the original counts in each source's info.""" result = [] episode_offset = frame_offset = 0 + previous_id = None for ordinal, dataset in enumerate(datasets): try: - if operator.index(dataset.get("index", dataset.get("dataset_index"))) != ordinal: - raise ValueError("index must be contiguous") + identifier = dataset["id"] + if not isinstance(identifier, str) or str(uuid.UUID(identifier)) != identifier \ + or uuid.UUID(identifier).version != 7: + raise ValueError("id must be a canonical UUIDv7 string") + if previous_id is not None and identifier <= previous_id: + raise ValueError("id must be unique and sorted") + previous_id = identifier info = json.loads(dataset["info"] if "info" in dataset else dataset["info_json"]) if not isinstance(info, dict): raise ValueError("info must be a JSON object") diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index e39c2bffb118..7e40fce747ca 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -261,9 +261,9 @@ def _commit_metadata( _require_initial_snapshot("datasets", snapshot_id) component_snapshots.append((tables["datasets"].identifier, snapshot_id)) frames_table = connection.catalog.get_table(frames_identifier) - if "dataset_index" in _target_schema(frames_table).names: + if "dataset_id" in _target_schema(frames_table).names: frames_snapshot_id = _build_initial_btree( - frames_table, "dataset_index", frames_snapshot_id) + frames_table, "dataset_id", frames_snapshot_id) frames_snapshot_id = _build_initial_btree( frames_table, "index", frames_snapshot_id) # Tag the root last so a failed component tag does not expose a root tag. @@ -568,6 +568,24 @@ def read_schema(path): return {"paths": paths, "schema": schema} +def _episode_task_indices(names, task_indices, episode_index): + names = names or [] + if isinstance(names, str): + names = [names] + if task_indices and not names: + raise ValueError( + "LeRobot Episode %d does not declare any task." % episode_index) + try: + result = [task_indices[str(name)] for name in names] + except (KeyError, TypeError) as error: + raise ValueError( + "LeRobot Episode %d refers to an unknown task." % episode_index + ) from error + if len(set(result)) != len(result): + raise ValueError("LeRobot Episode %d repeats a task." % episode_index) + return result + + def _validated_episode_tables(metadata, video_fields=(), check_episode_fps=False): episodes = _EpisodeIndex(video_fields) expected_begin = 0 @@ -619,23 +637,8 @@ def _validated_episode_tables(metadata, video_fields=(), check_episode_fps=False 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) + task_indices = _episode_task_indices( + columns["tasks"][offset].as_py(), metadata["task_indices"], index) video_values = {} for field in video_fields: prefix = "videos/%s/" % field diff --git a/paimon-python/pypaimon/multimodal/lerobot/ray_loader.py b/paimon-python/pypaimon/multimodal/lerobot/ray_loader.py index a67620aa9ccc..64900ba95660 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/ray_loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/ray_loader.py @@ -197,7 +197,7 @@ def arguments(): "subtask_indices": plan["metadata"]["subtask_indices"], }) mapping_ref = ray.put({key: plan[key] for key in ( - "dataset_index", "frame_offset", "episode_offset", + "dataset_id", "frame_offset", "episode_offset", "task_index", "subtask_index") if key in plan}) for group in _task_groups(plan["info"], plan["dataset"]._episodes, plan["videos"], parallelism): diff --git a/paimon-python/pypaimon/multimodal/lerobot/reader.py b/paimon-python/pypaimon/multimodal/lerobot/reader.py index 01f89321d46c..5cb7d865114f 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/reader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/reader.py @@ -49,18 +49,18 @@ def read_keys(self, keys, columns): from pypaimon.common.predicate_builder import PredicateBuilder by_dataset = {} - for dataset_index, index in keys: - by_dataset.setdefault(dataset_index, []).append(index) + for dataset_id, index in keys: + by_dataset.setdefault(dataset_id, []).append(index) builder = self._table.new_read_builder().new_predicate_builder() predicate = PredicateBuilder.or_predicates([ PredicateBuilder.and_predicates([ - builder.equal("dataset_index", dataset_index), + builder.equal("dataset_id", dataset_id), builder.is_in("index", indices), ]) - for dataset_index, indices in by_dataset.items() + for dataset_id, indices in by_dataset.items() ]) projection = list(columns) - for name in ("dataset_index", "index"): + for name in ("dataset_id", "index"): if name not in projection: projection.append(name) return self._read_predicate(predicate, projection) diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py b/paimon-python/pypaimon/multimodal/lerobot/writer.py index ebf7eb56f96f..fb71b3b79797 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/writer.py +++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py @@ -20,6 +20,7 @@ import copy import io import json +import uuid from typing import Mapping, Optional, Sequence import numpy as np @@ -32,7 +33,10 @@ from pypaimon.schema.arrow_schema import arrow_schemas_compatible from pypaimon.multimodal.arrow_utils import strict_arrow_table from pypaimon.multimodal.hdf5 import _SnapshotRecorder -from pypaimon.multimodal.lerobot.aggregate import _DATASETS_SCHEMA +from pypaimon.multimodal.lerobot.aggregate import ( + _DATASETS_SCHEMA, + _next_dataset_uuid, +) from pypaimon.multimodal.lerobot.metadata import ( _COMPANION_OPTION_KEYS, _EMPTY_EPISODES_SCHEMA, @@ -79,7 +83,7 @@ "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", } _EPISODE_STATE_COLUMNS = list(_EMPTY_EPISODES_SCHEMA.names) + [ - "stats/index/count", + "stats/index/count", "dataset_id", ] @@ -267,6 +271,8 @@ def __init__( raise ValueError("features must be a non-empty mapping.") if "task" in features: raise ValueError("task is managed by PaimonLeRobotWriter.") + if "dataset_id" in features: + raise ValueError("dataset_id is managed by PaimonLeRobotWriter.") video_fields = _video_feature_names({"features": features}) if video_fields: raise ValueError( @@ -282,7 +288,8 @@ def __init__( self._user_features = copy.deepcopy(dict(features)) self.features.update(copy.deepcopy(_DEFAULT_FEATURES)) _validate_v3_required_features({"features": self.features}) - self._source_schema = _schema_from_info({"features": self.features}) + self._source_schema = _schema_from_info({"features": self.features}).append( + pa.field("dataset_id", pa.string(), nullable=False)) metadata = self._writer_metadata(self.features) self._episodes_schema = metadata["episodes_schema"] create_options = dict(options or {}) @@ -327,7 +334,8 @@ def __init__( ) (self.num_frames, self.num_episodes, self._task_indices, - self._stats, stored_subtasks) = self._load_existing_state() + self._stats, stored_subtasks, self._dataset_uuid) = \ + self._load_existing_state() if stored_subtasks is not None: if requested_subtasks is not None \ and requested_subtasks != stored_subtasks: @@ -356,7 +364,8 @@ def __init__( @staticmethod def _writer_metadata(features): return { - "episodes_schema": _episode_schema(features), + "episodes_schema": _episode_schema(features).append( + pa.field("dataset_id", pa.string(), nullable=False)), "datasets_table": pa.Table.from_pylist([], schema=_DATASETS_SCHEMA), } @@ -388,7 +397,7 @@ def _load_existing_state(self): for table in self._metadata_tables.values()): raise ValueError( "Existing LeRobot table group state is inconsistent.") - return 0, 0, {}, None, None + return 0, 0, {}, None, None, _next_dataset_uuid() properties = snapshot.properties or {} if _STATE_VERSION in properties: state = self._state_from_snapshot_properties(properties) @@ -403,9 +412,17 @@ def _load_existing_state(self): def _state_from_companion_tables(self): sources = _read_arrow(self._metadata_tables["datasets"]).to_pylist() - if len(sources) != 1 or sources[0]["index"] != 0: + if len(sources) != 1: raise ValueError("Existing LeRobot Dataset metadata is invalid.") source_row = sources[0] + try: + source_id = source_row["id"] + if str(uuid.UUID(source_id)) != source_id \ + or uuid.UUID(source_id).version != 7: + raise ValueError("Expected a UUIDv7 Dataset ID.") + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + "Existing LeRobot Dataset metadata is invalid.") from error task_rows = source_row["tasks"] task_rows.sort(key=lambda row: row["task_index"]) task_indices = {} @@ -443,6 +460,7 @@ def _state_from_companion_tables(self): next_index = 0 for expected, row in enumerate(episode_rows): if row["episode_index"] != expected \ + or row["dataset_id"] != source_id \ or row["dataset_from_index"] != next_index \ or row["dataset_to_index"] <= next_index \ or row["length"] != ( @@ -490,7 +508,7 @@ def _state_from_companion_tables(self): raise ValueError( "Existing LeRobot stats metadata is inconsistent.") return (next_index, len(episode_rows), task_indices, numpy_stats, - subtasks) + subtasks, source_id) @staticmethod def _state_from_snapshot_properties(properties): @@ -630,6 +648,7 @@ def save_episode(self): self.num_episodes += 1 self.pending_episodes += 1 self._pending_episode_rows.append({ + "dataset_id": self._dataset_uuid, "episode_index": self.num_episodes - 1, "dataset_from_index": self.num_frames - episode.num_rows, "dataset_to_index": self.num_frames, @@ -689,7 +708,7 @@ def flush(self): {"subtask_index": index, "subtask": text} for index, text in enumerate(self.subtasks)]) dataset_row = { - "index": 0, + "id": self._dataset_uuid, "name": None, "info": json.dumps(_json_value(info), ensure_ascii=False), "tasks": task_rows, @@ -759,6 +778,7 @@ def _episode_table(self): field = self._source_schema.field(name) arrays.append(_safe_array( values, field, name, str(feature.get("dtype", "")))) + arrays.append(pa.array([self._dataset_uuid] * size, type=pa.string())) source = pa.Table.from_arrays(arrays, schema=self._source_schema) return strict_arrow_table( source, diff --git a/paimon-python/pypaimon/tests/lerobot_aggregate_test.py b/paimon-python/pypaimon/tests/lerobot_aggregate_test.py index a3bb9106835a..f62d17828acc 100644 --- a/paimon-python/pypaimon/tests/lerobot_aggregate_test.py +++ b/paimon-python/pypaimon/tests/lerobot_aggregate_test.py @@ -16,6 +16,7 @@ import json import sys +import uuid from contextlib import nullcontext from decimal import Decimal from pathlib import Path @@ -132,13 +133,49 @@ def rows(conn, name): return conn.get_table(name).scan().to_arrow() +def source_rows(conn, name): + """Project physical UUIDs to source ordinals for logical test assertions.""" + data = rows(conn, name) + prefix = name.rsplit("_", 1)[0] + datasets = rows(conn, prefix + "_datasets").sort_by("id") + ordinals = {identifier: ordinal + for ordinal, identifier in enumerate(datasets["id"].to_pylist())} + return data.set_column( + data.schema.get_field_index("dataset_id"), + pa.field("dataset_index", pa.int64()), + pa.array([ordinals[value] for value in data["dataset_id"].to_pylist()], + type=pa.int64())) + + +def assert_dataset_uuids(conn, prefix, count): + datasets = rows(conn, prefix + "_datasets").sort_by("id") + identifiers = datasets["id"].to_pylist() + assert len(identifiers) == count + assert len(set(identifiers)) == count + assert [str(uuid.UUID(value)) for value in identifiers] == identifiers + assert [uuid.UUID(value).version for value in identifiers] == [7] * count + for suffix in ("frames", "episodes"): + actual = rows(conn, prefix + "_" + suffix)["dataset_id"].to_pylist() + assert set(actual) == set(identifiers) + return datasets + + +def test_dataset_uuids_preserve_order_when_clock_does_not_advance(): + with patch.object(aggregate.time, "time_ns", return_value=1_700_000_000_000_000_000): + first = aggregate._next_dataset_uuid() + second = aggregate._next_dataset_uuid(first) + third = aggregate._next_dataset_uuid(second) + assert first < second < third + assert [uuid.UUID(value).version for value in (first, second, third)] == [7] * 3 + + def frame_rows(conn, name): - return rows(conn, name).sort_by([ + return source_rows(conn, name).sort_by([ ("dataset_index", "ascending"), ("index", "ascending")]) def episode_rows(conn, name): - return rows(conn, name).sort_by([ + return source_rows(conn, name).sort_by([ ("dataset_index", "ascending"), ("episode_index", "ascending")]) @@ -147,7 +184,7 @@ def metadata(conn, suffix): _merged_source_info, _merged_source_stats, _source_vocabulary, ) - sources = rows(conn, "merged_datasets").sort_by("index").to_pylist() + sources = rows(conn, "merged_datasets").sort_by("id").to_pylist() tasks, _ = _source_vocabulary(sources, "task") subtasks, _ = _source_vocabulary(sources, "subtask") if suffix == "info": @@ -172,6 +209,46 @@ def test_import_uses_group_prefix(setup, multiple): "frames", "episodes", "datasets"} +@pytest.mark.parametrize("source_as_list", [False, True]) +def test_single_source_accepts_string_episode_tasks(setup, source_as_list): + conn, sources = setup + episode_path = (Path(sources[0][7:]) / + "meta/episodes/chunk-000/file-000.parquet") + episodes = pq.read_table(episode_path) + names = [value[0] for value in episodes["tasks"].to_pylist()] + episodes = episodes.set_column( + episodes.schema.get_field_index("tasks"), "tasks", + pa.array(names, type=pa.string())) + pq.write_table(episodes, episode_path) + + conn.load_from_lerobot("robot", [sources[0]] if source_as_list else sources[0]) + assert rows(conn, "robot_episodes")["tasks"].to_pylist() == names + assert_dataset_uuids(conn, "robot", 1) + + +def test_single_source_preserves_feature_quantiles(setup): + conn, sources = setup + source_path = Path(sources[0][7:]) + stats_path = source_path / "meta/stats.json" + original = json.loads(stats_path.read_text()) + original["action"]["q01"] = [0.0, -4.0] + original["action"]["q99"] = [4.0, 0.0] + stats_path.write_text(json.dumps(original)) + + conn.load_from_lerobot("single", sources[0], tag_name="train") + reader = pm.PaimonLeRobotDataset( + conn.get_table("single_frames"), tag_name="train") + try: + for name in ("action", "observation.state"): + quantiles = ("q50", "q01", "q99") if name == "action" else ("q50",) + for quantile in quantiles: + np.testing.assert_allclose( + reader.meta.stats[name][quantile], original[name][quantile]) + assert "q50" not in reader.meta.stats["index"] + finally: + reader.close() + + @pytest.mark.parametrize("import_parallelism", [1, 2], indirect=True) def test_append_source_preserves_existing_indices_and_tags(setup, import_parallelism): conn, sources = setup @@ -182,15 +259,13 @@ def test_append_source_preserves_existing_indices_and_tags(setup, import_paralle assert set(conn.catalog.list_tables("default")) == { "robot_frames", "robot_episodes", "robot_datasets"} - frames = rows(conn, "robot_frames").sort_by([ - ("dataset_index", "ascending"), ("index", "ascending")]) + frames = frame_rows(conn, "robot_frames") assert frames["index"].to_pylist() == list(range(5)) * 2 assert frames["dataset_index"].to_pylist() == [0] * 5 + [1] * 5 episodes = episode_rows(conn, "robot_episodes") assert episodes["dataset_from_index"].to_pylist() == [0, 2, 0, 2] assert episodes["dataset_to_index"].to_pylist() == [2, 5, 2, 5] - datasets = rows(conn, "robot_datasets").sort_by("index") - assert datasets["index"].to_pylist() == [0, 1] + datasets = assert_dataset_uuids(conn, "robot", 2) assert datasets["name"].to_pylist() == [None, "second"] with pytest.raises(ValueError, match="already exists"): conn.append_to_lerobot("robot", sources[1], tag_name="before") @@ -347,7 +422,7 @@ def test_append_twice_keeps_source_keys_and_training_order(setup, tmp_path): frames = frame_rows(conn, "robot_frames") assert frames["dataset_index"].to_pylist() == [0] * 5 + [1] * 5 + [2] * 5 assert frames["index"].to_pylist() == list(range(5)) * 3 - assert rows(conn, "robot_datasets").sort_by("index")["index"].to_pylist() == [0, 1, 2] + assert_dataset_uuids(conn, "robot", 3) pytest.importorskip("datasets") pytest.importorskip("torch") @@ -360,34 +435,29 @@ def test_append_twice_keeps_source_keys_and_training_order(setup, tmp_path): dataset.close() -def test_single_source_feature_named_dataset_index(setup): +def test_single_source_feature_named_dataset_id(setup): pytest.importorskip("datasets") pytest.importorskip("torch") conn, sources = setup root = Path(sources[0][7:]) info_path = root / "meta/info.json" info = json.loads(info_path.read_text()) - info["features"]["dataset_index"] = { + info["features"]["dataset_id"] = { "dtype": "int64", "shape": [1], "names": None, } info_path.write_text(json.dumps(info)) data_path = root / "data/chunk-000/file-000.parquet" data = pq.read_table(data_path).append_column( - "dataset_index", pa.array([7] * 5, type=pa.int64())) + "dataset_id", pa.array([7] * 5, type=pa.int64())) pq.write_table(data, data_path) - conn.load_from_lerobot("single", sources[0]) - reader = pm.PaimonLeRobotDataset(conn.get_table("single_frames")) - try: - assert len(reader) == 5 - assert reader[0]["dataset_index"].item() == 7 - finally: - reader.close() + with pytest.raises(ValueError, match="reserved dataset_id column"): + conn.load_from_lerobot("single", sources[0]) def test_transform_frames_preserves_unlabelled_subtasks(): plan = { - "dataset_index": 1, + "dataset_id": str(uuid.uuid4()), "frame_offset": 5, "episode_offset": 2, "task_index": {0: 1}, @@ -405,7 +475,7 @@ def test_transform_frames_preserves_unlabelled_subtasks(): assert converted["index"].to_pylist() == [0, 1] assert converted["episode_index"].to_pylist() == [0, 0] - assert converted["dataset_index"].to_pylist() == [1, 1] + assert converted["dataset_id"].to_pylist() == [plan["dataset_id"]] * 2 assert "source_index" not in converted.column_names assert converted["task_index"].to_pylist() == [0, 0] assert converted["subtask_index"].to_pylist() == [0, None] @@ -454,8 +524,8 @@ def test_merge_numbering_stats_and_training(setup, import_parallelism): assert episodes["episode_index"].to_pylist() == [0, 1, 0, 1] assert "source_frame_from_index" not in episodes.column_names assert "source_frame_to_index" not in episodes.column_names - datasets = rows(conn, "merged_datasets").sort_by("index") - assert datasets.column_names == ["index", "name", "info", "tasks", "subtasks", "stats"] + datasets = rows(conn, "merged_datasets").sort_by("id") + assert datasets.column_names == ["id", "name", "info", "tasks", "subtasks", "stats"] assert datasets["name"].to_pylist() == [None, None] assert "fps" not in datasets.column_names assert [json.loads(value)["fps"] for value in datasets["info"].to_pylist()] == [10, 10] @@ -780,7 +850,7 @@ def test_single_source_metadata_opt_in(setup, import_parallelism): assert "fps_mode" not in info assert [json.loads(value)["fps"] for value in rows(conn, "single_datasets")["info"].to_pylist()] == [10] - assert rows(conn, "single_episodes")["dataset_index"].to_pylist() == [0, 0] + assert episode_rows(conn, "single_episodes")["dataset_index"].to_pylist() == [0, 0] pytest.importorskip("datasets") assert reader.fps == 10 assert reader.episode_fps == (10, 10) @@ -796,7 +866,7 @@ def test_source_info_remains_distinct_from_merged_info(setup): conn.load_from_lerobot("merged", sources, dataset_name=["move cup", "empty cup"]) assert "vendor_extension" not in metadata(conn, "info") - datasets = rows(conn, "merged_datasets").sort_by("index") + datasets = rows(conn, "merged_datasets").sort_by("id") assert datasets["name"].to_pylist() == ["move cup", "empty cup"] assert [json.loads(value)["vendor_extension"]["source"] for value in datasets["info"].to_pylist()] == [0, 1] @@ -867,8 +937,8 @@ def test_shared_business_label_does_not_merge_source_datasets(setup): conn, sources = setup conn.load_from_lerobot("merged", sources, dataset_name=["same task", "same task"]) - datasets = rows(conn, "merged_datasets").sort_by("index") - assert datasets["index"].to_pylist() == [0, 1] + datasets = rows(conn, "merged_datasets").sort_by("id") + assert_dataset_uuids(conn, "merged", 2) assert datasets["name"].to_pylist() == ["same task", "same task"] assert episode_rows(conn, "merged_episodes")["episode_index"].to_pylist() == [ 0, 1, 0, 1] @@ -969,7 +1039,7 @@ def test_mixed_fps_group_keeps_episode_rates(tmp_path, import_parallelism): info = metadata(conn, "info") assert info["fps"] is None assert "fps_mode" not in info - datasets = rows(conn, "merged_datasets").sort_by("index") + datasets = rows(conn, "merged_datasets").sort_by("id") assert "fps" not in datasets.column_names assert [json.loads(value)["fps"] for value in datasets["info"].to_pylist()] == [25, 30] @@ -1025,7 +1095,7 @@ def test_vocabulary_columns_without_pandas_index_are_preserved(setup): table = pq.read_table(path).replace_schema_metadata(None) pq.write_table(table, path) conn.load_from_lerobot("merged", sources) - dataset = rows(conn, "merged_datasets").sort_by("index").to_pylist()[0] + dataset = rows(conn, "merged_datasets").sort_by("id").to_pylist()[0] for component in ("tasks", "subtasks"): expected = pq.read_table(Path(sources[0][7:]) / ("meta/%s.parquet" % component)) actual = dataset[component] @@ -1127,13 +1197,13 @@ def test_ray_merge_multifile_sources_matches_serial(setup, ray_cluster): pq.write_table(episodes, episodes_path) conn.load_from_lerobot("serial", sources, batch_size=1) conn.load_from_lerobot("parallel", sources, batch_size=1, parallelism=2) - for suffix, key in ( - ("_frames", "index"), - ("_episodes", [("dataset_index", "ascending"), - ("episode_index", "ascending")]), - ("_datasets", "index")): - assert rows(conn, "serial" + suffix).sort_by(key).equals( - rows(conn, "parallel" + suffix).sort_by(key)) + for suffix, reader in (("_frames", frame_rows), ("_episodes", episode_rows)): + assert reader(conn, "serial" + suffix).equals(reader(conn, "parallel" + suffix)) + serial = rows(conn, "serial_datasets").sort_by("id").drop(["id"]) + parallel = rows(conn, "parallel_datasets").sort_by("id").drop(["id"]) + assert serial.equals(parallel) + assert_dataset_uuids(conn, "serial", 2) + assert_dataset_uuids(conn, "parallel", 2) def test_ray_worker_failure_does_not_commit(setup, ray_cluster): @@ -1241,26 +1311,6 @@ def read_indices(self, indices, columns): Reader(meta) -@pytest.mark.parametrize("legacy_fps", [False, True]) -def test_reader_accepts_legacy_source_info(setup, legacy_fps): - from pypaimon.multimodal.lerobot.dataset import _source_dataset_fps - - pytest.importorskip("datasets") - pytest.importorskip("torch") - conn, sources = setup - conn.load_from_lerobot("merged", sources) - reader = pm.PaimonLeRobotDataset(conn.get_table("merged_frames")) - try: - for dataset in reader.meta.datasets: - dataset["dataset_index"] = dataset.pop("index") - dataset["info_json"] = dataset.pop("info") - if legacy_fps: - dataset["fps"] = json.loads(dataset["info_json"])["fps"] - assert _source_dataset_fps(reader.meta, reader.fps, 4) == (10, 10, 10, 10) - finally: - reader.close() - - def test_reader_requires_fps_without_episode_column(setup): pytest.importorskip("datasets") pytest.importorskip("torch") @@ -1411,7 +1461,7 @@ def test_nullable_batch_type_stays_stable_after_import(setup, import_parallelism @pytest.mark.parametrize("mode", ["episode", "source"]) -def test_legacy_single_source_extensions(setup, mode): +def test_single_source_preserves_episode_extensions(setup, mode): pytest.importorskip("datasets") pytest.importorskip("torch") conn, sources = setup @@ -1427,7 +1477,7 @@ def test_legacy_single_source_extensions(setup, mode): reader = pm.PaimonLeRobotDataset(conn.get_table("single_frames"), tag_name="train") try: assert reader.episode_fps == (10, 10) - assert reader.meta.episodes["vendor_note"] == ["a", "b"] + assert episode_rows(conn, "single_episodes")["vendor_note"].to_pylist() == ["a", "b"] assert int(reader[2]["episode_index"]) == 1 finally: reader.close() @@ -1469,7 +1519,7 @@ def test_merge_archives_standard_vocabulary_columns(setup, import_parallelism): pq.write_table(table, path) originals[index, kind] = pq.read_table(path).select([kind[:-1] + "_index", kind[:-1]]) conn.load_from_lerobot("merged", sources, parallelism=import_parallelism, tag_name="train") - datasets = rows(conn, "merged_datasets").sort_by("index").to_pylist() + datasets = rows(conn, "merged_datasets").sort_by("id").to_pylist() for index, dataset in enumerate(datasets): for kind in ("tasks", "subtasks"): restored = dataset[kind] @@ -1624,10 +1674,10 @@ def test_merge_archives_vocabulary_without_arrow_schema(setup, import_parallelis pq.write_table(data.cast(schema), path) originals[index, kind] = pq.read_table(path) conn.load_from_lerobot("merged", sources, parallelism=import_parallelism) - for row in rows(conn, "merged_datasets").to_pylist(): + for ordinal, row in enumerate(rows(conn, "merged_datasets").sort_by("id").to_pylist()): for kind in ("task", "subtask"): restored = row[kind + "s"] - assert restored == originals[row["index"], kind].to_pylist() + assert restored == originals[ordinal, kind].to_pylist() reader = pm.PaimonLeRobotDataset(conn.get_table("merged_frames")) try: assert reader[0]["task"] == "pick" @@ -1760,7 +1810,7 @@ def test_merge_rejects_episode_json_group_column_collision(setup): @pytest.mark.parametrize("name,values,error", [ - ("dataset_index/vendor", pa.array([1, 2]), "reserved column dataset_index"), + ("dataset_id/vendor", pa.array([1, 2]), "reserved column dataset_id"), ("vendor_flag", pa.array([1, 2], type=pa.uint8()), "Unsupported pyarrow type: uint8"), ]) def test_merge_rejects_unwritable_episode_extensions_before_create(setup, name, values, error): @@ -1996,7 +2046,7 @@ def test_extra_dataset_text_fields(setup, multiple, import_parallelism): table_name="merged", source=sources if multiple else sources[0], dataset_name=["A", "B"] if multiple else "A", parallelism=import_parallelism, extra_dataset_fields={"config": "config/*.json", "notes": "notes.md"}, tag_name="train") - data = rows(conn, "merged_datasets").sort_by("index") + data = rows(conn, "merged_datasets").sort_by("id") assert pa.types.is_map(data.schema.field("config").type) assert data.schema.field("notes").type == pa.string() assert data["name"].to_pylist() == (["A", "B"] if multiple else ["A"]) @@ -2124,10 +2174,9 @@ def test_schema_evolution_unions_three_tables(setup, import_parallelism, reverse sources.reverse() conn.load_from_lerobot("merged", sources, parallelism=import_parallelism, extra_dataset_fields={"config": "config/*.json", "notes": "notes.md"}) - data = rows(conn, "merged_frames").sort_by([ - ("dataset_index", "ascending"), ("index", "ascending")]) + data = frame_rows(conn, "merged_frames") episodes = episode_rows(conn, "merged_episodes").to_pylist() - datasets = rows(conn, "merged_datasets").sort_by("index").to_pylist() + datasets = rows(conn, "merged_datasets").sort_by("id").to_pylist() for ordinal, uri in enumerate(sources): first = uri == original[0] assert data["temperature"].slice(ordinal * 5, 5).to_pylist() == ( @@ -2301,10 +2350,10 @@ def test_append_evolves_frame_and_dataset_columns(setup, tmp_path, import_parall current = conn.get_table("merged_frames").raw_table assert {f.name: f.id for f in current.fields if f.name in field_ids} == field_ids assert next(f.type.nullable for f in current.fields if f.name == "humidity") - datasets = rows(conn, "merged_datasets").sort_by("index") + datasets = rows(conn, "merged_datasets").sort_by("id") assert datasets["notes"].to_pylist() == [None, None, "# Third dataset"] assert datasets["config"].to_pylist() == [None, None, None] - assert datasets["index"].to_pylist() == [0, 1, 2] + assert_dataset_uuids(conn, "merged", 3) assert metadata(conn, "stats")["humidity"]["count"] == [5] assert metadata(conn, "stats")["temperature"]["count"] == [5] assert set(conn.catalog.list_tables("default")) == { diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_ray_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_ray_test.py index 7ae5921ff60a..902de822803e 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_ray_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_ray_test.py @@ -40,7 +40,7 @@ def _counted_prepare_group(trace_path, *args): from pypaimon.multimodal.lerobot import source as source_module - from pypaimon.multimodal.lerobot.ray_loader import _prepare_group + from pypaimon.multimodal.lerobot.ray_loader import _prepare_merged_source original = source_module._read_remote_parquet @@ -51,7 +51,7 @@ def counted(file_io, path, *read_args, **read_kwargs): source_module._read_remote_parquet = counted try: - return _prepare_group(*args) + return _prepare_merged_source(*args) finally: source_module._read_remote_parquet = original @@ -227,7 +227,7 @@ def counted_prepare(*args): return _counted_prepare_group(str(trace_path), *args) with patch( - "pypaimon.multimodal.lerobot.ray_loader._prepare_group", + "pypaimon.multimodal.lerobot.ray_loader._prepare_merged_source", counted_prepare): for data_files in (1, 4): with self.subTest(data_files=data_files): @@ -252,7 +252,7 @@ def counted_prepare(*args): def _import_serial(self): # FileIO sources use the native Parquet reader, not the LeRobot class. - with patch("pypaimon.multimodal.lerobot.api._import_lerobot_dataset", return_value=None): + with patch("pypaimon.multimodal.lerobot.aggregate._import_lerobot_dataset", return_value=None): self.conn.load_from_lerobot("serial", self.source.as_uri(), batch_size=1, parallelism=1) @@ -289,14 +289,19 @@ def read(name, tag=None): self.assertEqual( Blob.from_descriptor_bytes(original.as_py(), file_io=file_io).to_data(), Blob.from_descriptor_bytes(imported.as_py(), file_io=file_io).to_data()) - left, right = left.drop(["camera"]), right.drop(["camera"]) + left, right = left.drop(["camera", "dataset_id"]), right.drop( + ["camera", "dataset_id"]) else: left = left.sort_by(left.column_names[0]) right = right.sort_by(right.column_names[0]) + field = "id" if suffix == "_datasets" else "dataset_id" + left, right = left.drop([field]), right.drop([field]) self.assertTrue(left.equals(right), suffix) - # One data snapshot plus the initial BTree snapshot, regardless of task count. + # Parallel import publishes the same number of snapshots as serial import. table = self.conn.get_table("parallel_frames").raw_table - self.assertEqual(2, table.snapshot_manager().get_latest_snapshot().id) + serial = self.conn.get_table("serial_frames").raw_table + self.assertEqual(serial.snapshot_manager().get_latest_snapshot().id, + table.snapshot_manager().get_latest_snapshot().id) @unittest.skipUnless( importlib.util.find_spec("av") @@ -337,7 +342,8 @@ def __getattr__(self, name): def get(self, ref, **kwargs): result = ray.get(ref, **kwargs) - if isinstance(result, list) and result and hasattr(result[0], "new_files"): + prepared = result[0] if isinstance(result, tuple) else result + if prepared and hasattr(prepared[0], "new_files"): table = connection.get_table("broken_frames").raw_table files_before_failure.extend(_bucket_files(table)) received.touch() @@ -350,11 +356,12 @@ def prepare(*args): if time.monotonic() > deadline: raise TimeoutError("Successful task result was not received.") time.sleep(0.01) - return _prepare_group(*args) + from pypaimon.multimodal.lerobot.ray_loader import _prepare_merged_source + return _prepare_merged_source(*args) with patch("pypaimon.multimodal.lerobot.ray_loader._require_ray", return_value=MarkingRay()), \ - patch("pypaimon.multimodal.lerobot.ray_loader._prepare_group", + patch("pypaimon.multimodal.lerobot.ray_loader._prepare_merged_source", prepare), \ patch("pypaimon.multimodal.lerobot.ray_loader._abort_messages", wraps=_abort_messages) as abort: @@ -378,7 +385,7 @@ def test_writer_close_failure_aborts_prepared_files(self): trace = self.root / (table_name + ".json") def prepare(*args): - from pypaimon.multimodal.lerobot.ray_loader import _prepare_group + from pypaimon.multimodal.lerobot.ray_loader import _prepare_merged_source from pypaimon.write.table_write import BatchTableWrite original_close = BatchTableWrite.close @@ -391,11 +398,11 @@ def fail_close(write): BatchTableWrite.close = fail_close try: - return _prepare_group(*args) + return _prepare_merged_source(*args) finally: BatchTableWrite.close = original_close - with patch("pypaimon.multimodal.lerobot.ray_loader._prepare_group", + with patch("pypaimon.multimodal.lerobot.ray_loader._prepare_merged_source", prepare): with self.assertRaisesRegex(BaseException, "injected writer close"): self.conn.load_from_lerobot( @@ -451,9 +458,9 @@ def interrupted_source(*source_args): raise KeyboardInterrupt("interrupted source close") with patch.object(ray_loader, "_resolved_source", interrupted_source): - return ray_loader._prepare_group(*args) + return ray_loader._prepare_merged_source(*args) - with patch("pypaimon.multimodal.lerobot.ray_loader._prepare_group", prepare): + with patch("pypaimon.multimodal.lerobot.ray_loader._prepare_merged_source", prepare): with self.assertRaisesRegex(BaseException, "interrupted source close"): self.conn.load_from_lerobot("source_close", self.source, parallelism=2) self.assertTrue(json.loads(trace.read_text())) @@ -466,10 +473,10 @@ def test_invalid_commit_options_fail_before_preparing(self): for parallelism in (1, 2): with self.subTest(parallelism=parallelism): name = "invalid_commit_%d" % parallelism - with patch("pypaimon.multimodal.lerobot.api._import_lerobot_dataset", + with patch("pypaimon.multimodal.lerobot.aggregate._import_lerobot_dataset", return_value=None), \ patch("pypaimon.multimodal.lerobot.loader._prepare_dataset") as prepare, \ - patch("pypaimon.multimodal.lerobot.ray_loader._prepare_group") as group: + patch("pypaimon.multimodal.lerobot.ray_loader._prepare_merged_source") as group: with self.assertRaises(ValueError): self.conn.load_from_lerobot( name, self.source.as_uri(), parallelism=parallelism, @@ -523,7 +530,7 @@ def prepare(*args): try: with patch("pypaimon.multimodal.lerobot.ray_loader._require_ray", return_value=TrackingRay()), \ - patch("pypaimon.multimodal.lerobot.ray_loader._prepare_group", + patch("pypaimon.multimodal.lerobot.ray_loader._prepare_merged_source", prepare), \ patch("pypaimon.multimodal.lerobot.ray_loader._abort_messages", wraps=_abort_messages) as abort: diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 99b8ab5f9402..c36ace73b6ea 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1638,7 +1638,7 @@ def test_invalid_training_tag_fails_before_catalog_access(self): def test_import_validates_tag_before_source_access(self): connection = Mock() - with patch("pypaimon.multimodal.lerobot.api._resolved_source", + with patch("pypaimon.multimodal.lerobot.aggregate._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): @@ -1651,7 +1651,7 @@ def test_import_validates_tag_before_source_access(self): load_from_lerobot(connection, "robot", "source", tag_name=None) resolve.assert_called_once() - @patch("pypaimon.multimodal.lerobot.api._import_lerobot_dataset", + @patch("pypaimon.multimodal.lerobot.aggregate._import_lerobot_dataset", return_value=Mock()) def test_training_tag_uses_current_component_snapshots(self, _): import pandas as pd @@ -1680,7 +1680,9 @@ def test_training_tag_uses_current_component_snapshots(self, _): "custom": {"labels": ["pick", None], "enabled": True}, } (source / "meta" / "info.json").write_text(json.dumps(info)) - stats = {"timestamp": {"min": [0.0], "max": [0.0]}} + stats = {"timestamp": { + "min": [0.0], "max": [0.0], "mean": [0.0], + "std": [0.0], "count": [1]}} (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"), @@ -1715,15 +1717,20 @@ def test_training_tag_uses_current_component_snapshots(self, _): self.assertEqual("-1", schema.options.get("bucket", "-1")) self.assertFalse(schema.primary_keys) self.assertGreater(component.scan().to_arrow().num_rows, 0) + dataset_row = _catalog_rows(connection, "robot_datasets")[0] self.assertEqual([{"subtask_index": 0, "subtask": "reach"}], - _catalog_rows(connection, "robot_subtasks")) - self.assertEqual(info, _catalog_metadata(connection, "robot_info")) - self.assertEqual(stats, _catalog_metadata(connection, "robot_stats")) + dataset_row["subtasks"]) + self.assertEqual(info, json.loads(dataset_row["info"])) + self.assertEqual(stats, json.loads(dataset_row["stats"])) table = connection.get_table("robot_frames") + frames = frames.append_column( + "dataset_id", pa.array([dataset_row["id"]])) self.assertEqual([], table.raw_table.tag_manager().list_tags()) table.add(frames) snapshots = connection.create_lerobot_tag("robot_frames", "training") - self.assertEqual({"frames": 3, "episodes": 1, "datasets": 1}, snapshots) + self.assertEqual({"frames": table.raw_table.snapshot_manager() + .get_latest_snapshot().id, "episodes": 1, "datasets": 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) @@ -1751,7 +1758,7 @@ def fail_second_component(*args, **kwargs): with self.assertRaisesRegex(RuntimeError, "tag failed"): connection.create_lerobot_tag("robot_frames", "retry") self.assertFalse(table.raw_table.tag_manager().tag_exists("retry")) - self.assertEqual(4, connection.create_lerobot_tag( + self.assertEqual(snapshots["frames"] + 1, connection.create_lerobot_tag( "robot_frames", "retry")["frames"]) self.assertEqual(3, table.scan(tag_name="retry").to_arrow().num_rows) @@ -1768,8 +1775,10 @@ def fail_second_component(*args, **kwargs): connection.load_from_lerobot("no_stats", remote, tag_name="ready") with self.assertRaises(TableNotExistException): connection.get_table("no_stats_stats") - self.assertEqual({"frames": 2, "episodes": 1, "datasets": 1}, - connection.create_lerobot_tag("no_stats_frames", "training")) + no_stats = connection.get_table("no_stats_frames") + self.assertEqual( + no_stats.raw_table.snapshot_manager().get_latest_snapshot().id, + connection.create_lerobot_tag("no_stats_frames", "training")["frames"]) def test_self_contained_import_rejects_table_branches(self): with self.assertRaisesRegex(ValueError, "does not support"): @@ -2802,7 +2811,7 @@ def test_empty_local_dataset_is_rejected_before_opening_lerobot(self): "warehouse": str(temp_dir / "warehouse"), }) with patch( - "pypaimon.multimodal.lerobot.api._import_lerobot_dataset" + "pypaimon.multimodal.lerobot.aggregate._import_lerobot_dataset" ) as import_lerobot: with self.assertRaisesRegex(ValueError, "non-empty"): connection.load_from_lerobot("empty_frames", source) @@ -2913,7 +2922,7 @@ def test_empty_fast_path_validates_required_counts(self): (source / "meta" / "info.json").write_text( json.dumps(info)) with patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_import_lerobot_dataset" ) as import_lerobot: with self.assertRaisesRegex(ValueError, message): @@ -3326,11 +3335,11 @@ def sample_timestamps(unused_dataset, uri): ) with patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_import_lerobot_dataset", return_value=object, ), patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_open_resolved_dataset", return_value=Dataset(), ), patch( @@ -3396,7 +3405,7 @@ def sample_timestamps(unused_dataset, uri): "pypaimon.multimodal.lerobot.source._SourceFileIO", return_value=source_file_io, ), patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_import_lerobot_dataset", return_value=object, ), patch( @@ -3435,11 +3444,11 @@ def sample_timestamps(unused_dataset, uri): }) pq.write_table(pa.Table.from_pylist(episodes), episodes_path) with patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_import_lerobot_dataset", return_value=object, ), patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_open_resolved_dataset", return_value=Dataset(), ), patch( @@ -3686,11 +3695,11 @@ def open_file(path): return _prepare_dataset(table, *args, **kwargs) with patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_import_lerobot_dataset", return_value=object, ), patch( - "pypaimon.multimodal.lerobot.api." + "pypaimon.multimodal.lerobot.aggregate." "_open_resolved_dataset", return_value=Dataset(), ), patch( @@ -4285,7 +4294,7 @@ def test_table_dataset_pickle_preserves_episode_metadata_and_reads(self): def test_training_reads_previously_imported_vector_features(self): import torch - with patch("pypaimon.multimodal.lerobot.api._schema_from_info", + with patch("pypaimon.multimodal.lerobot.aggregate._schema_from_info", side_effect=lambda info: _schema_from_info(info, legacy_vectors=True)): self.connection.load_from_lerobot( "legacy_vectors", self.image_source, batch_size=2) @@ -4320,7 +4329,7 @@ def test_import_infers_schema_and_preserves_episodes(self): self.assertEqual("FLOAT NOT NULL", types["timestamp"]) self.assertEqual("BIGINT NOT NULL", types["episode_index"]) self.assertEqual("BLOB NOT NULL", types["observation.image"]) - self.assertNotIn("dataset_id", types) + self.assertEqual("STRING NOT NULL", types["dataset_id"]) self.assertNotIn("task", types) rows = table.scan().select([ @@ -4354,6 +4363,9 @@ def test_import_infers_schema_and_preserves_episodes(self): ) dataset_rows = _catalog_rows(self.connection, "robot_data_datasets") self.assertEqual(1, len(dataset_rows)) + dataset_id = dataset_rows[0]["id"] + frame_ids = table.scan().select(["dataset_id"]).to_arrow().to_pylist() + self.assertEqual({dataset_id}, {row["dataset_id"] for row in frame_ids}) self.assertEqual( json.loads((self.image_source / "meta" / "info.json").read_text()), json.loads(dataset_rows[0]["info"]), @@ -4366,12 +4378,10 @@ def test_import_infers_schema_and_preserves_episodes(self): self.connection._identifier( "robot_data_episodes")).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.assertTrue({"dataset_id", "episode_index", "dataset_from_index", + "dataset_to_index", "length", "tasks", "meta", "data", + "videos", "stats"}.issubset(episode_fields)) + self.assertEqual({dataset_id}, {row["dataset_id"] for row in episodes}) self.assertEqual([(0, 0, 2), (1, 2, 5)], [ (row["episode_index"], row["dataset_from_index"], row["dataset_to_index"]) @@ -4379,7 +4389,7 @@ def test_import_infers_schema_and_preserves_episodes(self): ]) self.assertEqual([["pick"], ["place"]], [ row["tasks"] for row in episodes]) - self.assertTrue(any( + self.assertFalse(any( name.startswith("stats/") for name in episode_fields)) tasks = _catalog_rows(self.connection, "robot_data_tasks") @@ -4388,7 +4398,7 @@ def test_import_infers_schema_and_preserves_episodes(self): ]) self.assertEqual(tasks, dataset_rows[0]["tasks"]) self.assertEqual( - 2, + 3, table.raw_table.snapshot_manager().get_latest_snapshot().id, ) @@ -4624,16 +4634,16 @@ def test_import_preserves_quoted_database_name(self): "robot_frames", ], sorted(table_names)) - def test_import_reuses_validated_episode_metadata(self): - from pypaimon.multimodal.lerobot import api + def test_import_rejects_episode_metadata_changed_after_validation(self): + from pypaimon.multimodal.lerobot import aggregate 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 + original_prepare = aggregate._prepare_dataset - def write_then_replace(*args, **kwargs): - snapshot_id = original_write(*args, **kwargs) + def prepare_then_replace(*args, **kwargs): + messages = original_prepare(*args, **kwargs) episodes = pq.read_table(episode_path) tasks = episodes.column("tasks").to_pylist() tasks[0] = ["place"] @@ -4642,15 +4652,14 @@ def write_then_replace(*args, **kwargs): "tasks", pa.array(tasks, type=episodes.schema.field("tasks").type), ), episode_path) - return snapshot_id + return messages 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"]) + aggregate, "_prepare_dataset", side_effect=prepare_then_replace): + with self.assertRaisesRegex(ValueError, "changed after validation"): + self.connection.load_from_lerobot("stable_episodes", source) + self.assertEqual([], _catalog_rows( + self.connection, "stable_episodes_episodes")) def test_frame_controls_must_match_published_episode_metadata(self): cases = [ @@ -4797,9 +4806,9 @@ def test_paimon_dataset_reads_lazy_batches_with_lerobot_metadata(self): frame_indexes = _catalog_rows( self.connection, "training_data_frames$table_indexes") self.assertEqual( - {"btree", "index"}, - {frame_indexes[0]["index_type"], - frame_indexes[0]["index_field_name"]}, + {("btree", "index"), ("btree", "dataset_id")}, + {(row["index_type"], row["index_field_name"]) + for row in frame_indexes}, ) from pypaimon.multimodal.blob_read import fetch_blob_bodies @@ -4901,6 +4910,8 @@ def counted_plan(scan): self.assertEqual(1, int(reordered[-1]["episode_index"])) table.add(pa.Table.from_pylist([{ + "dataset_id": _catalog_rows( + self.connection, "training_data_datasets")[0]["id"], "index": 999, "episode_index": 99, "frame_index": 0, @@ -5106,7 +5117,7 @@ def create_then_lose_response(*args, **kwargs): self.assertTrue(lost[0]) self.assertIsNone(result) - self.assertEqual(2, self.connection.catalog.get_tag( + self.assertEqual(3, self.connection.catalog.get_tag( self.connection._identifier("tag_response_loss_frames"), "training").snapshot.id) @@ -5171,9 +5182,9 @@ def fail_once(name): self.connection.get_table("failed_open_frames") def test_dataset_close_failure_does_not_override_success(self): - from pypaimon.multimodal.lerobot import api + from pypaimon.multimodal.lerobot import aggregate - original_open = api._open_resolved_dataset + original_open = aggregate._open_resolved_dataset def open_with_failing_close(*args, **kwargs): return _FailingCloseDataset(original_open(*args, **kwargs)) @@ -5181,7 +5192,7 @@ def open_with_failing_close(*args, **kwargs): with self.assertLogs( "pypaimon.multimodal.lerobot.source", level="WARNING"): with patch.object( - api, + aggregate, "_open_resolved_dataset", side_effect=open_with_failing_close): result = self.connection.load_from_lerobot( @@ -5220,9 +5231,9 @@ def test_existing_target_is_rejected(self): table.raw_table.snapshot_manager().get_latest_snapshot()) def test_concurrent_import_cannot_claim_the_same_target(self): - from pypaimon.multimodal.lerobot import api + from pypaimon.multimodal.lerobot import aggregate - original_prepare = api._prepare_metadata_tables + original_prepare = aggregate._prepare_metadata_tables root_created = threading.Event() release = threading.Event() @@ -5232,7 +5243,7 @@ def prepare_then_wait(*args, **kwargs): return original_prepare(*args, **kwargs) with patch.object( - api, + aggregate, "_prepare_metadata_tables", side_effect=prepare_then_wait): with ThreadPoolExecutor(max_workers=1) as executor: @@ -5259,11 +5270,11 @@ def prepare_then_wait(*args, **kwargs): ) def test_concurrent_append_rejects_initial_import(self): - from pypaimon.multimodal.lerobot import api + from pypaimon.multimodal.lerobot import aggregate - original_write = api._write_dataset + original_prepare = aggregate._prepare_dataset - def append_then_write( + def append_then_prepare( table, dataset, info, @@ -5271,27 +5282,24 @@ def append_then_write( source_schema, batch_size, metadata, - video_fields=()): - table.add(_read_batch( + video_fields=(), + **kwargs): + messages = original_prepare( + table, dataset, info, source, source_schema, + batch_size, metadata, video_fields, **kwargs) + batch = _read_batch( dataset, info, 0, 1, source_schema, - )) - return original_write( - table, - dataset, - info, - source, - source_schema, - batch_size, - metadata, - video_fields, ) + self.connection.get_table("concurrent_append_frames").add( + kwargs["transform"](batch, 0)) + return messages with patch.object( - api, "_write_dataset", side_effect=append_then_write): + aggregate, "_prepare_dataset", side_effect=append_then_prepare): with self.assertRaisesRegex(RuntimeError, "concurrent writes"): self.connection.load_from_lerobot( "concurrent_append", self.image_source) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py index 25166b4ce2b8..6815dfbe8e16 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -21,6 +21,7 @@ import shutil import tempfile import unittest +import uuid from pathlib import Path from unittest.mock import patch @@ -119,6 +120,13 @@ def test_creates_lerobot_table_group_on_finalize(self): "task_index": 0, "task": "pick", }], _dataset_row(self.connection, "table_group")["tasks"]) + dataset_id = _dataset_row(self.connection, "table_group")["id"] + self.assertEqual(7, uuid.UUID(dataset_id).version) + self.assertEqual([dataset_id], self.connection.get_table( + "table_group").scan().select(["dataset_id"]).to_arrow()[ + "dataset_id"].to_pylist()) + self.assertEqual([dataset_id], [row["dataset_id"] for row in + _catalog_rows(self.connection, "table_group_episodes")]) info = json.loads(_dataset_row(self.connection, "table_group")["info"]) self.assertEqual(10, info["fps"]) self.assertEqual(1, info["total_frames"]) @@ -312,6 +320,8 @@ def test_existing_table_resumes_global_and_task_indices(self): }) first.save_episode() first.finalize() + dataset_id = _dataset_row(self.connection, "resume")["id"] + self.assertEqual(7, uuid.UUID(dataset_id).version) manager = self.connection.get_table( "resume").raw_table.snapshot_manager() @@ -336,6 +346,10 @@ def test_existing_table_resumes_global_and_task_indices(self): }) resumed.save_episode() resumed.finalize() + self.assertEqual(dataset_id, _dataset_row(self.connection, "resume")["id"]) + self.assertEqual([dataset_id] * 4, self.connection.get_table( + "resume").scan().select(["dataset_id"]).to_arrow()[ + "dataset_id"].to_pylist()) rows = self.connection.get_table("resume").scan().select([ "episode_index", "frame_index", "index", "task_index" @@ -387,7 +401,7 @@ def test_empty_frames_rejects_nonempty_companion_state(self): datasets = self.connection.catalog.get_table( self.connection._identifier("inconsistent_datasets")) _append_arrow(datasets, pa.Table.from_pylist([{ - "index": 0, "name": None, "info": "{}", "tasks": [], + "id": str(uuid.uuid4()), "name": None, "info": "{}", "tasks": [], "subtasks": None, "stats": "{}", }], schema=_target_schema(datasets))) @@ -454,6 +468,19 @@ def test_video_features_are_rejected(self): }, ) + def test_dataset_id_feature_is_reserved(self): + with self.assertRaisesRegex(ValueError, "dataset_id is managed"): + PaimonLeRobotWriter( + self.connection, + "reserved_id", + fps=10, + features={ + "dataset_id": { + "dtype": "string", "shape": (1,), "names": None, + }, + }, + ) + def test_missing_lerobot_stats_dependency_fails_before_table_creation(self): self.connection.catalog.create_database( "default", ignore_if_exists=True) @@ -557,16 +584,23 @@ def test_multiple_flushes_append_rows_and_replace_info(self): }, }, ) + first_id = None for value, task in ((1.0, "pick"), (2.0, "place")): writer.add_frame({ "action": np.array([value], dtype=np.float32), "task": task, }) writer.save_episode() + current_id = _dataset_row(self.connection, "multiple_flushes")["id"] + if first_id is None: + first_id = current_id + else: + self.assertEqual(first_id, current_id) writer.finalize() self.assertEqual(2, len(_catalog_rows( self.connection, "multiple_flushes_episodes"))) + self.assertEqual(7, uuid.UUID(first_id).version) self.assertEqual(2, len(_dataset_row( self.connection, "multiple_flushes")["tasks"])) info = json.loads(_dataset_row(self.connection, "multiple_flushes")["info"])