From 70d0f043260e1f98db4d680e87ed3f7acb6a1a03 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 24 Jul 2026 14:32:06 +0200 Subject: [PATCH] analyze: normalize pdg dtype in open_side to fix rollout/reference concat Rollout output and the reference file's ROOT-derived parquet disagree on pdg's integer width (Int32 vs Int64), which only surfaced downstream as a pl.concat SchemaError in build_context's pdg-count merge. Cast to a canonical Int64 at the single scan entry point instead. --- giant/analysis/sources.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) 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: