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 c749d9c9d9 [python] Support video frames in contiguous window datasets
(#10511)
c749d9c9d9 is described below
commit c749d9c9d99aad8eeb1521c5b91f93e34d4733de
Author: Yann Byron <[email protected]>
AuthorDate: Wed Oct 14 10:01:53 2026 +0800
[python] Support video frames in contiguous window datasets (#10511)
---
docs/docs/pypaimon/multimodal-reading.md | 25 ++++-
docs/docs/pypaimon/pytorch.md | 10 +-
docs/docs/pypaimon/video.md | 27 +++++
paimon-python/pypaimon/multimodal/query.py | 6 +-
.../pypaimon/multimodal/window_dataset.py | 116 ++++++++++++++++-----
.../tests/contiguous_window_dataset_test.py | 109 ++++++++++++++++++-
6 files changed, 258 insertions(+), 35 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-reading.md
b/docs/docs/pypaimon/multimodal-reading.md
index e4b02f70dc..c53b9951a9 100644
--- a/docs/docs/pypaimon/multimodal-reading.md
+++ b/docs/docs/pypaimon/multimodal-reading.md
@@ -315,9 +315,28 @@ window. These calls preserve their single `is_pad` mask
and previous shapes.
`anchor_columns` and `tail` require `window_size`; do not combine `tail` and
`boundary`. New applications should use field offsets and per-column masks.
-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.
+Columns configured by `video-frame-field` keep their descriptors until all rows
+for an item or batch have been selected. The Dataset groups descriptors by
+physical video, decodes each requested frame once, and restores the window
+order. Install `pypaimon[torch,video]` for PyAV, or install TorchCodec
+separately. Use `video_backend="pyav"` or `"torchcodec"` to require one
+backend, or omit it to select an available one:
+
+```python
+video_windows = frames.scan().to_contiguous_window_dataset(
+ columns=["state", "video", "action"],
+ frame_offsets={"video": [-1, 0], "action": list(range(16))},
+ group_key="episode_index",
+ order_key="frame_index",
+ boundary="pad",
+ column_transforms={"video": torch.stack},
+ video_backend="pyav",
+)
+```
+
+Without a video column transform, each window contains a list of `CHW` tensors.
+The `torch.stack` transform above produces `TCHW`; DataLoader adds the batch
+dimension. Decoder sessions and range caches are process-local and bounded.
### Distributed BLOB processing with Ray
diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md
index 320b824c85..3915b5ab22 100644
--- a/docs/docs/pypaimon/pytorch.md
+++ b/docs/docs/pypaimon/pytorch.md
@@ -268,9 +268,13 @@ the sampler's `set_epoch(epoch)` each epoch. Its padding
or drop behavior still
applies when the sample count is not divisible by the number of ranks. Dataset
version pinning does not save sampler progress or random augmentation state.
-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.
+Columns configured by `video-frame-field` are decoded after the requested rows
+have been coalesced. Their descriptor metadata remains intact, overlapping
+windows reuse the same decoded rows, and `column_transforms={"video":
+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.
## File Format Metadata Cache
Reusable PyArrow Dataset metadata is cached across reads. Configure its
estimated
diff --git a/docs/docs/pypaimon/video.md b/docs/docs/pypaimon/video.md
index 01e3521e80..30b3477550 100644
--- a/docs/docs/pypaimon/video.md
+++ b/docs/docs/pypaimon/video.md
@@ -269,3 +269,30 @@ Define the decoder class and function at module scope when
DataLoader uses the
`shuffle=False`; row-buffer shuffle can turn monotonic frame access into seeks
and require more open decoder sessions. Paimon's worker sharding is
split-level,
so it does not duplicate a split across workers.
+
+### Read video training windows
+
+`ContiguousWindowDataset` decodes `video-frame-field` columns directly. It
+keeps descriptors through row selection, groups overlapping requests by
+physical video, and decodes each selected row once before assembling windows:
+
+```python
+video_windows = frames.scan().to_contiguous_window_dataset(
+ columns=["state", "video", "action"],
+ frame_offsets={"video": [-2, -1, 0], "action": list(range(16))},
+ group_key="episode_id",
+ order_key="frame_index",
+ boundary="pad",
+ column_transforms={"video": torch.stack},
+ video_backend="pyav",
+)
+
+loader = DataLoader(video_windows, batch_size=32, num_workers=4)
+for batch in loader:
+ # video is [N, T, C, H, W].
+ train(batch["video"], batch["state"], batch["action"])
+```
+
+Install `pypaimon[torch,video]` for the PyAV path, or install TorchCodec
+separately. Omit `video_backend` to select an available decoder. The decoder
+and range cache stay local to each DataLoader worker.
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index 195c05e9ac..13aa615b80 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -183,7 +183,8 @@ class ScanQuery:
delta_timestamps=None,
fps=None,
tolerance_s=None,
- boundary=None):
+ boundary=None,
+ video_backend=None):
"""Build snapshot-pinned training windows from this single-table scan.
The Dataset indexes only ``group_key``, ``order_key``, and Paimon row
@@ -220,6 +221,8 @@ class ScanQuery:
defaulting to 1e-4. Only valid with delta_timestamps.
boundary: Handling at both group ends: drop (default), pad, or
error.
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.
Returns:
A snapshot-pinned ``ContiguousWindowDataset``. See that class for
@@ -244,6 +247,7 @@ class ScanQuery:
fps=fps,
tolerance_s=tolerance_s,
boundary=boundary,
+ video_backend=video_backend,
)
def to_ray(
diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py
b/paimon-python/pypaimon/multimodal/window_dataset.py
index 7b158eba83..8e1dd443df 100644
--- a/paimon-python/pypaimon/multimodal/window_dataset.py
+++ b/paimon-python/pypaimon/multimodal/window_dataset.py
@@ -20,6 +20,7 @@
import copy
import math
import operator
+from functools import partial
import numpy as np
import pyarrow as pa
@@ -30,6 +31,7 @@ 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.multimodal.video import VideoFrameCollator
from pypaimon.read.datasource.torch_dataset import (
SplitRangeIndex,
row_ranges_for_split,
@@ -76,8 +78,8 @@ class ContiguousWindowDataset(Dataset):
every column uses ``[0]``. ``stride`` schedules anchors from each group's
first row; boundary handling never shifts that schedule.
``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.
+ Video frame columns retain their descriptors until the requested rows are
+ grouped and decoded, so overlapping windows reuse each physical frame read.
"""
_BOUNDARY_POLICIES = ("drop", "pad", "error")
@@ -101,7 +103,8 @@ class ContiguousWindowDataset(Dataset):
delta_timestamps=None,
fps=None,
tolerance_s=None,
- boundary=None):
+ boundary=None,
+ video_backend=None):
if not isinstance(query, ScanQuery) or isinstance(query,
_PreFilterQuery):
raise TypeError(
"ContiguousWindowDataset is only supported on scan(), "
@@ -135,7 +138,14 @@ class ContiguousWindowDataset(Dataset):
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._video_columns = [
+ name for name in self.columns
+ if name in query._table.options.video_frame_fields()
+ ]
+ if video_backend not in (None, "torchcodec", "pyav"):
+ raise ValueError(
+ "video_backend must be None, 'torchcodec', or 'pyav'.")
+ self.video_backend = video_backend
self.anchor_columns = _anchor_columns(anchor_columns, self.columns)
self.frame_offsets = _frame_offsets(
self.columns, frame_offsets, delta_timestamps, fps, tolerance_s,
@@ -170,6 +180,19 @@ class ContiguousWindowDataset(Dataset):
self.snapshot_id = snapshot_id
self._table = _pin_table(query._table, snapshot_id)
self._plans = {}
+ self._video_collators = {
+ name: VideoFrameCollator(
+ self._table,
+ video_column=name,
+ decoder_factory=partial(
+ _open_default_video_decoder, backend=video_backend),
+ decode_batch_fn=_decode_video_frames,
+ output_column=name,
+ collate_fn=list,
+ range_parallelism=self.blob_parallelism,
+ )
+ for name in self._video_columns
+ }
self._build_index(index)
@classmethod
@@ -223,6 +246,24 @@ class ContiguousWindowDataset(Dataset):
state["_plans"] = {}
return state
+ def close(self):
+ """Close video decoders opened by this Dataset process."""
+ first_error = None
+ for collator in self._video_collators.values():
+ try:
+ collator.close()
+ except Exception as error:
+ if first_error is None:
+ first_error = error
+ if first_error is not None:
+ raise first_error
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
def _resolve_window(self, index):
index = operator.index(index)
if index < 0:
@@ -402,11 +443,16 @@ class ContiguousWindowDataset(Dataset):
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.
+ The internal ``_ROW_ID`` field is removed. Ordinary BLOB
+ descriptors are resolved to their bodies, while video descriptors
+ are decoded to frames.
"""
columns = self.columns if columns is None else columns
rows = self._plan_for(columns).read(row_ids)
+ for name in columns:
+ collator = self._video_collators.get(name)
+ if collator is not None:
+ rows = collator(rows)
row_id_column = SpecialFields.ROW_ID.name
by_row_id = {}
for row in rows:
@@ -423,7 +469,12 @@ class ContiguousWindowDataset(Dataset):
plan = self._plans.get(key)
if plan is None:
plan = _PinnedRowIdPlan(
- self._table, columns, self.blob_parallelism, self.snapshot_id)
+ self._table,
+ columns,
+ self.blob_parallelism,
+ self.snapshot_id,
+ self._video_columns,
+ )
self._plans[key] = plan
return plan
@@ -436,7 +487,13 @@ class _PinnedRowIdPlan:
and batch reads never revisit the snapshot's manifests.
"""
- def __init__(self, table, columns, blob_parallelism, snapshot_id):
+ def __init__(
+ self,
+ table,
+ columns,
+ blob_parallelism,
+ snapshot_id,
+ video_columns=()):
self._blob_parallelism = blob_parallelism
self._snapshot_id = snapshot_id
self._blob_columns = [
@@ -444,13 +501,18 @@ class _PinnedRowIdPlan:
if field.name in columns
and is_blob_file_type(field.type)
]
+ self._body_blob_columns = [
+ name for name in self._blob_columns if name not in video_columns
+ ]
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)
+ if field.name in self._body_blob_columns
+ and is_map_blob_type(field.type)
}
self._array_blob_columns = {
field.name for field in table.fields
- if field.name in self._blob_columns and
is_array_blob_type(field.type)
+ if field.name in self._body_blob_columns
+ and is_array_blob_type(field.type)
}
blob_column_set = set(self._blob_columns)
self._projection = (
@@ -495,21 +557,21 @@ class _PinnedRowIdPlan:
arrow = arrow.filter(pc.is_in(
row_id_column,
value_set=pa.array(requested, type=row_id_column.type)))
- if not self._blob_columns:
+ if not self._body_blob_columns:
return arrow.to_pylist()
bodies = fetch_blob_bodies(
self._table.file_io,
- arrow.select(self._blob_columns).to_pydict(),
- self._blob_columns,
+ arrow.select(self._body_blob_columns).to_pydict(),
+ self._body_blob_columns,
self._blob_parallelism,
self._map_blob_columns,
self._array_blob_columns)
- blob_column_set = set(self._blob_columns)
+ blob_column_set = set(self._body_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:
+ for name in self._body_blob_columns:
values = bodies[name]
if len(values) != len(rows):
raise RuntimeError(
@@ -642,15 +704,21 @@ def _require_tag_snapshot(table, tag_name, snapshot_id):
"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 _open_default_video_decoder(stream, backend=None):
+ # Keep the codec implementation shared with PaimonLeRobotDataset while the
+ # generic collator owns grouping, range reads, and process-local caching.
+ from pypaimon.multimodal.lerobot.dataset import _open_video_decoder
+ return _open_video_decoder(stream, backend=backend)
+
+
+def _decode_video_frames(decoder, frame_indices, unused_rows):
+ get_frames_at = getattr(decoder, "get_frames_at", None)
+ if get_frames_at is not None:
+ frames = get_frames_at(indices=frame_indices).data
+ # Iterating a batched tensor returns views backed by the whole batch.
+ # Give each row compact storage before window assembly deep-copies it.
+ return [frame.clone() for frame in frames]
+ return [decoder[index] for index in frame_indices]
def _frame_offsets(columns, offsets, timestamps, fps, tolerance_s,
diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
index 0fd9d218cb..0b5cb96dd5 100644
--- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
+++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py
@@ -21,6 +21,7 @@ import pickle
import shutil
import tempfile
import unittest
+from types import SimpleNamespace
from unittest.mock import patch
import pyarrow as pa
@@ -892,7 +893,7 @@ class ContiguousWindowDatasetTest(unittest.TestCase):
window_size=2, columns=["value"],
group_key="episode", order_key="step")
- def test_rejects_video_frame_columns_that_would_lose_frame_metadata(self):
+ def test_decodes_video_windows_and_reuses_overlapping_frames(self):
table = self.conn.create_table(
"videos",
schema=pa.schema([
@@ -905,12 +906,106 @@ class ContiguousWindowDatasetTest(unittest.TestCase):
"blob-as-descriptor": "true",
}),
)
+ video_path = os.path.join(self.temp_dir, "episode-a.mp4")
+ with open(video_path, "wb") as output:
+ output.write(b"encoded-video")
+ table.add_video(
+ pmm.Blob.from_local(video_path),
+ [
+ {"episode": "episode-a", "step": step}
+ for step in range(3)
+ ],
+ first_frame=7,
+ )
+ decoded = []
- with self.assertRaisesRegex(
- ValueError, "video frame columns.*frame_index"):
- table.scan().to_contiguous_window_dataset(
+ class Decoder:
+
+ def get_frames_at(self, *, indices):
+ decoded.append(tuple(indices))
+ return SimpleNamespace(data=torch.tensor(indices))
+
+ with patch(
+ "pypaimon.multimodal.lerobot.dataset._open_video_decoder",
+ return_value=Decoder()):
+ dataset = table.scan().to_contiguous_window_dataset(
window_size=2, columns=["video"],
group_key="episode", order_key="step")
+ try:
+ samples = dataset.__getitems__([0, 1])
+ finally:
+ close = getattr(dataset, "close", None)
+ if close is not None:
+ close()
+
+ self.assertEqual([(7, 8, 9)], decoded)
+ self.assertEqual(
+ [[7, 8], [8, 9]],
+ [[int(frame) for frame in sample["video"]] for sample in samples],
+ )
+ for sample in samples:
+ for frame in sample["video"]:
+ self.assertEqual(
+ frame.numel() * frame.element_size(),
+ frame.untyped_storage().nbytes(),
+ )
+
+ def test_decodes_encoded_video_windows_with_multiple_workers(self):
+ try:
+ import av # noqa: F401
+ from PIL import Image
+ except ImportError:
+ self.skipTest("PyAV and Pillow are required for video decoding")
+
+ table = self.conn.create_table(
+ "encoded_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",
+ }),
+ )
+ table.add_images_as_video(
+ [
+ Image.new("RGB", (8, 6), color)
+ for color in ("red", "green", "blue")
+ ],
+ [
+ {"episode": "episode-a", "step": step}
+ for step in range(3)
+ ],
+ fps=10,
+ codec="mpeg4",
+ pixel_format="yuv420p",
+ gop_size=1,
+ codec_options={"qscale": "3"},
+ )
+ dataset = table.scan().to_contiguous_window_dataset(
+ window_size=2,
+ columns=["video"],
+ group_key="episode",
+ order_key="step",
+ column_transforms={"video": torch.stack},
+ video_backend="pyav",
+ )
+ try:
+ batches = list(torch.utils.data.DataLoader(
+ dataset, batch_size=2, num_workers=2, shuffle=False))
+ finally:
+ dataset.close()
+
+ self.assertEqual(1, len(batches))
+ self.assertEqual((2, 2, 3, 6, 8), tuple(batches[0]["video"].shape))
+ first = batches[0]["video"][0, 0].float().mean(dim=(1, 2))
+ last = batches[0]["video"][1, 1].float().mean(dim=(1, 2))
+ self.assertGreater(first[0], first[1] + 100)
+ self.assertGreater(first[0], first[2] + 100)
+ self.assertGreater(last[2], last[0] + 100)
+ self.assertGreater(last[2], last[1] + 100)
def test_validates_configuration_and_scan_only_contract(self):
table = self._table()
@@ -949,6 +1044,12 @@ class ContiguousWindowDatasetTest(unittest.TestCase):
window_size=2, columns=["value"],
group_key="episode", order_key="step")
+ with self.assertRaisesRegex(ValueError, "video_backend must be"):
+ table.scan().to_contiguous_window_dataset(
+ window_size=2, columns=["value"],
+ group_key="episode", order_key="step",
+ video_backend="unsupported")
+
if __name__ == "__main__":
unittest.main()
