This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 4bc8200f46 [python] Extend ContiguousWindowDataset for frame-level
training (#9875)
4bc8200f46 is described below
commit 4bc8200f46d129a7d484e712ede2f7354433d95e
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 267218a44d..1af50be1e5 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"',
