JingsongLi commented on code in PR #9466: URL: https://github.com/apache/paimon/pull/9466#discussion_r3912018310
########## docs/docs/pypaimon/robomind-act-benchmark.md: ########## @@ -0,0 +1,181 @@ +--- +title: "RoboMIND ACT Storage Benchmark" +sidebar_position: 8 +--- + +<!-- +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 ACT Storage Benchmark + +This benchmark measures the same CPU LeRobot ACT training workload over an +original RoboMIND AgileX HDF5 dataset or an already ingested and +canonical-action-backfilled Paimon warehouse. Ingestion and backfill are outside +the timed scope. + +The backends run independently. A resolved experiment document preserves the +shared configuration, normalization, seed, episode selection, Paimon snapshot, +and logical window sequence. Result comparison verifies that contract before it +calculates performance ratios. + +## Install + Review Comment: [P2] State the Python 3.10 floor in the installation instructions The documented command succeeds on PyPaimon's supported Python 3.8/3.9, but every dependency in the new `act` extra is guarded by `python_version >= 3.10`. Pip consequently installs none of LeRobot, datasets, or Pillow there, and the benchmark then fails at import time even though installation appeared successful. LeRobot 0.4.4 itself requires Python 3.10. Please state Python 3.10+ beside this command and provide an early actionable version check (or make requesting the extra fail clearly on older interpreters). ########## paimon-python/pypaimon/benchmark/act/runner.py: ########## @@ -0,0 +1,691 @@ +# 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. + +"""Prepare and run RoboMIND ACT benchmarks over HDF5 or Paimon. + +Both adapters consume one immutable :class:`BenchmarkConfig`, one train-only +normalization object, and one explicit window plan. The runner resets the same +seed before constructing the same LeRobot ACT policy and AdamW trainer for each +backend. Each backend runs independently without attempting OS cache control +and writes its tensor fingerprint, loss trace, timing metrics, and Python +allocation metrics to one result JSON document. +Ingestion and canonical-action backfill are deliberately outside the benchmark. +""" + +import gc +import hashlib +import json +import os +import platform +import subprocess +import time +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np +import torch +import pypaimon.multimodal as pmm +from pypaimon.benchmark.act.hdf5 import ( + compute_normalization as compute_hdf5_normalization, + create_datasets as create_hdf5_datasets, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.harness import ( + BenchmarkConfig, + WindowPlan, + build_window_plan, + run_backend, +) +from pypaimon.benchmark.act.compare import canonical_sha256 +from pypaimon.benchmark.act.paimon import ( + create_datasets as create_paimon_datasets, + latest_snapshot_id, + statistics_row, +) +from pypaimon.sample import robomind_agilex as agilex + + +@dataclass(frozen=True) +class _BenchmarkEpisode: + path: Path + source_key: str + episode_id: str + split: str + success: bool + frame_count: int + + +def prepare_experiment( + input_root, + warehouse, + output_path, + *, + definition=None, + database=agilex.DEFAULT_DATABASE): + """Resolve a benchmark definition against matching HDF5 and Paimon data. + + Preparation is outside timed benchmark execution. It verifies source + identity and Paimon statistics, selects eligible train/validation episodes, + computes train-only normalization, and fixes every logical window index. + + Args: + input_root: RoboMIND AgileX HDF5 root used as the source of episode + files and raw normalization moments. + warehouse: Existing Paimon warehouse containing the matching ingested + and canonical-action-backfilled dataset. + output_path: Destination for the resolved experiment JSON document. + definition: Optional decoded experiment definition. The packaged + defaults are used when omitted. + database: Paimon database containing the RoboMIND tables. + + Returns: + The resolved, JSON-compatible experiment dictionary written to + ``output_path``. + """ + definition = load_experiment() if definition is None else definition + if definition.get("schema_version") != "act-benchmark-experiment@1": + raise ValueError("Unsupported ACT benchmark experiment schema.") + config = BenchmarkConfig(**definition["config"]) + statistics_version = definition["statistics_version"] + input_root = Path(input_root).expanduser().resolve() + warehouse = Path(warehouse).expanduser().resolve() + output_path = Path(output_path).expanduser().resolve() + + discovered = agilex.discover_episodes(input_root) + connection = pmm.connect( + database=database, options={"warehouse": str(warehouse)}) + source_episodes, source_sha256 = _validate_source_identity( + discovered, _episode_rows(connection)) + source_by_id = {episode.episode_id: episode for episode in source_episodes} + frames = connection.get_table(agilex.FRAMES_TABLE) + frames_snapshot_id = latest_snapshot_id(frames) + normalization, normalization_metadata = _shared_normalization( + source_episodes, + connection, + frames_snapshot_id, + statistics_version, + ) + del normalization + train_episode = _select_episode( + source_by_id, + split="train", + requested=definition.get("train_episode_id"), + action_horizon=config.action_horizon, + ) + validation_episode = _select_episode( + source_by_id, + split="val", + requested=definition.get("validation_episode_id"), + action_horizon=config.action_horizon, + ) + plan = build_window_plan( + train_episode.frame_count - config.action_horizon + 1, + validation_episode.frame_count - config.action_horizon + 1, + config, + ) + sequence_sha256 = _sample_sequence_sha256( + train_episode.episode_id, validation_episode.episode_id, plan) + episodes = sorted(({ + "episode_id": episode.episode_id, + "source_key": episode.source_key, + "split": episode.split, + "success": episode.success, + "frame_count": episode.frame_count, + } for episode in source_episodes), key=lambda item: item["episode_id"]) + experiment = { + "schema_version": "act-benchmark-experiment@1", + "benchmark_id": definition.get("benchmark_id", "robomind-act"), + "dataset": definition.get("dataset", "RoboMIND AgileX"), + "config": config.to_dict(), + "statistics_version": statistics_version, + "train_episode_id": train_episode.episode_id, + "validation_episode_id": validation_episode.episode_id, + "source": { + "sha256": source_sha256, + "episodes": episodes, + }, + "normalization": normalization_metadata, + "window_plan": { + **plan.to_dict(), + "sha256": plan.sha256, + "sample_sequence_sha256": sequence_sha256, + }, + "paimon": { + "database": database, + "frames_table": agilex.FRAMES_TABLE, + "frames_snapshot_id": frames_snapshot_id, + }, + } + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(experiment, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return experiment + + +def run_experiment( + backend, + experiment_path, + output_path, + *, + input_root=None, + warehouse=None, + policy_factory=None): + """Run one storage backend against a resolved ACT experiment. + + Args: + backend: Result label and dataset implementation, either ``hdf5`` or + ``paimon``. + experiment_path: Resolved JSON produced by :func:`prepare_experiment`. + output_path: Destination JSON result path. + input_root: Required only for the HDF5 backend. + warehouse: Required only for the Paimon backend. + policy_factory: Optional test hook returning ``(policy, metadata)``. + + Returns: + A JSON-compatible single-backend result containing the resolved + experiment, runtime environment, tensor fingerprint, per-round raw + metrics, and median/min/max summary. + """ + if backend not in ("hdf5", "paimon"): + raise ValueError("backend must be 'hdf5' or 'paimon'.") + experiment = load_experiment(experiment_path) + _validate_resolved_experiment(experiment) + config = BenchmarkConfig(**experiment["config"]) + plan = _window_plan_from_experiment(experiment) + normalization = { + name: np.asarray(value, dtype=np.float32) + for name, value in experiment["normalization"]["values"].items() + } + sequence_sha256 = experiment["window_plan"]["sample_sequence_sha256"] + if backend == "hdf5": + if input_root is None: + raise ValueError("input_root is required for the HDF5 backend.") + episodes = _hdf5_episodes_from_experiment(input_root, experiment) + by_id = {episode.episode_id: episode for episode in episodes} + train_episode = by_id[experiment["train_episode_id"]] + validation_episode = by_id[experiment["validation_episode_id"]] + + def dataset_factory(): + return create_hdf5_datasets( + train_episode, validation_episode, normalization, config) + + source = {"input_root": str(Path(input_root).expanduser().resolve())} + else: + if warehouse is None: + raise ValueError("warehouse is required for the Paimon backend.") + dataset_factory, source = _paimon_factory_from_experiment( + warehouse, experiment, normalization, config) + + started_at = _utc_now() + started = time.monotonic() + fingerprint = _tensor_fingerprint(dataset_factory(), plan) + runs = [] + for round_number in range(1, config.rounds + 1): + runs.append(run_backend( + backend, + round_number, + dataset_factory, + plan, + config, + sequence_sha256, + policy_factory=policy_factory, + )) + gc.collect() + result = { + "schema_version": "act-benchmark-result@1", + "benchmark_id": experiment["benchmark_id"], + "run_id": "%s-%s" % ( + started_at.replace(":", "").replace("-", ""), + uuid.uuid4().hex[:8], + ), + "status": "SUCCEEDED", + "backend": backend, + "experiment": experiment, + "experiment_sha256": canonical_sha256(experiment), + "source": source, + "tensor_fingerprint": fingerprint, + "model": runs[0]["model"], + "runs": runs, + "summary": _summarize(runs), + "environment": { Review Comment: [P2] Fingerprint the performance-relevant runtime before comparing results `compare_results()` treats equality of this object as proof that independently produced results are compatible, but it omits the installed PyPaimon build, NumPy/PyArrow/h5py/Pillow versions, CPU identity/count, and Torch intra/inter-op thread settings. Two same-architecture hosts or wheel environments can therefore hash identically while having materially different storage costs; for wheel installs `_git_head(...)` is also commonly `UNKNOWN` even though the packaged build exposes its identity via `build_info.full_version()`. The comparator will still emit HDF5/Paimon ratios. Please record and compare the installed build identity, performance-sensitive dependency versions, CPU/count and Torch thread settings, and reject an unknown source identity when no equivalent package identity is available. ########## paimon-python/pypaimon/benchmark/act/harness.py: ########## @@ -0,0 +1,573 @@ +# 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. + +"""Shared deterministic ACT model, trainer, and window plan for benchmarks.""" + +import gc +import hashlib +import json +import math +import random +import time +import tracemalloc +from dataclasses import asdict, dataclass +from io import BytesIO + +import numpy as np +import torch +import torch.nn.functional as functional +from PIL import Image +from torch.utils.data import default_collate + + +CAMERA_KEYS = ( + "observation.images.front", + "observation.images.left_wrist", + "observation.images.right_wrist", +) + + +@dataclass(frozen=True) +class BenchmarkConfig: + """Immutable model, sampling, training, and measurement parameters. + + Every backend reconstructs this configuration from the resolved experiment + so tensor shapes, optimizer behavior, random seeds, and metric boundaries + remain comparable. + """ + + seed: int = 20260825 + action_horizon: int = 32 + batch_size: int = 2 + optimizer_steps: int = 2 + image_height: int = 64 + image_width: int = 80 + learning_rate: float = 1e-4 + weight_decay: float = 1e-4 + warmup_batches: int = 1 + timed_batches: int = 4 + fetch_batches: int = 8 + rounds: int = 3 + + def __post_init__(self): + positive_ints = ( + "action_horizon", + "batch_size", + "optimizer_steps", + "image_height", + "image_width", + "warmup_batches", + "timed_batches", + "fetch_batches", + ) + for name in positive_ints: + value = getattr(self, name) + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value <= 0): + raise ValueError("%s must be a positive int." % name) + if isinstance(self.seed, bool) or not isinstance(self.seed, int): + raise ValueError("seed must be an int.") + if isinstance(self.rounds, bool) or not isinstance(self.rounds, int): + raise ValueError("rounds must be an int.") + if self.rounds < 3: + raise ValueError("rounds must be at least 3.") + if self.learning_rate <= 0: + raise ValueError("learning_rate must be positive.") + if self.weight_decay < 0: + raise ValueError("weight_decay must not be negative.") + + def to_dict(self): + return asdict(self) + + +@dataclass(frozen=True) +class WindowPlan: + """Logical dataset-window indices consumed by one experiment. + + Measurement indices cover warm-up and timed reads, train indices cover + fixed optimizer steps, and validation indices cover the final loss. These + are map-style dataset indices, not Paimon row IDs. ``sha256`` identifies + the exact plan across independent backend processes. + """ + + seed: int + measurement_indices: tuple + train_indices: tuple + validation_indices: tuple + + @property + def sha256(self): + payload = json.dumps( + self.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def to_dict(self): + return { + "seed": self.seed, + "measurement_indices": list(self.measurement_indices), + "train_indices": list(self.train_indices), + "validation_indices": list(self.validation_indices), + } + + +def build_window_plan(train_window_count, validation_window_count, config): + """Build deterministic measurement, training, and validation indices. + + Args: + train_window_count: Number of complete windows in the train dataset. + validation_window_count: Number of complete validation windows. + config: Shared benchmark configuration supplying counts and the seed. + + Returns: + A :class:`WindowPlan`. When more samples are needed than a dataset + contains, consecutive seeded permutations are concatenated; sampling + does not become independent sampling with replacement. + """ + train_window_count = _positive_int( + train_window_count, "train_window_count") + validation_window_count = _positive_int( + validation_window_count, "validation_window_count") + batch_fetch_count = ( + config.warmup_batches + config.timed_batches) * config.batch_size + train_count = config.optimizer_steps * config.batch_size + return WindowPlan( + seed=config.seed, + measurement_indices=tuple(_repeat_permutations( + train_window_count, batch_fetch_count, config.seed + 1)), + train_indices=tuple(_repeat_permutations( + train_window_count, train_count, config.seed + 2)), + validation_indices=tuple(_repeat_permutations( + validation_window_count, config.batch_size, config.seed + 3)), + ) + + +def decode_rgb_image(payload): + """Decode JPEG/PNG bytes into an ``H x W x 3`` RGB NumPy array. + + Raises: + ValueError: If Pillow cannot decode the payload as an image. + """ + try: + return np.asarray(Image.open(BytesIO(payload)).convert("RGB")) + except Exception as error: + raise ValueError("Cannot decode ACT RGB image bytes.") from error + + +def decode_image_tensor(value): + """Decode bytes or an HDF5 uint8 value into normalized ``C x H x W``. + + The returned NumPy array is float32 with values in ``[0, 1]``. Both + storage backends call this function so image conversion is not part of the + performance difference being measured. + """ + payload = ( + bytes(value) + if isinstance(value, (bytes, bytearray, memoryview)) + else np.asarray(value, dtype=np.uint8).tobytes() + ) + image = decode_rgb_image(payload) + return np.transpose(image, (2, 0, 1)).astype(np.float32) / 255.0 + + +def validate_act_batch(batch, config): + """Validate a collated batch against the shared ACT tensor contract. + + Successful validation returns ``None``. It checks exact fields, tensor + shapes and dtypes, finite values, image range, complete unpadded windows, + and ``sample_id == episode_id#step_idx`` identity. + """ + required = { + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + } + if set(batch) != required: + raise ValueError( + "ACT batch fields differ: expected %s, got %s." + % (sorted(required), sorted(batch))) + batch_size = len(batch["sample_id"]) + expected = { + "qpos": ((batch_size, 14), torch.float32), + "action": ((batch_size, config.action_horizon, 14), torch.float32), + "images": ( + (batch_size, len(CAMERA_KEYS), 3) + + tuple(batch["images"].shape[-2:]), + torch.float32, + ), + "is_pad": ((batch_size, config.action_horizon), torch.bool), + "step_idx": ((batch_size,), torch.int64), + } + for name, (shape, dtype) in expected.items(): + value = batch[name] + if not isinstance(value, torch.Tensor): + raise ValueError("%s must be a torch.Tensor." % name) + if tuple(value.shape) != shape: + raise ValueError( + "%s has shape %s; expected %s." + % (name, tuple(value.shape), shape)) + if value.dtype != dtype: + raise ValueError( + "%s has dtype %s; expected %s." % (name, value.dtype, dtype)) + for name in ("qpos", "action", "images"): + if not torch.isfinite(batch[name]).all(): + raise ValueError("%s contains NaN or Inf." % name) + if torch.any(batch["images"] < 0) or torch.any(batch["images"] > 1): + raise ValueError("images must be normalized to [0, 1].") + if batch["is_pad"].any(): + raise ValueError("ACT benchmark windows must be complete and unpadded.") + for sample_id, episode_id, step_idx in zip( + batch["sample_id"], batch["episode_id"], + batch["step_idx"].tolist()): + if sample_id != "%s#%s" % (episode_id, step_idx): + raise ValueError( + "sample_id is not aligned with episode_id and step_idx.") + + +def build_lerobot_batch(batch, config): + """Map a shared ACT batch to LeRobot ``ACTPolicy`` feature names. + + Images are resized bilinearly to the configured height and width when + necessary. State, action, and padding retain their original semantics. + """ + validate_act_batch(batch, config) + images = batch["images"] + target_size = (config.image_height, config.image_width) + if tuple(images.shape[-2:]) != target_size: + flat = images.flatten(0, 1) + flat = functional.interpolate( + flat, size=target_size, mode="bilinear", align_corners=False) + images = flat.reshape(images.shape[:3] + target_size) + result = { + "observation.state": batch["qpos"], + "action": batch["action"], + "action_is_pad": batch["is_pad"], + } + for index, name in enumerate(CAMERA_KEYS): + result[name] = images[:, index] + return result + + +def build_act_policy(config): + """Build the reduced CPU ACT policy used only by this benchmark. + + Returns: + ``(policy, metadata)`` containing the LeRobot policy and a + JSON-compatible description of its architecture and parameter counts. + Pretrained weights are disabled, so this function performs no model + download and does not represent a production training configuration. + """ + try: + import importlib.metadata + from lerobot.configs.types import FeatureType, PolicyFeature + from lerobot.policies.act.configuration_act import ACTConfig + from lerobot.policies.act.modeling_act import ACTPolicy + except ImportError as error: + raise ImportError( + "ACT benchmark requires: " + "pip install -e '.[act]'.") from error + + inputs = { + "observation.state": PolicyFeature(FeatureType.STATE, (14,)), + } + inputs.update({ + name: PolicyFeature( + FeatureType.VISUAL, + (3, config.image_height, config.image_width), + ) + for name in CAMERA_KEYS + }) + act_config = ACTConfig( + input_features=inputs, + output_features={ + "action": PolicyFeature(FeatureType.ACTION, (14,)), + }, + device="cpu", + chunk_size=config.action_horizon, + n_action_steps=config.action_horizon, + vision_backbone="resnet18", + pretrained_backbone_weights=None, + dim_model=64, + n_heads=4, + dim_feedforward=256, + n_encoder_layers=1, + n_decoder_layers=1, + use_vae=True, + latent_dim=16, + n_vae_encoder_layers=1, + kl_weight=10.0, + ) + policy = ACTPolicy(act_config) + return policy, { + "implementation": "lerobot.ACTPolicy", + "lerobot_version": importlib.metadata.version("lerobot"), + "vision_backbone": act_config.vision_backbone, + "pretrained_backbone_weights": act_config.pretrained_backbone_weights, + "chunk_size": act_config.chunk_size, + "dim_model": act_config.dim_model, + "n_heads": act_config.n_heads, + "n_encoder_layers": act_config.n_encoder_layers, + "n_decoder_layers": act_config.n_decoder_layers, + "n_vae_encoder_layers": act_config.n_vae_encoder_layers, + "latent_dim": act_config.latent_dim, + "kl_weight": act_config.kl_weight, + "parameter_count": sum( + parameter.numel() for parameter in policy.parameters()), + "trainable_parameter_count": sum( + parameter.numel() + for parameter in policy.parameters() if parameter.requires_grad), + } + + +def run_backend( + backend, + round_number, + dataset_factory, + plan, + config, + sample_sequence_sha256, + policy_factory=None): + """Measure one backend with the shared plan, model, and trainer. + + ``backend`` is a result label and ``round_number`` identifies the repeat. + ``dataset_factory`` must return ``(train_dataset, validation_dataset)`` and + must be reusable: it is called for the timed run and again by the separate + Python-memory replay. ``policy_factory`` is an optional test hook returning + ``(policy, model_metadata)``. + + Returns: + A JSON-compatible metrics dictionary covering dataset construction, + first batch, timed batch fetch, fixed optimizer steps, validation loss, + and a separate ``tracemalloc`` peak replay. OS page cache is not + controlled and native Arrow/Torch allocations are outside tracemalloc. + """ + _seed_everything(config.seed) + policy_factory = policy_factory or build_act_policy + started = time.monotonic() + dataset_started = time.monotonic() + train_dataset, validation_dataset = dataset_factory() + dataset_build_s = time.monotonic() - dataset_started + + warmup_sample_count = config.warmup_batches * config.batch_size + warmup_iterator = _iter_logical_batches( + train_dataset, + plan.measurement_indices[:warmup_sample_count], + logical_batch_size=config.batch_size, + fetch_batches=1, + ) + first_batch_started = time.monotonic() + first_batch = next(warmup_iterator) + first_batch_s = time.monotonic() - first_batch_started + validate_act_batch(first_batch, config) + for _ in range(config.warmup_batches - 1): + validate_act_batch(next(warmup_iterator), config) + + batch_fetch_iterator = _iter_logical_batches( + train_dataset, + plan.measurement_indices[warmup_sample_count:], + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + ) + batch_fetch_started = time.monotonic() Review Comment: [P2] Exclude tensor validation from the fetch-throughput timer `batch_fetch_s` currently includes `validate_act_batch(batch)`, which scans every image/state/action tensor for finiteness and range. That work is independent of fetching and can dominate full-resolution image batches, so the reported `batch_fetch_samples_per_s` measures fetch plus tensor validation and compresses the backend difference. It is also inconsistent with `first_batch_s`, whose timer stops before validation at lines 371-374. Please accumulate only the time spent advancing `batch_fetch_iterator` (then validate outside that interval), or rename/document the metric as fetch-plus-validation. ########## paimon-python/pypaimon/multimodal/window_dataset.py: ########## @@ -0,0 +1,483 @@ +# 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. + +"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" + +import copy +import operator +from collections import defaultdict +from numbers import Integral + +import torch +from torch.utils.data import Dataset + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.query import ScanQuery +from pypaimon.schema.data_types import is_blob_type +from pypaimon.snapshot.time_travel_util import SCAN_KEYS +from pypaimon.table.special_fields import SpecialFields + + +class ContiguousWindowDataset(Dataset): + """Map-style Dataset which reads fixed row windows on demand. + + The in-memory index contains only group values, order values, and Paimon + row IDs. Each ``__getitem__`` reads the projected rows from the snapshot + resolved while the index was built. Within each group, ``order_key`` must + contain non-null integers that increase by exactly one; rows from different + groups never share a window. ``tail`` controls scheduled anchors whose + remaining rows are shorter than ``window_size``: + + * ``drop`` omits them; + * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``error`` rejects the dataset. + + The raw result mapping contains scalar group and order values, a + length-``window_size`` Boolean ``is_pad`` tensor, one-element lists for + ``anchor_columns``, and length-``window_size`` lists for other projected + columns. ``anchor_columns`` therefore avoids loading repeated context such + as observation images or initial robot state. ``column_transforms`` then + convert individual column lists before ``adapter`` adapts the complete + mapping to a model-specific contract. + ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. + """ + + _TAIL_POLICIES = ("drop", "pad", "error") + + def __init__( + self, + query, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + if getattr(query, "_result_factory", None) is not None: + raise TypeError( + "ContiguousWindowDataset is only supported on scan(), " + "not search queries.") + self.window_size = _positive_int(window_size, "window_size") + self.stride = _positive_int(stride, "stride") + if tail not in self._TAIL_POLICIES: + raise ValueError( + "tail must be one of %s; got %r." + % (self._TAIL_POLICIES, tail)) + self.tail = tail + self.group_key = _column(query, group_key, "group_key") + self.order_key = _column(query, order_key, "order_key") + if self.group_key == self.order_key: + raise ValueError("group_key and order_key must name different columns.") + if "is_pad" in (self.group_key, self.order_key): + raise ValueError("group_key and order_key must not be is_pad.") + self.columns = _columns( + query, columns, self.group_key, self.order_key) + self.anchor_columns = _anchor_columns(anchor_columns, self.columns) + anchor_column_set = set(self.anchor_columns) + self._window_columns = [ + name for name in self.columns if name not in anchor_column_set + ] + self.column_transforms = _column_transforms( + column_transforms, self.columns) + self.pad_values = _pad_values(pad_values, self.columns) + if adapter is not None and not callable(adapter): + raise TypeError("adapter must be callable or None.") + self.adapter = adapter + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + + if not query._table.options.row_tracking_enabled(): + raise ValueError( + "ContiguousWindowDataset requires row-tracking.enabled=true.") + + index, snapshot_id = _read_window_index( + query, self.group_key, self.order_key) + self.snapshot_id = snapshot_id + self._table = _pin_table(query._table, snapshot_id) + self._groups, self._anchors = self._build_index(index) + + @classmethod + def from_query(cls, query, **kwargs): + """Build a contiguous-window Dataset from a ``ScanQuery``.""" + return cls(query, **kwargs) + + def __len__(self): + return len(self._anchors) + + def __getitem__(self, index): + """Read one window by map-style Dataset index. + + Negative indices follow Python sequence semantics. The return value is + the pre-adapter mapping described by the class, or the adapter result + when an adapter is configured. + """ + anchor, row_ids = self._resolve_window(index) + rows = self._read_window_rows(row_ids) + anchor_row = ( + self._read_rows(row_ids[:1], self.anchor_columns)[0] + if self.anchor_columns else None + ) + return self._sample(anchor, rows, anchor_row) + + def __getitems__(self, indices): + """Read several Dataset indices while coalescing overlapping row IDs. + + The returned list preserves the requested index order and duplicates. + Coalescing affects only physical reads, not logical sample cardinality. + """ + windows = [self._resolve_window(index) for index in indices] + if not windows: + return [] + row_ids = list(dict.fromkeys( + row_id for _, window_row_ids in windows + for row_id in window_row_ids + )) + rows_by_id = dict(zip(row_ids, self._read_window_rows(row_ids))) + anchor_row_ids = list(dict.fromkeys( + window_row_ids[0] for _, window_row_ids in windows + )) + anchor_rows_by_id = ( + dict(zip( + anchor_row_ids, + self._read_rows(anchor_row_ids, self.anchor_columns), + )) + if self.anchor_columns else {} + ) + return [ + self._sample( + anchor, + [rows_by_id[row_id] for row_id in window_row_ids], + anchor_rows_by_id.get(window_row_ids[0]), + ) + for anchor, window_row_ids in windows + ] + + def _resolve_window(self, index): + index = operator.index(index) + if index < 0: + index += len(self._anchors) + if index < 0 or index >= len(self._anchors): + raise IndexError("window index out of range") + + anchor = self._anchors[index] + group_index, start, valid_count = anchor + row_ids = self._groups[group_index][2] + return anchor, row_ids[start:start + valid_count] + + def _sample(self, anchor, rows, anchor_row=None): + group_index, start, valid_count = anchor + group_key, order_values, _ = self._groups[group_index] + padding_count = self.window_size - valid_count + padding_mask = torch.zeros(self.window_size, dtype=torch.bool) + if padding_count: + padding_mask[valid_count:] = True + sample = { + self.group_key: group_key, + self.order_key: order_values[start], + "is_pad": padding_mask, + } + for name in self.columns: + if name in self.anchor_columns: + values = [anchor_row[name]] + else: + values = [row[name] for row in rows] + if padding_count and name not in self.anchor_columns: + pad_value = self.pad_values.get(name, values[-1]) + values.extend( + copy.deepcopy(pad_value) for _ in range(padding_count)) + transform = self.column_transforms.get(name) + sample[name] = transform(values) if transform is not None else values + if self.adapter is not None: + return self.adapter(sample) + return sample + + def _build_index(self, index): + """Validate index rows and return grouped row IDs plus window anchors. + + Args: + index: Arrow table containing ``group_key``, ``order_key``, and + Paimon's ``_ROW_ID`` for the resolved snapshot. + + Returns: + ``(groups, anchors)``. Each group stores its key, ordered positions, + and row IDs. Each anchor stores group index, start offset, and the + number of real rows available before optional padding. + """ + group_values = index.column(self.group_key).to_pylist() + order_values = index.column(self.order_key).to_pylist() + row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist() + grouped = defaultdict(list) + for group_key, order_value, row_id in zip( + group_values, order_values, row_ids): + if group_key is None: + raise ValueError("%s must not contain null values." % self.group_key) + if order_value is None: + raise ValueError("%s must not contain null values." % self.order_key) + if isinstance(order_value, bool) or not isinstance(order_value, Integral): + raise ValueError( + "%s must contain integer values." % self.order_key) + try: + grouped[group_key].append((int(order_value), int(row_id))) + except TypeError: + raise ValueError( + "%s values must be hashable." % self.group_key) + + groups = [] + anchors = [] + try: + sorted_groups = sorted(grouped.items(), key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values must be mutually orderable." % self.group_key) + for group_key, members in sorted_groups: + try: + members.sort(key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values in group %r must be mutually orderable." + % (self.order_key, group_key)) + for previous, current in zip(members, members[1:]): + if previous[0] == current[0]: + raise ValueError( + "Group %s has duplicate order value %r in %s." + % (group_key, current[0], self.order_key)) + if current[0] != previous[0] + 1: + raise ValueError( + "Group %s is not contiguous in %s: %s followed by %s." + % (group_key, self.order_key, + previous[0], current[0])) + + group_index = len(groups) + group_orders = [member[0] for member in members] + group_row_ids = [member[1] for member in members] + groups.append((group_key, group_orders, group_row_ids)) + for start in range(0, len(members), self.stride): + valid_count = min(self.window_size, len(members) - start) + if valid_count < self.window_size: + if self.tail == "drop": + continue + if self.tail == "error": + raise ValueError( + "Group %s has an incomplete window at %s: " + "window_size=%d, available=%d." + % (group_key, group_orders[start], + self.window_size, valid_count)) + anchors.append((group_index, start, valid_count)) + return groups, anchors + + def _read_window_rows(self, row_ids): + if not self._window_columns: + return [{} for _ in row_ids] + return self._read_rows(row_ids, self._window_columns) + + def _read_rows(self, row_ids, columns=None): + """Read projected rows by ID from the pinned snapshot. + + Args: + row_ids: Paimon row IDs to read. Their order and duplicates define + the returned row order. + columns: Projected value columns, or all Dataset columns when + omitted. + + Returns: + A list of row dictionaries aligned one-for-one with ``row_ids``. + The internal ``_ROW_ID`` field is removed, and BLOB descriptors are + resolved to their bodies. + """ + columns = self.columns if columns is None else columns + query = ScanQuery(self._table) + predicate_builder = ( + self._table.new_read_builder() + .with_projection( + [field.name for field in self._table.fields] + + [SpecialFields.ROW_ID.name]) + .new_predicate_builder() + ) + query._predicate = predicate_builder.is_in( + SpecialFields.ROW_ID.name, row_ids) + query._projection = list(columns) + query._include_row_id = True + + blob_columns = [ + field.name for field in self._table.fields + if field.name in columns and is_blob_type(field.type) + ] + if blob_columns: + scalar, blobs = query.read_blobs( + blob_columns, parallelism=self.blob_parallelism) + rows = scalar.to_pylist() + for name in blob_columns: + values = blobs[name] + if len(values) != len(rows): + raise RuntimeError( + "BLOB column %s is not row-aligned with a window read." + % name) + for row, value in zip(rows, values): + row[name] = value + else: + rows = query.to_arrow().to_pylist() + + by_row_id = {} + row_id_column = SpecialFields.ROW_ID.name + for row in rows: + row_id = int(row[row_id_column]) + del row[row_id_column] + by_row_id[row_id] = row + missing = [row_id for row_id in row_ids if row_id not in by_row_id] + if missing: + raise RuntimeError( + "Pinned snapshot %s did not return indexed row IDs %s." + % (self.snapshot_id, missing)) + return [by_row_id[row_id] for row_id in row_ids] + + +def _read_window_index(query, group_key, order_key): + index_query = copy.copy(query) + index_query._projection = [group_key, order_key] + index_query._include_row_id = True + read_builder = index_query._configured_read_builder() + plan = read_builder.new_scan().plan() + index = read_builder.new_read().to_arrow(plan.splits()) + if index.num_rows and plan.snapshot_id is None: + raise RuntimeError("Cannot pin the snapshot used to build the window index.") + return index, plan.snapshot_id + + +def _pin_table(table, snapshot_id): + """Pin a table copy to ``snapshot_id``, or reuse it when unresolved.""" + if snapshot_id is None: + return table + options = { + key: None for key in SCAN_KEYS Review Comment: [P2] Clear scan.mode when pinning the resolved snapshot `SCAN_KEYS` does not include `scan.mode` (nor the incremental/creation-time companion options), so this copy can retain an explicit startup mode and then add `scan.snapshot-id`. For example, a table with the valid `scan.mode=latest-full` builds the window index successfully, but its first `dataset[0]` read fails with `scan.mode 'latest-full' conflicts with: ['scan.snapshot-id']`; incremental modes have the same delayed conflict. Other pinning code such as `ray/join_common.py::pin_latest_snapshot` clears `scan.mode` and all of its companions before setting the snapshot. Please normalize these options here too, or reject scan modes whose semantics cannot safely be converted, and add a `latest-full` regression test. ########## paimon-python/pypaimon/multimodal/window_dataset.py: ########## @@ -0,0 +1,483 @@ +# 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. + +"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" + +import copy +import operator +from collections import defaultdict +from numbers import Integral + +import torch +from torch.utils.data import Dataset + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.query import ScanQuery +from pypaimon.schema.data_types import is_blob_type +from pypaimon.snapshot.time_travel_util import SCAN_KEYS +from pypaimon.table.special_fields import SpecialFields + + +class ContiguousWindowDataset(Dataset): + """Map-style Dataset which reads fixed row windows on demand. + + The in-memory index contains only group values, order values, and Paimon + row IDs. Each ``__getitem__`` reads the projected rows from the snapshot + resolved while the index was built. Within each group, ``order_key`` must + contain non-null integers that increase by exactly one; rows from different + groups never share a window. ``tail`` controls scheduled anchors whose + remaining rows are shorter than ``window_size``: + + * ``drop`` omits them; + * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``error`` rejects the dataset. + + The raw result mapping contains scalar group and order values, a + length-``window_size`` Boolean ``is_pad`` tensor, one-element lists for + ``anchor_columns``, and length-``window_size`` lists for other projected + columns. ``anchor_columns`` therefore avoids loading repeated context such + as observation images or initial robot state. ``column_transforms`` then + convert individual column lists before ``adapter`` adapts the complete + mapping to a model-specific contract. + ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. + """ + + _TAIL_POLICIES = ("drop", "pad", "error") + + def __init__( + self, + query, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + if getattr(query, "_result_factory", None) is not None: + raise TypeError( + "ContiguousWindowDataset is only supported on scan(), " + "not search queries.") + self.window_size = _positive_int(window_size, "window_size") + self.stride = _positive_int(stride, "stride") + if tail not in self._TAIL_POLICIES: + raise ValueError( + "tail must be one of %s; got %r." + % (self._TAIL_POLICIES, tail)) + self.tail = tail + self.group_key = _column(query, group_key, "group_key") + self.order_key = _column(query, order_key, "order_key") + if self.group_key == self.order_key: + raise ValueError("group_key and order_key must name different columns.") + if "is_pad" in (self.group_key, self.order_key): + raise ValueError("group_key and order_key must not be is_pad.") + self.columns = _columns( + query, columns, self.group_key, self.order_key) + self.anchor_columns = _anchor_columns(anchor_columns, self.columns) + anchor_column_set = set(self.anchor_columns) + self._window_columns = [ + name for name in self.columns if name not in anchor_column_set + ] + self.column_transforms = _column_transforms( + column_transforms, self.columns) + self.pad_values = _pad_values(pad_values, self.columns) + if adapter is not None and not callable(adapter): + raise TypeError("adapter must be callable or None.") + self.adapter = adapter + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + + if not query._table.options.row_tracking_enabled(): + raise ValueError( + "ContiguousWindowDataset requires row-tracking.enabled=true.") + + index, snapshot_id = _read_window_index( + query, self.group_key, self.order_key) + self.snapshot_id = snapshot_id + self._table = _pin_table(query._table, snapshot_id) + self._groups, self._anchors = self._build_index(index) + + @classmethod + def from_query(cls, query, **kwargs): + """Build a contiguous-window Dataset from a ``ScanQuery``.""" + return cls(query, **kwargs) + + def __len__(self): + return len(self._anchors) + + def __getitem__(self, index): + """Read one window by map-style Dataset index. + + Negative indices follow Python sequence semantics. The return value is + the pre-adapter mapping described by the class, or the adapter result + when an adapter is configured. + """ + anchor, row_ids = self._resolve_window(index) + rows = self._read_window_rows(row_ids) + anchor_row = ( + self._read_rows(row_ids[:1], self.anchor_columns)[0] + if self.anchor_columns else None + ) + return self._sample(anchor, rows, anchor_row) + + def __getitems__(self, indices): + """Read several Dataset indices while coalescing overlapping row IDs. + + The returned list preserves the requested index order and duplicates. + Coalescing affects only physical reads, not logical sample cardinality. + """ + windows = [self._resolve_window(index) for index in indices] + if not windows: + return [] + row_ids = list(dict.fromkeys( + row_id for _, window_row_ids in windows + for row_id in window_row_ids + )) + rows_by_id = dict(zip(row_ids, self._read_window_rows(row_ids))) + anchor_row_ids = list(dict.fromkeys( + window_row_ids[0] for _, window_row_ids in windows + )) + anchor_rows_by_id = ( + dict(zip( + anchor_row_ids, + self._read_rows(anchor_row_ids, self.anchor_columns), + )) Review Comment: [P2] Keep plural-window samples independent when rows overlap `rows_by_id` stores each physical row dict once, and overlapping/repeated windows then pass the same nested ARRAY/MAP objects into multiple calls to `_sample()`. A supported `column_transform` or `adapter` that mutates a value in place therefore changes another logical sample: with overlapping `[0,1]` windows and ARRAY cells, plural `__getitems__([0,1])` lets the first transform mutate the second sample's shared row, while two `__getitem__` calls remain independent. This violates the method's claim that coalescing changes only physical reads. Please copy mutable cells per logical sample before invoking transforms/adapters and cover overlapping plus duplicate indices. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
