From 7090767dd402e8df10f6159ce0f27ee24b507a13 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 01:24:21 -0700 Subject: [PATCH 01/10] [python] Read LeRobot video datasets for training --- docs/docs/pypaimon/lerobot.md | 17 +- .../pypaimon/multimodal/lerobot/dataset.py | 206 ++++++++++++++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 39 ++++ 3 files changed, 236 insertions(+), 26 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 722d63f3f2a2..e226f64048ac 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -110,8 +110,8 @@ Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested Video features map to `BLOB`. Frame rows reference MP4 payloads copied once per aligned file group. Video imports use the video grouping policy and check rolling before each Episode. They require a bucket-unaware table. Read them -with a Paimon scan and `VideoFrameCollator`; `PaimonLeRobotDataset` currently -supports image features only. +with a Paimon scan and `VideoFrameCollator`, or let `PaimonLeRobotDataset` +decode the referenced frames with TorchCodec. ## Capture LeRobot frames directly into Paimon @@ -208,10 +208,10 @@ pin one named snapshot on every component. ## Train with Paimon LeRobot data -For map-style training, read a tagged table group created by -`load_from_lerobot` directly from Paimon. `PaimonLeRobotDataset` requires the -complete table group; a frame-only table created by `PaimonLeRobotWriter` is -not sufficient. +For map-style training, read a tagged image- or video-backed table group +created by `load_from_lerobot` directly from Paimon. `PaimonLeRobotDataset` +requires the complete table group; a frame-only table created by +`PaimonLeRobotWriter` is not sufficient. ```python from torch.utils.data import DataLoader @@ -226,4 +226,7 @@ loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) If `tag_name` is omitted, the latest snapshots are used. Metadata is available through `dataset.meta`. Frame lookups use the BTree on `index`; payload columns -remain lazy. +remain lazy. Video decoder sessions are cached per DataLoader worker; set +`max_open_videos` to bound the number retained for each video feature. Video +decoding uses TorchCodec when available and otherwise PyAV; set +`video_backend="torchcodec"` or `"pyav"` to select one explicitly. diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 49e26097e0e9..b8f087c69a51 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -24,6 +24,7 @@ import operator import os import sys +from functools import partial import pyarrow as pa @@ -42,6 +43,7 @@ _validate_lerobot_schema, ) from pypaimon.multimodal.table import _target_schema, _time_travel_table +from pypaimon.multimodal.video import VideoFrameCollator from pypaimon.read.query_auth_split import QueryAuthSplit @@ -78,7 +80,7 @@ class PaimonLeRobotDataset: LeRobot metadata is resolved from the Paimon table group and remains available through :attr:`meta`. - Set ``return_uint8=True`` to keep 8-bit images in their decoded + Set ``return_uint8=True`` to keep 8-bit visual frames in their decoded ``torch.uint8`` representation instead of normalizing them to float32. Higher-bit-depth images retain the existing float32 behavior. """ @@ -93,6 +95,8 @@ def __init__( delta_timestamps=None, tolerance_s=1e-4, blob_parallelism=16, + video_backend=None, + max_open_videos=8, return_uint8=False): if sys.version_info < (3, 10): raise RuntimeError( @@ -109,6 +113,12 @@ def __init__( raise ValueError("tolerance_s must be finite and non-negative.") self.blob_parallelism = _positive_int( blob_parallelism, "blob_parallelism") + if video_backend not in (None, "torchcodec", "pyav"): + raise ValueError( + "video_backend must be None, 'torchcodec', or 'pyav'.") + self.video_backend = video_backend + self.max_open_videos = _positive_int( + max_open_videos, "max_open_videos") if not isinstance(return_uint8, bool): raise TypeError("return_uint8 must be a boolean.") self.return_uint8 = return_uint8 @@ -130,15 +140,11 @@ def _init_metadata(self): name for name, feature in self._features.items() if feature.get("dtype") == "image" ] - video_keys = [ + self._video_keys = [ name for name, feature in self._features.items() if feature.get("dtype") == "video" ] - if video_keys: - raise NotImplementedError( - "PaimonLeRobotDataset currently supports image-backed " - "features only; video features are not yet supported: %s" - % video_keys) + self._visual_keys = self._image_keys + self._video_keys self._total_frames = int( _metadata_member( @@ -230,6 +236,19 @@ def _init_reader(self, raw_table, info): self._read_table, snapshot, splits) self._validation_context = validation_context self._file_io = self._read_table.file_io + self._video_collators = [ + VideoFrameCollator( + self._read_table, + video_column=key, + decoder_factory=partial( + _open_video_decoder, backend=self.video_backend), + decode_fn=_decode_video_frame, + output_column=key, + max_open_videos=self.max_open_videos, + collate_fn=_identity, + ) + for key in self._video_keys + ] self._task_names = validation_context["task_names"] self._subtask_names = validation_context["subtask_names"] self._delta_projection = None @@ -319,21 +338,30 @@ def __getitems__(self, indices): self._image_keys, self.blob_parallelism, ) - converted = { - position: _torch_row( - row, self._features, self.return_uint8) - for position, row in base_rows.items() - } - converted.update({ - position: _torch_row( - row, self._features, self.return_uint8) - for position, row in delta_rows.items() - }) + _decode_image_rows( + row_groups, + self._image_keys, + self._features, + self.return_uint8, + ) break except OSError: if attempt + 1 == _IMAGE_READ_ATTEMPTS: raise + _decode_video_rows( + row_groups, getattr(self, "_video_collators", ())) + converted = { + position: _torch_row( + row, self._features, self.return_uint8) + for position, row in base_rows.items() + } + converted.update({ + position: _torch_row( + row, self._features, self.return_uint8) + for position, row in delta_rows.items() + }) + import torch duplicates = _duplicate_indices(plans) result = [] @@ -350,11 +378,34 @@ def __getitems__(self, indices): ]) item.update(plan["padding"]) if self.image_transforms is not None: - for key in self._image_keys: + for key in self._visual_keys: item[key] = self.image_transforms(item[key]) result.append(item) return result + def close(self): + first_error = None + locator = getattr(self, "_frame_locator", None) + if locator is not None: + try: + locator.close() + except Exception as error: + first_error = error + for collator in getattr(self, "_video_collators", ()): + try: + collator.close() + except Exception as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error + + def __del__(self): + try: + self.close() + except Exception: + pass + def _read_rows( self, indices, projection, splits=None, needs_filter=True): if not indices: @@ -1086,6 +1137,15 @@ def _resolve_image_blobs( row[key] = body +def _decode_image_rows(row_groups, image_keys, features, return_uint8): + for rows in row_groups: + for row in rows.values(): + for key in image_keys: + if key in row: + row[key] = _image_tensor( + row[key], features[key], return_uint8=return_uint8) + + def _image_blob_sources(row_groups, image_keys): return [ (row, key, row[key]) @@ -1118,9 +1178,12 @@ def _torch_row(row, features, return_uint8=False): if key not in result: continue value = result[key] - if feature.get("dtype") == "image": + if feature.get("dtype") == "image" and not torch.is_tensor(value): result[key] = _image_tensor( value, feature, return_uint8=return_uint8) + elif feature.get("dtype") == "video": + result[key] = _video_tensor( + value, feature, return_uint8=return_uint8) elif feature.get("dtype") != "string" and not torch.is_tensor(value): dtype = getattr(torch, _TORCH_DTYPE_NAMES[feature.get("dtype")]) result[key] = torch.tensor(value, dtype=dtype) @@ -1164,6 +1227,111 @@ def _image_tensor(payload, feature, return_uint8=False): return tensor.div_(255) if normalize else tensor +def _video_tensor(frame, feature, return_uint8=False): + import torch + + if not torch.is_tensor(frame): + raise ValueError("LeRobot video decoder must return a Torch tensor.") + expected_shape = _feature_shape(feature, "video") + if len(expected_shape) != 3: + raise ValueError("LeRobot video feature must have three dimensions.") + names = feature.get("names") or [] + output_shape = expected_shape if names and names[0] in ( + "channel", "channels" + ) else expected_shape[2:] + expected_shape[:2] + if tuple(frame.shape) != output_shape: + raise ValueError( + "LeRobot video frame has shape %s, expected %s." + % (tuple(frame.shape), output_shape) + ) + if frame.dtype == torch.uint8 and not return_uint8: + return frame.float().div_(255) + return frame + + +def _open_video_decoder(stream, backend=None): + if backend in (None, "torchcodec"): + try: + return _open_torchcodec_decoder(stream) + except (ImportError, RuntimeError): + if backend == "torchcodec": + raise + stream.seek(0) + return _PyAVVideoDecoder(stream) + + +def _open_torchcodec_decoder(stream): + try: + from torchcodec.decoders import VideoDecoder + except (ImportError, RuntimeError) as error: + raise ImportError( + "Video-backed PaimonLeRobotDataset requires TorchCodec from " + "'pypaimon[lerobot]'." + ) from error + try: + return VideoDecoder(stream, seek_mode="exact") + except TypeError: + # TorchCodec 0.2 accepts bytes but not seekable file-like objects. + stream.seek(0) + return VideoDecoder(stream.read(), seek_mode="exact") + + +class _PyAVVideoDecoder: + + def __init__(self, stream): + try: + import av + except ImportError as error: + raise ImportError( + "Video-backed PaimonLeRobotDataset requires PyAV from " + "'pypaimon[lerobot]'." + ) from error + self._container = av.open(stream) + self._next_index = 0 + self._frames = iter(self._container.decode(video=0)) + + def __getitem__(self, index): + if index < self._next_index: + self._container.seek(0) + self._next_index = 0 + self._frames = iter(self._container.decode(video=0)) + try: + while self._next_index <= index: + frame = next(self._frames) + self._next_index += 1 + except StopIteration as error: + raise IndexError( + "Video frame index %d is out of range." % index + ) from error + + import numpy as np + import torch + array = np.array(frame.to_ndarray(format="rgb24"), copy=True) + return torch.from_numpy(array).permute(2, 0, 1) + + def close(self): + self._container.close() + + +def _decode_video_frame(decoder, frame_index, unused_row): + return decoder[frame_index] + + +def _identity(values): + return values + + +def _decode_video_rows(row_groups, collators): + for collator in collators: + for rows in row_groups: + indices = list(rows) + if not indices: + continue + decoded = collator([rows[index] for index in indices]) + for index, row in zip(indices, decoded): + rows[index] = row + + def _normalize_index(index, size): index = operator.index(index) if index < 0: diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 82e2b7631c1a..e3ca808ea2ee 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1915,6 +1915,45 @@ def close(self): [0.5, 0.6, 0.1, 0.2, 0.3], atol=1e-6, ) + + dataset = pmm.PaimonLeRobotDataset( + table, + delta_timestamps={"camera": [0.0, 0.1]}, + max_open_videos=1, + ) + try: + last, first = dataset.__getitems__([4, 0]) + self.assertEqual( + [2, 3, 16, 16], list(last["camera"].shape)) + self.assertEqual( + [2, 3, 16, 16], list(first["camera"].shape)) + self.assertEqual("torch.float32", str(last["camera"].dtype)) + np.testing.assert_allclose( + [ + float(last["camera"][0].mean()) * 255, + float(first["camera"][0].mean()) * 255, + float(first["camera"][1].mean()) * 255, + ], + [120, 168, 216], + atol=5, + ) + self.assertEqual( + [False, True], last["camera_is_pad"].tolist()) + self.assertEqual( + 1, len(dataset._video_collators[0]._decoders)) + + from torch.utils.data import DataLoader + worker_indices = [] + for batch in DataLoader( + dataset, + batch_size=2, + shuffle=False, + num_workers=2, + multiprocessing_context="spawn"): + worker_indices.extend(batch["index"].tolist()) + self.assertEqual(list(range(5)), worker_indices) + finally: + dataset.close() finally: shutil.rmtree(temp_dir, ignore_errors=True) From dc788bd014e0b837ba6b640b9c6c689306190207 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 01:28:53 -0700 Subject: [PATCH 02/10] [docs] Tighten LeRobot video read documentation --- docs/docs/pypaimon/lerobot.md | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index e226f64048ac..6aedfb27c392 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -109,9 +109,8 @@ Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested Video features map to `BLOB`. Frame rows reference MP4 payloads copied once per aligned file group. Video imports use the video grouping policy and check -rolling before each Episode. They require a bucket-unaware table. Read them -with a Paimon scan and `VideoFrameCollator`, or let `PaimonLeRobotDataset` -decode the referenced frames with TorchCodec. +rolling before each Episode. They require a bucket-unaware table. Use +`VideoFrameCollator` for scans or `PaimonLeRobotDataset` for training. ## Capture LeRobot frames directly into Paimon @@ -208,9 +207,8 @@ pin one named snapshot on every component. ## Train with Paimon LeRobot data -For map-style training, read a tagged image- or video-backed table group -created by `load_from_lerobot` directly from Paimon. `PaimonLeRobotDataset` -requires the complete table group; a frame-only table created by +For map-style training, pass an image- or video-backed table group created by +`load_from_lerobot` to `PaimonLeRobotDataset`. A frame-only table created by `PaimonLeRobotWriter` is not sufficient. ```python @@ -224,9 +222,7 @@ dataset = PaimonLeRobotDataset( loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) ``` -If `tag_name` is omitted, the latest snapshots are used. Metadata is available -through `dataset.meta`. Frame lookups use the BTree on `index`; payload columns -remain lazy. Video decoder sessions are cached per DataLoader worker; set -`max_open_videos` to bound the number retained for each video feature. Video -decoding uses TorchCodec when available and otherwise PyAV; set -`video_backend="torchcodec"` or `"pyav"` to select one explicitly. +Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree +on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back +to PyAV, and caches at most `max_open_videos` decoders per worker and feature. +Set `video_backend` to force either decoder. From 23e4306c5637c6d0cf3b1fac27e9ad149b1d38a7 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 01:34:40 -0700 Subject: [PATCH 03/10] [python] Hide LeRobot decoder cache tuning --- docs/docs/pypaimon/lerobot.md | 4 ++-- paimon-python/pypaimon/multimodal/lerobot/dataset.py | 4 ---- paimon-python/pypaimon/tests/multimodal_lerobot_test.py | 1 - 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 6aedfb27c392..b1e487fff1f1 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -224,5 +224,5 @@ loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back -to PyAV, and caches at most `max_open_videos` decoders per worker and feature. -Set `video_backend` to force either decoder. +to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force +either decoder. diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index b8f087c69a51..4b3058486553 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -96,7 +96,6 @@ def __init__( tolerance_s=1e-4, blob_parallelism=16, video_backend=None, - max_open_videos=8, return_uint8=False): if sys.version_info < (3, 10): raise RuntimeError( @@ -117,8 +116,6 @@ def __init__( raise ValueError( "video_backend must be None, 'torchcodec', or 'pyav'.") self.video_backend = video_backend - self.max_open_videos = _positive_int( - max_open_videos, "max_open_videos") if not isinstance(return_uint8, bool): raise TypeError("return_uint8 must be a boolean.") self.return_uint8 = return_uint8 @@ -244,7 +241,6 @@ def _init_reader(self, raw_table, info): _open_video_decoder, backend=self.video_backend), decode_fn=_decode_video_frame, output_column=key, - max_open_videos=self.max_open_videos, collate_fn=_identity, ) for key in self._video_keys diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index e3ca808ea2ee..287d2586a68a 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -1919,7 +1919,6 @@ def close(self): dataset = pmm.PaimonLeRobotDataset( table, delta_timestamps={"camera": [0.0, 0.1]}, - max_open_videos=1, ) try: last, first = dataset.__getitems__([4, 0]) From 81e88303132de7da1b5f18da7856c5147a7b8a3c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 01:38:56 -0700 Subject: [PATCH 04/10] [python] Fall back from unavailable TorchCodec --- .../pypaimon/multimodal/lerobot/dataset.py | 2 +- .../pypaimon/tests/multimodal_lerobot_test.py | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 4b3058486553..2ce1351bb79f 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1249,7 +1249,7 @@ def _open_video_decoder(stream, backend=None): if backend in (None, "torchcodec"): try: return _open_torchcodec_decoder(stream) - except (ImportError, RuntimeError): + except (ImportError, OSError, RuntimeError): if backend == "torchcodec": raise stream.seek(0) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 287d2586a68a..147f58c6a6db 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -47,6 +47,7 @@ _arrow_rows, _image_tensor, _index_names, + _open_video_decoder, _selected_episodes, _torch_row, ) @@ -126,6 +127,26 @@ def _catalog_metadata(connection, name): class LeRobotValidationTest(unittest.TestCase): + def test_default_video_backend_falls_back_on_os_error(self): + stream = Mock() + decoder = object() + with patch( + "pypaimon.multimodal.lerobot.dataset." + "_open_torchcodec_decoder", + side_effect=OSError("unavailable")), patch( + "pypaimon.multimodal.lerobot.dataset._PyAVVideoDecoder", + return_value=decoder) as pyav: + self.assertIs(decoder, _open_video_decoder(stream)) + stream.seek.assert_called_once_with(0) + pyav.assert_called_once_with(stream) + + with patch( + "pypaimon.multimodal.lerobot.dataset." + "_open_torchcodec_decoder", + side_effect=OSError("unavailable")): + with self.assertRaises(OSError): + _open_video_decoder(stream, backend="torchcodec") + def test_dataset_requires_supported_python(self): with patch( "pypaimon.multimodal.lerobot.dataset.sys.version_info", From 80cc11059c9032fa13e4fc199b0edc6f03a9cde1 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 01:55:44 -0700 Subject: [PATCH 05/10] [python] Fix partial video windows and PyAV seeks --- .../pypaimon/multimodal/lerobot/dataset.py | 75 +++++++++++++-- .../pypaimon/tests/multimodal_lerobot_test.py | 96 +++++++++++++++++++ 2 files changed, 163 insertions(+), 8 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 2ce1351bb79f..5823ecec4666 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -24,6 +24,7 @@ import operator import os import sys +from collections import OrderedDict from functools import partial import pyarrow as pa @@ -1274,6 +1275,9 @@ def _open_torchcodec_decoder(stream): class _PyAVVideoDecoder: + # Reuse common overlapping delta windows without retaining a whole video. + _FRAME_CACHE_SIZE = 8 + def __init__(self, stream): try: import av @@ -1283,23 +1287,75 @@ def __init__(self, stream): "'pypaimon[lerobot]'." ) from error self._container = av.open(stream) + self._stream = self._container.streams.video[0] self._next_index = 0 - self._frames = iter(self._container.decode(video=0)) + self._timestamps = [] + self._cache = OrderedDict() + self._frames = iter(self._container.decode(self._stream)) def __getitem__(self, index): - if index < self._next_index: - self._container.seek(0) - self._next_index = 0 - self._frames = iter(self._container.decode(video=0)) + index = operator.index(index) + if index < 0: + raise IndexError("Video frame index %d is out of range." % index) + frame = self._cache.pop(index, None) + if frame is not None: + self._cache[index] = frame + return self._tensor(frame) + + at_frontier = self._next_index == len(self._timestamps) + if index != self._next_index and not ( + at_frontier and index >= self._next_index): + self._seek(index) try: - while self._next_index <= index: + while True: frame = next(self._frames) - self._next_index += 1 + if frame.pts is None: + continue + timestamp = frame.pts * ( + frame.time_base or self._stream.time_base) + position = bisect.bisect_left(self._timestamps, timestamp) + if ( + position < len(self._timestamps) + and self._timestamps[position] == timestamp + ): + frame_index = position + elif self._next_index == len(self._timestamps): + frame_index = self._next_index + self._timestamps.append(timestamp) + else: + continue + self._next_index = frame_index + 1 + self._remember(frame_index, frame) + if frame_index == index: + return self._tensor(frame) + if frame_index > index: + break except StopIteration as error: raise IndexError( "Video frame index %d is out of range." % index ) from error + raise IndexError("Video frame index %d is out of range." % index) + + def _seek(self, index): + anchor = min(index, len(self._timestamps) - 1) + timestamp = self._timestamps[anchor] + self._container.seek( + round(timestamp / self._stream.time_base), + backward=True, + any_frame=False, + stream=self._stream, + ) + self._next_index = None + self._frames = iter(self._container.decode(self._stream)) + def _remember(self, index, frame): + self._cache.pop(index, None) + self._cache[index] = frame + if len(self._cache) > self._FRAME_CACHE_SIZE: + self._cache.popitem(last=False) + + @staticmethod + def _tensor(frame): import numpy as np import torch array = np.array(frame.to_ndarray(format="rgb24"), copy=True) @@ -1320,7 +1376,10 @@ def _identity(values): def _decode_video_rows(row_groups, collators): for collator in collators: for rows in row_groups: - indices = list(rows) + 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]) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 147f58c6a6db..f5ddfeb5447b 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -44,6 +44,7 @@ from pypaimon.multimodal.connection import MultimodalConnection from pypaimon.multimodal.lerobot import load_from_lerobot from pypaimon.multimodal.lerobot.dataset import ( + _PyAVVideoDecoder, _arrow_rows, _image_tensor, _index_names, @@ -127,6 +128,63 @@ def _catalog_metadata(connection, name): class LeRobotValidationTest(unittest.TestCase): + @unittest.skipIf(av is None, "PyAV is not installed") + def test_pyav_decoder_reuses_windows_and_seeks_known_frames(self): + class Frame: + + time_base = Fraction(1, 10) + + def __init__(self, pts): + self.pts = pts + + def to_ndarray(self, format): + assert format == "rgb24" + return np.full((2, 2, 3), self.pts, dtype=np.uint8) + + class Container: + + def __init__(self): + self.stream = SimpleNamespace(time_base=Fraction(1, 10)) + self.streams = SimpleNamespace(video=[self.stream]) + self.position = 0 + self.decoded = 0 + self.seeks = [] + + def decode(self, stream): + assert stream is self.stream + while self.position < 120: + index = self.position + self.position += 1 + self.decoded += 1 + yield Frame(index) + + def seek(self, offset, *, backward, any_frame, stream): + assert backward + assert not any_frame + assert stream is self.stream + self.seeks.append(offset) + self.position = offset // 10 * 10 + + def close(self): + pass + + container = Container() + with patch("av.open", return_value=container): + decoder = _PyAVVideoDecoder(io.BytesIO()) + try: + for index in range(119): + decoder[index] + decoder[index + 1] + self.assertEqual(120, container.decoded) + self.assertEqual([], container.seeks) + + decoder[5] + decoder[90] + self.assertEqual(127, container.decoded) + self.assertEqual([5, 90], container.seeks) + finally: + decoder.close() + def test_default_video_backend_falls_back_on_os_error(self): stream = Mock() decoder = object() @@ -1776,11 +1834,17 @@ def test_imported_video_payload_can_be_decoded(self): "fps": 10.0, }, "task_index": {"dtype": "int64", "shape": [1]}, + "action": {"dtype": "float32", "shape": [1]}, "camera": { "dtype": "video", "shape": [16, 16, 3], "video_info": {"video.fps": 10.0}, }, + "camera_b": { + "dtype": "video", + "shape": [16, 16, 3], + "video_info": {"video.fps": 10.0}, + }, }, } episodes = [ @@ -1796,6 +1860,10 @@ def test_imported_video_payload_can_be_decoded(self): "videos/camera/file_index": 0, "videos/camera/from_timestamp": 0.5, "videos/camera/to_timestamp": 0.7, + "videos/camera_b/chunk_index": 0, + "videos/camera_b/file_index": 0, + "videos/camera_b/from_timestamp": 0.5, + "videos/camera_b/to_timestamp": 0.7, }, { "episode_index": 1, @@ -1809,6 +1877,10 @@ def test_imported_video_payload_can_be_decoded(self): "videos/camera/file_index": 0, "videos/camera/from_timestamp": 0.1, "videos/camera/to_timestamp": 0.4, + "videos/camera_b/chunk_index": 0, + "videos/camera_b/file_index": 0, + "videos/camera_b/from_timestamp": 0.1, + "videos/camera_b/to_timestamp": 0.4, }, ] physical_frame_values = [24, 56, 88, 120, 168, 216] @@ -1845,6 +1917,10 @@ def test_imported_video_payload_can_be_decoded(self): container.mux(packet) for packet in stream.encode(): container.mux(packet) + camera_b_path = ( + temp_dir / "videos/camera_b/chunk-000/file-000.mp4") + camera_b_path.parent.mkdir(parents=True) + shutil.copy2(video_path, camera_b_path) class Dataset: @@ -1862,6 +1938,8 @@ class Dataset: type=pa.float32(), ), "task_index": pa.array([0] * 5, type=pa.int64()), + "action": pa.array( + [0.0, 1.0, 2.0, 3.0, 4.0], type=pa.float32()), }) def __len__(self): @@ -1947,6 +2025,8 @@ def close(self): [2, 3, 16, 16], list(last["camera"].shape)) self.assertEqual( [2, 3, 16, 16], list(first["camera"].shape)) + self.assertEqual( + [3, 16, 16], list(first["camera_b"].shape)) self.assertEqual("torch.float32", str(last["camera"].dtype)) np.testing.assert_allclose( [ @@ -1974,6 +2054,22 @@ def close(self): self.assertEqual(list(range(5)), worker_indices) finally: dataset.close() + + action_dataset = pmm.PaimonLeRobotDataset( + table, + delta_timestamps={"action": [0.0, 0.1]}, + ) + try: + item = action_dataset[0] + self.assertEqual([2], list(item["action"].shape)) + np.testing.assert_allclose( + [0.0, 1.0], item["action"].tolist()) + self.assertEqual( + [3, 16, 16], list(item["camera"].shape)) + self.assertEqual( + [3, 16, 16], list(item["camera_b"].shape)) + finally: + action_dataset.close() finally: shutil.rmtree(temp_dir, ignore_errors=True) From e627d6a713f219e685a26083e3f70c891c468a22 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 02:09:32 -0700 Subject: [PATCH 06/10] [python] Seek PyAV before the target frame --- paimon-python/pypaimon/multimodal/lerobot/dataset.py | 2 +- paimon-python/pypaimon/tests/multimodal_lerobot_test.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 5823ecec4666..58df72f5e6ff 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1337,7 +1337,7 @@ def __getitem__(self, index): raise IndexError("Video frame index %d is out of range." % index) def _seek(self, index): - anchor = min(index, len(self._timestamps) - 1) + anchor = max(0, min(index, len(self._timestamps) - 1) - 1) timestamp = self._timestamps[anchor] self._container.seek( round(timestamp / self._stream.time_base), diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index f5ddfeb5447b..87d55f171e63 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -163,7 +163,8 @@ def seek(self, offset, *, backward, any_frame, stream): assert not any_frame assert stream is self.stream self.seeks.append(offset) - self.position = offset // 10 * 10 + # Model a B-frame seek that starts after an exact target PTS. + self.position = 3 if offset == 1 else offset // 10 * 10 def close(self): pass @@ -180,8 +181,9 @@ def close(self): decoder[5] decoder[90] - self.assertEqual(127, container.decoded) - self.assertEqual([5, 90], container.seeks) + decoder[1] + self.assertEqual(139, container.decoded) + self.assertEqual([4, 89, 0], container.seeks) finally: decoder.close() From 4183e4a637631f6092c068154708f4a97024fed8 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 17:22:55 +0800 Subject: [PATCH 07/10] [python] Seek PyAV from known keyframes --- .../pypaimon/multimodal/lerobot/dataset.py | 6 ++- .../pypaimon/tests/multimodal_lerobot_test.py | 46 +++++++++++++++++-- 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 58df72f5e6ff..90bb1abece2f 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1290,6 +1290,7 @@ def __init__(self, stream): self._stream = self._container.streams.video[0] self._next_index = 0 self._timestamps = [] + self._keyframes = [] self._cache = OrderedDict() self._frames = iter(self._container.decode(self._stream)) @@ -1322,6 +1323,8 @@ def __getitem__(self, index): elif self._next_index == len(self._timestamps): frame_index = self._next_index self._timestamps.append(timestamp) + if frame.key_frame: + self._keyframes.append(frame_index) else: continue self._next_index = frame_index + 1 @@ -1337,7 +1340,8 @@ def __getitem__(self, index): raise IndexError("Video frame index %d is out of range." % index) def _seek(self, index): - anchor = max(0, min(index, len(self._timestamps) - 1) - 1) + position = bisect.bisect_right(self._keyframes, index) + anchor = self._keyframes[position - 1] if position else 0 timestamp = self._timestamps[anchor] self._container.seek( round(timestamp / self._stream.time_base), diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 87d55f171e63..17a4519d5b25 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -136,6 +136,7 @@ class Frame: def __init__(self, pts): self.pts = pts + self.key_frame = pts % 10 == 0 def to_ndarray(self, format): assert format == "rgb24" @@ -163,8 +164,7 @@ def seek(self, offset, *, backward, any_frame, stream): assert not any_frame assert stream is self.stream self.seeks.append(offset) - # Model a B-frame seek that starts after an exact target PTS. - self.position = 3 if offset == 1 else offset // 10 * 10 + self.position = offset def close(self): pass @@ -181,9 +181,45 @@ def close(self): decoder[5] decoder[90] - decoder[1] - self.assertEqual(139, container.decoded) - self.assertEqual([4, 89, 0], container.seeks) + decoder[8] + self.assertEqual(136, container.decoded) + self.assertEqual([0, 90, 0], container.seeks) + finally: + decoder.close() + + @unittest.skipIf(av is None, "PyAV is not installed") + def test_pyav_decoder_seeks_before_b_frames(self): + output = io.BytesIO() + with av.open(output, mode="w", format="mp4") as container: + stream = container.add_stream("mpeg4", rate=30) + stream.width = 16 + stream.height = 16 + stream.pix_fmt = "yuv420p" + stream.gop_size = 12 + stream.codec_context.max_b_frames = 2 + for index in range(70): + image = np.full( + (16, 16, 3), index + 24, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(image, format="rgb24") + frame.pts = index + frame.time_base = Fraction(1, 30) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + + payload = output.getvalue() + with av.open(io.BytesIO(payload)) as container: + expected = [ + np.array(frame.to_ndarray(format="rgb24"), copy=True) + for frame in container.decode(video=0) + ] + + decoder = _PyAVVideoDecoder(io.BytesIO(payload)) + try: + for index in (69, 20, 35, 1, 68): + actual = decoder[index].permute(1, 2, 0).numpy() + np.testing.assert_array_equal(expected[index], actual) finally: decoder.close() From 1cc00e046ed98aa853544a24e623d121aedfd733 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 17:54:01 +0800 Subject: [PATCH 08/10] [python] Gate video training test dependencies --- .../pypaimon/tests/multimodal_lerobot_test.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 17a4519d5b25..c4282f99bd05 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -17,6 +17,7 @@ import builtins from array import array from fractions import Fraction +import importlib.util import io import json import pickle @@ -1844,6 +1845,19 @@ def sample_timestamps(unused_dataset, uri): @unittest.skipUnless(av is not None, "PyAV is required for MP4 decoding") def test_imported_video_payload_can_be_decoded(self): + self._assert_imported_video_payload_can_be_decoded(False) + + @unittest.skipUnless( + av is not None + and sys.version_info >= (3, 10) + and importlib.util.find_spec("datasets") is not None + and importlib.util.find_spec("torch") is not None, + "Video training reads require Python 3.10+, PyAV, datasets, and Torch", + ) + def test_imported_video_payload_supports_training_reads(self): + self._assert_imported_video_payload_can_be_decoded(True) + + def _assert_imported_video_payload_can_be_decoded(self, training_reads): import pandas as pd temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_mp4_")) @@ -2052,6 +2066,8 @@ def close(self): [0.5, 0.6, 0.1, 0.2, 0.3], atol=1e-6, ) + if not training_reads: + return dataset = pmm.PaimonLeRobotDataset( table, From 1858232317b261f65a1a60c917f41cfca4222e3c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 12 Sep 2026 21:47:09 +0800 Subject: [PATCH 09/10] [python] Gate PyAV decoder tests on Torch --- .../pypaimon/tests/multimodal_lerobot_test.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index c4282f99bd05..17bddc553dfa 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -129,7 +129,10 @@ def _catalog_metadata(connection, name): class LeRobotValidationTest(unittest.TestCase): - @unittest.skipIf(av is None, "PyAV is not installed") + @unittest.skipUnless( + av is not None and importlib.util.find_spec("torch") is not None, + "PyAV and Torch are required for video decoding", + ) def test_pyav_decoder_reuses_windows_and_seeks_known_frames(self): class Frame: @@ -188,7 +191,10 @@ def close(self): finally: decoder.close() - @unittest.skipIf(av is None, "PyAV is not installed") + @unittest.skipUnless( + av is not None and importlib.util.find_spec("torch") is not None, + "PyAV and Torch are required for video decoding", + ) def test_pyav_decoder_seeks_before_b_frames(self): output = io.BytesIO() with av.open(output, mode="w", format="mp4") as container: From 4cfabc7b340cd0d2f98f91f054e0a71e859e65c3 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sun, 13 Sep 2026 14:19:41 +0800 Subject: [PATCH 10/10] python: index PyAV packets for cold random reads --- .../pypaimon/multimodal/lerobot/dataset.py | 30 ++++++++- .../pypaimon/tests/multimodal_lerobot_test.py | 66 +++++++++++++++++++ 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 90bb1abece2f..7a13731e5fcd 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1303,8 +1303,14 @@ def __getitem__(self, index): self._cache[index] = frame return self._tensor(frame) + indexed = ( + index > 0 and not self._timestamps + and self._index_packets() + ) at_frontier = self._next_index == len(self._timestamps) - if index != self._next_index and not ( + if indexed: + self._seek(index) + elif index != self._next_index and not ( at_frontier and index >= self._next_index): self._seek(index) try: @@ -1339,6 +1345,28 @@ def __getitem__(self, index): ) from error raise IndexError("Video frame index %d is out of range." % index) + def _index_packets(self): + entries = [] + for packet in self._container.demux(self._stream): + if (packet.pts is None + or getattr(packet, "is_discard", False)): + continue + timestamp = packet.pts * ( + packet.time_base or self._stream.time_base) + entries.append((timestamp, packet.is_keyframe)) + if not entries: + self._container.seek( + 0, backward=True, any_frame=False, stream=self._stream) + self._frames = iter(self._container.decode(self._stream)) + return False + entries.sort(key=lambda entry: entry[0]) + self._timestamps = [timestamp for timestamp, unused in entries] + self._keyframes = [ + index for index, (unused, keyframe) in enumerate(entries) + if keyframe + ] + return True + def _seek(self, index): position = bisect.bisect_right(self._keyframes, index) anchor = self._keyframes[position - 1] if position else 0 diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 17bddc553dfa..97b180d39e65 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -191,6 +191,72 @@ def close(self): finally: decoder.close() + def test_pyav_decoder_indexes_cold_random_reads(self): + class Frame: + + time_base = Fraction(1, 10) + + def __init__(self, pts): + self.pts = pts + self.key_frame = pts % 10 == 0 + + def to_ndarray(self, format): + assert format == "rgb24" + return np.full((2, 2, 3), self.pts, dtype=np.uint8) + + class Container: + + def __init__(self): + self.stream = SimpleNamespace(time_base=Fraction(1, 10)) + self.streams = SimpleNamespace(video=[self.stream]) + self.position = 0 + self.decoded = 0 + self.demuxed = 0 + self.seeks = [] + + def demux(self, stream): + assert stream is self.stream + for index in range(120): + self.demuxed += 1 + yield SimpleNamespace( + pts=index, time_base=Fraction(1, 10), + is_discard=False, is_keyframe=index % 10 == 0, + ) + + def decode(self, stream): + assert stream is self.stream + while self.position < 120: + index = self.position + self.position += 1 + self.decoded += 1 + yield Frame(index) + + def seek(self, offset, *, backward, any_frame, stream): + assert backward + assert not any_frame + assert stream is self.stream + self.seeks.append(offset) + self.position = offset + + def close(self): + pass + + container = Container() + fake_av = SimpleNamespace(open=lambda unused_stream: container) + tensor = staticmethod( + lambda frame: frame.to_ndarray(format="rgb24")) + with patch.dict(sys.modules, {"av": fake_av}), patch.object( + _PyAVVideoDecoder, "_tensor", tensor): + decoder = _PyAVVideoDecoder(io.BytesIO()) + try: + frame = decoder[95] + self.assertEqual(120, container.demuxed) + self.assertEqual([90], container.seeks) + self.assertEqual(6, container.decoded) + self.assertTrue((frame == 95).all()) + finally: + decoder.close() + @unittest.skipUnless( av is not None and importlib.util.find_spec("torch") is not None, "PyAV and Torch are required for video decoding",