-
Notifications
You must be signed in to change notification settings - Fork 1.4k
[python][torch] Add lazy contiguous window dataset #9580
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -866,6 +866,64 @@ Notes: | |
| few large reads); scattered point reads coalesce less. | ||
| - Blob reads are available only on `scan()`, not on the `search()` queries. | ||
|
|
||
| ### 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`. | ||
|
|
||
| ```shell | ||
| pip install pypaimon[torch] | ||
| ``` | ||
|
|
||
| ```python | ||
| import torch | ||
|
|
||
|
|
||
| def float32_window(values): | ||
| return torch.tensor(values, dtype=torch.float32) | ||
|
|
||
|
|
||
| windows = ( | ||
| frames.scan() | ||
| .where("split = 'train'") | ||
| .to_contiguous_window_dataset( | ||
| window_size=16, | ||
| columns=["state", "action"], | ||
| group_key="episode_id", | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These defaults do not match the native LeRobot schema preserved by #9529, which uses episode_index and frame_index. |
||
| order_key="step_idx", | ||
| tail="pad", | ||
| column_transforms={ | ||
| "state": float32_window, | ||
| "action": float32_window, | ||
| }, | ||
| ) | ||
| ) | ||
|
|
||
| sample = windows[0] | ||
| assert sample["action"].shape == (16, action_size) | ||
| assert sample["is_pad"].shape == (16,) | ||
| ``` | ||
|
|
||
| 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. | ||
|
|
||
| `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. | ||
|
|
||
| ### Distributed BLOB processing with Ray | ||
|
|
||
| For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -164,6 +164,70 @@ def to_torch( | |
| max_buffer_input_splits=max_buffer_input_splits, | ||
| ) | ||
|
|
||
| def to_contiguous_window_dataset( | ||
| self, | ||
| *, | ||
| window_size, | ||
| columns=None, | ||
| anchor_columns=None, | ||
| group_key="episode_id", | ||
| order_key="step_idx", | ||
| stride=1, | ||
| tail="drop", | ||
| column_transforms=None, | ||
| pad_values=None, | ||
| adapter=None, | ||
| blob_parallelism=64): | ||
| """Build a snapshot-pinned, map-style Dataset of contiguous rows. | ||
|
|
||
| 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 | ||
| 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. | ||
| 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. | ||
| 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``. | ||
| column_transforms: Per-column callables applied to value lists. | ||
| pad_values: Optional replacement values used by ``tail='pad'``. | ||
| adapter: Callable that converts the complete sample mapping. | ||
| blob_parallelism: Maximum concurrent BLOB body reads per fetch. | ||
|
|
||
| Returns: | ||
| A snapshot-pinned ``ContiguousWindowDataset``. See that class for | ||
| padding, mask, transform, and adapter result semantics. | ||
| """ | ||
| if self._result_factory is not None: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P1] Reject batch-vector queries instead of silently scanning the base table This scan-only guard relies on |
||
| raise TypeError( | ||
| "to_contiguous_window_dataset is only supported on scan(), " | ||
| "not search queries.") | ||
| from pypaimon.multimodal.window_dataset import ContiguousWindowDataset | ||
| return ContiguousWindowDataset( | ||
| self, | ||
| window_size=window_size, | ||
| columns=columns, | ||
| anchor_columns=anchor_columns, | ||
| group_key=group_key, | ||
| order_key=order_key, | ||
| stride=stride, | ||
| tail=tail, | ||
| column_transforms=column_transforms, | ||
| pad_values=pad_values, | ||
| adapter=adapter, | ||
| blob_parallelism=blob_parallelism, | ||
| ) | ||
|
|
||
| def to_ray( | ||
| self, | ||
| *, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[P2] Quote the Torch extra in this installation command
Default zsh interprets
pypaimon[torch]as a glob and aborts withno matches foundbefore pip runs. Please match the existing PyTorch documentation and writepip install 'pypaimon[torch]'.