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 22b8d7bd5b [python] Decode LeRobot cameras in parallel (#9832)
22b8d7bd5b is described below
commit 22b8d7bd5b4fe7b85a92a96f03f8e8f93eb3bb12
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 = []
