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 222864e164 [python] Avoid duplicate LeRobot video batch copies (#10405)
222864e164 is described below
commit 222864e16489b787e7ae80eea44f1d1ba1070d02
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Oct 7 18:18:28 2026 +0800
[python] Avoid duplicate LeRobot video batch copies (#10405)
---
docs/docs/pypaimon/lerobot.md | 10 +-
.../pypaimon/multimodal/lerobot/dataset.py | 120 ++++++++++++++++++++-
.../pypaimon/tests/multimodal_lerobot_test.py | 77 +++++++++++++
3 files changed, 200 insertions(+), 7 deletions(-)
diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md
index e97ee770e7..518a5d0674 100644
--- a/docs/docs/pypaimon/lerobot.md
+++ b/docs/docs/pypaimon/lerobot.md
@@ -211,16 +211,22 @@ For map-style training, pass an image- or video-backed
table group created by
`PaimonLeRobotWriter` is not sufficient.
```python
-from torch.utils.data import DataLoader
from pypaimon.multimodal import PaimonLeRobotDataset
dataset = PaimonLeRobotDataset(
conn.get_table("robot_data"),
tag_name=tag,
)
-loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)
+loader = dataset.to_dataloader(
+ batch_size=32,
+ shuffle=True,
+ num_workers=4,
+)
```
+`to_dataloader` writes decoded video windows directly into shared worker
+batches, avoiding a second full-resolution copy during collation.
+
Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree
on `index`; payloads remain lazy. Indexed videos prefetch metadata and target
GOPs, then fetch
uncached PyAV reads on demand. Unindexed videos use the TorchCodec/PyAV scan
path. Set
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index f48ce7b446..665ceb1478 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -364,6 +364,13 @@ class PaimonDatasetReader(ABC):
def get_items(self, indices):
"""Return fully assembled frames for one batch."""
+ return self._get_items(indices, return_batch=False, share_memory=False)
+
+ def _get_batch(self, indices, share_memory):
+ return self._get_items(
+ indices, return_batch=True, share_memory=share_memory)
+
+ def _get_items(self, indices, return_batch, share_memory):
dataset_indices = [
_normalize_index(index, len(self)) for index in indices
]
@@ -432,7 +439,8 @@ class PaimonDatasetReader(ABC):
# Transforms may shrink frames; assemble those windows one sample at a
time.
video_windows = {} if self.image_transforms is not None else
_decode_video_windows(
plans, rows, getattr(self, "_video_collators", ()),
- self._features, self.return_uint8)
+ self._features, self.return_uint8,
+ return_batch=return_batch, share_memory=share_memory)
for group in row_groups:
for row in group.values():
for key in video_windows:
@@ -477,6 +485,8 @@ class PaimonDatasetReader(ABC):
for key in self._visual_keys:
item[key] = self.image_transforms(item[key])
result.append(item)
+ if return_batch:
+ return _collate_lerobot_batch(result, video_windows)
return result
def __getitem__(self, index):
@@ -663,6 +673,18 @@ class PaimonLeRobotDataset:
def __getitems__(self, indices):
return self.reader.get_items(indices)
+ def to_dataloader(self, **kwargs):
+ """Create a DataLoader that assembles video windows once per batch."""
+ collate_fn = kwargs.pop("collate_fn", None)
+ if collate_fn is not None:
+ raise ValueError(
+ "PaimonLeRobotDataset.to_dataloader manages collate_fn.")
+ from torch.utils.data import DataLoader
+ num_workers = kwargs.get("num_workers", 0)
+ batched = _LeRobotDataLoaderAdapter(
+ self, share_memory=bool(num_workers))
+ return DataLoader(batched, collate_fn=_identity, **kwargs)
+
@property
def return_uint8(self):
return self.reader.return_uint8
@@ -703,6 +725,46 @@ class PaimonLeRobotDataset:
self.reader.__class__.__name__, self.__class__.__name__, 1)
+class _LeRobotDataLoaderAdapter:
+
+ def __init__(self, dataset, share_memory):
+ self.dataset = dataset
+ self.share_memory = share_memory
+
+ def __len__(self):
+ return len(self.dataset)
+
+ def __getitem__(self, index):
+ return self.dataset[index]
+
+ def __getitems__(self, indices):
+ return self.dataset.reader._get_batch(indices, self.share_memory)
+
+ @property
+ def return_uint8(self):
+ return self.dataset.return_uint8
+
+ @return_uint8.setter
+ def return_uint8(self, value):
+ self.dataset.return_uint8 = value
+
+ @property
+ def image_transforms(self):
+ return self.dataset.image_transforms
+
+ @image_transforms.setter
+ def image_transforms(self, value):
+ self.dataset.image_transforms = value
+
+ def __getattr__(self, name):
+ if name.startswith("__") and name.endswith("__"):
+ raise AttributeError(name)
+ dataset = self.__dict__.get("dataset")
+ if dataset is None:
+ raise AttributeError(name)
+ return getattr(dataset, name)
+
+
class _PaimonLeRobotMetadata:
def __init__(
@@ -1347,7 +1409,9 @@ def _attach_task_labels(rows, task_names, subtask_names):
row["subtask"] = subtask_names[subtask_index]
-def _decode_video_windows(plans, rows, collators, features, return_uint8):
+def _decode_video_windows(
+ plans, rows, collators, features, return_uint8, *,
+ return_batch=False, share_memory=False):
import torch
tasks = [c for c in collators if c.video_column in plans[0]["windows"]]
@@ -1381,7 +1445,7 @@ def _decode_video_windows(plans, rows, collators,
features, return_uint8):
requests.setdefault((payload, keyframe_index), []).append((
offset, [d.frame_index for d in window]))
- output = [None] * len(plans)
+ output = None if return_batch else [None] * len(plans)
with torch.inference_mode(inference_enabled), \
torch.set_grad_enabled(grad_enabled):
for (payload, keyframe_index), windows in requests.items():
@@ -1399,17 +1463,49 @@ def _decode_video_windows(plans, rows, collators,
features, return_uint8):
if sum(len(window) for _, window in windows) >= 2 *
len(indices):
frames = frames.contiguous()
positions = {index: pos for pos, index in enumerate(indices)}
+ if return_batch:
+ window_sizes = {len(window) for _, window in windows}
+ if len(window_sizes) != 1:
+ raise ValueError(
+ "Video batch windows must have the same size.")
+ window_size = next(iter(window_sizes))
+ output_dtype = torch.float32 \
+ if frames.dtype == torch.uint8 and not return_uint8 \
+ else frames.dtype
+ shape = ((len(plans), window_size)
+ + tuple(frames.shape[1:]))
+ if output is None:
+ output = torch.empty(
+ shape, dtype=output_dtype, device=frames.device)
+ if share_memory:
+ output.share_memory_()
+ elif tuple(output.shape) != shape \
+ or output.dtype != output_dtype:
+ raise ValueError(
+ "Video batch payloads have incompatible frames.")
for offset, window in windows:
selection = [positions[index] for index in window]
start = selection[0]
- if selection == list(range(start, start + len(selection))):
+ contiguous = selection == list(
+ range(start, start + len(selection)))
+ if return_batch:
+ source = frames[start:start + len(selection)] \
+ if contiguous else frames.index_select(
+ 0, torch.tensor(
+ selection, dtype=torch.long,
+ device=frames.device))
+ output[offset].copy_(source)
+ elif contiguous:
output[offset] = frames[start:start +
len(selection)].clone(
memory_format=torch.contiguous_format)
else:
output[offset] = frames.index_select(0, torch.tensor(
selection, dtype=torch.long, device=frames.device))
if frames.dtype == torch.uint8 and not return_uint8:
- output[offset] = output[offset].float().div_(255)
+ if return_batch:
+ output[offset].div_(255)
+ else:
+ output[offset] = output[offset].float().div_(255)
return key, output
if len(tasks) == 1:
@@ -1421,6 +1517,20 @@ def _decode_video_windows(plans, rows, collators,
features, return_uint8):
return {key: windows for key, windows in decoded if windows is not None}
+def _collate_lerobot_batch(rows, video_windows):
+ from torch.utils.data import default_collate
+
+ if not video_windows:
+ return default_collate(rows)
+ remaining = [{
+ key: value for key, value in row.items()
+ if key not in video_windows
+ } for row in rows]
+ batch = default_collate(remaining)
+ batch.update(video_windows)
+ return batch
+
+
def _stack_visual_windows(plans, rows, visual_keys):
import torch
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index 1f5da1f922..f39d7e1add 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -150,6 +150,17 @@ def _catalog_metadata(connection, name):
}
+def _increment_image(value):
+ return value + 1
+
+
+def _configure_uint8_transform_worker(unused_worker_id):
+ from torch.utils.data import get_worker_info
+ dataset = get_worker_info().dataset
+ dataset.return_uint8 = True
+ dataset.image_transforms = _increment_image
+
+
class LeRobotValidationTest(unittest.TestCase):
def test_episode_metadata_pickle_stays_small_and_usable(self):
@@ -930,6 +941,12 @@ class LeRobotValidationTest(unittest.TestCase):
sample["action"], torch.tensor([0.0, 1.0, 2.0]))
self.assertEqual([False, False, False],
sample["action_is_pad"].tolist())
+ batch = next(iter(dataset.to_dataloader(
+ batch_size=2, shuffle=False, num_workers=0)))
+ self.assertEqual([0, 1], batch["index"].tolist())
+ self.assertEqual([2, 3], list(batch["action"].shape))
+ with self.assertRaisesRegex(ValueError, "manages collate_fn"):
+ dataset.to_dataloader(batch_size=2, collate_fn=lambda rows: rows)
dataset.return_uint8 = True
self.assertTrue(reader.return_uint8)
with self.assertRaisesRegex(TypeError, "return_uint8"):
@@ -1455,6 +1472,10 @@ class LeRobotValidationTest(unittest.TestCase):
item[key], expected.expand(3, 3, 2, 2).float()))
result[0][keys[0]].zero_()
self.assertEqual(1, result[0][keys[1]][-1, 0, 0, 0].item())
+ batch = reader._get_batch([0, 1], share_memory=False)
+ for key in keys:
+ self.assertEqual(
+ [2, 3, 3, 2, 2], list(batch[key].shape))
def test_video_windows_decode_directly_and_fall_back(self):
try:
@@ -1497,6 +1518,22 @@ class LeRobotValidationTest(unittest.TestCase):
self.assertTrue(torch.equal(batch, frames))
decoder.get_frames_at.assert_called_once_with(indices=[0, 1])
+ decoder.get_frames_at.reset_mock()
+ decoder.get_frames_at.return_value = SimpleNamespace(data=frames)
+ result = _decode_video_windows(
+ plans, rows, [collator], feature, False,
+ return_batch=True, share_memory=True)[key]
+ expected = torch.stack([
+ frames[[1, 0, 0]], frames[[1, 0, 0]]
+ ]).float().div(255)
+ self.assertTrue(result.is_shared())
+ self.assertTrue(torch.equal(result, expected))
+ result[0, 1].zero_()
+ self.assertTrue(torch.equal(result[0, 2], expected[0, 2]))
+ self.assertTrue(torch.equal(result[1], expected[1]))
+ self.assertTrue(torch.equal(frames, batch))
+ decoder.get_frames_at.assert_called_once_with(indices=[0, 1])
+
for indices in ([0, 1], [1, 0, 0]):
decoder.get_frames_at.reset_mock()
decoder.get_frames_at.return_value =
SimpleNamespace(data=frames)
@@ -1527,6 +1564,18 @@ class LeRobotValidationTest(unittest.TestCase):
self.assertTrue(torch.equal(result[1], batch))
self.assertTrue(torch.equal(batch, frames))
+ fixed = [{"windows": {key: window}}
+ for window in ([1, 0], [2, 3], [0, 1], [3, 2])]
+ result = _decode_video_windows(
+ fixed, separate, [collator], feature, True,
+ return_batch=True, share_memory=True)[key]
+ expected = torch.stack([
+ batch[[1, 0]], batch[[0, 1]],
+ batch[[0, 1]], batch[[1, 0]],
+ ])
+ self.assertTrue(result.is_shared())
+ self.assertTrue(torch.equal(result, expected))
+
open_decoder.reset_mock()
mixed = dict(rows)
mixed[1] = {key: VideoFrameDescriptor("b.video", 0, 1, 0, -1,
0).serialize()}
@@ -3096,6 +3145,16 @@ class LeRobotValidationTest(unittest.TestCase):
multiprocessing_context="spawn"):
worker_indices.extend(batch["index"].tolist())
self.assertEqual(list(range(5)), worker_indices)
+
+ worker_indices = []
+ for batch in dataset.to_dataloader(
+ batch_size=2,
+ shuffle=False,
+ num_workers=2,
+ multiprocessing_context="spawn"):
+ worker_indices.extend(batch["index"].tolist())
+ self.assertTrue(batch["camera"].is_shared())
+ self.assertEqual(list(range(5)), worker_indices)
finally:
dataset.close()
@@ -3926,6 +3985,24 @@ class LeRobotImportTest(unittest.TestCase):
worker_indices.extend(batch["index"].tolist())
self.assertEqual(list(range(5)), worker_indices)
+ worker_batches = list(dataset.to_dataloader(
+ batch_size=2,
+ shuffle=False,
+ num_workers=1,
+ multiprocessing_context="spawn",
+ worker_init_fn=_configure_uint8_transform_worker,
+ ))
+ self.assertEqual("torch.uint8", str(
+ worker_batches[0]["observation.image"].dtype))
+ self.assertEqual(
+ 1.0,
+ float(worker_batches[0]["observation.image"][0].float().mean()),
+ )
+ self.assertEqual(
+ 101.0,
+ float(worker_batches[-1]["observation.image"][-1].float().mean()),
+ )
+
uint8_dataset = pmm.PaimonLeRobotDataset(
table,
tag_name="training",
