perf: build the RNTuple form once and prune it per field - #1723
Open
kadircanyildirm-crypto wants to merge 2 commits into
Open
kadircanyildirm-crypto wants to merge 2 commits into
kadircanyildirm-crypto wants to merge 2 commits into
Conversation
`to_akform` rebuilt the form from the field records on every call. Reading a wide RNTuple one column at a time, which is what coffea does, therefore rebuilt the whole form once per column: 1579 builds for the 1491-field NanoAOD skim in scikit-hep#1637, and 252 s for a read that takes 4 s from the same data in a TTree. Build the unfiltered form once per RNTuple, keep it there, and serve later requests by pruning that form to the requested keys. A field takes the same cached form, walks down to its own path and prunes from there, so the per-column path comes off the cache as well. For a field under a collection the form is still built around the outermost collection on its path and `field_path` is returned as before; `ak_add_doc` still takes the original path. The read is now 32-36 s with a single form build, and all 1491 leaf columns are bit-identical to the ones built the old way. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B6b4SK8hGvDHGAx1W6KqGJ
Codecov Report❌ Patch coverage is
❌ Your patch check has failed because the patch coverage (80.89%) is below the target coverage (98.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files
|
codecov flagged the patch at 80.9%. Two of the uncovered paths turned out to be dead: pruning only ever returns "nothing survives" for the top-level call with an empty selection, never for a field below it, because a field is visited only when its own path is in the keep set. Handle the empty selection where it happens, in `to_akform`, and let `_prune_akform` always return a form. `exact` then has no reader left, so it goes too. The rest is covered by tests: an empty selection, the tuple and subfields-of-a-record shapes from the STL-container and nested-struct files, walking to a field that isn't in the form, and the fallback for a form the walk cannot follow. Swept for equivalence while doing it: every field of 11 RNTuple test files plus four filter shapes, 1453 comparisons of pruned against rebuilt, no differences. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B6b4SK8hGvDHGAx1W6KqGJ
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of #1637.
to_akformrebuilt the form from the field records on every call, so reading a wide RNTuple one column at a time — which is what coffea does — rebuilt the whole form once per column. For the 16-event, 1491-field NanoAOD skim in the issue that is 1579 builds.This builds the unfiltered form once per RNTuple, keeps it there, and serves later requests by pruning it to the requested keys. A field takes the same cached form, walks down to its own path and prunes from there, so the per-column path comes off the cache as well — that is the path coffea actually uses, since it asks fields rather than the ntuple. For a field under a collection the form is still built around the outermost collection on its path and
field_pathis returned as before;ak_add_docstill takes the original path.Measured on
main(c979803), Windows, Python 3.12, awkward 2.13.0, coffea 2026.7.0, reading the skim back throughNanoEventsFactory:The same data written as a TTree reads in 3.8–4.5 s, so this does not close the gap to TTree by itself. What is left is a different effect, measured and described in #1637 (comment).
Checks:
ak.to_buffers, compared byte for byte)pytest -k "rntuple and not dask": 146 passed, 0 failed — 143 onmain, plus the 3 added heremain: 1032 passed, 29 failed, 117 skipped either way, the 29 all in the dask test files becausedistributedis missing in my environmentak_add_docDisclosure, per CONTRIBUTING: this was written with Claude Code. The measurements above are from runs on my own machine against this branch.
🤖 Generated with Claude Code
https://claude.ai/code/session_01B6b4SK8hGvDHGAx1W6KqGJ