You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@paimon.apache.org by "YannByron (via GitHub)" <gi...@apache.org> on 2026/09/03 09:44:04 UTC
[PR] [python][torch] Add lazy contiguous window dataset [paimon]
YannByron opened a new pull request, #9580:
URL: https://github.com/apache/paimon/pull/9580
## 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
- [x] `PYTHONPATH=. /opt/homebrew/bin/python3.13 -m pytest pypaimon/tests/contiguous_window_dataset_test.py -q` (16 passed)
- [x] `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)
- [x] `/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`
- [x] `python3 dev/check_license_header.py`
- [x] `/opt/homebrew/bin/python3.13 -m compileall -q pypaimon/multimodal`
- [x] `git diff --check origin/master..HEAD`
## Notes
No ACT benchmark implementation, packaging, tests, or benchmark documentation is included in this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3930371181
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,707 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.blob_read import fetch_blob_bodies
+from pypaimon.multimodal.query import ScanQuery, _PreFilterQuery
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.utils.range import Range
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``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")
+
+ def __init__(
+ self,
+ query,
+ *,
+ window_size,
+ columns=None,
+ anchor_columns=None,
+ group_key="episode_index",
+ order_key="frame_index",
+ stride=1,
+ tail="drop",
+ column_transforms=None,
+ pad_values=None,
+ adapter=None,
+ blob_parallelism=64):
+ 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.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._plans = {}
+ self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return int(self._anchor_groups.size)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def __getstate__(self):
+ # Cached plans hold planning state of one process; each DataLoader
+ # worker plans the pinned snapshot once for itself.
+ state = self.__dict__.copy()
+ state["_plans"] = {}
+ return state
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self)
+ if index < 0 or index >= len(self):
+ raise IndexError("window index out of range")
+
+ 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
+ sample = {
+ self.group_key: self._group_keys[group_index],
+ self.order_key: int(self._group_first_orders[group_index]) + start,
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows, then store row IDs, groups, and window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ The index is kept as NumPy arrays of row IDs, per-group offsets, and
+ window anchors, plus one Python group value per group. Order values are
+ not stored per row: contiguity makes them the group's first value plus
+ the offset inside the group.
+ """
+ group_column = index.column(self.group_key)
+ order_column = index.column(self.order_key)
+ row_id_column = index.column(SpecialFields.ROW_ID.name)
+ if row_id_column.null_count:
+ raise ValueError(
+ "ContiguousWindowDataset requires readable Paimon row IDs, "
+ "but %s contains null values." % SpecialFields.ROW_ID.name)
+ if group_column.null_count:
+ raise ValueError(
+ "%s must not contain null values." % self.group_key)
+ if order_column.null_count:
+ raise ValueError(
+ "%s must not contain null values." % self.order_key)
+ if not pa.types.is_integer(order_column.type):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ if pa.types.is_floating(group_column.type) and pc.any(
+ pc.is_nan(group_column)).as_py():
+ raise ValueError(
+ "%s must not contain NaN values, which never compare equal to "
+ "themselves and would split one group." % self.group_key)
+
+ try:
+ ordered = index.sort_by([
+ (self.group_key, "ascending"),
+ (self.order_key, "ascending"),
+ ])
+ except pa.ArrowNotImplementedError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+
+ group_values = _contiguous_array(ordered.column(self.group_key))
+ order_values = _integer_numpy(
+ _contiguous_array(ordered.column(self.order_key)))
+ self._row_ids = _integer_numpy(
+ _contiguous_array(ordered.column(SpecialFields.ROW_ID.name)))
+ self._group_starts = _group_starts(group_values)
+ self._group_lengths = np.diff(
+ np.append(self._group_starts, len(self._row_ids)))
+ self._group_keys = group_values.take(
+ pa.array(self._group_starts, type=pa.int64())).to_pylist()
+ self._group_first_orders = order_values[self._group_starts]
+ self._validate_contiguity(order_values)
+ self._anchor_groups, self._anchor_starts = self._build_anchors()
+
+ def _validate_contiguity(self, order_values):
+ if len(order_values) < 2:
+ return
+ same_group = np.ones(len(order_values) - 1, dtype=bool)
+ same_group[self._group_starts[1:] - 1] = False
+ steps = np.diff(order_values)
+ duplicate = same_group & (steps == 0)
+ if duplicate.any():
+ position = int(np.flatnonzero(duplicate)[0])
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (self._group_of(position),
+ int(order_values[position]), self.order_key))
+ broken = same_group & (steps != 1)
+ if broken.any():
+ position = int(np.flatnonzero(broken)[0])
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (self._group_of(position), self.order_key,
+ int(order_values[position]), int(order_values[position + 1])))
+
+ def _group_of(self, position):
+ group_index = int(np.searchsorted(
+ self._group_starts, position, side="right")) - 1
+ return self._group_keys[group_index]
+
+ def _build_anchors(self):
+ groups = []
+ starts = []
+ 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)
+ if incomplete.size:
+ if self.tail == "error":
+ first = int(incomplete[0])
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%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]
+ if positions.size:
+ groups.append(np.full(positions.size, group_index, dtype=np.int64))
+ starts.append(positions)
+ if not groups:
+ empty = np.zeros(0, dtype=np.int64)
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ rows = self._plan_for(columns).read(row_ids)
+ row_id_column = SpecialFields.ROW_ID.name
+ by_row_id = {}
+ for row in rows:
+ by_row_id[int(row.pop(row_id_column))] = row
+ missing = [row_id for row_id in row_ids if row_id not in by_row_id]
+ if missing:
+ raise RuntimeError(
+ "Pinned snapshot %s did not return indexed row IDs %s."
+ % (self.snapshot_id, missing))
+ return [by_row_id[row_id] for row_id in row_ids]
+
+ def _plan_for(self, columns):
+ key = tuple(columns)
+ plan = self._plans.get(key)
+ if plan is None:
+ plan = _PinnedRowIdPlan(self._table, columns, self.blob_parallelism)
+ self._plans[key] = plan
+ return plan
+
+
+class _PinnedRowIdPlan:
+ """One authorized scan plan of the pinned snapshot, read by row ID.
+
+ Planning happens once per projection. Every read then narrows the cached
+ splits to the files covering the requested row IDs, so repeated item and
+ batch reads never revisit the snapshot's manifests.
+ """
+
+ def __init__(self, table, columns, blob_parallelism):
+ self._blob_parallelism = blob_parallelism
+ self._blob_columns = [
+ field.name for field in table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ self._map_blob_columns = {
+ field.name for field in table.fields
+ if field.name in self._blob_columns and is_map_blob_type(field.type)
+ }
+ blob_column_set = set(self._blob_columns)
+ self._projection = (
+ [name for name in columns if name not in blob_column_set]
+ + [SpecialFields.ROW_ID.name]
+ + self._blob_columns
+ )
+ self._table = (
+ table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"})
+ if self._blob_columns else table
+ )
+ self._splits = self._new_read_builder().new_scan().plan().splits()
Review Comment:
[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.
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,707 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.blob_read import fetch_blob_bodies
+from pypaimon.multimodal.query import ScanQuery, _PreFilterQuery
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.utils.range import Range
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``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")
+
+ def __init__(
+ self,
+ query,
+ *,
+ window_size,
+ columns=None,
+ anchor_columns=None,
+ group_key="episode_index",
+ order_key="frame_index",
+ stride=1,
+ tail="drop",
+ column_transforms=None,
+ pad_values=None,
+ adapter=None,
+ blob_parallelism=64):
+ 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.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._plans = {}
+ self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return int(self._anchor_groups.size)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def __getstate__(self):
+ # Cached plans hold planning state of one process; each DataLoader
+ # worker plans the pinned snapshot once for itself.
+ state = self.__dict__.copy()
+ state["_plans"] = {}
+ return state
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self)
+ if index < 0 or index >= len(self):
+ raise IndexError("window index out of range")
+
+ 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
+ sample = {
+ self.group_key: self._group_keys[group_index],
+ self.order_key: int(self._group_first_orders[group_index]) + start,
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows, then store row IDs, groups, and window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ The index is kept as NumPy arrays of row IDs, per-group offsets, and
+ window anchors, plus one Python group value per group. Order values are
+ not stored per row: contiguity makes them the group's first value plus
+ the offset inside the group.
+ """
+ group_column = index.column(self.group_key)
+ order_column = index.column(self.order_key)
+ row_id_column = index.column(SpecialFields.ROW_ID.name)
+ if row_id_column.null_count:
+ raise ValueError(
+ "ContiguousWindowDataset requires readable Paimon row IDs, "
+ "but %s contains null values." % SpecialFields.ROW_ID.name)
+ if group_column.null_count:
+ raise ValueError(
+ "%s must not contain null values." % self.group_key)
+ if order_column.null_count:
+ raise ValueError(
+ "%s must not contain null values." % self.order_key)
+ if not pa.types.is_integer(order_column.type):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ if pa.types.is_floating(group_column.type) and pc.any(
+ pc.is_nan(group_column)).as_py():
+ raise ValueError(
+ "%s must not contain NaN values, which never compare equal to "
+ "themselves and would split one group." % self.group_key)
+
+ try:
+ ordered = index.sort_by([
+ (self.group_key, "ascending"),
+ (self.order_key, "ascending"),
+ ])
+ except pa.ArrowNotImplementedError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+
+ group_values = _contiguous_array(ordered.column(self.group_key))
+ order_values = _integer_numpy(
+ _contiguous_array(ordered.column(self.order_key)))
+ self._row_ids = _integer_numpy(
+ _contiguous_array(ordered.column(SpecialFields.ROW_ID.name)))
+ self._group_starts = _group_starts(group_values)
+ self._group_lengths = np.diff(
+ np.append(self._group_starts, len(self._row_ids)))
+ self._group_keys = group_values.take(
+ pa.array(self._group_starts, type=pa.int64())).to_pylist()
+ self._group_first_orders = order_values[self._group_starts]
+ self._validate_contiguity(order_values)
+ self._anchor_groups, self._anchor_starts = self._build_anchors()
+
+ def _validate_contiguity(self, order_values):
+ if len(order_values) < 2:
+ return
+ same_group = np.ones(len(order_values) - 1, dtype=bool)
+ same_group[self._group_starts[1:] - 1] = False
+ steps = np.diff(order_values)
+ duplicate = same_group & (steps == 0)
+ if duplicate.any():
+ position = int(np.flatnonzero(duplicate)[0])
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (self._group_of(position),
+ int(order_values[position]), self.order_key))
+ broken = same_group & (steps != 1)
+ if broken.any():
+ position = int(np.flatnonzero(broken)[0])
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (self._group_of(position), self.order_key,
+ int(order_values[position]), int(order_values[position + 1])))
+
+ def _group_of(self, position):
+ group_index = int(np.searchsorted(
+ self._group_starts, position, side="right")) - 1
+ return self._group_keys[group_index]
+
+ def _build_anchors(self):
+ groups = []
+ starts = []
+ 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)
+ if incomplete.size:
+ if self.tail == "error":
+ first = int(incomplete[0])
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%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]
+ if positions.size:
+ groups.append(np.full(positions.size, group_index, dtype=np.int64))
+ starts.append(positions)
+ if not groups:
+ empty = np.zeros(0, dtype=np.int64)
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ rows = self._plan_for(columns).read(row_ids)
+ row_id_column = SpecialFields.ROW_ID.name
+ by_row_id = {}
+ for row in rows:
+ by_row_id[int(row.pop(row_id_column))] = row
+ missing = [row_id for row_id in row_ids if row_id not in by_row_id]
+ if missing:
+ raise RuntimeError(
+ "Pinned snapshot %s did not return indexed row IDs %s."
+ % (self.snapshot_id, missing))
+ return [by_row_id[row_id] for row_id in row_ids]
+
+ def _plan_for(self, columns):
+ key = tuple(columns)
+ plan = self._plans.get(key)
+ if plan is None:
+ plan = _PinnedRowIdPlan(self._table, columns, self.blob_parallelism)
+ self._plans[key] = plan
+ return plan
+
+
+class _PinnedRowIdPlan:
+ """One authorized scan plan of the pinned snapshot, read by row ID.
+
+ Planning happens once per projection. Every read then narrows the cached
+ splits to the files covering the requested row IDs, so repeated item and
+ batch reads never revisit the snapshot's manifests.
+ """
+
+ def __init__(self, table, columns, blob_parallelism):
+ self._blob_parallelism = blob_parallelism
+ self._blob_columns = [
+ field.name for field in table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ self._map_blob_columns = {
+ field.name for field in table.fields
+ if field.name in self._blob_columns and is_map_blob_type(field.type)
+ }
+ blob_column_set = set(self._blob_columns)
+ self._projection = (
+ [name for name in columns if name not in blob_column_set]
+ + [SpecialFields.ROW_ID.name]
+ + self._blob_columns
+ )
+ self._table = (
+ table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"})
+ if self._blob_columns else table
+ )
+ self._splits = self._new_read_builder().new_scan().plan().splits()
+ _reject_masked_row_ids(self._splits)
+
+ def read(self, row_ids):
+ """Return raw row dictionaries, including ``_ROW_ID``, for ``row_ids``."""
+ requested = list(dict.fromkeys(row_ids))
+ ranges = Range.to_ranges(requested)
+ splits = [
+ pruned for pruned in (
+ _prune_split_files(split, ranges) for split in self._splits)
+ if pruned is not None
+ ]
+ if not splits:
+ return []
+ 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)
+ # Row-ID pruning happens per file and per split; drop the remaining
+ # rows before BLOB bodies are fetched for them.
+ row_id_column = arrow.column(SpecialFields.ROW_ID.name)
+ arrow = arrow.filter(pc.is_in(
+ row_id_column,
+ value_set=pa.array(requested, type=row_id_column.type)))
+ if not self._blob_columns:
+ return arrow.to_pylist()
+
+ bodies = fetch_blob_bodies(
+ self._table.file_io,
+ arrow.select(self._blob_columns).to_pydict(),
+ self._blob_columns,
+ self._blob_parallelism,
+ self._map_blob_columns)
+ blob_column_set = set(self._blob_columns)
+ rows = arrow.select([
+ name for name in arrow.column_names if name not in blob_column_set
+ ]).to_pylist()
+ for name in self._blob_columns:
+ values = bodies[name]
+ if len(values) != len(rows):
+ raise RuntimeError(
+ "BLOB column %s is not row-aligned with a window read."
+ % name)
+ for row, value in zip(rows, values):
+ row[name] = value
+ return rows
+
+ def _new_read_builder(self):
+ return self._table.new_read_builder().with_projection(self._projection)
+
+
+def _prune_split_files(split, ranges):
+ """Drop files of ``split`` outside ``ranges``; ``None`` when nothing is left."""
+ auth_result = None
+ data_split = split
+ if isinstance(data_split, QueryAuthSplit):
+ auth_result = data_split.auth_result
+ data_split = data_split.split
+ filter_file = getattr(data_split, "filter_file", None)
+ if filter_file is None:
+ return split
+ pruned = filter_file(lambda meta: _file_overlaps(meta, ranges))
+ if pruned is None:
+ return None
+ if auth_result is None:
+ return pruned
+ return QueryAuthSplit(pruned, auth_result)
+
+
+def _file_overlaps(meta, ranges):
+ file_range = meta.row_id_range()
+ if file_range is None:
+ return True
+ return bool(Range.and_(ranges, [file_range]))
+
+
+def _reject_masked_row_ids(splits):
+ for split in splits:
+ if not isinstance(split, QueryAuthSplit):
+ continue
+ masking = getattr(split.auth_result, "column_masking", None)
+ if masking and SpecialFields.ROW_ID.name in masking:
+ raise ValueError(
+ "ContiguousWindowDataset requires readable Paimon row IDs, but "
+ "query authorization masks %s for this table; read the rows "
+ "with scan().to_arrow() instead."
+ % SpecialFields.ROW_ID.name)
+
+
+def _contiguous_array(column):
+ combined = column.combine_chunks()
+ if not isinstance(combined, pa.ChunkedArray):
+ return combined
+ if combined.num_chunks == 1:
+ return combined.chunk(0)
+ if not combined.num_chunks:
+ return pa.nulls(0, type=combined.type)
+ return pa.concat_arrays(list(combined.iterchunks()))
+
+
+def _integer_numpy(array):
+ # Cast narrow widths so differences of adjacent values cannot overflow.
+ if not pa.types.is_uint64(array.type):
+ array = array.cast(pa.int64())
+ return array.to_numpy(zero_copy_only=False)
+
+
+def _group_starts(group_values):
+ count = len(group_values)
+ if not count:
+ return np.zeros(0, dtype=np.int64)
+ changed = pc.not_equal(
+ group_values.slice(1), group_values.slice(0, count - 1))
+ return np.concatenate((
+ [0],
+ np.flatnonzero(changed.to_numpy(zero_copy_only=False)) + 1,
+ )).astype(np.int64, copy=False)
+
+
+def _read_window_index(query, group_key, order_key):
+ index_query = copy.copy(query)
+ index_query._projection = [group_key, order_key]
+ index_query._include_row_id = True
+ read_builder = index_query._configured_read_builder()
+ plan = read_builder.new_scan().plan()
+ splits = plan.splits()
Review Comment:
[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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org>.
XiaoHongbo-Hope commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3924363345
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
Review Comment:
Suggestion: WindowDataset, SequenceDataset or SequenceWindowDataset?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org>.
XiaoHongbo-Hope commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3924616350
##########
docs/docs/pypaimon/multimodal-api.mdx:
##########
@@ -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",
Review Comment:
These defaults do not match the native LeRobot schema preserved by #9529, which uses episode_index and frame_index.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "YannByron (via GitHub)" <gi...@apache.org>.
YannByron commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3925219184
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
Review Comment:
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__`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org>.
XiaoHongbo-Hope commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3924363345
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
Review Comment:
Suggestion: SequenceDataset or SequenceWindowDataset?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3930374955
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,707 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.blob_read import fetch_blob_bodies
+from pypaimon.multimodal.query import ScanQuery, _PreFilterQuery
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.utils.range import Range
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``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")
+
+ def __init__(
+ self,
+ query,
+ *,
+ window_size,
+ columns=None,
+ anchor_columns=None,
+ group_key="episode_index",
+ order_key="frame_index",
+ stride=1,
+ tail="drop",
+ column_transforms=None,
+ pad_values=None,
+ adapter=None,
+ blob_parallelism=64):
+ 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.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._plans = {}
+ self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return int(self._anchor_groups.size)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def __getstate__(self):
+ # Cached plans hold planning state of one process; each DataLoader
+ # worker plans the pinned snapshot once for itself.
+ state = self.__dict__.copy()
+ state["_plans"] = {}
+ return state
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self)
+ if index < 0 or index >= len(self):
+ raise IndexError("window index out of range")
+
+ 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
+ sample = {
+ self.group_key: self._group_keys[group_index],
+ self.order_key: int(self._group_first_orders[group_index]) + start,
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows, then store row IDs, groups, and window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ The index is kept as NumPy arrays of row IDs, per-group offsets, and
+ window anchors, plus one Python group value per group. Order values are
+ not stored per row: contiguity makes them the group's first value plus
+ the offset inside the group.
+ """
+ group_column = index.column(self.group_key)
+ order_column = index.column(self.order_key)
+ row_id_column = index.column(SpecialFields.ROW_ID.name)
+ if row_id_column.null_count:
+ raise ValueError(
+ "ContiguousWindowDataset requires readable Paimon row IDs, "
+ "but %s contains null values." % SpecialFields.ROW_ID.name)
+ if group_column.null_count:
+ raise ValueError(
+ "%s must not contain null values." % self.group_key)
+ if order_column.null_count:
+ raise ValueError(
+ "%s must not contain null values." % self.order_key)
+ if not pa.types.is_integer(order_column.type):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ if pa.types.is_floating(group_column.type) and pc.any(
+ pc.is_nan(group_column)).as_py():
+ raise ValueError(
+ "%s must not contain NaN values, which never compare equal to "
+ "themselves and would split one group." % self.group_key)
+
+ try:
+ ordered = index.sort_by([
+ (self.group_key, "ascending"),
+ (self.order_key, "ascending"),
+ ])
+ except pa.ArrowNotImplementedError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+
+ group_values = _contiguous_array(ordered.column(self.group_key))
+ order_values = _integer_numpy(
+ _contiguous_array(ordered.column(self.order_key)))
+ self._row_ids = _integer_numpy(
+ _contiguous_array(ordered.column(SpecialFields.ROW_ID.name)))
+ self._group_starts = _group_starts(group_values)
+ self._group_lengths = np.diff(
+ np.append(self._group_starts, len(self._row_ids)))
+ self._group_keys = group_values.take(
+ pa.array(self._group_starts, type=pa.int64())).to_pylist()
+ self._group_first_orders = order_values[self._group_starts]
+ self._validate_contiguity(order_values)
+ self._anchor_groups, self._anchor_starts = self._build_anchors()
+
+ def _validate_contiguity(self, order_values):
+ if len(order_values) < 2:
+ return
+ same_group = np.ones(len(order_values) - 1, dtype=bool)
+ same_group[self._group_starts[1:] - 1] = False
+ steps = np.diff(order_values)
+ duplicate = same_group & (steps == 0)
+ if duplicate.any():
+ position = int(np.flatnonzero(duplicate)[0])
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (self._group_of(position),
+ int(order_values[position]), self.order_key))
+ broken = same_group & (steps != 1)
+ if broken.any():
+ position = int(np.flatnonzero(broken)[0])
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (self._group_of(position), self.order_key,
+ int(order_values[position]), int(order_values[position + 1])))
+
+ def _group_of(self, position):
+ group_index = int(np.searchsorted(
+ self._group_starts, position, side="right")) - 1
+ return self._group_keys[group_index]
+
+ def _build_anchors(self):
+ groups = []
+ starts = []
+ 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)
+ if incomplete.size:
+ if self.tail == "error":
+ first = int(incomplete[0])
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%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]
+ if positions.size:
+ groups.append(np.full(positions.size, group_index, dtype=np.int64))
+ starts.append(positions)
+ if not groups:
+ empty = np.zeros(0, dtype=np.int64)
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ rows = self._plan_for(columns).read(row_ids)
+ row_id_column = SpecialFields.ROW_ID.name
+ by_row_id = {}
+ for row in rows:
+ by_row_id[int(row.pop(row_id_column))] = row
+ missing = [row_id for row_id in row_ids if row_id not in by_row_id]
+ if missing:
+ raise RuntimeError(
+ "Pinned snapshot %s did not return indexed row IDs %s."
+ % (self.snapshot_id, missing))
+ return [by_row_id[row_id] for row_id in row_ids]
+
+ def _plan_for(self, columns):
+ key = tuple(columns)
+ plan = self._plans.get(key)
+ if plan is None:
+ plan = _PinnedRowIdPlan(self._table, columns, self.blob_parallelism)
+ self._plans[key] = plan
+ return plan
+
+
+class _PinnedRowIdPlan:
+ """One authorized scan plan of the pinned snapshot, read by row ID.
+
+ Planning happens once per projection. Every read then narrows the cached
+ splits to the files covering the requested row IDs, so repeated item and
+ batch reads never revisit the snapshot's manifests.
+ """
+
+ def __init__(self, table, columns, blob_parallelism):
+ self._blob_parallelism = blob_parallelism
+ self._blob_columns = [
+ field.name for field in table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ self._map_blob_columns = {
+ field.name for field in table.fields
+ if field.name in self._blob_columns and is_map_blob_type(field.type)
+ }
+ blob_column_set = set(self._blob_columns)
+ self._projection = (
+ [name for name in columns if name not in blob_column_set]
+ + [SpecialFields.ROW_ID.name]
+ + self._blob_columns
+ )
+ self._table = (
+ table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"})
+ if self._blob_columns else table
+ )
+ self._splits = self._new_read_builder().new_scan().plan().splits()
+ _reject_masked_row_ids(self._splits)
+
+ def read(self, row_ids):
+ """Return raw row dictionaries, including ``_ROW_ID``, for ``row_ids``."""
+ requested = list(dict.fromkeys(row_ids))
+ ranges = Range.to_ranges(requested)
+ splits = [
+ pruned for pruned in (
+ _prune_split_files(split, ranges) for split in self._splits)
+ if pruned is not None
+ ]
+ if not splits:
+ return []
+ 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)
Review Comment:
[P2] Preserve exact row ranges when reading cached splits
The cached plan contains unfiltered `DataSplit`s, 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 `IndexedSplit`s 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org>.
XiaoHongbo-Hope commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3924543164
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
Review Comment:
We can use Arrow or NumPy arrays instead to avoid memory issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org>.
XiaoHongbo-Hope commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3924616350
##########
docs/docs/pypaimon/multimodal-api.mdx:
##########
@@ -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",
Review Comment:
These defaults do not match the native LeRobot schema preserved by #9529, which uses episode_index and frame_index. Should these arguments be required for a generic Dataset, or default to the LeRobot field names?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3923920282
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
+ order_values = index.column(self.order_key).to_pylist()
+ row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist()
+ grouped = defaultdict(list)
+ for group_key, order_value, row_id in zip(
+ group_values, order_values, row_ids):
+ if group_key is None:
+ raise ValueError("%s must not contain null values." % self.group_key)
+ if order_value is None:
+ raise ValueError("%s must not contain null values." % self.order_key)
+ if isinstance(order_value, bool) or not isinstance(order_value, Integral):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ try:
+ grouped[group_key].append((int(order_value), int(row_id)))
Review Comment:
[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.
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
+ order_values = index.column(self.order_key).to_pylist()
+ row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist()
+ grouped = defaultdict(list)
+ for group_key, order_value, row_id in zip(
+ group_values, order_values, row_ids):
+ if group_key is None:
+ raise ValueError("%s must not contain null values." % self.group_key)
+ if order_value is None:
+ raise ValueError("%s must not contain null values." % self.order_key)
+ if isinstance(order_value, bool) or not isinstance(order_value, Integral):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ try:
+ grouped[group_key].append((int(order_value), int(row_id)))
+ except TypeError:
+ raise ValueError(
+ "%s values must be hashable." % self.group_key)
+
+ groups = []
+ anchors = []
+ try:
+ sorted_groups = sorted(grouped.items(), key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+ for group_key, members in sorted_groups:
+ try:
+ members.sort(key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values in group %r must be mutually orderable."
+ % (self.order_key, group_key))
+ for previous, current in zip(members, members[1:]):
+ if previous[0] == current[0]:
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (group_key, current[0], self.order_key))
+ if current[0] != previous[0] + 1:
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (group_key, self.order_key,
+ previous[0], current[0]))
+
+ group_index = len(groups)
+ group_orders = [member[0] for member in members]
+ group_row_ids = [member[1] for member in members]
+ groups.append((group_key, group_orders, group_row_ids))
+ for start in range(0, len(members), self.stride):
+ valid_count = min(self.window_size, len(members) - start)
+ if valid_count < self.window_size:
+ if self.tail == "drop":
+ continue
+ if self.tail == "error":
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%d."
+ % (group_key, group_orders[start],
+ self.window_size, valid_count))
+ anchors.append((group_index, start, valid_count))
+ return groups, anchors
+
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ query = ScanQuery(self._table)
+ predicate_builder = (
+ self._table.new_read_builder()
+ .with_projection(
+ [field.name for field in self._table.fields]
+ + [SpecialFields.ROW_ID.name])
+ .new_predicate_builder()
+ )
+ query._predicate = predicate_builder.is_in(
+ SpecialFields.ROW_ID.name, row_ids)
+ query._projection = list(columns)
+ query._include_row_id = True
+
+ blob_columns = [
+ field.name for field in self._table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ if blob_columns:
+ scalar, blobs = query.read_blobs(
+ blob_columns, parallelism=self.blob_parallelism)
+ rows = scalar.to_pylist()
+ for name in blob_columns:
+ values = blobs[name]
+ if len(values) != len(rows):
+ raise RuntimeError(
+ "BLOB column %s is not row-aligned with a window read."
+ % name)
+ for row, value in zip(rows, values):
+ row[name] = value
+ else:
+ rows = query.to_arrow().to_pylist()
+
+ by_row_id = {}
+ row_id_column = SpecialFields.ROW_ID.name
+ for row in rows:
+ row_id = int(row[row_id_column])
+ del row[row_id_column]
+ by_row_id[row_id] = row
+ missing = [row_id for row_id in row_ids if row_id not in by_row_id]
+ if missing:
+ raise RuntimeError(
+ "Pinned snapshot %s did not return indexed row IDs %s."
+ % (self.snapshot_id, missing))
+ return [by_row_id[row_id] for row_id in row_ids]
+
+
+def _read_window_index(query, group_key, order_key):
+ index_query = copy.copy(query)
+ index_query._projection = [group_key, order_key]
+ 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())
+ if index.num_rows and plan.snapshot_id is None:
+ raise RuntimeError("Cannot pin the snapshot used to build the window index.")
+ return index, plan.snapshot_id
+
+
+def _pin_table(table, snapshot_id):
+ """Pin a table copy to ``snapshot_id``, or reuse it when unresolved."""
+ if snapshot_id is None:
+ return table
+ scan_keys = set(SCAN_KEYS)
+ scan_keys.update(option.key() for option in (
+ CoreOptions.SCAN_MODE,
+ CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP,
+ CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS,
+ CoreOptions.SCAN_CREATION_TIME_MILLIS,
+ ))
+ options = {
+ key: None for key in scan_keys
+ if table.options.options.contains_key(key)
+ }
+ options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id)
Review Comment:
[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.
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
+ order_values = index.column(self.order_key).to_pylist()
+ row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist()
+ grouped = defaultdict(list)
+ for group_key, order_value, row_id in zip(
+ group_values, order_values, row_ids):
+ if group_key is None:
+ raise ValueError("%s must not contain null values." % self.group_key)
+ if order_value is None:
+ raise ValueError("%s must not contain null values." % self.order_key)
+ if isinstance(order_value, bool) or not isinstance(order_value, Integral):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ try:
+ grouped[group_key].append((int(order_value), int(row_id)))
+ except TypeError:
+ raise ValueError(
+ "%s values must be hashable." % self.group_key)
+
+ groups = []
+ anchors = []
+ try:
+ sorted_groups = sorted(grouped.items(), key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+ for group_key, members in sorted_groups:
+ try:
+ members.sort(key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values in group %r must be mutually orderable."
+ % (self.order_key, group_key))
+ for previous, current in zip(members, members[1:]):
+ if previous[0] == current[0]:
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (group_key, current[0], self.order_key))
+ if current[0] != previous[0] + 1:
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (group_key, self.order_key,
+ previous[0], current[0]))
+
+ group_index = len(groups)
+ group_orders = [member[0] for member in members]
+ group_row_ids = [member[1] for member in members]
+ groups.append((group_key, group_orders, group_row_ids))
+ for start in range(0, len(members), self.stride):
+ valid_count = min(self.window_size, len(members) - start)
+ if valid_count < self.window_size:
+ if self.tail == "drop":
+ continue
+ if self.tail == "error":
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%d."
+ % (group_key, group_orders[start],
+ self.window_size, valid_count))
+ anchors.append((group_index, start, valid_count))
+ return groups, anchors
+
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ query = ScanQuery(self._table)
+ predicate_builder = (
+ self._table.new_read_builder()
+ .with_projection(
+ [field.name for field in self._table.fields]
+ + [SpecialFields.ROW_ID.name])
+ .new_predicate_builder()
+ )
+ query._predicate = predicate_builder.is_in(
+ SpecialFields.ROW_ID.name, row_ids)
+ query._projection = list(columns)
+ query._include_row_id = True
+
+ blob_columns = [
+ field.name for field in self._table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ if blob_columns:
+ scalar, blobs = query.read_blobs(
+ blob_columns, parallelism=self.blob_parallelism)
+ rows = scalar.to_pylist()
+ for name in blob_columns:
+ values = blobs[name]
+ if len(values) != len(rows):
+ raise RuntimeError(
+ "BLOB column %s is not row-aligned with a window read."
+ % name)
+ for row, value in zip(rows, values):
+ row[name] = value
+ else:
+ rows = query.to_arrow().to_pylist()
+
+ by_row_id = {}
+ row_id_column = SpecialFields.ROW_ID.name
+ for row in rows:
+ row_id = int(row[row_id_column])
+ del row[row_id_column]
+ by_row_id[row_id] = row
+ missing = [row_id for row_id in row_ids if row_id not in by_row_id]
+ if missing:
+ raise RuntimeError(
+ "Pinned snapshot %s did not return indexed row IDs %s."
+ % (self.snapshot_id, missing))
+ return [by_row_id[row_id] for row_id in row_ids]
+
+
+def _read_window_index(query, group_key, order_key):
+ index_query = copy.copy(query)
+ index_query._projection = [group_key, order_key]
+ 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())
Review Comment:
[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.
##########
paimon-python/pypaimon/multimodal/query.py:
##########
@@ -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:
Review Comment:
[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.
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
+ order_values = index.column(self.order_key).to_pylist()
+ row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist()
+ grouped = defaultdict(list)
+ for group_key, order_value, row_id in zip(
+ group_values, order_values, row_ids):
+ if group_key is None:
+ raise ValueError("%s must not contain null values." % self.group_key)
+ if order_value is None:
+ raise ValueError("%s must not contain null values." % self.order_key)
+ if isinstance(order_value, bool) or not isinstance(order_value, Integral):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ try:
+ grouped[group_key].append((int(order_value), int(row_id)))
+ except TypeError:
+ raise ValueError(
+ "%s values must be hashable." % self.group_key)
+
+ groups = []
+ anchors = []
+ try:
+ sorted_groups = sorted(grouped.items(), key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+ for group_key, members in sorted_groups:
+ try:
+ members.sort(key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values in group %r must be mutually orderable."
+ % (self.order_key, group_key))
+ for previous, current in zip(members, members[1:]):
+ if previous[0] == current[0]:
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (group_key, current[0], self.order_key))
+ if current[0] != previous[0] + 1:
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (group_key, self.order_key,
+ previous[0], current[0]))
+
+ group_index = len(groups)
+ group_orders = [member[0] for member in members]
+ group_row_ids = [member[1] for member in members]
+ groups.append((group_key, group_orders, group_row_ids))
+ for start in range(0, len(members), self.stride):
+ valid_count = min(self.window_size, len(members) - start)
+ if valid_count < self.window_size:
+ if self.tail == "drop":
+ continue
+ if self.tail == "error":
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%d."
+ % (group_key, group_orders[start],
+ self.window_size, valid_count))
+ anchors.append((group_index, start, valid_count))
+ return groups, anchors
+
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ query = ScanQuery(self._table)
Review Comment:
[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.
##########
docs/docs/pypaimon/multimodal-api.mdx:
##########
@@ -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]
Review Comment:
[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]'`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "YannByron (via GitHub)" <gi...@apache.org>.
YannByron commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3925266238
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
+ order_values = index.column(self.order_key).to_pylist()
+ row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist()
+ grouped = defaultdict(list)
+ for group_key, order_value, row_id in zip(
+ group_values, order_values, row_ids):
+ if group_key is None:
+ raise ValueError("%s must not contain null values." % self.group_key)
+ if order_value is None:
+ raise ValueError("%s must not contain null values." % self.order_key)
+ if isinstance(order_value, bool) or not isinstance(order_value, Integral):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ try:
+ grouped[group_key].append((int(order_value), int(row_id)))
+ except TypeError:
+ raise ValueError(
+ "%s values must be hashable." % self.group_key)
+
+ groups = []
+ anchors = []
+ try:
+ sorted_groups = sorted(grouped.items(), key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+ for group_key, members in sorted_groups:
+ try:
+ members.sort(key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values in group %r must be mutually orderable."
+ % (self.order_key, group_key))
+ for previous, current in zip(members, members[1:]):
+ if previous[0] == current[0]:
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (group_key, current[0], self.order_key))
+ if current[0] != previous[0] + 1:
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (group_key, self.order_key,
+ previous[0], current[0]))
+
+ group_index = len(groups)
+ group_orders = [member[0] for member in members]
+ group_row_ids = [member[1] for member in members]
+ groups.append((group_key, group_orders, group_row_ids))
+ for start in range(0, len(members), self.stride):
+ valid_count = min(self.window_size, len(members) - start)
+ if valid_count < self.window_size:
+ if self.tail == "drop":
+ continue
+ if self.tail == "error":
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%d."
+ % (group_key, group_orders[start],
+ self.window_size, valid_count))
+ anchors.append((group_index, start, valid_count))
+ return groups, anchors
+
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ query = ScanQuery(self._table)
+ predicate_builder = (
+ self._table.new_read_builder()
+ .with_projection(
+ [field.name for field in self._table.fields]
+ + [SpecialFields.ROW_ID.name])
+ .new_predicate_builder()
+ )
+ query._predicate = predicate_builder.is_in(
+ SpecialFields.ROW_ID.name, row_ids)
+ query._projection = list(columns)
+ query._include_row_id = True
+
+ blob_columns = [
+ field.name for field in self._table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ if blob_columns:
+ scalar, blobs = query.read_blobs(
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on PR #9580:
URL: https://github.com/apache/paimon/pull/9580#issuecomment-5526153791
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org>.
XiaoHongbo-Hope commented on code in PR #9580:
URL: https://github.com/apache/paimon/pull/9580#discussion_r3924457305
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -0,0 +1,491 @@
+# 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.
+
+"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
+
+import copy
+import operator
+from collections import defaultdict
+from numbers import Integral
+
+import torch
+from torch.utils.data import Dataset
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.schema.data_types import is_blob_type, is_map_blob_type
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.table.special_fields import SpecialFields
+
+
+class ContiguousWindowDataset(Dataset):
+ """Map-style Dataset which reads fixed row windows on demand.
+
+ The in-memory index contains only group values, order values, and Paimon
+ row IDs. Each ``__getitem__`` reads the projected rows from the snapshot
+ resolved while the index was built. 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``:
+
+ * ``drop`` omits them;
+ * ``pad`` repeats final values and marks repeats in ``is_pad``;
+ * ``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.
+ ``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
+ """
+
+ _TAIL_POLICIES = ("drop", "pad", "error")
+
+ def __init__(
+ self,
+ query,
+ *,
+ 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):
+ if getattr(query, "_result_factory", None) is not None:
+ raise TypeError(
+ "ContiguousWindowDataset is only supported on scan(), "
+ "not search queries.")
+ self.window_size = _positive_int(window_size, "window_size")
+ self.stride = _positive_int(stride, "stride")
+ if tail not in self._TAIL_POLICIES:
+ raise ValueError(
+ "tail must be one of %s; got %r."
+ % (self._TAIL_POLICIES, tail))
+ self.tail = tail
+ 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):
+ raise ValueError("group_key and order_key must not be is_pad.")
+ self.columns = _columns(
+ query, columns, self.group_key, self.order_key)
+ 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.column_transforms = _column_transforms(
+ column_transforms, self.columns)
+ self.pad_values = _pad_values(pad_values, self.columns)
+ if adapter is not None and not callable(adapter):
+ raise TypeError("adapter must be callable or None.")
+ self.adapter = adapter
+ self.blob_parallelism = _positive_int(
+ blob_parallelism, "blob_parallelism")
+
+ if not query._table.options.row_tracking_enabled():
+ raise ValueError(
+ "ContiguousWindowDataset requires row-tracking.enabled=true.")
+
+ index, snapshot_id = _read_window_index(
+ query, self.group_key, self.order_key)
+ self.snapshot_id = snapshot_id
+ self._table = _pin_table(query._table, snapshot_id)
+ self._groups, self._anchors = self._build_index(index)
+
+ @classmethod
+ def from_query(cls, query, **kwargs):
+ """Build a contiguous-window Dataset from a ``ScanQuery``."""
+ return cls(query, **kwargs)
+
+ def __len__(self):
+ return len(self._anchors)
+
+ def __getitem__(self, index):
+ """Read one window by map-style Dataset index.
+
+ 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)
+
+ def __getitems__(self, indices):
+ """Read several Dataset indices while coalescing overlapping row IDs.
+
+ The returned list preserves the requested index order and duplicates.
+ Coalescing affects only physical reads, not logical sample cardinality.
+ """
+ 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 {}
+ )
+ 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
+ ]
+
+ def _resolve_window(self, index):
+ index = operator.index(index)
+ if index < 0:
+ index += len(self._anchors)
+ if index < 0 or index >= len(self._anchors):
+ raise IndexError("window index out of range")
+
+ anchor = self._anchors[index]
+ group_index, start, valid_count = anchor
+ row_ids = self._groups[group_index][2]
+ return anchor, row_ids[start:start + valid_count]
+
+ def _sample(self, anchor, rows, anchor_row=None):
+ group_index, start, valid_count = anchor
+ group_key, order_values, _ = self._groups[group_index]
+ 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
+ sample = {
+ self.group_key: group_key,
+ self.order_key: order_values[start],
+ "is_pad": padding_mask,
+ }
+ 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))
+ transform = self.column_transforms.get(name)
+ sample[name] = transform(values) if transform is not None else values
+ if self.adapter is not None:
+ return self.adapter(sample)
+ return sample
+
+ def _build_index(self, index):
+ """Validate index rows and return grouped row IDs plus window anchors.
+
+ Args:
+ index: Arrow table containing ``group_key``, ``order_key``, and
+ Paimon's ``_ROW_ID`` for the resolved snapshot.
+
+ Returns:
+ ``(groups, anchors)``. Each group stores its key, ordered positions,
+ 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()
+ order_values = index.column(self.order_key).to_pylist()
+ row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist()
+ grouped = defaultdict(list)
+ for group_key, order_value, row_id in zip(
+ group_values, order_values, row_ids):
+ if group_key is None:
+ raise ValueError("%s must not contain null values." % self.group_key)
+ if order_value is None:
+ raise ValueError("%s must not contain null values." % self.order_key)
+ if isinstance(order_value, bool) or not isinstance(order_value, Integral):
+ raise ValueError(
+ "%s must contain integer values." % self.order_key)
+ try:
+ grouped[group_key].append((int(order_value), int(row_id)))
+ except TypeError:
+ raise ValueError(
+ "%s values must be hashable." % self.group_key)
+
+ groups = []
+ anchors = []
+ try:
+ sorted_groups = sorted(grouped.items(), key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values must be mutually orderable." % self.group_key)
+ for group_key, members in sorted_groups:
+ try:
+ members.sort(key=lambda item: item[0])
+ except TypeError:
+ raise ValueError(
+ "%s values in group %r must be mutually orderable."
+ % (self.order_key, group_key))
+ for previous, current in zip(members, members[1:]):
+ if previous[0] == current[0]:
+ raise ValueError(
+ "Group %s has duplicate order value %r in %s."
+ % (group_key, current[0], self.order_key))
+ if current[0] != previous[0] + 1:
+ raise ValueError(
+ "Group %s is not contiguous in %s: %s followed by %s."
+ % (group_key, self.order_key,
+ previous[0], current[0]))
+
+ group_index = len(groups)
+ group_orders = [member[0] for member in members]
+ group_row_ids = [member[1] for member in members]
+ groups.append((group_key, group_orders, group_row_ids))
+ for start in range(0, len(members), self.stride):
+ valid_count = min(self.window_size, len(members) - start)
+ if valid_count < self.window_size:
+ if self.tail == "drop":
+ continue
+ if self.tail == "error":
+ raise ValueError(
+ "Group %s has an incomplete window at %s: "
+ "window_size=%d, available=%d."
+ % (group_key, group_orders[start],
+ self.window_size, valid_count))
+ anchors.append((group_index, start, valid_count))
+ return groups, anchors
+
+ 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.
+
+ Args:
+ row_ids: Paimon row IDs to read. Their order and duplicates define
+ the returned row order.
+ columns: Projected value columns, or all Dataset columns when
+ omitted.
+
+ Returns:
+ A list of row dictionaries aligned one-for-one with ``row_ids``.
+ The internal ``_ROW_ID`` field is removed, and BLOB descriptors are
+ resolved to their bodies.
+ """
+ columns = self.columns if columns is None else columns
+ query = ScanQuery(self._table)
+ predicate_builder = (
+ self._table.new_read_builder()
+ .with_projection(
+ [field.name for field in self._table.fields]
+ + [SpecialFields.ROW_ID.name])
+ .new_predicate_builder()
+ )
+ query._predicate = predicate_builder.is_in(
+ SpecialFields.ROW_ID.name, row_ids)
+ query._projection = list(columns)
+ query._include_row_id = True
+
+ blob_columns = [
+ field.name for field in self._table.fields
+ if field.name in columns
+ and (is_blob_type(field.type) or is_map_blob_type(field.type))
+ ]
+ if blob_columns:
+ scalar, blobs = query.read_blobs(
Review Comment:
This loses the `frame_index` of `.video` values.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "YannByron (via GitHub)" <gi...@apache.org>.
YannByron commented on PR #9580:
URL: https://github.com/apache/paimon/pull/9580#issuecomment-5527387796
> 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](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/sample/robomind_agilex.py) 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 current AgileX contract does not yet model `tasks`, `annotations`, or `calibrations`. These can be added when they are present in the AgileX source and required by a concrete consumer. I definitely agree that different datasets and variants should define their own table groups and table schemas according to their native contracts.
The current `contiguous-window` dataset API addresses the requirements of the HDF5-based RoboMIND AgileX dataset while remaining extensible for future sources and modalities. I will first address the comments above. Afterwards, I will use PR https://github.com/apache/paimon/pull/9529 as a reference for adding a multi-table publication contract.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on PR #9580:
URL: https://github.com/apache/paimon/pull/9580#issuecomment-5582068275
+1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi merged PR #9580:
URL: https://github.com/apache/paimon/pull/9580
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python][torch] Add lazy contiguous window dataset [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on PR #9580:
URL: https://github.com/apache/paimon/pull/9580#issuecomment-5525222831
Regarding the schema design, we need to align our design and trade-offs with https://github.com/apache/paimon/pull/9529.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org