Repository navigation
perf: project RNTuple subfields in uproot.dask and uproot.graphed - #1783
Merged
ariostas merged 9 commits intoOct 28, 2026
Merged
Conversation
An RNTuple was projected by top-level field, so reading Muon.pt read every leaf of Muon. A key now covers the subfields beneath it: necessary_columns keeps a mapping's subfield key when a field above it is readable, and _read_tree loads a buffer whose RNTuple key lies under a requested key. The trivial mapping names an RNTuple buffer by its leaf's dotted path (a collection's offsets by its first leaf, a union by its field), and load_buffers reads only the selected leaves. With known_base_form, which opens no file, the trivial mapping still reads whole fields. Assisted-by: ClaudeCode:claude-opus-5-5
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files
|
No key reaches beneath a leaf or a union, so the separate fall-through return was unreachable. Assisted-by: ClaudeCode:claude-opus-5-5
Contributor
Author
|
@ariostas this one is good to go - further incremental fixes I'm running into. |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Offsets-only projection fails when a collection’s first member is an empty record.
1 open issue
What changed in this PR
This PR reduces unnecessary RNTuple reads by projecting nested subfields in uproot.dask and uproot.graphed, rather than reading entire top-level fields.
Changes:
- Map RNTuple buffers to leaf paths and prune reads to requested subfields.
- Support subfield keys from custom form mappings while preserving whole-field defaults when the file type is unknown.
- Add column-count and correctness tests for projected and unprojected reads.
| File | Description |
|---|---|
tests/test_1720_rntuple_nested_fields.py |
Tests subfield projection, column counts, custom mappings, and buffer ownership. |
src/uproot/_graphed.py |
Enables subfield mapping for detected RNTuples. |
src/uproot/_dask.py |
Implements leaf-level mapping, projection pruning, and subfield key coverage. |
Review effort: Balanced
ariostas
reviewed
Oct 27, 2026
ariostas
left a comment
ariostas
left a comment
Member
There was a problem hiding this comment.
@lgray I ran an AI review and the findings are below. I'm happy to push fixes and the alternative design if you're okay with it.
🤖 AI text below 🤖
Bugs
ak.numcrashes when a collection's first field is an empty record (_dask.py:1009, same as Copilot's comment). Forjets: var * {e: {}, x: float64},ak.num(ev.jets, axis=1).compute()raisesFieldNotFoundError: no field 'jets' in record with 0 fields, withopen_files=TrueandFalse. It works onmain.- The first leaf can be expensive (
_dask.py:1009). Forjets: var * {hits: var * float64, pt: float64, eta: float64},ev.jets.ptreads 4 columns, all ofhitsincluded, andak.num(ev.jets)reads 3. - A read without the first leaf fails (
_dask.py:1290). Calling_read_treewithkeys={"jets.eta"}:load_buffersdoes return thejetsoffsets, but the gate maps the offsets buffer tojets.pt, sees it wasn't requested, and puts in a placeholder, so accessingjets.etafails. Through dask the first leaf is always requested, so this stays hidden, but it is why the first-leaf rule is needed.
Inconsistencies (not regressions)
necessary_columnsaccepts subfield keys for any tree type, but_read_treetreats dotted keys as subfields only for RNTuples. A custom mapping that namesjetp4.fEover a TTree (uproot-HZZ-objects.root) now fails withno field 'jetp4.fE'. Onmainit fails too, with a placeholder error.- Subfield mode is decided from the first file only (
_dask.py:2081and2158), so a dataset mixing RNTuple and TTree files would send subfield keys to the TTree partitions. (Not run.) FormMappingInfoWithVirtualArraysinheritssubfields=, but itsbuffer_replacementsstill callsself._form.content(key), which fails for a dotted key. (Not run.)
Simplifications
_select_keysre-implements awkward'sForm.select_columns, which gives the same result for nested records, lists, options and unions.necessary_columnsrebuilds_subfield_paths(self.base_form)on every call, and_select_keysruns for every partition, which adds up on NanoAOD-sized forms.- Per AGENTS.md, the new tests should go in a new file named after this PR (
tests/test_1783_<description>.py) rather than intest_1720_rntuple_nested_fields.py.
Alternative design
Items 1–3 come from one limitation: the gate's rule (a buffer is used only if all the keys it maps to were read) can't express that a collection's offsets come with any field under it. coffea relies on that rule, so it should stay. I prototyped this locally:
- With
subfields=True, each buffer is keyed by its own node's path: thejetsoffsets byjetsand the data byjets.pt. A union is still read whole. There is no first-leaf choice. - One rule for both loading and the gate: a requested key brings everything beneath it and the structure above it.
load_buffersusesself._form.select_columns(keys), and the gate (RNTuple only) accepts a buffer whose key is a requested key, under one, or above one.select_columnskeeps exactly those nodes, so the two can't disagree. TrivialFormMappingInfo.keys_for_buffer_keysdrops a collection key when a field under it is also needed, soev.jets.etaasks for{"jets.eta"}.
Columns read (open_files=True and False give the same numbers):
| access | main | this PR | prototype |
|---|---|---|---|
ak.num(jets), first field {} |
ok | crash | ok |
jets.pt, first field a list |
5 | 4 | 2 |
jets.eta + ak.num(jets) |
5 | 4 | 2 |
only jets.eta requested (item 3) |
— | fails | ok |
ak.num(jets) alone |
5 | 3 | 5 |
The prototype is about 30 lines smaller than this PR and removes the extra leaf per collection mentioned in the description. The 640 dask, graphed and RNTuple tests pass, with two assertions that encode the first-leaf rule updated.
Open points:
ak.num(ev.Muon)with no Muon field used reads the whole collection, as onmain. Turning a lone collection key into one cheap scalar leaf would recover this PR's 2-column case, but it's only safe for the trivial mapping, since a custom mapping may useMuonto mean the whole field. Reading only the offsets would need RNTuple reader support: selecting just the collection field currently reads every field under it.- The gate now accepts a buffer keyed
Muonwhen onlyMuon.ptwas requested. This PR's custom-mapping test passes, but it should be checked against scikit-hep/coffea#1648.
With subfields=True each buffer is keyed by its own node's path (a collection's offsets by the collection, data by the leaf; unions whole). A requested key brings everything under it and the structure above it: load_buffers uses Form.select_columns, and the RNTuple gate accepts a buffer whose key is requested, under one, or above one. keys_for_buffer_keys drops a collection key when a field under it is also needed, so ev.jets.eta reads only the jets offsets and eta. - an empty record as a collection's first field no longer breaks ak.num - reading jets.pt no longer reads the first leaf (e.g. a nested list) - a TTree refuses subfield keys with a clear error; subfield mode is chosen from every opened tree, so open_files=True mixed datasets read whole fields - virtual-array replacements go through load_buffers, so dotted keys work - new tests move to tests/test_1783_rntuple_subfield_projection.py Assisted-by: ClaudeCode:claude-opus-5-5
- load_buffers selects by escaped field-name tuples, so dots, braces and glob characters in a field name are not reinterpreted by select_columns - a dtype mismatch raises TypeError instead of an assert - with a subfield mapping, a TTree reads the whole branches its keys are under and the gate applies the path rule, so a dataset that mixes an RNTuple first and a TTree later reads with open_files=False, as on main Assisted-by: ClaudeCode:claude-opus-5-5
Assisted-by: ClaudeCode:claude-opus-5-5
Contributor
Author
|
🤖 AI text below 🤖 @ariostas thanks for the prototype design. It's implemented here as you described: buffers keyed by their own node's path,
On the open point about the gate accepting Each new test fails on f7704f4 or on |
Member
|
🤖 AI text below 🤖 @lgray I have two small commits for this PR, but I can't push them: 1. CI fix. 2. Nits from my last look.
The dask, graphed and RNTuple test files give 664 passed, 3 skipped. With 1. test: skip the literal field-name test without dask-awkward and cover the dtype checkdiff --git a/tests/test_1783_rntuple_subfield_projection.py b/tests/test_1783_rntuple_subfield_projection.py
index bd817a9b..aa87d344 100644
--- a/tests/test_1783_rntuple_subfield_projection.py
+++ b/tests/test_1783_rntuple_subfield_projection.py
@@ -285,6 +285,7 @@ def test_rntuple_buffers_are_keyed_by_their_own_node():
@pytest.mark.parametrize("name", ["[ab]", "{a,b}", "d.e", "x*"])
def test_a_field_name_is_selected_literally(tmp_path, monkeypatch, name):
+ pytest.importorskip("dask_awkward")
where, arr = _rntuple(
tmp_path,
ak.Array(
@@ -295,3 +296,24 @@ def test_a_field_name_is_selected_literally(tmp_path, monkeypatch, name):
read = _columns_read(monkeypatch)
assert ak.array_equal(ev["rec", name].compute(), arr["rec", name])
assert len(read) == 1
+
+
+def test_a_buffer_of_the_wrong_dtype_is_an_error(tmp_path):
+ from uproot._dask import TrivialFormMappingInfo
+ from uproot.source.futures import TrivialExecutor
+
+ where, _ = _nested_rntuple(tmp_path)
+ nt = uproot.open(where)
+ form = nt.to_akform()[0].to_dict()
+ form["contents"][form["fields"].index("z")]["primitive"] = "int32"
+ expected = ak.forms.form_with_unique_keys(ak.forms.from_dict(form), ("<root>",))
+ with pytest.raises(TypeError, match="Data type mismatch"):
+ TrivialFormMappingInfo(expected, subfields=True).load_buffers(
+ nt,
+ frozenset({"z"}),
+ 0,
+ nt.num_entries,
+ TrivialExecutor(),
+ TrivialExecutor(),
+ {"ak_add_doc": False},
+ )2. refactor: make TrivialFormMappingInfo.subfields public and drop dotted subfield namesdiff --git a/src/uproot/_dask.py b/src/uproot/_dask.py
index 3af9cbe8..b653e94c 100644
--- a/src/uproot/_dask.py
+++ b/src/uproot/_dask.py
@@ -979,9 +979,9 @@ class TrivialFormMappingInfo(ImplementsFormMappingInfo):
assert isinstance(form, awkward.forms.RecordForm)
self._form = form
- self._subfields = subfields
+ # whether buffers are keyed by their dotted subfield paths (an RNTuple)
+ self.subfields = subfields
self._form_key_to_key = self.build_form_key_to_key(form, subfields)
- self._field_paths = _field_paths(form)
@property
def behavior(self) -> None:
@@ -1024,7 +1024,7 @@ class TrivialFormMappingInfo(ImplementsFormMappingInfo):
# Identify key from form_key
keys.add(self._form_key_to_key[form_key])
# reading a subfield reads the collections above it, so they need not be read whole
- if self._subfields:
+ if self.subfields:
keys -= _ancestors(keys)
return frozenset(keys)
@@ -1041,9 +1041,12 @@ class TrivialFormMappingInfo(ImplementsFormMappingInfo):
# An RNTuple reads every leaf beneath a requested key, selected by key because a name
# like "x" would also select the fields named "x" in other records.
if isinstance(tree, HasFields):
- # by field-name tuple, so a name's dots, braces, and glob characters are literal
+ # by field-name tuple, so braces and glob characters in a name are literal; a
+ # top-level field is selected whole, as a dot in its name (which ROOT forbids,
+ # but uproot can write) is only supported there
+ names = ((k,) if self._form.has_field(k) else k.split(".") for k in keys)
form = self._form.select_columns(
- [tuple(map(glob.escape, self._field_paths[k])) for k in keys]
+ [tuple(map(glob.escape, n)) for n in names]
)
projection_subforms = {f: form.content(f) for f in form.fields}
leaf_keys = [
@@ -1054,7 +1057,7 @@ class TrivialFormMappingInfo(ImplementsFormMappingInfo):
selection = {"filter_field": _filter_field_by_key(tree, leaf_keys)}
else:
fields = list(keys)
- if self._subfields: # a TTree reads the whole branches the keys are under
+ if self.subfields: # a TTree reads the whole branches the keys are under
above = _ancestors(keys)
fields = [f for f in self._form.fields if f in keys or f in above]
projection_subforms = {f: self._form.content(f) for f in fields}
@@ -1108,18 +1111,16 @@ def _ancestors(keys):
return {key[:i] for key in keys for i, c in enumerate(key) if c == "."}
-def _field_paths(form, path=None, names=()):
- """Each record field's dotted key, as ``build_form_key_to_key`` names it, mapped to its
- field names, which may themselves contain dots."""
- out = {}
+def _subfield_paths(form, path=None):
+ """The dotted paths of the record fields beneath ``form``'s top-level fields."""
if form.is_record:
for field in form.fields:
subpath = field if path is None else f"{path}.{field}"
- out[subpath] = (*names, field)
- out.update(_field_paths(form.content(field), subpath, (*names, field)))
+ if path is not None:
+ yield subpath
+ yield from _subfield_paths(form.content(field), subpath)
elif form.is_list or form.is_option or form.is_indexed:
- out.update(_field_paths(form.content, path, names))
- return out
+ yield from _subfield_paths(form.content, path)
def _covered(key, keys):
@@ -1249,7 +1250,7 @@ def _read_tree(
keys = frozenset(keys)
# a dotted RNTuple key names a subfield, read with everything beneath it and the
# structure above it; a dotted TTree branch name is just a name
- if isinstance(tree, HasFields) or getattr(form_mapping_info, "_subfields", False):
+ if isinstance(tree, HasFields) or getattr(form_mapping_info, "subfields", False):
above = _ancestors(keys)
read = lambda key: key in above or _covered(key, keys) # noqa: E731
else:
@@ -1407,9 +1408,7 @@ class UprootReadMixin:
# Determine which TTree keys need to be read; an RNTuple's are read by subfield too
common_keys = frozenset(self.common_keys)
- subfields = {
- k for k, names in _field_paths(self.base_form).items() if len(names) > 1
- }
+ subfields = frozenset(_subfield_paths(self.base_form))
return frozenset(
key
for key in form_info.keys_for_buffer_keys(data_buffers)
diff --git a/tests/test_1783_rntuple_subfield_projection.py b/tests/test_1783_rntuple_subfield_projection.py
index aa87d344..88a920f2 100644
--- a/tests/test_1783_rntuple_subfield_projection.py
+++ b/tests/test_1783_rntuple_subfield_projection.py
@@ -283,7 +283,7 @@ def test_rntuple_buffers_are_keyed_by_their_own_node():
) == {"jets.eta"}
-@pytest.mark.parametrize("name", ["[ab]", "{a,b}", "d.e", "x*"])
+@pytest.mark.parametrize("name", ["[ab]", "{a,b}", "x*"])
def test_a_field_name_is_selected_literally(tmp_path, monkeypatch, name):
pytest.importorskip("dask_awkward")
where, arr = _rntuple(
@@ -298,6 +298,18 @@ def test_a_field_name_is_selected_literally(tmp_path, monkeypatch, name):
assert len(read) == 1
+def test_a_top_level_field_name_may_contain_a_dot(tmp_path, monkeypatch):
+ pytest.importorskip("dask_awkward")
+ # ROOT forbids dots in RNTuple field names, but uproot can write them
+ where, arr = _rntuple(
+ tmp_path, ak.Array([{"a.b": k, "c": {"d": 2 * k}} for k in range(4)])
+ )
+ ev = uproot.dask(where)
+ read = _columns_read(monkeypatch)
+ assert ak.array_equal(ev["a.b"].compute(), arr["a.b"])
+ assert len(read) == 1
+
+
def test_a_buffer_of_the_wrong_dtype_is_an_error(tmp_path):
from uproot._dask import TrivialFormMappingInfo
from uproot.source.futures import TrivialExecutor |
…bfield-projection # Conflicts: # src/uproot/_graphed.py
Contributor
Author
|
🤖 AI text below 🤖 @ariostas thanks. Both diffs are applied unchanged, as commits authored by you: a1e6d23 (CI fix and dtype coverage) and 6e03078 ( |
Member
|
🤖 AI text below 🤖 @lgray thanks for applying those. The PR now conflicts with There is one conflict, in base_form = _base_form_of(
obj, common_keys, interp_options.get("ak_add_doc"), form_mapping
)
is_rntuple = isinstance(obj, HasFields)I tried that resolution locally: the dask, graphed, RNTuple and |
…bfield-projection # Conflicts: # src/uproot/_dask.py
Contributor
Author
ariostas
approved these changes
Oct 28, 2026
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.
🤖 AI text below 🤖
uproot.daskanduproot.graphedproject an RNTuple by top-level field, so readingevents.Muon.ptreads every leaf ofMuon. This PR projects by subfield instead. The design is @ariostas's (review); thanks for the prototype.Columns read from a nested RNTuple whose
Muoncollection has 14 leaves (open_files=TrueandFalsealike):ev.Muon.ptev.Muon.p4.xev.MET.pt(record of 2)ak.num(ev.Muon, axis=1)ev.MuonDesign (
TrivialFormMappingInfo(form, subfields=True)):Muon's offsets byMuon, its data byMuon.pt. A union is read whole.load_buffersselects that withForm.select_columns, using glob-escaped field-name tuples so names are taken literally._read_tree's gate, for an RNTuple or a subfield mapping, accepts a buffer whose key is requested, under a requested key, or above one.keys_for_buffer_keysdrops a collection key when a field under it is also needed, soev.Muon.ptasks for{"Muon.pt"}.uproot.daskanduproot.graphedturn subfields on when every opened tree is an RNTuple. Withknown_base_formno file is opened, so whole fields are read as before. The publicTrivialFormMapping/FormMappingWithVirtualArrayskeepsubfields=False.open_files=Falsedataset whose first file is an RNTuple and a later one a TTree.form_mappingwhosekeys_for_buffer_keysreturns subfield keys gets them read; today they are dropped and become placeholders.FormMappingInfoWithVirtualArraysmaterializes unread buffers throughload_buffers, so dotted keys work there too.Not changed:
ak.num(ev.Muon)alone still reads the whole collection, as onmain. Reading only the offsets would need RNTuple reader support, because selecting the collection field reads every field under it.Tests, in
tests/test_1783_rntuple_subfield_projection.py. Reads are witnessed throughread_cluster_pages.uproot.dask(open_filesTrue/False) anduproot.graphed, including the cases from @ariostas's table:jets.eta+ak.num(jets)jets.etarequested alone.,*?[and braces.test_1720'snecessary_columns == {"rec"}becomes{"rec.x"}. Every new test fails onmainor on the earlier first-leaf version of this PR. scikit-hep/coffea#1648's nested-RNTuple test and coffea's nanoevents tests (565 passed) pass against this branch.