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 3ed90bc68d [python] Reduce LeRobot DataLoader worker startup time
(#9888)
3ed90bc68d is described below
commit 3ed90bc68d9a814ee09e409d267c7c10715339fe
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Sep 20 10:09:10 2026 +0800
[python] Reduce LeRobot DataLoader worker startup time (#9888)
---
.../pypaimon/multimodal/lerobot/dataset.py | 19 ++++++
.../pypaimon/tests/multimodal_lerobot_test.py | 77 ++++++++++++++++++++++
2 files changed, 96 insertions(+)
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 062192150a..494c111597 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -21,7 +21,9 @@ import bisect
import json
import math
import operator
+import pickle
import sys
+import zlib
from abc import ABC, abstractmethod
from collections import OrderedDict
from collections.abc import Mapping
@@ -709,6 +711,22 @@ class _PaimonLeRobotMetadata:
self.episodes = episodes
self.tasks = tasks
self.subtasks = subtasks
+ self._compress_episodes = False
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ if state.get("_compress_episodes", False):
+ # Keep worker-startup payloads small without changing Dataset
state.
+ state["episodes"] = zlib.compress(
+ pickle.dumps(self.episodes, protocol=pickle.HIGHEST_PROTOCOL),
+ level=1)
+ state["_episodes_zlib"] = True
+ return state
+
+ def __setstate__(self, state):
+ if state.pop("_episodes_zlib", False):
+ state["episodes"] =
pickle.loads(zlib.decompress(state["episodes"]))
+ self.__dict__.update(state)
def __getattr__(self, name):
info = self.__dict__.get("info", {})
@@ -822,6 +840,7 @@ def _load_dataset(table, tag_name):
metadata = _PaimonLeRobotMetadata(
str(table.identifier), tag_name, info, stats, episodes, tasks,
subtasks)
+ metadata._compress_episodes = True
return frames, metadata
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index cdd027c53c..6495381850 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -46,6 +46,7 @@ from pypaimon.multimodal.source_utils import _SourceFileIO
from pypaimon.multimodal.connection import MultimodalConnection
from pypaimon.multimodal.lerobot import load_from_lerobot
from pypaimon.multimodal.lerobot.dataset import (
+ _PaimonLeRobotMetadata,
_PyAVVideoDecoder,
_arrow_rows,
_decode_video_frames,
@@ -149,6 +150,59 @@ def _catalog_metadata(connection, name):
class LeRobotValidationTest(unittest.TestCase):
+ def test_episode_metadata_pickle_stays_small_and_usable(self):
+ try:
+ from datasets import Dataset
+ except ImportError:
+ self.skipTest("datasets is not installed")
+
+ rows = [{
+ "episode_index": index,
+ "dataset_from_index": index * 400,
+ "dataset_to_index": (index + 1) * 400,
+ "length": 400,
+ "tasks": ["pick", "place"],
+ } for index in range(50)]
+ episodes_arrow = pa.Table.from_pylist(rows)
+ fingerprint = "0123456789abcdef"
+ episodes = Dataset(episodes_arrow, fingerprint=fingerprint)
+ metadata = _PaimonLeRobotMetadata(
+ "robot", "tag", {"fps": 50}, None, episodes, ["pick", "place"],
+ None)
+ metadata._compress_episodes = True
+
+ payload = pickle.dumps(metadata)
+ self.assertLess(len(payload), len(pickle.dumps(episodes)) * 3 // 4)
+ restored = pickle.loads(payload)
+ self.assertIsInstance(restored.episodes, Dataset)
+ self.assertEqual(episodes[:], restored.episodes[:])
+ self.assertEqual(episodes.features, restored.episodes.features)
+ self.assertEqual(episodes._fingerprint, restored.episodes._fingerprint)
+ self.assertEqual("tag", restored.revision)
+ self.assertEqual(50, restored.fps)
+
+ episodes.set_format("numpy")
+ restored = pickle.loads(pickle.dumps(metadata))
+ self.assertEqual("numpy", restored.episodes.format["type"])
+ episodes.reset_format()
+
+ metadata.episodes = episodes.with_format("numpy")
+ restored = pickle.loads(pickle.dumps(metadata))
+ self.assertEqual("numpy", restored.episodes.format["type"])
+
+ with tempfile.TemporaryDirectory() as directory:
+ path = str(Path(directory) / "episodes.arrow")
+ with pa.OSFile(path, "wb") as output:
+ with pa.ipc.new_stream(
+ output, episodes.data.table.schema) as writer:
+ writer.write_table(episodes.data.table)
+ metadata.episodes = Dataset.from_file(path)
+ restored = pickle.loads(pickle.dumps(metadata))
+ self.assertEqual(
+ metadata.episodes.cache_files,
+ restored.episodes.cache_files,
+ )
+
def test_video_columns_decode_in_parallel(self):
barrier = threading.Barrier(2)
@@ -3068,6 +3122,29 @@ class LeRobotImportTest(unittest.TestCase):
dataset.save_episode()
dataset.finalize()
+ def test_table_dataset_pickle_preserves_episode_metadata_and_reads(self):
+ import torch
+
+ self.connection.load_from_lerobot("worker_pickle", self.image_source)
+ table = self.connection.get_table("worker_pickle")
+ dataset = pmm.PaimonLeRobotDataset(table, return_uint8=True)
+ restored = pickle.loads(pickle.dumps(dataset))
+
+ self.assertEqual(dataset.meta.episodes[:], restored.meta.episodes[:])
+ self.assertEqual(
+ dataset.meta.episodes._fingerprint,
+ restored.meta.episodes._fingerprint,
+ )
+ for index in (0, 2, 4):
+ original = dataset[index]
+ reread = restored[index]
+ self.assertEqual(original.keys(), reread.keys())
+ for key in original:
+ if torch.is_tensor(original[key]):
+ self.assertTrue(torch.equal(original[key], reread[key]))
+ else:
+ self.assertEqual(original[key], reread[key])
+
def test_import_infers_schema_and_preserves_episodes(self):
import pandas as pd
