[python] Speed up LeRobot episode metadata initialization (#10754)

Author: XiaoHongbo-Hope

paimon-python/pypaimon/multimodal/lerobot/dataset.py

@@ -1032,9 +1032,30 @@ def _metadata_member(metadata, name, default=None):
     return default if value is None else value
 
 
-def _episode_row(episodes, ordinal):
-    return episodes.iloc[ordinal] if hasattr(episodes, "iloc") \
-        else episodes[ordinal]
+def _episode_rows(episodes, columns):
+    if hasattr(episodes, "select_columns"):
+        dataset_format = getattr(episodes, "format", None)
+        if isinstance(dataset_format, Mapping) and (
+                dataset_format.get("type") == "custom"
+                or dataset_format.get("columns") != episodes.column_names):
+            # Custom transforms can read or create columns.
+            for ordinal in range(len(episodes)):
+                yield episodes[ordinal]
+            return
+        columns = [name for name in columns if name in episodes.column_names]
+        if columns:
+            # Avoid formatting wide statistics columns for every episode.
+            projected = episodes.select_columns(columns)
+            for batch in projected.iter(batch_size=1024):
+                for values in zip(*(batch[name] for name in columns)):
+                    yield dict(zip(columns, values))
+        else:
+            for _ in range(len(episodes)):
+                yield {}
+    else:
+        for ordinal in range(len(episodes)):
+            yield episodes.iloc[ordinal] if hasattr(episodes, "iloc") \
+                else episodes[ordinal]
 
 
 def _episode_ranges(metadata, total_frames, total_episodes):
@@ -1047,8 +1068,8 @@ def _episode_ranges(metadata, total_frames, total_episodes):
             % (len(episodes), total_episodes))
     ranges = []
     expected = 0
-    for ordinal in range(total_episodes):
-        row = _episode_row(episodes, ordinal)
+    columns = ("episode_index", "dataset_from_index", "dataset_to_index", "length")
+    for ordinal, row in enumerate(_episode_rows(episodes, columns)):
         try:
             index = operator.index(row["episode_index"])
             begin = operator.index(row["dataset_from_index"])
@@ -1097,7 +1118,7 @@ def _build_frame_validation_context(
         metadata, episode_ranges, fps, tasks, subtasks, timestamp_type):
     task_names = _index_names(tasks, "task_index")
     subtask_names = _index_names(subtasks, "subtask_index")
-    episode_tasks = _episode_tasks(metadata, len(episode_ranges)) \
+    episode_tasks = _episode_tasks(metadata) \
         if episode_ranges is not None else None
     return {
         "episode_ranges": episode_ranges,
@@ -1145,13 +1166,12 @@ def _index_names(values, index_field):
     return result
 
 
-def _episode_tasks(metadata, total_episodes):
+def _episode_tasks(metadata):
     episodes = _metadata_member(metadata, "episodes")
     if episodes is None:
         return None
     result = []
-    for ordinal in range(total_episodes):
-        row = _episode_row(episodes, ordinal)
+    for row in _episode_rows(episodes, ("tasks",)):
         tasks = row.get("tasks") if hasattr(row, "get") else None
         if tasks is None:
             result.append(None)

paimon-python/pypaimon/tests/multimodal_lerobot_test.py

@@ -54,6 +54,8 @@
     _decode_video_frames,
     _decode_video_rows,
     _decode_video_windows,
+    _episode_ranges,
+    _episode_tasks,
     _image_tensor,
     _index_names,
     _open_video_decoder,
@@ -165,6 +167,81 @@ def _configure_uint8_transform_worker(unused_worker_id):
 
 class LeRobotValidationTest(unittest.TestCase):
 
+    def test_episode_validation_projects_batched_hf_metadata(self):
+        try:
+            from datasets import Dataset
+        except ImportError:
+            self.skipTest("datasets is not installed")
+        from pypaimon.multimodal.lerobot.dataset import (
+            _episode_ranges, _episode_tasks,
+        )
+
+        rows = [{
+            "episode_index": index,
+            "dataset_from_index": index * 2,
+            "dataset_to_index": index * 2 + 2,
+            "length": 2,
+            "tasks": ["place", "pick"],
+            "stats/camera/mean": [0.5] * 32,
+        } for index in range(2050)]
+        episodes = Dataset.from_list(rows)
+        expected = [(index * 2, index * 2 + 2) for index in range(2050)]
+        with patch.object(
+                Dataset, "__getitem__",
+                side_effect=AssertionError("Do not format individual rows")):
+            self.assertEqual(expected, _episode_ranges(
+                {"episodes": episodes}, 4100, 2050))
+            self.assertEqual((("pick", "place"),) * 2050,
+                             _episode_tasks({"episodes": episodes}))
+        self.assertIn("stats/camera/mean", episodes.column_names)
+        self.assertEqual(rows[0], episodes[0])
+        self.assertEqual((None,) * 2050, _episode_tasks({
+            "episodes": episodes.remove_columns("tasks"),
+        }))
+        formatted = episodes.select(range(2)).with_format(
+            "numpy",
+            columns=["episode_index", "dataset_from_index",
+                     "dataset_to_index", "length"],
+        )
+        self.assertNotIn("tasks", formatted[0])
+        self.assertEqual([(0, 2), (2, 4)], _episode_ranges(
+            {"episodes": formatted}, 4, 2))
+        self.assertEqual((None, None), _episode_tasks(
+            {"episodes": formatted}))
+        visible = formatted.with_format(
+            "numpy", columns=formatted.format["columns"],
+            output_all_columns=True)
+        self.assertEqual((("pick", "place"),) * 2, _episode_tasks(
+            {"episodes": visible}))
+        # Preserve a selected Dataset's logical row order and range validation.
+        with self.assertRaisesRegex(ValueError, "episode_index=1"):
+            _episode_ranges({"episodes": episodes.select([1, 0])}, 4, 2)
+        with self.assertRaisesRegex(ValueError, "integer controls"):
+            _episode_ranges({"episodes": episodes.remove_columns("length")},
+                            4100, 2050)
+
+    def test_episode_validation_preserves_metadata_representations(self):
+        from pypaimon.multimodal.lerobot.dataset import (
+            _episode_ranges, _episode_tasks,
+        )
+        import pandas as pd
+
+        rows = [{
+            "episode_index": 0, "dataset_from_index": 0,
+            "dataset_to_index": 2, "length": 2, "tasks": "pick",
+        }, {
+            "episode_index": 1, "dataset_from_index": 2,
+            "dataset_to_index": 3, "length": 1, "tasks": None,
+        }]
+        for episodes in (rows, pd.DataFrame(rows)):
+            with self.subTest(type=type(episodes).__name__):
+                metadata = {"episodes": episodes}
+                self.assertEqual([(0, 2), (2, 3)],
+                                 _episode_ranges(metadata, 3, 2))
+                self.assertEqual((("pick",), None), _episode_tasks(metadata))
+                with self.assertRaisesRegex(ValueError, "expected 4"):
+                    _episode_ranges(metadata, 4, 2)
+
     def test_episode_metadata_pickle_stays_small_and_usable(self):
         try:
             from datasets import Dataset
@@ -931,6 +1008,35 @@ def test_torchcodec_batch_matches_single_frame_decoding(self):
         finally:
             collator.close()
 
+    def test_custom_episode_transform_keeps_cross_column_inputs(self):
+        try:
+            from datasets import Dataset
+        except ImportError as error:
+            self.skipTest(str(error))
+
+        episodes = Dataset.from_dict({
+            "episode_index": [0, 1],
+            "dataset_from_index": [0, 1],
+            "dataset_to_index": [1, 3],
+            "length": [0, 0],
+            "stats": [1, 2],
+        })
+
+        def transform(batch):
+            return dict(batch, length=batch["stats"], tasks=[
+                ["pick" if index == 0 else "place"]
+                for index in batch["episode_index"]
+            ])
+
+        for source in (episodes, episodes.add_column(
+                "tasks", [["old"], ["old"]])):
+            with self.subTest(physical_tasks="tasks" in source.column_names):
+                metadata = {"episodes": source.with_transform(transform)}
+                self.assertEqual([(0, 1), (1, 3)],
+                                 _episode_ranges(metadata, 3, 2))
+                self.assertEqual((("pick",), ("place",)),
+                                 _episode_tasks(metadata))
+
     def test_dataset_requires_supported_python(self):
         with patch(
                 "pypaimon.multimodal.lerobot.dataset.sys.version_info",
@@ -1038,6 +1144,19 @@ def close(self):
         self.assertEqual([1.0], dataset.meta.stats["action"]["mean"].tolist())
         self.assertEqual(0, dataset.meta.get_task_index("pick"))
         self.assertEqual("pick", sample["task"])
+        try:
+            from datasets import Dataset
+        except ImportError:
+            pass
+        else:
+            formatted_metadata = dict(metadata)
+            formatted_metadata["episodes"] = Dataset.from_list(
+                metadata["episodes"]).with_format(
+                    "numpy", columns=["episode_index", "dataset_from_index",
+                                      "dataset_to_index", "length"])
+            formatted_reader = Reader(formatted_metadata)
+            self.assertEqual("pick", formatted_reader[1]["task"])
+            formatted_reader.close()
         torch.testing.assert_close(
             sample["observation.state"], torch.tensor([1.0, 2.0]))
         torch.testing.assert_close(