[python][torch] Add lazy contiguous window dataset by YannByron · Pull Request #9580 · apache/paimon

YannByron · 2026-09-03T09:44:02Z

Summary

Add a generic, snapshot-pinned ContiguousWindowDataset for exposing fixed-size Paimon row windows to PyTorch without loading projected rows or BLOB payloads into memory up front.

This is the contiguous-window portion split from #9466 following review feedback. The ACT benchmark workflow will be submitted separately.

Changes

Add row-ID-backed window indexing with deterministic group/order semantics, configurable stride, and drop/pad/error tail policies.
Read projected values lazily from the pinned snapshot, including scalar and MAP<..., BLOB> payloads and anchor-only columns.
Coalesce overlapping plural reads while keeping logical samples independent for transforms and adapters.
Expose the Dataset through ScanQuery.to_contiguous_window_dataset and document its PyTorch usage and contracts.
Cover snapshot pinning, BLOB handling, batching, mutable values, padding, pickling, and multi-worker DataLoader behavior.

Testing

PYTHONPATH=. /opt/homebrew/bin/python3.13 -m pytest pypaimon/tests/contiguous_window_dataset_test.py -q (16 passed)
PYTHONPATH=. /opt/homebrew/bin/python3.13 -m pytest pypaimon/tests/multimodal_table_test.py -k 'scan_read_blobs or scan_read_and_stream_map_blobs' -q (6 passed, 74 deselected)
/opt/homebrew/bin/python3.13 -m flake8 --config=dev/cfg.ini pypaimon/multimodal/query.py pypaimon/multimodal/window_dataset.py pypaimon/tests/contiguous_window_dataset_test.py
python3 dev/check_license_header.py
/opt/homebrew/bin/python3.13 -m compileall -q pypaimon/multimodal
git diff --check origin/master..HEAD

Notes

No ACT benchmark implementation, packaging, tests, or benchmark documentation is included in this PR.

Expose snapshot-pinned, row-ID-backed windows with lazy BLOB reads and coalesced batch access for PyTorch workloads. AI-Contributed/Feature: 0/655 AI-Contributed/UT: 0/472

JingsongLi

Requirement fit is supported, but the inline findings below should be addressed before merging.

JingsongLi · 2026-09-03T11:14:16Z

+            A snapshot-pinned ``ContiguousWindowDataset``. See that class for
+            padding, mask, transform, and adapter result semantics.
+        """
+        if self._result_factory is not None:


[P1] Reject batch-vector queries instead of silently scanning the base table

This scan-only guard relies on _result_factory, but BatchVectorQuery deliberately leaves _result_factory=None because it implements search through its overridden to_arrow(). Consequently, search_vectors(...).to_contiguous_window_dataset(...) passes this guard, _read_window_index() uses the ordinary scan builder, and the resulting Dataset contains full-table rows while ignoring the vector results and _pre_filter. Please reject this method in _PreFilterQuery, alongside the existing scan-only API overrides, and add a search_vectors() regression test. The public from_query path should enforce the same query-kind check.

JingsongLi · 2026-09-03T11:14:16Z

+        key: None for key in scan_keys
+        if table.options.options.contains_key(key)
+    }
+    options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id)


[P1] Preserve reads through tag-retained snapshot metadata

Replacing every resolved scan.tag-name with scan.snapshot-id loses the tag's retention namespace. After the main snapshot file expires, the tagged scan remains valid and can still build this index, but the first deferred item read fails with ValueError: Snapshot id '1' doesn't exist. Please keep subsequent reads on the tag-retained snapshot (with resolved-ID validation), or introduce a read path that can use the resolved Tag/Snapshot directly, and cover a retained tag whose main snapshot file has expired.

JingsongLi · 2026-09-03T11:14:16Z

+            resolved to their bodies.
+        """
+        columns = self.columns if columns is None else columns
+        query = ScanQuery(self._table)


[P2] Reuse the pinned scan plan instead of replanning every fetch

Every _read_rows() call constructs a new query and executes plan() through either to_arrow() or read_blobs(). This means one full snapshot/manifest planning pass per direct item or modern DataLoader batch, and two when anchor_columns is configured; the work repeats on every epoch and in every worker. I instrumented a two-row table and each repeated item read added two manifest-list reads and one manifest-file read. Please retain authorized pinned-snapshot splits and route requested row-ID ranges to cached/IndexedSplit splits, as the existing lazy TorchDataset does.

JingsongLi · 2026-09-03T11:14:16Z

+    index_query._include_row_id = True
+    read_builder = index_query._configured_read_builder()
+    plan = read_builder.new_scan().plan()
+    index = read_builder.new_read().to_arrow(plan.splits())


[P2] Do not treat authorization-masked row IDs as physical row IDs

This table is the authorization-processed plan output, so _ROW_ID may be masked. With the supported _ROW_ID -> NULL rule, ordinary projected scans still work, but Dataset construction reaches int(None) and misleadingly reports that the group key is unhashable. The existing lazy TorchDataset explicitly disables row-ID-backed access when a QueryAuthSplit masks _ROW_ID. Please inspect the plan for this condition and either use a safe authorized materialized fallback or reject it with an explicit unsupported/permission error; do not bypass the mask to obtain raw IDs.

JingsongLi · 2026-09-03T11:14:16Z

+                raise ValueError(
+                    "%s must contain integer values." % self.order_key)
+            try:
+                grouped[group_key].append((int(order_value), int(row_id)))


[P2] Reject or canonicalize NaN group keys before dictionary grouping

Separate FLOAT/DOUBLE NaN values are non-null and hashable, but nan != nan, so using them directly as dictionary keys splits rows carrying the same NaN group value into separate groups. With two consecutive rows and window_size=2, I reproduced tail='drop' returning no windows and tail='pad' returning two separately padded groups, while Arrow group_by groups those NaNs together. Please reject non-reflexive group values with a clear error, or canonicalize NaNs before grouping, and add FLOAT/DOUBLE coverage.

JingsongLi · 2026-09-03T11:14:16Z

+the snapshot recorded in `dataset.snapshot_id`.
+
+```shell
+pip install pypaimon[torch]


[P2] Quote the Torch extra in this installation command

Default zsh interprets pypaimon[torch] as a glob and aborts with no matches found before pip runs. Please match the existing PyTorch documentation and write pip install 'pypaimon[torch]'.

JingsongLi · 2026-09-03T11:45:36Z

Regarding the schema design, we need to align our design and trade-offs with #9529.

XiaoHongbo-Hope · 2026-09-03T12:11:28Z

+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):


Suggestion: WindowDataset, SequenceDataset or SequenceWindowDataset?

I prefer to keep ContiguousWindowDataset: strict contiguity within each group is its defining guarantee. WindowDataset could imply arbitrary, non-contiguous windows, while SequenceDataset does not describe the map-style window sample returned by __getitem__.

XiaoHongbo-Hope · 2026-09-03T12:22:42Z

+            and (is_blob_type(field.type) or is_map_blob_type(field.type))
+        ]
+        if blob_columns:
+            scalar, blobs = query.read_blobs(


This loses the frame_index of .video values.

Thanks for pointing out this additional case. The current RoboMIND HDF5 ingestion is image-based: it stores individual RGB/depth frames and has no video asset or video-frame-offset semantics. Therefore, this PR has not yet designed the window representation around .video descriptors.

We will treat video support as a follow-up extension and preserve VideoFrameDescriptor.frame_index and the other descriptor metadata when adding it. Until then, .video columns should be explicitly unsupported rather than silently losing metadata.

XiaoHongbo-Hope · 2026-09-03T12:32:17Z

+            and row IDs. Each anchor stores group index, start offset, and the
+            number of real rows available before optional padding.
+        """
+        group_values = index.column(self.group_key).to_pylist()


We can use Arrow or NumPy arrays instead to avoid memory issue?

XiaoHongbo-Hope · 2026-09-03T12:41:06Z

+    .to_contiguous_window_dataset(
+        window_size=16,
+        columns=["state", "action"],
+        group_key="episode_id",


These defaults do not match the native LeRobot schema preserved by #9529, which uses episode_index and frame_index.

JingsongLi · 2026-09-03T12:59:10Z

I suggest defining RoboMIND's native Paimon storage contract early, following the approach in #9529:

Preserve original master/puppet fields, dtypes and image/depth payloads in frames; keep trajectory metadata in episodes.
Add tasks, annotations and calibrations when provided by the source.
Publish components under the same BIGINT version_id tag, with a versions manifest and READY written last.
Treat canonical state/action and normalization statistics as explicitly derived data, not as the native schema.

Start with AgileX without assuming other RoboMIND variants share its layout or sampling axis. This can be a separate schema/ingestion change; the window API should stay generic and be tested against that native contract.

YannByron · 2026-09-03T14:36:41Z

I suggest defining RoboMIND's native Paimon storage contract early, following the approach in #9529:

Preserve original master/puppet fields, dtypes and image/depth payloads in frames; keep trajectory metadata in episodes.

Add tasks, annotations and calibrations when provided by the source.

Publish components under the same BIGINT version_id tag, with a versions manifest and READY written last.

Treat canonical state/action and normalization statistics as explicitly derived data, not as the native schema.

Start with AgileX without assuming other RoboMIND variants share its layout or sampling axis. This can be a separate schema/ingestion change; the window API should stay generic and be tested against that native contract.

robomind_agilex already defines an explicit Paimon storage contract for the RoboMIND AgileX dataset through episode_schema, frame_schema, and feature_stats_schema. It covers dataset metadata, ordered frame data, image/depth payloads, and the normalization data required for training. The AgileX contract does not yet model tasks, annotations, or calibrations.

The current contiguous-window dataset API addresses the requirements of the HDF5-based RoboMIND AgileX dataset while remaining extensible for future sources and modalities.

Reuse the pinned scan plan per projection and prune files by row-id range, so repeated window reads stop replanning the snapshot. Keep tag-pinned reads working after the snapshot file expires, store the window index in Arrow/NumPy arrays, and reject inputs the window contract cannot honor: search queries, masked _ROW_ID, NaN group keys and video frame columns whose frame metadata a window read would drop. Default keys now follow the LeRobot-native episode_index/frame_index. AI-Contributed/Feature: 0/496 AI-Contributed/UT: 0/140

JingsongLi · 2026-09-04T02:29:17Z

+            table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"})
+            if self._blob_columns else table
+        )
+        self._splits = self._new_read_builder().new_scan().plan().splits()


[P1] Revalidate the retained tag on every deferred plan

_pin_table verifies the tag only during dataset construction, but this lazy plan resolves the mutable tag name again and immediately discards plan.snapshot_id. FileStoreTable.replace_tag can therefore move the tag between those two points. I reproduced a dataset whose recorded snapshot_id stayed at snapshot 1 while the first lazy value read returned data updated only in snapshot 2. Pickled/DataLoader workers are also exposed because __getstate__ clears the plans.

Please pass the expected snapshot ID into _PinnedRowIdPlan, keep the full Plan, and fail closed unless plan.snapshot_id still equals the indexed snapshot before caching its splits (or plan from an immutable resolved snapshot). A regression test should replace the tag after construction and before the first projection read, including the post-pickle path.

JingsongLi · 2026-09-04T02:29:20Z

+    index_query._include_row_id = True
+    read_builder = index_query._configured_read_builder()
+    plan = read_builder.new_scan().plan()
+    splits = plan.splits()


[P2] Reject authorization masks on the grouping and ordering keys

The index scan applies column masking, but the validation below rejects only a masked _ROW_ID. A non-injective mask on group_key can merge real groups before sorting and anchor construction; a mask on order_key can similarly change contiguity. I reproduced two groups a:[0,1] and b:[2,3] with a constant mask on episode, which produced a window crossing the group boundary ([0,1,102]) despite the API guarantee that windows never cross groups.

Please inspect QueryAuthSplit.auth_result.column_masking and reject masks on group_key or order_key unless the index can be built from unmasked keys without weakening authorization. Add a regression test using a non-null constant mask; null-only tests would be caught by the existing null validation and miss this case.

JingsongLi · 2026-09-04T02:30:18Z

+        read_builder = self._new_read_builder()
+        predicate = read_builder.new_predicate_builder().is_in(
+            SpecialFields.ROW_ID.name, requested)
+        arrow = read_builder.with_filter(predicate).new_read().to_arrow(splits)


[P2] Preserve exact row ranges when reading cached splits

The cached plan contains unfiltered DataSplits, and _prune_split_files only drops whole files. This later _ROW_ID IN (...) predicate cannot recover native row-range pushdown for a data-evolution read: DataEvolutionSplitRead._push_down_predicate() deliberately returns None, and its row_ranges remains None unless the supplied split is an IndexedSplit. As a result, every overlapping file is decoded before FilterRecordBatchReader (and the Arrow filter below) removes unrelated rows.

That is on the documented shuffled training hot path: one batch can request many scattered windows and repeatedly decode large files. Please build a split-range index once, intersect requested ranges with matching splits, and pass IndexedSplits while preserving QueryAuthSplit; TorchDataset._SplitRangeIndex / _select_splits is an existing implementation to reuse. Add a multi-file test that observes format-reader row ranges or decoded-row counts, since the six-row logical-call tests do not expose full-file decoding.

[python][torch] Add lazy contiguous window dataset

b9f1c02

Expose snapshot-pinned, row-ID-backed windows with lazy BLOB reads and coalesced batch access for PyTorch workloads. AI-Contributed/Feature: 0/655 AI-Contributed/UT: 0/472

JingsongLi reviewed Sep 3, 2026

View reviewed changes

XiaoHongbo-Hope reviewed Sep 3, 2026

View reviewed changes

JingsongLi reviewed Sep 4, 2026

View reviewed changes

		from pypaimon.table.special_fields import SpecialFields


		class ContiguousWindowDataset(Dataset):

Uh oh!

Conversation

YannByron commented Sep 3, 2026

Summary

Changes

Testing

Notes

Uh oh!

JingsongLi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JingsongLi commented Sep 3, 2026

Uh oh!

XiaoHongbo-Hope Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

XiaoHongbo-Hope Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JingsongLi commented Sep 3, 2026

Uh oh!

YannByron commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

XiaoHongbo-Hope Sep 3, 2026 •

edited

Loading

XiaoHongbo-Hope Sep 3, 2026 •

edited

Loading

YannByron commented Sep 3, 2026 •

edited

Loading