This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 0353fb7a78 [python] Reuse CWD context plans and select window anchors
(#10567)
0353fb7a78 is described below
commit 0353fb7a782b5cade226c88466c6022cd5bbdc88
Author: Yann Byron <[email protected]>
AuthorDate: Fri Oct 16 13:02:09 2026 +0800
[python] Reuse CWD context plans and select window anchors (#10567)
---
docs/docs/pypaimon/multimodal-reading.md | 14 ++
docs/docs/pypaimon/pytorch.md | 45 ++++
paimon-python/pypaimon/benchmark/cwd_read.md | 66 ++++++
paimon-python/pypaimon/benchmark/cwd_read.py | 234 ++++++++++++++++++++
paimon-python/pypaimon/multimodal/query.py | 9 +-
.../pypaimon/multimodal/window_dataset.py | 111 ++++++----
.../tests/contiguous_window_dataset_test.py | 106 ++++++++-
.../pypaimon/tests/contiguous_window_plan_test.py | 241 +++++++++++++++++++++
8 files changed, 785 insertions(+), 41 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-reading.md
b/docs/docs/pypaimon/multimodal-reading.md
index c53b9951a9..83aed1a6f9 100644
--- a/docs/docs/pypaimon/multimodal-reading.md
+++ b/docs/docs/pypaimon/multimodal-reading.md
@@ -256,6 +256,20 @@ The stride controls anchor spacing, not the spacing inside
a window:
| `pad` | Keep scheduled anchors and repeat the nearest endpoint for missing
positions. |
| `error` | Fail construction if any scheduled anchor has an incomplete
window. |
+`anchor_row_ids` optionally selects scheduled anchors by Paimon `_ROW_ID`
+without filtering their readable context. It preserves the original stride
+and group/order sequence; duplicates do not repeat samples. Boundary handling
+then applies to selected anchors only. IDs must come from the same fixed
+snapshot and belong to the context query. The query's filters and limit still
+define context boundaries, so restrict it to complete episodes when their
+original history, future and stride must be preserved. See
+[task-local context and anchor
selection](pytorch#task-local-context-and-anchor-selection)
+for a metadata-only selection example and distributed ownership guidance.
+
+The Dataset plans and authorizes all requested columns once, reads only the
+index columns at construction, and reuses the resulting split metadata for
+lazy payload reads across projections and workers.
+
`pad_values={"action": [0.0, 0.0]}` overrides endpoint repetition for that
column. Supply values in the same raw representation as a table cell. Padding
happens before column transforms, so normalization also affects constant
padding.
diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md
index 3915b5ab22..5deedfe8da 100644
--- a/docs/docs/pypaimon/pytorch.md
+++ b/docs/docs/pypaimon/pytorch.md
@@ -275,6 +275,51 @@ torch.stack}` converts the raw list of `CHW` frames to
`TCHW`. Install
`pypaimon[torch,video]` for PyAV, or install TorchCodec separately. Select the
backend with `video_backend="pyav"` or `"torchcodec"`; omitting the option
selects an available backend.
+
+### Task-local context and anchor selection
+
+The scan query defines the readable context, including its filters and limit.
+For a task or rank-local dataset, restrict that query to complete episodes.
+Use `anchor_row_ids` to select training anchors without removing the history
+or future frames needed by their windows:
+
+```python
+context = frames.scan(tag_name="train-v1").where("task_id = 'pick'")
+eligible = (
+ frames.scan(tag_name="train-v1")
+ .where("task_id = 'pick' AND eligible = true")
+ .select(["episode_index", "frame_index"])
+ .with_row_id()
+ .to_arrow()
+)
+dataset = context.to_contiguous_window_dataset(
+ columns=["state", "image", "action"],
+ frame_offsets={"state": [-2, -1, 0], "action": list(range(16))},
+ anchor_row_ids=eligible.column("_ROW_ID").to_pylist(),
+ boundary="pad",
+)
+```
+
+Both scans must resolve to the same fixed snapshot. Row IDs outside the
+context raise `ValueError`. Selection intersects the original stride schedule;
+it neither changes its starting position nor orders or repeats samples.
+Duplicates in `anchor_row_ids` are ignored. An empty selection produces an
+empty dataset. Boundary handling applies only to selected scheduled anchors.
+
+The context is planned once for the index and all requested value columns.
+Only index columns are read at construction; payloads remain lazy. All
+projections and DataLoader workers reuse that plan's split metadata and apply
+the authorization rules resolved for the complete request at construction.
+There is no per-projection full-table planning when the first batch is read.
+
+A context filter or limit that cuts an episode changes its boundaries. It
+does not automatically supply missing history or future frames. Use anchor
+selection for frame eligibility; keep complete episodes in the context when
+preserving their original boundaries and stride. Windows can still read across
+physical file/split boundaries inside those episodes. If context has already
+been divided among ranks, use a local sampler instead of sharding it again
+with a world-size `DistributedSampler`; coordinate rank step counts separately.
+
## File Format Metadata Cache
Reusable PyArrow Dataset metadata is cached across reads. Configure its
estimated
diff --git a/paimon-python/pypaimon/benchmark/cwd_read.md
b/paimon-python/pypaimon/benchmark/cwd_read.md
new file mode 100644
index 0000000000..844df197b7
--- /dev/null
+++ b/paimon-python/pypaimon/benchmark/cwd_read.md
@@ -0,0 +1,66 @@
+<!--
+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.
+-->
+
+# CWD planning and batch-read benchmark
+
+This CPU microbenchmark measures construction, first batch, and steady batch
+reads of the same snapshot and sample sequence. It does not run a model, decode
+images, compare HDF5, or establish GPU training performance. OS caches are
+uncontrolled; a fresh process is not a cold-cache run.
+
+The generated Parquet table is partitioned by task. Every episode is written in
+a separate commit, with scalar values and deterministic BLOB bytes. Samples
read
+scalar offsets `[-2, 0, 1, ..., 15]` and only the anchor BLOB. Every returned
ID,
+scalar, BLOB and padding mask is checked against an independent generator.
+Preparation and sample verification are outside the read timers.
+
+```bash
+python -m pypaimon.benchmark.cwd_read selfcheck
+python -m pypaimon.benchmark.cwd_read prepare --warehouse /tmp/cwd-warehouse
+python -m pypaimon.benchmark.cwd_read run --warehouse /tmp/cwd-warehouse
--scope task
+python -m pypaimon.benchmark.cwd_read run --warehouse /tmp/cwd-warehouse
--scope all
+```
+
+Preparation refuses an existing directory. The manifest pins its snapshot and
+data dimensions. The run seed shuffles the full eligible index space; requests
+cycle that order if the requested count exceeds the dataset. The first batch is
+reported separately and excluded from steady throughput. Construction includes
+opening the catalog/table and building the CWD index. `build_and_read_s` is
their
+sum, excluding sample verification and process/import startup.
+
+For a source comparison, invoke **this same script by its absolute path** with
+`PYTHONPATH` set separately to each checkout's `paimon-python` directory. This
+also works when the baseline does not contain this benchmark. Check the
reported
+`source` and `window_dataset_sha256`; record the corresponding commit and diff
+alongside the results. Use one immutable prepared warehouse, the same CLI
+arguments and interpreter, and fresh processes in alternating
baseline/candidate
+order. Perform one discarded warm-up pair and at least five measured pairs per
+scope. Compare median and range and require identical `sample_digest` and
+`manifest_sha256` before comparing throughput. Task-local and all-task results
+must be reported separately.
+
+Run `--diagnostic` separately to record every `TableScan.plan` call's phase,
+split count and distinct data-file count. These are **planned** files, not
opened
+files or transferred bytes. Diagnostic timings are not performance results.
+
+A 30% throughput increase means `candidate samples_per_s / baseline >= 1.30`;
+it is not a 30% latency reduction. Report construction, first batch and steady
+reads separately: reducing planning time alone does not establish higher steady
+throughput. This small synthetic workload and its scalar/BLOB sizes are
explicit
+inputs, not a substitute for a real workload or a claim about object storage.
diff --git a/paimon-python/pypaimon/benchmark/cwd_read.py
b/paimon-python/pypaimon/benchmark/cwd_read.py
new file mode 100644
index 0000000000..7ba28bd663
--- /dev/null
+++ b/paimon-python/pypaimon/benchmark/cwd_read.py
@@ -0,0 +1,234 @@
+# 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.
+
+"""CPU benchmark of pinned CWD planning and ordered scalar/BLOB batch reads.
+
+Prepare once, then run the same script against different source checkouts.
+OS caches are uncontrolled. Diagnostic instrumentation is a separate run.
+"""
+
+import argparse
+from contextlib import nullcontext
+import hashlib
+import json
+from pathlib import Path
+import platform
+import random
+import struct
+import tempfile
+import time
+from unittest.mock import patch
+
+import pyarrow as pa
+import torch
+
+import pypaimon
+import pypaimon.multimodal as pmm
+from pypaimon.multimodal import window_dataset
+from pypaimon.read.table_scan import TableScan
+
+
+_TABLE = "benchmark.cwd_frames"
+_MANIFEST = "cwd-benchmark.json"
+_OFFSETS = [-2] + list(range(16))
+
+
+def _payload(seed, episode, step, length):
+ key = "%d:%d:%d" % (seed, episode, step)
+ return hashlib.shake_256(key.encode()).digest(length)
+
+
+def prepare(warehouse, tasks, episodes, frames, blob_bytes, seed):
+ """Write one episode per commit; never append to an existing benchmark."""
+ if warehouse.exists():
+ raise ValueError("prepare requires a new warehouse directory")
+ if min(tasks, episodes, blob_bytes) < 1 or frames <= 17:
+ raise ValueError("positive dimensions and at least 18 frames required")
+ table = pmm.connect(options={"warehouse": str(warehouse)}).create_table(
+ _TABLE,
+ schema=pa.schema([
+ pa.field("task", pa.int32(), nullable=False),
+ pa.field("episode", pa.int64(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("value", pa.int64(), nullable=False),
+ pa.field("payload", pa.large_binary(), nullable=False),
+ ]),
+ partitioned=["task"],
+ options={
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "true",
+ "deletion-vectors.enabled": "true",
+ "file.format": "parquet",
+ "vector.file.format": "parquet",
+ "blob-as-descriptor": "true",
+ },
+ )
+ for task in range(tasks):
+ for local_episode in range(episodes):
+ episode = task * episodes + local_episode
+ table.add([
+ {"task": task, "episode": episode, "step": step,
+ "value": episode * frames + step,
+ "payload": _payload(seed, episode, step, blob_bytes)}
+ for step in range(frames)
+ ])
+ snapshot = table.raw_table.snapshot_manager().get_latest_snapshot()
+ manifest = dict(version=1, table=_TABLE, tasks=tasks, episodes=episodes,
+ frames=frames, blob_bytes=blob_bytes, seed=seed,
+ snapshot_id=snapshot.id)
+ (warehouse / _MANIFEST).write_text(json.dumps(manifest, indent=2) + "\n")
+ return manifest
+
+
+def _indices(length, count, seed):
+ order = list(range(length))
+ random.Random(seed).shuffle(order)
+ return [order[index % length] for index in range(count)]
+
+
+def _check_and_hash(samples, indices, manifest, first_episode, digest):
+ windows_per_episode = manifest["frames"] - 17
+ for sample, index in zip(samples, indices):
+ episode = first_episode + index // windows_per_episode
+ step = index % windows_per_episode + 2
+ values = [episode * manifest["frames"] + step + offset
+ for offset in _OFFSETS]
+ payload = _payload(manifest["seed"], episode, step,
manifest["blob_bytes"])
+ if ((sample["episode"], sample["step"]) != (episode, step)
+ or sample["value"] != values
+ or sample["payload"] != [payload]
+ or any(sample["value_is_pad"])
+ or any(sample["payload_is_pad"])):
+ raise AssertionError("CWD output differs from generated sample
oracle")
+ digest.update(struct.pack("<qq", episode, step))
+ digest.update(struct.pack("<%dq" % len(values), *values))
+ digest.update(payload)
+ if len(samples) != len(indices):
+ raise AssertionError("CWD returned a different batch size")
+
+
+def run(warehouse, scope, task, batch_size, batches, seed, diagnostic=False):
+ """Measure one fresh dataset; verify every sample outside fetch timers."""
+ manifest_bytes = (warehouse / _MANIFEST).read_bytes()
+ manifest = json.loads(manifest_bytes)
+ if manifest["version"] != 1:
+ raise ValueError("unsupported benchmark manifest version")
+ if not 0 <= task < manifest["tasks"] or min(batch_size, batches) < 1:
+ raise ValueError("invalid task, batch size or steady batch count")
+ events = []
+ phase = ["build"]
+ original = TableScan.plan
+
+ def recording_plan(scan):
+ plan = original(scan)
+ splits = plan.splits()
+ files = {file.file_name for split in splits for file in split.files}
+ events.append(dict(phase=phase[0], splits=len(splits),
files=len(files)))
+ return plan
+
+ instrumentation = (patch.object(TableScan, "plan", recording_plan)
+ if diagnostic else nullcontext())
+ with instrumentation:
+ started = time.perf_counter()
+ table = pmm.connect(options={"warehouse":
str(warehouse)}).get_table(manifest["table"])
+ query = table.scan(snapshot_id=manifest["snapshot_id"])
+ if scope == "task":
+ query = query.where("task = %d" % task)
+ dataset = query.to_contiguous_window_dataset(
+ columns=["value", "payload"], frame_offsets={"value": _OFFSETS},
+ group_key="episode", order_key="step", boundary="drop")
+ build_s = time.perf_counter() - started
+ expected = (manifest["frames"] - 17) * manifest["episodes"]
+ if scope == "all":
+ expected *= manifest["tasks"]
+ if len(dataset) != expected or dataset.snapshot_id !=
manifest["snapshot_id"]:
+ raise AssertionError("CWD cardinality or snapshot changed")
+ indices = _indices(len(dataset), (batches + 1) * batch_size, seed)
+ digest = hashlib.sha256()
+ first_episode = task * manifest["episodes"] if scope == "task" else 0
+ elapsed = []
+ for batch in range(batches + 1):
+ phase[0] = "first_batch" if batch == 0 else "steady"
+ selected = indices[batch * batch_size:(batch + 1) * batch_size]
+ started = time.perf_counter()
+ samples = dataset.__getitems__(selected)
+ elapsed.append(time.perf_counter() - started)
+ _check_and_hash(samples, selected, manifest, first_episode, digest)
+ steady_s = sum(elapsed[1:])
+ result = dict(
+ scope=scope, task=task if scope == "task" else None,
+ snapshot_id=manifest["snapshot_id"], dataset_samples=len(dataset),
+ manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(),
+ seed=seed, batch_size=batch_size, steady_batches=batches,
+ build_s=build_s, first_batch_s=elapsed[0], steady_s=steady_s,
+ steady_samples=batches * batch_size,
+ samples_per_s=batches * batch_size / steady_s,
+ build_and_read_s=build_s + sum(elapsed),
+ sample_digest=digest.hexdigest(), verified_samples=len(indices),
+ cache_control="uncontrolled", device="cpu", num_workers=0,
+ diagnostic=diagnostic, plan_events=events,
+ source=str(Path(pypaimon.__file__).resolve()),
+
window_dataset_sha256=hashlib.sha256(Path(window_dataset.__file__).read_bytes()).hexdigest(),
+ python=platform.python_version(), platform=platform.platform(),
+ pyarrow=pa.__version__, torch=torch.__version__,
+ )
+ return result
+
+
+def selfcheck():
+ """Exercise a real partitioned table and independently check every
value."""
+ assert _indices(3, 8, 7) == _indices(3, 8, 7)
+ assert sorted(_indices(3, 3, 7)) == [0, 1, 2]
+ with tempfile.TemporaryDirectory(prefix="cwd_benchmark_") as directory:
+ warehouse = Path(directory) / "warehouse"
+ prepare(warehouse, 2, 2, 24, 32, 7)
+ for scope in ("task", "all"):
+ plain = run(warehouse, scope, 1, 4, 2, 9)
+ traced = run(warehouse, scope, 1, 4, 2, 9, True)
+ assert plain["sample_digest"] == traced["sample_digest"]
+ assert traced["plan_events"]
+ return {"selfcheck": "passed"}
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ commands = parser.add_subparsers(dest="command", required=True)
+ commands.add_parser("selfcheck")
+ create = commands.add_parser("prepare")
+ create.add_argument("--warehouse", type=Path, required=True)
+ create.add_argument("--tasks", type=int, default=16)
+ create.add_argument("--episodes", type=int, default=8, help="episodes per
task")
+ create.add_argument("--frames", type=int, default=128, help="frames per
episode")
+ create.add_argument("--blob-bytes", type=int, default=1024)
+ create.add_argument("--seed", type=int, default=20261013)
+ read = commands.add_parser("run")
+ read.add_argument("--warehouse", type=Path, required=True)
+ read.add_argument("--scope", choices=("task", "all"), required=True)
+ read.add_argument("--task", type=int, default=0)
+ read.add_argument("--batch-size", type=int, default=16)
+ read.add_argument("--batches", type=int, default=32, help="steady batches
after first batch")
+ read.add_argument("--seed", type=int, default=20261013)
+ read.add_argument("--diagnostic", action="store_true")
+ args = vars(parser.parse_args())
+ command = args.pop("command")
+ if "warehouse" in args:
+ args["warehouse"] = args["warehouse"].expanduser().resolve()
+ result = {"prepare": prepare, "run": run, "selfcheck":
selfcheck}[command](**args)
+ print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index bba7073da9..4471d841ef 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -185,7 +185,8 @@ class ScanQuery:
fps=None,
tolerance_s=None,
boundary=None,
- video_backend=None):
+ video_backend=None,
+ anchor_row_ids=None):
"""Build snapshot-pinned training windows from this single-table scan.
The Dataset indexes only ``group_key``, ``order_key``, and Paimon row
@@ -224,6 +225,11 @@ class ScanQuery:
Drop retains only anchors valid for every requested offset.
video_backend: Decoder for video-frame-field columns: ``pyav``,
``torchcodec``, or ``None`` to select the available backend.
+ anchor_row_ids: Optional iterable of Paimon row IDs selecting
+ anchors within this scan's context. Does not remove history or
+ future frames, change stride, reorder or repeat samples. IDs
+ outside the context raise ValueError. This scan, including its
+ filters and limit, still defines the context's group bounds.
Returns:
A snapshot-pinned ``ContiguousWindowDataset``. See that class for
@@ -249,6 +255,7 @@ class ScanQuery:
tolerance_s=tolerance_s,
boundary=boundary,
video_backend=video_backend,
+ anchor_row_ids=anchor_row_ids,
)
def to_ray(
diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py
b/paimon-python/pypaimon/multimodal/window_dataset.py
index 8e1dd443df..4a8209c955 100644
--- a/paimon-python/pypaimon/multimodal/window_dataset.py
+++ b/paimon-python/pypaimon/multimodal/window_dataset.py
@@ -77,6 +77,15 @@ class ContiguousWindowDataset(Dataset):
Only one of the three window definitions may be supplied. Without any,
every column uses ``[0]``. ``stride`` schedules anchors from each group's
first row; boundary handling never shifts that schedule.
+ ``anchor_row_ids`` optionally selects scheduled anchors by Paimon row ID
+ without removing their context. IDs must belong to the context query;
+ duplicates do not repeat samples and selection order does not change the
+ group/order schedule. The query, including its limit, defines group bounds.
+ Restrict it to complete groups to retain the original episode boundaries.
+
+ One scan plans and authorizes the index and all requested value columns.
+ Its split metadata is reused by every projection and DataLoader worker;
+ payloads remain lazy. Authorization is resolved at Dataset construction.
``blob_parallelism`` controls concurrent BLOB reads for each item or batch.
Video frame columns retain their descriptors until the requested rows are
grouped and decoded, so overlapping windows reuse each physical frame read.
@@ -104,7 +113,8 @@ class ContiguousWindowDataset(Dataset):
fps=None,
tolerance_s=None,
boundary=None,
- video_backend=None):
+ video_backend=None,
+ anchor_row_ids=None):
if not isinstance(query, ScanQuery) or isinstance(query,
_PreFilterQuery):
raise TypeError(
"ContiguousWindowDataset is only supported on scan(), "
@@ -175,10 +185,11 @@ class ContiguousWindowDataset(Dataset):
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)
+ index, plan = _read_window_index(
+ query, self.group_key, self.order_key, self.columns)
+ self.snapshot_id = plan.snapshot_id
+ self._table = _pin_table(query._table, self.snapshot_id)
+ self._splits = plan.splits()
self._plans = {}
self._video_collators = {
name: VideoFrameCollator(
@@ -193,7 +204,7 @@ class ContiguousWindowDataset(Dataset):
)
for name in self._video_columns
}
- self._build_index(index)
+ self._build_index(index, anchor_row_ids)
@classmethod
def from_query(cls, query, **kwargs):
@@ -240,8 +251,7 @@ class ContiguousWindowDataset(Dataset):
]
def __getstate__(self):
- # Cached plans hold planning state of one process; each DataLoader
- # worker plans the pinned snapshot once for itself.
+ # Reuse split metadata in workers, but rebuild process-local readers.
state = self.__dict__.copy()
state["_plans"] = {}
return state
@@ -311,12 +321,13 @@ class ContiguousWindowDataset(Dataset):
return self.adapter(sample)
return sample
- def _build_index(self, index):
+ def _build_index(self, index, anchor_row_ids):
"""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.
+ anchor_row_ids: Optional selection within the context's row IDs.
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
@@ -366,7 +377,9 @@ class ContiguousWindowDataset(Dataset):
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()
+ anchor_mask = _anchor_mask(
+ ordered.column(SpecialFields.ROW_ID.name), anchor_row_ids)
+ self._anchor_groups, self._anchor_starts =
self._build_anchors(anchor_mask)
def _validate_contiguity(self, order_values):
if len(order_values) < 2:
@@ -394,7 +407,7 @@ class ContiguousWindowDataset(Dataset):
self._group_starts, position, side="right")) - 1
return self._group_keys[group_index]
- def _build_anchors(self):
+ def _build_anchors(self, anchor_mask):
groups = []
starts = []
minimum = min(min(offsets) for offsets in self.frame_offsets.values())
@@ -406,6 +419,9 @@ class ContiguousWindowDataset(Dataset):
for group_index, length in enumerate(self._group_lengths):
length = int(length)
positions = np.arange(0, length, self.stride, dtype=np.int64)
+ if anchor_mask is not None:
+ positions = positions[
+ anchor_mask[self._group_starts[group_index] + positions]]
lower = min(length, max(0, -minimum))
upper = max(0, min(length, length - maximum))
valid = (positions >= lower) & (positions < upper)
@@ -468,11 +484,14 @@ class ContiguousWindowDataset(Dataset):
key = tuple(columns)
plan = self._plans.get(key)
if plan is None:
+ tag_name = self._table.options.scan_tag_name()
+ if tag_name is not None and self.snapshot_id is not None:
+ _require_tag_snapshot(self._table, tag_name, self.snapshot_id)
plan = _PinnedRowIdPlan(
self._table,
columns,
self.blob_parallelism,
- self.snapshot_id,
+ self._splits,
self._video_columns,
)
self._plans[key] = plan
@@ -480,11 +499,10 @@ class ContiguousWindowDataset(Dataset):
class _PinnedRowIdPlan:
- """One authorized scan plan of the pinned snapshot, read by row ID.
+ """Read a projection from the window context's authorized splits by row ID.
- Planning happens once per projection. Every read then narrows the cached
- splits to the row ranges covering the requested row IDs, so repeated item
- and batch reads never revisit the snapshot's manifests.
+ Every read narrows the shared splits to the requested row ranges. Neither
+ new projections nor workers need to revisit the snapshot's manifests.
"""
def __init__(
@@ -492,10 +510,9 @@ class _PinnedRowIdPlan:
table,
columns,
blob_parallelism,
- snapshot_id,
+ splits,
video_columns=()):
self._blob_parallelism = blob_parallelism
- self._snapshot_id = snapshot_id
self._blob_columns = [
field.name for field in table.fields
if field.name in columns
@@ -524,16 +541,7 @@ class _PinnedRowIdPlan:
table.copy({CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true"})
if self._blob_columns else table
)
- plan = self._new_read_builder().new_scan().plan()
- # A tag can be replaced between indexing and this deferred plan, so the
- # plan is only usable while it still resolves the indexed snapshot.
- if snapshot_id is not None and plan.snapshot_id != snapshot_id:
- raise RuntimeError(
- "The window index was built from snapshot %s, but reading its "
- "rows now resolves snapshot %s; the pinned snapshot moved."
- % (snapshot_id, plan.snapshot_id))
- self._splits = plan.splits()
- _reject_masked_row_ids(self._splits)
+ self._splits = splits
self._split_ranges = [
row_ranges_for_split(split) for split in self._splits
]
@@ -547,10 +555,9 @@ class _PinnedRowIdPlan:
self._splits, self._split_ranges, self._split_range_index, ranges)
if not splits:
return []
- read_builder = self._new_read_builder()
- predicate = read_builder.new_predicate_builder().is_in(
- SpecialFields.ROW_ID.name, requested)
- arrow = read_builder.with_filter(predicate).new_read().to_arrow(splits)
+ # IndexedSplit already pushes the row ranges into the readers. Avoid
+ # running a second predicate scanner over the same selected rows.
+ arrow = self._new_read_builder().new_read().to_arrow(splits)
# Row ranges are merged per split, so a read can still return rows
# between requested IDs; drop them before BLOB bodies are fetched.
row_id_column = arrow.column(SpecialFields.ROW_ID.name)
@@ -645,19 +652,48 @@ def _group_starts(group_values):
)).astype(np.int64, copy=False)
-def _read_window_index(query, group_key, order_key):
+def _anchor_mask(row_ids, selection):
+ if selection is None:
+ return None
+ selected = []
+ maximum = np.iinfo(np.int64).max
+ for value in selection:
+ if isinstance(value, (bool, np.bool_)):
+ raise TypeError("anchor_row_ids must contain integers, not
booleans.")
+ try:
+ row_id = operator.index(value)
+ except TypeError:
+ raise TypeError("anchor_row_ids must contain integer row IDs.")
+ if not 0 <= row_id <= maximum:
+ raise ValueError("anchor_row_ids must contain nonnegative int64
row IDs.")
+ selected.append(row_id)
+ selected = pa.array(selected, type=pa.int64())
+ mask = pc.is_in(row_ids, value_set=selected)
+ missing = selected.filter(pc.invert(pc.is_in(
+ selected, value_set=row_ids.filter(mask))))
+ if len(missing):
+ raise ValueError(
+ "anchor_row_ids contains IDs outside the context query: %s."
+ % missing.slice(0, 10).to_pylist())
+ return mask.to_numpy(zero_copy_only=False)
+
+
+def _read_window_index(query, group_key, order_key, columns):
index_query = copy.copy(query)
- index_query._projection = [group_key, order_key]
+ index_columns = [group_key, order_key, SpecialFields.ROW_ID.name]
+ # Plan all requested columns: native planning may prune DE column files,
+ # and authorization must cover payload columns as well as the index.
+ index_query._projection = index_columns + list(columns)
index_query._include_row_id = True
read_builder = index_query._configured_read_builder()
plan = read_builder.new_scan().plan()
splits = plan.splits()
_reject_masked_row_ids(splits)
_reject_masked_index_keys(splits, group_key, order_key)
- index = read_builder.new_read().to_arrow(splits)
+ index =
read_builder.with_projection(index_columns).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
+ return index, plan
def _pin_table(table, snapshot_id):
@@ -701,7 +737,8 @@ def _require_tag_snapshot(table, tag_name, 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))
+ "built from snapshot %s; the pinned snapshot moved."
+ % (tag_name, resolved, snapshot_id))
def _open_default_video_decoder(stream, backend=None):
diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
index 0b5cb96dd5..ac3296295d 100644
--- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
+++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
@@ -33,6 +33,7 @@ from pypaimon.catalog.table_query_auth import
TableQueryAuthResult
from pypaimon.multimodal import window_dataset
from pypaimon.multimodal.query import ScanQuery
from pypaimon.multimodal.window_dataset import ContiguousWindowDataset
+from pypaimon.read.reader.filter_record_batch_reader import
FilterRecordBatchReader
from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
from pypaimon.read.table_scan import TableScan
@@ -145,6 +146,80 @@ class ContiguousWindowDatasetTest(unittest.TestCase):
self.assertEqual([False], dataset[0]["value_is_pad"].tolist())
self.assertEqual(103, dataset[-1]["value"][0])
+ def test_anchor_selection_preserves_context_across_files(self):
+ table = self.conn.create_table(
+ "anchor_context", schema=self._schema(), options=_TABLE_OPTIONS)
+ for start in (0, 3, 6):
+ table.add([self._row("episode-a", step)
+ for step in range(start, min(start + 3, 8))])
+ row_ids = {row["step"]: row["_ROW_ID"] for row in
+ table.scan().select(["step"]).with_row_id().to_list()}
+ dataset = table.scan().to_contiguous_window_dataset(
+ columns=["value", "payload"], group_key="episode",
order_key="step",
+ frame_offsets={"value": [-1, 0, 1], "payload": [0, 1]},
+ anchor_row_ids=[row_ids[5], row_ids[2], row_ids[5]])
+
+ self.assertEqual(2, len(dataset))
+ samples = dataset.__getitems__([1, 0, 1])
+ self.assertEqual([5, 2, 5], [sample["step"] for sample in samples])
+ self.assertEqual([[4, 5, 6], [1, 2, 3], [4, 5, 6]],
+ [sample["value"] for sample in samples])
+ self.assertEqual([[b"episode-a-5", b"episode-a-6"],
+ [b"episode-a-2", b"episode-a-3"],
+ [b"episode-a-5", b"episode-a-6"]],
+ [sample["payload"] for sample in samples])
+ self.assertTrue(all(not sample["value_is_pad"].any() for sample in
samples))
+
+ def
test_anchor_selection_keeps_stride_and_checks_only_selected_boundaries(self):
+ table = self._table()
+ row_ids = {row["step"]: row["_ROW_ID"] for row in
+ table.scan().where("episode = 'episode-b'")
+ .select(["step"]).with_row_id().to_list()}
+ kwargs = dict(
+ columns=["value"], group_key="episode", order_key="step",
+ frame_offsets={"value": [-1, 0, 1]}, stride=2, boundary="error")
+ dataset = table.scan().to_contiguous_window_dataset(
+ anchor_row_ids=(row_ids[step] for step in (1, 2)), **kwargs)
+
+ self.assertEqual(1, len(dataset))
+ self.assertEqual(2, dataset[0]["step"])
+ self.assertEqual([101, 102, 103], dataset[0]["value"])
+ with self.assertRaisesRegex(ValueError, "episode-b.*incomplete window
at 0"):
+ table.scan().to_contiguous_window_dataset(
+ anchor_row_ids=[row_ids[0]], **kwargs)
+
+ def test_empty_anchor_selection_has_no_windows_or_payload_reads(self):
+ with patch("pypaimon.multimodal.window_dataset.fetch_blob_bodies") as
fetch:
+ dataset = self._table().scan().to_contiguous_window_dataset(
+ columns=["value", "payload"], group_key="episode",
order_key="step",
+ frame_offsets={"value": [-2, 0, 2]}, boundary="error",
+ anchor_row_ids=[])
+ self.assertEqual(0, len(dataset))
+ self.assertEqual([], dataset[:])
+ with self.assertRaises(IndexError):
+ dataset[0]
+ fetch.assert_not_called()
+
+ def test_anchor_selection_rejects_invalid_or_out_of_context_ids(self):
+ table = self._table()
+ for selected, error, message in (
+ ([True], TypeError, "booleans"),
+ ([1.5], TypeError, "integer row IDs"),
+ (["1"], TypeError, "integer row IDs"),
+ ([-1], ValueError, "nonnegative int64"),
+ ([2 ** 63], ValueError, "nonnegative int64"),
+ ([999], ValueError, "outside the context query")):
+ with self.subTest(selected=selected):
+ with self.assertRaisesRegex(error, message):
+ self._dataset(table, anchor_row_ids=selected)
+
+ excluded_id = (table.scan().where("episode = 'episode-b'")
+ .select(["step"]).with_row_id().to_list()[0]["_ROW_ID"])
+ with self.assertRaisesRegex(ValueError, "outside the context query"):
+ table.scan().where("episode =
'episode-a'").to_contiguous_window_dataset(
+ columns=["value"], group_key="episode", order_key="step",
+ anchor_row_ids=[excluded_id])
+
def test_drop_offsets_keep_stride_anchored_at_group_start(self):
table = self._table()
dataset = table.scan().to_contiguous_window_dataset(
@@ -772,16 +847,17 @@ class ContiguousWindowDatasetTest(unittest.TestCase):
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):
+ def test_reuses_context_plan_across_projections_and_pickle(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)))))
+ restored = pickle.loads(pickle.dumps(dataset))
+ samples.extend(restored.__getitems__(list(range(len(dataset)))))
- self.assertEqual(2, plan.call_count)
+ self.assertEqual(0, plan.call_count)
self.assertEqual(
[[100, 101, 102], [101, 102, 103]] * 2,
[sample["value"] for sample in samples])
@@ -789,6 +865,30 @@ class ContiguousWindowDatasetTest(unittest.TestCase):
[[b"episode-b-0"], [b"episode-b-1"]] * 2,
[sample["payload"] for sample in samples])
+ @pytest.mark.python_read
+ def test_window_ranges_do_not_need_a_second_predicate_scan(self):
+ table = self.conn.create_table(
+ "sparse_reads", schema=self._schema(), options=_TABLE_OPTIONS)
+ table.add([self._row("episode-a", step) for step in range(12)])
+ dataset = table.scan().to_contiguous_window_dataset(
+ window_size=2, columns=["value"],
+ group_key="episode", order_key="step")
+ filtered_rows = []
+ original_filter = FilterRecordBatchReader._filter_batch
+
+ def record_filter(reader, batch):
+ filtered_rows.append(batch.num_rows)
+ return original_filter(reader, batch)
+
+ with patch.object(FilterRecordBatchReader, "_filter_batch",
record_filter):
+ samples = dataset.__getitems__([5, 1, 5])
+
+ self.assertEqual([[5, 6], [1, 2], [5, 6]],
+ [sample["value"] for sample in samples])
+ # IndexedSplit already selects these sparse rows; a predicate scanner
+ # would evaluate the same membership again for every fetched batch.
+ self.assertEqual(0, sum(filtered_rows))
+
def test_rejects_query_authorization_that_masks_row_ids(self):
table = self._table()
auth = TableQueryAuthResult(
diff --git a/paimon-python/pypaimon/tests/contiguous_window_plan_test.py
b/paimon-python/pypaimon/tests/contiguous_window_plan_test.py
new file mode 100644
index 0000000000..8dcd6e8315
--- /dev/null
+++ b/paimon-python/pypaimon/tests/contiguous_window_plan_test.py
@@ -0,0 +1,241 @@
+# 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 shutil
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+import pypaimon.multimodal as pmm
+from pypaimon import Schema
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.multimodal.table import MultimodalTable
+from pypaimon.read.native_plan import native_reader_available
+from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader
+from pypaimon.read.table_scan import TableScan
+
+
+class ContiguousWindowPlanTest(unittest.TestCase):
+
+ def setUp(self):
+ self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_window_plan_")
+ self.conn = pmm.connect(options={
+ "warehouse": os.path.join(self.temp_dir, "warehouse"),
+ })
+ self.schema = pa.schema([
+ pa.field("episode", pa.string(), nullable=False),
+ pa.field("step", pa.int32(), nullable=False),
+ pa.field("value", pa.int32()),
+ ])
+
+ def tearDown(self):
+ shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+ def _table(self, partitioned=None):
+ return self.conn.create_table(
+ "frames", schema=self.schema, partitioned=partitioned, options={
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "true",
+ "deletion-vectors.enabled": "true",
+ "file.format": "parquet",
+ })
+
+ @staticmethod
+ def _rows(episode, steps):
+ return [
+ {"episode": episode, "step": step, "value": step + 10}
+ for step in steps
+ ]
+
+ @staticmethod
+ def _windows(query):
+ return query.to_contiguous_window_dataset(
+ window_size=2, columns=["value"],
+ group_key="episode", order_key="step")
+
+ @pytest.mark.python_plan
+ @pytest.mark.python_read
+ def test_task_query_keeps_its_pruned_files_for_window_reads(self):
+ table = self._table(partitioned=["episode"])
+ table.add(self._rows("wanted", range(4)))
+ table.add(self._rows("wanted", range(4, 8)))
+ table.add(self._rows("unrelated", range(8)))
+ plans = []
+ opened = []
+ original_plan = TableScan.plan
+ original_reader = FormatPyArrowReader.__init__
+
+ def record_plan(scan):
+ plan = original_plan(scan)
+ plans.append(plan)
+ return plan
+
+ def record_reader(reader, *args, **kwargs):
+ opened.append(args[2])
+ original_reader(reader, *args, **kwargs)
+
+ with patch.object(TableScan, "plan", record_plan), \
+ patch.object(FormatPyArrowReader, "__init__", record_reader):
+ dataset = self._windows(table.scan().where(
+ "episode = 'wanted' AND step >= 4"))
+ samples = dataset.__getitems__(list(range(len(dataset))))
+
+ self.assertEqual([[14, 15], [15, 16], [16, 17]], [
+ sample["value"] for sample in samples])
+ # A deferred unfiltered payload plan would visit the other task too.
+ self.assertEqual(1, len(plans))
+ self.assertEqual([("wanted",)], [
+ tuple(split.partition.values) for split in plans[0].splits()])
+ planned_files = {
+ file.file_name for split in plans[0].splits() for file in
split.files
+ }
+ # File statistics may prune the earlier file in the wanted partition.
+ self.assertEqual(planned_files, {os.path.basename(path) for path in
opened})
+
+ @pytest.mark.python_plan
+ @pytest.mark.python_read
+ def test_payload_authorization_controls_index_membership_and_masking(self):
+ table = self._table()
+ table.add(self._rows("denied", range(4)))
+ table.add(self._rows("allowed", range(4)))
+ auth = TableQueryAuthResult(
+ filter=[json.dumps({
+ "kind": "LEAF",
+ "transform": {
+ "name": "FIELD_REF",
+ "fieldRef": {"name": "episode", "type": "STRING NOT NULL"},
+ },
+ "function": "EQUAL",
+ "literals": ["allowed"],
+ })],
+ column_masking={"value": json.dumps({"name": "NULL"})},
+ )
+
+ def authorize(select):
+ return auth if "value" in select else None
+
+ table.raw_table.catalog_environment.table_query_auth = (
+ lambda options, identifier: authorize)
+ dataset = self._windows(table.scan())
+ samples = dataset.__getitems__(list(range(len(dataset))))
+
+ # An index-only plan misses both rules; losing its wrapper leaks
values.
+ self.assertEqual(3, len(samples))
+ self.assertEqual(["allowed"] * 3, [sample["episode"] for sample in
samples])
+ self.assertEqual([0, 1, 2], [sample["step"] for sample in samples])
+ self.assertEqual([[None, None]] * 3, [sample["value"] for sample in
samples])
+
+ @pytest.mark.python_plan
+ @pytest.mark.python_read
+ def test_independent_payload_column_files_are_read(self):
+ self._check_independent_column_files(native=False)
+
+ @unittest.skipUnless(native_reader_available(),
+ "pypaimon-rust native reader API not installed")
+ def test_native_plan_retains_independent_payload_column_files(self):
+ self._check_independent_column_files(native=True)
+
+ def _check_independent_column_files(self, native):
+ table = self._table()
+ builder = table.raw_table.new_batch_write_builder()
+ keys = builder.new_write().with_write_type(["episode", "step"])
+ values = builder.new_write().with_write_type(["value"])
+ commit = builder.new_commit()
+ try:
+ keys.write_arrow(pa.Table.from_pylist(
+ self._rows("episode", range(4)),
+ schema=pa.schema([self.schema.field("episode"),
self.schema.field("step")]),
+ ))
+ values.write_arrow(pa.table(
+ {"value": [100, 200, 300, 400]},
+ schema=pa.schema([self.schema.field("value")]),
+ ))
+ messages = keys.prepare_commit() + values.prepare_commit()
+ for message in messages:
+ for file in message.new_files:
+ file.first_row_id = 0
+ commit.commit(messages)
+ finally:
+ keys.close()
+ values.close()
+ commit.close()
+
+ raw = table.raw_table.copy({
+ "scan.native-plan.enabled": str(native).lower(),
+ "read.native.enabled": str(native).lower(),
+ })
+ native_plans = []
+ original_native_plan = TableScan._try_native_plan
+
+ def record_native_plan(scan):
+ plan = original_native_plan(scan)
+ native_plans.append(plan)
+ return plan
+
+ with patch.object(TableScan, "_try_native_plan", record_native_plan):
+ dataset = self._windows(ScanQuery(raw))
+ samples = dataset.__getitems__([2, 0, 1])
+
+ self.assertEqual([[300, 400], [100, 200], [200, 300]], [
+ sample["value"] for sample in samples])
+ if native:
+ # Do not count a silent fallback as coverage of projection pruning.
+ self.assertEqual(1, len(native_plans))
+ self.assertIsNotNone(native_plans[0])
+
+ def test_predicate_limit_and_deleted_rows_define_the_context(self):
+ table = self._table()
+ table.add(self._rows("episode", range(7)))
+ table.delete(where="step = 1")
+
+ dataset = self._windows(table.scan().where("value >= 11").limit(4))
+ samples = dataset.__getitems__(list(range(len(dataset))))
+
+ # The deleted first match must not consume the limit. Visible steps
+ # are 2, 3, 4, 5; step 6 is outside the context despite sharing its
file.
+ self.assertEqual([2, 3, 4], [sample["step"] for sample in samples])
+ self.assertEqual([[12, 13], [13, 14], [14, 15]], [
+ sample["value"] for sample in samples])
+
+ @pytest.mark.python_read
+ def test_non_de_sparse_windows_keep_file_offsets(self):
+ # The multimodal creation facade requires DE. CWD also accepts raw
+ # row-tracked tables, so create this append-only fixture via catalog.
+ catalog = self.conn.catalog
+ catalog.create_database("default", True)
+ identifier = "default.frames"
+ catalog.create_table(identifier, Schema.from_pyarrow_schema(
+ self.schema, options={
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "false",
+ "file.format": "parquet",
+ }), False)
+ table = MultimodalTable(catalog, identifier,
catalog.get_table(identifier))
+ for start in (0, 4, 8):
+ table.add(self._rows("episode", range(start, start + 4)))
+ dataset = self._windows(table.scan())
+
+ samples = dataset.__getitems__([8, 2, 8])
+ self.assertEqual([8, 2, 8], [sample["step"] for sample in samples])
+ self.assertEqual([[18, 19], [12, 13], [18, 19]], [
+ sample["value"] for sample in samples])
