This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 1fc0f17ad2dbb94c17372a348cfbc4a0a9ec133b Author: Yann Byron <[email protected]> AuthorDate: Thu Sep 17 11:39:54 2026 +0800 [python] Extend ContiguousWindowDataset for frame-level training (#9875) --- docs/docs/pypaimon/multimodal-reading.md | 119 +++++++-- docs/docs/pypaimon/pytorch.md | 55 ++-- paimon-python/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 +++++++ paimon-python/pypaimon/tests/act_runner_test.py | 20 +- .../tests/contiguous_window_dataset_test.py | 236 +++++++++++++++++ .../pypaimon/tests/window_transforms_test.py | 111 ++++++++ paimon-python/setup.py | 1 + 11 files changed, 828 insertions(+), 182 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-reading.md b/docs/docs/pypaimon/multimodal-reading.md index 234dd4f6e2..d3df05a9bd 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 `<column>_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 6f973ade3f..26770bdbba 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 `<column>_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 18b50664b6..161edadbbc 100644 --- a/paimon-python/pypaimon/benchmark/act/paimon.py +++ b/paimon-python/pypaimon/benchmark/act/paimon.py @@ -52,7 +52,7 @@ class PaimonACTAdapter: 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 @@ class PaimonACTAdapter: "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 eca82cbfd6..006ca6357b 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 @@ -1438,14 +1437,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: @@ -1455,21 +1450,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 489ff64152..e21edb2f1e 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -167,42 +167,58 @@ class ScanQuery: 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 @@ class ScanQuery: 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 2a63845962..ec5f639969 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 @@ from pypaimon.utils.range import Range 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 + ``<column>_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 @@ class ContiguousWindowDataset(Dataset): 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 @@ class ContiguousWindowDataset(Dataset): 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 @@ class ContiguousWindowDataset(Dataset): 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 @@ class ContiguousWindowDataset(Dataset): 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 @@ class ContiguousWindowDataset(Dataset): 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 0000000000..3385b8525d --- /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 57d64bf8c1..a7fbc69e92 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 c3582ea7b5..c0e5d4d7dd 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 @@ class ContiguousWindowDatasetTest(unittest.TestCase): ) ) + 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 0000000000..45df0766bb --- /dev/null +++ b/paimon-python/pypaimon/tests/window_transforms_test.py @@ -0,0 +1,111 @@ +# 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 pytest + +torch = pytest.importorskip("torch") +Image = pytest.importorskip("PIL.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 60cfe979d3..246f77714c 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -263,6 +263,7 @@ setup( ], 'torch': [ 'torch', + 'Pillow', ], 'act': LEROBOT_DEPENDENCIES + [ 'Pillow; python_version>="3.10"',
