diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index b3c59043a8ac..c90a018dd615 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -866,6 +866,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 af9189bbae5e..1769f67dd100 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -157,7 +157,53 @@ 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. + +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 d4491651d8e4..489ff641525b 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -164,6 +164,66 @@ 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_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 @@ def to_arrow_batch_reader(self, *args, **kwargs): "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 000000000000..9c62792d94d9 --- /dev/null +++ b/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() + _reject_masked_row_ids(splits) + 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/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py new file mode 100644 index 000000000000..b0ecf0a7ae0a --- /dev/null +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -0,0 +1,592 @@ +# 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.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_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()