From 7b376628d5165764d0f61299a4f79f97d921b941 Mon Sep 17 00:00:00 2001 From: Jarred Green Date: Fri, 13 Mar 2026 16:05:28 +0100 Subject: [PATCH] Save `event_no` as a non-index column in ParquetWriter Fixes #871 --- src/graphnet/data/writers/parquet_writer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/graphnet/data/writers/parquet_writer.py b/src/graphnet/data/writers/parquet_writer.py index ce4085430..eaab096cb 100644 --- a/src/graphnet/data/writers/parquet_writer.py +++ b/src/graphnet/data/writers/parquet_writer.py @@ -55,7 +55,7 @@ def _save_file( table_dir = os.path.join(save_path, f"{table}") os.makedirs(table_dir, exist_ok=True) - df = data[table].set_index(self._index_column) + df = data[table] df.to_parquet( os.path.join(table_dir, file_name + f"_{table}.parquet") ) @@ -199,8 +199,8 @@ def _process_shard(self, settings: List[Any]) -> None: id = split[index_column][split["file_name"] == unique_file] # Filter out indices that point to empty events - idx = [i for i in id if i in df.index] - table_shards.append(df.loc[idx, :]) + idx = [i for i in id if i in set(df[index_column])] + table_shards.append(df[df[index_column].isin(idx)]) os.makedirs(os.path.join(outdir, table), exist_ok=True) if len(table_shards) > 0: