This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new d53103a442 [python] Add RoboMIND AgileX HDF5 pipeline (#9445)
d53103a442 is described below
commit d53103a44286aabcb4e3be162cd5f16d43c66bdd
Author: Yann Byron <[email protected]>
AuthorDate: Sat Aug 29 21:04:09 2026 +0800
[python] Add RoboMIND AgileX HDF5 pipeline (#9445)
---
.github/workflows/paimon-python-checks.yml | 1 +
docs/docs/pypaimon/robomind-agilex.md | 134 ++++
paimon-python/conftest.py | 7 +
paimon-python/pypaimon/multimodal/hdf5.py | 75 +-
paimon-python/pypaimon/ray/__init__.py | 2 +
paimon-python/pypaimon/ray/hdf5.py | 146 ++++
paimon-python/pypaimon/ray/ray_paimon.py | 9 +-
paimon-python/pypaimon/sample/robomind_agilex.py | 797 +++++++++++++++++++++
paimon-python/pypaimon/tests/ray_hdf5_test.py | 78 ++
.../tests/robomind_agilex_pipeline_test.py | 520 ++++++++++++++
paimon-python/pypaimon/tests/table_update_test.py | 34 +
paimon-python/pypaimon/write/ray_datasink.py | 81 ++-
paimon-python/pypaimon/write/table_update.py | 29 +-
13 files changed, 1865 insertions(+), 48 deletions(-)
diff --git a/.github/workflows/paimon-python-checks.yml
b/.github/workflows/paimon-python-checks.yml
index 5b35fff4f8..a7bdd6aa29 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 0000000000..e8aa3308d6
--- /dev/null
+++ b/docs/docs/pypaimon/robomind-agilex.md
@@ -0,0 +1,134 @@
+---
+title: "RoboMIND AgileX"
+sidebar_position: 7
+---
+
+<!--
+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
+
+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 f6adcd31ca..4ea23aa02e 100644
--- a/paimon-python/conftest.py
+++ b/paimon-python/conftest.py
@@ -24,6 +24,13 @@ _native_plan_count = 0
_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 68389cf4e2..acd85fb5d6 100644
--- a/paimon-python/pypaimon/multimodal/hdf5.py
+++ b/paimon-python/pypaimon/multimodal/hdf5.py
@@ -72,10 +72,14 @@ class Hdf5File:
@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 63141ecd41..e29d845061 100644
--- a/paimon-python/pypaimon/ray/__init__.py
+++ b/paimon-python/pypaimon/ray/__init__.py
@@ -31,6 +31,7 @@ from pypaimon.ray.data_evolution_merge_transform import (
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 @@ __all__ = [
"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 0000000000..2871c3b885
--- /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 2a5bbf0a26..26968d899d 100644
--- a/paimon-python/pypaimon/ray/ray_paimon.py
+++ b/paimon-python/pypaimon/ray/ray_paimon.py
@@ -33,6 +33,7 @@ from pypaimon.common.predicate import Predicate
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 0000000000..aa56133e25
--- /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 0000000000..3aac6f79d9
--- /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 0000000000..92013a9017
--- /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
+
+
[email protected]
+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
+
+
[email protected]
+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)
+
+
[email protected]("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 8c8a7d00b8..9b183be85d 100644
--- a/paimon-python/pypaimon/tests/table_update_test.py
+++ b/paimon-python/pypaimon/tests/table_update_test.py
@@ -31,12 +31,46 @@ from pypaimon.tests.data_evolution_test_helpers import (
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 e121bd2432..baf8d49597 100644
--- a/paimon-python/pypaimon/write/ray_datasink.py
+++ b/paimon-python/pypaimon/write/ray_datasink.py
@@ -21,6 +21,7 @@ Module to write a Paimon table from a Ray Dataset, by using
the Ray Datasink API
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.block import BlockAccessor, Block
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 @@ if TYPE_CHECKING:
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 @@ class PaimonDatasink(_DatasinkBase):
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 @@ class PaimonDatasink(_DatasinkBase):
ctx: TaskContext,
) -> List["CommitMessage"]:
commit_messages_list: List["CommitMessage"] = []
+ row_count = 0
table_write = None
try:
@@ -139,6 +169,7 @@ class PaimonDatasink(_DatasinkBase):
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 @@ class PaimonDatasink(_DatasinkBase):
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 @@ class PaimonDatasink(_DatasinkBase):
"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 @@ class PaimonDatasink(_DatasinkBase):
)
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_postpone_primary_key_blocks(
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_primary_key_groups(
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 175af0367e..a9404b32fe 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 @@ class TableUpdate:
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 @@ class BatchTableUpdate(TableUpdate):
"""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]:
