Repository navigation
fix(rntuple): support arrays how="zip" and numpy.asarray (#1782) - #1837
sumitsaraswat362 wants to merge 3 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files
|
ariostas
left a comment
There was a problem hiding this comment.
🤖 AI text below 🤖
Thanks for this! The np.asarray part matches TTree. I checked that np.asarray(tree) and np.asarray(ntuple) both return (n_branches/fields, n_entries), and both raise the same "inhomogeneous shape" error when field shapes differ. Small nit on the description: "returns an array of rows" reads like one row per entry; "one row per field" would be clearer.
The main issue is in the how="zip" context: string and fixed-size array fields get treated as jagged (details and reproducers inline).
|
|
||
| expression_context = [(f, None) for f in arrays.fields] | ||
| expression_context = [ | ||
| (f, {"is_jagged": arrays[f].layout.is_list}) for f in arrays.fields |
There was a problem hiding this comment.
🤖 AI text below 🤖
layout.is_list is also True for strings (a list of chars) and for fixed-size arrays (RegularArray). TTree treats neither as jagged (is_jagged there is only isinstance(interpretation, AsJagged)), so how="zip" restructures these fields in ways TTree doesn't.
The most common case is a single string field next to a scalar:
import awkward as ak, numpy as np, uproot
with uproot.recreate("t.root") as f:
f["ntuple"] = {"x": np.arange(2), "dataset": ["abc", "de"]}
with uproot.open("t.root") as f:
print(f["ntuple"].arrays(how="zip").type)
# 2 * {x: int64, jagged0: {dataset: string}}
# expected: 2 * {x: int64, dataset: string}Other inputs (TTree output shown where uproot can write the same data to a TTree):
| fields | RNTuple how="zip" |
TTree how="zip" |
|---|---|---|
name_a, name_b (strings) |
{name: {a: string}, jagged1: {name_b: string}} |
n/a |
x, pos shape (n, 3) |
{x, jagged0: 3 * {pos}} |
{x, pos: 3 * float64} |
x, arr_a, arr_b shape (n, 2) |
{x, arr: 2 * {a, b}} |
{x, arr_a: 2 * int64, arr_b: 2 * int64} |
jet_pt=[[1.,2.],[3.]], jet_name=["ab","c"] |
jet: var * {name, pt} (event string copied onto every jet) |
n/a |
The last row only happens when every string's length equals that entry's jet count, but then the string gets zipped into the jet collection.
A check on the layout type that excludes strings fixes all of the above. I tried it locally: vector<string> still zips into jet.tag correctly, and the tests in this PR still pass:
expression_context = []
for f in arrays.fields:
layout = arrays[f].layout
is_jagged = isinstance(
layout, (ak.contents.ListOffsetArray, ak.contents.ListArray)
) and layout.parameter("__array__") not in ("string", "bytestring")
expression_context.append((f, {"is_jagged": is_jagged}))Could you add a test with a string field and a fixed-size array field as well?
| is_jagged = ( | ||
| context["is_jagged"] | ||
| if (context is not None and "is_jagged" in context) | ||
| else getattr(array.layout, "is_list", False) | ||
| ) |
There was a problem hiding this comment.
🤖 AI text below 🤖
This fallback has the same is_list problem as the RNTuple context (strings and fixed-size arrays count as jagged). Once HasFields.arrays passes a real context, nothing in the codebase calls group with None. I'd drop the fallback and keep context["is_jagged"], so a caller that builds a bad context fails loudly instead of zipping with a guess. The same applies to test_library_awkward_group_defensive_none_context, which locks this behaviour in.
| @@ -0,0 +1,162 @@ | |||
| from __future__ import annotations | |||
There was a problem hiding this comment.
🤖 AI text below 🤖
Nit: the other test files start with the license line # BSD 3-Clause License; see https://github.com/scikit-hep/uproot/blob/main/LICENSE.
|
Good catch on the string and RegularArray list handling. I updated the layout check, dropped the fallback in library.py, added the license header, and added tests covering strings, fixed-size arrays, and vector of strings. |
Fixes #1782.
This PR addresses two related issues with RNTuple array conversion:
arrays(how="zip")failed becauseHasFields.arrayspassedNoneas the expression context, causingAwkward.groupto raise aTypeErrorwhen checkingcontext["is_jagged"]. We now pass{"is_jagged": ...}inexpression_context, checking forListOffsetArrayandListArraywhile excluding strings and fixed-size arrays.numpy.asarray(ntuple)raisedValueError: object __array__ method not producing an arraybecauseHasFields.__array__returned a dictionary fromarrays(library="np"). Following the TTree and TBranch design,__array__is moved toRFieldso fields produce NumPy arrays directly, whileRNTupledelegates to sequence iteration sonumpy.asarray(ntuple)returns an array of fields (one row per field) matching TTree.