JingsongLi commented on code in PR #9529: URL: https://github.com/apache/paimon/pull/9529#discussion_r3920645057
########## paimon-python/pypaimon/multimodal/lerobot/metadata.py: ########## @@ -0,0 +1,656 @@ +# 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.""" + +import json +import numbers +import uuid +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq + +from pypaimon import Schema as PaimonSchema +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableAlreadyExistException, + TableNotExistException, +) +from pypaimon.common.identifier import Identifier +from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.table import _target_schema + + +_VERSION_ID = "version_id" +_OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" +_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("status", pa.string(), nullable=False), + pa.field("info_json", pa.string(), nullable=False), + pa.field("stats_json", pa.string()), + pa.field("has_subtasks", pa.bool_(), nullable=False), +]) +_EMPTY_TASKS_SCHEMA = pa.schema([ + pa.field("task_index", pa.int64(), nullable=False), + 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", +] + + +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.to_pylist(), int(info["total_tasks"])) + subtasks_table = _source_subtasks(dataset, source) + subtask_indices = _subtask_indices(subtasks_table, info) + total_episodes = int(info["total_episodes"]) + episode_source = ( + _source_episodes(dataset, source) + if total_episodes > 0 + else {"paths": [], "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 _new_owner_id(): + return uuid.uuid4().hex + + +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, owner_id): + identifier = Identifier.from_string(str(frames_identifier)) + if identifier.get_branch_name() is not None: + raise ValueError( + "LeRobot import does not support table branches.") + result = {_OWNER_ID_OPTION: owner_id} + for name, suffix in _TABLE_SUFFIXES.items(): + result[_COMPANION_OPTION_KEYS[name]] = _companion_identifier( + frames_identifier, suffix) + return result + + +def _is_managed_root(options): + return _OWNER_ID_OPTION in options and all( + key in options for key in _COMPANION_OPTION_KEYS.values()) + + +def _companion_table_identifiers(frames_table): + options = frames_table.table_schema.options + identifiers = {} + for name, key in _COMPANION_OPTION_KEYS.items(): + value = options.get(key) + if not value: + raise ValueError( + "LeRobot table %s is missing managed option %s." + % (frames_table.identifier, key)) + identifiers[name] = value + return identifiers + + +def _prepare_metadata_tables(connection, frames_table, owner_id, 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] + try: + table = connection.catalog.get_table(identifier) + except (DatabaseNotExistException, TableNotExistException): + paimon_schema = PaimonSchema.from_pyarrow_schema( + schema, + options={ + "bucket": "-1", + _OWNER_ID_OPTION: owner_id, + }, + ) + try: + connection.catalog.create_table( + identifier, paimon_schema, False) + except TableAlreadyExistException: + pass + table = connection.catalog.get_table(identifier) + if table.table_schema.primary_keys: + raise ValueError( + "LeRobot metadata table %s must be append-only." % identifier) + actual = _target_schema(table) + if not actual.equals(schema, check_metadata=False): Review Comment: Ignoring schema metadata here hides a round-trip incompatibility. LeRobot writes tasks.parquet and subtasks.parquet from Pandas DataFrames, with the task or subtask text encoded as a named Pandas index. Schema.from_pyarrow_schema does not persist the top-level pandas schema metadata, so scanning or exporting this Paimon table and calling to_pandas() produces a RangeIndex instead of restoring that text index; LeRobot lookups such as tasks.iloc[task_idx].name then return the wrong value. Please either persist and reattach the component Arrow schema metadata, or explicitly reconstruct the index from the physical text column in the Paimon LeRobot reader and cover that round trip in a test. ########## paimon-python/pypaimon/multimodal/lerobot/schema.py: ########## @@ -47,22 +47,15 @@ def _require_v3(info, source): % (source, version or None)) -def _schema_from_info(info, include_task): +def _schema_from_info(info): features = info.get("features") if not isinstance(features, dict) or not features: raise ValueError("LeRobot metadata features must be a non-empty object.") - fields = [] - for name, feature in features.items(): - fields.append(_feature_field(name, feature)) - if include_task: - fields.append(pa.field( - "task", - pa.string(), - nullable=False, - metadata={b"description": b"LeRobot task"}, - )) - return pa.schema(fields) + return pa.schema([ + _feature_field(name, feature) Review Comment: This derives the frame schema from arbitrary info.json descriptors but never enforces the mandatory LeRobot V3 control schema. LeRobot V3 defines timestamp as scalar float32 and frame_index, episode_index, index, and task_index as scalar int64; subtask_index should likewise be scalar int64 when present. Today a source declaring timestamp as float64 or frame_index as int32 is accepted and published as a non-native frame schema because the row validator only checks values. Please validate these required feature descriptors against the V3 defaults before creating the table. ########## paimon-python/pypaimon/multimodal/lerobot/metadata.py: ########## @@ -0,0 +1,656 @@ +# 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.""" + +import json +import numbers +import uuid +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq + +from pypaimon import Schema as PaimonSchema +from pypaimon.catalog.catalog_exception import ( + DatabaseNotExistException, + TableAlreadyExistException, + TableNotExistException, +) +from pypaimon.common.identifier import Identifier +from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.table import _target_schema + + +_VERSION_ID = "version_id" +_OWNER_ID_OPTION = "pypaimon.lerobot.owner-id" +_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("status", pa.string(), nullable=False), + pa.field("info_json", pa.string(), nullable=False), + pa.field("stats_json", pa.string()), + pa.field("has_subtasks", pa.bool_(), nullable=False), +]) +_EMPTY_TASKS_SCHEMA = pa.schema([ + pa.field("task_index", pa.int64(), nullable=False), + 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", +] + + +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.to_pylist(), int(info["total_tasks"])) + subtasks_table = _source_subtasks(dataset, source) + subtask_indices = _subtask_indices(subtasks_table, info) + total_episodes = int(info["total_episodes"]) + episode_source = ( + _source_episodes(dataset, source) + if total_episodes > 0 + else {"paths": [], "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 _new_owner_id(): + return uuid.uuid4().hex + + +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, owner_id): + identifier = Identifier.from_string(str(frames_identifier)) + if identifier.get_branch_name() is not None: + raise ValueError( + "LeRobot import does not support table branches.") + result = {_OWNER_ID_OPTION: owner_id} + for name, suffix in _TABLE_SUFFIXES.items(): + result[_COMPANION_OPTION_KEYS[name]] = _companion_identifier( + frames_identifier, suffix) + return result + + +def _is_managed_root(options): + return _OWNER_ID_OPTION in options and all( + key in options for key in _COMPANION_OPTION_KEYS.values()) + + +def _companion_table_identifiers(frames_table): + options = frames_table.table_schema.options + identifiers = {} + for name, key in _COMPANION_OPTION_KEYS.items(): + value = options.get(key) + if not value: + raise ValueError( + "LeRobot table %s is missing managed option %s." + % (frames_table.identifier, key)) + identifiers[name] = value + return identifiers + + +def _prepare_metadata_tables(connection, frames_table, owner_id, 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] + try: + table = connection.catalog.get_table(identifier) + except (DatabaseNotExistException, TableNotExistException): + paimon_schema = PaimonSchema.from_pyarrow_schema( + schema, + options={ + "bucket": "-1", + _OWNER_ID_OPTION: owner_id, + }, + ) + try: + connection.catalog.create_table( + identifier, paimon_schema, False) + except TableAlreadyExistException: + pass + table = connection.catalog.get_table(identifier) + if table.table_schema.primary_keys: + raise ValueError( + "LeRobot metadata table %s must be append-only." % identifier) + actual = _target_schema(table) + if not actual.equals(schema, check_metadata=False): + raise ValueError( + "LeRobot metadata table %s has schema %s; expected %s." + % (identifier, actual, schema)) + actual_owner_id = table.table_schema.options.get(_OWNER_ID_OPTION) + if actual_owner_id != owner_id: + raise ValueError( + "LeRobot metadata table %s belongs to a different target " + "table. Drop the stale companion tables before importing." + % identifier) + tables[name] = table + return tables + + +def _reserve_dataset_version( + versions_table, + version_id, + metadata): + pending = _manifest_row(version_id, "PENDING", metadata) + snapshot_id = _append_arrow( + versions_table, + pa.Table.from_pylist([pending], schema=_VERSIONS_SCHEMA), + ) + _require_initial_snapshot("versions", snapshot_id) + + +def _publish_dataset( + connection, + tables, + 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, "READY", 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, + status, + metadata): + return { + _VERSION_ID: version_id, + "status": status, + "info_json": metadata["info_json"], + "stats_json": metadata["stats_json"], + "has_subtasks": metadata["subtasks_table"] is not None, + } + + +def _drop_import_tables(catalog, frames_table, owner_id): + identifiers = list( + _companion_table_identifiers(frames_table).values()) + identifiers.append(frames_table.identifier) + for identifier in identifiers: + try: + table = catalog.get_table(identifier) + except (DatabaseNotExistException, TableNotExistException): + continue + if table.table_schema.options.get(_OWNER_ID_OPTION) == owner_id: + catalog.drop_table(identifier, ignore_if_not_exists=True) + + +def _append_arrow(table, data): + return _append_arrow_tables(table, [data]) + + +def _append_arrow_tables(table, tables, flush_each=False): + builder = ( + table.new_stream_write_builder() + if flush_each else table.new_batch_write_builder() + ) + 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 + messages = [] + target_schema = _target_schema(table) + for data in tables: + if data.num_rows == 0: + continue + if not data.schema.equals(target_schema, check_metadata=False): + raise ValueError( + "LeRobot component schema %s does not match target %s." + % (data.schema, target_schema)) + table_write.write_arrow(data) + row_count += data.num_rows + if flush_each: + messages = table_write.prepare_commit(0) + del data + if row_count == 0: + table_write.abort() + return None + if not flush_each: + messages = table_write.prepare_commit() + commit_started = True + if flush_each: + table_commit.commit(messages, 0) + else: + table_commit.commit(messages) + if recorder.snapshot_id is None: + raise RuntimeError("LeRobot metadata commit has no snapshot id.") + return recorder.snapshot_id + 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: + catalog.create_tag( + identifier, tag_name, snapshot_id=snapshot_id) + except NotImplementedError: + catalog.get_table(identifier).create_tag( + tag_name, snapshot_id=snapshot_id) + + +def _source_stats(dataset, source): + if source.file_io is not None: + from pypaimon.multimodal.lerobot.source import ( + _read_remote_json, + _remote_path, + ) + path = _remote_path(source.path, "meta/stats.json") + try: + source.file_io.get_file_status(path) + except FileNotFoundError: + return None + return _read_remote_json(source.file_io, path) + root = _metadata_root(dataset, source) + path = root / "meta" / "stats.json" + if not path.is_file(): + return None + with path.open("r", encoding="utf-8") as file: + return json.load(file) + + +def _source_tasks(dataset, source, total_tasks): + if total_tasks == 0: + return 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): + rows = [] + for table in _source_episode_tables(metadata): + rows.extend(table.select(_EPISODE_CONTROL_COLUMNS).to_pylist()) Review Comment: Although each Arrow shard is released before the next one is read, this still materializes every Episode control row as Python dictionaries in rows, and _episode_rows then builds a second full Python list retained in metadata. The FileIO path has already materialized another locator list in _RemoteLeRobotDataset, so peak heap remains O(total_episodes) and can reach multiple GB at the million-Episode scale that V3 targets. Please validate ordered Episode rows incrementally and keep the boundaries in Arrow, mmap, or compact arrays, or consume them directly during frame import, instead of retaining duplicate to_pylist() structures. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
