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 bad0c3e011706a4ec1e6394e0e6b70bff6dee4df Author: XiaoHongbo <[email protected]> AuthorDate: Tue Sep 15 17:29:49 2026 +0800 [python] Decode LeRobot cameras in parallel (#9832) --- .../pypaimon/multimodal/lerobot/dataset.py | 20 +++++++- .../pypaimon/tests/multimodal_lerobot_test.py | 56 ++++++++++++++++++++++ .../pypaimon/tests/multimodal_video_test.py | 1 + 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index b9220df4cb..1332b870cf 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -26,6 +26,7 @@ import sys from abc import ABC, abstractmethod from collections import OrderedDict from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor from functools import partial import pyarrow as pa @@ -64,6 +65,7 @@ _TORCH_DTYPE_NAMES = { } _IMAGE_READ_ATTEMPTS = 3 +_MAX_VIDEO_DECODE_WORKERS = 8 _CONTROL_FEATURES = frozenset({ "index", @@ -1546,6 +1548,7 @@ def _identity(values): def _decode_video_rows(row_groups, collators): + tasks = [] for collator in collators: targets = [] input_rows = [] @@ -1556,9 +1559,22 @@ def _decode_video_rows(row_groups, collators): input_rows.append(row) if not input_rows: continue - decoded = collator(input_rows) + tasks.append((collator, targets, input_rows)) + + if not tasks: + return + if len(tasks) == 1: + decoded_groups = [tasks[0][0](tasks[0][2])] + else: + with ThreadPoolExecutor( + max_workers=min( + len(tasks), _MAX_VIDEO_DECODE_WORKERS)) as executor: + decoded_groups = list(executor.map( + lambda task: task[0](task[2]), tasks)) + + for (collator, targets, _), decoded in zip(tasks, decoded_groups): for (rows, index), row in zip(targets, decoded): - rows[index] = row + rows[index][collator.output_column] = row[collator.output_column] def _normalize_index(index, size): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index a9790b4e16..a9d0ff6533 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -47,6 +47,7 @@ from pypaimon.multimodal.lerobot import load_from_lerobot from pypaimon.multimodal.lerobot.dataset import ( _PyAVVideoDecoder, _arrow_rows, + _decode_video_rows, _image_tensor, _index_names, _open_video_decoder, @@ -144,6 +145,61 @@ def _catalog_metadata(connection, name): class LeRobotValidationTest(unittest.TestCase): + def test_video_columns_decode_in_parallel(self): + barrier = threading.Barrier(2) + + class Collator: + + def __init__(self, video_column): + self.video_column = video_column + self.output_column = video_column + "_decoded" + + def __call__(self, rows): + barrier.wait(timeout=5) + return [dict( + row, + **{self.output_column: row[self.video_column] + 10}, + ) for row in rows] + + rows = { + 0: {"camera_a": 1, "camera_b": 2}, + 1: {"camera_a": 3, "camera_b": 4}, + } + _decode_video_rows( + [rows], [Collator("camera_a"), Collator("camera_b")]) + + self.assertEqual(11, rows[0]["camera_a_decoded"]) + self.assertEqual(12, rows[0]["camera_b_decoded"]) + self.assertEqual(13, rows[1]["camera_a_decoded"]) + self.assertEqual(14, rows[1]["camera_b_decoded"]) + + def test_video_column_decode_parallelism_is_bounded(self): + worker_counts = [] + + def executor(*args, **kwargs): + worker_counts.append(kwargs["max_workers"]) + return ThreadPoolExecutor(*args, **kwargs) + + class Collator: + + def __init__(self, index): + self.video_column = "camera_%d" % index + self.output_column = self.video_column + + def __call__(self, rows): + return rows + + rows = { + 0: {"camera_%d" % index: index for index in range(9)}, + } + with patch( + "pypaimon.multimodal.lerobot.dataset.ThreadPoolExecutor", + side_effect=executor): + _decode_video_rows( + [rows], [Collator(index) for index in range(9)]) + + self.assertEqual([8], worker_counts) + @unittest.skipUnless( av is not None and importlib.util.find_spec("torch") is not None, "PyAV and Torch are required for video decoding", diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py b/paimon-python/pypaimon/tests/multimodal_video_test.py index 8aaedfccc5..06c19cb053 100644 --- a/paimon-python/pypaimon/tests/multimodal_video_test.py +++ b/paimon-python/pypaimon/tests/multimodal_video_test.py @@ -155,6 +155,7 @@ class VideoFrameCollatorTest(unittest.TestCase): class Collator: video_column = "video" + output_column = "video" def __init__(self): self.calls = []
