From a79cddfef835d8b03fdef1ee120538794a6bbfe6 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 16 Sep 2026 14:25:21 +0800 Subject: [PATCH 1/2] feat(python): support per-column training windows in CWD Build training samples from one materialized frame table with per-column frame or time offsets, two-sided padding and per-column masks. Add tensor and image transforms, adapt the ACT benchmark, and preserve legacy windows. Co-Authored-By: Codex AI-Model: gpt-6 Co-Authored-By: Codex AI-Contributed/Feature: 643/643 AI-Contributed/UT: 365/365 --- docs/docs/pypaimon/multimodal-reading.md | 119 +++++-- docs/docs/pypaimon/pytorch.md | 55 ++-- .../pypaimon/benchmark/act/paimon.py | 11 +- .../pypaimon/multimodal/lerobot/dataset.py | 26 +- paimon-python/pypaimon/multimodal/query.py | 51 ++- .../pypaimon/multimodal/window_dataset.py | 292 ++++++++++++------ .../pypaimon/multimodal/window_transforms.py | 88 ++++++ .../pypaimon/tests/act_runner_test.py | 20 +- .../tests/contiguous_window_dataset_test.py | 236 ++++++++++++++ .../pypaimon/tests/window_transforms_test.py | 109 +++++++ paimon-python/setup.py | 1 + 11 files changed, 826 insertions(+), 182 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/window_transforms.py create mode 100644 paimon-python/pypaimon/tests/window_transforms_test.py diff --git a/docs/docs/pypaimon/multimodal-reading.md b/docs/docs/pypaimon/multimodal-reading.md index 234dd4f6e2ff..d3df05a9bddc 100644 --- a/docs/docs/pypaimon/multimodal-reading.md +++ b/docs/docs/pypaimon/multimodal-reading.md @@ -182,35 +182,37 @@ Notes: ### Contiguous windows for PyTorch -Install the `torch` extra, then use `to_contiguous_window_dataset` to expose -map-style windows without loading the selected rows or BLOB payloads into Python -memory up front. The Dataset builds a compact index from the group column, order -column, and Paimon row IDs. Each `__getitem__` call fetches only that window from -the snapshot recorded in `dataset.snapshot_id`. +Use `to_contiguous_window_dataset` to build map-style training samples from a +single frame table. Each row represents one time step; task text, labels, and +other context must already be materialized in that table. The Dataset does not +read companion tables. It builds a compact index of group/order values and +Paimon row IDs, then fetches requested columns and BLOB payloads on demand from +the snapshot recorded in `dataset.snapshot_id`. Row tracking must be enabled. ```shell pip install 'pypaimon[torch]' ``` ```python -import torch - +from functools import partial -def float32_window(values): - return torch.tensor(values, dtype=torch.float32) +import torch +from pypaimon.multimodal.window_transforms import images_to_tensor, to_tensor +float32_window = partial(to_tensor, dtype=torch.float32) windows = ( - frames.scan() + frames.scan(tag_name="train-v1") .where("split = 'train'") .to_contiguous_window_dataset( - window_size=16, - columns=["state", "action"], + columns=["state", "image", "action"], group_key="episode_index", order_key="frame_index", - tail="pad", + frame_offsets={"state": [-2, -1, 0], "action": list(range(16))}, + boundary="pad", column_transforms={ "state": float32_window, + "image": images_to_tensor, "action": float32_window, }, ) @@ -218,25 +220,86 @@ windows = ( sample = windows[0] assert sample["action"].shape == (16, action_size) -assert sample["is_pad"].shape == (16,) +assert sample["action_is_pad"].shape == (16,) +assert sample["state"].shape == (3, state_size) +assert sample["image"].shape == (1, channels, height, width) ``` -The group and order keys in a sample identify the window anchor. Every projected -column contains the whole window. With `tail="drop"`, only full windows are -exposed. With `tail="pad"`, every real row is an anchor; missing suffix values -repeat the last real value by default and `is_pad` is `True` exactly at those -positions. With `tail="error"`, construction fails if any scheduled anchor is -incomplete. Use `pad_values` to override the repeated value for individual -columns. Anchors advance by `stride`, which defaults to one row. +The group and order keys identify the anchor. `frame_offsets` defines each +column's relative frame positions, including history, future, and sparse +sampling. Offset order and duplicates are preserved. A selected column omitted +from the mapping uses `[0]`, so the image above is read only at the anchor. +Without any window definition, every selected column uses `[0]`. Each column +has its own Boolean `_is_pad` tensor, where `True` marks an out-of-group +position. These generated names must not collide with other output fields. + +Anchors start at each group's first row and advance by `stride` (default one). +The stride controls anchor spacing, not the spacing inside a window: + +| `boundary` | Behavior at either end of a group | +|------------|----------------------------------| +| `drop` (default) | Keep only anchors for which all requested offsets are valid. | +| `pad` | Keep scheduled anchors and repeat the nearest endpoint for missing positions. | +| `error` | Fail construction if any scheduled anchor has an incomplete window. | + +`pad_values={"action": [0.0, 0.0]}` overrides endpoint repetition for that +column. Supply values in the same raw representation as a table cell. Padding +happens before column transforms, so normalization also affects constant padding. +Use the mask when excluding padded actions from a loss. + +As an alternative, specify seconds with `delta_timestamps`, an explicit positive +`fps`, and optional `tolerance_s` (default `1e-4`). For example, +`delta_timestamps={"state": [-0.1, 0.0]}, fps=10` selects offsets `[-1, 0]`. +Both `fps` and `tolerance_s` are only accepted with `delta_timestamps`. +Offsets must align to the regular frame grid within the tolerance; this API +does not search actual timestamps or interpolate missing frames. Do not combine +`delta_timestamps`, `frame_offsets`, and `window_size`. `column_transforms` receive one padded Python list per projected column. This is -where applications define tensor dtype and shape or decode BLOB bytes. The -optional `adapter` receives the resulting sample mapping and can rename or -combine fields for a model-specific batch contract. The core Dataset does not -know model field names, image formats, or normalization rules. Top-level -functions and callable classes are recommended for transforms and adapters so -the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader` -instances. +where applications define tensor dtype and shape or decode BLOB bytes. Even +singleton windows retain the time dimension. `images_to_tensor` returns TCHW +images with EXIF orientation applied. Eight-bit pixels become float32 in +`[0, 1]`; use `partial(images_to_tensor, return_uint8=True)` to keep uint8. +Higher-bit-depth images become float32 in their original units. All images in +a column window must have the same shape. + +The optional `adapter` receives the transformed sample and can normalize values, +rename/combine fields, remove singleton time dimensions, or coordinate random +augmentation across cameras and time. Statistics must be supplied explicitly; +the Dataset never computes or discovers normalization statistics. For example: + +```python +class NormalizeAction: + def __init__(self, mean, std): + self.mean = torch.as_tensor(mean, dtype=torch.float32) + self.std = torch.as_tensor(std, dtype=torch.float32) + if not torch.isfinite(self.std).all() or not (self.std > 0).all(): + raise ValueError("Action standard deviations must be finite and positive.") + + def __call__(self, sample): + sample["action"] = (sample["action"] - self.mean) / self.std + return sample +``` + +Pass `adapter=NormalizeAction(train_action_mean, train_action_std)` when +constructing the Dataset. Use training-set statistics, with the same values for +validation. Keep callbacks picklable (top-level functions, `partial`, or callable +classes) for multi-worker DataLoader use. Physical reads are coalesced in +`__getitems__`, while samples and transforms retain their requested order and +duplicates. Negative indices and slices follow Python sequence semantics. + +Rows are sorted by group and order. Order values must be non-null integers +increasing by exactly one within each group. Duplicate steps and internal gaps, +including gaps introduced by a scan filter, fail validation. Windows never +cross groups. A snapshot pin fixes the data being read; use retained tags for +long-running training. Sampler progress and augmentation RNG state remain the +training application's responsibility. + +Existing `window_size=16` calls remain supported, with `anchor_columns` selecting +singleton fields and `tail="drop"/"pad"/"error"` controlling the old forward +window. These calls preserve their single `is_pad` mask and previous shapes. +`anchor_columns` and `tail` require `window_size`; do not combine `tail` and +`boundary`. New applications should use field offsets and per-column masks. Columns configured by `video-frame-field` are rejected: a window read would drop the `frame_index` and other metadata carried by their `VideoFrameDescriptor` diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 6f973ade3f17..26770bdbba53 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -159,42 +159,55 @@ physical video ranges and cache decoder sessions per worker. See for the write path and a complete decoder example. ## Contiguous Windows -Use a map-style `ContiguousWindowDataset` when training samples are fixed-size -windows which must not cross a sequence boundary. The dataset builds an index +Use a map-style `ContiguousWindowDataset` to construct training samples from +one materialized frame table. Each field can have its own history or future +window, and windows never cross a sequence boundary. The dataset builds an index from only the group column, order column, and Paimon row IDs. Projected values, including BLOB payloads, are read from the pinned snapshot when a sample is requested; they are not retained in the index. ```python +from functools import partial + +import torch from torch.utils.data import DataLoader +from pypaimon.multimodal.window_transforms import images_to_tensor, to_tensor dataset = ( frames.scan() .to_contiguous_window_dataset( - window_size=16, - columns=["state", "image"], - anchor_columns=["image"], + columns=["state", "image", "action"], + frame_offsets={"state": [-2, -1, 0], "action": list(range(16))}, group_key="episode_index", order_key="frame_index", - tail="pad", + boundary="pad", + column_transforms={ + "state": partial(to_tensor, dtype=torch.float32), + "action": partial(to_tensor, dtype=torch.float32), + "image": images_to_tensor, + }, ) ) loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True) ``` -Each item contains the group and order keys, one list for each requested -column, and a boolean `is_pad` tensor where `True` marks padding. Padding -repeats the final real value by default; `pad_values` can override individual -columns. Columns named in `anchor_columns` contain only the first row's value, -which is useful when an observation applies to a full action window. Use -`column_transforms` to convert column lists to tensors and -`adapter` to produce a model-specific sample mapping. Keep these callbacks -picklable when using multiple DataLoader workers. - -Scheduled anchors start at row zero and advance by `stride` (default `1`). -`tail="drop"` omits incomplete windows, `tail="pad"` includes and pads them, -and `tail="error"` rejects a sequence with any scheduled incomplete window. +Each raw item contains scalar group/order keys, a list for every selected column, +and a Boolean `_is_pad` tensor where `True` marks padding. Unspecified +field offsets default to `[0]`. The example produces state/action tensors with +time lengths 3/16 and an image tensor of shape `(1, C, H, W)`; DataLoader adds the +batch dimension. `pad_values` can replace endpoint repetition for individual +columns, before any transforms. `adapter` can normalize or rename fields and +produce a model-specific mapping. Keep callbacks picklable for multiple workers. + +Scheduled anchors start at each group's first row and advance by `stride` +(default `1`). `boundary="drop"` (default) omits anchors incomplete for any +field, `boundary="pad"` pads either end, and `boundary="error"` rejects any +scheduled incomplete window. `delta_timestamps` with an explicit `fps` is an +alternative to integer offsets, with frame-grid alignment checked against +`tolerance_s`. Existing `window_size`/`anchor_columns`/`tail` calls retain their +single `is_pad` output. See [multimodal reading](multimodal-reading#contiguous-windows-for-pytorch) +for the full sample and conversion contract. Rows are sorted by `order_key` inside each `group_key` value. Order values must be integers which increase by exactly one; duplicates and missing steps are rejected, and windows never cross groups. The resolved Paimon @@ -203,6 +216,12 @@ change its index or sample contents. A dataset pinned through `tag_name` fails its reads if the tag is moved to another snapshot, rather than mixing rows from the two snapshots. +Use normal PyTorch samplers for shuffling and distributed training. For DDP, +pass `DistributedSampler(dataset)` to DataLoader, omit `shuffle=True`, and call +the sampler's `set_epoch(epoch)` each epoch. Its padding or drop behavior still +applies when the sample count is not divisible by the number of ranks. Dataset +version pinning does not save sampler progress or random augmentation state. + Columns configured by `video-frame-field` are rejected: a window read would drop the `frame_index` and other metadata carried by their `VideoFrameDescriptor` values. Read those columns with `to_torch()` instead. diff --git a/paimon-python/pypaimon/benchmark/act/paimon.py b/paimon-python/pypaimon/benchmark/act/paimon.py index 18b50664b6d4..161edadbbc92 100644 --- a/paimon-python/pypaimon/benchmark/act/paimon.py +++ b/paimon-python/pypaimon/benchmark/act/paimon.py @@ -52,7 +52,7 @@ def __call__(self, sample): The persisted ``frame_index`` is forwarded as the shared ACT sample position. State and image columns are singleton lists; action retains the full - horizon and ``is_pad`` is forwarded unchanged. + horizon and ``action_is_pad`` becomes the ACT ``is_pad`` mask. """ qpos = np.concatenate([ np.asarray(sample[name][0], dtype=np.float32) @@ -80,7 +80,7 @@ def __call__(self, sample): "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), "action": torch.from_numpy(np.ascontiguousarray(action)), "images": torch.from_numpy(np.ascontiguousarray(images)), - "is_pad": sample["is_pad"], + "is_pad": sample["action_is_pad"], } @@ -114,13 +114,14 @@ def create_datasets( frames.scan(snapshot_id=snapshot_id).where( "episode_id = '%s'" % episode_id.replace("'", "''") ).to_contiguous_window_dataset( - window_size=config.action_horizon, + frame_offsets={ + name: range(config.action_horizon) for name in ACTION_COLUMNS + }, columns=QPOS_COLUMNS + ACTION_COLUMNS + IMAGE_COLUMNS, - anchor_columns=QPOS_COLUMNS + IMAGE_COLUMNS, group_key="episode_id", order_key="frame_index", stride=1, - tail="drop", + boundary="drop", adapter=PaimonACTAdapter(normalization), ) for episode_id in (train_episode_id, validation_episode_id) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 1332b870cf4e..37ddd814b230 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -18,7 +18,6 @@ """LeRobot-compatible map-style reads from a multimodal Paimon table.""" import bisect -import io import json import math import operator @@ -1332,14 +1331,10 @@ def _torch_row(row, features, return_uint8=False): def _image_tensor(payload, feature, return_uint8=False): if payload is None: raise ValueError("LeRobot image feature contains a null frame.") - import numpy as np - import torch - try: - from PIL import Image, ImageOps - except ImportError as error: - raise ImportError( - "PaimonLeRobotDataset requires Pillow from " - "'pypaimon[lerobot]'.") from error + from pypaimon.multimodal.window_transforms import ( + _decode_image, + _image_array_to_tensor, + ) expected_shape = _feature_shape(feature, "image") if len(expected_shape) != 3: @@ -1349,21 +1344,12 @@ def _image_tensor(payload, feature, return_uint8=False): payload_shape = expected_shape[1:] + expected_shape[:1] \ if names and names[0] in ("channel", "channels") \ else expected_shape - with Image.open(io.BytesIO(payload)) as image: - array = np.array(ImageOps.exif_transpose(image), copy=True) - if array.ndim == 2: - array = array[:, :, None] + array = _decode_image(payload) if array.shape != payload_shape: raise ValueError( "LeRobot image payload has shape %s, expected %s." % (array.shape, payload_shape)) - normalize = array.dtype == np.uint8 - tensor = torch.from_numpy(array).permute(2, 0, 1) - if normalize and return_uint8: - return tensor - # Preserve high-bit-depth and floating-point images in native units. - tensor = tensor.float() - return tensor.div_(255) if normalize else tensor + return _image_array_to_tensor(array, return_uint8) def _video_tensor(frame, feature, return_uint8=False): diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 489ff641525b..e21edb2f1ea0 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -167,42 +167,58 @@ def to_torch( def to_contiguous_window_dataset( self, *, - window_size, + window_size=None, columns=None, anchor_columns=None, group_key="episode_index", order_key="frame_index", stride=1, - tail="drop", + tail=None, column_transforms=None, pad_values=None, adapter=None, - blob_parallelism=64): - """Build a snapshot-pinned, map-style Dataset of contiguous rows. + blob_parallelism=64, + frame_offsets=None, + delta_timestamps=None, + fps=None, + tolerance_s=None, + boundary=None): + """Build snapshot-pinned training windows from this single-table scan. The Dataset indexes only ``group_key``, ``order_key``, and Paimon row - IDs, then reads projected values on demand. Columns listed in - ``anchor_columns`` are provided to ``column_transforms`` as one-element - lists read from the first row of each window; ``adapter`` receives the - transformed values. ``order_key`` must contain non-null integers that + IDs, then reads projected values on demand. Each column has its own + relative frame offsets; unspecified columns use ``[0]``. Transforms + receive padded lists, including singleton windows, before ``adapter`` + receives the sample. ``order_key`` must contain non-null integers that increase by exactly one within each group. The Dataset sorts rows within each group and never creates a window across groups. Args: - window_size: Number of rows in a complete window. + window_size: Legacy forward window size. Mutually exclusive with + frame_offsets and delta_timestamps. Retains a single is_pad + mask instead of the new per-column masks. columns: Value columns to return, excluding the group and order keys. The scan projection is used when omitted. - anchor_columns: Subset of ``columns`` read only from the window's - first row. + anchor_columns: With window_size, columns using only the anchor. group_key: Column identifying an independent row sequence. order_key: Integer position column within each group. - stride: Distance between scheduled window starts. - tail: Handling for incomplete final windows: ``drop``, ``pad``, or - ``error``. + stride: Distance between anchors, starting at each group's first row. + tail: Legacy boundary alias requiring window_size; cannot be used + together with boundary. column_transforms: Per-column callables applied to value lists. - pad_values: Optional replacement values used by ``tail='pad'``. + pad_values: Per-column raw padding values; otherwise repeat the + nearest endpoint. Padding happens before column transforms. adapter: Callable that converts the complete sample mapping. blob_parallelism: Maximum concurrent BLOB body reads per fetch. + frame_offsets: Mapping of columns to nonempty integer offset + sequences. Supports history, future, sparse and repeated frames. + delta_timestamps: Alternative per-column offsets in seconds, aligned + to the regular frame grid defined by fps within tolerance_s. + fps: Finite positive frame rate, required with delta_timestamps. + tolerance_s: Finite nonnegative seconds-conversion tolerance, + defaulting to 1e-4. Only valid with delta_timestamps. + boundary: Handling at both group ends: drop (default), pad, or error. + Drop retains only anchors valid for every requested offset. Returns: A snapshot-pinned ``ContiguousWindowDataset``. See that class for @@ -222,6 +238,11 @@ def to_contiguous_window_dataset( pad_values=pad_values, adapter=adapter, blob_parallelism=blob_parallelism, + frame_offsets=frame_offsets, + delta_timestamps=delta_timestamps, + fps=fps, + tolerance_s=tolerance_s, + boundary=boundary, ) def to_ray( diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index 2a6384596281..ec5f63996976 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -18,6 +18,7 @@ """Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" import copy +import math import operator import numpy as np @@ -42,77 +43,116 @@ class ContiguousWindowDataset(Dataset): - """Map-style Dataset which reads fixed row windows on demand. + """Map-style training samples from a single table of contiguous frames. The in-memory index contains only group values, order bounds, and Paimon row IDs, stored in Arrow and NumPy arrays. Each ``__getitem__`` reads the projected rows from the snapshot resolved while the index was built, reusing that snapshot's authorized scan plan instead of planning again. 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``: + one; rows from different groups never share a window. ``frame_offsets`` + selects relative frames independently for each projected column. Unspecified + columns use ``[0]``. Alternatively, ``delta_timestamps`` specifies offsets in + seconds on a regular grid defined by ``fps``. ``boundary`` handles scheduled + anchors whose offsets extend beyond either end of a group: * ``drop`` omits them; - * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``pad`` repeats the closest endpoint, or uses the column's ``pad_values``; * ``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. + The result contains scalar group/order values identifying the anchor, a + list for each value column (including singleton windows), and a Boolean + ``_is_pad`` tensor per column. Offsets retain their order and + duplicates. Padding happens before ``column_transforms`` convert each list; + ``adapter`` then adapts the complete mapping to a model-specific contract. + Both transforms operate on independent copies of the requested values. + + For compatibility, ``window_size`` selects forward consecutive frames for + all columns except ``anchor_columns``, which use ``[0]``. These calls retain + the original ``tail`` argument and single ``is_pad`` mask in their output. + Only one of the three window definitions may be supplied. Without any, + every column uses ``[0]``. ``stride`` schedules anchors from each group's + first row; boundary handling never shifts that schedule. ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. Video frame columns are not supported yet, because a window read would drop the frame metadata carried by their descriptors. """ - _TAIL_POLICIES = ("drop", "pad", "error") + _BOUNDARY_POLICIES = ("drop", "pad", "error") def __init__( self, query, *, - window_size, + window_size=None, columns=None, anchor_columns=None, group_key="episode_index", order_key="frame_index", stride=1, - tail="drop", + tail=None, column_transforms=None, pad_values=None, adapter=None, - blob_parallelism=64): + blob_parallelism=64, + frame_offsets=None, + delta_timestamps=None, + fps=None, + tolerance_s=None, + boundary=None): if not isinstance(query, ScanQuery) or isinstance(query, _PreFilterQuery): raise TypeError( "ContiguousWindowDataset is only supported on scan(), " "not search queries.") - self.window_size = _positive_int(window_size, "window_size") + self.window_size = ( + _positive_int(window_size, "window_size") + if window_size is not None else None) self.stride = _positive_int(stride, "stride") - if tail not in self._TAIL_POLICIES: + if sum(value is not None for value in ( + window_size, frame_offsets, delta_timestamps)) > 1: raise ValueError( - "tail must be one of %s; got %r." - % (self._TAIL_POLICIES, tail)) - self.tail = tail + "Specify only one of window_size, frame_offsets, " + "or delta_timestamps.") + if window_size is None and (anchor_columns is not None or tail is not None): + raise ValueError("anchor_columns and tail require window_size.") + if tail is not None and boundary is not None: + raise ValueError("Specify boundary or tail, not both.") + self.boundary = boundary if boundary is not None else ( + tail if tail is not None else "drop") + if self.boundary not in self._BOUNDARY_POLICIES: + raise ValueError( + "boundary must be one of %s; got %r." + % (self._BOUNDARY_POLICIES, self.boundary)) + if self.window_size is not None: + self.tail = self.boundary 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): + if self.window_size is not None and "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) _reject_video_columns(query._table, self.columns) 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.frame_offsets = _frame_offsets( + self.columns, frame_offsets, delta_timestamps, fps, tolerance_s, + self.window_size, self.anchor_columns) + reserved = {"is_pad"} if self.window_size is not None else { + "%s_is_pad" % name for name in self.columns} + collisions = reserved.intersection( + self.columns + [self.group_key, self.order_key]) + if collisions: + raise ValueError( + "Generated padding mask names conflict with output " + "columns: %s." % sorted(collisions)) self.column_transforms = _column_transforms( column_transforms, self.columns) self.pad_values = _pad_values(pad_values, self.columns) + self._read_groups = {} + for name, offsets in self.frame_offsets.items(): + key = (offsets, self.boundary == "pad" and name in self.pad_values) + self._read_groups.setdefault(key, []).append(name) if adapter is not None and not callable(adapter): raise TypeError("adapter must be callable or None.") self.adapter = adapter @@ -139,19 +179,15 @@ def __len__(self): return int(self._anchor_groups.size) def __getitem__(self, index): - """Read one window by map-style Dataset index. + """Read one sample, or a list of samples for a Python slice. 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) + if isinstance(index, slice): + return self.__getitems__(range(*index.indices(len(self)))) + return self.__getitems__([index])[0] def __getitems__(self, indices): """Read several Dataset indices while coalescing overlapping row IDs. @@ -162,28 +198,20 @@ def __getitems__(self, indices): 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 {} - ) + reads = {} + for (offsets, constant_pad), columns in self._read_groups.items(): + plans = [self._column_window(anchor, offsets, constant_pad) + for anchor in windows] + row_ids = list(dict.fromkeys( + row_id for ids, _ in plans for row_id in ids + if row_id is not None)) + rows = self._read_rows(row_ids, columns) if row_ids else [] + by_id = dict(zip(row_ids, rows)) + for name in columns: + reads[name] = (plans, by_id) 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 + self._sample(anchor, position, reads) + for position, anchor in enumerate(windows) ] def __getstate__(self): @@ -202,32 +230,38 @@ def _resolve_window(self, index): group_index = int(self._anchor_groups[index]) start = int(self._anchor_starts[index]) - valid_count = min( - self.window_size, int(self._group_lengths[group_index]) - start) - offset = int(self._group_starts[group_index]) + start - row_ids = self._row_ids[offset:offset + valid_count].tolist() - return (group_index, start, valid_count), row_ids - - def _sample(self, anchor, rows, anchor_row=None): - group_index, start, valid_count = anchor - 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 + return group_index, start + + def _column_window(self, anchor, offsets, constant_pad): + group_index, start = anchor + length = int(self._group_lengths[group_index]) + group_start = int(self._group_starts[group_index]) + positions = [start + delta for delta in offsets] + padding = [position < 0 or position >= length for position in positions] + row_ids = [ + None if padded and constant_pad else int(self._row_ids[ + group_start + min(max(position, 0), length - 1)]) + for position, padded in zip(positions, padding) + ] + return row_ids, padding + + def _sample(self, anchor, position, reads): + group_index, start = anchor sample = { self.group_key: self._group_keys[group_index], self.order_key: int(self._group_first_orders[group_index]) + start, - "is_pad": padding_mask, } + if self.window_size is not None: + available = int(self._group_lengths[group_index]) - start + sample["is_pad"] = torch.arange(self.window_size) >= available for name in self.columns: - if name in self.anchor_columns: - values = [copy.deepcopy(anchor_row[name])] - else: - values = [copy.deepcopy(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)) + plans, rows = reads[name] + row_ids, padding = plans[position] + values = [copy.deepcopy( + self.pad_values[name] if row_id is None else rows[row_id][name]) + for row_id in row_ids] + if self.window_size is None: + sample["%s_is_pad" % name] = torch.tensor(padding, dtype=torch.bool) transform = self.column_transforms.get(name) sample[name] = transform(values) if transform is not None else values if self.adapter is not None: @@ -320,24 +354,33 @@ def _group_of(self, position): def _build_anchors(self): groups = [] starts = [] + minimum = min(min(offsets) for offsets in self.frame_offsets.values()) + maximum = max(max(offsets) for offsets in self.frame_offsets.values()) + # Legacy windows schedule their full horizon even with anchor-only + # columns, so their cardinality and global mask remain unchanged. + if self.window_size is not None: + maximum = max(maximum, self.window_size - 1) for group_index, length in enumerate(self._group_lengths): length = int(length) positions = np.arange(0, length, self.stride, dtype=np.int64) - valid_counts = np.minimum(self.window_size, length - positions) - incomplete = np.flatnonzero(valid_counts < self.window_size) + lower = min(length, max(0, -minimum)) + upper = max(0, min(length, length - maximum)) + valid = (positions >= lower) & (positions < upper) + incomplete = np.flatnonzero(~valid) if incomplete.size: - if self.tail == "error": + if self.boundary == "error": first = int(incomplete[0]) + size = "window_size=%s, " % self.window_size \ + if self.window_size is not None else "" raise ValueError( "Group %s has an incomplete window at %s: " - "window_size=%d, available=%d." + "%soffset bounds=[%d, %d]." % (self._group_keys[group_index], int(self._group_first_orders[group_index]) + int(positions[first]), - self.window_size, - int(valid_counts[first]))) - if self.tail == "drop": - positions = positions[valid_counts == self.window_size] + size, minimum, maximum)) + if self.boundary == "drop": + positions = positions[valid] if positions.size: groups.append(np.full(positions.size, group_index, dtype=np.int64)) starts.append(positions) @@ -346,11 +389,6 @@ def _build_anchors(self): return empty, empty.copy() return np.concatenate(groups), np.concatenate(starts) - 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. @@ -608,6 +646,80 @@ def _reject_video_columns(table, columns): "frame_index." % video_columns) +def _frame_offsets(columns, offsets, timestamps, fps, tolerance_s, + window_size, anchor_columns): + """Normalize the public window forms to nonempty integer offset tuples.""" + if timestamps is not None: + tolerance = _finite_number( + 1e-4 if tolerance_s is None else tolerance_s, "tolerance_s") + if tolerance < 0: + raise ValueError("tolerance_s must be non-negative.") + if fps is None: + raise ValueError("delta_timestamps requires an explicit fps.") + fps = _finite_number(fps, "fps") + if fps <= 0: + raise ValueError("fps must be positive.") + elif fps is not None or tolerance_s is not None: + raise ValueError("fps and tolerance_s are only used with delta_timestamps.") + if window_size is not None: + forward = tuple(range(window_size)) + return { + name: (0,) if name in anchor_columns else forward + for name in columns + } + parameter = "delta_timestamps" if timestamps is not None else "frame_offsets" + supplied = _mapping(timestamps if timestamps is not None else offsets, parameter) + _validate_mapping_columns(supplied, columns, parameter) + result = {} + for name in columns: + if name not in supplied: + result[name] = (0,) + continue + values = supplied[name] + if isinstance(values, (str, bytes, dict)): + raise TypeError("%s[%r] must be a sequence of offsets." % (parameter, name)) + try: + values = list(values) + except TypeError: + raise TypeError("%s[%r] must be a sequence of offsets." % (parameter, name)) + if not values: + raise ValueError("%s[%r] must not be empty." % (parameter, name)) + converted = [] + for value in values: + if timestamps is not None: + seconds = _finite_number(value, "delta_timestamps[%r]" % name) + scaled = seconds * fps + if not math.isfinite(scaled): + raise ValueError("delta_timestamps * fps must be finite.") + delta = round(scaled) + if abs(seconds - delta / fps) > tolerance: + raise ValueError( + "delta_timestamps for %s must align to multiples of 1/fps." + % name) + else: + if isinstance(value, (bool, np.bool_)): + raise TypeError("frame_offsets must contain integers, not booleans.") + try: + delta = operator.index(value) + except TypeError: + raise TypeError("frame_offsets must contain integer offsets.") + converted.append(delta) + result[name] = tuple(converted) + return result + + +def _finite_number(value, name): + if isinstance(value, (bool, np.bool_)): + raise ValueError("%s must be a finite number, not a boolean." % name) + try: + number = float(value) + except (TypeError, ValueError, OverflowError): + raise ValueError("%s must be a finite number." % name) + if not math.isfinite(number): + raise ValueError("%s must be a finite number." % name) + return number + + def _columns(query, columns, group_key, order_key): available = {field.name for field in query._table.fields} if columns is None: @@ -635,10 +747,10 @@ def _columns(query, columns, group_key, order_key): if invalid: raise ValueError("columns do not exist: %s." % invalid) reserved = [name for name in columns - if name in (group_key, order_key, "is_pad")] + if name in (group_key, order_key)] if reserved: raise ValueError( - "columns must not include group_key, order_key, or is_pad: %s." + "columns must not include group_key or order_key: %s." % reserved) return columns diff --git a/paimon-python/pypaimon/multimodal/window_transforms.py b/paimon-python/pypaimon/multimodal/window_transforms.py new file mode 100644 index 000000000000..3385b8525df2 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/window_transforms.py @@ -0,0 +1,88 @@ +# 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. + +"""Column transforms for materialized training windows.""" + +import io + + +def to_tensor(values, dtype=None): + """Convert numeric window values to a tensor with a leading time axis. + + Use ``partial(to_tensor, dtype=torch.float32)`` from ``functools`` + to choose a dtype. + This conversion does not normalize values. + """ + import torch + + result = torch.as_tensor(values, dtype=dtype) + if result.ndim == 0: + raise ValueError("Window values must have a time axis.") + return result + + +def images_to_tensor(values, return_uint8=False): + """Decode a non-empty sequence of image bytes to a TCHW tensor. + + Apply EXIF orientation and preserve grayscale as one channel. Eight-bit + pixels become float32 in [0, 1], or stay uint8 with ``return_uint8=True``. + Higher-bit-depth pixels always become float32 in their original units. + All decoded frames must have the same shape. + """ + import torch + + if not isinstance(return_uint8, bool): + raise TypeError("return_uint8 must be a boolean.") + frames = [] + for payload in values: + if not isinstance(payload, (bytes, bytearray, memoryview)): + raise ValueError("Image window values must contain image bytes.") + frame = _image_array_to_tensor(_decode_image(payload), return_uint8) + if frames and frame.shape != frames[0].shape: + raise ValueError("Image window frames must have the same shape.") + frames.append(frame) + if not frames: + raise ValueError("Image window must contain at least one frame.") + return torch.stack(frames) + + +def _decode_image(payload): + import numpy as np + try: + from PIL import Image, ImageOps + except ImportError as error: + raise ImportError( + "Image decoding requires Pillow; install 'pypaimon[torch]' " + "or Pillow.") from error + + with Image.open(io.BytesIO(payload)) as image: + array = np.array(ImageOps.exif_transpose(image), copy=True) + if array.ndim == 2: + array = array[:, :, None] + return array + + +def _image_array_to_tensor(array, return_uint8=False): + import numpy as np + import torch + + normalize = array.dtype == np.uint8 + tensor = torch.from_numpy(array).permute(2, 0, 1) + if normalize and return_uint8: + return tensor + tensor = tensor.float() + return tensor.div_(255) if normalize else tensor diff --git a/paimon-python/pypaimon/tests/act_runner_test.py b/paimon-python/pypaimon/tests/act_runner_test.py index 57d64bf8c110..a7fbc69e9299 100644 --- a/paimon-python/pypaimon/tests/act_runner_test.py +++ b/paimon-python/pypaimon/tests/act_runner_test.py @@ -418,7 +418,16 @@ def test_backends_match_the_golden_act_window_contract(benchmark_input): ) expected = hdf5[1] - actual = paimon[1] + with patch.object(paimon, "adapter", wraps=paimon.adapter) as adapt: + actual = paimon[1] + generic_sample = adapt.call_args.args[0] + assert "is_pad" not in generic_sample + assert len(generic_sample["action"]) == 3 + assert generic_sample["action_is_pad"].tolist() == [False] * 3 + for name in QPOS_COLUMNS + IMAGE_COLUMNS: + assert len(generic_sample[name]) == 1 + assert generic_sample[name + "_is_pad"].tolist() == [False] + assert actual["is_pad"] is generic_sample["action_is_pad"] assert set(expected) == { "sample_id", "episode_id", "frame_index", "qpos", "action", @@ -605,11 +614,10 @@ def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( with patch.object( train, "_read_rows", wraps=train._read_rows) as read_rows: sample_before_append = train[0] - assert [call.args[1] for call in read_rows.call_args_list] == [ - list(ACTION_COLUMNS), - list(QPOS_COLUMNS + IMAGE_COLUMNS), - ] - assert [len(call.args[0]) for call in read_rows.call_args_list] == [3, 1] + assert { + tuple(call.args[1]): len(call.args[0]) + for call in read_rows.call_args_list + } == {ACTION_COLUMNS: 3, QPOS_COLUMNS + IMAGE_COLUMNS: 1} assert fetch.call_count == 1 assert { name: len(fetch.call_args.args[1][name]) diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index c3582ea7b576..c0e5d4d7dd64 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -116,6 +116,242 @@ def _dataset(table, **kwargs): ) ) + def test_per_column_offsets_combine_history_future_and_anchor(self): + dataset = self._table().scan().to_contiguous_window_dataset( + columns=["value", "payload"], group_key="episode", order_key="step", + frame_offsets={"value": [-1, 0, 2]}, boundary="pad") + + self.assertEqual(6, len(dataset)) + sample = dataset[0] + self.assertEqual(("episode-a", 0), (sample["episode"], sample["step"])) + self.assertEqual([0, 0, 1], sample["value"]) + self.assertEqual([True, False, True], sample["value_is_pad"].tolist()) + self.assertEqual([b"episode-a-0"], sample["payload"]) + self.assertEqual([False], sample["payload_is_pad"].tolist()) + self.assertEqual(torch.bool, sample["value_is_pad"].dtype) + self.assertNotIn("is_pad", sample) + self.assertEqual([101, 102, 103], dataset[4]["value"]) + self.assertEqual([False, False, True], dataset[4]["value_is_pad"].tolist()) + + def test_frame_mode_defaults_to_single_frame_with_time_dimension(self): + dataset = self._table().scan().select(["value"]).to_contiguous_window_dataset( + group_key="episode", order_key="step") + + self.assertEqual(6, len(dataset)) + self.assertEqual({"episode", "step", "value", "value_is_pad"}, set(dataset[0])) + self.assertEqual([0], dataset[0]["value"]) + self.assertEqual([False], dataset[0]["value_is_pad"].tolist()) + self.assertEqual(103, dataset[-1]["value"][0]) + + def test_drop_offsets_keep_stride_anchored_at_group_start(self): + table = self._table() + dataset = table.scan().to_contiguous_window_dataset( + columns=["value", "payload"], group_key="episode", order_key="step", + frame_offsets={"value": [-1, 0], "payload": [1]}, stride=2) + + self.assertEqual(1, len(dataset)) + self.assertEqual(("episode-b", 2), (dataset[0]["episode"], dataset[0]["step"])) + self.assertEqual([101, 102], dataset[0]["value"]) + self.assertEqual([b"episode-b-3"], dataset[0]["payload"]) + self.assertFalse(dataset[0]["value_is_pad"].any()) + with self.assertRaisesRegex( + ValueError, r"episode-a.*at 0: offset bounds=\[-1, 0\]"): + table.scan().to_contiguous_window_dataset( + columns=["value"], group_key="episode", order_key="step", + frame_offsets={"value": [-1, 0]}, boundary="error") + + def test_sparse_repeated_offsets_preserve_batch_and_slice_order(self): + dataset = self._table().scan().to_contiguous_window_dataset( + columns=["value"], group_key="episode", order_key="step", + frame_offsets={"value": [2, -1, 2, 0]}, boundary="pad") + + samples = dataset.__getitems__([4, 0, 4, -1]) + self.assertEqual([[103, 101, 103, 102], [1, 0, 1, 0], + [103, 101, 103, 102], [103, 102, 103, 103]], + [sample["value"] for sample in samples]) + self.assertEqual([True, False, True, False], samples[0]["value_is_pad"].tolist()) + self.assertEqual([dataset[i]["value"] for i in (5, 3, 1)], + [sample["value"] for sample in dataset[::-2]]) + self.assertEqual([], dataset.__getitems__([])) + with self.assertRaises(IndexError): + dataset.__getitems__([0, len(dataset)]) + + def test_custom_padding_precedes_transforms_even_for_all_padding(self): + dataset = self._table().scan().to_contiguous_window_dataset( + columns=["value", "payload"], group_key="episode", order_key="step", + frame_offsets={"value": [-10, 10], "payload": [0, 10]}, + boundary="pad", pad_values={"value": -7, "payload": b"missing"}, + column_transforms={"value": _TensorColumnTransform()}) + + with patch.object(dataset, "_read_rows", wraps=dataset._read_rows) as read: + sample = dataset[0] + self.assertEqual(1, read.call_count) + self.assertEqual(["payload"], read.call_args.args[1]) + self.assertTrue(torch.equal(torch.tensor([-7, -7]), sample["value"])) + self.assertEqual([True, True], sample["value_is_pad"].tolist()) + self.assertEqual([b"episode-a-0", b"missing"], sample["payload"]) + self.assertEqual([False, True], sample["payload_is_pad"].tolist()) + + def test_seconds_offsets_match_frame_offsets_with_alignment_tolerance(self): + table = self._table() + kwargs = dict(columns=["value", "payload"], group_key="episode", + order_key="step", boundary="pad") + frames = table.scan().to_contiguous_window_dataset( + frame_offsets={"value": [-1, 0, 2]}, **kwargs) + seconds = table.scan().to_contiguous_window_dataset( + delta_timestamps={"value": [-0.1, 0, 0.20001]}, fps=10, + tolerance_s=0.0001, **kwargs) + + self.assertEqual(len(frames), len(seconds)) + for expected, actual in zip(frames[:], seconds[:]): + self.assertEqual(expected["value"], actual["value"]) + self.assertEqual(expected["payload"], actual["payload"]) + self.assertTrue(torch.equal(expected["value_is_pad"], actual["value_is_pad"])) + + def test_offset_projections_do_not_fetch_anchor_images_for_action_rows(self): + dataset = self._table().scan().to_contiguous_window_dataset( + columns=["value", "payload"], group_key="episode", order_key="step", + frame_offsets={"value": [-1, 0, 1], "payload": [0]}, boundary="pad") + + with patch.object(dataset, "_read_rows", wraps=dataset._read_rows) as read, \ + patch("pypaimon.multimodal.window_dataset.fetch_blob_bodies", + side_effect=window_dataset.fetch_blob_bodies) as fetch: + samples = dataset.__getitems__([3, 4, 3]) + + projections = {tuple(call.args[1]): len(call.args[0]) for call in read.call_args_list} + self.assertEqual({("value",): 4, ("payload",): 2}, projections) + self.assertEqual(2, sum(len(call.args[1]["payload"]) for call in fetch.call_args_list)) + self.assertEqual([[100, 101, 102], [101, 102, 103], [100, 101, 102]], + [sample["value"] for sample in samples]) + + def test_unused_pad_values_do_not_split_reads(self): + table = self._table() + for boundary in ("drop", "error"): + with self.subTest(boundary=boundary): + dataset = table.scan().to_contiguous_window_dataset( + columns=["value", "payload"], group_key="episode", order_key="step", + frame_offsets={"value": [0, 1], "payload": [0, 1]}, + stride=2, boundary=boundary, pad_values={"value": -1}) + with patch.object(dataset, "_read_rows", wraps=dataset._read_rows) as read: + samples = dataset.__getitems__([0, 1, 0]) + self.assertEqual(1, read.call_count) + self.assertEqual(["value", "payload"], read.call_args.args[1]) + self.assertEqual(4, len(read.call_args.args[0])) + self.assertEqual([[0, 1], [100, 101], [0, 1]], + [sample["value"] for sample in samples]) + self.assertEqual([b"episode-a-0", b"episode-a-1"], samples[0]["payload"]) + + def test_offset_padding_and_repeated_cells_are_mutably_isolated(self): + table = self.conn.create_table( + "offset_mutable", schema=pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("values", pa.list_(pa.int32()), nullable=False), + ]), options=_TABLE_OPTIONS) + table.add([{"episode": "a", "step": 0, "values": [0]}]) + pad = [-1] + dataset = table.scan().to_contiguous_window_dataset( + columns=["values"], group_key="episode", order_key="step", + frame_offsets={"values": [-1, 0, 0, 1]}, boundary="pad", + pad_values={"values": pad}) + + first, second = dataset.__getitems__([0, 0]) + first["values"][0].append(99) + first["values"][1].append(99) + self.assertEqual([0], first["values"][2]) + self.assertEqual([-1], first["values"][3]) + self.assertEqual([[-1], [0], [0], [-1]], second["values"]) + self.assertEqual([[-1], [0], [0], [-1]], dataset[0]["values"]) + self.assertEqual([-1], pad) + + def test_offset_dataloader_workers_and_distributed_sampler(self): + dataset = self._table().scan().to_contiguous_window_dataset( + columns=["value"], group_key="episode", order_key="step", + frame_offsets={"value": [-1, 0, 1]}, boundary="pad", + column_transforms={"value": _TensorColumnTransform()}) + restored = pickle.loads(pickle.dumps(dataset)) + self.assertEqual(dataset.snapshot_id, restored.snapshot_id) + batches = list(torch.utils.data.DataLoader( + restored, batch_size=2, num_workers=2, shuffle=False)) + self.assertEqual([[0, 0, 1], [0, 1, 1], [100, 100, 101], + [100, 101, 102], [101, 102, 103], [102, 103, 103]], + [row for batch in batches for row in batch["value"].tolist()]) + self.assertEqual((2, 3), tuple(batches[0]["value_is_pad"].shape)) + self.assertEqual(torch.bool, batches[0]["value_is_pad"].dtype) + partitions = [list(torch.utils.data.DistributedSampler( + dataset, num_replicas=2, rank=rank, shuffle=False)) for rank in range(2)] + self.assertEqual([[0, 2, 4], [1, 3, 5]], partitions) + self.assertEqual([0, 100, 102], [dataset[i]["value"][1].item() for i in partitions[0]]) + + def test_rejects_invalid_offset_configuration_at_construction(self): + table = self._table() + cases = [ + {"frame_offsets": {"missing": [0]}}, + {"frame_offsets": {"value": []}}, + {"frame_offsets": {"value": [1.0]}}, + {"frame_offsets": {"value": [True]}}, + {"frame_offsets": {"value": [0]}, "window_size": 2}, + {"delta_timestamps": {"value": [0]}, "fps": 10, "window_size": 2}, + {"frame_offsets": {"value": [0]}, "delta_timestamps": {"value": [0]}, "fps": 10}, + {"delta_timestamps": {"value": [0]}}, + {"delta_timestamps": {"value": []}, "fps": 10}, + {"delta_timestamps": {"value": [float("nan")]}, "fps": 10}, + {"delta_timestamps": {"value": [float("inf")]}, "fps": 10}, + {"delta_timestamps": {"value": [0]}, "fps": 0}, + {"delta_timestamps": {"value": [0]}, "fps": -1}, + {"delta_timestamps": {"value": [0]}, "fps": float("inf")}, + {"delta_timestamps": {"value": [0]}, "fps": float("nan")}, + {"delta_timestamps": {"value": [0.05]}, "fps": 10}, + {"frame_offsets": {"value": [0]}, "anchor_columns": ["value"]}, + {"frame_offsets": {"value": [0]}, "tail": "pad"}, + {"frame_offsets": {"value": [0]}, "tolerance_s": 0.01}, + {"window_size": 2, "tolerance_s": 0.01}, + {"window_size": 2, "tail": "pad", "boundary": "pad"}, + {"boundary": "unknown"}, + ] + for kwargs in cases: + with self.subTest(kwargs=kwargs), self.assertRaises((TypeError, ValueError)): + table.scan().to_contiguous_window_dataset( + columns=["value"], group_key="episode", order_key="step", **kwargs) + + def test_rejects_generated_padding_mask_collisions(self): + table = self.conn.create_table( + "mask_collision", schema=pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + pa.field("value_is_pad", pa.int32(), nullable=False), + ]), options=_TABLE_OPTIONS) + table.add([{"episode": "a", "step": 0, "value": 1, "value_is_pad": 2}]) + with self.assertRaises(ValueError): + table.scan().to_contiguous_window_dataset( + columns=["value", "value_is_pad"], group_key="episode", order_key="step") + with self.assertRaises(ValueError): + table.scan().to_contiguous_window_dataset( + columns=["value"], group_key="episode", order_key="value_is_pad") + + def test_offset_mode_allows_is_pad_as_a_stored_column(self): + table = self.conn.create_table( + "stored_is_pad", schema=pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("is_pad", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + ]), options=_TABLE_OPTIONS) + table.add([{"episode": "a", "step": i, "is_pad": i, "value": 100 + i} + for i in range(3)]) + for group, order, column in (("episode", "step", "is_pad"), + ("is_pad", "step", "value"), + ("episode", "is_pad", "value")): + with self.subTest(group=group, order=order, column=column): + kwargs = dict(group_key=group, order_key=order, columns=[column]) + dataset = table.scan().to_contiguous_window_dataset(**kwargs) + self.assertEqual(3, len(dataset)) + self.assertEqual([0] if column == "is_pad" else [100], dataset[0][column]) + self.assertEqual([False], dataset[0][column + "_is_pad"].tolist()) + with self.assertRaises(ValueError): + table.scan().to_contiguous_window_dataset(window_size=1, **kwargs) + def test_sorts_rows_and_never_crosses_episode_boundaries(self): dataset = self._dataset(self._table()) diff --git a/paimon-python/pypaimon/tests/window_transforms_test.py b/paimon-python/pypaimon/tests/window_transforms_test.py new file mode 100644 index 000000000000..18493c5f89b4 --- /dev/null +++ b/paimon-python/pypaimon/tests/window_transforms_test.py @@ -0,0 +1,109 @@ +# 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 io +import unittest +from functools import partial +from unittest.mock import patch + +import numpy as np +import torch +from PIL import Image + +from pypaimon.multimodal.window_transforms import images_to_tensor, to_tensor + + +def _png(array, **kwargs): + buffer = io.BytesIO() + Image.fromarray(array).save(buffer, format="PNG", **kwargs) + return buffer.getvalue() + + +class WindowTransformsTest(unittest.TestCase): + + def test_numeric_values_keep_time_axis_and_explicit_dtype(self): + convert = partial(to_tensor, dtype=torch.float64) + values = convert([[1, 2], [3, 4]]) + self.assertEqual((2, 2), tuple(values.shape)) + self.assertEqual(torch.float64, values.dtype) + self.assertEqual([1, 2], to_tensor([1, 2]).tolist()) + with self.assertRaisesRegex(ValueError, "time axis"): + to_tensor(1) + + def test_rgb_and_gray_pixels_and_time_axis(self): + rgb = _png(np.full((2, 3, 3), 128, dtype=np.uint8)) + values = images_to_tensor([rgb, rgb]) + self.assertEqual((2, 3, 2, 3), tuple(values.shape)) + self.assertEqual(torch.float32, values.dtype) + torch.testing.assert_close(values, torch.full_like(values, 128 / 255)) + raw = images_to_tensor([rgb], return_uint8=True) + self.assertEqual(torch.uint8, raw.dtype) + self.assertEqual(128, raw[0, 0, 0, 0].item()) + gray = _png(np.array([[0, 255]], dtype=np.uint8)) + self.assertEqual([[[[0., 1.]]]], images_to_tensor([gray]).tolist()) + + def test_high_bit_depth_keeps_native_units(self): + payload = _png(np.array([[0, 1024, 65535]], dtype=np.uint16)) + for return_uint8 in (False, True): + values = images_to_tensor([payload], return_uint8=return_uint8) + self.assertEqual(torch.float32, values.dtype) + self.assertEqual([[[[0., 1024., 65535.]]]], values.tolist()) + + def test_exif_orientation_is_applied_before_stacking(self): + exif = Image.Exif() + exif[274] = 6 + payload = _png(np.array([[1, 2], [3, 4], [5, 6]], dtype=np.uint8), + exif=exif) + values = images_to_tensor([payload], return_uint8=True) + self.assertEqual([[[[5, 3, 1], [6, 4, 2]]]], values.tolist()) + + def test_invalid_image_inputs_fail_clearly(self): + for values in ([], [None], ["not bytes"]): + with self.subTest(values=values), self.assertRaises(ValueError): + images_to_tensor(values) + with self.assertRaises(OSError): + images_to_tensor([b"not an image"]) + with self.assertRaises(TypeError): + images_to_tensor([b"unused"], return_uint8=1) + with self.assertRaises(ValueError): + images_to_tensor([_png(np.zeros((2, 2), dtype=np.uint8)), + _png(np.zeros((3, 2), dtype=np.uint8))]) + + def test_image_import_errors_identify_the_failing_dependency(self): + from pypaimon.multimodal.lerobot.dataset import _image_tensor + + payload = _png(np.zeros((2, 3, 3), dtype=np.uint8)) + feature = {"dtype": "image", "shape": (2, 3, 3)} + readers = (lambda: images_to_tensor([payload]), + lambda: _image_tensor(payload, feature)) + for read in readers: + with patch.dict("sys.modules", {"PIL": None}): + with self.assertRaisesRegex(ImportError, "requires Pillow"): + read() + with patch.dict("sys.modules", {"numpy": None}): + with self.assertRaises(ImportError) as raised: + read() + self.assertEqual("numpy", raised.exception.name) + failure = ImportError("decoder plugin unavailable") + with patch.object(Image, "open", side_effect=failure): + with self.assertRaises(ImportError) as raised: + read() + self.assertIs(failure, raised.exception) + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/setup.py b/paimon-python/setup.py index a706fe0661af..91e1a3d933da 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -262,6 +262,7 @@ def read_requirements(): ], 'torch': [ 'torch', + 'Pillow', ], 'act': LEROBOT_DEPENDENCIES + [ 'Pillow; python_version>="3.10"', From 0aa48f4051d73654193dff2f6e6253aeacd2bf9c Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 16 Sep 2026 14:38:41 +0800 Subject: [PATCH 2/2] fix(python): guard optional dependencies in window transform tests Use importorskip for torch and Pillow, matching the other optional training suites. CI environments without Pillow can collect the remaining tests. Verified 44 related tests, collection without Pillow or torch, and flake8. Co-Authored-By: Codex AI-Model: gpt-6 Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 0/0 AI-Contributed/UT: 6/6 --- paimon-python/pypaimon/tests/window_transforms_test.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/tests/window_transforms_test.py b/paimon-python/pypaimon/tests/window_transforms_test.py index 18493c5f89b4..45df0766bb54 100644 --- a/paimon-python/pypaimon/tests/window_transforms_test.py +++ b/paimon-python/pypaimon/tests/window_transforms_test.py @@ -21,8 +21,10 @@ from unittest.mock import patch import numpy as np -import torch -from PIL import Image +import pytest + +torch = pytest.importorskip("torch") +Image = pytest.importorskip("PIL.Image") from pypaimon.multimodal.window_transforms import images_to_tensor, to_tensor