diff --git a/src/uproot/writing/writable.py b/src/uproot/writing/writable.py index b99642cfe..21db2ddf2 100644 --- a/src/uproot/writing/writable.py +++ b/src/uproot/writing/writable.py @@ -1014,9 +1014,8 @@ def _get(self, name, cycle): if self._file._has_tree(key.seek_location): return self._file._get_tree(key.seek_location) else: - raise TypeError( - "WritableDirectory cannot view preexisting TTrees; open the file with uproot.open instead of uproot.recreate or uproot.update" - ) + # load existing TTree and reconstruct cascade + return self._load_existing_ttree(key) elif key.classname.string == "ROOT::RNTuple": if self._file._has_ntuple(key.seek_location): return self._file._get_ntuple(key.seek_location) @@ -1054,6 +1053,221 @@ def get_chunk(start, stop): return readonlykey.get() + def _load_existing_ttree(self, key): + """ + Loads an existing TTree from disk and reconstructs a writable + :doc:`uproot.writing.writable.WritableTree` object with a proper + cascade object, enabling extend via existing machinery. + """ + import struct as _struct + + import uproot.writing._cascadetree as ct + + if self.file_path is None: + raise TypeError( + "uproot.update() on a file-like object does not support accessing " + "existing TTrees; use uproot.update() with a file path instead." + ) + + name = key.name.string + + _dtype_to_struct = { + "f4": "f", + "f8": "d", + "i4": "i", + "i8": "q", + "i2": "h", + "i1": "b", + "u4": "I", + "u8": "Q", + "u2": "H", + "u1": "B", + } + + # read using sink.read + Chunk.wrap + _ReadForUpdate (same as _get) + # avoids loading entire file into memory + self._file.sink.flush() + + def _get_chunk(start, stop): + raw_bytes = self._file.sink.read(start, stop - start) + return uproot.source.chunk.Chunk.wrap( + _readforupdate, raw_bytes, start=start + ) + + _readforupdate = uproot.writing._cascade._ReadForUpdate( + self._file.file_path, + self._file.uuid, + _get_chunk, + self._file._cascading.tlist_of_streamers, + ) + _readforupdate.options = dict(uproot.reading.open.defaults) + _readforupdate.options["minimal_ttree_metadata"] = False + + _raw_bytes = self._file.sink.read( + key.seek_location, + key.num_bytes + key.compressed_bytes, + ) + _chunk = uproot.source.chunk.Chunk.wrap( + _readforupdate, _raw_bytes, start=key.seek_location + ) + _cursor = uproot.source.cursor.Cursor(key.seek_location, origin=key.num_bytes) + _readonlykey = uproot.reading.ReadOnlyKey( + _chunk, _cursor, {}, _readforupdate, self, read_strings=True + ) + tree = _readonlykey.get() + branches = list(tree.branches) + _rkey_chunk, _rkey_cursor = _readonlykey.get_uncompressed_chunk_cursor() + raw = bytearray(_rkey_chunk.raw_data.tobytes()) + + fEntries = tree.member("fEntries") + fTotBytes = tree.member("fTotBytes") + fZipBytes_val = tree.member("fZipBytes") + seq = ( + _struct.pack(">q", fEntries) + + _struct.pack(">q", fTotBytes) + + _struct.pack(">q", fZipBytes_val) + ) + metadata_start = raw.find(seq) + if metadata_start == -1: + raise RuntimeError(f"Could not find TTree metadata position in {name!r}") + + branch_data = [] + branch_lookup = {} + for branch_idx, b in enumerate(branches): + refs_list = list(b.cursor._refs.keys()) + try: + dtype = b.interpretation.numpy_dtype.newbyteorder(">") + except AttributeError: + # TBranchElement or other complex branch — skip + continue + sc = _dtype_to_struct.get(dtype.kind + str(dtype.itemsize), "f") + # detect counter branches (e.g. njets for jagged jets array) + _leaves = b.member("fLeaves") + _is_counter = bool(_leaves) and bool(_leaves[0].member("fIsRange")) + bd = { + "fName": b.name, + "branch_type": dtype, + "kind": "counter" if _is_counter else "normal", + "counter": None, + "dtype": dtype, + "shape": (), + "fTitle": b.member("fTitle"), + "compression": b.compression, + "fBasketSize": b.member("fBasketSize"), + "fEntryOffsetLen": b.member("fEntryOffsetLen"), + "fOffset": b.member("fOffset"), + "fSplitLevel": b.member("fSplitLevel"), + "fFirstEntry": b.member("fFirstEntry"), + "fTotBytes": b.member("fTotBytes"), + "fZipBytes": b.member("fZipBytes"), + "fBasketBytes": b.member("fBasketBytes").copy(), + "fBasketEntry": b.member("fBasketEntry").copy(), + "fBasketSeek": b.member("fBasketSeek").copy(), + "arrays_write_start": b.member("fWriteBasket"), + "arrays_write_stop": b.member("fWriteBasket"), + "metadata_start": ( + # find by searching for fBasketSize + fEntryOffsetLen + fWriteBasket pattern + raw.find( + _struct.pack( + ">iii", + b.member("fBasketSize"), + b.member("fEntryOffsetLen"), + b.member("fWriteBasket"), + ), + b.cursor.index, + ) + - 4 # -4 for fCompress field before fBasketSize + ), + "basket_metadata_start": ( + # fBasketSeek[0] is preceded by: + # speedbump(1) + fBasketBytes(fMaxBaskets*4) + speedbump(1) + fBasketEntry(fMaxBaskets*8) + speedbump(1) + raw.find( + _struct.pack(">q", b.member("fBasketSeek")[0]), + b.cursor.index, + ) + - ( + 1 + + b.member("fMaxBaskets") * 4 + + 1 + + b.member("fMaxBaskets") * 8 + + 1 + ) + ), + "tleaf_reference_number": ( + refs_list[2 + branch_idx * 4] + if 2 + branch_idx * 4 < len(refs_list) + else 0 + ), + "tleaf_maximum_value": ( + int(b.member("fLeaves")[0].member("fMaximum")) + if b.member("fLeaves") + else 0 + ), + "tleaf_special_struct": _struct.Struct(">" + sc + sc), + } + branch_data.append(bd) + branch_lookup[b.name] = branch_idx + + # fix counter references for jagged branches + for bd in branch_data: + if bd.get("fEntryOffsetLen", 0) > 0 and bd["counter"] is None: + counter_nm = "n" + bd["fName"] + counter_bd = next( + (x for x in branch_data if x["fName"] == counter_nm), None + ) + if counter_bd is not None: + bd["counter"] = counter_bd + + fWriteBasket = branches[0].member("fWriteBasket") if branches else 0 + metadata = { + k: tree.member(k) + for k in [ + "fTotBytes", + "fZipBytes", + "fSavedBytes", + "fFlushedBytes", + "fWeight", + "fTimerInterval", + "fScanField", + "fUpdate", + "fDefaultEntryOffsetLen", + "fNClusterRange", + "fMaxEntries", + "fMaxEntryLoop", + "fMaxVirtualSize", + "fAutoSave", + "fAutoFlush", + "fEstimate", + ] + } + + dir_key = self._cascading.data.get_key(name, key.cycle) + freesegments = self._file._cascading.freesegments + + casc = ct.Tree.__new__(ct.Tree) + casc._directory = self._cascading + casc._name = name + casc._title = tree.title + casc._freesegments = freesegments + casc._branch_data = branch_data + casc._branch_lookup = branch_lookup + casc._basket_capacity = ( + next(iter(branches)).member("fMaxBaskets") if branches else 10 + ) + casc._resize_factor = 10.0 + casc._counter_name = lambda counted: "n" + counted + casc._field_name = None + casc._metadata_start = metadata_start + casc._num_baskets = fWriteBasket + casc._num_entries = fEntries + casc._metadata = metadata + casc._key = dir_key + + path = (*self._path, name) + writable_tree = WritableTree(path, self._file, casc) + self._file._trees[key.seek_location] = writable_tree + return writable_tree + def _del(self, name, cycle): key = self._cascading.data.get_key(name, cycle) if key is None: @@ -1883,10 +2097,116 @@ def num_baskets(self) -> int: """ return self._cascading.num_baskets - def extend(self, data): + def add_branches(self, branches): + """ + Args: + branches (dict of str -> array): Names and data of new branches. + + Adds new branches to this TTree in-place. Only the new branch data and + an updated TTree header are written; existing data is never touched. + Works with both simple TBranch and TBranchElement files. + + .. code-block:: python + + with uproot.update("file.root") as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + """ + if self._file.sink.closed: + raise ValueError("cannot modify a TTree in a closed file") + + source = self._path[-1] + + # navigate to the correct directory (handles subdirectories) + directory = self._file.root_directory + for part in self._path[:-1]: + directory = directory[part] + + # validate all branches have same length as existing tree + key = directory._cascading.data.get_key(source) + casc = directory._load_existing_ttree(key)._cascading + num_entries = casc._num_entries + + for branch_name, branch_data in branches.items(): + arr = numpy.asarray(branch_data) + if len(arr) != num_entries: + raise ValueError( + f"branch {branch_name!r} has {len(arr)} entries but TTree has " + f"{num_entries} entries; all new branches must match the tree length" + ) + if branch_name in casc._branch_lookup: + raise ValueError(f"branch {branch_name!r} already exists in this TTree") + + # check if file has TBranchElement branches (object dtype) + # _load_existing_ttree skips them, so we detect by checking dtype + if any( + bd.get("dtype") is not None and bd.get("dtype") == numpy.dtype("O") + for bd in casc._branch_data + ): + raise NotImplementedError( + "add_branches for files with TBranchElement branches is not yet " + "supported via the cascade approach" + ) + + # add new branch dicts to cascade + compression = casc._freesegments.fileheader.compression + for branch_name, branch_data in branches.items(): + arr = numpy.asarray(branch_data) + if arr.dtype.kind == "O": + raise TypeError( + f"branch {branch_name!r} has object dtype — only simple numeric " + f"types are supported for add_branches" + ) + dtype = arr.dtype.newbyteorder(">") + new_bd = casc._branch_np(branch_name, arr.dtype, dtype) + new_bd["compression"] = compression + casc._branch_data.append(new_bd) + casc._branch_lookup[branch_name] = len(casc._branch_data) - 1 + + # rewrite TTree metadata blob with new branches included + casc.write_anew(self._file.sink) + + # write one basket per new branch + old_num_baskets = casc._num_baskets + casc._num_baskets = 0 + for branch_name, branch_data in branches.items(): + arr = numpy.asarray(branch_data).astype( + casc._branch_data[casc._branch_lookup[branch_name]]["dtype"] + ) + totbytes, zipbytes, location = casc.write_np_basket( + self._file.sink, branch_name, compression, arr + ) + datum = casc._branch_data[casc._branch_lookup[branch_name]] + datum["fTotBytes"] += totbytes + datum["fZipBytes"] += zipbytes + datum["fBasketBytes"][0] = zipbytes + datum["fBasketSeek"][0] = location + datum["fBasketEntry"][1] = num_entries + datum["arrays_write_start"] = 0 + datum["arrays_write_stop"] = 1 + casc._metadata["fTotBytes"] += totbytes + casc._metadata["fZipBytes"] += zipbytes + + casc._num_baskets = old_num_baskets + casc.write_updates(self._file.sink) + self._file.sink.flush() + + # update in-memory directory cache + dir_key_obj = directory._cascading.data.get_key(source) + dir_key_obj._seek_location = casc._key.seek_location + + # update self._cascading so subsequent extend uses correct metadata + writable_tree = uproot.writing.writable.WritableTree( + self._path, self._file, casc + ) + self._file._trees[casc._key.seek_location] = writable_tree + self._cascading = casc + + def extend(self, data, *, accept_new_fields=False): """ Args: data (dict of str \u2192 arrays): More array data to add to the TTree. + accept_new_fields (bool): If True, new fields in data are automatically added + with zeros back-filled for existing entries before extending. This method adds data to an existing TTree, whether it was created through assignment or :doc:`uproot.writing.writable.WritableDirectory.mktree`. @@ -1910,6 +2230,81 @@ def extend(self, data): **As a word of warning,** be sure that each call to :ref:`uproot.writing.writable.WritableTree.extend` includes at least 100 kB per branch/array. (NumPy and Awkward Arrays have an `nbytes `__ property; you want at least ``100000`` per array.) If you ask Uproot to write very small TBaskets, it will spend more time working on TBasket overhead than actually writing data. The absolute worst case is one-entry-per-:ref:`uproot.writing.writable.WritableTree.extend`. See `#428 (comment) `__. """ + if self._cascading is None: + raise RuntimeError( + "_cascading is None — this should not happen; please report this bug" + ) + # validate branches + # get user-facing branch names (exclude auto-generated counter and record parent branches) + # get record parent names to exclude their sub-fields + _record_names = { + bd.get("name", "") + for bd in self._cascading._branch_data + if bd.get("kind") == "record" + } + _user_branch_names = [ + bd["fName"] + for bd in self._cascading._branch_data + if bd.get("kind") not in ("counter", "record") + and "fName" in bd + and not any( + bd["fName"].startswith(rn + "_") or bd["fName"].startswith(rn + ".") + for rn in _record_names + if rn + ) + ] + # include counter branches that user explicitly provides in data + _counter_branch_names = [ + bd["fName"] + for bd in self._cascading._branch_data + if bd.get("kind") == "counter" and "fName" in bd + ] + # check if data looks like a flat dict of branch arrays (not a record/awkward array) + _data_is_flat_dict = isinstance(data, dict) and all( + not hasattr(v, "fields") for v in data.values() + ) + if isinstance(data, dict) and _data_is_flat_dict: + existing_names = _user_branch_names + # also get record parent names that the user passes as dicts + _record_parent_names = { + bd.get("name") + for bd in self._cascading._branch_data + if bd.get("kind") == "record" and bd.get("name") + } + new_fields = { + k: v + for k, v in data.items() + if k not in existing_names + and k not in _record_parent_names + and k not in _counter_branch_names + } + # skip counter branches not provided by user (auto-generated) + missing = [ + b + for b in existing_names + if b not in data and b not in _counter_branch_names + ] + # add counter branches to existing_names if user provides them + existing_names = _user_branch_names + [ + c for c in _counter_branch_names if c in data + ] + if missing: + raise ValueError( + f"'extend' must fill every branch with the same number of entries; missing: {missing}" + ) + if new_fields: + if not accept_new_fields: + raise ValueError( + "'extend' was given data that do not correspond to any branch: " + + repr(next(iter(new_fields))) + ) + zeros = { + k: numpy.zeros( + self._cascading._num_entries, dtype=numpy.asarray(v).dtype + ) + for k, v in new_fields.items() + } + self.add_branches(zeros) self._cascading.extend(self._file, self._file.sink, data) def show( diff --git a/tests/test_0406_write_a_ttree.py b/tests/test_0406_write_a_ttree.py index 91a5ce48a..a5697a21f 100644 --- a/tests/test_0406_write_a_ttree.py +++ b/tests/test_0406_write_a_ttree.py @@ -342,9 +342,6 @@ def test_writable_vs_readable_tree(tmp_path): b2 = [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9] with uproot.update(newfile) as fin: - with pytest.raises(TypeError): - oldtree = fin["t1"] - fin.mktree("t2", {"b1": np.int32, "b2": np.float64}, "title") for _ in range(5): diff --git a/tests/test_1690_ttree_inplace.py b/tests/test_1690_ttree_inplace.py new file mode 100644 index 000000000..b74fd6b4d --- /dev/null +++ b/tests/test_1690_ttree_inplace.py @@ -0,0 +1,394 @@ +import os +import shutil + +import numpy as np +import pytest + +import uproot +from skhep_testdata import data_path + +try: + import ROOT + + has_root = True +except ImportError: + has_root = False + +skip_no_root = pytest.mark.skipif(not has_root, reason="ROOT is not installed") + + +# ── add_branches tests ──────────────────────────────────────────────────────── + + +def test_add_branch_simple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 2 + assert "new_branch" in [b.name for b in f["tree"].branches] + assert np.all(f["tree"]["new_branch"].array() == 1.0) + + +def test_add_branch_multiple_branches(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches( + { + "branch_a": np.ones(100, dtype=np.float32), + "branch_b": np.zeros(100, dtype=np.int32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 3 + assert np.all(f["tree"]["branch_a"].array() == 1.0) + assert np.all(f["tree"]["branch_b"].array() == 0) + + +def test_add_branch_int32(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_int": np.arange(100, dtype=np.int32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert np.all(f["tree"]["new_int"].array() == np.arange(100, dtype=np.int32)) + + +def test_add_branch_preserves_existing(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend( + { + "x": np.arange(100, dtype=np.float32), + "y": np.arange(100, dtype=np.int32), + } + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert np.all(f["tree"]["x"].array() == np.arange(100, dtype=np.float32)) + assert np.all(f["tree"]["y"].array() == np.arange(100, dtype=np.int32)) + assert np.all(f["tree"]["new_branch"].array() == 1.0) + + +def test_add_branch_tbranchelement(tmp_path): + # add_branches for TBranchElement files is not supported + # due to internal reference numbers that break when blob is rewritten + shutil.copy( + data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root") + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + with pytest.raises(Exception): + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) + + +def test_add_branch_wrong_length(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="entries"): + f["tree"].add_branches({"new_branch": np.ones(50, dtype=np.float32)}) + + +def test_add_branch_nonexistent_tree(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["nonexistent"].add_branches( + {"new_branch": np.ones(100, dtype=np.float32)} + ) + + +@skip_no_root +def test_add_branch_root_readable(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.ones(100, dtype=np.float32)}) + + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree;1") + tree.SetCacheSize(0) + tree.GetEntry(0) + assert tree.new_branch == pytest.approx(1.0) + f.Close() + + +@skip_no_root +def test_add_branch_tbranchelement_root_readable(tmp_path): + shutil.copy( + data_path("uproot-HZZ-objects.root"), os.path.join(tmp_path, "HZZ.root") + ) + + with uproot.update(os.path.join(tmp_path, "HZZ.root")) as f: + with pytest.raises((NotImplementedError, TypeError, KeyError)): + f["events"].add_branches({"new_branch": np.ones(2421, dtype=np.float32)}) + + +# ── extend tests ────────────────────────────────────────────────────────────── + + +def test_extend_simple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "y": np.ones(50, dtype=np.int32) * 3, + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["x"].array()[:100] == 1.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + assert np.all(f["tree"]["y"].array()[:100] == 0) + assert np.all(f["tree"]["y"].array()[100:] == 3) + + +def test_extend_preserves_existing(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.arange(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.arange(100, dtype=np.float32) + 100}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + arr = f["tree"]["x"].array() + assert len(arr) == 200 + assert np.all(arr[:100] == np.arange(100, dtype=np.float32)) + assert np.all(arr[100:] == np.arange(100, dtype=np.float32) + 100) + + +def test_extend_missing_branch(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="missing"): + f["tree"].extend({"x": np.ones(50, dtype=np.float32)}) + + +def test_extend_mismatched_lengths(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32, "y": np.int32}) + f["tree"].extend( + {"x": np.ones(100, dtype=np.float32), "y": np.zeros(100, dtype=np.int32)} + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError): + f["tree"].extend( + {"x": np.ones(50, dtype=np.float32), "y": np.ones(30, dtype=np.int32)} + ) + + +def test_extend_nonexistent_branch(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(Exception): + f["tree"].extend({"nonexistent": np.ones(100, dtype=np.float32)}) + + +def test_extend_accept_new_fields(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + }, + accept_new_fields=True, + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert "new_branch" in [b.name for b in f["tree"].branches] + assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) + assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + + +def test_extend_new_fields_error_without_flag(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError): + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32), + "new_branch": np.ones(50, dtype=np.float32), + } + ) + + +@skip_no_root +def test_extend_root_readable(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) + + ROOT.gROOT.ProcessLine("gErrorIgnoreLevel = kError;") + f = ROOT.TFile.Open(str(os.path.join(tmp_path, "test.root")), "READ") + tree = f.Get("tree") + assert tree.GetEntries() == 150 + tree.SetCacheSize(0) + tree.GetEntry(149) + assert tree.x == pytest.approx(2.0) + f.Close() + + +def test_add_branch_sequential(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"branch_a": np.ones(100, dtype=np.float32) * 2}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"branch_b": np.ones(100, dtype=np.int32) * 3}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert len(f["tree"].branches) == 3 + assert np.all(f["tree"]["branch_a"].array() == 2.0) + assert np.all(f["tree"]["branch_b"].array() == 3) + + +def test_add_branch_then_extend_same_session(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + t = f["tree"] + t.add_branches({"new_branch": np.zeros(100, dtype=np.float32)}) + t.extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) + assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) + assert np.all(f["tree"]["x"].array()[100:] == 2.0) + + +def test_extend_multiple_sessions(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 2}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.ones(50, dtype=np.float32) * 3}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 200 + assert np.all(f["tree"]["x"].array()[:100] == 1.0) + assert np.all(f["tree"]["x"].array()[100:150] == 2.0) + assert np.all(f["tree"]["x"].array()[150:] == 3.0) + + +def test_extend_after_add_branch_new_session(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + f["tree"].extend({"x": np.ones(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].add_branches({"new_branch": np.zeros(100, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend( + { + "x": np.ones(50, dtype=np.float32) * 2, + "new_branch": np.ones(50, dtype=np.float32) * 99, + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].member("fEntries") == 150 + assert np.all(f["tree"]["new_branch"].array()[:100] == 0.0) + assert np.all(f["tree"]["new_branch"].array()[100:] == 99.0) + + +def test_extend_jagged_array(tmp_path): + """Counter branches should not be required from the user when extending.""" + ak = pytest.importorskip("awkward") + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"jets": "var * float32", "x": np.float32}) + f["tree"].extend( + { + "jets": ak.Array([[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]]), + "x": np.array([1.0, 2.0, 3.0], dtype=np.float32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].num_entries == 3 + assert f["tree"]["jets"].array().tolist() == [ + [1.0, 2.0], + [3.0], + [4.0, 5.0, 6.0], + ] + + +def test_extend_after_many_extends(tmp_path): + """Extending a tree that already has more than 10 baskets (fMaxBaskets expansion).""" + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mktree("tree", {"x": np.float32}) + for i in range(12): + f["tree"].extend({"x": np.full(5, i, dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tree"].extend({"x": np.full(5, 99, dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["tree"].num_entries == 65 + assert f["tree"]["x"].array()[-5:].tolist() == [99.0] * 5