From 307a1adf470c91a5c9856e5f9607c4dca5b05dbd Mon Sep 17 00:00:00 2001 From: Andres Rios Tascon Date: Wed, 5 Aug 2026 16:21:56 -0400 Subject: [PATCH] fix: keep TTree branch lookup indices valid when a counter collides Tree._branch_lookup maps a branch name to its index in Tree._branch_data. When a jagged branch's generated counter name matched a branch that had already been declared, the colliding datum was removed with 'del self._branch_data[...]', which shifts every later datum down by one while _branch_lookup keeps pointing at the old indices. mktree with scalar 'nx', scalar 'y' and jagged 'x' produced the lookup {'nx': 1, 'y': 1, 'x': 2}: 'y' aliased the generated counter, and extending the tree failed. Replace the datum in place instead, which also keeps the counter ahead of the jagged branch it counts. Assisted-by: claude-code:claude-opus-5[1m] --- src/uproot/writing/_cascadetree.py | 12 ++- tests/test_1688_cascadetree_counter_lookup.py | 94 +++++++++++++++++++ 2 files changed, 102 insertions(+), 4 deletions(-) create mode 100644 tests/test_1688_cascadetree_counter_lookup.py diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index 205caa8ed..cfde3580d 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -201,10 +201,14 @@ def __init__( counter_name, counter_dtype, counter_dtype, kind="counter" ) if counter_name in self._branch_lookup: - # counters always replace non-counters - del self._branch_data[self._branch_lookup[counter_name]] - self._branch_lookup[counter_name] = len(self._branch_data) - self._branch_data.append(counter) + # counters always replace non-counters; replace the datum + # in place, because deleting it would shift every later + # datum down by one and invalidate the indices that + # self._branch_lookup already holds for them + self._branch_data[self._branch_lookup[counter_name]] = counter + else: + self._branch_lookup[counter_name] = len(self._branch_data) + self._branch_data.append(counter) if type(content).__name__ == "RecordType": if hasattr(content, "contents"): diff --git a/tests/test_1688_cascadetree_counter_lookup.py b/tests/test_1688_cascadetree_counter_lookup.py new file mode 100644 index 000000000..8fc6e81dc --- /dev/null +++ b/tests/test_1688_cascadetree_counter_lookup.py @@ -0,0 +1,94 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE + +"""Regression tests for issue #1688: generated TTree counters vs. branch lookup. + +When a jagged branch's generated counter name collides with a branch that was +already declared, the colliding datum was deleted from ``_branch_data`` without +reindexing ``_branch_lookup``, so every branch declared after it pointed at the +wrong datum. +""" + +from __future__ import annotations + +import awkward as ak +import numpy as np +import pytest + +import uproot + + +def test_branch_lookup_indices_stay_consistent(tmp_path): + path = str(tmp_path / "file.root") + with uproot.recreate(path) as f: + tree = f.mktree( + "t", + { + "nx": np.dtype("int32"), + "y": np.dtype("float64"), + "x": ak.Array([[1.1]]).type, + }, + ) + lookup = tree._cascading._branch_lookup + data = tree._cascading._branch_data + + # every name must resolve to the datum that carries that name + assert len(set(lookup.values())) == len(lookup) + for name, index in lookup.items(): + assert data[index]["fName"] == name + + # the generated counter replaced the scalar 'nx' declared earlier + assert data[lookup["nx"]]["kind"] == "counter" + + +def test_write_and_read_back_with_colliding_counter(tmp_path): + path = str(tmp_path / "file.root") + x = ak.Array([[1.1], [2.2, 3.3], [4.4, 5.5, 6.6]]) + y = np.array([10.0, 20.0, 30.0]) + + with uproot.recreate(path) as f: + tree = f.mktree( + "t", + {"nx": np.dtype("int32"), "y": np.dtype("float64"), "x": x.type}, + ) + tree.extend({"nx": np.array([1, 2, 3], dtype=np.int32), "y": y, "x": x}) + + with uproot.open(path) as f: + result = f["t"].arrays() + assert result["y"].tolist() == y.tolist() + assert result["x"].tolist() == x.tolist() + assert result["nx"].tolist() == [1, 2, 3] + + +def test_counter_disagreement_still_raises(tmp_path): + path = str(tmp_path / "file.root") + x = ak.Array([[1.1], [2.2, 3.3], [4.4, 5.5, 6.6]]) + + with uproot.recreate(path) as f: + tree = f.mktree( + "t", + {"nx": np.dtype("int32"), "y": np.dtype("float64"), "x": x.type}, + ) + with pytest.raises(ValueError, match="disagree"): + tree.extend( + { + "nx": np.array([9, 9, 9], dtype=np.int32), + "y": np.array([10.0, 20.0, 30.0]), + "x": x, + } + ) + + +def test_no_collision_is_unaffected(tmp_path): + path = str(tmp_path / "file.root") + x = ak.Array([[1.1], [2.2, 3.3], [4.4, 5.5, 6.6]]) + y = np.array([10.0, 20.0, 30.0]) + + with uproot.recreate(path) as f: + tree = f.mktree("t", {"y": np.dtype("float64"), "x": x.type}) + assert list(tree._cascading._branch_lookup) == ["y", "nx", "x"] + tree.extend({"y": y, "x": x}) + + with uproot.open(path) as f: + result = f["t"].arrays() + assert result["y"].tolist() == y.tolist() + assert result["x"].tolist() == x.tolist()