Incremental parquet writing

src/lerobot/datasets/aggregate.py

@@ -31,8 +31,8 @@
     DEFAULT_EPISODES_PATH,
     DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_PATH,
+    get_file_size_in_mb,
     get_parquet_file_size_in_mb,
-    get_video_size_in_mb,
     to_parquet_with_hf_images,
     update_chunk_file_indices,
     write_info,
@@ -217,6 +217,7 @@ def aggregate_datasets(
         robot_type=robot_type,
         features=features,
         root=aggr_root,
+        use_videos=len(video_keys) > 0,
         chunks_size=chunk_size,
         data_files_size_in_mb=data_files_size_in_mb,
         video_files_size_in_mb=video_files_size_in_mb,
@@ -307,8 +308,9 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
                 current_offset += src_duration
                 continue
 
-            src_size = get_video_size_in_mb(src_path)
-            dst_size = get_video_size_in_mb(dst_path)
+            # Check file sizes before appending
+            src_size = get_file_size_in_mb(src_path)
+            dst_size = get_file_size_in_mb(dst_path)
 
             if dst_size + src_size >= video_files_size_in_mb:
                 # Rotate to a new file, this source becomes start of new destination

src/lerobot/datasets/dataset_tools.py

@@ -42,6 +42,7 @@
     DEFAULT_DATA_PATH,
     DEFAULT_EPISODES_PATH,
     get_parquet_file_size_in_mb,
+    load_episodes,
     to_parquet_with_hf_images,
     update_chunk_file_indices,
     write_info,
@@ -436,6 +437,9 @@ def _copy_and_reindex_data(
     Returns:
         dict mapping episode index to its data file metadata (chunk_index, file_index, etc.)
     """
+    if src_dataset.meta.episodes is None:
+        src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
+
     file_to_episodes: dict[Path, set[int]] = {}
     for old_idx in episode_mapping:
         file_path = src_dataset.meta.get_data_file_path(old_idx)
@@ -645,6 +649,8 @@ def _copy_and_reindex_videos(
     Returns:
         dict mapping episode index to its video metadata (chunk_index, file_index, timestamps)
     """
+    if src_dataset.meta.episodes is None:
+        src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
 
     episodes_video_metadata: dict[int, dict] = {new_idx: {} for new_idx in episode_mapping.values()}
 
@@ -770,6 +776,9 @@ def _copy_and_reindex_episodes_metadata(
     """
     from lerobot.datasets.utils import flatten_dict
 
+    if src_dataset.meta.episodes is None:
+        src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
+
     all_stats = []
     total_frames = 0
 
@@ -831,6 +840,8 @@ def _copy_and_reindex_episodes_metadata(
 
         total_frames += src_episode["length"]
 
+    dst_meta._close_writer()
+
     dst_meta.info.update(
         {
             "total_episodes": len(episode_mapping),