giant/giant/data/scan.py

"""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", empty_as_null=False)
            .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,
    )