diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index 5b35fff4f83a..a7bdd6aa291f 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -141,6 +141,7 @@ 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. # 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..e8aa3308d614 --- /dev/null +++ b/docs/docs/pypaimon/robomind-agilex.md @@ -0,0 +1,134 @@ +--- +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, 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,vortex]' +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 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 +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..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] @@ -188,33 +192,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 +229,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..2871c3b88502 --- /dev/null +++ b/paimon-python/pypaimon/ray/hdf5.py @@ -0,0 +1,146 @@ +# 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): + """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. 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.") + + 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, + _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 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) + 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_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: + + 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/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 new file mode 100644 index 000000000000..aa56133e253e --- /dev/null +++ b/paimon-python/pypaimon/sample/robomind_agilex.py @@ -0,0 +1,797 @@ +# 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 = { + "deletion-vectors.enabled": "true", + "blob-as-descriptor": "false", + "vector.file.format": "vortex", +} + +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: + 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( + 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=resolved_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] = pa.FixedSizeListArray.from_arrays( + pa.array(values.reshape(-1), type=pa.float64()), 7) + 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) + _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] + episode_result = load_from_hdf5( + "%s.%s" % (database, EPISODES_TABLE), paths, catalog_options, + transform=RoboMindAgileXEpisodeTransform(episodes), + concurrency=concurrency, + ) + frame_result = load_from_hdf5( + "%s.%s" % (database, FRAMES_TABLE), paths, catalog_options, + transform=RoboMindAgileXFrameTransform( + episodes, batch_size=batch_size), + concurrency=concurrency, + ) + return IngestResult( + mode="ray", + 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: + 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 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) + + from pypaimon.schema.data_types import AtomicType, VectorType + from pypaimon.schema.schema_change import SchemaChange + + 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) + 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, 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) + 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() + row_count = 0 + + def updates(): + nonlocal row_count + for source in _iter_raw( + table, [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"]): + if source.num_rows == 0: + continue + 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_batches_with_row_id(all_updates()) + ) + commit.commit(messages) + finally: + commit.close() + return row_count + + +def _stream_action_statistics(table, train_episode_ids): + """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 + 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()) + if value in train_ids + ] + if not selected: + continue + action = np.asarray( + source[_ACTION_COLUMN].take(pa.array(selected)).to_pylist(), + dtype=np.float64, + ) + 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.") + variance = np.maximum(m2 / count, 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: + 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): + 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: + 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) + 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..3aac6f79d9ba --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_hdf5_test.py @@ -0,0 +1,78 @@ +# 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 + + +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 new file mode 100644 index 000000000000..92013a9017c2 --- /dev/null +++ b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py @@ -0,0 +1,520 @@ +# 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 importlib.util +import json +import subprocess +import sys +from unittest.mock import MagicMock + +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") + +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] + + +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 + + +@requires_vortex +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 _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_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) + + 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): + 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 + 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", + "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 + 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") + 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): + 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() + 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") + + 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_stream_action_statistics_are_stable_and_order_independent(monkeypatch): + frame_count = 257 + 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, + "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())) + + 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 + + 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/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 ) 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]: