From b9f1c026ce93de9b485d4900ab79805cc156ce37 Mon Sep 17 00:00:00 2001 From: Yann Date: Thu, 3 Sep 2026 17:28:17 +0800 Subject: [PATCH 1/2] [python][torch] Add lazy contiguous window dataset Expose snapshot-pinned, row-ID-backed windows with lazy BLOB reads and coalesced batch access for PyTorch workloads. AI-Contributed/Feature: 0/655 AI-Contributed/UT: 0/472 --- docs/docs/pypaimon/multimodal-api.mdx | 58 +++ docs/docs/pypaimon/pytorch.md | 42 ++ paimon-python/pypaimon/multimodal/query.py | 64 +++ .../pypaimon/multimodal/window_dataset.py | 491 ++++++++++++++++++ .../tests/contiguous_window_dataset_test.py | 472 +++++++++++++++++ 5 files changed, 1127 insertions(+) create mode 100644 paimon-python/pypaimon/multimodal/window_dataset.py create mode 100644 paimon-python/pypaimon/tests/contiguous_window_dataset_test.py diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index b3c59043a8ac..987b3ab42ea7 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -866,6 +866,64 @@ Notes: few large reads); scattered point reads coalesce less. - Blob reads are available only on `scan()`, not on the `search()` queries. +### Contiguous windows for PyTorch + +Install the `torch` extra, then use `to_contiguous_window_dataset` to expose +map-style windows without loading the selected rows or BLOB payloads into Python +memory up front. The Dataset builds a compact index from the group column, order +column, and Paimon row IDs. Each `__getitem__` call fetches only that window from +the snapshot recorded in `dataset.snapshot_id`. + +```shell +pip install pypaimon[torch] +``` + +```python +import torch + + +def float32_window(values): + return torch.tensor(values, dtype=torch.float32) + + +windows = ( + frames.scan() + .where("split = 'train'") + .to_contiguous_window_dataset( + window_size=16, + columns=["state", "action"], + group_key="episode_id", + order_key="step_idx", + tail="pad", + column_transforms={ + "state": float32_window, + "action": float32_window, + }, + ) +) + +sample = windows[0] +assert sample["action"].shape == (16, action_size) +assert sample["is_pad"].shape == (16,) +``` + +The group and order keys in a sample identify the window anchor. Every projected +column contains the whole window. With `tail="drop"`, only full windows are +exposed. With `tail="pad"`, every real row is an anchor; missing suffix values +repeat the last real value by default and `is_pad` is `True` exactly at those +positions. With `tail="error"`, construction fails if any scheduled anchor is +incomplete. Use `pad_values` to override the repeated value for individual +columns. Anchors advance by `stride`, which defaults to one row. + +`column_transforms` receive one padded Python list per projected column. This is +where applications define tensor dtype and shape or decode BLOB bytes. The +optional `adapter` receives the resulting sample mapping and can rename or +combine fields for a model-specific batch contract. The core Dataset does not +know model field names, image formats, or normalization rules. Top-level +functions and callable classes are recommended for transforms and adapters so +the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader` +instances. + ### Distributed BLOB processing with Ray For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index af9189bbae5e..f6e7c87617e1 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -157,7 +157,49 @@ 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_id", + order_key="step_idx", + 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. ## 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..33f0483b519a 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -164,6 +164,70 @@ def to_torch( max_buffer_input_splits=max_buffer_input_splits, ) + def to_contiguous_window_dataset( + self, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + """Build a snapshot-pinned, map-style Dataset of contiguous rows. + + The Dataset indexes only ``group_key``, ``order_key``, and Paimon row + IDs, then reads projected values on demand. Columns listed in + ``anchor_columns`` are provided to ``column_transforms`` as one-element + lists read from the first row of each window; ``adapter`` receives the + transformed values. ``order_key`` must contain non-null integers that + increase by exactly one within each group. The Dataset sorts rows within + each group and never creates a window across groups. + + Args: + window_size: Number of rows in a complete window. + columns: Value columns to return, excluding the group and order + keys. The scan projection is used when omitted. + anchor_columns: Subset of ``columns`` read only from the window's + first row. + group_key: Column identifying an independent row sequence. + order_key: Integer position column within each group. + stride: Distance between scheduled window starts. + tail: Handling for incomplete final windows: ``drop``, ``pad``, or + ``error``. + column_transforms: Per-column callables applied to value lists. + pad_values: Optional replacement values used by ``tail='pad'``. + adapter: Callable that converts the complete sample mapping. + blob_parallelism: Maximum concurrent BLOB body reads per fetch. + + Returns: + A snapshot-pinned ``ContiguousWindowDataset``. See that class for + padding, mask, transform, and adapter result semantics. + """ + if self._result_factory is not None: + raise TypeError( + "to_contiguous_window_dataset is only supported on scan(), " + "not search queries.") + from pypaimon.multimodal.window_dataset import ContiguousWindowDataset + return ContiguousWindowDataset( + self, + window_size=window_size, + columns=columns, + anchor_columns=anchor_columns, + group_key=group_key, + order_key=order_key, + stride=stride, + tail=tail, + column_transforms=column_transforms, + pad_values=pad_values, + adapter=adapter, + blob_parallelism=blob_parallelism, + ) + def to_ray( self, *, diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py new file mode 100644 index 000000000000..6eab515a2cbb --- /dev/null +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -0,0 +1,491 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" + +import copy +import operator +from collections import defaultdict +from numbers import Integral + +import torch +from torch.utils.data import Dataset + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.query import ScanQuery +from pypaimon.schema.data_types import is_blob_type, is_map_blob_type +from pypaimon.snapshot.time_travel_util import SCAN_KEYS +from pypaimon.table.special_fields import SpecialFields + + +class ContiguousWindowDataset(Dataset): + """Map-style Dataset which reads fixed row windows on demand. + + The in-memory index contains only group values, order values, and Paimon + row IDs. Each ``__getitem__`` reads the projected rows from the snapshot + resolved while the index was built. Within each group, ``order_key`` must + contain non-null integers that increase by exactly one; rows from different + groups never share a window. ``tail`` controls scheduled anchors whose + remaining rows are shorter than ``window_size``: + + * ``drop`` omits them; + * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``error`` rejects the dataset. + + The raw result mapping contains scalar group and order values, a + length-``window_size`` Boolean ``is_pad`` tensor, one-element lists for + ``anchor_columns``, and length-``window_size`` lists for other projected + columns. ``anchor_columns`` therefore avoids loading repeated context such + as observation images or initial robot state. ``column_transforms`` then + convert individual column lists before ``adapter`` adapts the complete + mapping to a model-specific contract. + ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. + """ + + _TAIL_POLICIES = ("drop", "pad", "error") + + def __init__( + self, + query, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + if getattr(query, "_result_factory", None) is not None: + raise TypeError( + "ContiguousWindowDataset is only supported on scan(), " + "not search queries.") + self.window_size = _positive_int(window_size, "window_size") + self.stride = _positive_int(stride, "stride") + if tail not in self._TAIL_POLICIES: + raise ValueError( + "tail must be one of %s; got %r." + % (self._TAIL_POLICIES, tail)) + self.tail = tail + self.group_key = _column(query, group_key, "group_key") + self.order_key = _column(query, order_key, "order_key") + if self.group_key == self.order_key: + raise ValueError("group_key and order_key must name different columns.") + if "is_pad" in (self.group_key, self.order_key): + raise ValueError("group_key and order_key must not be is_pad.") + self.columns = _columns( + query, columns, self.group_key, self.order_key) + self.anchor_columns = _anchor_columns(anchor_columns, self.columns) + anchor_column_set = set(self.anchor_columns) + self._window_columns = [ + name for name in self.columns if name not in anchor_column_set + ] + self.column_transforms = _column_transforms( + column_transforms, self.columns) + self.pad_values = _pad_values(pad_values, self.columns) + if adapter is not None and not callable(adapter): + raise TypeError("adapter must be callable or None.") + self.adapter = adapter + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + + if not query._table.options.row_tracking_enabled(): + raise ValueError( + "ContiguousWindowDataset requires row-tracking.enabled=true.") + + index, snapshot_id = _read_window_index( + query, self.group_key, self.order_key) + self.snapshot_id = snapshot_id + self._table = _pin_table(query._table, snapshot_id) + self._groups, self._anchors = self._build_index(index) + + @classmethod + def from_query(cls, query, **kwargs): + """Build a contiguous-window Dataset from a ``ScanQuery``.""" + return cls(query, **kwargs) + + def __len__(self): + return len(self._anchors) + + def __getitem__(self, index): + """Read one window by map-style Dataset index. + + Negative indices follow Python sequence semantics. The return value is + the pre-adapter mapping described by the class, or the adapter result + when an adapter is configured. + """ + anchor, row_ids = self._resolve_window(index) + rows = self._read_window_rows(row_ids) + anchor_row = ( + self._read_rows(row_ids[:1], self.anchor_columns)[0] + if self.anchor_columns else None + ) + return self._sample(anchor, rows, anchor_row) + + def __getitems__(self, indices): + """Read several Dataset indices while coalescing overlapping row IDs. + + The returned list preserves the requested index order and duplicates. + Coalescing affects only physical reads, not logical sample cardinality. + """ + windows = [self._resolve_window(index) for index in indices] + if not windows: + return [] + row_ids = list(dict.fromkeys( + row_id for _, window_row_ids in windows + for row_id in window_row_ids + )) + rows_by_id = dict(zip(row_ids, self._read_window_rows(row_ids))) + anchor_row_ids = list(dict.fromkeys( + window_row_ids[0] for _, window_row_ids in windows + )) + anchor_rows_by_id = ( + dict(zip( + anchor_row_ids, + self._read_rows(anchor_row_ids, self.anchor_columns), + )) + if self.anchor_columns else {} + ) + return [ + self._sample( + anchor, + [rows_by_id[row_id] for row_id in window_row_ids], + anchor_rows_by_id.get(window_row_ids[0]), + ) + for anchor, window_row_ids in windows + ] + + def _resolve_window(self, index): + index = operator.index(index) + if index < 0: + index += len(self._anchors) + if index < 0 or index >= len(self._anchors): + raise IndexError("window index out of range") + + anchor = self._anchors[index] + group_index, start, valid_count = anchor + row_ids = self._groups[group_index][2] + return anchor, row_ids[start:start + valid_count] + + def _sample(self, anchor, rows, anchor_row=None): + group_index, start, valid_count = anchor + group_key, order_values, _ = self._groups[group_index] + padding_count = self.window_size - valid_count + padding_mask = torch.zeros(self.window_size, dtype=torch.bool) + if padding_count: + padding_mask[valid_count:] = True + sample = { + self.group_key: group_key, + self.order_key: order_values[start], + "is_pad": padding_mask, + } + for name in self.columns: + if name in self.anchor_columns: + values = [copy.deepcopy(anchor_row[name])] + else: + values = [copy.deepcopy(row[name]) for row in rows] + if padding_count and name not in self.anchor_columns: + pad_value = self.pad_values.get(name, values[-1]) + values.extend( + copy.deepcopy(pad_value) for _ in range(padding_count)) + transform = self.column_transforms.get(name) + sample[name] = transform(values) if transform is not None else values + if self.adapter is not None: + return self.adapter(sample) + return sample + + def _build_index(self, index): + """Validate index rows and return grouped row IDs plus window anchors. + + Args: + index: Arrow table containing ``group_key``, ``order_key``, and + Paimon's ``_ROW_ID`` for the resolved snapshot. + + Returns: + ``(groups, anchors)``. Each group stores its key, ordered positions, + and row IDs. Each anchor stores group index, start offset, and the + number of real rows available before optional padding. + """ + group_values = index.column(self.group_key).to_pylist() + order_values = index.column(self.order_key).to_pylist() + row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist() + grouped = defaultdict(list) + for group_key, order_value, row_id in zip( + group_values, order_values, row_ids): + if group_key is None: + raise ValueError("%s must not contain null values." % self.group_key) + if order_value is None: + raise ValueError("%s must not contain null values." % self.order_key) + if isinstance(order_value, bool) or not isinstance(order_value, Integral): + raise ValueError( + "%s must contain integer values." % self.order_key) + try: + grouped[group_key].append((int(order_value), int(row_id))) + except TypeError: + raise ValueError( + "%s values must be hashable." % self.group_key) + + groups = [] + anchors = [] + try: + sorted_groups = sorted(grouped.items(), key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values must be mutually orderable." % self.group_key) + for group_key, members in sorted_groups: + try: + members.sort(key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values in group %r must be mutually orderable." + % (self.order_key, group_key)) + for previous, current in zip(members, members[1:]): + if previous[0] == current[0]: + raise ValueError( + "Group %s has duplicate order value %r in %s." + % (group_key, current[0], self.order_key)) + if current[0] != previous[0] + 1: + raise ValueError( + "Group %s is not contiguous in %s: %s followed by %s." + % (group_key, self.order_key, + previous[0], current[0])) + + group_index = len(groups) + group_orders = [member[0] for member in members] + group_row_ids = [member[1] for member in members] + groups.append((group_key, group_orders, group_row_ids)) + for start in range(0, len(members), self.stride): + valid_count = min(self.window_size, len(members) - start) + if valid_count < self.window_size: + if self.tail == "drop": + continue + if self.tail == "error": + raise ValueError( + "Group %s has an incomplete window at %s: " + "window_size=%d, available=%d." + % (group_key, group_orders[start], + self.window_size, valid_count)) + anchors.append((group_index, start, valid_count)) + return groups, anchors + + def _read_window_rows(self, row_ids): + if not self._window_columns: + return [{} for _ in row_ids] + return self._read_rows(row_ids, self._window_columns) + + def _read_rows(self, row_ids, columns=None): + """Read projected rows by ID from the pinned snapshot. + + Args: + row_ids: Paimon row IDs to read. Their order and duplicates define + the returned row order. + columns: Projected value columns, or all Dataset columns when + omitted. + + Returns: + A list of row dictionaries aligned one-for-one with ``row_ids``. + The internal ``_ROW_ID`` field is removed, and BLOB descriptors are + resolved to their bodies. + """ + columns = self.columns if columns is None else columns + query = ScanQuery(self._table) + predicate_builder = ( + self._table.new_read_builder() + .with_projection( + [field.name for field in self._table.fields] + + [SpecialFields.ROW_ID.name]) + .new_predicate_builder() + ) + query._predicate = predicate_builder.is_in( + SpecialFields.ROW_ID.name, row_ids) + query._projection = list(columns) + query._include_row_id = True + + blob_columns = [ + field.name for field in self._table.fields + if field.name in columns + and (is_blob_type(field.type) or is_map_blob_type(field.type)) + ] + if blob_columns: + scalar, blobs = query.read_blobs( + blob_columns, parallelism=self.blob_parallelism) + rows = scalar.to_pylist() + for name in blob_columns: + values = blobs[name] + if len(values) != len(rows): + raise RuntimeError( + "BLOB column %s is not row-aligned with a window read." + % name) + for row, value in zip(rows, values): + row[name] = value + else: + rows = query.to_arrow().to_pylist() + + by_row_id = {} + row_id_column = SpecialFields.ROW_ID.name + for row in rows: + row_id = int(row[row_id_column]) + del row[row_id_column] + by_row_id[row_id] = row + missing = [row_id for row_id in row_ids if row_id not in by_row_id] + if missing: + raise RuntimeError( + "Pinned snapshot %s did not return indexed row IDs %s." + % (self.snapshot_id, missing)) + return [by_row_id[row_id] for row_id in row_ids] + + +def _read_window_index(query, group_key, order_key): + index_query = copy.copy(query) + index_query._projection = [group_key, order_key] + index_query._include_row_id = True + read_builder = index_query._configured_read_builder() + plan = read_builder.new_scan().plan() + index = read_builder.new_read().to_arrow(plan.splits()) + if index.num_rows and plan.snapshot_id is None: + raise RuntimeError("Cannot pin the snapshot used to build the window index.") + return index, plan.snapshot_id + + +def _pin_table(table, snapshot_id): + """Pin a table copy to ``snapshot_id``, or reuse it when unresolved.""" + if snapshot_id is None: + return table + scan_keys = set(SCAN_KEYS) + scan_keys.update(option.key() for option in ( + CoreOptions.SCAN_MODE, + CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP, + CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS, + CoreOptions.SCAN_CREATION_TIME_MILLIS, + )) + options = { + key: None for key in scan_keys + if table.options.options.contains_key(key) + } + options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id) + return table.copy(options) + + +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..389d4bb5f36a --- /dev/null +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -0,0 +1,472 @@ +# 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 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.multimodal.query import ScanQuery +from pypaimon.multimodal.window_dataset import ContiguousWindowDataset + + +_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() + original = ScanQuery._fetch_bodies + with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) 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() + original = ScanQuery._fetch_bodies + with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) 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_id", pa.string(), nullable=False), + pa.field("step_idx", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + ]), + options=_TABLE_OPTIONS, + ) + table.add([ + {"episode_id": "episode-a", "step_idx": 0, "value": 10}, + {"episode_id": "episode-a", "step_idx": 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_id"]) + self.assertEqual(0, dataset[0]["step_idx"]) + self.assertEqual([10, 11], dataset[0]["value"]) + + 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() From 8d56ac87d57b3aab6175005fae9e4230c15f9f94 Mon Sep 17 00:00:00 2001 From: Yann Date: Thu, 3 Sep 2026 23:52:31 +0800 Subject: [PATCH 2/2] [python][torch] Harden contiguous window dataset per review Reuse the pinned scan plan per projection and prune files by row-id range, so repeated window reads stop replanning the snapshot. Keep tag-pinned reads working after the snapshot file expires, store the window index in Arrow/NumPy arrays, and reject inputs the window contract cannot honor: search queries, masked _ROW_ID, NaN group keys and video frame columns whose frame metadata a window read would drop. Default keys now follow the LeRobot-native episode_index/frame_index. AI-Contributed/Feature: 0/496 AI-Contributed/UT: 0/140 --- docs/docs/pypaimon/multimodal-api.mdx | 10 +- docs/docs/pypaimon/pytorch.md | 8 +- paimon-python/pypaimon/multimodal/query.py | 14 +- .../pypaimon/multimodal/window_dataset.py | 464 +++++++++++++----- .../tests/contiguous_window_dataset_test.py | 140 +++++- 5 files changed, 491 insertions(+), 145 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index 987b3ab42ea7..c90a018dd615 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -875,7 +875,7 @@ column, and Paimon row IDs. Each `__getitem__` call fetches only that window fro the snapshot recorded in `dataset.snapshot_id`. ```shell -pip install pypaimon[torch] +pip install 'pypaimon[torch]' ``` ```python @@ -892,8 +892,8 @@ windows = ( .to_contiguous_window_dataset( window_size=16, columns=["state", "action"], - group_key="episode_id", - order_key="step_idx", + group_key="episode_index", + order_key="frame_index", tail="pad", column_transforms={ "state": float32_window, @@ -924,6 +924,10 @@ 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 f6e7c87617e1..1769f67dd100 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -174,8 +174,8 @@ dataset = ( window_size=16, columns=["state", "image"], anchor_columns=["image"], - group_key="episode_id", - order_key="step_idx", + group_key="episode_index", + order_key="frame_index", tail="pad", ) ) @@ -200,6 +200,10 @@ 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 33f0483b519a..489ff641525b 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -170,8 +170,8 @@ def to_contiguous_window_dataset( window_size, columns=None, anchor_columns=None, - group_key="episode_id", - order_key="step_idx", + group_key="episode_index", + order_key="frame_index", stride=1, tail="drop", column_transforms=None, @@ -208,10 +208,6 @@ def to_contiguous_window_dataset( A snapshot-pinned ``ContiguousWindowDataset``. See that class for padding, mask, transform, and adapter result semantics. """ - if self._result_factory is not None: - raise TypeError( - "to_contiguous_window_dataset is only supported on scan(), " - "not search queries.") from pypaimon.multimodal.window_dataset import ContiguousWindowDataset return ContiguousWindowDataset( self, @@ -482,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 index 6eab515a2cbb..9c62792d94d9 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -19,28 +19,33 @@ import copy import operator -from collections import defaultdict -from numbers import Integral +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.query import ScanQuery +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 values, and Paimon - row IDs. Each ``__getitem__`` reads the projected rows from the snapshot - resolved while the index was built. Within each group, ``order_key`` must - contain non-null integers that increase by exactly one; rows from different - groups never share a window. ``tail`` controls scheduled anchors whose - remaining rows are shorter than ``window_size``: + 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``; @@ -54,6 +59,8 @@ class ContiguousWindowDataset(Dataset): 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") @@ -65,15 +72,15 @@ def __init__( window_size, columns=None, anchor_columns=None, - group_key="episode_id", - order_key="step_idx", + group_key="episode_index", + order_key="frame_index", stride=1, tail="drop", column_transforms=None, pad_values=None, adapter=None, blob_parallelism=64): - if getattr(query, "_result_factory", None) is not None: + if not isinstance(query, ScanQuery) or isinstance(query, _PreFilterQuery): raise TypeError( "ContiguousWindowDataset is only supported on scan(), " "not search queries.") @@ -92,6 +99,7 @@ def __init__( 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 = [ @@ -114,7 +122,8 @@ def __init__( query, self.group_key, self.order_key) self.snapshot_id = snapshot_id self._table = _pin_table(query._table, snapshot_id) - self._groups, self._anchors = self._build_index(index) + self._plans = {} + self._build_index(index) @classmethod def from_query(cls, query, **kwargs): @@ -122,7 +131,7 @@ def from_query(cls, query, **kwargs): return cls(query, **kwargs) def __len__(self): - return len(self._anchors) + return int(self._anchor_groups.size) def __getitem__(self, index): """Read one window by map-style Dataset index. @@ -172,28 +181,37 @@ def __getitems__(self, indices): 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._anchors) - if index < 0 or index >= len(self._anchors): + index += len(self) + if index < 0 or index >= len(self): raise IndexError("window index out of range") - anchor = self._anchors[index] - group_index, start, valid_count = anchor - row_ids = self._groups[group_index][2] - return anchor, row_ids[start:start + valid_count] + 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 - group_key, order_values, _ = self._groups[group_index] padding_count = self.window_size - valid_count padding_mask = torch.zeros(self.window_size, dtype=torch.bool) if padding_count: padding_mask[valid_count:] = True sample = { - self.group_key: group_key, - self.order_key: order_values[start], + 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: @@ -212,78 +230,116 @@ def _sample(self, anchor, rows, anchor_row=None): return sample def _build_index(self, index): - """Validate index rows and return grouped row IDs plus window anchors. + """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. - Returns: - ``(groups, anchors)``. Each group stores its key, ordered positions, - and row IDs. Each anchor stores group index, start offset, and the - number of real rows available before optional padding. + 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_values = index.column(self.group_key).to_pylist() - order_values = index.column(self.order_key).to_pylist() - row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist() - grouped = defaultdict(list) - for group_key, order_value, row_id in zip( - group_values, order_values, row_ids): - if group_key is None: - raise ValueError("%s must not contain null values." % self.group_key) - if order_value is None: - raise ValueError("%s must not contain null values." % self.order_key) - if isinstance(order_value, bool) or not isinstance(order_value, Integral): - raise ValueError( - "%s must contain integer values." % self.order_key) - try: - grouped[group_key].append((int(order_value), int(row_id))) - except TypeError: - raise ValueError( - "%s values must be hashable." % self.group_key) + 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) - groups = [] - anchors = [] try: - sorted_groups = sorted(grouped.items(), key=lambda item: item[0]) - except TypeError: + 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) - for group_key, members in sorted_groups: - try: - members.sort(key=lambda item: item[0]) - except TypeError: - raise ValueError( - "%s values in group %r must be mutually orderable." - % (self.order_key, group_key)) - for previous, current in zip(members, members[1:]): - if previous[0] == current[0]: - raise ValueError( - "Group %s has duplicate order value %r in %s." - % (group_key, current[0], self.order_key)) - if current[0] != previous[0] + 1: + + 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 is not contiguous in %s: %s followed by %s." - % (group_key, self.order_key, - previous[0], current[0])) - - group_index = len(groups) - group_orders = [member[0] for member in members] - group_row_ids = [member[1] for member in members] - groups.append((group_key, group_orders, group_row_ids)) - for start in range(0, len(members), self.stride): - valid_count = min(self.window_size, len(members) - start) - if valid_count < self.window_size: - if self.tail == "drop": - continue - if self.tail == "error": - raise ValueError( - "Group %s has an incomplete window at %s: " - "window_size=%d, available=%d." - % (group_key, group_orders[start], - self.window_size, valid_count)) - anchors.append((group_index, start, valid_count)) - return groups, anchors + "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: @@ -305,45 +361,11 @@ def _read_rows(self, row_ids, columns=None): resolved to their bodies. """ columns = self.columns if columns is None else columns - query = ScanQuery(self._table) - predicate_builder = ( - self._table.new_read_builder() - .with_projection( - [field.name for field in self._table.fields] - + [SpecialFields.ROW_ID.name]) - .new_predicate_builder() - ) - query._predicate = predicate_builder.is_in( - SpecialFields.ROW_ID.name, row_ids) - query._projection = list(columns) - query._include_row_id = True - - blob_columns = [ - field.name for field in self._table.fields - if field.name in columns - and (is_blob_type(field.type) or is_map_blob_type(field.type)) - ] - if blob_columns: - scalar, blobs = query.read_blobs( - blob_columns, parallelism=self.blob_parallelism) - rows = scalar.to_pylist() - for name in blob_columns: - values = blobs[name] - if len(values) != len(rows): - raise RuntimeError( - "BLOB column %s is not row-aligned with a window read." - % name) - for row, value in zip(rows, values): - row[name] = value - else: - rows = query.to_arrow().to_pylist() - - by_row_id = {} + rows = self._plan_for(columns).read(row_ids) row_id_column = SpecialFields.ROW_ID.name + by_row_id = {} for row in rows: - row_id = int(row[row_id_column]) - del row[row_id_column] - by_row_id[row_id] = row + 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( @@ -351,6 +373,162 @@ def _read_rows(self, row_ids, columns=None): % (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) @@ -358,16 +536,26 @@ def _read_window_index(query, group_key, order_key): index_query._include_row_id = True read_builder = index_query._configured_read_builder() plan = read_builder.new_scan().plan() - index = read_builder.new_read().to_arrow(plan.splits()) + 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.""" + """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, @@ -375,14 +563,42 @@ def _pin_table(table, snapshot_id): 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) } - options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id) + 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: diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index 389d4bb5f36a..b0ecf0a7ae0a 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import json import os import pickle import shutil @@ -26,8 +27,11 @@ 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 = { @@ -132,8 +136,9 @@ def test_sorts_rows_and_never_crosses_episode_boundaries(self): def test_reads_blob_payloads_only_when_a_window_is_requested(self): table = self._table() - original = ScanQuery._fetch_bodies - with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + 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) @@ -198,8 +203,9 @@ def window(columns): def test_anchor_columns_read_only_the_window_anchor(self): table = self._table() - original = ScanQuery._fetch_bodies - with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + 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] @@ -411,25 +417,139 @@ def test_default_keys_and_public_from_query_entry_point(self): table = self.conn.create_table( "default_keys", schema=pa.schema([ - pa.field("episode_id", pa.string(), nullable=False), - pa.field("step_idx", pa.int32(), nullable=False), + 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_id": "episode-a", "step_idx": 0, "value": 10}, - {"episode_id": "episode-a", "step_idx": 1, "value": 11}, + {"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_id"]) - self.assertEqual(0, dataset[0]["step_idx"]) + 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()