Skip to content

Commit 7b37662

Browse files
committed
Save event_no as a non-index column in ParquetWriter
Fixes #871
1 parent ad36dc6 commit 7b37662

1 file changed

Lines changed: 3 additions & 3 deletions

File tree

src/graphnet/data/writers/parquet_writer.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ def _save_file(
5555

5656
table_dir = os.path.join(save_path, f"{table}")
5757
os.makedirs(table_dir, exist_ok=True)
58-
df = data[table].set_index(self._index_column)
58+
df = data[table]
5959
df.to_parquet(
6060
os.path.join(table_dir, file_name + f"_{table}.parquet")
6161
)
@@ -199,8 +199,8 @@ def _process_shard(self, settings: List[Any]) -> None:
199199
id = split[index_column][split["file_name"] == unique_file]
200200

201201
# Filter out indices that point to empty events
202-
idx = [i for i in id if i in df.index]
203-
table_shards.append(df.loc[idx, :])
202+
idx = [i for i in id if i in set(df[index_column])]
203+
table_shards.append(df[df[index_column].isin(idx)])
204204

205205
os.makedirs(os.path.join(outdir, table), exist_ok=True)
206206
if len(table_shards) > 0:

0 commit comments

Comments
 (0)