[python] Speed up LeRobot single-video window decoding

docs/docs/pypaimon/lerobot.md

@@ -259,6 +259,11 @@ Each open decoder uses memory and a file handle, so usage grows with cameras and
 `read.video.decode-parallelism` limits concurrent video decodes across all cameras in each
 Dataset worker (default: 16). Set it to 1 for serial decoding. Parallel decoding
 applies to video windows without `image_transforms` and uses more CPU and memory.
+`read.video.decoders-per-payload` (default: 4) also splits single-video windows across
+independent cached decoders, within those limits. Set it to 1 to disable this splitting;
+the reader automatically reduces the decoder count when overlapping windows would add
+more than 50% duplicate frame requests. Additional decoders can increase initialization
+cost and repeated reads.
 
 Subclass `PaimonDatasetReader` for a custom logical frame layout:
 

paimon-python/pypaimon/common/options/core_options.py

@@ -1181,6 +1181,15 @@ class CoreOptions:
             "PyPaimon LeRobot Dataset worker.")
     )
 
+    READ_VIDEO_DECODERS_PER_PAYLOAD: ConfigOption[int] = (
+        ConfigOptions.key("read.video.decoders-per-payload")
+        .int_type()
+        .default_value(4)
+        .with_description(
+            "Maximum independent decoders for single-payload LeRobot video windows, "
+            "bounded by read.video.decode-parallelism and read.video.max-open-decoders.")
+    )
+
     PARQUET_COLUMN_INDEX_ENABLED: ConfigOption[bool] = (
         ConfigOptions.key("parquet.filter.columnindex.enabled")
         .boolean_type()
@@ -1951,6 +1960,9 @@ def read_video_max_open_decoders(self) -> int:
     def read_video_decode_parallelism(self) -> int:
         return self._positive_video_option(CoreOptions.READ_VIDEO_DECODE_PARALLELISM)
 
+    def read_video_decoders_per_payload(self) -> int:
+        return self._positive_video_option(CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD)
+
     def _positive_video_option(self, option) -> int:
         raw = self.options.to_map().get(option.key())
         if isinstance(raw, (bool, float)):

paimon-python/pypaimon/multimodal/lerobot/dataset.py

@@ -279,6 +279,11 @@ def _set_frame_rows(
             if isinstance(table_options, CoreOptions)
             else CoreOptions.READ_VIDEO_DECODE_PARALLELISM.default_value()
         )
+        self._video_decoders_per_payload = (
+            table_options.read_video_decoders_per_payload()
+            if isinstance(table_options, CoreOptions)
+            else CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD.default_value()
+        )
         self._video_decode_pool = _VideoDecodePool()
         self._video_collators = [
             VideoFrameCollator(
@@ -455,6 +460,9 @@ def _get_items(self, indices, return_batch, share_memory):
             plans, rows, getattr(self, "_video_collators", ()),
             self._features, self.return_uint8,
             return_batch=return_batch, share_memory=share_memory,
+            decoders_per_payload=getattr(
+                self, "_video_decoders_per_payload",
+                CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD.default_value()),
             decode_pool=getattr(self, "_video_decode_pool", None),
             decode_parallelism=getattr(
                 self, "_video_decode_parallelism",
@@ -1439,21 +1447,56 @@ def _attach_task_labels(rows, task_names, subtask_names):
 
 def _decode_video_windows(
         plans, rows, collators, features, return_uint8, *,
-        return_batch=False, share_memory=False, decode_parallelism=1, decode_pool=None):
+        return_batch=False, share_memory=False, decode_parallelism=1, decode_pool=None,
+        decoders_per_payload=1):
     pool = decode_pool if decode_pool is not None else _VideoDecodePool()
     try:
         return _decode_video_windows_with_pool(
             plans, rows, collators, features, return_uint8,
             return_batch=return_batch, share_memory=share_memory,
-            decode_parallelism=decode_parallelism, executor=pool.get(decode_parallelism))
+            decode_parallelism=decode_parallelism, executor=pool.get(decode_parallelism),
+            decoders_per_payload=decoders_per_payload)
     finally:
         if decode_pool is None:
             pool.close()
 
 
+def _partition_video_windows(windows, count):
+    """Split windows while bounding duplicate frame requests across decoders."""
+    grouped = {}
+    for window in windows:
+        # Windows with the same requested frames should always share a decoder,
+        # even when their output order differs.
+        grouped.setdefault(tuple(sorted(set(window[1]))), []).append(window)
+    ordered = sorted(grouped.values(), key=lambda group: min(group[0][1]))
+    count = min(count, len(ordered))
+    serial_frames = len({
+        index for _, window in windows for index in window
+    })
+    while count > 1:
+        partitions = [
+            [
+                window
+                for group in ordered[len(ordered) * replica // count:
+                                     len(ordered) * (replica + 1) // count]
+                for window in group
+            ]
+            for replica in range(count)
+        ]
+        parallel_frames = sum(len({
+            index for _, window in partition for index in window
+        }) for partition in partitions)
+        # Bound additional frame requests to 50%. This retains parallelism for
+        # mostly independent windows without multiplying highly overlapping IO.
+        if 2 * parallel_frames <= 3 * serial_frames:
+            return partitions
+        count -= 1
+    return [windows]
+
+
 def _decode_video_windows_with_pool(
         plans, rows, collators, features, return_uint8, *,
-        return_batch, share_memory, decode_parallelism, executor):
+        return_batch, share_memory, decode_parallelism, executor, decoders_per_payload):
     import torch
 
     tasks = [c for c in collators if c.video_column in plans[0]["windows"]]
@@ -1487,6 +1530,18 @@ def decode(collator):
             requests.setdefault((payload, keyframe_index), []).append((
                 offset, [d.frame_index for d in window]))
 
+        if len(requests) == 1 and executor is not None:
+            request_key, windows = next(iter(requests.items()))
+            count = min(decoders_per_payload, decode_parallelism,
+                        collator.max_open_videos, len(windows))
+            if count > 1:
+                partitions = _partition_video_windows(windows, count)
+                if len(partitions) > 1:
+                    requests = {
+                        request_key + (replica,): partition
+                        for replica, partition in enumerate(partitions)
+                    }
+
         def read_frames(decoder, windows):
             get_frames = getattr(decoder, "get_frames_at", None)
             if not callable(get_frames):

paimon-python/pypaimon/multimodal/video.py

@@ -366,8 +366,8 @@ def _prepare_row(self, row):
             )
         return output, descriptor
 
-    def _decoder(self, descriptor, keyframe_index_descriptor=None):
-        cache_key = (descriptor, keyframe_index_descriptor)
+    def _decoder(self, descriptor, keyframe_index_descriptor=None, replica=0):
+        cache_key = (descriptor, keyframe_index_descriptor, replica)
         resource = self._decoders.pop(cache_key, None)
         if resource is not None:
             self._decoders[cache_key] = resource

paimon-python/pypaimon/tests/video_decode_parallelism_test.py

@@ -442,8 +442,171 @@ def blob(file_io, uri, offset, length):
             finally:
                 collator.close()
         self.assertTrue(torch.equal(*outputs))
+        outputs = []
+        single_plans = [{"windows": {"camera": indices}} for indices in (
+            [6, 1, 6], [0, 1, 2], [2, 3, 4], [4, 5, 7])]
+        for replicas in (1, 4):
+            collator = VideoFrameCollator(
+                SimpleNamespace(file_io="test"), video_column="camera",
+                decoder_factory=partial(_open_video_decoder, backend="pyav"),
+                decode_fn=identity, max_open_videos=4)
+            try:
+                with patch("pypaimon.multimodal.video.Blob.from_file", blob), \
+                        patch("torch.get_num_threads", return_value=1):
+                    outputs.append(_decode_video_windows(
+                        single_plans, rows, [collator],
+                        {"camera": {"dtype": "video", "shape": [16, 16, 3]}}, True,
+                        return_batch=True, decode_parallelism=4,
+                        decoders_per_payload=replicas)["camera"])
+                    self.assertEqual(len(collator._decoders), replicas)
+            finally:
+                collator.close()
+        self.assertTrue(torch.equal(*outputs))
+
+    def test_single_payload_windows_reuse_independent_decoders(self):
+        try:
+            import torch
+        except ImportError as error:
+            self.skipTest(str(error))
+        from pypaimon.multimodal.lerobot.dataset import _decode_video_windows
+        from pypaimon.table.row.blob import VideoFrameDescriptor
+
+        windows = [[7, 2, 7], [0, 1, 2], [5, 6, 7], [2, 3, 4]]
+        rows = {i: {"camera": VideoFrameDescriptor(
+            "one-video", 0, 1, i, -1, 0).serialize()} for i in range(8)}
+        plans = [{"windows": {"camera": window}} for window in windows]
+        features = {"camera": {"dtype": "video", "shape": [4, 5, 3]}}
+        for parallelism, capacity, replicas, expected in (
+                (1, 4, 4, 1), (4, 1, 4, 1), (2, 4, 4, 2),
+                (4, 4, 1, 1), (4, 4, 4, 4), (8, 8, 16, 4)):
+            with self.subTest(parallelism=parallelism, capacity=capacity, replicas=replicas):
+                collator = self.collator(capacity)
+                created = []
+                barrier = threading.Barrier(expected)
+
+                class Decoder:
+                    def __init__(self, stream):
+                        self.stream = stream
+                        self.closed = False
+                        created.append(self)
+
+                    def get_frames_at(self, indices):
+                        barrier.wait(timeout=10)
+                        # Non-contiguous decoder output exercises layout conversion.
+                        data = torch.tensor(indices, dtype=torch.uint8)[:, None, None, None]
+                        return SimpleNamespace(data=data.expand(-1, 4, 5, 3).permute(0, 3, 1, 2))
+
+                    def close(self):
+                        self.closed = True
+
+                collator.decoder_factory = Decoder
+                try:
+                    with patch("pypaimon.multimodal.video.Blob.from_file", self.streams), \
+                            patch("torch.get_num_threads", return_value=1):
+                        for return_batch in (False, True):
+                            for uint8 in (True, False):
+                                result = _decode_video_windows(
+                                    plans, rows, [collator], features, uint8,
+                                    return_batch=return_batch, share_memory=return_batch,
+                                    decode_parallelism=parallelism,
+                                    decoders_per_payload=replicas)["camera"]
+                                target = torch.tensor(windows, dtype=torch.uint8)[:, :, None, None, None]
+                                target = target.expand(-1, -1, 3, 4, 5)
+                                if not uint8:
+                                    target = target.float() / 255
+                                actual = result if return_batch else torch.stack(result)
+                                self.assertTrue(torch.equal(actual, target))
+                                self.assertTrue(actual.is_contiguous())
+                                if return_batch:
+                                    self.assertTrue(result.is_shared())
+                                result[0].zero_()
+                                self.assertTrue(torch.equal(result[2], target[2]))
+                        self.assertEqual(len(created), expected)
+                        self.assertEqual(len({id(d.stream) for d in created}), expected)
+                finally:
+                    collator.close()
+                self.assertTrue(all(d.closed and d.stream.closed for d in created))
+
+    def test_single_payload_split_bounds_duplicate_frame_requests(self):
+        try:
+            import torch
+        except ImportError as error:
+            self.skipTest(str(error))
+        from pypaimon.multimodal.lerobot.dataset import _decode_video_windows
+        from pypaimon.table.row.blob import VideoFrameDescriptor
+
+        features = {"camera": {"dtype": "video", "shape": [4, 5, 3]}}
+
+        def decode(windows):
+            rows = {i: {"camera": VideoFrameDescriptor(
+                "one-video", 0, 1, i, -1, 0).serialize()}
+                for i in {index for window in windows for index in window}}
+            plans = [{"windows": {"camera": window}} for window in windows]
+            collator = self.collator(4)
+            calls = []
+
+            class Decoder:
+                def __init__(self, stream):
+                    self.stream = stream
+
+                def get_frames_at(self, indices):
+                    calls.append(tuple(indices))
+                    data = torch.tensor(
+                        indices, dtype=torch.uint8)[:, None, None, None]
+                    return SimpleNamespace(
+                        data=data.expand(-1, 4, 5, 3).permute(0, 3, 1, 2))
+
+                def close(self):
+                    pass
+
+            collator.decoder_factory = Decoder
+            try:
+                with patch("pypaimon.multimodal.video.Blob.from_file", self.streams), \
+                        patch("torch.get_num_threads", return_value=1):
+                    result = _decode_video_windows(
+                        plans, rows, [collator], features, True,
+                        return_batch=True, decode_parallelism=4,
+                        decoders_per_payload=4)["camera"]
+                target = torch.tensor(
+                    windows, dtype=torch.uint8)[:, :, None, None, None]
+                target = target.expand(-1, -1, 3, 4, 5)
+                self.assertTrue(torch.equal(result, target))
+                return calls
+            finally:
+                collator.close()
+
+        repeated = decode([[0, 1, 2]] * 4)
+        self.assertEqual(repeated, [(0, 1, 2)])
+
+        overlapping = decode([list(range(i, i + 8)) for i in range(8)])
+        self.assertEqual(len(overlapping), 2)
+        self.assertEqual(sum(len(indices) for indices in overlapping), 22)
+        self.assertEqual(set().union(*map(set, overlapping)), set(range(15)))
+
+    def test_replica_cache_eviction_and_pickle(self):
+        collator = self.collator(2)
+        descriptor = BlobDescriptor("one-video", 0, 1)
+        with patch("pypaimon.multimodal.video.Blob.from_file", self.streams):
+            first = collator._decoder(descriptor, None, 0)
+            second = collator._decoder(descriptor, None, 1)
+            self.assertIs(collator._decoder(descriptor), first)
+            third = collator._decoder(descriptor, None, 2)
+            self.assertTrue(second.closed)
+            self.assertFalse(first.closed)
+            child = pickle.loads(pickle.dumps(collator))
+            self.assertFalse(child._decoders)
+            child.close()
+            collator.close()
+            self.assertTrue(first.closed and third.closed)
 
     def test_options(self):
+        self.assertEqual(CoreOptions(Options({})).read_video_decoders_per_payload(), 4)
+        self.assertEqual(CoreOptions(Options({"read.video.decoders-per-payload": "1"}))
+                         .read_video_decoders_per_payload(), 1)
+        for value in (0, -1, True, 1.5, "no", "1.5"):
+            with self.subTest(replicas=value), self.assertRaises(ValueError):
+                CoreOptions(Options({"read.video.decoders-per-payload": value})) \
+                    .read_video_decoders_per_payload()
         self.assertEqual(CoreOptions(Options({})).read_video_decode_parallelism(), 16)
         self.assertEqual(CoreOptions(Options({"read.video.decode-parallelism": "1"}))
                          .read_video_decode_parallelism(), 1)