From 9683e98704b7867a45a2666d73626b486eecc32e Mon Sep 17 00:00:00 2001 From: Yann Date: Fri, 28 Aug 2026 14:09:17 +0800 Subject: [PATCH 1/6] feat(python): add RoboMIND AgileX pipeline Reuse the HDF5 transform contract for local and Ray ingestion, then materialize canonical actions and independently refresh train statistics. Co-Authored-By: Codex Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 1112/1112 AI-Contributed/UT: 342/342 --- .github/workflows/paimon-python-checks.yml | 2 + docs/docs/pypaimon/robomind-agilex.md | 129 +++ paimon-python/conftest.py | 7 + paimon-python/pypaimon/multimodal/hdf5.py | 67 +- paimon-python/pypaimon/ray/__init__.py | 2 + paimon-python/pypaimon/ray/hdf5.py | 130 +++ .../pypaimon/sample/robomind_agilex.py | 775 ++++++++++++++++++ paimon-python/pypaimon/tests/ray_hdf5_test.py | 48 ++ .../tests/robomind_agilex_pipeline_test.py | 294 +++++++ 9 files changed, 1427 insertions(+), 27 deletions(-) create mode 100644 docs/docs/pypaimon/robomind-agilex.md create mode 100644 paimon-python/pypaimon/ray/hdf5.py create mode 100644 paimon-python/pypaimon/sample/robomind_agilex.py create mode 100644 paimon-python/pypaimon/tests/ray_hdf5_test.py create mode 100644 paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index 5b35fff4f83a..52e6e94eb3f4 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -141,6 +141,8 @@ jobs: python -c "import h5py; print('h5py', h5py.__version__)" if [[ "${{ matrix.python-version }}" == "3.11" ]]; then + # Run the RoboMIND pipeline tests with synthetic local HDF5 data. + python -m pip install 'h5py>=3.10,<4' # Exercise the 0.4 API in one lane until its wheel is published. python -m pip install "git+https://github.com/apache/paimon-rust.git@${PYPAIMON_RUST_REV}#subdirectory=bindings/python" python -m pip install "./paimon-python[sql]" diff --git a/docs/docs/pypaimon/robomind-agilex.md b/docs/docs/pypaimon/robomind-agilex.md new file mode 100644 index 000000000000..a1417839c549 --- /dev/null +++ b/docs/docs/pypaimon/robomind-agilex.md @@ -0,0 +1,129 @@ +--- +title: "RoboMIND AgileX" +sidebar_position: 7 +--- + + + +# RoboMIND AgileX + +The RoboMIND AgileX sample turns a downloaded HDF5 directory into three +Paimon tables: + +- `episodes_agilex` stores episode metadata derived from the dataset layout; +- `frames_agilex` stores ordered robot state, raw action, RGB, and depth rows; +- `feature_stats_agilex` versions the train-split statistics consumed by + policy training. + +The local and Ray paths use the same `RoboMindAgileXEpisodeTransform` and +`RoboMindAgileXFrameTransform` contracts and table schemas. Ray assigns each +complete HDF5 file to one transform task, while the Paimon sink performs one +coordinated commit. Discovery does not open or hash HDF5 contents; validation +and frame counting happen after a transform task has opened the file. + +`split` is RoboMIND dataset metadata derived from the `train` or `val` +directory component. It is not a required field of every Paimon multimodal +table. This sample uses successful train episodes to select frame rows for +normalization statistics. + +## Run the local pipeline + +After downloading RoboMIND, install the HDF5 extra and provide the source and +warehouse directories to one command: + +```bash +pip install 'pypaimon[hdf5]' +python -m pypaimon.sample.robomind_agilex \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --warehouse /data/warehouse +``` + +The command discovers and validates every `**/data/trajectory.hdf5`, ingests +the episode and frame tables locally, materializes the canonical action, and +writes versioned train-split normalization statistics. It prints a JSON result +with row counts and committed snapshot IDs. The input must already be present +locally; the command does not download RoboMIND or contact Hugging Face. + +Use a new warehouse for each run. Ingestion is append-only, so repeating the +same input against existing tables would create duplicate rows; a completed +warehouse also rejects a second canonical-action backfill. + +Pytest generates several small HDF5 episodes with the real AgileX field names, +shapes, dtypes, split layout, and success layout, so the default test needs no +download. To exercise a downloaded customer dataset explicitly, run: + +```bash +pytest -q pypaimon/tests/robomind_agilex_pipeline_test.py \ + --robomind-agilex-input /data/RoboMIND/h5_agilex_3rgb +``` + +## Python API + +```python +from pypaimon.sample.robomind_agilex import ( + backfill_canonical_action, + ingest_local, + ingest_ray, + run_local_pipeline, +) + +# Run local ingestion and canonical-action backfill together. +pipeline = run_local_pipeline( + "/data/RoboMIND/h5_agilex_3rgb", + "/data/warehouse", +) + +# Or compose the lower-level operations explicitly. Ray chooses distributed +# task placement; concurrency is only an optional upper bound. +ingest = ingest_ray( + "/data/RoboMIND/h5_agilex_3rgb", + "/data/warehouse", + concurrency=8, +) + +backfill = backfill_canonical_action( + "/data/warehouse", + statistics_version="robomind-agilex-joint-position@1", +) +``` + +Episode and frame ingestion commit separately and use the generic +`pypaimon.ray.load_from_hdf5` API in Ray mode. Canonical action materialization +and statistics refresh also commit separately. If statistics need to be +regenerated, call `refresh_action_statistics` without repeating ingestion or +the row-id update. + +The canonical `action` is `float32(concat(master/joint_position_left, +master/joint_position_right))`. The backfill materializes only this consumed +14-dimensional column. It does not materialize normalized actions. Instead, +the stats table stores the train-only population mean and standard deviation, +the `1e-2` standard-deviation floor, the train split manifest digest, and the +source `frames_agilex` snapshot. A training reader normalizes `action` at read +time with that versioned row. + +The tables are non-primary-key append tables. Repeating ingestion therefore +appends duplicate rows by design; it does not mean row-level update/delete is +disabled. The sample keeps deletion vectors enabled, stores vectors with +Vortex, and sets `blob-as-descriptor=false` because its transforms emit raw +image/depth bytes rather than external BLOB descriptors. Parquet data format, +dynamic bucket mode, and global-index search mode are inherited defaults and +are not repeated in the sample options. + +Run local and Ray modes against separate new warehouses when comparing them. diff --git a/paimon-python/conftest.py b/paimon-python/conftest.py index f6adcd31cac8..4ea23aa02e36 100644 --- a/paimon-python/conftest.py +++ b/paimon-python/conftest.py @@ -24,6 +24,13 @@ _force_native_for_test = False +def pytest_addoption(parser): + parser.addoption( + "--robomind-agilex-input", + help="Downloaded RoboMIND AgileX directory for the optional sample test.", + ) + + def _native_plan_enabled(): return os.environ.get(_NATIVE_PLAN_ENV) == "1" diff --git a/paimon-python/pypaimon/multimodal/hdf5.py b/paimon-python/pypaimon/multimodal/hdf5.py index 68389cf4e28c..c136e0e804d2 100644 --- a/paimon-python/pypaimon/multimodal/hdf5.py +++ b/paimon-python/pypaimon/multimodal/hdf5.py @@ -188,33 +188,17 @@ def _load_hdf5_files(table, files, transform, source_file_io, h5py): table_commit.add_commit_callback(snapshot_recorder) for source in files: - source_row_count = 0 - with closing(source_file_io.new_input_stream(source.path)) as stream: - _require_seekable(stream, source) - with h5py.File(stream, "r") as h5: - transformed = transform(h5, source) - batches = None - try: - batches = _arrow_batches(transformed) - for value in batches: - arrow_table = _strict_arrow_table( - value, - target_schema, - source, - batch_count, - ) - batch_count += 1 - row_count += arrow_table.num_rows - source_row_count += arrow_table.num_rows - if arrow_table.num_rows: - table_write.write_arrow(arrow_table) - finally: - _close_transform_iterator( - batches if batches is not None else transformed) - - if source_row_count == 0: - raise ValueError( - "HDF5 source %s produced no rows." % source.path) + for arrow_table in _transform_hdf5_file( + source, + transform, + source_file_io, + h5py, + target_schema, + batch_index=batch_count): + batch_count += 1 + row_count += arrow_table.num_rows + if arrow_table.num_rows: + table_write.write_arrow(arrow_table) commit_messages = table_write.prepare_commit() commit_started = True @@ -241,6 +225,35 @@ def _load_hdf5_files(table, files, transform, source_file_io, h5py): table_commit.close() +def _transform_hdf5_file( + source, + transform, + source_file_io, + h5py, + target_schema, + *, + batch_index=0): + """Yield validated Arrow tables for one HDF5 source.""" + produced_rows = 0 + with closing(source_file_io.new_input_stream(source.path)) as stream: + _require_seekable(stream, source) + with h5py.File(stream, "r") as h5: + transformed = transform(h5, source) + batches = None + try: + batches = _arrow_batches(transformed) + for index, value in enumerate(batches, start=batch_index): + arrow_table = _strict_arrow_table( + value, target_schema, source, index) + produced_rows += arrow_table.num_rows + yield arrow_table + finally: + _close_transform_iterator( + batches if batches is not None else transformed) + if produced_rows == 0: + raise ValueError("HDF5 source %s produced no rows." % source.path) + + def _discover_hdf5_files(paths, source_file_io): values = _path_values(paths) normalized = {} diff --git a/paimon-python/pypaimon/ray/__init__.py b/paimon-python/pypaimon/ray/__init__.py index 63141ecd41f0..e29d8450614d 100644 --- a/paimon-python/pypaimon/ray/__init__.py +++ b/paimon-python/pypaimon/ray/__init__.py @@ -31,6 +31,7 @@ from pypaimon.ray.update_by_row_id import update_by_row_id from pypaimon.ray.read_by_row_id import read_by_row_id from pypaimon.ray.process_row_id_ranges import process_row_id_ranges +from pypaimon.ray.hdf5 import load_from_hdf5 __all__ = [ "read_paimon", @@ -42,6 +43,7 @@ "update_by_row_id", "read_by_row_id", "process_row_id_ranges", + "load_from_hdf5", "WhenMatched", "WhenNotMatched", "source_col", diff --git a/paimon-python/pypaimon/ray/hdf5.py b/paimon-python/pypaimon/ray/hdf5.py new file mode 100644 index 000000000000..114914aaa104 --- /dev/null +++ b/paimon-python/pypaimon/ray/hdf5.py @@ -0,0 +1,130 @@ +# 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. + +"""Distributed HDF5 ingestion using the multimodal transform contract.""" + +from typing import Any, Dict, Mapping, Optional + + +def load_from_hdf5( + table_identifier: str, + paths, + catalog_options: Dict[str, str], + *, + transform, + source_options: Optional[Mapping[str, object]] = None, + concurrency: Optional[int] = None, + ray_remote_args: Optional[Dict[str, Any]] = None) -> None: + """Transform complete HDF5 files on Ray and append them in one commit. + + The transform has the same ``(h5py.File, Hdf5File)`` contract as + :meth:`MultimodalConnection.load_from_hdf5`. Discovery runs on the driver; + workers open and transform complete files, and the Paimon Ray sink commits + all worker messages once. + """ + if not callable(transform): + raise ValueError("transform must be callable.") + + from pypaimon.catalog.catalog_factory import CatalogFactory + from pypaimon.common.options import Options + from pypaimon.multimodal.hdf5 import ( + _Hdf5SourceFileIO, + _discover_hdf5_files, + _path_values, + _validate_source_kerberos, + _validated_source_options, + ) + from pypaimon.multimodal.table import _target_schema + from pypaimon.ray.ray_paimon import _require_ray_data, write_paimon + + ray_data = _require_ray_data() + validated_options = _validated_source_options(source_options) + path_values = _path_values(paths) + _validate_source_kerberos(path_values, validated_options) + source_file_io = _Hdf5SourceFileIO(Options(validated_options)) + try: + files = _discover_hdf5_files(path_values, source_file_io) + finally: + source_file_io.close() + if not files: + return + + table = CatalogFactory.create(catalog_options).get_table(table_identifier) + target_schema = _target_schema(table) + inputs = ray_data.from_items( + [{"path": source.path} for source in files], + override_num_blocks=len(files), + ) + transformed = inputs.map_batches( + _TransformHdf5File, + fn_constructor_kwargs={ + "transform": transform, + "source_options": validated_options, + "target_schema": target_schema, + }, + batch_format="pyarrow", + batch_size=1, + concurrency=concurrency, + **dict(ray_remote_args or {}), + ) + write_paimon( + transformed, + table_identifier, + catalog_options, + concurrency=concurrency, + ray_remote_args=ray_remote_args, + ) + + +class _TransformHdf5File: + + def __init__(self, *, transform, source_options, target_schema): + self.transform = transform + self.source_options = source_options + self.target_schema = target_schema + + def __call__(self, batch): + if batch.num_rows != 1: + raise ValueError("Ray HDF5 transform requires one source per batch.") + + from pypaimon.common.options import Options + from pypaimon.multimodal.hdf5 import ( + Hdf5File, + _Hdf5SourceFileIO, + _transform_hdf5_file, + ) + + try: + import h5py + except ImportError as error: + raise ImportError( + "load_from_hdf5 requires h5py; install 'pypaimon[ray,hdf5]'." + ) from error + + source = Hdf5File(path=batch["path"][0].as_py()) + source_file_io = _Hdf5SourceFileIO(Options(self.source_options)) + try: + for table in _transform_hdf5_file( + source, + self.transform, + source_file_io, + h5py, + self.target_schema): + if table.num_rows: + yield table + finally: + source_file_io.close() diff --git a/paimon-python/pypaimon/sample/robomind_agilex.py b/paimon-python/pypaimon/sample/robomind_agilex.py new file mode 100644 index 000000000000..008445ee77fe --- /dev/null +++ b/paimon-python/pypaimon/sample/robomind_agilex.py @@ -0,0 +1,775 @@ +# 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. + +"""RoboMIND AgileX HDF5 ingestion and canonical action backfill.""" + +import argparse +import hashlib +import json +from dataclasses import asdict, dataclass +from pathlib import Path + +import numpy as np +import pyarrow as pa + +import pypaimon.multimodal as pmm + + +DEFAULT_DATABASE = "robomind" +EPISODES_TABLE = "episodes_agilex" +FRAMES_TABLE = "frames_agilex" +FEATURE_STATS_TABLE = "feature_stats_agilex" +DEFAULT_STATISTICS_VERSION = "robomind-agilex-joint-position@1" + +TABLE_OPTIONS = { + "blob-as-descriptor": "false", +} + +NUMERIC_FIELDS = ( + ("state_end_effector_left", "puppet/end_effector_left"), + ("state_end_effector_right", "puppet/end_effector_right"), + ("state_joint_effort_left", "puppet/joint_effort_left"), + ("state_joint_effort_right", "puppet/joint_effort_right"), + ("state_joint_position_left", "puppet/joint_position_left"), + ("state_joint_position_right", "puppet/joint_position_right"), + ("state_joint_velocity_left", "puppet/joint_velocity_left"), + ("state_joint_velocity_right", "puppet/joint_velocity_right"), + ("action_end_effector_left", "master/end_effector_left"), + ("action_end_effector_right", "master/end_effector_right"), + ("action_joint_effort_left", "master/joint_effort_left"), + ("action_joint_effort_right", "master/joint_effort_right"), + ("action_joint_position_left", "master/joint_position_left"), + ("action_joint_position_right", "master/joint_position_right"), + ("action_joint_velocity_left", "master/joint_velocity_left"), + ("action_joint_velocity_right", "master/joint_velocity_right"), +) +IMAGE_FIELDS = ( + ("rgb_front", "observations/rgb_images/camera_front"), + ("rgb_left_wrist", "observations/rgb_images/camera_left_wrist"), + ("rgb_right_wrist", "observations/rgb_images/camera_right_wrist"), + ("depth_front", "observations/depth_images/camera_front"), + ("depth_left_wrist", "observations/depth_images/camera_left_wrist"), + ("depth_right_wrist", "observations/depth_images/camera_right_wrist"), +) + +_ACTION_LEFT = "action_joint_position_left" +_ACTION_RIGHT = "action_joint_position_right" +_ACTION_COLUMN = "action" +_ACTION_VECTOR_TYPE = pa.list_(pa.float32(), 14) +_STANDARD_DEVIATION_FLOOR = 1e-2 +@dataclass(frozen=True) +class EpisodeSource: + """RoboMIND metadata derived without opening the HDF5 source.""" + + path: Path + source_key: str + episode_id: str + split: str + success: bool + + +@dataclass(frozen=True) +class IngestResult: + """Small control-plane result returned by an AgileX ingestion.""" + + mode: str + episode_count: int + frame_count: int + episodes_snapshot_id: int + frames_snapshot_id: int + + +@dataclass(frozen=True) +class BackfillResult: + """Result of materializing canonical action and its statistics row.""" + + row_count: int + frames_snapshot_id: int + statistics_snapshot_id: int + statistics_version: str + + +@dataclass(frozen=True) +class LocalPipelineResult: + """Result of the complete local ingestion and backfill pipeline.""" + + ingest: IngestResult + backfill: BackfillResult + + +def episode_schema(): + """Return the shared AgileX episode business schema.""" + return pa.schema([ + pa.field("episode_id", pa.string(), nullable=False), + pa.field("source_key", pa.string(), nullable=False), + pa.field("split", pa.string(), nullable=False), + pa.field("success", pa.bool_(), nullable=False), + pa.field("instruction", pa.string()), + pa.field("instruction_embedding", pa.list_(pa.float32(), 768)), + pa.field("frame_count", pa.int32(), nullable=False), + pa.field("hdf5_compress", pa.bool_()), + pa.field("hdf5_sim", pa.bool_()), + ]) + + +def frame_schema(): + """Return the shared AgileX frame schema before canonical backfill.""" + fields = [ + pa.field("episode_id", pa.string(), nullable=False), + pa.field("frame_index", pa.int32(), nullable=False), + ] + fields.extend( + pa.field(name, pa.large_binary(), nullable=False) + for name, _ in IMAGE_FIELDS + ) + fields.extend( + pa.field(name, pa.list_(pa.float64(), 7), nullable=False) + for name, _ in NUMERIC_FIELDS + ) + return pa.schema(fields) + + +def backfilled_frame_schema(): + """Return the frame schema after canonical action is added.""" + return frame_schema().append(pa.field(_ACTION_COLUMN, _ACTION_VECTOR_TYPE)) + + +def feature_stats_schema(): + """Return the versioned normalization-statistics schema.""" + return pa.schema([ + pa.field("statistics_version", pa.string(), nullable=False), + pa.field("source_table", pa.string(), nullable=False), + pa.field("source_snapshot_id", pa.int64(), nullable=False), + pa.field("source_split", pa.string(), nullable=False), + pa.field("split_manifest_sha256", pa.string(), nullable=False), + pa.field("feature_name", pa.string(), nullable=False), + pa.field("frame_count", pa.int64(), nullable=False), + pa.field("action_mean", pa.list_(pa.float64(), 14), nullable=False), + pa.field("action_std", pa.list_(pa.float64(), 14), nullable=False), + pa.field("standard_deviation_floor", pa.float64(), nullable=False), + ]) + + +def discover_episodes(input_root): + """Discover RoboMIND episode paths without reading HDF5 contents.""" + root = Path(input_root).expanduser().resolve() + if not root.is_dir(): + raise ValueError("RoboMIND input root does not exist: %s" % root) + paths = sorted(root.glob("**/data/trajectory.hdf5")) + if not paths: + raise ValueError("No RoboMIND trajectory.hdf5 files found below %s." % root) + + episodes = [] + episode_ids = set() + for path in paths: + source_key = path.relative_to(root).as_posix() + split = _path_component(source_key, ("train", "val"), "split") + status = _path_component( + source_key, ("success_episodes", "failed_episodes"), "status") + episode_id = path.parent.parent.name + if episode_id in episode_ids: + raise ValueError("Duplicate RoboMIND episode_id %r." % episode_id) + episode_ids.add(episode_id) + episodes.append(EpisodeSource( + path=path, + source_key=source_key, + episode_id=episode_id, + split=split, + success=status == "success_episodes", + )) + return episodes + + +class _RoboMindAgileXTransform: + + def __init__(self, episodes): + self._episodes = { + episode.path: episode for episode in episodes + } + if not self._episodes: + raise ValueError("episodes must not be empty.") + + def _source(self, source): + source_path = source.local_path + if source_path is None: + raise ValueError( + "RoboMIND AgileX requires a local HDF5 source: %s" + % source.path + ) + source_path = source_path.resolve() + episode = self._episodes.get(source_path) + if episode is None: + raise ValueError("Unknown RoboMIND source path %r." % source_path) + return episode + + +class RoboMindAgileXEpisodeTransform(_RoboMindAgileXTransform): + """Validate one AgileX file and emit its episode metadata row.""" + + def __call__(self, h5, source): + episode = self._source(source) + frame_count = _validate_source(h5, episode.source_key) + yield self._episode_batch(h5, episode, frame_count) + + @staticmethod + def _episode_batch(h5, episode, frame_count): + instruction = _instruction(h5, episode.source_key) + embedding = _instruction_embedding(h5, episode.source_key) + return pa.RecordBatch.from_pydict({ + "episode_id": [episode.episode_id], + "source_key": [episode.source_key], + "split": [episode.split], + "success": [episode.success], + "instruction": [instruction], + "instruction_embedding": [embedding], + "frame_count": [frame_count], + "hdf5_compress": [_optional_bool(h5.attrs.get("compress"))], + "hdf5_sim": [_optional_bool(h5.attrs.get("sim"))], + }, schema=episode_schema()) + + +class RoboMindAgileXFrameTransform(_RoboMindAgileXTransform): + """Validate one AgileX file and stream its frame rows.""" + + def __init__(self, episodes, *, batch_size=64): + super().__init__(episodes) + self.batch_size = _positive_int(batch_size, "batch_size") + + def __call__(self, h5, source): + episode = self._source(source) + frame_count = _validate_source(h5, episode.source_key) + for begin in range(0, frame_count, self.batch_size): + end = min(begin + self.batch_size, frame_count) + count = end - begin + columns = { + "episode_id": [episode.episode_id] * count, + "frame_index": np.arange(begin, end, dtype=np.int32), + } + for name, hdf5_path in IMAGE_FIELDS: + columns[name] = [ + np.asarray(value, dtype=np.uint8).tobytes() + for value in h5[hdf5_path][begin:end] + ] + for name, hdf5_path in NUMERIC_FIELDS: + values = np.asarray(h5[hdf5_path][begin:end], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError( + "%s: /%s contains NaN or Inf." + % (episode.source_key, hdf5_path) + ) + columns[name] = values.tolist() + yield pa.RecordBatch.from_pydict(columns, schema=frame_schema()) + + +def ingest_local( + input_root, + warehouse, + *, + database=DEFAULT_DATABASE, + batch_size=64): + """Ingest AgileX episodes locally through strict ``load_from_hdf5``.""" + episodes = discover_episodes(input_root) + connection, _, _ = _create_tables(warehouse, database) + paths = [episode.path for episode in episodes] + episode_result = connection.load_from_hdf5( + EPISODES_TABLE, + paths, + transform=RoboMindAgileXEpisodeTransform(episodes), + ) + frame_result = connection.load_from_hdf5( + FRAMES_TABLE, + paths, + transform=RoboMindAgileXFrameTransform( + episodes, batch_size=batch_size), + ) + del connection + return IngestResult( + mode="local", + episode_count=episode_result.row_count, + frame_count=frame_result.row_count, + episodes_snapshot_id=episode_result.snapshot_id, + frames_snapshot_id=frame_result.snapshot_id, + ) + + +def run_local_pipeline( + input_root, + warehouse, + *, + database=DEFAULT_DATABASE, + batch_size=64, + statistics_version=DEFAULT_STATISTICS_VERSION): + """Run local AgileX ingestion and canonical-action backfill.""" + ingest = ingest_local( + input_root, + warehouse, + database=database, + batch_size=batch_size, + ) + backfill = backfill_canonical_action( + warehouse, + database=database, + statistics_version=statistics_version, + ) + return LocalPipelineResult(ingest=ingest, backfill=backfill) + + +def ingest_ray( + input_root, + warehouse, + *, + database=DEFAULT_DATABASE, + batch_size=64, + concurrency=None, + ray_address=None): + """Ingest AgileX through the public distributed HDF5 loader.""" + if concurrency is not None: + concurrency = _positive_int(concurrency, "concurrency") + episodes = discover_episodes(input_root) + _, episodes_table, frames_table = _create_tables(warehouse, database) + + try: + import ray + except ImportError: + raise ImportError( + "Ray ingestion requires ray; install pypaimon[ray,hdf5].") + + initialized_here = not ray.is_initialized() + if initialized_here: + init_args = { + "include_dashboard": False, + "ignore_reinit_error": True, + } + if ray_address is None: + init_args["num_cpus"] = 2 + else: + init_args["address"] = ray_address + ray.init(**init_args) + elif ray_address is not None: + raise ValueError( + "ray_address cannot be set after Ray has already been initialized.") + + try: + from pypaimon.ray import load_from_hdf5 + catalog_options = { + "warehouse": str(Path(warehouse).expanduser().resolve())} + paths = [episode.path for episode in episodes] + load_from_hdf5( + "%s.%s" % (database, EPISODES_TABLE), paths, catalog_options, + transform=RoboMindAgileXEpisodeTransform(episodes), + concurrency=concurrency, + ) + load_from_hdf5( + "%s.%s" % (database, FRAMES_TABLE), paths, catalog_options, + transform=RoboMindAgileXFrameTransform( + episodes, batch_size=batch_size), + concurrency=concurrency, + ) + episode_rows = _read_raw(episodes_table.raw_table, ["episode_id"]) + frame_rows = _read_raw(frames_table.raw_table, ["frame_index"]) + return IngestResult( + mode="ray", + episode_count=episode_rows.num_rows, + frame_count=frame_rows.num_rows, + episodes_snapshot_id=_snapshot_id(episodes_table), + frames_snapshot_id=_snapshot_id(frames_table), + ) + finally: + if initialized_here: + ray.shutdown() + + +def build_canonical_action_backfill(source): + """Build canonical actions keyed by the physical row ID.""" + required = [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"] + missing = [name for name in required if name not in source.column_names] + if missing: + raise ValueError("Source is missing required columns: %s." % missing) + left = np.asarray(source[_ACTION_LEFT].to_pylist(), dtype=np.float64) + right = np.asarray(source[_ACTION_RIGHT].to_pylist(), dtype=np.float64) + if left.ndim != 2 or left.shape[1:] != (7,): + raise ValueError( + "%s must have shape (rows, 7), got %s." + % (_ACTION_LEFT, left.shape) + ) + if right.shape != left.shape: + raise ValueError( + "%s must have shape %s, got %s." + % (_ACTION_RIGHT, left.shape, right.shape) + ) + action64 = np.concatenate([left, right], axis=1) + if not np.isfinite(action64).all(): + raise ValueError("Canonical action input contains NaN or Inf.") + action = action64.astype(np.float32) + return pa.table({ + "_ROW_ID": source["_ROW_ID"], + _ACTION_COLUMN: pa.array(action.tolist(), type=_ACTION_VECTOR_TYPE), + }) + + +def build_action_statistics(source, train_episode_ids): + """Compute train-only action population statistics.""" + train_ids = set(train_episode_ids) + if not train_ids: + raise ValueError("Cannot compute action statistics without train episodes.") + episode_ids = np.asarray(source["episode_id"].to_pylist(), dtype=object) + frame_indices = np.asarray(source["frame_index"].to_pylist(), dtype=np.int64) + train = np.asarray([value in train_ids for value in episode_ids], dtype=bool) + if not train.any(): + raise ValueError("No frame rows belong to the train episodes.") + action = np.asarray(source[_ACTION_COLUMN].to_pylist(), dtype=np.float64) + train_rows = sorted( + np.flatnonzero(train), + key=lambda index: (episode_ids[index], int(frame_indices[index]))) + train_action = action[train_rows] + mean = train_action.mean(axis=0) + std = np.maximum(train_action.std(axis=0), _STANDARD_DEVIATION_FLOOR) + statistics = { + "frame_count": int(train.sum()), + "action_mean": mean.tolist(), + "action_std": std.tolist(), + "standard_deviation_floor": _STANDARD_DEVIATION_FLOOR, + } + return statistics + + +def backfill_canonical_action( + warehouse, + *, + database=DEFAULT_DATABASE, + statistics_version=DEFAULT_STATISTICS_VERSION): + """Run the independently recoverable action and statistics stages.""" + row_count, frames_snapshot_id = materialize_canonical_action( + warehouse, database=database) + statistics_snapshot_id = refresh_action_statistics( + warehouse, + database=database, + statistics_version=statistics_version, + ) + return BackfillResult( + row_count=row_count, + frames_snapshot_id=frames_snapshot_id, + statistics_snapshot_id=statistics_snapshot_id, + statistics_version=statistics_version, + ) + + +def materialize_canonical_action(warehouse, *, database=DEFAULT_DATABASE): + """Stage one: add and populate canonical action, then commit it.""" + connection = pmm.connect( + database=database, + options={"warehouse": str(Path(warehouse).expanduser().resolve())}, + ) + frames_table = connection.get_table(FRAMES_TABLE) + _validate_backfill_target(frames_table) + + from pypaimon.schema.data_types import AtomicType, VectorType + from pypaimon.schema.schema_change import SchemaChange + + connection.catalog.alter_table( + frames_table.identifier, + [SchemaChange.add_column( + _ACTION_COLUMN, + VectorType(True, AtomicType("FLOAT"), 14), + comment=( + "Canonical AgileX action: master joint position left " + "followed by right."), + )], + False, + ) + frames_table = connection.get_table(FRAMES_TABLE) + row_count = _update_canonical_action_batches(frames_table.raw_table) + frames_table = connection.get_table(FRAMES_TABLE) + frames_snapshot_id = _snapshot_id(frames_table) + return row_count, frames_snapshot_id + + +def refresh_action_statistics( + warehouse, + *, + database=DEFAULT_DATABASE, + statistics_version=DEFAULT_STATISTICS_VERSION): + """Stage two: recompute statistics from committed episode/frame tables.""" + if not isinstance(statistics_version, str) or not statistics_version: + raise ValueError("statistics_version must be a non-empty string.") + connection = pmm.connect( + database=database, + options={"warehouse": str(Path(warehouse).expanduser().resolve())}, + ) + episodes_table = connection.get_table(EPISODES_TABLE) + frames_table = connection.get_table(FRAMES_TABLE) + if _ACTION_COLUMN not in frames_table.raw_table.field_names: + raise ValueError("Canonical action column does not exist.") + + episode_rows = _read_raw( + episodes_table.raw_table, ["episode_id", "split", "success"]) + train_episode_ids = sorted( + row["episode_id"] + for row in episode_rows.to_pylist() + if row["split"] == "train" and row["success"] + ) + frames_snapshot_id = _snapshot_id(frames_table) + statistics = _stream_action_statistics( + frames_table.raw_table, train_episode_ids) + + split_manifest_sha256 = hashlib.sha256( + "".join("%s\n" % value for value in train_episode_ids) + .encode("utf-8") + ).hexdigest() + stats_table = connection.create_table( + FEATURE_STATS_TABLE, + schema=feature_stats_schema(), + options=TABLE_OPTIONS, + ignore_if_exists=True, + ) + stats_table.add(pa.Table.from_pylist([{ + "statistics_version": statistics_version, + "source_table": "%s.%s" % (database, FRAMES_TABLE), + "source_snapshot_id": frames_snapshot_id, + "source_split": "train", + "split_manifest_sha256": split_manifest_sha256, + "feature_name": _ACTION_COLUMN, + "frame_count": statistics["frame_count"], + "action_mean": statistics["action_mean"], + "action_std": statistics["action_std"], + "standard_deviation_floor": statistics[ + "standard_deviation_floor"], + }], schema=feature_stats_schema())) + return _snapshot_id(stats_table) + + +def _create_tables(warehouse, database): + connection = pmm.connect( + database=database, + options={"warehouse": str(Path(warehouse).expanduser().resolve())}, + ) + episodes_table = connection.create_table( + EPISODES_TABLE, schema=episode_schema(), options=TABLE_OPTIONS) + frames_table = connection.create_table( + FRAMES_TABLE, schema=frame_schema(), options=TABLE_OPTIONS) + return connection, episodes_table, frames_table + + +def _validate_backfill_target(frames_table): + missing = [ + name for name in (_ACTION_LEFT, _ACTION_RIGHT) + if name not in frames_table.raw_table.field_names + ] + if missing: + raise ValueError("Frames table is missing raw action columns: %s." % missing) + if _ACTION_COLUMN in frames_table.raw_table.field_names: + raise ValueError("Canonical action column already exists.") + + +def _update_canonical_action_batches(table): + """Transform one planned Paimon split at a time and commit all updates once.""" + builder = table.new_batch_write_builder() + commit = builder.new_commit() + messages = [] + row_count = 0 + try: + for source in _iter_raw( + table, [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"]): + updates = build_canonical_action_backfill(source) + messages.extend( + builder.new_update() + .with_update_type([_ACTION_COLUMN]) + .update_by_arrow_with_row_id(updates) + ) + row_count += len(updates) + commit.commit(messages) + finally: + commit.close() + return row_count + + +def _stream_action_statistics(table, train_episode_ids): + """Accumulate only fixed-size count/sum/sum-of-squares on the driver.""" + train_ids = set(train_episode_ids) + if not train_ids: + raise ValueError("Cannot compute action statistics without train episodes.") + count = 0 + total = np.zeros(14, dtype=np.float64) + total_square = np.zeros(14, dtype=np.float64) + for source in _iter_raw(table, ["episode_id", _ACTION_COLUMN]): + selected = [ + index for index, value in enumerate(source["episode_id"].to_pylist()) + if value in train_ids + ] + if not selected: + continue + action = np.asarray( + source[_ACTION_COLUMN].take(pa.array(selected)).to_pylist(), + dtype=np.float64, + ) + count += len(action) + total += action.sum(axis=0) + total_square += np.square(action).sum(axis=0) + if count == 0: + raise ValueError("No frame rows belong to the train episodes.") + mean = total / count + variance = np.maximum(total_square / count - np.square(mean), 0.0) + return { + "frame_count": count, + "action_mean": mean.tolist(), + "action_std": np.maximum( + np.sqrt(variance), _STANDARD_DEVIATION_FLOOR).tolist(), + "standard_deviation_floor": _STANDARD_DEVIATION_FLOOR, + } + + +def _iter_raw(table, columns): + builder = table.new_read_builder().with_projection(columns) + read = builder.new_read() + for split in builder.new_scan().plan().splits(): + yield read.to_arrow([split]) + + +def _read_raw(table, columns): + builder = table.new_read_builder().with_projection(columns) + plan = builder.new_scan().plan() + return builder.new_read().to_arrow(plan.splits()) + + +def _snapshot_id(table): + raw_table = table.raw_table if hasattr(table, "raw_table") else table + snapshot = raw_table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + raise RuntimeError("Expected a committed Paimon snapshot.") + return snapshot.id + + +def _validate_source(h5, source_key): + lengths = set() + for _, hdf5_path in NUMERIC_FIELDS: + if hdf5_path not in h5 or h5[hdf5_path].shape[1:] != (7,): + raise ValueError( + "%s: invalid /%s shape." % (source_key, hdf5_path)) + if h5[hdf5_path].dtype != np.dtype("float64"): + raise ValueError( + "%s: invalid /%s dtype." % (source_key, hdf5_path)) + lengths.add(int(h5[hdf5_path].shape[0])) + for _, hdf5_path in IMAGE_FIELDS: + if hdf5_path not in h5 or len(h5[hdf5_path].shape) != 1: + raise ValueError( + "%s: invalid /%s shape." % (source_key, hdf5_path)) + lengths.add(int(h5[hdf5_path].shape[0])) + if len(lengths) != 1: + raise ValueError("%s: frame lengths differ." % source_key) + frame_count = lengths.pop() + if frame_count <= 0: + raise ValueError("%s: episode has no frames." % source_key) + _instruction(h5, source_key) + _instruction_embedding(h5, source_key) + return frame_count + + +def _instruction(h5, source_key): + if "language_raw" not in h5 or h5["language_raw"].shape != (1,): + raise ValueError("%s: invalid /language_raw shape." % source_key) + value = h5["language_raw"][0] + if isinstance(value, bytes): + return value.decode("utf-8") + if isinstance(value, str): + return value + raise ValueError("%s: /language_raw is not UTF-8 text." % source_key) + + +def _instruction_embedding(h5, source_key): + if ("language_distilbert" not in h5 + or h5["language_distilbert"].shape != (1, 1, 768)): + raise ValueError( + "%s: invalid /language_distilbert shape." % source_key) + values = np.asarray(h5["language_distilbert"][0, 0], dtype=np.float32) + if not np.isfinite(values).all(): + raise ValueError("%s: language embedding contains NaN or Inf." % source_key) + return values.tolist() + + +def _path_component(source_key, candidates, label): + matches = [value for value in Path(source_key).parts if value in candidates] + if len(matches) != 1: + raise ValueError( + "Cannot derive RoboMIND %s from %s." % (label, source_key)) + return matches[0] + + +def _optional_bool(value): + return None if value is None else bool(value) + + +def _positive_int(value, name): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("%s must be a positive int." % name) + return value + + +def _pipeline_summary(result): + """Return the fixed-size control-plane result printed by the CLI.""" + return { + "ingest": { + "mode": result.ingest.mode, + "episode_count": result.ingest.episode_count, + "frame_count": result.ingest.frame_count, + "episodes_snapshot_id": result.ingest.episodes_snapshot_id, + "frames_snapshot_id": result.ingest.frames_snapshot_id, + }, + "backfill": asdict(result.backfill), + } + + +def main(argv=None): + """Run the complete local RoboMIND AgileX pipeline from the command line.""" + parser = argparse.ArgumentParser( + description=( + "Ingest a downloaded RoboMIND AgileX HDF5 directory and " + "materialize canonical actions and normalization statistics." + ) + ) + parser.add_argument( + "--input", required=True, metavar="DIRECTORY", + help="downloaded RoboMIND AgileX HDF5 root", + ) + parser.add_argument( + "--warehouse", required=True, metavar="DIRECTORY", + help="new local Paimon warehouse directory", + ) + parser.add_argument( + "--database", default=DEFAULT_DATABASE, + help="Paimon database name (default: %(default)s)", + ) + parser.add_argument( + "--batch-size", default=64, type=int, metavar="ROWS", + help="frame rows per transform batch (default: %(default)s)", + ) + parser.add_argument( + "--statistics-version", default=DEFAULT_STATISTICS_VERSION, + help="version stored with train-split action statistics", + ) + args = parser.parse_args(argv) + result = run_local_pipeline( + args.input, + args.warehouse, + database=args.database, + batch_size=args.batch_size, + statistics_version=args.statistics_version, + ) + print(json.dumps(_pipeline_summary(result), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/paimon-python/pypaimon/tests/ray_hdf5_test.py b/paimon-python/pypaimon/tests/ray_hdf5_test.py new file mode 100644 index 000000000000..8215f20a364e --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_hdf5_test.py @@ -0,0 +1,48 @@ +# 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. + +from unittest.mock import patch + +import pyarrow as pa +import pytest + +from pypaimon.ray.hdf5 import _TransformHdf5File + + +h5py = pytest.importorskip("h5py") + + +def test_ray_hdf5_worker_uses_shared_file_transform(): + schema = pa.schema([pa.field("value", pa.int64(), nullable=False)]) + expected = pa.table({"value": [1]}, schema=schema) + transform = object() + worker = _TransformHdf5File( + transform=transform, + source_options={}, + target_schema=schema, + ) + + with patch( + "pypaimon.multimodal.hdf5._transform_hdf5_file", + return_value=iter([expected])) as shared_transform: + actual = list(worker(pa.table({"path": ["file:///tmp/source.h5"]}))) + + assert len(actual) == 1 + assert actual[0].equals(expected) + args = shared_transform.call_args.args + assert args[0].path == "file:///tmp/source.h5" + assert args[1] is transform + assert args[4] == schema diff --git a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py new file mode 100644 index 000000000000..c714edb99d8d --- /dev/null +++ b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py @@ -0,0 +1,294 @@ +# 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. + +import json +import subprocess +import sys + +import numpy as np +import pyarrow as pa +import pytest + +import pypaimon.multimodal as pmm +from pypaimon.sample import robomind_agilex as agilex + + +h5py = pytest.importorskip("h5py") + + +_NUMERIC_PATHS = [path for _, path in agilex.NUMERIC_FIELDS] +_IMAGE_PATHS = [path for _, path in agilex.IMAGE_FIELDS] + + +def _write_episode( + root, split, name, offset, frames=3, status="success_episodes"): + path = (root / "13_packbowl" / status / split / name + / "data" / "trajectory.hdf5") + path.parent.mkdir(parents=True) + with h5py.File(path, "w") as h5: + h5.attrs["compress"] = True + h5.attrs["sim"] = False + h5.create_dataset("language_raw", data=[b"pack the bowl"]) + h5.create_dataset( + "language_distilbert", + data=np.full((1, 1, 768), offset, dtype=np.float16), + ) + for index, hdf5_path in enumerate(_NUMERIC_PATHS): + values = np.arange(frames * 7, dtype=np.float64).reshape(frames, 7) + h5.create_dataset(hdf5_path, data=values + offset + index * 100) + variable = h5py.vlen_dtype(np.dtype("uint8")) + for index, hdf5_path in enumerate(_IMAGE_PATHS): + dataset = h5.create_dataset(hdf5_path, (frames,), dtype=variable) + for frame_index in range(frames): + payload = "%s:%s:%s" % (name, index, frame_index) + dataset[frame_index] = np.frombuffer( + payload.encode("utf-8"), dtype=np.uint8) + return path + + +@pytest.fixture +def agilex_input(tmp_path): + paths = [ + _write_episode(tmp_path, "train", "train-a", 0), + _write_episode(tmp_path, "train", "train-b", 10), + _write_episode(tmp_path, "val", "val-a", 20), + _write_episode( + tmp_path, "val", "val-b", 30, status="failed_episodes"), + ] + return tmp_path / "13_packbowl", paths + + +@pytest.fixture +def customer_agilex_input(request): + value = request.config.getoption("--robomind-agilex-input") + if not value: + pytest.skip("use --robomind-agilex-input to test downloaded data") + return value + + +def test_explicit_customer_input_uses_downloaded_episodes( + customer_agilex_input, tmp_path): + episodes = agilex.discover_episodes(customer_agilex_input) + result = agilex.ingest_local( + customer_agilex_input, tmp_path / "customer-warehouse") + assert result.episode_count == len(episodes) + assert result.frame_count > 0 + + +def _read(warehouse, table_name, columns=None): + connection = pmm.connect( + database=agilex.DEFAULT_DATABASE, + options={"warehouse": str(warehouse)}, + ) + table = connection.get_table(table_name) + query = table.scan() + if columns is not None: + query = query.select(columns) + return table, query.to_arrow() + + +def _logical_rows(warehouse, table_name, schema): + _, rows = _read(warehouse, table_name, schema.names) + if "frame_index" in schema.names: + sort_keys = [ + ("episode_id", "ascending"), + ("frame_index", "ascending"), + ] + elif "episode_id" in schema.names: + sort_keys = [("episode_id", "ascending")] + else: + sort_keys = [("statistics_version", "ascending")] + return rows.sort_by(sort_keys) + + +def _contains_payload(value): + if isinstance(value, (bytes, pa.Table, pa.RecordBatch, pa.Array)): + return True + if isinstance(value, dict): + return any(_contains_payload(item) for item in value.values()) + if isinstance(value, (list, tuple)): + return any(_contains_payload(item) for item in value) + return False + + +def test_shared_transform_streams_complete_agilex_business_schema( + agilex_input): + root, paths = agilex_input + episodes = agilex.discover_episodes(root) + transform = agilex.RoboMindAgileXFrameTransform( + episodes, batch_size=2) + source = pmm.Hdf5File(path=paths[0].as_uri()) + + with h5py.File(paths[0], "r") as h5: + batches = list(transform(h5, source)) + + assert [batch.num_rows for batch in batches] == [2, 1] + assert all(batch.schema == agilex.frame_schema() for batch in batches) + frames = pa.Table.from_batches(batches) + assert frames["episode_id"].to_pylist() == ["train-a"] * 3 + assert frames["frame_index"].to_pylist() == [0, 1, 2] + assert frames["rgb_front"][0].as_py() == b"train-a:0:0" + assert frames["depth_right_wrist"][2].as_py() == b"train-a:5:2" + assert frames["action_joint_position_left"][1].as_py() == [ + float(value) for value in range(1207, 1214) + ] + + +def test_local_ingest_and_backfill_materialize_only_canonical_action( + agilex_input, tmp_path): + root, paths = agilex_input + warehouse = tmp_path / "local-warehouse" + + completed = subprocess.run( + [ + sys.executable, + "-m", "pypaimon.sample.robomind_agilex", + "--input", str(root), + "--warehouse", str(warehouse), + "--batch-size", "2", + "--statistics-version", "synthetic-actions@1", + ], + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + ) + assert completed.returncode == 0, completed.stderr + result = json.loads(completed.stdout) + ingest = result["ingest"] + backfill = result["backfill"] + + assert ingest["mode"] == "local" + assert ingest["episode_count"] == len(paths) + assert ingest["frame_count"] == 12 + assert set(ingest) == { + "mode", + "episode_count", + "frame_count", + "episodes_snapshot_id", + "frames_snapshot_id", + } + assert not _contains_payload(result) + assert backfill["row_count"] == 12 + assert backfill["statistics_version"] == "synthetic-actions@1" + + episodes, episode_rows = _read( + warehouse, agilex.EPISODES_TABLE, agilex.episode_schema().names) + frames, frame_rows = _read(warehouse, agilex.FRAMES_TABLE) + stats, stats_rows = _read(warehouse, agilex.FEATURE_STATS_TABLE) + assert episode_rows.num_rows == 4 + assert set(episode_rows["split"].to_pylist()) == {"train", "val"} + assert set(episode_rows["success"].to_pylist()) == {True, False} + for table in (episodes, frames, stats): + options = table.raw_table.table_schema.options + assert options["deletion-vectors.enabled"] == "true" + assert options["vector.file.format"] == "vortex" + assert options["blob-as-descriptor"] == "false" + assert "file.format" not in agilex.TABLE_OPTIONS + assert "action" in frames.raw_table.field_names + assert "act_action_normalized" not in frames.raw_table.field_names + + ordered = frame_rows.select([ + "episode_id", + "frame_index", + "action_joint_position_left", + "action_joint_position_right", + "action", + ]).sort_by([ + ("episode_id", "ascending"), + ("frame_index", "ascending"), + ]) + for row in ordered.to_pylist(): + expected = np.asarray( + row["action_joint_position_left"] + + row["action_joint_position_right"], dtype=np.float32) + assert np.array_equal(np.asarray(row["action"], dtype=np.float32), expected) + + assert stats_rows.num_rows == 1 + stats_row = stats_rows.to_pylist()[0] + assert stats_row["statistics_version"] == "synthetic-actions@1" + assert stats_row["source_snapshot_id"] == backfill["frames_snapshot_id"] + assert stats_row["source_split"] == "train" + assert stats_row["frame_count"] == 6 + assert stats_row["standard_deviation_floor"] == 0.01 + assert len(stats_row["action_mean"]) == 14 + assert len(stats_row["action_std"]) == 14 + + refreshed_snapshot = agilex.refresh_action_statistics( + warehouse, statistics_version="synthetic-actions-refresh@1") + assert refreshed_snapshot > backfill["statistics_snapshot_id"] + + +def test_ray_ingest_matches_local_schema_rows_and_backfill( + agilex_input, tmp_path): + ray = pytest.importorskip("ray") + + root, paths = agilex_input + local_warehouse = tmp_path / "local-comparison" + ray_warehouse = tmp_path / "ray-comparison" + agilex.ingest_local(root, local_warehouse, batch_size=2) + local_backfill = agilex.backfill_canonical_action( + local_warehouse, statistics_version="synthetic-actions@1") + + ray.init(num_cpus=2, include_dashboard=False) + try: + ray_result = agilex.ingest_ray( + root, ray_warehouse, batch_size=2, concurrency=2) + finally: + ray.shutdown() + ray_backfill = agilex.backfill_canonical_action( + ray_warehouse, statistics_version="synthetic-actions@1") + + assert ray_result.mode == "ray" + assert ray_result.episode_count == len(paths) + assert ray_result.frame_count == 12 + assert ray_backfill.row_count == local_backfill.row_count == 12 + + for table_name, schema in ( + (agilex.EPISODES_TABLE, agilex.episode_schema()), + (agilex.FRAMES_TABLE, agilex.backfilled_frame_schema()), + (agilex.FEATURE_STATS_TABLE, agilex.feature_stats_schema())): + local_table, _ = _read(local_warehouse, table_name) + ray_table, _ = _read(ray_warehouse, table_name) + assert local_table.raw_table.table_schema.options == ( + ray_table.raw_table.table_schema.options) + assert _logical_rows(local_warehouse, table_name, schema).equals( + _logical_rows(ray_warehouse, table_name, schema)) + + +def test_action_statistics_are_independent_of_ingest_row_order(): + frame_count = 257 + base = np.linspace(-0.7, 1.3, frame_count * 14).reshape( + frame_count, 14) + values = base + np.sin(base * 17.0) * 0.003 + source = pa.table({ + "episode_id": ["train-a"] * frame_count, + "frame_index": np.arange(frame_count, dtype=np.int32), + "action_joint_position_left": values[:, :7].tolist(), + "action_joint_position_right": values[:, 7:].tolist(), + "_ROW_ID": np.arange(frame_count, dtype=np.int64), + }) + reversed_source = source.take( + pa.array(np.arange(frame_count - 1, -1, -1), type=pa.int64())) + + expected = agilex.build_action_statistics(source.append_column( + "action", pa.array(values.astype(np.float32).tolist(), + type=pa.list_(pa.float32(), 14))), ["train-a"]) + actual = agilex.build_action_statistics(reversed_source.append_column( + "action", pa.array(values[::-1].astype(np.float32).tolist(), + type=pa.list_(pa.float32(), 14))), ["train-a"]) + + assert actual == expected From f34afbc77e32b105172f65beec72df787df8c795 Mon Sep 17 00:00:00 2001 From: Yann Date: Fri, 28 Aug 2026 14:41:41 +0800 Subject: [PATCH 2/6] fix(python): satisfy RoboMIND Flake8 spacing Separate the module constants from the EpisodeSource class declaration so the Python CI matrix passes E302. Co-Authored-By: Codex Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 2/2 AI-Contributed/UT: 0/0 --- paimon-python/pypaimon/sample/robomind_agilex.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/paimon-python/pypaimon/sample/robomind_agilex.py b/paimon-python/pypaimon/sample/robomind_agilex.py index 008445ee77fe..e5b5b8c81a34 100644 --- a/paimon-python/pypaimon/sample/robomind_agilex.py +++ b/paimon-python/pypaimon/sample/robomind_agilex.py @@ -70,6 +70,8 @@ _ACTION_COLUMN = "action" _ACTION_VECTOR_TYPE = pa.list_(pa.float32(), 14) _STANDARD_DEVIATION_FLOOR = 1e-2 + + @dataclass(frozen=True) class EpisodeSource: """RoboMIND metadata derived without opening the HDF5 source.""" From 0157af005acbb9fda3a790fbc5fa763b686ae083 Mon Sep 17 00:00:00 2001 From: Yann Date: Fri, 28 Aug 2026 15:07:45 +0800 Subject: [PATCH 3/6] test(python): gate RoboMIND ingestion on Vortex Skip only the end-to-end ingestion tests when vortex-data is unavailable while retaining transform and statistics coverage on older Python versions. Co-Authored-By: Codex Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex AI-Contributed/Feature: 0/0 AI-Contributed/UT: 9/9 --- .../pypaimon/tests/robomind_agilex_pipeline_test.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py index c714edb99d8d..b23ed48a2378 100644 --- a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py +++ b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py @@ -15,6 +15,7 @@ # limitations under the License. import json +import importlib.util import subprocess import sys @@ -28,6 +29,11 @@ h5py = pytest.importorskip("h5py") +requires_vortex = pytest.mark.skipif( + importlib.util.find_spec("vortex") is None, + reason="RoboMIND ingestion uses Vortex, which requires Python >= 3.11", +) + _NUMERIC_PATHS = [path for _, path in agilex.NUMERIC_FIELDS] _IMAGE_PATHS = [path for _, path in agilex.IMAGE_FIELDS] @@ -79,6 +85,7 @@ def customer_agilex_input(request): return value +@requires_vortex def test_explicit_customer_input_uses_downloaded_episodes( customer_agilex_input, tmp_path): episodes = agilex.discover_episodes(customer_agilex_input) @@ -147,6 +154,7 @@ def test_shared_transform_streams_complete_agilex_business_schema( ] +@requires_vortex def test_local_ingest_and_backfill_materialize_only_canonical_action( agilex_input, tmp_path): root, paths = agilex_input @@ -232,6 +240,7 @@ def test_local_ingest_and_backfill_materialize_only_canonical_action( assert refreshed_snapshot > backfill["statistics_snapshot_id"] +@requires_vortex def test_ray_ingest_matches_local_schema_rows_and_backfill( agilex_input, tmp_path): ray = pytest.importorskip("ray") From 0722bf2cc6075cddaa1b09fbacf79130dcb056c7 Mon Sep 17 00:00:00 2001 From: Yann Date: Fri, 28 Aug 2026 17:00:49 +0800 Subject: [PATCH 4/6] fix(python): harden RoboMIND Ray ingestion Return exact Ray commit metadata without rescanning tables, stabilize action statistics, and strengthen validation and pipeline tests. Co-Authored-By: Codex AI-Model: gpt-5 Co-Authored-By: Codex AI-Contributed/Feature: 204/204 AI-Contributed/UT: 193/193 --- .github/workflows/paimon-python-checks.yml | 1 - paimon-python/pypaimon/multimodal/hdf5.py | 8 +- paimon-python/pypaimon/ray/hdf5.py | 24 ++- paimon-python/pypaimon/ray/ray_paimon.py | 9 +- .../pypaimon/sample/robomind_agilex.py | 81 ++++----- paimon-python/pypaimon/tests/ray_hdf5_test.py | 30 ++++ .../tests/robomind_agilex_pipeline_test.py | 163 ++++++++++++++++-- paimon-python/pypaimon/write/ray_datasink.py | 81 +++++++-- 8 files changed, 309 insertions(+), 88 deletions(-) diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index 52e6e94eb3f4..a7bdd6aa291f 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -142,7 +142,6 @@ jobs: if [[ "${{ matrix.python-version }}" == "3.11" ]]; then # Run the RoboMIND pipeline tests with synthetic local HDF5 data. - python -m pip install 'h5py>=3.10,<4' # Exercise the 0.4 API in one lane until its wheel is published. python -m pip install "git+https://github.com/apache/paimon-rust.git@${PYPAIMON_RUST_REV}#subdirectory=bindings/python" python -m pip install "./paimon-python[sql]" diff --git a/paimon-python/pypaimon/multimodal/hdf5.py b/paimon-python/pypaimon/multimodal/hdf5.py index c136e0e804d2..acd85fb5d697 100644 --- a/paimon-python/pypaimon/multimodal/hdf5.py +++ b/paimon-python/pypaimon/multimodal/hdf5.py @@ -72,10 +72,14 @@ def stem(self) -> str: @dataclass(frozen=True) class Hdf5LoadResult: - """Counts and optional snapshot for one ``load_from_hdf5`` call.""" + """Counts and optional snapshot for one ``load_from_hdf5`` call. + + ``batch_count`` is unavailable for Ray loads because counting lazy output + batches would execute the transform a second time. + """ file_count: int - batch_count: int + batch_count: Optional[int] row_count: int snapshot_id: Optional[int] diff --git a/paimon-python/pypaimon/ray/hdf5.py b/paimon-python/pypaimon/ray/hdf5.py index 114914aaa104..2871c3b88502 100644 --- a/paimon-python/pypaimon/ray/hdf5.py +++ b/paimon-python/pypaimon/ray/hdf5.py @@ -28,13 +28,15 @@ def load_from_hdf5( transform, source_options: Optional[Mapping[str, object]] = None, concurrency: Optional[int] = None, - ray_remote_args: Optional[Dict[str, Any]] = None) -> None: + ray_remote_args: Optional[Dict[str, Any]] = None): """Transform complete HDF5 files on Ray and append them in one commit. The transform has the same ``(h5py.File, Hdf5File)`` contract as :meth:`MultimodalConnection.load_from_hdf5`. Discovery runs on the driver; workers open and transform complete files, and the Paimon Ray sink commits - all worker messages once. + all worker messages once. The result reports logical output rows and the + exact committed snapshot; its batch count is ``None`` because Ray does not + expose that count without re-executing the lazy transform. """ if not callable(transform): raise ValueError("transform must be callable.") @@ -42,6 +44,7 @@ def load_from_hdf5( from pypaimon.catalog.catalog_factory import CatalogFactory from pypaimon.common.options import Options from pypaimon.multimodal.hdf5 import ( + Hdf5LoadResult, _Hdf5SourceFileIO, _discover_hdf5_files, _path_values, @@ -61,7 +64,12 @@ def load_from_hdf5( finally: source_file_io.close() if not files: - return + return Hdf5LoadResult( + file_count=0, + batch_count=None, + row_count=0, + snapshot_id=None, + ) table = CatalogFactory.create(catalog_options).get_table(table_identifier) target_schema = _target_schema(table) @@ -81,13 +89,21 @@ def load_from_hdf5( concurrency=concurrency, **dict(ray_remote_args or {}), ) - write_paimon( + write_result = write_paimon( transformed, table_identifier, catalog_options, concurrency=concurrency, ray_remote_args=ray_remote_args, ) + return Hdf5LoadResult( + file_count=len(files), + batch_count=None, + row_count=0 if write_result is None else write_result.row_count, + snapshot_id=( + None if write_result is None else write_result.snapshot_id + ), + ) class _TransformHdf5File: diff --git a/paimon-python/pypaimon/ray/ray_paimon.py b/paimon-python/pypaimon/ray/ray_paimon.py index 2a5bbf0a26b0..26968d899d58 100644 --- a/paimon-python/pypaimon/ray/ray_paimon.py +++ b/paimon-python/pypaimon/ray/ray_paimon.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: import ray.data + from pypaimon.write.ray_datasink import PaimonWriteResult def _require_ray_data(): @@ -290,7 +291,7 @@ def write_paimon( concurrency: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, hash_fixed_precluster: str = "auto", -) -> None: +) -> "Optional[PaimonWriteResult]": """Write a Ray Dataset to a Paimon table. HASH_FIXED rows are assigned to the correct bucket by the Paimon @@ -312,6 +313,10 @@ def write_paimon( hash_fixed_precluster: Pre-clustering mode. ``"auto"`` follows table options, ``"off"`` disables it, and ``"map_groups"`` explicitly enables HASH_FIXED grouping. + + Returns: + Metadata for the exact committed snapshot, or ``None`` when no + snapshot was committed. """ _require_ray_data() @@ -321,7 +326,7 @@ def write_paimon( catalog = CatalogFactory.create(catalog_options) table = catalog.get_table(table_identifier) - write_paimon_dataset( + return write_paimon_dataset( dataset, table, overwrite=overwrite, diff --git a/paimon-python/pypaimon/sample/robomind_agilex.py b/paimon-python/pypaimon/sample/robomind_agilex.py index e5b5b8c81a34..e3ccf993bfe5 100644 --- a/paimon-python/pypaimon/sample/robomind_agilex.py +++ b/paimon-python/pypaimon/sample/robomind_agilex.py @@ -177,6 +177,15 @@ def discover_episodes(input_root): episodes = [] episode_ids = set() for path in paths: + resolved_path = path.resolve() + try: + resolved_path.relative_to(root) + except ValueError: + raise ValueError( + "RoboMIND trajectory path escapes input root: %s" % path) + if path.is_symlink(): + raise ValueError( + "RoboMIND trajectory path must not be a symlink: %s" % path) source_key = path.relative_to(root).as_posix() split = _path_component(source_key, ("train", "val"), "split") status = _path_component( @@ -186,7 +195,7 @@ def discover_episodes(input_root): raise ValueError("Duplicate RoboMIND episode_id %r." % episode_id) episode_ids.add(episode_id) episodes.append(EpisodeSource( - path=path, + path=resolved_path, source_key=source_key, episode_id=episode_id, split=split, @@ -272,7 +281,8 @@ def __call__(self, h5, source): "%s: /%s contains NaN or Inf." % (episode.source_key, hdf5_path) ) - columns[name] = values.tolist() + columns[name] = pa.FixedSizeListArray.from_arrays( + pa.array(values.reshape(-1), type=pa.float64()), 7) yield pa.RecordBatch.from_pydict(columns, schema=frame_schema()) @@ -341,7 +351,7 @@ def ingest_ray( if concurrency is not None: concurrency = _positive_int(concurrency, "concurrency") episodes = discover_episodes(input_root) - _, episodes_table, frames_table = _create_tables(warehouse, database) + _create_tables(warehouse, database) try: import ray @@ -369,25 +379,23 @@ def ingest_ray( catalog_options = { "warehouse": str(Path(warehouse).expanduser().resolve())} paths = [episode.path for episode in episodes] - load_from_hdf5( + episode_result = load_from_hdf5( "%s.%s" % (database, EPISODES_TABLE), paths, catalog_options, transform=RoboMindAgileXEpisodeTransform(episodes), concurrency=concurrency, ) - load_from_hdf5( + frame_result = load_from_hdf5( "%s.%s" % (database, FRAMES_TABLE), paths, catalog_options, transform=RoboMindAgileXFrameTransform( episodes, batch_size=batch_size), concurrency=concurrency, ) - episode_rows = _read_raw(episodes_table.raw_table, ["episode_id"]) - frame_rows = _read_raw(frames_table.raw_table, ["frame_index"]) return IngestResult( mode="ray", - episode_count=episode_rows.num_rows, - frame_count=frame_rows.num_rows, - episodes_snapshot_id=_snapshot_id(episodes_table), - frames_snapshot_id=_snapshot_id(frames_table), + episode_count=episode_result.row_count, + frame_count=frame_result.row_count, + episodes_snapshot_id=episode_result.snapshot_id, + frames_snapshot_id=frame_result.snapshot_id, ) finally: if initialized_here: @@ -422,32 +430,6 @@ def build_canonical_action_backfill(source): }) -def build_action_statistics(source, train_episode_ids): - """Compute train-only action population statistics.""" - train_ids = set(train_episode_ids) - if not train_ids: - raise ValueError("Cannot compute action statistics without train episodes.") - episode_ids = np.asarray(source["episode_id"].to_pylist(), dtype=object) - frame_indices = np.asarray(source["frame_index"].to_pylist(), dtype=np.int64) - train = np.asarray([value in train_ids for value in episode_ids], dtype=bool) - if not train.any(): - raise ValueError("No frame rows belong to the train episodes.") - action = np.asarray(source[_ACTION_COLUMN].to_pylist(), dtype=np.float64) - train_rows = sorted( - np.flatnonzero(train), - key=lambda index: (episode_ids[index], int(frame_indices[index]))) - train_action = action[train_rows] - mean = train_action.mean(axis=0) - std = np.maximum(train_action.std(axis=0), _STANDARD_DEVIATION_FLOOR) - statistics = { - "frame_count": int(train.sum()), - "action_mean": mean.tolist(), - "action_std": std.tolist(), - "standard_deviation_floor": _STANDARD_DEVIATION_FLOOR, - } - return statistics - - def backfill_canonical_action( warehouse, *, @@ -585,6 +567,8 @@ def _update_canonical_action_batches(table): try: for source in _iter_raw( table, [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"]): + if source.num_rows == 0: + continue updates = build_canonical_action_backfill(source) messages.extend( builder.new_update() @@ -599,13 +583,13 @@ def _update_canonical_action_batches(table): def _stream_action_statistics(table, train_episode_ids): - """Accumulate only fixed-size count/sum/sum-of-squares on the driver.""" + """Accumulate fixed-size count, running mean, and M2 on the driver.""" train_ids = set(train_episode_ids) if not train_ids: raise ValueError("Cannot compute action statistics without train episodes.") count = 0 - total = np.zeros(14, dtype=np.float64) - total_square = np.zeros(14, dtype=np.float64) + mean = np.zeros(14, dtype=np.float64) + m2 = np.zeros(14, dtype=np.float64) for source in _iter_raw(table, ["episode_id", _ACTION_COLUMN]): selected = [ index for index, value in enumerate(source["episode_id"].to_pylist()) @@ -617,13 +601,20 @@ def _stream_action_statistics(table, train_episode_ids): source[_ACTION_COLUMN].take(pa.array(selected)).to_pylist(), dtype=np.float64, ) - count += len(action) - total += action.sum(axis=0) - total_square += np.square(action).sum(axis=0) + batch_count = len(action) + batch_mean = action.mean(axis=0) + batch_m2 = np.square(action - batch_mean).sum(axis=0) + delta = batch_mean - mean + combined_count = count + batch_count + mean += delta * batch_count / combined_count + m2 += ( + batch_m2 + + np.square(delta) * count * batch_count / combined_count + ) + count = combined_count if count == 0: raise ValueError("No frame rows belong to the train episodes.") - mean = total / count - variance = np.maximum(total_square / count - np.square(mean), 0.0) + variance = np.maximum(m2 / count, 0.0) return { "frame_count": count, "action_mean": mean.tolist(), diff --git a/paimon-python/pypaimon/tests/ray_hdf5_test.py b/paimon-python/pypaimon/tests/ray_hdf5_test.py index 8215f20a364e..3aac6f79d9ba 100644 --- a/paimon-python/pypaimon/tests/ray_hdf5_test.py +++ b/paimon-python/pypaimon/tests/ray_hdf5_test.py @@ -46,3 +46,33 @@ def test_ray_hdf5_worker_uses_shared_file_transform(): assert args[0].path == "file:///tmp/source.h5" assert args[1] is transform assert args[4] == schema + + +def test_ray_hdf5_worker_requires_one_source_per_batch(): + worker = _TransformHdf5File( + transform=object(), + source_options={}, + target_schema=pa.schema([pa.field("value", pa.int64())]), + ) + + with pytest.raises(ValueError, match="requires one source per batch"): + list(worker(pa.table({"path": ["a.h5", "b.h5"]}))) + + +def test_ray_hdf5_worker_filters_empty_transform_batches(): + schema = pa.schema([pa.field("value", pa.int64(), nullable=False)]) + empty = pa.table({"value": pa.array([], type=pa.int64())}, schema=schema) + expected = pa.table({"value": [1]}, schema=schema) + worker = _TransformHdf5File( + transform=object(), + source_options={}, + target_schema=schema, + ) + + with patch( + "pypaimon.multimodal.hdf5._transform_hdf5_file", + return_value=iter([empty, expected, empty])): + actual = list(worker(pa.table({"path": ["file:///tmp/source.h5"]}))) + + assert len(actual) == 1 + assert actual[0].equals(expected) diff --git a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py index b23ed48a2378..84cbb4996e50 100644 --- a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py +++ b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py @@ -14,10 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import importlib.util +import json import subprocess import sys +from unittest.mock import MagicMock import numpy as np import pyarrow as pa @@ -154,6 +155,86 @@ def test_shared_transform_streams_complete_agilex_business_schema( ] +def _consume_frame_transform(root, path): + episodes = agilex.discover_episodes(root) + transform = agilex.RoboMindAgileXFrameTransform(episodes) + with h5py.File(path, "r") as h5: + return list(transform(h5, pmm.Hdf5File(path=path.as_uri()))) + + +def test_discover_rejects_duplicate_episode_ids(tmp_path): + _write_episode(tmp_path / "task-a", "train", "duplicate", 0) + _write_episode(tmp_path / "task-b", "val", "duplicate", 10) + + with pytest.raises(ValueError, match="Duplicate RoboMIND episode_id"): + agilex.discover_episodes(tmp_path) + + +def test_discover_rejects_trajectory_symlink_outside_root(tmp_path): + target = _write_episode(tmp_path / "outside", "train", "target", 0) + root = tmp_path / "input" + link = (root / "success_episodes" / "train" / "linked" + / "data" / "trajectory.hdf5") + link.parent.mkdir(parents=True) + link.symlink_to(target) + + with pytest.raises(ValueError, match="escapes input root"): + agilex.discover_episodes(root) + + +@pytest.mark.parametrize("invalid_value", [np.nan, np.inf]) +def test_frame_transform_rejects_non_finite_numeric_values( + tmp_path, invalid_value): + path = _write_episode(tmp_path, "train", "invalid-action", 0) + with h5py.File(path, "r+") as h5: + h5["master/joint_position_left"][0, 0] = invalid_value + + with pytest.raises(ValueError, match="contains NaN or Inf"): + _consume_frame_transform(tmp_path / "13_packbowl", path) + + +def test_transform_rejects_invalid_numeric_dtype(tmp_path): + path = _write_episode(tmp_path, "train", "invalid-dtype", 0) + numeric_path = _NUMERIC_PATHS[0] + with h5py.File(path, "r+") as h5: + values = h5[numeric_path][...].astype(np.float32) + del h5[numeric_path] + h5.create_dataset(numeric_path, data=values) + + with pytest.raises(ValueError, match="invalid /.+ dtype"): + _consume_frame_transform(tmp_path / "13_packbowl", path) + + +def test_transform_rejects_missing_camera(tmp_path): + path = _write_episode(tmp_path, "train", "missing-camera", 0) + with h5py.File(path, "r+") as h5: + del h5[_IMAGE_PATHS[0]] + + with pytest.raises(ValueError, match="invalid /observations/rgb_images"): + _consume_frame_transform(tmp_path / "13_packbowl", path) + + +def test_transform_rejects_different_frame_lengths(tmp_path): + path = _write_episode(tmp_path, "train", "length-mismatch", 0) + image_path = _IMAGE_PATHS[0] + with h5py.File(path, "r+") as h5: + del h5[image_path] + h5.create_dataset(image_path, (2,), dtype=h5py.vlen_dtype(np.uint8)) + + with pytest.raises(ValueError, match="frame lengths differ"): + _consume_frame_transform(tmp_path / "13_packbowl", path) + + +def test_transform_rejects_empty_episode_but_accepts_one_frame(tmp_path): + empty = _write_episode(tmp_path, "train", "empty", 0, frames=0) + with pytest.raises(ValueError, match="episode has no frames"): + _consume_frame_transform(tmp_path / "13_packbowl", empty) + + one = _write_episode(tmp_path, "train", "one", 1, frames=1) + batches = _consume_frame_transform(tmp_path / "13_packbowl", one) + assert [batch.num_rows for batch in batches] == [1] + + @requires_vortex def test_local_ingest_and_backfill_materialize_only_canonical_action( agilex_input, tmp_path): @@ -208,6 +289,12 @@ def test_local_ingest_and_backfill_materialize_only_canonical_action( assert "file.format" not in agilex.TABLE_OPTIONS assert "action" in frames.raw_table.field_names assert "act_action_normalized" not in frames.raw_table.field_names + assert frame_rows.num_rows == 12 + frame_keys = list(zip( + frame_rows["episode_id"].to_pylist(), + frame_rows["frame_index"].to_pylist(), + )) + assert len(frame_keys) == len(set(frame_keys)) == 12 ordered = frame_rows.select([ "episode_id", @@ -232,8 +319,15 @@ def test_local_ingest_and_backfill_materialize_only_canonical_action( assert stats_row["source_split"] == "train" assert stats_row["frame_count"] == 6 assert stats_row["standard_deviation_floor"] == 0.01 - assert len(stats_row["action_mean"]) == 14 - assert len(stats_row["action_std"]) == 14 + expected_mean = np.concatenate([ + np.arange(1212, 1219), + np.arange(1312, 1319), + ]) + expected_std = np.full(14, np.sqrt(173.0 / 3.0)) + np.testing.assert_allclose( + stats_row["action_mean"], expected_mean, rtol=0, atol=1e-12) + np.testing.assert_allclose( + stats_row["action_std"], expected_std, rtol=1e-12, atol=1e-12) refreshed_snapshot = agilex.refresh_action_statistics( warehouse, statistics_version="synthetic-actions-refresh@1") @@ -258,6 +352,8 @@ def test_ray_ingest_matches_local_schema_rows_and_backfill( root, ray_warehouse, batch_size=2, concurrency=2) finally: ray.shutdown() + assert ray_result.episodes_snapshot_id == 1 + assert ray_result.frames_snapshot_id == 1 ray_backfill = agilex.backfill_canonical_action( ray_warehouse, statistics_version="synthetic-actions@1") @@ -278,26 +374,57 @@ def test_ray_ingest_matches_local_schema_rows_and_backfill( _logical_rows(ray_warehouse, table_name, schema)) -def test_action_statistics_are_independent_of_ingest_row_order(): +def test_stream_action_statistics_are_stable_and_order_independent(monkeypatch): frame_count = 257 - base = np.linspace(-0.7, 1.3, frame_count * 14).reshape( - frame_count, 14) - values = base + np.sin(base * 17.0) * 0.003 + row = np.arange(frame_count, dtype=np.float64).reshape(-1, 1) + feature = np.arange(14, dtype=np.float64).reshape(1, -1) + values = ( + 1e6 + row * 0.25 + feature * 0.5 + + ((row % 7) - 3) * 0.125 + ).astype(np.float32) source = pa.table({ "episode_id": ["train-a"] * frame_count, - "frame_index": np.arange(frame_count, dtype=np.int32), - "action_joint_position_left": values[:, :7].tolist(), - "action_joint_position_right": values[:, 7:].tolist(), - "_ROW_ID": np.arange(frame_count, dtype=np.int64), + "action": pa.array(values.tolist(), type=pa.list_(pa.float32(), 14)), }) reversed_source = source.take( pa.array(np.arange(frame_count - 1, -1, -1), type=pa.int64())) - expected = agilex.build_action_statistics(source.append_column( - "action", pa.array(values.astype(np.float32).tolist(), - type=pa.list_(pa.float32(), 14))), ["train-a"]) - actual = agilex.build_action_statistics(reversed_source.append_column( - "action", pa.array(values[::-1].astype(np.float32).tolist(), - type=pa.list_(pa.float32(), 14))), ["train-a"]) + def statistics(batches): + monkeypatch.setattr( + agilex, "_iter_raw", lambda table, columns: iter(batches)) + return agilex._stream_action_statistics(object(), ["train-a"]) + + forward = statistics([source.slice(0, 100), source.slice(100)]) + reverse = statistics([ + reversed_source.slice(0, 57), + reversed_source.slice(57, 100), + reversed_source.slice(157), + ]) + expected_mean = values.astype(np.float64).mean(axis=0) + expected_std = values.astype(np.float64).std(axis=0) + + for actual in (forward, reverse): + assert actual["frame_count"] == frame_count + np.testing.assert_allclose( + actual["action_mean"], expected_mean, rtol=0, atol=1e-9) + np.testing.assert_allclose( + actual["action_std"], expected_std, rtol=0, atol=1e-9) + + +def test_canonical_action_update_skips_empty_planned_split(monkeypatch): + empty = pa.table({ + "action_joint_position_left": pa.array([], type=pa.list_(pa.float64(), 7)), + "action_joint_position_right": pa.array([], type=pa.list_(pa.float64(), 7)), + "_ROW_ID": pa.array([], type=pa.int64()), + }) + table = MagicMock() + builder = table.new_batch_write_builder.return_value + commit = builder.new_commit.return_value + monkeypatch.setattr( + agilex, "_iter_raw", lambda raw_table, columns: iter([empty])) + + assert agilex._update_canonical_action_batches(table) == 0 - assert actual == expected + builder.new_update.assert_not_called() + commit.commit.assert_called_once_with([]) + commit.close.assert_called_once_with() diff --git a/paimon-python/pypaimon/write/ray_datasink.py b/paimon-python/pypaimon/write/ray_datasink.py index e121bd2432b5..baf8d4959714 100644 --- a/paimon-python/pypaimon/write/ray_datasink.py +++ b/paimon-python/pypaimon/write/ray_datasink.py @@ -21,6 +21,7 @@ import logging import traceback +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional from ray.data.datasource.datasink import Datasink @@ -30,6 +31,8 @@ from ray.data._internal.execution.interfaces import TaskContext import pyarrow as pa +from pypaimon.write.commit_callback import CommitCallback + if TYPE_CHECKING: from pypaimon.table.table import Table from pypaimon.write.write_builder import WriteBuilder @@ -38,6 +41,31 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class PaimonWriteResult: + """Metadata reported by the exact coordinator commit.""" + + row_count: int + snapshot_id: int + + +class _SnapshotIdRecorder(CommitCallback): + + def __init__(self): + self.snapshot_id = None + + def call(self, context): + self.snapshot_id = context.snapshot.id + + +class _TaskCommitMessages(list): + """Commit messages plus the logical rows written by one Ray task.""" + + def __init__(self, messages=(), row_count=0): + super().__init__(messages) + self.row_count = row_count + + def _cast_binary_to_table_schema(table: pa.Table, target_schema: pa.Schema) -> pa.Table: """Cast binary to large_binary for BLOB fields. @@ -82,6 +110,7 @@ def __init__( self._postpone_bucket_plan = postpone_bucket_plan self._table_name = table.identifier.get_full_name() self._writer_builder: Optional["WriteBuilder"] = None + self.commit_result: Optional[PaimonWriteResult] = None def _is_overwrite(self) -> bool: return self.overwrite or self.static_partition is not None @@ -115,6 +144,7 @@ def write( ctx: TaskContext, ) -> List["CommitMessage"]: commit_messages_list: List["CommitMessage"] = [] + row_count = 0 table_write = None try: @@ -139,6 +169,7 @@ def write( if block_arrow.num_rows == 0: continue + row_count += block_arrow.num_rows block_arrow = _cast_binary_to_table_schema(block_arrow, target_pa_schema) @@ -149,7 +180,7 @@ def write( table_write.close() table_write = None - return commit_messages_list + return _TaskCommitMessages(commit_messages_list, row_count) except Exception: if table_write is not None: try: @@ -175,6 +206,15 @@ def _extract_write_returns(write_result: Any): "lists. Refusing to proceed to avoid silent data loss." ) + @staticmethod + def _extract_row_count(write_result: Any, write_returns) -> int: + if hasattr(write_result, "num_rows"): + return int(write_result.num_rows) + return sum( + int(getattr(messages, "row_count", 0)) + for messages in write_returns + ) + def on_write_complete( self, write_result: Any ): @@ -205,7 +245,15 @@ def on_write_complete( ) table_commit = self._writer_builder.new_commit() + recorder = _SnapshotIdRecorder() + table_commit.add_commit_callback(recorder) table_commit.commit(non_empty_messages) + if recorder.snapshot_id is not None: + self.commit_result = PaimonWriteResult( + row_count=self._extract_row_count( + write_result, write_returns), + snapshot_id=recorder.snapshot_id, + ) logger.info(f"Successfully committed write job for table {self._table_name}") except Exception as e: @@ -243,7 +291,7 @@ def write_paimon_dataset( ray_remote_args: Optional[Dict[str, Any]] = None, hash_fixed_precluster: str = "auto", postpone_bucket_planner=None, -) -> None: +) -> Optional[PaimonWriteResult]: """Write a Ray Dataset through the safe path for the table's bucket mode.""" from pypaimon.ray.shuffle import ( HASH_FIXED_PRECLUSTER_MAP_GROUPS, @@ -294,7 +342,7 @@ def write_paimon_dataset( overwrite or static_partition is not None ), ) - _write_postpone_primary_key_blocks( + return _write_postpone_primary_key_blocks( dataset, table, overwrite=overwrite, @@ -304,14 +352,13 @@ def write_paimon_dataset( bucket_extractor=PostponeFixedBucketRowKeyExtractor(table, plan), postpone_bucket_plan=plan, ) - return if ( hash_fixed_precluster == HASH_FIXED_PRECLUSTER_MAP_GROUPS and table.bucket_mode() == BucketMode.HASH_FIXED and getattr(table, "is_primary_key_table", False) ): - _write_primary_key_groups( + return _write_primary_key_groups( dataset, table, overwrite=overwrite, @@ -319,18 +366,19 @@ def write_paimon_dataset( concurrency=concurrency, ray_remote_args=ray_remote_args, ) - return dataset = maybe_apply_repartition(dataset, table, hash_fixed_precluster) + datasink = PaimonDatasink( + table, + overwrite=overwrite, + static_partition=static_partition, + ) dataset.write_datasink( - PaimonDatasink( - table, - overwrite=overwrite, - static_partition=static_partition, - ), + datasink, concurrency=concurrency, ray_remote_args=ray_remote_args, ) + return datasink.commit_result def _write_postpone_primary_key_blocks( @@ -343,7 +391,7 @@ def _write_postpone_primary_key_blocks( ray_remote_args: Optional[Dict[str, Any]], bucket_extractor, postpone_bucket_plan, -) -> None: +) -> Optional[PaimonWriteResult]: import pickle from pypaimon.ray.shuffle import ( @@ -405,7 +453,7 @@ def _write_block(batch: pa.Table) -> pa.Table: static_partition=static_partition, ) coordinator.on_write_start() - _consume_write_results( + return _consume_write_results( results, coordinator, message_col, error_col ) @@ -435,7 +483,7 @@ def _consume_write_results( coordinator, message_col, error_col=None, -) -> None: +) -> Optional[PaimonWriteResult]: import pickle write_returns = [] @@ -459,6 +507,7 @@ def _consume_write_results( ) ) coordinator.on_write_complete(write_returns) + return coordinator.commit_result except Exception as error: coordinator.on_write_failed(error) raise @@ -517,7 +566,7 @@ def _write_primary_key_groups( ray_remote_args: Optional[Dict[str, Any]], bucket_extractor=None, postpone_bucket_plan=None, -) -> None: +) -> Optional[PaimonWriteResult]: import pickle from pypaimon.ray.shuffle import ( @@ -577,6 +626,6 @@ def _write_group(group: pa.Table) -> pa.Table: static_partition=static_partition, ) coordinator.on_write_start() - _consume_write_results( + return _consume_write_results( messages, coordinator, message_col, error_col ) From 1f5da94846805aae733cb02e5561cf6d86c046b6 Mon Sep 17 00:00:00 2001 From: Yann Date: Fri, 28 Aug 2026 19:25:13 +0800 Subject: [PATCH 5/6] fix(python): pin RoboMIND vector format Keep the sample's recommended table contract stable when multimodal facade defaults change on the PR merge base. Co-Authored-By: Codex AI-Model: gpt-5 Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 2/2 AI-Contributed/UT: 0/0 --- paimon-python/pypaimon/sample/robomind_agilex.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/paimon-python/pypaimon/sample/robomind_agilex.py b/paimon-python/pypaimon/sample/robomind_agilex.py index e3ccf993bfe5..ceaa7ed8e8fa 100644 --- a/paimon-python/pypaimon/sample/robomind_agilex.py +++ b/paimon-python/pypaimon/sample/robomind_agilex.py @@ -35,7 +35,9 @@ DEFAULT_STATISTICS_VERSION = "robomind-agilex-joint-position@1" TABLE_OPTIONS = { + "deletion-vectors.enabled": "true", "blob-as-descriptor": "false", + "vector.file.format": "vortex", } NUMERIC_FIELDS = ( From 09fc9636401511dcf81114fdef575df31beb820d Mon Sep 17 00:00:00 2001 From: Yann Date: Fri, 28 Aug 2026 21:18:02 +0800 Subject: [PATCH 6/6] fix(python): address RoboMIND review feedback Treat unofficial language datasets as optional, reuse one row-id file index across canonical-action batches, and allow compatible backfill retries. Document the Vortex runtime requirements. Co-Authored-By: Codex AI-Model: gpt-5 Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 123/123 AI-Contributed/UT: 124/124 --- docs/docs/pypaimon/robomind-agilex.md | 17 ++-- .../pypaimon/sample/robomind_agilex.py | 77 ++++++++++------ .../tests/robomind_agilex_pipeline_test.py | 90 +++++++++++++++++++ .../pypaimon/tests/table_update_test.py | 34 +++++++ paimon-python/pypaimon/write/table_update.py | 29 +++++- 5 files changed, 215 insertions(+), 32 deletions(-) diff --git a/docs/docs/pypaimon/robomind-agilex.md b/docs/docs/pypaimon/robomind-agilex.md index a1417839c549..e8aa3308d614 100644 --- a/docs/docs/pypaimon/robomind-agilex.md +++ b/docs/docs/pypaimon/robomind-agilex.md @@ -45,11 +45,11 @@ normalization statistics. ## Run the local pipeline -After downloading RoboMIND, install the HDF5 extra and provide the source and -warehouse directories to one command: +After downloading RoboMIND, use Python 3.11 or later, install the HDF5 and +Vortex extras, and provide the source and warehouse directories to one command: ```bash -pip install 'pypaimon[hdf5]' +pip install 'pypaimon[hdf5,vortex]' python -m pypaimon.sample.robomind_agilex \ --input /data/RoboMIND/h5_agilex_3rgb \ --warehouse /data/warehouse @@ -61,9 +61,14 @@ writes versioned train-split normalization statistics. It prints a JSON result with row counts and committed snapshot IDs. The input must already be present locally; the command does not download RoboMIND or contact Hugging Face. -Use a new warehouse for each run. Ingestion is append-only, so repeating the -same input against existing tables would create duplicate rows; a completed -warehouse also rejects a second canonical-action backfill. +Use a new warehouse for each ingestion run. Ingestion is append-only, so +repeating the same input against existing tables would create duplicate rows. +Canonical-action backfill is independently retryable after schema creation. + +The current Hugging Face example data includes `language_raw` and +`language_distilbert`, but these datasets are not part of the published AgileX +HDF5 schema. The episode transform therefore validates and stores them when +present, and writes null instruction metadata when they are absent. Pytest generates several small HDF5 episodes with the real AgileX field names, shapes, dtypes, split layout, and success layout, so the default test needs no diff --git a/paimon-python/pypaimon/sample/robomind_agilex.py b/paimon-python/pypaimon/sample/robomind_agilex.py index ceaa7ed8e8fa..aa56133e253e 100644 --- a/paimon-python/pypaimon/sample/robomind_agilex.py +++ b/paimon-python/pypaimon/sample/robomind_agilex.py @@ -460,23 +460,24 @@ def materialize_canonical_action(warehouse, *, database=DEFAULT_DATABASE): options={"warehouse": str(Path(warehouse).expanduser().resolve())}, ) frames_table = connection.get_table(FRAMES_TABLE) - _validate_backfill_target(frames_table) from pypaimon.schema.data_types import AtomicType, VectorType from pypaimon.schema.schema_change import SchemaChange - connection.catalog.alter_table( - frames_table.identifier, - [SchemaChange.add_column( - _ACTION_COLUMN, - VectorType(True, AtomicType("FLOAT"), 14), - comment=( - "Canonical AgileX action: master joint position left " - "followed by right."), - )], - False, - ) - frames_table = connection.get_table(FRAMES_TABLE) + action_type = VectorType(True, AtomicType("FLOAT"), 14) + if _validate_backfill_target(frames_table, action_type): + connection.catalog.alter_table( + frames_table.identifier, + [SchemaChange.add_column( + _ACTION_COLUMN, + action_type, + comment=( + "Canonical AgileX action: master joint position left " + "followed by right."), + )], + False, + ) + frames_table = connection.get_table(FRAMES_TABLE) row_count = _update_canonical_action_batches(frames_table.raw_table) frames_table = connection.get_table(FRAMES_TABLE) frames_snapshot_id = _snapshot_id(frames_table) @@ -549,35 +550,58 @@ def _create_tables(warehouse, database): return connection, episodes_table, frames_table -def _validate_backfill_target(frames_table): +def _validate_backfill_target(frames_table, action_type): missing = [ name for name in (_ACTION_LEFT, _ACTION_RIGHT) if name not in frames_table.raw_table.field_names ] if missing: raise ValueError("Frames table is missing raw action columns: %s." % missing) - if _ACTION_COLUMN in frames_table.raw_table.field_names: - raise ValueError("Canonical action column already exists.") + action_field = next( + (field for field in frames_table.raw_table.table_schema.fields + if field.name == _ACTION_COLUMN), + None, + ) + if action_field is None: + return True + if action_field.type != action_type: + raise ValueError( + "Canonical action column has incompatible type: %s." + % action_field.type) + return False def _update_canonical_action_batches(table): """Transform one planned Paimon split at a time and commit all updates once.""" builder = table.new_batch_write_builder() commit = builder.new_commit() - messages = [] row_count = 0 - try: + + def updates(): + nonlocal row_count for source in _iter_raw( table, [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"]): if source.num_rows == 0: continue - updates = build_canonical_action_backfill(source) - messages.extend( + update = build_canonical_action_backfill(source) + row_count += len(update) + yield update + + try: + update_batches = updates() + first = next(update_batches, None) + if first is None: + messages = [] + else: + def all_updates(): + yield first + yield from update_batches + + messages = ( builder.new_update() .with_update_type([_ACTION_COLUMN]) - .update_by_arrow_with_row_id(updates) + .update_by_arrow_batches_with_row_id(all_updates()) ) - row_count += len(updates) commit.commit(messages) finally: commit.close() @@ -673,7 +697,9 @@ def _validate_source(h5, source_key): def _instruction(h5, source_key): - if "language_raw" not in h5 or h5["language_raw"].shape != (1,): + if "language_raw" not in h5: + return None + if h5["language_raw"].shape != (1,): raise ValueError("%s: invalid /language_raw shape." % source_key) value = h5["language_raw"][0] if isinstance(value, bytes): @@ -684,8 +710,9 @@ def _instruction(h5, source_key): def _instruction_embedding(h5, source_key): - if ("language_distilbert" not in h5 - or h5["language_distilbert"].shape != (1, 1, 768)): + if "language_distilbert" not in h5: + return None + if h5["language_distilbert"].shape != (1, 1, 768): raise ValueError( "%s: invalid /language_distilbert shape." % source_key) values = np.asarray(h5["language_distilbert"][0, 0], dtype=np.float32) diff --git a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py index 84cbb4996e50..92013a9017c2 100644 --- a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py +++ b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py @@ -162,6 +162,23 @@ def _consume_frame_transform(root, path): return list(transform(h5, pmm.Hdf5File(path=path.as_uri()))) +def test_episode_transform_accepts_published_shape_without_language(tmp_path): + path = _write_episode(tmp_path, "train", "no-language", 0) + with h5py.File(path, "r+") as h5: + del h5["language_raw"] + del h5["language_distilbert"] + + episodes = agilex.discover_episodes(tmp_path / "13_packbowl") + transform = agilex.RoboMindAgileXEpisodeTransform(episodes) + with h5py.File(path, "r") as h5: + batches = list( + transform(h5, pmm.Hdf5File(path=path.as_uri()))) + + assert len(batches) == 1 + assert batches[0]["instruction"].to_pylist() == [None] + assert batches[0]["instruction_embedding"].to_pylist() == [None] + + def test_discover_rejects_duplicate_episode_ids(tmp_path): _write_episode(tmp_path / "task-a", "train", "duplicate", 0) _write_episode(tmp_path / "task-b", "val", "duplicate", 10) @@ -334,6 +351,41 @@ def test_local_ingest_and_backfill_materialize_only_canonical_action( assert refreshed_snapshot > backfill["statistics_snapshot_id"] +@requires_vortex +def test_canonical_action_backfill_resumes_after_schema_change( + agilex_input, tmp_path, monkeypatch): + root, _ = agilex_input + warehouse = tmp_path / "retry-warehouse" + agilex.ingest_local(root, warehouse) + + original_update = agilex._update_canonical_action_batches + + def fail_after_alter(table): + raise RuntimeError("injected update failure") + + monkeypatch.setattr( + agilex, "_update_canonical_action_batches", fail_after_alter) + with pytest.raises(RuntimeError, match="injected update failure"): + agilex.materialize_canonical_action(warehouse) + + monkeypatch.setattr( + agilex, "_update_canonical_action_batches", original_update) + row_count, snapshot_id = agilex.materialize_canonical_action(warehouse) + + frames, rows = _read( + warehouse, + agilex.FRAMES_TABLE, + ["action_joint_position_left", "action_joint_position_right", "action"], + ) + assert row_count == rows.num_rows == 12 + assert snapshot_id == agilex._snapshot_id(frames) + for row in rows.to_pylist(): + expected = np.asarray( + row["action_joint_position_left"] + + row["action_joint_position_right"], dtype=np.float32) + np.testing.assert_array_equal(row["action"], expected) + + @requires_vortex def test_ray_ingest_matches_local_schema_rows_and_backfill( agilex_input, tmp_path): @@ -428,3 +480,41 @@ def test_canonical_action_update_skips_empty_planned_split(monkeypatch): builder.new_update.assert_not_called() commit.commit.assert_called_once_with([]) commit.close.assert_called_once_with() + + +def test_canonical_action_update_reuses_one_row_id_updater(monkeypatch): + def source(row_id, offset): + return pa.table({ + "action_joint_position_left": pa.array( + [[offset + value for value in range(7)]], + type=pa.list_(pa.float64(), 7), + ), + "action_joint_position_right": pa.array( + [[offset + value for value in range(7, 14)]], + type=pa.list_(pa.float64(), 7), + ), + "_ROW_ID": pa.array([row_id], type=pa.int64()), + }) + + sources = [source(0, 0), source(1, 100)] + table = MagicMock() + builder = table.new_batch_write_builder.return_value + update = builder.new_update.return_value + update.with_update_type.return_value = update + captured = [] + + def update_batches(batches): + captured.extend(batches) + return ["message"] + + update.update_by_arrow_batches_with_row_id.side_effect = update_batches + monkeypatch.setattr( + agilex, "_iter_raw", lambda raw_table, columns: iter(sources)) + + assert agilex._update_canonical_action_batches(table) == 2 + + builder.new_update.assert_called_once_with() + assert len(captured) == 2 + assert captured[0]["_ROW_ID"].to_pylist() == [0] + assert captured[1]["_ROW_ID"].to_pylist() == [1] + builder.new_commit.return_value.commit.assert_called_once_with(["message"]) diff --git a/paimon-python/pypaimon/tests/table_update_test.py b/paimon-python/pypaimon/tests/table_update_test.py index 8c8a7d00b840..9b183be85d9a 100644 --- a/paimon-python/pypaimon/tests/table_update_test.py +++ b/paimon-python/pypaimon/tests/table_update_test.py @@ -31,12 +31,46 @@ DataEvolutionTestBase, StreamModeMixin, ) +from pypaimon.write.table_update import BatchTableUpdate # ====================================================================== # Shared base for batch & stream table-update tests # ====================================================================== + +def test_batch_row_id_update_batches_reuse_file_index(): + table = mock.MagicMock() + table.field_names = ["value"] + batches = [ + pa.table({"_ROW_ID": [0], "value": [10]}), + pa.table({"_ROW_ID": [1], "value": [20]}), + ] + + with mock.patch( + "pypaimon.write.table_update.TableUpdateByRowId") as factory: + updater = factory.return_value + updater.commit_messages = [] + + def update_columns(batch, columns): + updater.commit_messages.append((batch, columns)) + return updater.commit_messages + + updater.update_columns.side_effect = update_columns + messages = ( + BatchTableUpdate(table, "user") + .with_update_type(["value"]) + .update_by_arrow_batches_with_row_id(iter(batches)) + ) + + factory.assert_called_once() + assert updater.update_columns.call_count == 2 + assert messages == [ + (batches[0], ["value"]), + (batches[1], ["value"]), + ] + + class _TableUpdateTestBase(DataEvolutionTestBase): """Shared tests for ``TableUpdate.update_by_arrow_with_row_id``. diff --git a/paimon-python/pypaimon/write/table_update.py b/paimon-python/pypaimon/write/table_update.py index 175af0367ed7..a9404b32fecb 100644 --- a/paimon-python/pypaimon/write/table_update.py +++ b/paimon-python/pypaimon/write/table_update.py @@ -16,7 +16,7 @@ # under the License. from collections import defaultdict -from typing import Any, List, Mapping, Optional, Sequence, Tuple +from typing import Any, Iterable, List, Mapping, Optional, Sequence, Tuple import pyarrow import pyarrow as pa @@ -166,6 +166,26 @@ def _update_by_arrow_with_row_id( self.table, self.commit_user, commit_identifier, ).update_columns(table, cols) + def _update_by_arrow_batches_with_row_id( + self, tables: Iterable[pa.Table], commit_identifier: int + ) -> List[CommitMessage]: + updater = None + try: + for table in tables: + cols = self.update_cols if self.update_cols is not None else [ + c for c in table.column_names + if c != SpecialFields.ROW_ID.name + ] + if updater is None: + updater = TableUpdateByRowId( + self.table, self.commit_user, commit_identifier) + updater.update_columns(table, cols) + return [] if updater is None else updater.commit_messages + except Exception: + if updater is not None: + _abort_commit_messages(self.table, updater.commit_messages) + raise + def _upsert_by_arrow_with_key( self, table: pa.Table, @@ -624,6 +644,13 @@ def update_by_arrow_with_row_id(self, table: pa.Table) -> List[CommitMessage]: """Apply column updates keyed by ``_ROW_ID`` to existing rows.""" return self._update_by_arrow_with_row_id(table, BATCH_COMMIT_IDENTIFIER) + def update_by_arrow_batches_with_row_id( + self, tables: Iterable[pa.Table] + ) -> List[CommitMessage]: + """Apply row-id updates from batches using one target-file index.""" + return self._update_by_arrow_batches_with_row_id( + tables, BATCH_COMMIT_IDENTIFIER) + def upsert_by_arrow_with_key( self, table: pa.Table, upsert_keys: List[str] ) -> List[CommitMessage]: