JingsongLi commented on code in PR #9580: URL: https://github.com/apache/paimon/pull/9580#discussion_r3930374955
########## paimon-python/pypaimon/multimodal/window_dataset.py: ########## @@ -0,0 +1,707 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" + +import copy +import operator + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import torch +from torch.utils.data import Dataset + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.blob_read import fetch_blob_bodies +from pypaimon.multimodal.query import ScanQuery, _PreFilterQuery +from pypaimon.read.query_auth_split import QueryAuthSplit +from pypaimon.schema.data_types import is_blob_type, is_map_blob_type +from pypaimon.snapshot.time_travel_util import SCAN_KEYS +from pypaimon.table.special_fields import SpecialFields +from pypaimon.utils.range import Range + + +class ContiguousWindowDataset(Dataset): + """Map-style Dataset which reads fixed row windows on demand. + + The in-memory index contains only group values, order bounds, and Paimon + row IDs, stored in Arrow and NumPy arrays. Each ``__getitem__`` reads the + projected rows from the snapshot resolved while the index was built, reusing + that snapshot's authorized scan plan instead of planning again. Within each + group, ``order_key`` must contain non-null integers that increase by exactly + one; rows from different groups never share a window. ``tail`` controls + scheduled anchors whose remaining rows are shorter than ``window_size``: + + * ``drop`` omits them; + * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``error`` rejects the dataset. + + The raw result mapping contains scalar group and order values, a + length-``window_size`` Boolean ``is_pad`` tensor, one-element lists for + ``anchor_columns``, and length-``window_size`` lists for other projected + columns. ``anchor_columns`` therefore avoids loading repeated context such + as observation images or initial robot state. ``column_transforms`` then + convert individual column lists before ``adapter`` adapts the complete + mapping to a model-specific contract. + ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. + Video frame columns are not supported yet, because a window read would drop + the frame metadata carried by their descriptors. + """ + + _TAIL_POLICIES = ("drop", "pad", "error") + + def __init__( + self, + query, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_index", + order_key="frame_index", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + if not isinstance(query, ScanQuery) or isinstance(query, _PreFilterQuery): + raise TypeError( + "ContiguousWindowDataset is only supported on scan(), " + "not search queries.") + self.window_size = _positive_int(window_size, "window_size") + self.stride = _positive_int(stride, "stride") + if tail not in self._TAIL_POLICIES: + raise ValueError( + "tail must be one of %s; got %r." + % (self._TAIL_POLICIES, tail)) + self.tail = tail + self.group_key = _column(query, group_key, "group_key") + self.order_key = _column(query, order_key, "order_key") + if self.group_key == self.order_key: + raise ValueError("group_key and order_key must name different columns.") + if "is_pad" in (self.group_key, self.order_key): + raise ValueError("group_key and order_key must not be is_pad.") + self.columns = _columns( + query, columns, self.group_key, self.order_key) + _reject_video_columns(query._table, self.columns) + self.anchor_columns = _anchor_columns(anchor_columns, self.columns) + anchor_column_set = set(self.anchor_columns) + self._window_columns = [ + name for name in self.columns if name not in anchor_column_set + ] + self.column_transforms = _column_transforms( + column_transforms, self.columns) + self.pad_values = _pad_values(pad_values, self.columns) + if adapter is not None and not callable(adapter): + raise TypeError("adapter must be callable or None.") + self.adapter = adapter + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + + if not query._table.options.row_tracking_enabled(): + raise ValueError( + "ContiguousWindowDataset requires row-tracking.enabled=true.") + + index, snapshot_id = _read_window_index( + query, self.group_key, self.order_key) + self.snapshot_id = snapshot_id + self._table = _pin_table(query._table, snapshot_id) + self._plans = {} + self._build_index(index) + + @classmethod + def from_query(cls, query, **kwargs): + """Build a contiguous-window Dataset from a ``ScanQuery``.""" + return cls(query, **kwargs) + + def __len__(self): + return int(self._anchor_groups.size) + + def __getitem__(self, index): + """Read one window by map-style Dataset index. + + Negative indices follow Python sequence semantics. The return value is + the pre-adapter mapping described by the class, or the adapter result + when an adapter is configured. + """ + anchor, row_ids = self._resolve_window(index) + rows = self._read_window_rows(row_ids) + anchor_row = ( + self._read_rows(row_ids[:1], self.anchor_columns)[0] + if self.anchor_columns else None + ) + return self._sample(anchor, rows, anchor_row) + + def __getitems__(self, indices): + """Read several Dataset indices while coalescing overlapping row IDs. + + The returned list preserves the requested index order and duplicates. + Coalescing affects only physical reads, not logical sample cardinality. + """ + windows = [self._resolve_window(index) for index in indices] + if not windows: + return [] + row_ids = list(dict.fromkeys( + row_id for _, window_row_ids in windows + for row_id in window_row_ids + )) + rows_by_id = dict(zip(row_ids, self._read_window_rows(row_ids))) + anchor_row_ids = list(dict.fromkeys( + window_row_ids[0] for _, window_row_ids in windows + )) + anchor_rows_by_id = ( + dict(zip( + anchor_row_ids, + self._read_rows(anchor_row_ids, self.anchor_columns), + )) + if self.anchor_columns else {} + ) + return [ + self._sample( + anchor, + [rows_by_id[row_id] for row_id in window_row_ids], + anchor_rows_by_id.get(window_row_ids[0]), + ) + for anchor, window_row_ids in windows + ] + + def __getstate__(self): + # Cached plans hold planning state of one process; each DataLoader + # worker plans the pinned snapshot once for itself. + state = self.__dict__.copy() + state["_plans"] = {} + return state + + def _resolve_window(self, index): + index = operator.index(index) + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError("window index out of range") + + group_index = int(self._anchor_groups[index]) + start = int(self._anchor_starts[index]) + valid_count = min( + self.window_size, int(self._group_lengths[group_index]) - start) + offset = int(self._group_starts[group_index]) + start + row_ids = self._row_ids[offset:offset + valid_count].tolist() + return (group_index, start, valid_count), row_ids + + def _sample(self, anchor, rows, anchor_row=None): + group_index, start, valid_count = anchor + padding_count = self.window_size - valid_count + padding_mask = torch.zeros(self.window_size, dtype=torch.bool) + if padding_count: + padding_mask[valid_count:] = True + sample = { + self.group_key: self._group_keys[group_index], + self.order_key: int(self._group_first_orders[group_index]) + start, + "is_pad": padding_mask, + } + for name in self.columns: + if name in self.anchor_columns: + values = [copy.deepcopy(anchor_row[name])] + else: + values = [copy.deepcopy(row[name]) for row in rows] + if padding_count and name not in self.anchor_columns: + pad_value = self.pad_values.get(name, values[-1]) + values.extend( + copy.deepcopy(pad_value) for _ in range(padding_count)) + transform = self.column_transforms.get(name) + sample[name] = transform(values) if transform is not None else values + if self.adapter is not None: + return self.adapter(sample) + return sample + + def _build_index(self, index): + """Validate index rows, then store row IDs, groups, and window anchors. + + Args: + index: Arrow table containing ``group_key``, ``order_key``, and + Paimon's ``_ROW_ID`` for the resolved snapshot. + + The index is kept as NumPy arrays of row IDs, per-group offsets, and + window anchors, plus one Python group value per group. Order values are + not stored per row: contiguity makes them the group's first value plus + the offset inside the group. + """ + group_column = index.column(self.group_key) + order_column = index.column(self.order_key) + row_id_column = index.column(SpecialFields.ROW_ID.name) + if row_id_column.null_count: + raise ValueError( + "ContiguousWindowDataset requires readable Paimon row IDs, " + "but %s contains null values." % SpecialFields.ROW_ID.name) + if group_column.null_count: + raise ValueError( + "%s must not contain null values." % self.group_key) + if order_column.null_count: + raise ValueError( + "%s must not contain null values." % self.order_key) + if not pa.types.is_integer(order_column.type): + raise ValueError( + "%s must contain integer values." % self.order_key) + if pa.types.is_floating(group_column.type) and pc.any( + pc.is_nan(group_column)).as_py(): + raise ValueError( + "%s must not contain NaN values, which never compare equal to " + "themselves and would split one group." % self.group_key) + + try: + ordered = index.sort_by([ + (self.group_key, "ascending"), + (self.order_key, "ascending"), + ]) + except pa.ArrowNotImplementedError: + raise ValueError( + "%s values must be mutually orderable." % self.group_key) + + group_values = _contiguous_array(ordered.column(self.group_key)) + order_values = _integer_numpy( + _contiguous_array(ordered.column(self.order_key))) + self._row_ids = _integer_numpy( + _contiguous_array(ordered.column(SpecialFields.ROW_ID.name))) + self._group_starts = _group_starts(group_values) + self._group_lengths = np.diff( + np.append(self._group_starts, len(self._row_ids))) + self._group_keys = group_values.take( + pa.array(self._group_starts, type=pa.int64())).to_pylist() + self._group_first_orders = order_values[self._group_starts] + self._validate_contiguity(order_values) + self._anchor_groups, self._anchor_starts = self._build_anchors() + + def _validate_contiguity(self, order_values): + if len(order_values) < 2: + return + same_group = np.ones(len(order_values) - 1, dtype=bool) + same_group[self._group_starts[1:] - 1] = False + steps = np.diff(order_values) + duplicate = same_group & (steps == 0) + if duplicate.any(): + position = int(np.flatnonzero(duplicate)[0]) + raise ValueError( + "Group %s has duplicate order value %r in %s." + % (self._group_of(position), + int(order_values[position]), self.order_key)) + broken = same_group & (steps != 1) + if broken.any(): + position = int(np.flatnonzero(broken)[0]) + raise ValueError( + "Group %s is not contiguous in %s: %s followed by %s." + % (self._group_of(position), self.order_key, + int(order_values[position]), int(order_values[position + 1]))) + + def _group_of(self, position): + group_index = int(np.searchsorted( + self._group_starts, position, side="right")) - 1 + return self._group_keys[group_index] + + def _build_anchors(self): + groups = [] + starts = [] + for group_index, length in enumerate(self._group_lengths): + length = int(length) + positions = np.arange(0, length, self.stride, dtype=np.int64) + valid_counts = np.minimum(self.window_size, length - positions) + incomplete = np.flatnonzero(valid_counts < self.window_size) + if incomplete.size: + if self.tail == "error": + first = int(incomplete[0]) + raise ValueError( + "Group %s has an incomplete window at %s: " + "window_size=%d, available=%d." + % (self._group_keys[group_index], + int(self._group_first_orders[group_index]) + + int(positions[first]), + self.window_size, + int(valid_counts[first]))) + if self.tail == "drop": + positions = positions[valid_counts == self.window_size] + if positions.size: + groups.append(np.full(positions.size, group_index, dtype=np.int64)) + starts.append(positions) + if not groups: + empty = np.zeros(0, dtype=np.int64) + return empty, empty.copy() + return np.concatenate(groups), np.concatenate(starts) + + def _read_window_rows(self, row_ids): + if not self._window_columns: + return [{} for _ in row_ids] + return self._read_rows(row_ids, self._window_columns) + + def _read_rows(self, row_ids, columns=None): + """Read projected rows by ID from the pinned snapshot. + + Args: + row_ids: Paimon row IDs to read. Their order and duplicates define + the returned row order. + columns: Projected value columns, or all Dataset columns when + omitted. + + Returns: + A list of row dictionaries aligned one-for-one with ``row_ids``. + The internal ``_ROW_ID`` field is removed, and BLOB descriptors are + resolved to their bodies. + """ + columns = self.columns if columns is None else columns + rows = self._plan_for(columns).read(row_ids) + row_id_column = SpecialFields.ROW_ID.name + by_row_id = {} + for row in rows: + by_row_id[int(row.pop(row_id_column))] = row + missing = [row_id for row_id in row_ids if row_id not in by_row_id] + if missing: + raise RuntimeError( + "Pinned snapshot %s did not return indexed row IDs %s." + % (self.snapshot_id, missing)) + return [by_row_id[row_id] for row_id in row_ids] + + def _plan_for(self, columns): + key = tuple(columns) + plan = self._plans.get(key) + if plan is None: + plan = _PinnedRowIdPlan(self._table, columns, self.blob_parallelism) + self._plans[key] = plan + return plan + + +class _PinnedRowIdPlan: + """One authorized scan plan of the pinned snapshot, read by row ID. + + Planning happens once per projection. Every read then narrows the cached + splits to the files covering the requested row IDs, so repeated item and + batch reads never revisit the snapshot's manifests. + """ + + def __init__(self, table, columns, blob_parallelism): + self._blob_parallelism = blob_parallelism + self._blob_columns = [ + field.name for field in table.fields + if field.name in columns + and (is_blob_type(field.type) or is_map_blob_type(field.type)) + ] + self._map_blob_columns = { + field.name for field in table.fields + if field.name in self._blob_columns and is_map_blob_type(field.type) + } + blob_column_set = set(self._blob_columns) + self._projection = ( + [name for name in columns if name not in blob_column_set] + + [SpecialFields.ROW_ID.name] + + self._blob_columns + ) + self._table = ( + table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"}) + if self._blob_columns else table + ) + self._splits = self._new_read_builder().new_scan().plan().splits() + _reject_masked_row_ids(self._splits) + + def read(self, row_ids): + """Return raw row dictionaries, including ``_ROW_ID``, for ``row_ids``.""" + requested = list(dict.fromkeys(row_ids)) + ranges = Range.to_ranges(requested) + splits = [ + pruned for pruned in ( + _prune_split_files(split, ranges) for split in self._splits) + if pruned is not None + ] + if not splits: + return [] + read_builder = self._new_read_builder() + predicate = read_builder.new_predicate_builder().is_in( + SpecialFields.ROW_ID.name, requested) + arrow = read_builder.with_filter(predicate).new_read().to_arrow(splits) Review Comment: [P2] Preserve exact row ranges when reading cached splits The cached plan contains unfiltered `DataSplit`s, and `_prune_split_files` only drops whole files. This later `_ROW_ID IN (...)` predicate cannot recover native row-range pushdown for a data-evolution read: `DataEvolutionSplitRead._push_down_predicate()` deliberately returns `None`, and its `row_ranges` remains `None` unless the supplied split is an `IndexedSplit`. As a result, every overlapping file is decoded before `FilterRecordBatchReader` (and the Arrow filter below) removes unrelated rows. That is on the documented shuffled training hot path: one batch can request many scattered windows and repeatedly decode large files. Please build a split-range index once, intersect requested ranges with matching splits, and pass `IndexedSplit`s while preserving `QueryAuthSplit`; `TorchDataset._SplitRangeIndex` / `_select_splits` is an existing implementation to reuse. Add a multi-file test that observes format-reader row ranges or decoded-row counts, since the six-row logical-call tests do not expose full-file decoding. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
