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 3a858842a4 [python] Group LeRobot video frame decoding by payload
(#9799)
3a858842a4 is described below
commit 3a858842a498c61f1c385291873c81002bb90856
Author: chaoyang <[email protected]>
AuthorDate: Tue Sep 15 11:33:08 2026 +0800
[python] Group LeRobot video frame decoding by payload (#9799)
---
.../pypaimon/multimodal/lerobot/dataset.py | 20 +++---
paimon-python/pypaimon/multimodal/video.py | 38 ++++++++---
.../pypaimon/tests/multimodal_video_test.py | 79 ++++++++++++++++++++++
3 files changed, 120 insertions(+), 17 deletions(-)
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 6abbd34bab..b9220df4cb 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -1547,16 +1547,18 @@ def _identity(values):
def _decode_video_rows(row_groups, collators):
for collator in collators:
+ targets = []
+ input_rows = []
for rows in row_groups:
- indices = [
- index for index, row in rows.items()
- if collator.video_column in row
- ]
- if not indices:
- continue
- decoded = collator([rows[index] for index in indices])
- for index, row in zip(indices, decoded):
- rows[index] = row
+ for index, row in rows.items():
+ if collator.video_column in row:
+ targets.append((rows, index))
+ input_rows.append(row)
+ if not input_rows:
+ continue
+ decoded = collator(input_rows)
+ for (rows, index), row in zip(targets, decoded):
+ rows[index] = row
def _normalize_index(index, size):
diff --git a/paimon-python/pypaimon/multimodal/video.py
b/paimon-python/pypaimon/multimodal/video.py
index ebaaab74f4..c027a91e41 100644
--- a/paimon-python/pypaimon/multimodal/video.py
+++ b/paimon-python/pypaimon/multimodal/video.py
@@ -33,6 +33,10 @@ class VideoFrameCollator:
keeps Paimon independent of a particular video codec library while allowing
PyAV, TorchCodec, or an application decoder to be plugged in.
+ Batch inputs are grouped by physical video and decoded in ascending frame
+ order to avoid unnecessary decoder seeks. Rows are restored to their input
+ order before ``collate_fn`` is called.
+
The cache is process-local and keyed by physical video payload identity.
``collate_fn`` defaults to PyTorch's ``default_collate`` and may be
replaced
for decoders that already return batched objects.
@@ -82,7 +86,7 @@ class VideoFrameCollator:
self._ensure_process_local_cache()
single_row = isinstance(rows, Mapping)
input_rows = [rows] if single_row else list(rows)
- decoded_rows = [self._decode_row(row) for row in input_rows]
+ decoded_rows = self._decode_rows(input_rows)
if single_row:
return decoded_rows[0]
return self._collate(decoded_rows)
@@ -108,7 +112,29 @@ class VideoFrameCollator:
state["_owner_pid"] = None
return state
- def _decode_row(self, row):
+ def _decode_rows(self, rows):
+ decoded = [None] * len(rows)
+ grouped = OrderedDict()
+ for position, row in enumerate(rows):
+ output, descriptor = self._prepare_row(row)
+ if descriptor is None:
+ decoded[position] = output
+ continue
+ grouped.setdefault(descriptor.payload_descriptor, []).append(
+ (descriptor.frame_index, position, output)
+ )
+
+ for payload, frames in grouped.items():
+ decoder = self._decoder(payload)
+ for frame_index, position, output in sorted(
+ frames, key=lambda frame: frame[0]):
+ output[self.output_column] = self.decode_fn(
+ decoder, frame_index, output
+ )
+ decoded[position] = output
+ return decoded
+
+ def _prepare_row(self, row):
if not isinstance(row, Mapping):
raise ValueError("VideoFrameCollator expects row dictionaries.")
if self.video_column not in row:
@@ -120,7 +146,7 @@ class VideoFrameCollator:
output = dict(row)
if raw is None:
output[self.output_column] = None
- return output
+ return output, None
if hasattr(raw, "as_py"):
raw = raw.as_py()
if not VideoFrameDescriptor.is_video_frame_descriptor(raw):
@@ -138,11 +164,7 @@ class VideoFrameCollator:
"VideoFrameDescriptor without trailing bytes."
% self.video_column
)
- decoder = self._decoder(descriptor.payload_descriptor)
- output[self.output_column] = self.decode_fn(
- decoder, descriptor.frame_index, output
- )
- return output
+ return output, descriptor
def _decoder(self, descriptor):
resource = self._decoders.pop(descriptor, None)
diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py
b/paimon-python/pypaimon/tests/multimodal_video_test.py
index 080cbc2a9a..8aaedfccc5 100644
--- a/paimon-python/pypaimon/tests/multimodal_video_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_video_test.py
@@ -23,6 +23,7 @@ from types import SimpleNamespace
from pypaimon.common.file_io import FileIO
from pypaimon.multimodal import VideoFrameCollator
+from pypaimon.multimodal.lerobot.dataset import _decode_video_rows
from pypaimon.table.row.blob import VideoFrameDescriptor
@@ -94,6 +95,84 @@ class VideoFrameCollatorTest(unittest.TestCase):
)
self.assertEqual(descriptors[0], result[0]["video"])
+ def test_groups_and_sorts_frames_while_restoring_row_order(self):
+ descriptors = {
+ (video, frame): self._descriptor(
+ "episode-%s.mp4" % video,
+ ("video-%s" % video).encode(),
+ frame,
+ )
+ for video in ("one", "two")
+ for frame in (0, 1, 2, 3)
+ }
+ rows = [
+ {"request": "a", "video": descriptors["one", 3]},
+ {"request": "b", "video": descriptors["two", 2]},
+ {"request": "c", "video": descriptors["one", 1]},
+ {"request": "d", "video": descriptors["two", 0]},
+ ]
+ calls = []
+
+ def decode(decoder, frame, row):
+ video, decoded_frame = decoder.decode(frame)
+ calls.append((video, decoded_frame))
+ return row["request"], video, decoded_frame
+
+ collator = VideoFrameCollator(
+ self.table,
+ video_column="video",
+ decoder_factory=lambda stream: _Decoder(stream, []),
+ decode_fn=decode,
+ collate_fn=lambda decoded_rows: decoded_rows,
+ )
+ try:
+ result = collator(rows)
+ finally:
+ collator.close()
+
+ self.assertEqual(
+ [
+ (b"video-one", 1),
+ (b"video-one", 3),
+ (b"video-two", 0),
+ (b"video-two", 2),
+ ],
+ calls,
+ )
+ self.assertEqual(["a", "b", "c", "d"], [
+ row["request"] for row in result
+ ])
+ self.assertEqual(
+ [3, 2, 1, 0],
+ [row["frame"][2] for row in result],
+ )
+
+ def test_decodes_dataset_row_groups_in_one_batch(self):
+ row_groups = [
+ {4: {"request": "base", "video": b"base"}},
+ {1: {"request": "delta", "video": b"delta"}},
+ ]
+
+ class Collator:
+ video_column = "video"
+
+ def __init__(self):
+ self.calls = []
+
+ def __call__(self, rows):
+ self.calls.append([row["request"] for row in rows])
+ return [
+ dict(row, video="decoded-" + row["request"])
+ for row in rows
+ ]
+
+ collator = Collator()
+ _decode_video_rows(row_groups, [collator])
+
+ self.assertEqual([["base", "delta"]], collator.calls)
+ self.assertEqual("decoded-base", row_groups[0][4]["video"])
+ self.assertEqual("decoded-delta", row_groups[1][1]["video"])
+
def test_evicts_least_recently_used_decoder(self):
descriptors = [
self._descriptor("episode-%d.mp4" % index, bytes([index]), index)
