From 3466848de6e0362923a62fdb8aba9b3a7c19ef07 Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Mon, 15 Jun 2026 17:49:05 +0200 Subject: [PATCH] Update src/lerobot/datasets/io_utils.py Co-authored-by: Caroline Pascal Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> --- src/lerobot/datasets/io_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/lerobot/datasets/io_utils.py b/src/lerobot/datasets/io_utils.py index 5e3f10155..4beee4686 100644 --- a/src/lerobot/datasets/io_utils.py +++ b/src/lerobot/datasets/io_utils.py @@ -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: