fix(nanoevents): correct parquet entry_start jagged reads and cache column reads

[NJManganelli]

Part of #1578 — critical bug 4 (parquet reads with entry_start > 0 return shifted jagged data) and the performance bullet (parquet mapping reads the entire column per buffer access with no cache). Two commits, one per bullet.

Correctness: reading a jagged parquet column with entry_start > 0 returned silently-corrupted data. The mapping sliced the column with pyarrow (which records a logical offset rather than copying buffers) and extracted the content offset-aware via flatten(), but read the list offsets from the start of the offsets buffer, ignoring aspa.offset. Per-event values (e.g. Muon.pt) were reassigned to the wrong events, and NanoAOD cross-references crashed with a broadcast length error (reproduced with entry_start=5). The fix reads len(aspa)+1 offsets starting at aspa.offset and rebases them to zero, for both ListArray and LargeListArray; any entry_start/entry_stop subrange now equals the full-read slice exactly.

Performance: each buffer-key access re-read the entire column from the file, so a jagged column (offsets + content keys) was read at least twice. A small per-source LRU cache keyed by column name collapses these into a single read with no public-API change (measured 2 → 1 reads per jagged column via the raw from_buffers path; 1599 → 1516 total reads in the standard eager path). A true partial row-range read would require row-group iteration and a larger refactor, noted in the commit body.

Regression tests cover subrange-vs-full-slice equality (2 encodings × 4 ranges, jagged + flat branches) and a monkeypatched read-count assertion; both fail before their respective commits and pass after. tests/test_nanoevents.py (28) and tests/test_local_executors.py -k parquet (98) pass; pre-commit clean.

🤖 Generated with Claude Code

[lgray]

Seems to be a useful performance optimization and correctness improvement for reading parquet files, will catch bugs in the future that we haven't hit yet due to low uptake of parquet across all analyses (and usually it's fairly flat tables).

[lgray]

🤖 AI text below 🤖

This PR fixes a real, silent data-corruption bug: jagged parquet reads with entry_start > 0 read list offsets from the start of the offsets buffer while flatten() returned offset-aware content, reassigning per-event values to the wrong events (this hits the production Runner parquet path too — every chunk after the first, via executor.py:1657-1663 → from_parquet(mode="virtual", entry_start=...)). I verified the offset arithmetic against pyarrow's buffer semantics (raw_offsets[aspa.offset : aspa.offset + len(aspa) + 1] is exactly the n+1 offsets for the logical slice, and the rebase matches what flatten() returns), confirmed both new tests fail before their respective commits and pass after (merge-base: 7/9 fail with the described broadcast ValueError; commit-1-only: cache test fails with assert 2 == 1), probed the cache with back-to-back different subranges through one shared source (no poisoning — reads are always full-column, sliced after, so keying by column name alone is sound), and ran the full test_nanoevents.py (28 passed) plus test_local_executors.py -k parquet (98 passed). Virtual mode (the executor path) also verified correct. The branch merges cleanly onto current master. LGTM overall — one test-coverage gap worth closing before merge.

Should fix

• The int32 ListArray branch of the fix is untested. Both sample files (nano_dy.parquet and nano_dy.extensionarray.parquet) decode to LargeListArray (int64 offsets, verified by inspecting the decoded chunks), so the two parametrized "encodings" both exercise only the numpy.int64 side of the dtype selection at src/coffea/nanoevents/mapping/parquet.py:143-147 — the numpy.int32 path never runs in CI. I probed it manually (synthetic parquet written with pa.list_(pa.float32()) decodes as ListArray; all four subranges correct), so the code is right, but a regression there would currently go unnoticed. Suggest a small tmp_path fixture writing a pa.list_() (non-large) column and running it through the same subrange-vs-full-slice assertion.

Nits / Optional

• src/coffea/nanoevents/mapping/parquet.py:134 — offsets = None is dead now that both branches were unified; it's immediately reassigned at line 149. Leftover from the old structure; drop it.
• The subrange tests parametrize only mode="eager". Virtual mode goes through the identical extraction code (base.py:170-172 only defers the call) and I verified it passes, but the production executor path uses mode="virtual" — parametrizing mode over both is nearly free and would pin the path users actually hit.

Test coverage: The new tests are non-vacuous and discriminating, proven per commit by revert runs (subrange tests: 6/8 fail at merge-base, exactly the entry_start=0 cases passing; read-count test fails at both merge-base and commit-1-only with assert 2 == 1). Coverage spans jagged + flat branches, boundary ranges including start=0, full range, and a 3-event tail; the read-count test monkeypatches the stable pq.ParquetFile.read seam, which is reasonably robust. Gaps: the int32 offsets branch (above) and eager-only mode. The claimed "2 encodings" are plain-vs-extensionarray, not int32-vs-int64 — worth being aware of. Side note: on pyarrow ≥ 24, ParquetFile.read returns single-chunk columns even for multi-row-group files, so the pre-existing .chunk(0) at parquet.py:127 doesn't truncate here; true row-group-level partial reads remain deferred as the PR states.

Pydantic: the diff touches no pydantic code (grep of the full diff: no matches) — not applicable.

[ikrommyd]

I'm pretty sure the same happens in root files with base schema too but uproot's cache protects us I think from double reading. Should be investigated to make sure though. I also have a vague memory of adding a cache of size 1 somewhere but it may have been in awkward for a different reason. Just saying this as it should be investigated to make sure my assumptions are right.

[NJManganelli]

🤖 Opus 4.8 Output 🤖

Investigated and verified empirically. ikrommyd's assumption is correct, and the fix is in the right place for each format — no ROOT-side change is needed.

What I found

I built a genuinely jagged single ROOT branch (a std::vector-like ListOffset column — the base-schema case ikrommyd meant, not NanoAOD where offsets live in a separate nMuon counts branch and there's no double-read at all) and traced it through coffea's UprootSourceMapping in virtual mode with a spy on uproot's array_cache:

EXTRACT_COLUMN #1 (jag) ← offsets buffer key
MISS ...0:0 → SET ...0:0 physical read + decompress
MISS ...0:1 → SET ...0:1 physical read + decompress
EXTRACT_COLUMN #2 (jag) ← content buffer key
HIT ...0:0
HIT ...0:1 100% served from uproot's array_cache — no disk/decompress

So the double request does happen on the ROOT path (coffea calls extract_column once per buffer key — offsets and content), exactly as ikrommyd suspected. But uproot's built-in array_cache (a LRUArrayCache, default 100 MB, present on every default uproot.open — confirmed) serves the second call entirely from memory, so there is no double physical read/decompress. Coffea's own LRUCache(1) source cache in base.py is what keeps the opened file (and thus its array_cache) alive across the two buffer-key reads — likely the "cache of size 1" ikrommyd half-remembered, though it caches the source, not arrays.

This is precisely why parquet was different and needed our fix: the parquet mapping had no uproot-equivalent cache, so each buffer access called ParquetFile.read() and physically re-read the whole column. PR #1583's per-source column cache is the ROOT-array_cache analog, added in the one place parquet lacked it.

Answer to "addressed in the appropriate place?"

Yes:

• ROOT / base schema — nothing to add; uproot's array_cache already dedupes the double request at the physical-read level (verified: call Add .csv and .json interpretation to weights extractor #2 is all cache hits). The only way to lose that is opening with array_cache=None, which is uproot's documented knob and not the default.
• Parquet — the appropriate place was the parquet mapping itself, which is exactly where fix(nanoevents): correct parquet entry_start jagged reads and cache column reads #1583 put the per-source column cache.

[nsmith-]

good bot: yeah, there's a LRUCache(1) in the nanoevents mapping to keep the file alive

[NJManganelli]

This fixes a clear potential bug, which hopefully not many users have hit due to not frequently exercising parquet (and ~zero usage through NanoEvents due to broken mapping for such a long time). Someone else please also have a look @lgray @nsmith- @ikrommyd