Skip to content

Commit ba04065

Browse files
authored
[python][torch] Add lazy contiguous window dataset (#9580)
1 parent 0f320ef commit ba04065

7 files changed

Lines changed: 1613 additions & 46 deletions

File tree

‎docs/docs/pypaimon/multimodal-api.mdx‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -905,6 +905,68 @@ Notes:
905905
few large reads); scattered point reads coalesce less.
906906
- Blob reads are available only on `scan()`, not on the `search()` queries.
907907

908+
### Contiguous windows for PyTorch
909+
910+
Install the `torch` extra, then use `to_contiguous_window_dataset` to expose
911+
map-style windows without loading the selected rows or BLOB payloads into Python
912+
memory up front. The Dataset builds a compact index from the group column, order
913+
column, and Paimon row IDs. Each `__getitem__` call fetches only that window from
914+
the snapshot recorded in `dataset.snapshot_id`.
915+
916+
```shell
917+
pip install 'pypaimon[torch]'
918+
```
919+
920+
```python
921+
import torch
922+
923+
924+
def float32_window(values):
925+
return torch.tensor(values, dtype=torch.float32)
926+
927+
928+
windows = (
929+
frames.scan()
930+
.where("split = 'train'")
931+
.to_contiguous_window_dataset(
932+
window_size=16,
933+
columns=["state", "action"],
934+
group_key="episode_index",
935+
order_key="frame_index",
936+
tail="pad",
937+
column_transforms={
938+
"state": float32_window,
939+
"action": float32_window,
940+
},
941+
)
942+
)
943+
944+
sample = windows[0]
945+
assert sample["action"].shape == (16, action_size)
946+
assert sample["is_pad"].shape == (16,)
947+
```
948+
949+
The group and order keys in a sample identify the window anchor. Every projected
950+
column contains the whole window. With `tail="drop"`, only full windows are
951+
exposed. With `tail="pad"`, every real row is an anchor; missing suffix values
952+
repeat the last real value by default and `is_pad` is `True` exactly at those
953+
positions. With `tail="error"`, construction fails if any scheduled anchor is
954+
incomplete. Use `pad_values` to override the repeated value for individual
955+
columns. Anchors advance by `stride`, which defaults to one row.
956+
957+
`column_transforms` receive one padded Python list per projected column. This is
958+
where applications define tensor dtype and shape or decode BLOB bytes. The
959+
optional `adapter` receives the resulting sample mapping and can rename or
960+
combine fields for a model-specific batch contract. The core Dataset does not
961+
know model field names, image formats, or normalization rules. Top-level
962+
functions and callable classes are recommended for transforms and adapters so
963+
the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader`
964+
instances.
965+
966+
Columns configured by `video-frame-field` are rejected: a window read would drop
967+
the `frame_index` and other metadata carried by their `VideoFrameDescriptor`
968+
values. Read those columns with `to_torch()` instead.
969+
908970
### Distributed BLOB processing with Ray
909971

910972
For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB

‎docs/docs/pypaimon/pytorch.md‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,55 @@ embedded frame ordinals keep frame mapping out of the normal data file. Use
157157
physical video ranges and cache decoder sessions per worker. See
158158
[Multimodal API: Video Frame Storage](multimodal-api#video-frame-storage)
159159
for the write path and a complete decoder example.
160+
## Contiguous Windows
160161

162+
Use a map-style `ContiguousWindowDataset` when training samples are fixed-size
163+
windows which must not cross a sequence boundary. The dataset builds an index
164+
from only the group column, order column, and Paimon row IDs. Projected values,
165+
including BLOB payloads, are read from the pinned snapshot when a sample is
166+
requested; they are not retained in the index.
167+
168+
```python
169+
from torch.utils.data import DataLoader
170+
171+
dataset = (
172+
frames.scan()
173+
.to_contiguous_window_dataset(
174+
window_size=16,
175+
columns=["state", "image"],
176+
anchor_columns=["image"],
177+
group_key="episode_index",
178+
order_key="frame_index",
179+
tail="pad",
180+
)
181+
)
182+
183+
loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True)
184+
```
185+
186+
Each item contains the group and order keys, one list for each requested
187+
column, and a boolean `is_pad` tensor where `True` marks padding. Padding
188+
repeats the final real value by default; `pad_values` can override individual
189+
columns. Columns named in `anchor_columns` contain only the first row's value,
190+
which is useful when an observation applies to a full action window. Use
191+
`column_transforms` to convert column lists to tensors and
192+
`adapter` to produce a model-specific sample mapping. Keep these callbacks
193+
picklable when using multiple DataLoader workers.
194+
195+
Scheduled anchors start at row zero and advance by `stride` (default `1`).
196+
`tail="drop"` omits incomplete windows, `tail="pad"` includes and pads them,
197+
and `tail="error"` rejects a sequence with any scheduled incomplete window.
198+
Rows are sorted by `order_key` inside each `group_key` value. Order values must
199+
be integers which increase by exactly one; duplicates and missing steps are
200+
rejected, and windows never cross groups. The resolved Paimon
201+
snapshot is pinned for the lifetime of the dataset, so later commits cannot
202+
change its index or sample contents. A dataset pinned through `tag_name` fails
203+
its reads if the tag is moved to another snapshot, rather than mixing rows from
204+
the two snapshots.
205+
206+
Columns configured by `video-frame-field` are rejected: a window read would drop
207+
the `frame_index` and other metadata carried by their `VideoFrameDescriptor`
208+
values. Read those columns with `to_torch()` instead.
161209
## File Format Metadata Cache
162210

163211
Reusable PyArrow Dataset metadata is cached across reads. Configure its estimated

‎paimon-python/pypaimon/multimodal/query.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,66 @@ def to_torch(
164164
max_buffer_input_splits=max_buffer_input_splits,
165165
)
166166

167+
def to_contiguous_window_dataset(
168+
self,
169+
*,
170+
window_size,
171+
columns=None,
172+
anchor_columns=None,
173+
group_key="episode_index",
174+
order_key="frame_index",
175+
stride=1,
176+
tail="drop",
177+
column_transforms=None,
178+
pad_values=None,
179+
adapter=None,
180+
blob_parallelism=64):
181+
"""Build a snapshot-pinned, map-style Dataset of contiguous rows.
182+
183+
The Dataset indexes only ``group_key``, ``order_key``, and Paimon row
184+
IDs, then reads projected values on demand. Columns listed in
185+
``anchor_columns`` are provided to ``column_transforms`` as one-element
186+
lists read from the first row of each window; ``adapter`` receives the
187+
transformed values. ``order_key`` must contain non-null integers that
188+
increase by exactly one within each group. The Dataset sorts rows within
189+
each group and never creates a window across groups.
190+
191+
Args:
192+
window_size: Number of rows in a complete window.
193+
columns: Value columns to return, excluding the group and order
194+
keys. The scan projection is used when omitted.
195+
anchor_columns: Subset of ``columns`` read only from the window's
196+
first row.
197+
group_key: Column identifying an independent row sequence.
198+
order_key: Integer position column within each group.
199+
stride: Distance between scheduled window starts.
200+
tail: Handling for incomplete final windows: ``drop``, ``pad``, or
201+
``error``.
202+
column_transforms: Per-column callables applied to value lists.
203+
pad_values: Optional replacement values used by ``tail='pad'``.
204+
adapter: Callable that converts the complete sample mapping.
205+
blob_parallelism: Maximum concurrent BLOB body reads per fetch.
206+
207+
Returns:
208+
A snapshot-pinned ``ContiguousWindowDataset``. See that class for
209+
padding, mask, transform, and adapter result semantics.
210+
"""
211+
from pypaimon.multimodal.window_dataset import ContiguousWindowDataset
212+
return ContiguousWindowDataset(
213+
self,
214+
window_size=window_size,
215+
columns=columns,
216+
anchor_columns=anchor_columns,
217+
group_key=group_key,
218+
order_key=order_key,
219+
stride=stride,
220+
tail=tail,
221+
column_transforms=column_transforms,
222+
pad_values=pad_values,
223+
adapter=adapter,
224+
blob_parallelism=blob_parallelism,
225+
)
226+
167227
def to_ray(
168228
self,
169229
*,
@@ -418,6 +478,12 @@ def to_arrow_batch_reader(self, *args, **kwargs):
418478
"not search queries."
419479
)
420480

481+
def to_contiguous_window_dataset(self, *args, **kwargs):
482+
raise TypeError(
483+
"to_contiguous_window_dataset is only supported on scan(), "
484+
"not search queries."
485+
)
486+
421487

422488
class VectorQuery(_PreFilterQuery):
423489
"""Chainable query wrapper for vector global-index search."""

0 commit comments

Comments
 (0)