Merge pull request 'Offset event_id across multi-shard reference reads in giant analyze (gitea #22)' (#69) from fix/issue-22 into master

Reviewed-on: #69

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))
 
 

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]