diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py index d71b129..2248c90 100644 --- a/giant/analysis/sources.py +++ b/giant/analysis/sources.py @@ -106,18 +106,27 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame: can push their own narrow projection into the parquet read — the single biggest lever on a larger-than-RAM file. ``pl.LazyFrame`` inputs pass straight through (used by tests). + + ``pdg`` is cast to a canonical ``Int64`` here: the rollout writer and the + 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). """ if isinstance(source, pl.LazyFrame): - return source + return source.with_columns(pl.col("pdg").cast(pl.Int64)) path = Path(source) if side is Side.rollout: _check_rollout_metadata(path) - return pl.scan_parquet(path) - # The reference (a rollout's seed `dataset`) may be a directory of parquet - # shards rather than a single file — scan them all. - if path.is_dir(): - return pl.scan_parquet(str(path / "**/*.parquet")) - return pl.scan_parquet(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) + ) + return lf.with_columns(pl.col("pdg").cast(pl.Int64)) def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame: