This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 9555f0e79c7d233f12f7a8776eeae0e0d406f7d8 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)
