This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new ba040653d2 [python][torch] Add lazy contiguous window dataset (#9580)
ba040653d2 is described below
commit ba040653d2c3d21c8b317c8a7546c73840880ac8
Author: Yann Byron <[email protected]>
AuthorDate: Tue Sep 8 16:49:12 2026 +0800
[python][torch] Add lazy contiguous window dataset (#9580)
---
docs/docs/pypaimon/multimodal-api.mdx | 62 ++
docs/docs/pypaimon/pytorch.md | 48 ++
paimon-python/pypaimon/multimodal/query.py | 66 ++
.../pypaimon/multimodal/window_dataset.py | 716 +++++++++++++++++++++
.../pypaimon/read/datasource/torch_dataset.py | 99 +--
.../tests/contiguous_window_dataset_test.py | 664 +++++++++++++++++++
paimon-python/pypaimon/tests/torch_read_test.py | 4 +-
7 files changed, 1613 insertions(+), 46 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-api.mdx
b/docs/docs/pypaimon/multimodal-api.mdx
index bdf780859f..ddef5d1475 100644
--- a/docs/docs/pypaimon/multimodal-api.mdx
+++ b/docs/docs/pypaimon/multimodal-api.mdx
@@ -905,6 +905,68 @@ 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_index",
+ order_key="frame_index",
+ tail="pad",
+ column_transforms={
+ "state": float32_window,
+ "action": float32_window,
+ },
+ )
+)
+
+sample = windows[0]
+assert sample["action"].shape == (16, action_size)
+assert sample["is_pad"].shape == (16,)
+```
+
+The group and order keys in a sample identify the window anchor. Every
projected
+column contains the whole window. With `tail="drop"`, only full windows are
+exposed. With `tail="pad"`, every real row is an anchor; missing suffix values
+repeat the last real value by default and `is_pad` is `True` exactly at those
+positions. With `tail="error"`, construction fails if any scheduled anchor is
+incomplete. Use `pad_values` to override the repeated value for individual
+columns. Anchors advance by `stride`, which defaults to one row.
+
+`column_transforms` receive one padded Python list per projected column. This
is
+where applications define tensor dtype and shape or decode BLOB bytes. The
+optional `adapter` receives the resulting sample mapping and can rename or
+combine fields for a model-specific batch contract. The core Dataset does not
+know model field names, image formats, or normalization rules. Top-level
+functions and callable classes are recommended for transforms and adapters so
+the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader`
+instances.
+
+Columns configured by `video-frame-field` are rejected: a window read would
drop
+the `frame_index` and other metadata carried by their `VideoFrameDescriptor`
+values. Read those columns with `to_torch()` instead.
+
### Distributed BLOB processing with Ray
For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB
diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md
index af9189bbae..33dcb6f3c3 100644
--- a/docs/docs/pypaimon/pytorch.md
+++ b/docs/docs/pypaimon/pytorch.md
@@ -157,7 +157,55 @@ embedded frame ordinals keep frame mapping out of the
normal data file. Use
physical video ranges and cache decoder sessions per worker. See
[Multimodal API: Video Frame Storage](multimodal-api#video-frame-storage)
for the write path and a complete decoder example.
+## Contiguous Windows
+Use a map-style `ContiguousWindowDataset` when training samples are fixed-size
+windows which must not cross a sequence boundary. The dataset builds an index
+from only the group column, order column, and Paimon row IDs. Projected values,
+including BLOB payloads, are read from the pinned snapshot when a sample is
+requested; they are not retained in the index.
+
+```python
+from torch.utils.data import DataLoader
+
+dataset = (
+ frames.scan()
+ .to_contiguous_window_dataset(
+ window_size=16,
+ columns=["state", "image"],
+ anchor_columns=["image"],
+ group_key="episode_index",
+ order_key="frame_index",
+ tail="pad",
+ )
+)
+
+loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True)
+```
+
+Each item contains the group and order keys, one list for each requested
+column, and a boolean `is_pad` tensor where `True` marks padding. Padding
+repeats the final real value by default; `pad_values` can override individual
+columns. Columns named in `anchor_columns` contain only the first row's value,
+which is useful when an observation applies to a full action window. Use
+`column_transforms` to convert column lists to tensors and
+`adapter` to produce a model-specific sample mapping. Keep these callbacks
+picklable when using multiple DataLoader workers.
+
+Scheduled anchors start at row zero and advance by `stride` (default `1`).
+`tail="drop"` omits incomplete windows, `tail="pad"` includes and pads them,
+and `tail="error"` rejects a sequence with any scheduled incomplete window.
+Rows are sorted by `order_key` inside each `group_key` value. Order values must
+be integers which increase by exactly one; duplicates and missing steps are
+rejected, and windows never cross groups. The resolved Paimon
+snapshot is pinned for the lifetime of the dataset, so later commits cannot
+change its index or sample contents. A dataset pinned through `tag_name` fails
+its reads if the tag is moved to another snapshot, rather than mixing rows from
+the two snapshots.
+
+Columns configured by `video-frame-field` are rejected: a window read would
drop
+the `frame_index` and other metadata carried by their `VideoFrameDescriptor`
+values. Read those columns with `to_torch()` instead.
## File Format Metadata Cache
Reusable PyArrow Dataset metadata is cached across reads. Configure its
estimated
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index d4491651d8..489ff64152 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -164,6 +164,66 @@ class ScanQuery:
max_buffer_input_splits=max_buffer_input_splits,
)
+ def to_contiguous_window_dataset(
+ self,
+ *,
+ 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):
+ """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.
+ """
+ from pypaimon.multimodal.window_dataset import ContiguousWindowDataset
+ return ContiguousWindowDataset(
+ self,
+ window_size=window_size,
+ columns=columns,
+ anchor_columns=anchor_columns,
+ group_key=group_key,
+ order_key=order_key,
+ stride=stride,
+ tail=tail,
+ column_transforms=column_transforms,
+ pad_values=pad_values,
+ adapter=adapter,
+ blob_parallelism=blob_parallelism,
+ )
+
def to_ray(
self,
*,
@@ -418,6 +478,12 @@ class _PreFilterQuery(ScanQuery):
"not search queries."
)
+ def to_contiguous_window_dataset(self, *args, **kwargs):
+ raise TypeError(
+ "to_contiguous_window_dataset is only supported on scan(), "
+ "not search queries."
+ )
+
class VectorQuery(_PreFilterQuery):
"""Chainable query wrapper for vector global-index search."""
diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py
b/paimon-python/pypaimon/multimodal/window_dataset.py
new file mode 100644
index 0000000000..2a63845962
--- /dev/null
+++ b/paimon-python/pypaimon/multimodal/window_dataset.py
@@ -0,0 +1,716 @@
+# 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.datasource.torch_dataset import (
+ SplitRangeIndex,
+ row_ranges_for_split,
+ select_indexed_splits,
+)
+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.snapshot_id)
+ 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 row ranges covering the requested row IDs, so repeated item
+ and batch reads never revisit the snapshot's manifests.
+ """
+
+ def __init__(self, table, columns, blob_parallelism, snapshot_id):
+ self._blob_parallelism = blob_parallelism
+ self._snapshot_id = snapshot_id
+ 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
+ )
+ plan = self._new_read_builder().new_scan().plan()
+ # A tag can be replaced between indexing and this deferred plan, so the
+ # plan is only usable while it still resolves the indexed snapshot.
+ if snapshot_id is not None and plan.snapshot_id != snapshot_id:
+ raise RuntimeError(
+ "The window index was built from snapshot %s, but reading its "
+ "rows now resolves snapshot %s; the pinned snapshot moved."
+ % (snapshot_id, plan.snapshot_id))
+ self._splits = plan.splits()
+ _reject_masked_row_ids(self._splits)
+ self._split_ranges = [
+ row_ranges_for_split(split) for split in self._splits
+ ]
+ self._split_range_index = SplitRangeIndex(self._split_ranges)
+
+ 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 = select_indexed_splits(
+ self._splits, self._split_ranges, self._split_range_index, ranges)
+ 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 ranges are merged per split, so a read can still return rows
+ # between requested IDs; drop them before BLOB bodies are fetched.
+ 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 _masked_columns(split):
+ if not isinstance(split, QueryAuthSplit):
+ return ()
+ masking = getattr(split.auth_result, "column_masking", None)
+ return tuple(masking) if masking else ()
+
+
+def _reject_masked_row_ids(splits):
+ for split in splits:
+ if SpecialFields.ROW_ID.name in _masked_columns(split):
+ 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 _reject_masked_index_keys(splits, group_key, order_key):
+ for split in splits:
+ masked = _masked_columns(split)
+ keys = [name for name in (group_key, order_key) if name in masked]
+ if keys:
+ raise ValueError(
+ "ContiguousWindowDataset groups and orders rows by %s, but "
+ "query authorization masks %s for this table; a mask can merge
"
+ "distinct groups or break contiguity, so its windows would no "
+ "longer follow the stored rows."
+ % ([group_key, order_key], keys))
+
+
+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()
+ _reject_masked_row_ids(splits)
+ _reject_masked_index_keys(splits, group_key, order_key)
+ index = read_builder.new_read().to_arrow(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.
+
+ A scan already pinned by ``scan.tag-name`` keeps that tag: the tag retains
+ its snapshot's metadata after the main snapshot file expires, which reading
+ by raw snapshot ID cannot.
+ """
+ if snapshot_id is None:
+ return table
+ tag_name = table.options.scan_tag_name()
+ if tag_name is not None:
+ _require_tag_snapshot(table, tag_name, snapshot_id)
+ 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,
+ ))
+ if tag_name is not None:
+ scan_keys.discard(CoreOptions.SCAN_TAG_NAME.key())
+ options = {
+ key: None for key in scan_keys
+ if table.options.options.contains_key(key)
+ }
+ if tag_name is None:
+ options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id)
+ if not options:
+ return table
+ return table.copy(options)
+
+
+def _require_tag_snapshot(table, tag_name, snapshot_id):
+ tag = table.tag_manager().get(tag_name)
+ if tag is None:
+ raise RuntimeError(
+ "Tag %r used to build the window index no longer exists." %
tag_name)
+ resolved = tag.trim_to_snapshot().id
+ if resolved != snapshot_id:
+ raise RuntimeError(
+ "Tag %r now resolves to snapshot %s, but the window index was "
+ "built from snapshot %s." % (tag_name, resolved, snapshot_id))
+
+
+def _reject_video_columns(table, columns):
+ video_columns = [
+ name for name in columns if name in table.options.video_frame_fields()
+ ]
+ if video_columns:
+ raise ValueError(
+ "ContiguousWindowDataset does not support video frame columns %s "
+ "yet: a window read would drop their frame metadata, such as "
+ "frame_index." % video_columns)
+
+
+def _columns(query, columns, group_key, order_key):
+ available = {field.name for field in query._table.fields}
+ if columns is None:
+ if query._projection is None:
+ columns = [field.name for field in query._table.fields]
+ else:
+ columns = list(query._projection)
+ columns = [name for name in columns
+ if name not in (group_key, order_key)]
+ elif isinstance(columns, str):
+ columns = [columns]
+ else:
+ try:
+ columns = list(columns)
+ except TypeError:
+ raise TypeError(
+ "columns must be a non-empty sequence of column names.")
+ if not columns:
+ raise ValueError("columns must contain at least one value column.")
+ if any(not isinstance(name, str) or not name for name in columns):
+ raise TypeError("columns must contain only non-empty column names.")
+ if len(set(columns)) != len(columns):
+ raise ValueError("columns must not contain duplicates.")
+ invalid = [name for name in columns if name not in available]
+ if invalid:
+ raise ValueError("columns do not exist: %s." % invalid)
+ reserved = [name for name in columns
+ if name in (group_key, order_key, "is_pad")]
+ if reserved:
+ raise ValueError(
+ "columns must not include group_key, order_key, or is_pad: %s."
+ % reserved)
+ return columns
+
+
+def _anchor_columns(value, columns):
+ if value is None:
+ return []
+ if isinstance(value, str):
+ value = [value]
+ else:
+ try:
+ value = list(value)
+ except TypeError:
+ raise TypeError(
+ "anchor_columns must be a sequence of projected column names.")
+ if any(not isinstance(name, str) or not name for name in value):
+ raise TypeError(
+ "anchor_columns must contain only non-empty column names.")
+ if len(set(value)) != len(value):
+ raise ValueError("anchor_columns must not contain duplicates.")
+ invalid = [name for name in value if name not in columns]
+ if invalid:
+ raise ValueError(
+ "anchor_columns must be included in columns: %s." % invalid)
+ return value
+
+
+def _column_transforms(value, columns):
+ transforms = _mapping(value, "column_transforms")
+ _validate_mapping_columns(transforms, columns, "column_transforms")
+ invalid = [name for name, transform in transforms.items()
+ if not callable(transform)]
+ if invalid:
+ raise TypeError(
+ "column_transforms values must be callable: %s." % invalid)
+ return transforms
+
+
+def _pad_values(value, columns):
+ values = _mapping(value, "pad_values")
+ _validate_mapping_columns(values, columns, "pad_values")
+ return values
+
+
+def _mapping(value, name):
+ if value is None:
+ return {}
+ try:
+ return dict(value)
+ except (TypeError, ValueError):
+ raise TypeError("%s must be a mapping or None." % name)
+
+
+def _validate_mapping_columns(value, columns, name):
+ invalid = [column for column in value if column not in columns]
+ if invalid:
+ raise ValueError("%s contains unknown columns: %s." % (name, invalid))
+
+
+def _column(query, value, name):
+ if not isinstance(value, str) or not value:
+ raise TypeError("%s must be a non-empty column name." % name)
+ available = {field.name for field in query._table.fields}
+ if value not in available:
+ raise ValueError("%s column %r does not exist." % (name, value))
+ return value
+
+
+def _positive_int(value, name):
+ if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+ raise ValueError("%s must be a positive int." % name)
+ return value
+
+
+__all__ = ["ContiguousWindowDataset"]
diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py
b/paimon-python/pypaimon/read/datasource/torch_dataset.py
index de4bb2cace..45c1286fe4 100644
--- a/paimon-python/pypaimon/read/datasource/torch_dataset.py
+++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py
@@ -152,7 +152,7 @@ class _RowIdRangeIndex:
return row_ids
-class _SplitRangeIndex:
+class SplitRangeIndex:
def __init__(self, ranges_by_split):
self.intervals = sorted(
@@ -180,6 +180,56 @@ class _SplitRangeIndex:
return sorted(split_indices)
+def row_ranges_for_split(split) -> List[Range]:
+ """Return the merged global row-ID ranges covered by ``split``."""
+ if isinstance(split, QueryAuthSplit):
+ split = split.split
+ if isinstance(split, IndexedSplit):
+ ranges = split.row_ranges()
+ else:
+ ranges = [
+ data_file.row_id_range()
+ for data_file in split.files
+ if data_file.first_row_id is not None
+ ]
+ return Range.sort_and_merge_overlap(ranges, True)
+
+
+def select_indexed_splits(
+ splits, split_ranges, split_range_index, ranges) -> List[Split]:
+ """Return ``IndexedSplit``s restricting ``splits`` to ``ranges``.
+
+ Row ranges reach the format readers through ``IndexedSplit``, which lets a
+ data-evolution read skip files and row blocks that no requested row ID
+ touches. Authorization wrappers are preserved.
+ """
+ selected = []
+ for split_index in split_range_index.find(ranges):
+ original = splits[split_index]
+ auth_result = None
+ split = original
+ if isinstance(split, QueryAuthSplit):
+ auth_result = split.auth_result
+ split = split.split
+
+ if isinstance(split, IndexedSplit):
+ split = split.data_split()
+ allowed = Range.and_(ranges, split_ranges[split_index])
+ if not allowed:
+ continue
+
+ indexed = IndexedSplit(
+ split,
+ allowed,
+ exact_merged_row_count=sum(r.count() for r in allowed),
+ )
+ selected.append(
+ QueryAuthSplit(indexed, auth_result)
+ if auth_result is not None else indexed
+ )
+ return selected
+
+
class TorchDataset(Dataset):
"""
Map-style PyTorch Dataset for Paimon table data.
@@ -205,9 +255,9 @@ class TorchDataset(Dataset):
self._split_range_index = None
if self._supports_lazy_row_id_read():
self._split_ranges = [
- self._row_ranges_for_split(split) for split in splits
+ row_ranges_for_split(split) for split in splits
]
- self._split_range_index = _SplitRangeIndex(self._split_ranges)
+ self._split_range_index = SplitRangeIndex(self._split_ranges)
self._row_ids = self._compact_row_id_index()
if self._row_ids is None:
row_id_read = TableRead(
@@ -274,20 +324,6 @@ class TorchDataset(Dataset):
return None
return _RowIdRangeIndex(ranges, self.table_read.limit)
- @staticmethod
- def _row_ranges_for_split(split):
- if isinstance(split, QueryAuthSplit):
- split = split.split
- if isinstance(split, IndexedSplit):
- ranges = split.row_ranges()
- else:
- ranges = [
- data_file.row_id_range()
- for data_file in split.files
- if data_file.first_row_id is not None
- ]
- return Range.sort_and_merge_overlap(ranges, True)
-
def _materialize(self):
self._row_ids = None
self._data = self.table_read.to_arrow(self.splits)
@@ -377,33 +413,8 @@ class TorchDataset(Dataset):
return index
def _select_splits(self, ranges) -> List[Split]:
- selected = []
- split_indices = self._split_range_index.find(ranges)
- for split_index in split_indices:
- original = self.splits[split_index]
- auth_result = None
- split = original
- if isinstance(split, QueryAuthSplit):
- auth_result = split.auth_result
- split = split.split
-
- if isinstance(split, IndexedSplit):
- split = split.data_split()
- allowed = Range.and_(
- ranges, self._split_ranges[split_index])
- if not allowed:
- continue
-
- indexed = IndexedSplit(
- split,
- allowed,
- exact_merged_row_count=sum(r.count() for r in allowed),
- )
- selected.append(
- QueryAuthSplit(indexed, auth_result)
- if auth_result is not None else indexed
- )
- return selected
+ return select_indexed_splits(
+ self.splits, self._split_ranges, self._split_range_index, ranges)
class _BaseTorchIterDataset(IterableDataset):
diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
new file mode 100644
index 0000000000..c3582ea7b5
--- /dev/null
+++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
@@ -0,0 +1,664 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import json
+import os
+import pickle
+import shutil
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import pyarrow as pa
+import torch
+
+import pypaimon.multimodal as pmm
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+from pypaimon.multimodal import window_dataset
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.multimodal.window_dataset import ContiguousWindowDataset
+from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
+from pypaimon.read.table_scan import TableScan
+
+
+_TABLE_OPTIONS = {
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "true",
+ "deletion-vectors.enabled": "true",
+ "file.format": "parquet",
+ "vector.file.format": "parquet",
+}
+
+
+class _TensorColumnTransform:
+
+ def __call__(self, values):
+ return torch.tensor(values, dtype=torch.int64)
+
+
+class _WindowAdapter:
+
+ def __call__(self, sample):
+ return {
+ "episode": sample["episode"],
+ "start": sample["step"],
+ "values": sample["value"],
+ "padding_mask": sample["is_pad"],
+ }
+
+
+class ContiguousWindowDatasetTest(unittest.TestCase):
+
+ def setUp(self):
+ self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_windows_")
+ self.conn = pmm.connect(options={
+ "warehouse": os.path.join(self.temp_dir, "warehouse"),
+ })
+
+ def tearDown(self):
+ shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+ @staticmethod
+ def _schema():
+ return pa.schema([
+ pa.field("episode", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("value", pa.int32(), nullable=False),
+ pa.field("payload", pa.large_binary(), nullable=False),
+ ])
+
+ @staticmethod
+ def _row(episode, step):
+ return {
+ "episode": episode,
+ "step": step,
+ "value": step + (100 if episode == "episode-b" else 0),
+ "payload": ("%s-%d" % (episode, step)).encode(),
+ }
+
+ def _table(self, name="frames"):
+ table = self.conn.create_table(
+ name, schema=self._schema(), options=_TABLE_OPTIONS)
+ table.add([
+ self._row("episode-b", 2),
+ self._row("episode-a", 1),
+ self._row("episode-b", 0),
+ self._row("episode-a", 0),
+ self._row("episode-b", 3),
+ self._row("episode-b", 1),
+ ])
+ return table
+
+ @staticmethod
+ def _dataset(table, **kwargs):
+ return (
+ table.scan()
+ .to_contiguous_window_dataset(
+ window_size=3,
+ columns=["value", "payload"],
+ group_key="episode",
+ order_key="step",
+ **kwargs,
+ )
+ )
+
+ def test_sorts_rows_and_never_crosses_episode_boundaries(self):
+ dataset = self._dataset(self._table())
+
+ self.assertIsInstance(dataset, torch.utils.data.Dataset)
+ self.assertEqual(2, len(dataset))
+ self.assertIsInstance(dataset.snapshot_id, int)
+ self.assertNotIn("_episodes", vars(dataset))
+
+ first = dataset[0]
+ second = dataset[1]
+ self.assertEqual("episode-b", first["episode"])
+ self.assertEqual(0, first["step"])
+ self.assertEqual([100, 101, 102], first["value"])
+ self.assertEqual([101, 102, 103], second["value"])
+ self.assertFalse(first["is_pad"].any())
+ self.assertEqual({"episode-b"}, {
+ window["episode"] for window in (first, second)
+ })
+
+ def test_reads_blob_payloads_only_when_a_window_is_requested(self):
+ table = self._table()
+ with patch(
+ "pypaimon.multimodal.window_dataset.fetch_blob_bodies",
+ side_effect=window_dataset.fetch_blob_bodies) as fetch:
+ dataset = self._dataset(table)
+ self.assertEqual(0, fetch.call_count)
+
+ sample = dataset[0]
+
+ self.assertEqual(1, fetch.call_count)
+ self.assertEqual(3, len(fetch.call_args.args[1]["payload"]))
+ self.assertEqual(
+ [b"episode-b-0", b"episode-b-1", b"episode-b-2"],
+ sample["payload"],
+ )
+
+ def test_reads_map_blob_payloads_for_map_only_and_mixed_windows(self):
+ schema = pa.schema([
+ pa.field("episode", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("payload", pa.large_binary()),
+ pa.field(
+ "attachments",
+ pa.map_(pa.string(), pa.large_binary()),
+ ),
+ ])
+ table = self.conn.create_table(
+ "map_blobs",
+ schema=schema,
+ options=_TABLE_OPTIONS,
+ )
+ table.add(pa.Table.from_pylist([
+ {
+ "episode": "episode-a",
+ "step": 0,
+ "payload": b"scalar-0",
+ "attachments": {
+ "body": b"map-0", "empty": b"", "null": None},
+ },
+ {
+ "episode": "episode-a",
+ "step": 1,
+ "payload": b"scalar-1",
+ "attachments": None,
+ },
+ ], schema=schema))
+
+ def window(columns):
+ return table.scan().to_contiguous_window_dataset(
+ window_size=2,
+ columns=columns,
+ group_key="episode",
+ order_key="step",
+ )[0]
+
+ map_only = window(["attachments"])
+ mixed = window(["payload", "attachments"])
+
+ self.assertEqual(
+ {"body": b"map-0", "empty": b"", "null": None},
+ dict(map_only["attachments"][0]),
+ )
+ self.assertIsNone(map_only["attachments"][1])
+ self.assertEqual([b"scalar-0", b"scalar-1"], mixed["payload"])
+ self.assertEqual(map_only["attachments"], mixed["attachments"])
+
+ def test_anchor_columns_read_only_the_window_anchor(self):
+ table = self._table()
+ with patch(
+ "pypaimon.multimodal.window_dataset.fetch_blob_bodies",
+ side_effect=window_dataset.fetch_blob_bodies) as fetch:
+ dataset = self._dataset(table, anchor_columns=["payload"])
+
+ sample = dataset[0]
+
+ self.assertEqual([100, 101, 102], sample["value"])
+ self.assertEqual([b"episode-b-0"], sample["payload"])
+ self.assertEqual(1, fetch.call_count)
+ self.assertEqual(1, len(fetch.call_args.args[1]["payload"]))
+
+ def test_plural_access_coalesces_overlapping_window_reads(self):
+ dataset = self._dataset(
+ self._table(), anchor_columns=["payload"])
+
+ with patch.object(
+ dataset, "_read_rows", wraps=dataset._read_rows) as read:
+ actual = dataset.__getitems__([1, 0, 1])
+
+ self.assertEqual(2, read.call_count)
+ self.assertEqual(4, len(read.call_args_list[0].args[0]))
+ self.assertEqual(["value"], read.call_args_list[0].args[1])
+ self.assertEqual(2, len(read.call_args_list[1].args[0]))
+ self.assertEqual(["payload"], read.call_args_list[1].args[1])
+ self.assertEqual(
+ [("episode-b", 1), ("episode-b", 0), ("episode-b", 1)],
+ [(sample["episode"], sample["step"]) for sample in actual],
+ )
+ self.assertEqual(
+ [[101, 102, 103], [100, 101, 102], [101, 102, 103]],
+ [sample["value"] for sample in actual],
+ )
+ self.assertEqual(
+ [[b"episode-b-1"], [b"episode-b-0"], [b"episode-b-1"]],
+ [sample["payload"] for sample in actual],
+ )
+
+ def test_plural_access_isolates_mutable_cells_between_samples(self):
+ table = self.conn.create_table(
+ "mutable_cells",
+ schema=pa.schema([
+ pa.field("episode", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("values", pa.list_(pa.int32()), nullable=False),
+ ]),
+ options=_TABLE_OPTIONS,
+ )
+ table.add([
+ {"episode": "episode-a", "step": step, "values": [step]}
+ for step in range(3)
+ ])
+
+ def mutate(values):
+ for value in values:
+ value.append(99)
+ return values
+
+ dataset = table.scan().to_contiguous_window_dataset(
+ window_size=2,
+ columns=["values"],
+ group_key="episode",
+ order_key="step",
+ column_transforms={"values": mutate},
+ )
+
+ batched = dataset.__getitems__([0, 1, 0])
+ singles = [dataset[index] for index in (0, 1, 0)]
+
+ self.assertEqual(
+ [sample["values"] for sample in singles],
+ [sample["values"] for sample in batched],
+ )
+
+ def test_pad_tail_repeats_last_row_and_marks_real_padding(self):
+ dataset = self._dataset(
+ self._table(), tail="pad", pad_values={"value": -1})
+
+ self.assertEqual(6, len(dataset))
+ short_tail = dataset[1]
+ long_tail = dataset[-1]
+ self.assertEqual("episode-a", short_tail["episode"])
+ self.assertEqual([1, -1, -1], short_tail["value"])
+ self.assertEqual(
+ [b"episode-a-1"] * 3, short_tail["payload"])
+ self.assertEqual([False, True, True], short_tail["is_pad"].tolist())
+ self.assertEqual("episode-b", long_tail["episode"])
+ self.assertEqual([103, -1, -1], long_tail["value"])
+ self.assertEqual([False, True, True], long_tail["is_pad"].tolist())
+
+ def test_error_tail_rejects_an_incomplete_scheduled_window(self):
+ with self.assertRaisesRegex(
+ ValueError, "episode-a.*incomplete.*window_size=3"):
+ self._dataset(self._table(), tail="error")
+
+ def test_stride_controls_scheduled_window_anchors(self):
+ dataset = self._dataset(self._table(), stride=2, tail="pad")
+
+ self.assertEqual(
+ [("episode-a", 0), ("episode-b", 0), ("episode-b", 2)],
+ [(dataset[index]["episode"], dataset[index]["step"])
+ for index in range(len(dataset))],
+ )
+ self.assertEqual(
+ [False, False, True], dataset[-1]["is_pad"].tolist())
+
+ def test_rejects_missing_and_duplicate_order_keys_within_a_group(self):
+ gapped = self.conn.create_table(
+ "gapped", schema=self._schema(), options=_TABLE_OPTIONS)
+ gapped.add([
+ self._row("episode-a", 0),
+ self._row("episode-a", 2),
+ ])
+
+ with self.assertRaisesRegex(
+ ValueError, "episode-a.*not contiguous.*0.*2"):
+ self._dataset(gapped)
+
+ table = self.conn.create_table(
+ "duplicates", schema=self._schema(), options=_TABLE_OPTIONS)
+ table.add([
+ self._row("episode-a", 0),
+ self._row("episode-a", 0),
+ self._row("episode-a", 1),
+ ])
+
+ with self.assertRaisesRegex(
+ ValueError, "episode-a.*duplicate.*order.*0"):
+ self._dataset(table)
+
+ def test_pins_snapshot_for_later_on_demand_reads(self):
+ table = self._table()
+ dataset = self._dataset(table)
+ snapshot_id = dataset.snapshot_id
+
+ table.add([self._row("episode-b", 4)])
+
+ self.assertEqual(snapshot_id, dataset.snapshot_id)
+ self.assertNotEqual(
+ snapshot_id,
table.raw_table.snapshot_manager().get_latest_snapshot().id)
+ self.assertEqual(2, len(dataset))
+ self.assertEqual([101, 102, 103], dataset[-1]["value"])
+
+ def test_snapshot_pin_clears_scan_mode_before_on_demand_reads(self):
+ table = self._table()
+ query = ScanQuery(table.raw_table.copy({"scan.mode": "latest-full"}))
+
+ dataset = query.to_contiguous_window_dataset(
+ window_size=3,
+ columns=["value", "payload"],
+ group_key="episode",
+ order_key="step",
+ )
+
+ self.assertEqual([100, 101, 102], dataset[0]["value"])
+
+ def test_pickle_round_trip_preserves_snapshot_and_window(self):
+ dataset = self._dataset(
+ self._table(), anchor_columns=["payload"])
+ expected = dataset[-1]
+
+ restored = pickle.loads(pickle.dumps(dataset))
+
+ self.assertEqual(dataset.snapshot_id, restored.snapshot_id)
+ self.assertEqual(
+ dataset.snapshot_id,
+ restored._table.options.scan_snapshot_id(),
+ )
+ actual = restored[-1]
+ self.assertEqual(expected["episode"], actual["episode"])
+ self.assertEqual(expected["step"], actual["step"])
+ self.assertEqual(expected["value"], actual["value"])
+ self.assertEqual(expected["payload"], actual["payload"])
+ self.assertTrue(torch.equal(expected["is_pad"], actual["is_pad"]))
+
+ def test_projection_filter_transform_and_dataloader_workers(self):
+ table = self._table()
+ dataset = (
+ table.scan()
+ .where("episode = 'episode-b'")
+ .select(["value"])
+ .to_contiguous_window_dataset(
+ window_size=2,
+ group_key="episode",
+ order_key="step",
+ column_transforms={"value": _TensorColumnTransform()},
+ adapter=_WindowAdapter(),
+ )
+ )
+
+ loader = torch.utils.data.DataLoader(
+ dataset, batch_size=2, shuffle=False, num_workers=2)
+ batches = list(loader)
+
+ self.assertEqual(2, len(batches))
+ self.assertEqual(torch.int64, batches[0]["values"].dtype)
+ self.assertEqual((2, 2), tuple(batches[0]["values"].shape))
+ self.assertEqual(torch.bool, batches[0]["padding_mask"].dtype)
+ self.assertEqual([0, 1, 2], [
+ start for batch in batches for start in batch["start"].tolist()
+ ])
+ self.assertEqual(
+ [[100, 101], [101, 102], [102, 103]],
+ [values for batch in batches for values in
batch["values"].tolist()],
+ )
+ self.assertTrue(all(
+ episode == "episode-b"
+ for batch in batches for episode in batch["episode"]
+ ))
+
+ def test_default_keys_and_public_from_query_entry_point(self):
+ table = self.conn.create_table(
+ "default_keys",
+ schema=pa.schema([
+ pa.field("episode_index", pa.string(), nullable=False),
+ pa.field("frame_index", pa.int32(), nullable=False),
+ pa.field("value", pa.int32(), nullable=False),
+ ]),
+ options=_TABLE_OPTIONS,
+ )
+ table.add([
+ {"episode_index": "episode-a", "frame_index": 0, "value": 10},
+ {"episode_index": "episode-a", "frame_index": 1, "value": 11},
+ ])
+
+ dataset = ContiguousWindowDataset.from_query(
+ table.scan().select(["value"]), window_size=2)
+
+ self.assertEqual(1, len(dataset))
+ self.assertEqual("episode-a", dataset[0]["episode_index"])
+ self.assertEqual(0, dataset[0]["frame_index"])
+ self.assertEqual([10, 11], dataset[0]["value"])
+
+ def test_rejects_scan_and_batch_vector_search_queries(self):
+ table = self.conn.create_table(
+ "vectors",
+ schema=pa.schema([
+ pa.field("episode", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("embedding", pa.list_(pa.float32(), 2)),
+ ]),
+ options=_TABLE_OPTIONS,
+ )
+ table.add([
+ {"episode": "episode-a", "step": 0, "embedding": [1.0, 0.0]},
+ {"episode": "episode-a", "step": 1, "embedding": [0.0, 1.0]},
+ ])
+ kwargs = {
+ "window_size": 2,
+ "columns": ["embedding"],
+ "group_key": "episode",
+ "order_key": "step",
+ }
+
+ for query in (table.search([1.0, 0.0]),
+ table.search_vectors([[1.0, 0.0]])):
+ with self.subTest(query=type(query).__name__):
+ with self.assertRaisesRegex(TypeError, "only supported on
scan"):
+ query.to_contiguous_window_dataset(**kwargs)
+ with self.assertRaisesRegex(TypeError, "only supported on
scan"):
+ ContiguousWindowDataset.from_query(query, **kwargs)
+
+ def test_reads_a_pinned_tag_after_its_snapshot_file_is_removed(self):
+ table = self._table()
+ table.raw_table.create_tag("v1")
+
+ dataset = table.scan(tag_name="v1").to_contiguous_window_dataset(
+ window_size=3, columns=["value"],
+ group_key="episode", order_key="step")
+
+ raw_table = table.raw_table
+ raw_table.file_io.delete_quietly(
+
raw_table.snapshot_manager().get_snapshot_path(dataset.snapshot_id))
+
+ self.assertEqual("v1", dataset._table.options.scan_tag_name())
+ self.assertIsNone(dataset._table.options.scan_snapshot_id())
+ self.assertEqual([100, 101, 102], dataset[0]["value"])
+
+ def test_plans_the_pinned_snapshot_once_per_projection(self):
+ dataset = self._dataset(self._table(), anchor_columns=["payload"])
+
+ with patch.object(
+ TableScan, "plan", autospec=True,
+ side_effect=TableScan.plan) as plan:
+ samples = [dataset[index] for index in range(len(dataset))]
+ samples.extend(dataset.__getitems__(list(range(len(dataset)))))
+
+ self.assertEqual(2, plan.call_count)
+ self.assertEqual(
+ [[100, 101, 102], [101, 102, 103]] * 2,
+ [sample["value"] for sample in samples])
+ self.assertEqual(
+ [[b"episode-b-0"], [b"episode-b-1"]] * 2,
+ [sample["payload"] for sample in samples])
+
+ def test_rejects_query_authorization_that_masks_row_ids(self):
+ table = self._table()
+ auth = TableQueryAuthResult(
+ filter=None,
+ column_masking={"_ROW_ID": json.dumps({"name": "NULL"})},
+ )
+ table.raw_table.catalog_environment.table_query_auth = (
+ lambda options, table_identifier: lambda select: auth)
+
+ with self.assertRaisesRegex(ValueError, "masks _ROW_ID"):
+ self._dataset(table)
+
+ def
test_rejects_reads_after_the_pinned_tag_moves_to_another_snapshot(self):
+ table = self._table()
+ raw_table = table.raw_table
+ raw_table.create_tag("v1")
+ dataset = table.scan(tag_name="v1").to_contiguous_window_dataset(
+ window_size=3, columns=["value"],
+ group_key="episode", order_key="step")
+ restored = pickle.loads(pickle.dumps(dataset))
+
+ table.add([self._row("episode-b", 4)])
+ raw_table.replace_tag("v1")
+
+ for pinned in (dataset, restored):
+ with self.assertRaisesRegex(RuntimeError, "pinned snapshot moved"):
+ pinned[0]
+
+ def test_rejects_query_authorization_that_masks_the_group_key(self):
+ table = self._table()
+ auth = TableQueryAuthResult(
+ filter=None,
+ column_masking={"episode": json.dumps({
+ "name": "SUBSTRING",
+ "inputs": [
+ {"name": "episode", "type": "STRING NOT NULL"}, 1, 1,
+ ],
+ })},
+ )
+ table.raw_table.catalog_environment.table_query_auth = (
+ lambda options, table_identifier: lambda select: auth)
+
+ with self.assertRaisesRegex(ValueError, r"masks \['episode'\]"):
+ self._dataset(table)
+
+ def test_reads_only_the_files_and_row_ranges_a_window_touches(self):
+ table = self.conn.create_table(
+ "many_files", schema=self._schema(), options=_TABLE_OPTIONS)
+ for batch in range(4):
+ table.add([
+ self._row("episode-a", batch * 8 + offset)
+ for offset in range(8)
+ ])
+ dataset = table.scan().to_contiguous_window_dataset(
+ window_size=3, columns=["value"],
+ group_key="episode", order_key="step")
+
+ opened = []
+ decoded = []
+ original_init = FormatPyArrowReader.__init__
+ original_read = FormatPyArrowReader.read_arrow_batch
+
+ def recording_init(reader, *args, **kwargs):
+ original_init(reader, *args, **kwargs)
+ opened.append((args[2], kwargs.get("row_ranges")))
+
+ def recording_read(reader):
+ batch = original_read(reader)
+ decoded.append(0 if batch is None else batch.num_rows)
+ return batch
+
+ with patch.object(FormatPyArrowReader, "__init__", recording_init), \
+ patch.object(
+ FormatPyArrowReader, "read_arrow_batch", recording_read):
+ sample = dataset[len(dataset) - 1]
+
+ self.assertEqual([29, 30, 31], sample["value"])
+ self.assertEqual(1, len({file_path for file_path, _ in opened}))
+ # The reader receives positions inside the file, so global rows 29-31
+ # of the fourth eight-row file arrive as 5-7.
+ self.assertEqual([[(5, 7)]], [ranges for _, ranges in opened])
+ self.assertEqual(3, sum(decoded))
+
+ def test_rejects_nan_group_keys_that_never_compare_equal(self):
+ table = self.conn.create_table(
+ "nan_groups",
+ schema=pa.schema([
+ pa.field("episode", pa.float64()),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("value", pa.int32(), nullable=False),
+ ]),
+ options=_TABLE_OPTIONS,
+ )
+ table.add([
+ {"episode": float("nan"), "step": 0, "value": 0},
+ {"episode": float("nan"), "step": 1, "value": 1},
+ ])
+
+ with self.assertRaisesRegex(ValueError, "episode must not contain
NaN"):
+ table.scan().to_contiguous_window_dataset(
+ window_size=2, columns=["value"],
+ group_key="episode", order_key="step")
+
+ def test_rejects_video_frame_columns_that_would_lose_frame_metadata(self):
+ table = self.conn.create_table(
+ "videos",
+ schema=pa.schema([
+ pa.field("episode", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("video", pa.large_binary()),
+ ]),
+ options=dict(_TABLE_OPTIONS, **{
+ "video-frame-field": "video",
+ "blob-as-descriptor": "true",
+ }),
+ )
+
+ with self.assertRaisesRegex(
+ ValueError, "video frame columns.*frame_index"):
+ table.scan().to_contiguous_window_dataset(
+ window_size=2, columns=["video"],
+ group_key="episode", order_key="step")
+
+ def test_validates_configuration_and_scan_only_contract(self):
+ table = self._table()
+ query = table.scan()
+ for name, value in (
+ ("window_size", 0),
+ ("stride", 0),
+ ("tail", "unknown"),
+ ("group_key", "missing"),
+ ("order_key", "missing")):
+ kwargs = {
+ "window_size": 2,
+ "columns": ["value"],
+ "stride": 1,
+ "tail": "drop",
+ "group_key": "episode",
+ "order_key": "step",
+ }
+ kwargs[name] = value
+ with self.subTest(name=name), self.assertRaises((TypeError,
ValueError)):
+ query.to_contiguous_window_dataset(**kwargs)
+
+ reserved_table = self.conn.create_table(
+ "reserved", schema=pa.schema([
+ pa.field("is_pad", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("value", pa.int32(), nullable=False),
+ ]), options=_TABLE_OPTIONS)
+ with self.assertRaisesRegex(ValueError, "must not be is_pad"):
+ reserved_table.scan().to_contiguous_window_dataset(
+ window_size=2, columns=["value"],
+ group_key="is_pad", order_key="step")
+
+ with self.assertRaisesRegex(TypeError, "only supported on scan"):
+ table.search("anything",
column="episode").to_contiguous_window_dataset(
+ window_size=2, columns=["value"],
+ group_key="episode", order_key="step")
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/paimon-python/pypaimon/tests/torch_read_test.py
b/paimon-python/pypaimon/tests/torch_read_test.py
index 0dc269110d..bb14bd0320 100644
--- a/paimon-python/pypaimon/tests/torch_read_test.py
+++ b/paimon-python/pypaimon/tests/torch_read_test.py
@@ -37,7 +37,7 @@ from pypaimon.catalog.table_query_auth import
TableQueryAuthResult
from pypaimon.multimodal.table import MultimodalTable
from pypaimon.read.datasource.torch_dataset import (
- _SplitRangeIndex,
+ SplitRangeIndex,
TorchIterDataset,
TorchShuffledIterDataset,
_resolve_distributed_context,
@@ -666,7 +666,7 @@ class TorchReadTest(unittest.TestCase):
self.assertEqual(dataset[0], restored[0])
def test_split_range_index(self):
- index = _SplitRangeIndex([
+ index = SplitRangeIndex([
[Range(split * 10, split * 10 + 9)]
for split in range(10000)
])
