chore: update changelog for v0.3.17

Reviewed-on: #93

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.16"
+current_version = "0.3.17"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## [0.3.17] - 2026-09-02
+
+### Changed
+
+- Feat: add WGAN + AR stop-token config variant
+
+- Perf: replace pandas with polars in the setup-stage scan
+
 ## [0.3.16] - 2026-08-31
 
 ### Changed

configs/wgan_ar_stoptoken.toml

@@ -0,0 +1,122 @@
+# GIANT WGAN-GP + AR stop-token variant of configs/baseline.toml.
+#
+# Two roadmap axes, combined into one run: WGAN-GP generators for both
+# stages (unbenchmarked since the 2026-08-03 pre-v0.3.0 failure, which was
+# secondary-species mode collapse — the failure v0.3.0's AR/categorical
+# pivot exists to fix) and the AR stop-token multiplicity mode
+# (stage2_model.n_sec.mode = "stop_token", never benchmarked at all).
+# Everything else is byte-identical to baseline.toml so a rollout compared
+# against baseline's analysis_341dfb14 is attributable to these two axes
+# alone: conditioning (physical/physical), hidden_dim 512 / n_res_blocks 6 /
+# dropout 0.0 per stage, k_max 15, history "markov", teacher_forcing
+# "always", particle_type.target "onehot" (n_classes 32, other_policy
+# "sample"), lr 3e-4, warmup_epochs 3, weight_decay 0.01, ema_decay 0.9999,
+# val_fraction 0.1, num_workers 4, seed 0, validate_steps 10, W&B on.
+#
+# No [stage1_model.wgan] / [stage2_model.wgan] block: the dataclass defaults
+# (noise_dim 64, n_critic 5, gp_weight 10.0, critic_lr 0.0 = inherit
+# train.lr, critic_hidden_dim/critic_n_res_blocks 0 = inherit the stage's
+# 512/6, stage 2's gumbel_tau_start/_end 1.0/0.1) are what the earlier WGAN
+# runs used — writing them out would add keys that don't vary.
+#
+# particle_type.class_weighting stays "none" (the default): config.py's
+# validate_config rejects any other value under stage2_model.generator =
+# "wgan", since that path feeds the type slice to the critic via a
+# straight-through Gumbel relaxation instead of a weighted cross-entropy.
+#
+# Prior WGAN writeup (pre-v0.3.0, describes the failure this run re-tests):
+# /home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to baseline.toml.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "wgan"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# Autoregressive in descending-energy order, as baseline.toml — this variant
+# only swaps the generator (flow -> wgan) and the multiplicity mode
+# (head -> stop_token), not the decoder shape.
+decoder      = "autoregressive"
+generator    = "wgan"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.n_sec]
+# EOS-style per-slot stop head on the AR secondary decoder, replacing the
+# n_sec classifier entirely (mutually exclusive — see NSecConfig's
+# docstring in giant/config.py). Requires decoder = "autoregressive" and
+# owner = "stage2" (both already true above/by default); validate_config
+# enforces this.
+mode = "stop_token"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs = 30
+# Halved from baseline's 36864. That figure came from a measured linear fit
+# of the *flow-AR* training step (peak reserved MiB = 0.9736 * batch_size +
+# 115); WGAN invalidates it twice over — each stage gains a critic that by
+# default inherits the stage's own 512/6 body, and gradient_penalty
+# (giant/model/wgan.py, forced fp32 internally) runs a double-backward every
+# batch. 18432 is a conservative choice pending a real memory measurement on
+# this exact config, not a re-derived fit. Throughput is already flat above
+# bs~4096 on the 4070, so this costs occupancy on the L40S, not step
+# efficiency.
+batch_size    = 18432
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# Tightened from baseline's 10: WGANStageTrainer.supports_val_loss = False,
+# and with both stages adversarial there is no per-epoch val loss at all, so
+# validate_every's marginal-KL pass (giant/training/trainers.py's
+# val_objective) is the only comparable-across-epochs best-checkpoint
+# selection signal available. 5 gives 6 evaluations over 30 epochs instead
+# of baseline's 3, at ~6x5000s of extra walltime.
+validate_every = 5
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"

giant/data/loader.py

@@ -1,13 +1,16 @@
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Iterator
+from typing import TYPE_CHECKING, Any, Iterator, Mapping
 
 import numpy as np
-import pandas as pd
+import polars as pl
 import pyarrow.parquet as pq
 
 from giant.constants import K_MAX
 
+if TYPE_CHECKING:
+    from giant.data.scan import ValueStat
+
 # A manifest is a plain text file listing one parquet path per line, used to
 # name a curated subset of files (e.g. a train/holdout pool) without copying
 # or symlinking the underlying parquet files. Lines are resolved relative to
@@ -77,88 +80,75 @@ def find_parquet_files(path: str | Path) -> list[Path]:
     return [p]
 
 
-def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
+def _pad_list_column(df: pl.DataFrame, col: str, k: int, fill, dtype: type[pl.DataType] | pl.DataType) -> np.ndarray:
-    """Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
+    """Pad / truncate a list-valued column to fixed width `k` → (N, k) numpy array.
-    out = np.full((len(series), K), fill, dtype=np.float32)
-    for i, lst in enumerate(series):
-        if lst is not None and len(lst) > 0:
-            n = min(len(lst), K)
-            out[i, :n] = lst[:n]
-    return out
 
-
+    Concatenating `k` fill values before truncating to `k` guarantees every
-def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
+    row ends up with exactly `k` non-null elements regardless of how short
-    """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
+    (including empty) or long the original list was, so `list.to_array(k)`
-    out = np.full((len(series), K), fill, dtype=np.int64)
+    (a fixed-size-array dtype) converts to a plain 2D numpy array with a
-    for i, lst in enumerate(series):
+    single vectorized expression — no per-row Python loop.
-        if lst is not None and len(lst) > 0:
-            n = min(len(lst), K)
-            out[i, :n] = lst[:n]
-    return out
-
-
-def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
-    """Pad three list-valued direction columns → (N, K, 3) float32.
-
-    Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
     """
-    N = len(dx)
+    fill_tail = pl.lit([fill] * k, dtype=pl.List(dtype))
-    out = np.zeros((N, K, 3), dtype=np.float32)
+    out = df.select(pl.col(col).cast(pl.List(dtype)).list.concat(fill_tail).list.head(k).list.to_array(k).alias("_p"))
-    out[:, :, 2] = 1.0
+    return out["_p"].to_numpy()
-    for i in range(N):
-        lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
-        if lx is not None and len(lx) > 0:
-            n = min(len(lx), K)
-            out[i, :n, 0] = lx[:n]
-            out[i, :n, 1] = ly[:n]
-            out[i, :n, 2] = lz[:n]
-    return out
 
 
-def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
+def _pad_dir_col(df: pl.DataFrame, dx: str, dy: str, dz: str, k: int) -> np.ndarray:
+    """Pad three list-valued direction columns → (N, k, 3) float32.
+
+    Padding direction defaults to (0, 0, 1) (forward) so it is a valid unit vector.
+    """
+    px = _pad_list_column(df, dx, k, 0.0, pl.Float64)
+    py = _pad_list_column(df, dy, k, 0.0, pl.Float64)
+    pz = _pad_list_column(df, dz, k, 1.0, pl.Float64)
+    return np.stack([px, py, pz], axis=-1).astype(np.float32)
+
+
+def _df_to_dict(df: pl.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
         "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
-        "pdg": df["pdg"].to_numpy(dtype=np.int32),
+        "pdg": df["pdg"].to_numpy().astype(np.int32),
-        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
+        "pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
-        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
+        "pre_E": df["pre_E"].to_numpy().astype(np.float32),
-        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
+        "pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
-        "material": df["material"].to_numpy(dtype=object),
+        "material": df["material"].to_numpy().astype(object),
-        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
+        "layer_id": df["layer_id"].to_numpy().astype(np.int32),
-        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
+        "n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
-        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
+        "e_sec": df["e_sec"].to_numpy().astype(np.float32),
         # The physics process that ended the step (e.g. "compt", "phot",
         # "eBrem") — a post-step outcome, so it's a router/classifier
         # supervision label only, never conditioning (see build_process_map*
         # / ProcessRouter). Guarded like has_sec_lists: older parquet
         # conversions predating this column still load fine.
         "process": (
-            df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
+            df["process"].to_numpy().astype(object) if "process" in df.columns else np.full(len(df), "", dtype=object)
         ),
-        "step_length": df["step_length"].to_numpy(dtype=np.float32),
+        "step_length": df["step_length"].to_numpy().astype(np.float32),
-        "post_E": df["post_E"].to_numpy(dtype=np.float32),
+        "post_E": df["post_E"].to_numpy().astype(np.float32),
-        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
+        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy().astype(np.float32),
-        "edep": df["edep"].to_numpy(dtype=np.float32),
+        "edep": df["edep"].to_numpy().astype(np.float32),
-        "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
+        "post_dir": df.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32),
-        "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
+        "post_pos": df.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32),
     }
 
     if has_sec_lists:
-        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
+        d["sec_E_list"] = _pad_list_column(df, "sec_E_list", k_max, 0.0, pl.Float64).astype(np.float32)
-        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
+        d["sec_pdg_list"] = _pad_list_column(df, "sec_pdg_list", k_max, 0, pl.Int64).astype(np.int64)
-        d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
+        d["sec_dir_list"] = _pad_dir_col(df, "sec_dx_list", "sec_dy_list", "sec_dz_list", k_max)
 
     return d
 
 
 def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
-    return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
+    return _df_to_dict(pl.read_parquet(path), offset=offset, k_max=k_max)
 
 
 def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
     """Read only the event_id column — cheap scan for split assignment."""
-    ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
+    ids = pl.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
     return _offset_event_id(ids, offset)
 
 
@@ -170,7 +160,7 @@ def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> I
     module constant for callers that don't care (e.g. Stage-1-only reads)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)
+        yield _df_to_dict(pl.DataFrame(pf.read_row_group(i)), offset=offset, k_max=k_max)
 
 
 _COND_COLS = [
@@ -190,17 +180,17 @@ _COND_COLS = [
 ]
 
 
-def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
+def _cond_df_to_dict(df: pl.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     return {
         "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
-        "pdg": df["pdg"].to_numpy(dtype=np.int32),
+        "pdg": df["pdg"].to_numpy().astype(np.int32),
-        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
+        "pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
-        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
+        "pre_E": df["pre_E"].to_numpy().astype(np.float32),
-        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
+        "pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
-        "material": df["material"].to_numpy(dtype=object),
+        "material": df["material"].to_numpy().astype(object),
-        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
+        "layer_id": df["layer_id"].to_numpy().astype(np.int32),
-        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
+        "n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
-        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
+        "e_sec": df["e_sec"].to_numpy().astype(np.float32),
     }
 
 
@@ -208,7 +198,7 @@ def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np
     """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(), offset=offset)
+        yield _cond_df_to_dict(pl.DataFrame(pf.read_row_group(i, columns=_COND_COLS)), offset=offset)
 
 
 def build_index_maps(
@@ -225,42 +215,28 @@ def build_index_maps(
 def build_index_maps_from_files(
     files: list[Path],
 ) -> tuple[dict[int, int], dict[str, int]]:
-    """Scan only pdg and material columns across all files (2-column read)."""
+    """Scan only pdg and material columns across all files (fused single-pass scan)."""
-    pdg_vals: set[int] = set()
+    from giant.data.scan import ScanRequest, scan_metadata
-    mat_vals: set[str] = set()
+
-    for path in files:
+    result = scan_metadata(files, ScanRequest(pdg=True, material=True))
-        df = pd.read_parquet(path, columns=["pdg", "material"])
+    assert result.pdg is not None and result.material is not None
-        pdg_vals.update(int(v) for v in df["pdg"].unique())
-        mat_vals.update(str(v) for v in df["material"].unique())
     return (
-        {v: i for i, v in enumerate(sorted(pdg_vals))},
+        {v: i for i, v in enumerate(sorted(result.pdg))},
-        {v: i for i, v in enumerate(sorted(mat_vals))},
+        {v: i for i, v in enumerate(sorted(result.material))},
     )
 
 
-def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
+def _topn_plus_other_map(counts: "Mapping[Any, ValueStat]", n_classes: int) -> tuple[dict, dict, dict]:
-    for name, count in series.value_counts().items():
-        name = cast(name)
-        counts[name] = counts.get(name, 0) + int(count)
-
-
-def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
-    """Scan `column` across `files` and return `{cast(value): total_count}`,
-    accumulated in file order (see `fingerprint_files`'s docstring on why
-    scan order — not a normalized/sorted order — is preserved: it drives
-    tie-breaking in the frequency ranking below)."""
-    counts: dict = {}
-    for path in files:
-        df = pd.read_parquet(path, columns=[column])
-        _accumulate_value_counts(counts, df[column], cast)
-    return counts
-
-
-def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
     """Frequency-capped value->index map: the `n_classes - 1` most frequent
     keys get their own index; every rarer key is bucketed into a shared
     "other" index (`n_classes - 1`).
 
+    `counts` maps each key to something with `.count` and `.first_seen`
+    attributes (`giant.data.scan.ValueStat`) — ties in `.count` are broken by
+    `.first_seen` (whichever value was scanned first: file order, then row
+    order within a file — see `giant.data.scan`'s module docstring). This is
+    an explicit, documented contract, not an accident of iteration order.
+
     Returns `(class_map, other_members, class_counts)` — `other_members` is
     `{key: count}` for every key bucketed into "other" (the empirical
     within-bucket distribution, for `other_policy = "sample"` at rollout);
@@ -270,15 +246,15 @@ def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict
     (gitea #44) needs and that would otherwise be dropped once `counts` is
     collapsed into `class_map`.
     """
-    ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
+    ranked = sorted(counts, key=lambda k: (-counts[k].count, counts[k].first_seen))
     keep = ranked[: max(n_classes - 1, 0)]
     class_map = {k: i for i, k in enumerate(keep)}
-    class_counts = {i: counts[k] for i, k in enumerate(keep)}
+    class_counts = {i: counts[k].count for i, k in enumerate(keep)}
     other_idx = n_classes - 1
     other_members: dict = {}
     for k in ranked[len(keep) :]:
         class_map[k] = other_idx
-        other_members[k] = counts[k]
+        other_members[k] = counts[k].count
     if other_members:
         class_counts[other_idx] = sum(other_members.values())
     return class_map, other_members, class_counts
@@ -294,8 +270,11 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
     mirrors how `build_features` clamps the n_sec label to K_MAX for the
     fixed-width n_sec_head classifier.
     """
-    counts = _rank_by_frequency_from_files(files, "process", str)
+    from giant.data.scan import ScanRequest, scan_metadata
-    class_map, _, _ = _topn_plus_other_map(counts, n_experts)
+
+    result = scan_metadata(files, ScanRequest(process=True))
+    assert result.process is not None
+    class_map, _, _ = _topn_plus_other_map(result.process, n_experts)
     return class_map
 
 
@@ -327,8 +306,13 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca
     later for `other_policy = "sample"` at rollout — computed now since it's
     free during this same scan.
     """
-    counts = _rank_by_frequency_from_files(files, column, cast)
+    from giant.data.scan import ScanRequest, scan_metadata
-    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+
+    if column != "material":
+        raise ValueError(f"build_topn_map_from_files only supports column='material', got {column!r}")
+    result = scan_metadata(files, ScanRequest(material=True))
+    assert result.material is not None
+    class_map, other_members, class_counts = _topn_plus_other_map(result.material, n_classes)
     return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
 
 
@@ -349,16 +333,9 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
     join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for
     those, same convention as elsewhere in this module.
     """
-    counts: dict = {}
+    from giant.data.scan import ScanRequest, scan_metadata
-    for path in files:
+
-        columns = ["pdg"]
+    result = scan_metadata(files, ScanRequest(pooled_pdg=True))
-        has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
+    assert result.pooled_pdg is not None
-        if has_sec:
+    class_map, other_members, class_counts = _topn_plus_other_map(result.pooled_pdg, n_classes)
-            columns.append("sec_pdg_list")
-        df = pd.read_parquet(path, columns=columns)
-        _accumulate_value_counts(counts, df["pdg"], int)
-        if has_sec:
-            exploded = df["sec_pdg_list"].explode().dropna()
-            _accumulate_value_counts(counts, exploded, int)
-    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
     return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)

giant/data/scan.py

@@ -0,0 +1,171 @@
+"""Fused metadata scan over one or more parquet files.
+
+`giant.pipeline.run_setup_stage` needs several distinct frequency summaries
+before training can start — the event-id → row-count index (for the train/val
+split), the pdg/material vocabularies, an optional physics-process count, and
+a pooled pdg count (primary + secondary species, for onehot conditioning).
+Each of those used to be its own full `pd.read_parquet(path, columns=[...])`
+per file (`giant.data.loader`'s old `_rank_by_frequency_from_files` /
+`build_index_maps_from_files` / `build_pdg_topn_map_from_files`) — up to five
+separate reads of the same file. `scan_metadata` answers all of them in one
+`pl.collect_all` per file instead, sharing the file open/decompress cost.
+
+Every requested count comes back keyed by value, as a `ValueStat(count,
+first_seen)`. `first_seen` is the value's row ordinal — file order (as given
+in `files`), then row order within a file — via `row_index_name` on the
+per-file lazy scan plus a running row offset across files. This is what
+`giant.data.loader._topn_plus_other_map`'s frequency-ranking tie-break keys
+on: among equally-frequent values, whichever was scanned first wins its own
+class slot. That is an explicit, documented contract (this module is where
+it's implemented), not an accident of iteration order.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+
+from giant.data.loader import _offset_event_id, event_id_offset
+
+
+@dataclass(frozen=True)
+class ScanRequest:
+    """Which aggregations to compute. Every field defaults off so a caller
+    only pays for what it actually needs."""
+
+    event_index: bool = False
+    pdg: bool = False
+    material: bool = False
+    process: bool = False
+    pooled_pdg: bool = False
+    """pdg ∪ exploded sec_pdg_list — both roles a PDG code plays (primary
+    species and secondary species), pooled into one count per code. See
+    `giant.data.loader.build_pdg_topn_map_from_files`'s docstring for why."""
+
+
+@dataclass(frozen=True)
+class ValueStat:
+    count: int
+    first_seen: int
+
+
+@dataclass
+class MetadataScan:
+    event_index: tuple[np.ndarray, np.ndarray] | None = None
+    """(unique_ids, counts), ids ascending — matches
+    `setup_cache.compute_event_index_from_files`'s return shape."""
+    pdg: dict[int, ValueStat] | None = None
+    material: dict[str, ValueStat] | None = None
+    process: dict[str, ValueStat] | None = None
+    pooled_pdg: dict[int, ValueStat] | None = None
+
+
+def _group_lazy(path: Path, column: str) -> pl.LazyFrame:
+    return (
+        pl.scan_parquet(path, row_index_name="__row")
+        .select(column, "__row")
+        .group_by(column)
+        .agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row"))
+    )
+
+
+def _pooled_pdg_lazy(path: Path, has_sec_pdg_list: bool) -> pl.LazyFrame:
+    lf = pl.scan_parquet(path, row_index_name="__row")
+    parts = [lf.select(pl.col("pdg").alias("__val"), "__row")]
+    if has_sec_pdg_list:
+        parts.append(lf.select(pl.col("sec_pdg_list").alias("__val"), "__row").explode("__val").drop_nulls("__val"))
+    combined = pl.concat(parts)
+    return combined.group_by("__val").agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row"))
+
+
+def _merge_counts(acc: dict, df: pl.DataFrame, column: str, row_offset: int, cast) -> None:
+    for key, count, first_row in zip(
+        df[column].to_list(), df["__count"].to_list(), df["__first_row"].to_list(), strict=True
+    ):
+        key = cast(key)
+        first_seen = row_offset + int(first_row)
+        if key in acc:
+            prev_count, prev_first = acc[key]
+            acc[key] = (prev_count + int(count), min(prev_first, first_seen))
+        else:
+            acc[key] = (int(count), first_seen)
+
+
+def scan_metadata(files: list[Path], request: ScanRequest) -> MetadataScan:
+    """Scan `files` once (one `pl.collect_all` per file) and return every
+    aggregation `request` asks for. Files with zero rows contribute nothing
+    but still advance nothing (no row_offset change, nothing to merge)."""
+    event_id_parts: list[tuple[np.ndarray, np.ndarray]] = []
+    pdg_acc: dict[int, tuple[int, int]] = {}
+    material_acc: dict[str, tuple[int, int]] = {}
+    process_acc: dict[str, tuple[int, int]] = {}
+    pooled_pdg_acc: dict[int, tuple[int, int]] = {}
+
+    row_offset = 0
+    for file_idx, path in enumerate(files):
+        keys: list[str] = []
+        lazies: list[pl.LazyFrame] = []
+
+        if request.event_index:
+            keys.append("event_id")
+            lazies.append(_group_lazy(path, "event_id"))
+        if request.pdg:
+            keys.append("pdg")
+            lazies.append(_group_lazy(path, "pdg"))
+        if request.material:
+            keys.append("material")
+            lazies.append(_group_lazy(path, "material"))
+        if request.process:
+            keys.append("process")
+            lazies.append(_group_lazy(path, "process"))
+        if request.pooled_pdg:
+            has_sec = "sec_pdg_list" in pl.scan_parquet(path).collect_schema().names()
+            keys.append("pooled_pdg")
+            lazies.append(_pooled_pdg_lazy(path, has_sec))
+
+        keys.append("__n")
+        lazies.append(pl.scan_parquet(path).select(pl.len().alias("__n")))
+
+        results = dict(zip(keys, pl.collect_all(lazies, engine="streaming"), strict=True))
+        n_rows = int(results["__n"].item()) if len(results["__n"]) else 0
+
+        if request.event_index:
+            df = results["event_id"]
+            ids = _offset_event_id(df["event_id"].to_numpy(), event_id_offset(file_idx))
+            counts = df["__count"].to_numpy().astype(np.int64)
+            if ids.size:
+                event_id_parts.append((ids, counts))
+        if request.pdg:
+            _merge_counts(pdg_acc, results["pdg"], "pdg", row_offset, int)
+        if request.material:
+            _merge_counts(material_acc, results["material"], "material", row_offset, str)
+        if request.process:
+            _merge_counts(process_acc, results["process"], "process", row_offset, str)
+        if request.pooled_pdg:
+            _merge_counts(pooled_pdg_acc, results["pooled_pdg"], "__val", row_offset, int)
+
+        row_offset += n_rows
+
+    event_index = None
+    if request.event_index:
+        if event_id_parts:
+            all_ids = np.concatenate([p[0] for p in event_id_parts])
+            all_counts = np.concatenate([p[1] for p in event_id_parts])
+            order = np.argsort(all_ids, kind="stable")
+            event_index = (all_ids[order], all_counts[order])
+        else:
+            event_index = (np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64))
+
+    def _to_stats(acc: dict) -> dict:
+        return {k: ValueStat(*v) for k, v in acc.items()}
+
+    return MetadataScan(
+        event_index=event_index,
+        pdg=_to_stats(pdg_acc) if request.pdg else None,
+        material=_to_stats(material_acc) if request.material else None,
+        process=_to_stats(process_acc) if request.process else None,
+        pooled_pdg=_to_stats(pooled_pdg_acc) if request.pooled_pdg else None,
+    )

giant/data/setup_cache.py

@@ -23,7 +23,7 @@ import numpy as np
 
 from giant import config
 from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
-from giant.data.loader import TopNMap, event_id_offset, load_event_ids
+from giant.data.loader import TopNMap
 from giant.data.transforms import Normalizer, sorted_membership
 
 # Bump manually on a change to the data-encoding semantics (e.g. a future
@@ -349,12 +349,21 @@ def save(
 
 
 def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]:
-    """Unique event ids + per-event row (step) counts, across all `files`."""
+    """Unique event ids + per-event row (step) counts, across all `files`.
+
+    Computed via a streaming per-file `group_by("event_id")` (see
+    `giant.data.scan.scan_metadata`) rather than concatenating every row's
+    raw event_id across every file before `np.unique` — the latter's peak
+    memory is 8 bytes x total row count; this is bounded by the (much
+    smaller) unique event count instead.
+    """
+    from giant.data.scan import ScanRequest, scan_metadata
+
     if not files:
         return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
-    all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
+    result = scan_metadata(files, ScanRequest(event_index=True))
-    unique_ids, counts = np.unique(all_ids, return_counts=True)
+    assert result.event_index is not None
-    return unique_ids, counts
+    return result.event_index
 
 
 def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:

giant/geometry.py

@@ -28,7 +28,7 @@ from pathlib import Path
 from typing import Any, Iterable
 
 import numpy as np
-import pandas as pd
+import polars as pl
 import pyarrow.parquet as pq
 
 _INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
@@ -301,14 +301,18 @@ def _fit_slab_lookup(
     edges = np.linspace(z_min, z_max, n_bins + 1)
     bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
 
+    # pandas' groupby(...).size() sorts group keys ascending by default, so
+    # a tie in `n` for the same bin (equal counts split between two
+    # material/layer_id combos) resolves to the lexicographically-first
+    # combo — matched here by sorting on the keys first, then a
+    # maintain_order-stable sort on `n` so ties keep that key order.
     counts = (
-        pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
+        pl.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
-        .groupby(["bin", "material", "layer_id"])
+        .group_by(["bin", "material", "layer_id"])
-        .size()
+        .agg(pl.len().alias("n"))
-        .to_frame("n")
+        .sort(["bin", "material", "layer_id"])
-        .reset_index()
+        .sort("n", descending=True, maintain_order=True)
-        .sort_values("n", ascending=False)
+        .unique(subset="bin", keep="first", maintain_order=True)
-        .drop_duplicates("bin")
     )
 
     bin_material = np.full(n_bins, "", dtype=object)

giant/pipeline.py

@@ -15,14 +15,12 @@ from giant.constants import (
 from giant.data import setup_cache
 from giant.data.loader import (
     TopNMap,
+    _topn_plus_other_map,
     event_id_offset,
     find_parquet_files,
     iter_file_chunks,
-    build_index_maps_from_files,
-    build_pdg_topn_map_from_files,
-    build_process_map_from_files,
-    build_topn_map_from_files,
 )
+from giant.data.scan import MetadataScan, ScanRequest, scan_metadata
 from giant.data.transforms import (
     Normalizer,
     build_features,
@@ -127,12 +125,71 @@ def run_setup_stage(
             loaded = setup_cache.load(data, files, echo=echo)
         cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
 
-    if cache is not None and cache.event_index is not None:
+    # Every section below first asks the cache; whatever's missing is
+    # collected into one ScanRequest and answered by a single fused scan
+    # (giant.data.scan.scan_metadata), instead of a separate full pass per
+    # section (event index, vocab, process counts, pdg/material top-N counts
+    # used to each re-open and re-read every file on their own).
+    particle_cfg = cfg["conditioning"]["particle"]
+    material_cfg = cfg["conditioning"]["material"]
+    particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
+    particle_type_target = particle_type_cfg.target
+
+    # A process map is needed if either stage's router reads the physics
+    # process label (type="process"). Only one map is built even if both
+    # stages want one — see the module-level note in giant/cli.py's
+    # _router_total_experts for why composed-router n_experts isn't a plain
+    # int; process routers are never composed in practice, so this doesn't
+    # need that generality.
+    process_router_cfg = next(
+        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
+        None,
+    )
+    process_n_experts = process_router_cfg["n_experts"] if process_router_cfg is not None else None
+
+    need_pdg_onehot = particle_cfg["type"] == "onehot"
+    need_sec_type_onehot = particle_type_target == "onehot"
+    sec_type_n_classes = (
+        resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) if need_sec_type_onehot else None
+    )
+    need_material_onehot = material_cfg["type"] == "onehot"
+    material_n_classes = material_cfg["emb_dim"] if need_material_onehot else None
+
+    def _topn_cached(axis: str, n_classes: int) -> TopNMap | None:
+        return cache.topn_maps.get(setup_cache.topn_key(axis, n_classes)) if cache is not None else None
+
+    need_event_index = cache is None or cache.event_index is None
+    need_vocab = cache is None or cache.vocab is None
+    need_process = process_n_experts is not None and (cache is None or cache.proc_maps.get(process_n_experts) is None)
+    # The PDG axis is used independently by conditioning.particle.type="onehot"
+    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
+    # (secondary-species decode) — their class counts can now differ (gitea
+    # #29: stage2_model.particle_type.n_classes, 0 = inherit
+    # conditioning.particle.emb_dim), but both are built from the same
+    # pooled pdg-count scan, so a cache miss on either one asks for it.
+    need_pdg_pooled = (need_pdg_onehot and _topn_cached("pdg", particle_cfg["emb_dim"]) is None) or (
+        need_sec_type_onehot and sec_type_n_classes is not None and _topn_cached("pdg", sec_type_n_classes) is None
+    )
+    need_material_topn = (
+        need_material_onehot and material_n_classes is not None and _topn_cached("material", material_n_classes) is None
+    )
+
+    request = ScanRequest(
+        event_index=need_event_index,
+        pdg=need_vocab,
+        material=need_vocab or need_material_topn,
+        process=need_process,
+        pooled_pdg=need_pdg_pooled,
+    )
+    scan = scan_metadata(files, request) if request != ScanRequest() else MetadataScan()
+
+    if not need_event_index:
         unique_ids, counts = cache.event_index
         echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
     else:
         echo("scanning event IDs …")
-        unique_ids, counts = setup_cache.compute_event_index_from_files(files)
+        assert scan.event_index is not None
+        unique_ids, counts = scan.event_index
         if cache is not None:
             cache.event_index = (unique_ids, counts)
 
@@ -141,55 +198,34 @@ def run_setup_stage(
     n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
     echo(f"  {int(counts.sum()):,} steps  |  {len(train_events)} train events  |  {len(val_events)} val events")
 
-    if cache is not None and cache.vocab is not None:
+    if not need_vocab:
         pdg_map, mat_map = cache.vocab
         echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
     else:
         echo("building vocabulary maps …")
-        pdg_map, mat_map = build_index_maps_from_files(files)
+        assert scan.pdg is not None and scan.material is not None
+        pdg_map = {v: i for i, v in enumerate(sorted(scan.pdg))}
+        mat_map = {v: i for i, v in enumerate(sorted(scan.material))}
         echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
         if cache is not None:
             cache.vocab = (pdg_map, mat_map)
 
-    # A process map is needed if either stage's router reads the physics
-    # process label (type="process"). Only one map is built even if both
-    # stages want one — see the module-level note in giant/cli.py's
-    # _router_total_experts for why composed-router n_experts isn't a plain
-    # int; process routers are never composed in practice, so this doesn't
-    # need that generality.
     proc_map: dict[str, int] | None = None
-    process_router_cfg = next(
-        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
-        None,
-    )
     if process_router_cfg is not None:
-        n_experts = process_router_cfg["n_experts"]
+        assert process_n_experts is not None
-        cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
+        if not need_process:
-        if cached_proc_map is not None:
+            assert cache is not None
+            cached_proc_map = cache.proc_maps.get(process_n_experts)
+            assert cached_proc_map is not None
             proc_map = cached_proc_map
-            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
+            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {process_n_experts} experts)")
         else:
             echo("building process vocabulary …")
-            proc_map = build_process_map_from_files(files, n_experts=n_experts)
+            assert scan.process is not None
-            echo(f"  {len(proc_map)} process labels mapped to {n_experts} experts")
+            proc_map, _, _ = _topn_plus_other_map(scan.process, process_n_experts)
+            echo(f"  {len(proc_map)} process labels mapped to {process_n_experts} experts")
             if cache is not None:
-                cache.proc_maps[n_experts] = proc_map
+                cache.proc_maps[process_n_experts] = proc_map
-
-    # Top-N-plus-other maps for onehot conditioning/type axes.
-    # The PDG axis is used independently by conditioning.particle.type="onehot"
-    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
-    # (secondary-species decode) — their class counts can now differ (gitea
-    # #29: stage2_model.particle_type.n_classes, 0 = inherit
-    # conditioning.particle.emb_dim), so each is resolved and built
-    # independently via _pdg_topn below. cache.topn_maps is keyed by
-    # (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
-    # the same N the second call is a cache hit against the first — no extra
-    # scan in the common case where they still match. The material axis is
-    # independent of both.
-    particle_cfg = cfg["conditioning"]["particle"]
-    material_cfg = cfg["conditioning"]["material"]
-    particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
-    particle_type_target = particle_type_cfg.target
 
     def _pdg_topn(n_classes: int) -> TopNMap:
         cache_key = setup_cache.topn_key("pdg", n_classes)
@@ -198,33 +234,36 @@ def run_setup_stage(
             echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
             return cached
         echo("building pdg top-N map …")
-        topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
+        assert scan.pooled_pdg is not None
+        class_map, other_members, class_counts = _topn_plus_other_map(scan.pooled_pdg, n_classes)
+        topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
         echo(f"  {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
         if cache is not None:
             cache.topn_maps[cache_key] = topn_map
         return topn_map
 
-    pdg_topn_map: TopNMap | None = None
+    pdg_topn_map: TopNMap | None = _pdg_topn(particle_cfg["emb_dim"]) if need_pdg_onehot else None
-    if particle_cfg["type"] == "onehot":
-        pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
-
     sec_type_topn_map: TopNMap | None = None
-    if particle_type_target == "onehot":
+    if need_sec_type_onehot:
-        sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+        assert sec_type_n_classes is not None
         sec_type_topn_map = _pdg_topn(sec_type_n_classes)
 
     mat_topn_map: TopNMap | None = None
-    if material_cfg["type"] == "onehot":
+    if need_material_onehot:
-        n_classes = material_cfg["emb_dim"]
+        assert material_n_classes is not None
-        cache_key = setup_cache.topn_key("material", n_classes)
+        cache_key = setup_cache.topn_key("material", material_n_classes)
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
             mat_topn_map = cached
-            echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
+            echo(
+                f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {material_n_classes} classes)"
+            )
         else:
             echo("building material top-N map …")
-            mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
+            assert scan.material is not None
-            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
+            class_map, other_members, class_counts = _topn_plus_other_map(scan.material, material_n_classes)
+            mat_topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
+            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {material_n_classes} classes")
             if cache is not None:
                 cache.topn_maps[cache_key] = mat_topn_map
 
@@ -456,17 +495,30 @@ def run_train_job(
     )
 
     pin = device.type == "cuda"
+    # DataLoader worker subprocesses default to fork() on Linux, but by the
+    # time they're created this process has already run polars queries
+    # (run_setup_stage's fused metadata scan, above) — polars' native
+    # (rayon) thread pool doesn't survive a fork: a worker that inherits it
+    # mid-fork deadlocks the instant it touches polars itself, which
+    # StreamingStepsDataset's iter_file_chunks now does on every row group.
+    # "spawn" starts each worker as a fresh interpreter with no inherited
+    # thread-pool state, avoiding that hazard entirely. Only matters when
+    # workers actually exist — num_workers=0 runs the dataset in-process and
+    # never forks.
+    mp_context = "spawn" if num_workers > 0 else None
     train_loader = DataLoader(
         train_ds,
         batch_size=None,
         num_workers=num_workers,
         pin_memory=pin,
+        multiprocessing_context=mp_context,
     )
     val_loader = DataLoader(
         val_ds,
         batch_size=None,
         num_workers=num_workers,
         pin_memory=pin,
+        multiprocessing_context=mp_context,
     )
 
     model_config = {

giant/tools/profile_setup_scan.py

@@ -0,0 +1,162 @@
+"""Benchmark `giant.pipeline.run_setup_stage`'s cold-cache scan against synthetic data.
+
+Generates a schema-complete synthetic steps parquet (matching
+`tests/test_pipeline.py`'s `_make_synthetic_steps`, but built with vectorized
+numpy instead of a per-row Python loop so it scales to millions of rows) at a
+few row counts, times `run_setup_stage` with `cache_setup=False` (so every
+call is a genuine cold scan, never served from the sidecar), and prints a
+before/after-style table. Run this on `master` before a change and again
+after to see what a step actually bought — see the "speed up dwarf
+warm-cache" plan for the pass-by-pass breakdown this benchmark is meant to
+attribute (giant/data/loader.py, giant/data/scan.py, giant/pipeline.py).
+
+Usage: ``uv run python giant/tools/profile_setup_scan.py``
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import polars as pl
+
+from giant import config as gconfig
+from giant.pipeline import run_setup_stage
+
+ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]
+
+_MATERIALS = ["G4_AIR", "G4_Fe"]
+_PDGS = [11, 22]
+_PROCESSES = ["eIoni", "phot", "compt"]
+
+
+def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray:
+    v = rng.normal(size=(n, 3))
+    return v / np.linalg.norm(v, axis=1, keepdims=True)
+
+
+def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float) -> list[list[float]]:
+    total = int(k.sum())
+    flat = rng.uniform(lo, hi, size=total)
+    idx = np.cumsum(k)[:-1]
+    return [arr.tolist() for arr in np.split(flat, idx)]
+
+
+def _make_synthetic_steps(n: int, seed: int = 0) -> pl.DataFrame:
+    """Vectorized equivalent of tests/test_pipeline.py's `_make_synthetic_steps`.
+
+    event_id is assigned so each event gets 2-3 steps (matching that
+    fixture's structure), and pdg/material/process cycle deterministically
+    by row index rather than being drawn at random, same as the original.
+    """
+    rng = np.random.default_rng(seed)
+    n_events = max(n // 3, 1)
+
+    pre_E = rng.uniform(50.0, 500.0, size=n)
+    n_sec = rng.integers(0, 3, size=n)
+    frac_dep = rng.uniform(0.05, 0.3, size=n)
+    frac_sec = np.where(n_sec > 0, rng.uniform(0.05, 0.2, size=n), 0.0)
+    frac_post = 1.0 - frac_dep - frac_sec
+    edep = pre_E * frac_dep
+    e_sec = pre_E * frac_sec
+    post_E = pre_E * frac_post
+    pre_pos = rng.uniform(-10, 10, size=(n, 3))
+    step_length = rng.uniform(0.1, 5.0, size=n)
+    pre_dir = np.zeros((n, 3))
+    pre_dir[:, 2] = 1.0
+    post_dir = _unit_vectors(n, rng)
+    post_pos = pre_pos + step_length[:, None] * pre_dir
+
+    row_idx = np.arange(n)
+    event_id = row_idx % n_events
+
+    sec_E = _ragged_lists(n_sec, rng, 0.1, 1.0)  # placeholder magnitude, rescaled below
+    sec_dx = _ragged_lists(n_sec, rng, -1.0, 1.0)
+    sec_dy = _ragged_lists(n_sec, rng, -1.0, 1.0)
+    sec_dz = _ragged_lists(n_sec, rng, -1.0, 1.0)
+    total_sec = int(n_sec.sum())
+    flat_pdg = [_PDGS[(row_idx[i] + j) % 2] for i in range(n) for j in range(n_sec[i])]
+    idx = np.cumsum(n_sec)[:-1]
+    sec_pdg = (
+        [list(x) for x in np.split(np.array(flat_pdg, dtype=np.int64), idx)] if total_sec else [[] for _ in range(n)]
+    )
+    # Rescale each row's secondary energies to sum to that row's e_sec (a
+    # Dirichlet split, like the original fixture) rather than the raw
+    # uniform placeholder.
+    sec_E_scaled = []
+    for i in range(n):
+        vals = np.array(sec_E[i])
+        if vals.size:
+            sec_E_scaled.append((vals / vals.sum() * e_sec[i]).tolist())
+        else:
+            sec_E_scaled.append([])
+
+    return pl.DataFrame(
+        {
+            "event_id": event_id,
+            "pdg": np.array(_PDGS)[row_idx % 2],
+            "pre_x": pre_pos[:, 0],
+            "pre_y": pre_pos[:, 1],
+            "pre_z": pre_pos[:, 2],
+            "pre_E": pre_E,
+            "pre_dx": pre_dir[:, 0],
+            "pre_dy": pre_dir[:, 1],
+            "pre_dz": pre_dir[:, 2],
+            "material": np.array(_MATERIALS)[row_idx % 2],
+            "layer_id": row_idx % 5,
+            "child_track_ids": [list(range(int(k))) for k in n_sec],
+            "e_sec": e_sec,
+            "process": np.array(_PROCESSES)[row_idx % 3],
+            "step_length": step_length,
+            "post_E": post_E,
+            "edep": edep,
+            "post_dx": post_dir[:, 0],
+            "post_dy": post_dir[:, 1],
+            "post_dz": post_dir[:, 2],
+            "post_x": post_pos[:, 0],
+            "post_y": post_pos[:, 1],
+            "post_z": post_pos[:, 2],
+            "sec_E_list": sec_E_scaled,
+            "sec_pdg_list": sec_pdg,
+            "sec_dx_list": sec_dx,
+            "sec_dy_list": sec_dy,
+            "sec_dz_list": sec_dz,
+        }
+    )
+
+
+def _time_setup_stage(data: Path) -> float:
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {})
+    gconfig.validate_config(cfg)
+    t0 = time.perf_counter()
+    run_setup_stage(
+        data,
+        val_fraction=cfg["train"]["val_fraction"],
+        seed=cfg["train"]["seed"],
+        cfg=cfg,
+        cache_setup=False,
+        echo=lambda *a, **k: None,
+    )
+    return time.perf_counter() - t0
+
+
+def main() -> None:
+    with TemporaryDirectory(prefix="giant-setup-scan-profile-") as tmp:
+        tmp_path = Path(tmp)
+        print(f"{'n_rows':>10s}  {'time (s)':>10s}  {'rows/s':>12s}")
+        for n in ROW_COUNTS:
+            path = tmp_path / f"steps_{n}.parquet"
+            _make_synthetic_steps(n).write_parquet(path)
+            # warm the OS page cache so the timed pass measures compute, not
+            # the one-time cold read of a freshly-written file.
+            pl.scan_parquet(path).select(pl.len()).collect()
+
+            dt = _time_setup_stage(path)
+            print(f"{n:>10,d}  {dt:>10.3f}  {n / dt:>12,.0f}")
+            path.unlink()
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "giant"
-version = "0.3.16"
+version = "0.3.17"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "numpy>=1.26,<3",
-    "pandas>=2.2,<4",
+    "polars>=1.0,<2",
     "pyarrow>=16,<25",
     "tqdm>=4.60,<5",
     "typer>=0.12,<1",
@@ -28,6 +28,10 @@ dev = [
     "ty>=0.0.50,<0.1",
     "bump-my-version>=1.2,<2",
     "git-cliff>=2,<3",
+    # Only used by test fixtures (writing small parquet files) — not a
+    # runtime dependency of giant itself since the pandas -> polars
+    # data-loading rewrite.
+    "pandas>=2.2,<4",
     "giant[convert,analysis,geometry,wandb]",
 ]
 geometry = [

tests/test_loader.py

@@ -108,12 +108,12 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path):
 
 def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path):
     """When two processes end up with equal total counts, ranking falls back
-    to whichever was accumulated first (`sorted(..., reverse=True)` is stable,
+    to whichever was scanned first — file order, then row order within a
-    and `counts` is built in file/row-scan order) — this is implementation-
+    file (`giant.data.scan`'s `first_seen` ordinal, ranked by
-    defined, not a documented contract, so pin it explicitly: a future
+    `giant.data.loader._topn_plus_other_map`'s `(-count, first_seen)` key).
-    rewrite (e.g. a polars-based single-scan) that ties differently would
+    This is an explicit, documented contract (not an accident of iteration
-    silently reshuffle which processes get their own expert slot across a
+    order), pinned here so a future change to the ranking can't silently
-    retrain, and this test is what should catch that."""
+    reshuffle which processes get their own expert slot across a retrain."""
     path = tmp_path / "a.parquet"
     pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)
 
@@ -233,6 +233,23 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
     assert m.class_counts == {0: 11, 1: 5}
 
 
+def test_build_pdg_topn_map_from_files_pooled_tie_breaks_by_row_position(tmp_path):
+    """Pooled pdg counting merges the primary `pdg` column and the exploded
+    `sec_pdg_list` column via one `group_by` over both (see
+    `giant.data.scan._pooled_pdg_lazy`), keyed by row position regardless of
+    which role (primary or secondary) a code was seen in — not "all
+    primaries before all secondaries" the way a two-pass accumulation would.
+    11 (primary, row 0), 33 (primary, row 1) and 22 (secondary, row 1) all
+    end up with count 1; 11's strictly earlier row wins the tie over both,
+    whatever order 33/22 (tied with each other, same row) land in."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [11, 33], "sec_pdg_list": [[], [22]]}).to_parquet(path)
+
+    m = build_pdg_topn_map_from_files([path], n_classes=4)
+
+    assert m.class_map[11] == 0
+
+
 def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
     """Files predating the parent->child join have no sec_pdg_list column —
     must not raise, just count the primary pdg column alone."""

tests/test_pipeline.py

@@ -142,7 +142,7 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
     def _forbidden(*a, **k):
         raise AssertionError("should be served from cache, not recomputed")
 
-    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+    monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden)
     monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden)
 
     echo2 = _run(data, tmp_path / "out2")
@@ -283,7 +283,7 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
     def _forbidden(*a, **k):
         raise AssertionError("vocab should be served from cache")
 
-    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+    monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden)
 
     echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3))
     joined = "\n".join(echo2)

uv.lock

@@ -675,12 +675,12 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.16"
+version = "0.3.17"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },
-    { name = "pandas" },
     { name = "particle" },
+    { name = "polars" },
     { name = "pyarrow" },
     { name = "pyyaml" },
     { name = "tqdm" },
@@ -712,6 +712,7 @@ dev = [
     { name = "git-cliff" },
     { name = "ipykernel" },
     { name = "matplotlib" },
+    { name = "pandas" },
     { name = "plotstyle" },
     { name = "polars" },
     { name = "pytest" },
@@ -738,9 +739,10 @@ requires-dist = [
     { name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
     { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
     { name = "numpy", specifier = ">=1.26,<3" },
-    { name = "pandas", specifier = ">=2.2,<4" },
+    { name = "pandas", marker = "extra == 'dev'", specifier = ">=2.2,<4" },
     { name = "particle", specifier = ">=1.0,<2" },
     { name = "plotstyle", marker = "extra == 'analysis'", specifier = ">=1.0.0", index = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" },
+    { name = "polars", specifier = ">=1.0,<2" },
     { name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" },
     { name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
     { name = "pyarrow", specifier = ">=16,<25" },