From c1e6ffd8c6e06a0dc47b879f44a4b0a31a141a6f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 17 Aug 2026 15:53:02 +0200 Subject: [PATCH] Offset event_id across multi-shard reference reads in giant analyze (gitea #22) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit giant/analysis/sources.py's open_side scanned a reference directory of parquet shards with a bare glob and never offset event_id across them. Each shard is a separate Geant4 job whose own event_id numbering restarts from 0, so events from different shards collided on the same event_id, corrupting every downstream per-event grouping and the event_id % n_chunks condor chunking — the same root cause already fixed on the training/rollout side via giant/data/loader.py's per-file event_id_offset. open_side's reference branch now uses find_parquet_files (the same deterministically ordered file lister giant rollout's _seed_from_data uses) and offsets each shard's event_id via a join on polars' include_file_paths, so both sides of a comparison agree on what an event_id means. Two incidental behaviour changes come along for free: .manifest references now work (they crashed before), and the directory glob narrows from recursive **/*.parquet to top-level *.parquet, matching the file list rollout itself used to assign offsets — a deliberate choice, since a differing file list would make the two sides' offsets disagree again in a subtler way. No overflow guard on the per-shard offset stride (unlike loader's _offset_event_id): checking it here would cost an eager event_id-column read per shard in every condor compute job, and giant rollout already runs that check over the same file list when producing the seed. Co-Authored-By: Claude Sonnet 5 --- giant/analysis/sources.py | 36 ++++++++++++++++++++++++++++++++--- tests/test_analysis_reduce.py | 34 +++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py index abedc74..250a8e5 100644 --- a/giant/analysis/sources.py +++ b/giant/analysis/sources.py @@ -40,6 +40,11 @@ from giant.constants import ( TERM_MAX_STEPS, TERM_UNKNOWN_PDG, ) +from giant.data.loader import event_id_offset, find_parquet_files + +# Helper column name for the per-shard offset join in open_side; dropped before +# the LazyFrame is returned, so it never leaks into a caller's schema. +_SOURCE_PATH_COL = "__source_path" # The world-frame physical columns both sides share under identical names. PHYS_COLS: tuple[str, ...] = ( @@ -108,6 +113,22 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame: reference file's upstream ROOT→parquet conversion don't agree on integer width, and an uncast mismatch only surfaces later as a ``pl.concat`` ``SchemaError`` (e.g. in ``build_context``'s pdg-count merge). + + The reference (a rollout's seed ``dataset``) may be a directory of parquet + shards, or a ``.manifest`` naming a subset, rather than a single file — each + such shard is a separate Geant4 job whose own ``event_id`` numbering + restarts from 0, so a multi-shard load offsets every shard's ids by + ``giant.data.loader.event_id_offset(file_index)`` to keep them globally + unique, exactly as the training/rollout data pipeline already does + (``giant/data/loader.py``). ``file_index`` comes from + ``find_parquet_files``'s deterministic ordering — the same list and + ordering ``giant rollout`` used (via ``_seed_from_data``) to offset the + rollout side's own ``event_id``s, so both sides agree on what an + ``event_id`` means. There is no overflow guard here (unlike + ``loader._offset_event_id``): checking it would cost an eager + ``event_id``-column read per shard in every condor compute job, and + ``giant rollout`` already ran that check over this exact file list when it + produced the seed. """ if isinstance(source, pl.LazyFrame): return source.with_columns(pl.col("pdg").cast(pl.Int64)) @@ -116,9 +137,18 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame: _check_rollout_metadata(path) lf = pl.scan_parquet(path) else: - # The reference (a rollout's seed `dataset`) may be a directory of - # parquet shards rather than a single file — scan them all. - lf = pl.scan_parquet(str(path / "**/*.parquet")) if path.is_dir() else pl.scan_parquet(path) + files = find_parquet_files(path) + if len(files) == 1: + lf = pl.scan_parquet(files[0]) + else: + offsets = {str(p): event_id_offset(i) for i, p in enumerate(files)} + lf = ( + pl.scan_parquet(files, include_file_paths=_SOURCE_PATH_COL) + .with_columns( + pl.col("event_id") + pl.col(_SOURCE_PATH_COL).replace_strict(offsets, return_dtype=pl.Int64) + ) + .drop(_SOURCE_PATH_COL) + ) return lf.with_columns(pl.col("pdg").cast(pl.Int64)) diff --git a/tests/test_analysis_reduce.py b/tests/test_analysis_reduce.py index 614233c..1241716 100644 --- a/tests/test_analysis_reduce.py +++ b/tests/test_analysis_reduce.py @@ -10,9 +10,11 @@ from giant.analysis import reduce as R from giant.analysis.sources import ( SYNTHETIC_TERMINATION_REASONS, Side, + open_side, physical_steps, secondaries, ) +from giant.data.loader import EVENT_ID_FILE_STRIDE def _rollout_frame() -> pl.LazyFrame: @@ -194,3 +196,35 @@ def test_pdg_and_material_labels(): assert G.pdg_label(22) == "gamma" assert G.pdg_label(999999) == "999999" assert G.material_label("G4_PbWO4") == "PbWO4" + + +def _write_shard(path, event_ids, edeps): + pl.DataFrame({"event_id": event_ids, "pdg": [11] * len(event_ids), "edep": edeps}).write_parquet(path) + + +def test_open_side_reference_offsets_event_ids_across_shards(tmp_path): + # Each shard is a separate Geant4 job whose own event_id numbering restarts + # from 0 — a naive multi-shard scan collides on event_id across shards. + _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0]) + _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0]) + df = open_side(tmp_path, Side.reference).sort("event_id").collect() + assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1] + assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0] + assert "__source_path" not in df.columns + + +def test_open_side_reference_single_file_unchanged(tmp_path): + _write_shard(tmp_path / "only.parquet", [0, 1], [1.0, 2.0]) + df = open_side(tmp_path / "only.parquet", Side.reference).sort("event_id").collect() + assert df["event_id"].to_list() == [0, 1] + assert "__source_path" not in df.columns + + +def test_open_side_reference_manifest(tmp_path): + _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0]) + _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0]) + manifest = tmp_path / "shards.manifest" + manifest.write_text("a.parquet\nb.parquet\n") + df = open_side(manifest, Side.reference).sort("event_id").collect() + assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1] + assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]