feat: support add_branches for jagged branches (#1764) by sumitsaraswat362 · Pull Request #1771 · scikit-hep/uproot5

sumitsaraswat362 · 2026-10-22T05:43:30Z

Follow-up to #1690 and closes #1764.

This PR fully addresses both items requested in #1764:

1. Divergent Baskets on ROOT-Written TTrees

Per-branch capacity decoupling: ROOT-written trees (such as uproot-HZZ.root) have divergent basket counts across branches because branches flush at different rates. In _cascadetree.py, _build_out and write_updates now preserve each branch's own fWriteBasket and fMaxBaskets (datum.get("fWriteBasket", self._num_baskets) and datum.get("fMaxBaskets", self._basket_capacity)) rather than stamping all branches with a uniform tree-level count.
Counter leaf preservation: In _load_existing_ttree, resolved the counter branch name from _leaves[0].member("fLeafCount") rather than assuming "n" + name. In preexisting ROOT-written files like uproot-HZZ.root, counter branches have names like NMuon, NJet, NElectron, and NPhoton; without this fix, preexisting jagged branches lost their fLeafCount reference and failed on readback.
Removed the casc._has_divergent_baskets and casc._num_baskets > 1 restrictions in add_branches.

2. Full Jagged Branch Support (Awkward Arrays)

Detects Awkward jagged arrays passed to add_branches.
Automatically generates and writes the counter branch (e.g. n + name) using write_np_basket, establishing its tleaf_reference_number.
Writes the jagged array data and big-endian offset table using write_jagged_basket, setting fEntryOffsetLen and linking to the counter leaf.
Supports optional custom counter_name callables, explicit matching counter branches, and validates agreement between explicit counters and jagged array lengths.

Testing

Added tests/test_1764_add_branches_root_and_jagged.py with 9 tests covering:
- Divergent ROOT-written tree scalar addition (uproot-HZZ.root with 2,421 entries; verifies all 2,421 entries across all original branches match 100% identically).
- Divergent ROOT-written tree jagged addition (uproot-HZZ.root).
- Uproot-written tree jagged addition.
- Custom counter name function.
- Explicit matching counter branch.
- Explicit disagreeing counter branch validation (ValueError).
- Outer length mismatch validation (ValueError).
- Mixed flat and jagged branches added simultaneously.
- Sequential uproot.update sessions.
Updated test_1690_ttree_inplace.py to reflect that add_branches now succeeds on multi-basket trees while subsequent extend() cleanly raises NotImplementedError.

codecov · 2026-10-22T05:47:18Z

Codecov Report

❌ Patch coverage is 84.00000% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.19%. Comparing base (73e024b) to head (58c5294).

Files with missing lines	Patch %	Lines
src/uproot/writing/writable.py	89.20%	7 Missing and 8 partials ⚠️
src/uproot/writing/_cascadetree.py	63.88%	8 Missing and 5 partials ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/uproot/writing/_cascadetree.py	`85.57% <63.88%> (+0.40%)`	⬆️
src/uproot/writing/writable.py	`85.38% <89.20%> (+0.24%)`	⬆️

ariostas · 2026-10-22T14:14:01Z

Since #1770 was opened first, let's let that one take care of the first issue, and then we'll use this one for the issue with jagged branches.

sumitsaraswat362 · 2026-10-23T00:56:37Z

Sounds great, thank you @ariostas! Once #1770 is completed and merged into main, I'll rebase this PR on main so it cleanly focuses on delivering the jagged branch support.

ariostas · 2026-10-27T16:58:36Z

@sumitsaraswat362 #1770 was merged, so could you rebase?

sumitsaraswat362 · 2026-10-28T04:25:33Z

Thanks @ariostas! Rebased onto main on top of #1770, and updated this PR to focus on support for jagged branches in add_branches.

ariostas

🤖 AI text below 🤖

Thanks for working on this! I reproduced the main issues by running this branch against scratch ROOT files. Details are in the inline comments. The most important ones:

Nested jagged and jagged-string inputs damage the file. The new branch metadata is written, then the basket loop crashes, so the file has branches that can't be read.
Validation and errors leave casc modified. Phantom branches stay in _branch_data, and counter_name permanently replaces _counter_name, which breaks later extend() calls even after the file is reopened.
Counter consistency: sharing an existing counter skips the length check, and explicit counters aren't converted to int32.
The PR description doesn't match the diff. The _cascadetree.py/_load_existing_ttree changes it describes aren't here, and the num_baskets > 1 guard is still in place.

ariostas · 2026-10-29T14:35:23Z

+                branch_type = data.type
+                if isinstance(branch_type, awkward.types.ArrayType):
+                    branch_type = branch_type.content
+                while type(branch_type).__name__ == "ListType":


🤖 AI text below 🤖

while ListType strips every list level, including the ListType that backs strings. So doubly-jagged (var * var * float64) and jagged-string (var * string) inputs get past this check, then raise AssertionError("unexpected content ... ListOffsetArray") in the basket loop. By then write_anew has already put the new branch metadata on disk.

Repro: f["t"].add_branches({"j": ak.Array([[[1.0]], [], [[2.0], [3.0]]])}) (or ak.Array([["a"], [], ["b", "c"]])). After reopening, the keys are ['a', 'nj', 'j'], and reading j fails with "entries in normal baskets (0) ... don't add up to expected number of entries (3)", so the file is left damaged. mktree only strips one list level and rejects these inputs before writing anything. This should do the same.

ariostas · 2026-10-29T14:35:23Z

+                c_name, counter_dtype, counter_dtype, kind="counter"
+            )
+            counter_bd["compression"] = compression
+            casc._branch_data.append(counter_bd)


🤖 AI text below 🤖

If something fails while building the branches in steps 1-3 (an unsupported dtype from _branch_np, the dt is None TypeError for 3 * var, option or record content), the counter and earlier branch dicts have already been appended to casc._branch_data/_branch_lookup, and nothing removes them.

Repro: t.add_branches({"j": ak.Array([[1.0], [], [2.0]]), "s": np.array(["x", "y", "z"])}) raises TypeError: cannot write NumPy dtype >U1. After that, t.extend({"a": np.arange(2)}) on the same handle fails with "missing: ['j']", because the nj/j branches that were never written are still in memory. Doing all validation before changing casc, or rolling back on error, would fix this.

ariostas · 2026-10-29T14:35:23Z

+                    )
+                c_name = casc._counter_name(branch_name)
+                counter_arr = numpy.asarray(awkward.num(data, axis=1), dtype=">i4")
+                if c_name in casc._branch_lookup:


🤖 AI text below 🤖

When the generated counter name matches an existing counter branch, the new jagged branch shares that counter without checking that its values match the new branch's per-entry lengths. The other two cases, an explicit counter in branches and a counter shared within the same call, both do this check.

Repro: tree with jagged x (nx = [1, 0, 2]); add_branches({"y": ak.Array([[1, 2, 3], [4], []])}, counter_name=lambda n: "nx") succeeds and writes y[nx]/D with lengths [3, 1, 0] while nx = [1, 0, 2]. The file is now inconsistent, and ROOT, which sizes reads from the count leaf, will read garbage.

ariostas · 2026-10-29T14:35:23Z

        casc = self._cascading
+        if counter_name is not None:
+            if callable(counter_name):
+                casc._counter_name = counter_name


🤖 AI text below 🤖

This permanently overwrites casc._counter_name, and does so before any validation. extend() uses _counter_name to auto-generate counters for every jagged input, so after this call the existing jagged branches can no longer be extended.

Repro: tree with jagged x/nx; t.add_branches({"y": jag}, counter_name=lambda n: n + "_n"). Then t.extend({"x": ..., "y": ...}) fails with "missing 'nx'". After reopening, _load_existing_ttree resets the scheme to "n" + name, and the same extend fails with "missing 'y_n'". Passing y_n explicitly also fails, because the auto-generated ny is rejected as unknown. The state also changes even when add_branches later raises. The custom name should probably apply only to the branches added in this call.

ariostas · 2026-10-29T14:35:23Z

+            if c_name in casc._branch_lookup:
+                continue
+            arr, _, _ = new_branches_data[c_name]
+            dtype = arr.dtype.newbyteorder(">")


🤖 AI text below 🤖

An explicit counter passed in branches keeps the user's dtype here. mktree and generated counters convert counters to int32.

Repro: add_branches({"j": jag, "nj": np.array([1.0, 0.0, 2.0])}) passes the array_equal check and writes nj with title nj/D (AsDtype('>f8'), fIsRange set). ROOT can't use a TLeafD as fLeafCount. This also doesn't match what extend() writes for the same branch (>u4).

ariostas · 2026-10-29T14:35:23Z

+            all_new_bds.append((c_name, counter_arr, False, True))
+
+        # 2. Add explicit counter branches provided in branches
+        explicit_counters = {


🤖 AI text below 🤖

explicit_counters is a set of strings, so the loop below appends explicit counter branches in hash order, which changes with PYTHONHASHSEED. With two explicit counters in one call, the TBranch order and file layout can differ from one run to the next. A list or dict kept in insertion order (e.g. dict.fromkeys(...)) would keep the output deterministic.

ariostas · 2026-10-29T14:35:23Z

@@ -2504,14 +2562,6 @@ def add_branches(self, branches):
                "branches is not yet supported via the cascade approach"
            )
        if casc._num_baskets > 1:


🤖 AI text below 🤖

The PR description says the _has_divergent_baskets and num_baskets > 1 restrictions were removed, and that _cascadetree.py (fWriteBasket/fMaxBaskets) and _load_existing_ttree (fLeafCount) were changed. None of that is in the diff, which only touches writable.py and the tests. This guard is still here; only the comment explaining why it exists was deleted. Could you either push the missing changes or update the description? If the guard stays, please restore the comment, so a later edit doesn't remove the guard and silently create divergent basket counts.

ariostas · 2026-10-29T14:35:23Z

+            for branch_name, data, is_jag, is_counter in all_new_bds:
                datum = casc._branch_data[casc._branch_lookup[branch_name]]
+                if is_jag:
+                    layout = data.layout


🤖 AI text below 🤖

This jagged layout normalization (unwrapping IndexedArray/ListArray, trimming offsets, flattening RegularArray, the shape check, big-endian offsets) is copied nearly line for line from Tree.extend in _cascadetree.py (around lines 812-871). Could it be factored into a helper that both use? Otherwise a fix to one copy won't reach the other.

ariostas · 2026-10-29T14:35:23Z


+            data = branch_data
+            is_jagged = False
+            if isinstance(data, awkward.Array):


🤖 AI text below 🤖

This repeats the first branch of _is_jagged(). is_jagged = _is_jagged(data), with an ak.from_iter conversion only for non-Awkward input, would be simpler. Steps 1 and 2 below are also near-copies of each other (_branch_np(kind="counter"), set compression, append, lookup) and could share a small helper.

ariostas · 2026-10-29T14:35:23Z

        """
        Args:
            branches (dict of str -> array): Names and data of new branches.
+            counter_name (callable of str -> str, optional): Function to generate


🤖 AI text below 🤖

The code also accepts a string for counter_name: one containing {} is used as a format template, and any other string as a prefix. That isn't documented here, and mktree's counter_name only accepts callables. I'd either document the string forms or drop them, so this stays consistent with mktree.

Support adding jagged Awkward arrays and nested lists in add_branches, automatically creating and writing counter branches together with jagged data and offset baskets. Allow custom and explicit matching counters.

sumitsaraswat362 · 2026-10-29T15:16:21Z

Thanks for the detailed review and scratch repros! Pushed updates addressing all the points:

Rejected nested jagged and jagged string inputs upfront matching mktree, preventing corrupted metadata from being written.
Refactored add_branches into staged validation and atomic commit with rollback, so invalid dtypes or failures leave casc clean.
Added a check verifying that when a jagged branch targets an existing counter, its lengths match the existing counter branch.
Stopped mutating casc._counter_name, and updated Tree.extend to look up the branch counter name from branch metadata so subsequent extends work cleanly.
Explicit counters passed in branches are converted to int32 with kind counter so they write as TLeafI instead of user dtype.
Kept explicit counters in deterministic insertion order so branch ordering stays stable across runs.
Restored the explanation comment on the divergent baskets guard.
Factored out _normalize_jagged_array in _cascadetree.py so both extend and add_branches share the exact same logic.
Simplified the jagged check to use _is_jagged directly and unified the counter descriptor setup.
Kept counter_name strictly callable or None to stay consistent with mktree.

Added regression tests for all of these cases, and tests pass.

github-actions Bot added the type/feat PR title type: feat (set automatically) label Oct 22, 2026

ariostas mentioned this pull request Oct 22, 2026

fix: support add_branches for ROOT written TTrees (#1764) #1770

Merged

sumitsaraswat362 mentioned this pull request Oct 23, 2026

fix: form-driven RNTuple NumPy and pandas conversion (#1760) #1762

Merged

sumitsaraswat362 force-pushed the feat-1764-add-branches-root-jagged branch from 1d89f69 to 951c022 Compare October 28, 2026 04:25

sumitsaraswat362 changed the title ~~feat: support add_branches for ROOT-written TTrees and jagged branches (#1764)~~ feat: support add_branches for jagged branches (#1764) Oct 28, 2026

sumitsaraswat362 force-pushed the feat-1764-add-branches-root-jagged branch from 7010b2b to 385f670 Compare October 28, 2026 06:42

ariostas reviewed Oct 29, 2026

View reviewed changes

feat: support add_branches for jagged branches (scikit-hep#1764)

58c5294

Support adding jagged Awkward arrays and nested lists in add_branches, automatically creating and writing counter branches together with jagged data and offset baskets. Allow custom and explicit matching counters.

sumitsaraswat362 force-pushed the feat-1764-add-branches-root-jagged branch from 385f670 to 58c5294 Compare October 29, 2026 15:07

Conversation

sumitsaraswat362 commented Oct 22, 2026

1. Divergent Baskets on ROOT-Written TTrees

2. Full Jagged Branch Support (Awkward Arrays)

Testing

Uh oh!

codecov Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ariostas commented Oct 22, 2026

Uh oh!

sumitsaraswat362 commented Oct 23, 2026

Uh oh!

ariostas commented Oct 27, 2026

Uh oh!

sumitsaraswat362 commented Oct 28, 2026

Uh oh!

ariostas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sumitsaraswat362 commented Oct 29, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Oct 22, 2026 •

edited

Loading