Update src/lerobot/datasets/io_utils.py

Co-authored-by: Caroline Pascal <caroline8.pascal@gmail.com>
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
This commit is contained in:
Pepijn
2026-06-15 17:49:05 +02:00
committed by GitHub
parent 99840ebef3
commit 3466848de6
+4 -1
View File
@@ -312,7 +312,10 @@ def to_parquet_with_hf_images(
def to_parquet_one_row_group_per_episode(df: pandas.DataFrame, path: Path) -> None:
"""Write a (non-image) DataFrame to parquet with one row group per episode."""
table = pa.Table.from_pandas(df, preserve_index=False)
write_table_one_row_group_per_episode(table, path)
if "episode_index" in table.column_names:
write_table_one_row_group_per_episode(table, path)
else:
pq.write_table(table, str(path))
def item_to_torch(item: dict) -> dict: