Offset event_id per file to avoid cross-file collisions
CI / Format (ruff format) (push) Successful in 25s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 23s
CI / Lint (ruff check) (pull_request) Successful in 27s
CI / Tests (push) Successful in 1m14s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (pull_request) Successful in 59s

Each input parquet file is one Geant4 job (scripts/steps_to_parquet.py),
and a job's event_id numbering always restarts from 0 — so loading
multiple files together (a directory or .manifest) let same-numbered
events from different files collapse into one during the event index
scan and train/val split, corrupting both. Every per-file event_id now
gets offset by file index * EVENT_ID_FILE_STRIDE (giant/data/loader.py),
threaded through the setup-cache event index, the streaming dataset,
and predict/rollout seeding. Bumps the setup-cache format version so
stale sidecars computed pre-fix are invalidated.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-30 11:36:48 +02:00
co-authored by Claude Sonnet 5
parent e288c3fe21
commit b944bba8fb
8 changed files with 268 additions and 24 deletions
+33 -12
View File
@@ -12,6 +12,20 @@ import pyarrow.parquet as pq
# dataset tree is moved or copied elsewhere intact.
MANIFEST_SUFFIX = ".manifest"
# Each input parquet file is a separate Geant4 job converted 1:1 from its own
# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
# always restarts from 0 — so when multiple files are loaded together (a
# directory or .manifest), raw event_id values collide across files even
# though they refer to unrelated events. Every per-file event_id column gets
# offset by its file's index in the (deterministically ordered) files list
# so ids stay globally unique across a multi-file load; the stride is far
# larger than any realistic per-file event count.
EVENT_ID_FILE_STRIDE = 1_000_000
def event_id_offset(file_index: int) -> int:
return file_index * EVENT_ID_FILE_STRIDE
def _read_manifest(path: Path) -> list[Path]:
files = []
@@ -78,13 +92,13 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
return out
def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
from giant.constants import K_MAX
has_sec_lists = "sec_E_list" in df.columns
d: dict[str, np.ndarray] = {
"event_id": df["event_id"].to_numpy(),
"event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
"pdg": df["pdg"].to_numpy(dtype=np.int32),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -121,20 +135,23 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
return d
def load_steps(path: str | Path) -> dict[str, np.ndarray]:
return _df_to_dict(pd.read_parquet(path))
def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
return _df_to_dict(pd.read_parquet(path), offset=offset)
def load_event_ids(path: str | Path) -> np.ndarray:
def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
"""Read only the event_id column — cheap scan for split assignment."""
return pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
return ids.astype(np.int64) + offset
def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
def iter_file_chunks(
path: str | Path, offset: int = 0
) -> Iterator[dict[str, np.ndarray]]:
"""Yield one parquet row-group at a time so a large file never fully loads."""
pf = pq.ParquetFile(path)
for i in range(pf.num_row_groups):
yield _df_to_dict(pf.read_row_group(i).to_pandas())
yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)
_COND_COLS = [
@@ -154,9 +171,9 @@ _COND_COLS = [
]
def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
return {
"event_id": df["event_id"].to_numpy(),
"event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
"pdg": df["pdg"].to_numpy(dtype=np.int32),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -168,11 +185,15 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
}
def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
def iter_cond_chunks(
path: str | Path, offset: int = 0
) -> Iterator[dict[str, np.ndarray]]:
"""Yield conditioning-only row-groups (no post-step columns read from disk)."""
pf = pq.ParquetFile(path)
for i in range(pf.num_row_groups):
yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas())
yield _cond_df_to_dict(
pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
)
def build_index_maps(