diff --git a/src/uproot/writing/_cascade.py b/src/uproot/writing/_cascade.py index e3ad033c0..cb200b90c 100644 --- a/src/uproot/writing/_cascade.py +++ b/src/uproot/writing/_cascade.py @@ -1059,6 +1059,7 @@ def __init__(self, file_path, uuid, get_chunk=None, tlist_of_streamers=None): self._get_chunk = get_chunk self._tlist_of_streamers = tlist_of_streamers self._custom_classes = None + self.source = self @property def detached(self): @@ -2145,7 +2146,7 @@ def deserialize(cls, raw_bytes, location): assert compression_code >= 0 assert info_location >= 0 assert info_num_bytes >= 0 - assert uuid_version <= 8 + assert 0 <= uuid_version <= 8 out = FileHeader( end, diff --git a/src/uproot/writing/_cascadentuple.py b/src/uproot/writing/_cascadentuple.py index ebca2a72f..f28deb1ee 100644 --- a/src/uproot/writing/_cascadentuple.py +++ b/src/uproot/writing/_cascadentuple.py @@ -242,21 +242,33 @@ def serialize(self): # https://github.com/root-project/root/blob/master/tree/ntuple/v7/doc/specifications.md#column-description class NTuple_Column_Description: - def __init__(self, type_num, bits_on_disk, field_id, flags, repr_index): + def __init__( + self, type_num, bits_on_disk, field_id, flags, repr_index, first_element_index=0 + ): self.type_num = type_num self.bits_on_disk = bits_on_disk self.field_id = field_id self.flags = flags self.repr_index = repr_index + self.first_element_index = first_element_index def serialize(self): + import uproot.const + + flags = self.flags + if self.first_element_index > 0: + flags = flags | uproot.const.RNTupleColumnFlags.DEFERRED header_bytes = _rntuple_column_record_format.pack( self.type_num, self.bits_on_disk, self.field_id, - self.flags, + int(flags), self.repr_index, ) + if self.first_element_index > 0: + import struct + + header_bytes += struct.pack("".format( @@ -2155,10 +2356,18 @@ def num_entries(self) -> int: """ return self._cascading.num_entries - def extend(self, data): + def extend(self, data, accept_new_fields=False): + if self._column_encoding_error is not None: + raise ValueError(self._column_encoding_error) """ Args: data (dict of str \u2192 arrays): More array data to add to the RNTuple. + accept_new_fields (bool): If False (default), raises ValueError if + data contains fields not already in the RNTuple, forcing the user + to call :ref:`uproot.writing.writable.WritableNTuple.add_fields` + first. If True, new fields are automatically added back-filled + with zeros for existing entries, then extended with the provided + values for new entries. This method adds data to an existing RNTuple, whether it was created through assignment or :doc:`uproot.writing.writable.WritableDirectory.mkrntuple`. @@ -2173,17 +2382,310 @@ def extend(self, data): .. code-block:: python - my_directory.mkrntuple("ntuple6", {"branch1": numpy_dtype, "branch2": awkward_type}) + with uproot.update("file.root") as f: + f["mytuple"].extend({ + "x": np.array([4, 5, 6]), + "y": np.array([40, 50, 60]), + }) - my_directory["ntuple6"].extend({"branch1": another_numpy_array, - "branch2": another_awkward_array}) + # automatically add new field and extend + with uproot.update("file.root") as f: + f["mytuple"].extend({"x": np.array([7, 8]), "z": np.array([70, 80])}, + accept_new_fields=True) .. warning:: **As a word of warning,** be sure that each call to :ref:`uproot.writing.writable.WritableNTuple.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 _column_counts has more columns than _column_keys, reload cascading + # this happens after add_fields adds extension columns + if len(self._cascading._column_counts) > len( + self._cascading._header._column_keys + ): + key = self._file.root_directory._cascading.data.get_key(self._path[-1], 1) + reloaded = self._file.root_directory._load_existing_ntuple(key) + self._cascading = reloaded._cascading + + if isinstance(data, dict): + existing_keys = set(self._cascading._header._akform.fields) + if hasattr(self._cascading, "_existing_field_records"): + existing_keys.update( + fr.field_name for fr in self._cascading._existing_field_records + ) + new_field_names = {k for k in data.keys() if k not in existing_keys} + if new_field_names and not accept_new_fields: + raise ValueError( + f"Data contains fields not in this RNTuple: {sorted(new_field_names)}. " + f"Call add_fields() first, or pass accept_new_fields=True to add them automatically." + ) + elif new_field_names and accept_new_fields: + self.add_fields( + { + k: numpy.asarray( + awkward.flatten(awkward.Array(data[k]), axis=None) + ).dtype + for k in new_field_names + } + ) + key = self._file.root_directory._cascading.data.get_key( + self._path[-1], 1 + ) + reloaded = self._file.root_directory._load_existing_ntuple(key) + self._cascading = reloaded._cascading + self._cascading.extend(self._file, self._file.sink, data) + def add_fields(self, new_fields): + """ + Args: + new_fields (dict of str -> numpy dtype): New field names and types. + Only scalar numeric types are accepted (e.g. ``np.int32``, + ``np.float64``). Variable-length or nested types are not supported. + + Adds new fields to this RNTuple, back-filled with zeros for existing entries. + Only top-level fields or subfields of existing untyped record fields are supported. + + Note: when using ``accept_new_fields=True`` in :meth:`extend`, new fields + must be flat (scalar) types. Jagged or nested new fields will fail the + form compatibility check in ``extend``. + + Raises: + TypeError: if a field type is not a simple scalar numeric type. + ValueError: if a field already exists, if the subfield parent is not + found or is a typed struct (C++ typename), or if the file was + opened without a file path (file-like objects are not supported). + + For example, + + .. code-block:: python + + with uproot.update("file.root") as f: + f["mytuple"].add_fields({"z": np.int32, "w": np.float32}) + """ + + import uproot.writing._cascadentuple as cnt + + compression = self._cascading._freesegments.fileheader.compression + num_entries = self._cascading._num_entries + header = self._cascading._header + footer = self._cascading._footer + existing_footer = self._cascading._existing_footer + existing_page_list_envelopes = self._cascading._existing_page_list_envelopes + existing_field_records = self._cascading._existing_field_records + + next_field_id = len(existing_field_records) + + existing_field_names = {fr.field_name for fr in existing_field_records} + for field_name in new_fields: + if field_name in existing_field_names: + raise ValueError(f"Field {field_name!r} already exists in this RNTuple") + + for field_name, field_dtype_raw in new_fields.items(): + ak_form = _type_specification_to_awkward_form(field_dtype_raw) + if not isinstance(ak_form, awkward.forms.NumpyForm): + raise TypeError( + f"add_fields only supports simple numeric types, got {field_dtype_raw!r}" + ) + ak_primitive = ak_form.primitive + type_name = cnt._ak_primitive_to_typename_dict[ak_primitive] + type_num = cnt._ak_primitive_to_num_dict[ak_primitive] + type_size = uproot.const.rntuple_col_num_to_size_dict[type_num] + + if "." in field_name: + parts = field_name.split(".") + actual_field_name = parts[-1] + parent_field_id = None + # walk the full dotted path against the parent-id chain + # for single-level paths (e.g. "parent.field"), match by name only + # for multi-level paths (e.g. "p2.track.field"), walk full chain + if len(parts) == 2: + # single-level: match immediate parent by name only (backward compat) + parent_field_id = None + for i, fr in enumerate(existing_field_records): + if fr.field_name == parts[0]: + parent_field_id = i + break + if parent_field_id is None: + raise ValueError( + f"Field {parts[0]!r} not found in this RNTuple" + ) + fr = existing_field_records[parent_field_id] + else: + # multi-level: walk full dotted path using parent-id chain + current_parent_id = None # None means root + for part_idx, part in enumerate(parts[:-1]): + found = None + for i, fr in enumerate(existing_field_records): + is_root_field = ( + fr.parent_field_id == 0 or fr.parent_field_id == i + ) + if fr.field_name == part: + if current_parent_id is None and is_root_field: + found = (i, fr) + break + elif ( + current_parent_id is not None + and fr.parent_field_id == current_parent_id + ): + found = (i, fr) + break + if found is None: + raise ValueError( + f"Field {'.'.join(parts[:part_idx+1])!r} not found in this RNTuple" + ) + current_parent_id = found[0] + parent_field_id = found[0] + fr = found[1] + if fr.type_name != "": + raise ValueError( + f"Field {'.'.join(parts[:-1])!r} has type {fr.type_name!r} and cannot be extended. " + f"Only untyped records (empty type_name) can have subfields added." + ) + if fr.struct_role != uproot.const.RNTupleFieldRole.RECORD: + raise ValueError( + f"Field {'.'.join(parts[:-1])!r} is not a record and cannot have subfields added." + ) + else: + actual_field_name = field_name + parent_field_id = next_field_id + + new_field = cnt.NTuple_Field_Description( + parent_field_id, + uproot.const.RNTupleFieldRole.LEAF, + actual_field_name, + type_name, + ) + footer.extension_field_record_frames.append(new_field) + if self._column_encoding_error is not None: + raise ValueError(self._column_encoding_error) + # use deferred column — first_element_index marks where new data starts + new_col = cnt.NTuple_Column_Description( + type_num, + type_size, + next_field_id, + 0, + 0, + first_element_index=num_entries, + ) + footer.extension_column_record_frames.append(new_col) + next_field_id += 1 + + footer.cluster_group_record_frames = [] + for cg_idx, cg in enumerate(existing_footer.cluster_group_records): + ple = existing_page_list_envelopes[cg_idx] + all_cluster_page_data = [] + for _cluster_idx, cluster_col_pages in enumerate(ple.pagelinklist): + new_cluster_page_data = [] + for col_pages in cluster_col_pages: + existing_pages = [ + cnt.NTuple_PageDescription( + p.num_elements, + cnt.NTuple_Locator(p.locator.num_bytes, p.locator.offset), + ) + for p in col_pages.pages + ] + new_cluster_page_data.append( + cnt.NTuple_ColumnPageListDescription( + existing_pages, col_pages.element_offset, compression.code + ) + ) + # deferred columns: no pages needed for existing cluster groups + all_cluster_page_data.append(new_cluster_page_data) + cluster_summaries = [ + cnt.NTuple_ClusterSummary(s.num_first_entry, s.num_entries) + for s in ple.cluster_summaries + ] + pagelistenv = cnt.NTuple_PageListEnvelope( + header._checksum, cluster_summaries, all_cluster_page_data + ) + pagelistenv_raw = pagelistenv.serialize() + pagelistenv_key = self._cascading.add_rblob( + self._file.sink, pagelistenv_raw, len(pagelistenv_raw) + ) + pagelistenv_locator = cnt.NTuple_Locator( + len(pagelistenv_raw), + pagelistenv_key.location + pagelistenv_key.allocation, + ) + pagelistenv_envlink = cnt.NTuple_EnvLink( + len(pagelistenv_raw), pagelistenv_locator + ) + footer.cluster_group_record_frames.append( + cnt.NTuple_ClusterGroupRecord( + cg.min_entry_num, + cg.entry_span, + cg.num_clusters, + pagelistenv_envlink, + ) + ) + + footer_raw = footer.serialize() + new_footer_key = self._cascading.add_rblob( + self._file.sink, footer_raw, len(footer_raw) + ) + self._cascading._anchor.seek_footer = ( + new_footer_key.location + new_footer_key.allocation + ) + self._cascading._anchor.nbytes_footer = len(footer_raw) + self._cascading._anchor.len_footer = len(footer_raw) + anchor_raw = self._cascading._anchor.serialize() + self._file.sink.write(self._cascading._anchor._location, anchor_raw) + self._cascading._freesegments.write(self._file.sink) + self._file.sink.flush() + + # update in-memory state without full reload + # update field records and column counts directly from what we just wrote + self._cascading._existing_field_records = list(existing_field_records) + list( + footer.extension_field_record_frames + ) + self._cascading._column_counts = numpy.append( + self._cascading._column_counts, numpy.zeros(len(new_fields), dtype=int) + ) + # reload footer, page list envelopes and akform using _ReadForUpdate + self._file.sink.flush() + + def _get_chunk_reload(start, stop): + raw_bytes = self._file.sink.read(start, stop - start) + return uproot.source.chunk.Chunk.wrap( + _readforupdate_reload, raw_bytes, start=start + ) + + _readforupdate_reload = uproot.writing._cascade._ReadForUpdate( + self._file.file_path, + self._file.uuid, + _get_chunk_reload, + self._file._cascading.tlist_of_streamers, + ) + _readforupdate_reload.options = dict(uproot.reading.open.defaults) + _readforupdate_reload.options["minimal_ttree_metadata"] = False + + _ntuple_key = self._file._cascading.rootdirectory.data.get_key(self._path[-1]) + _raw_reload = self._file.sink.read( + _ntuple_key.seek_location, + _ntuple_key.num_bytes + _ntuple_key.compressed_bytes, + ) + _chunk_reload = uproot.source.chunk.Chunk.wrap( + _readforupdate_reload, _raw_reload, start=_ntuple_key.seek_location + ) + _cursor_reload = uproot.source.cursor.Cursor( + _ntuple_key.seek_location, origin=_ntuple_key.num_bytes + ) + _key_reload = uproot.reading.ReadOnlyKey( + _chunk_reload, + _cursor_reload, + {}, + _readforupdate_reload, + self, + read_strings=True, + ) + existing_reload = _key_reload.get() + _ = existing_reload.keys() + self._cascading._existing_footer = existing_reload._footer + self._cascading._existing_page_list_envelopes = ( + existing_reload.page_list_envelopes + ) + full_akform_reload, _ = existing_reload.to_akform() + self._cascading._header._akform = full_akform_reload + def _is_type_specification(obj): to_check = [obj] diff --git a/tests/test_1687_rntuple_update.py b/tests/test_1687_rntuple_update.py new file mode 100644 index 000000000..4f06c08bd --- /dev/null +++ b/tests/test_1687_rntuple_update.py @@ -0,0 +1,665 @@ +import os +import shutil + +import awkward as ak +import numpy as np +import pytest +import skhep_testdata +import uproot + +try: + import ROOT + + has_root = True +except ImportError: + has_root = False + + +def test_extend_existing_ntuple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = { + "x": np.array([1, 2, 3, 4, 5], dtype=np.float32), + "y": np.array([10, 20, 30, 40, 50], dtype=np.int32), + } + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend( + { + "x": np.array([6, 7, 8], dtype=np.float32), + "y": np.array([60, 70, 80], dtype=np.int32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all( + nt["x"].array() == np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32) + ) + assert ak.all( + nt["y"].array() + == np.array([10, 20, 30, 40, 50, 60, 70, 80], dtype=np.int32) + ) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 8 + entry = reader.CreateEntry() + vals = [] + for i in range(reader.GetNEntries()): + reader.LoadEntry(i, entry) + vals.append(entry["x"]) + assert vals == pytest.approx([1, 2, 3, 4, 5, 6, 7, 8]) + + +def test_add_field_ntuple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3, 4, 5], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"z": np.int32}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all(nt["x"].array() == np.array([1, 2, 3, 4, 5], dtype=np.float32)) + assert ak.all(nt["z"].array() == np.zeros(5, dtype=np.int32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 5 + entry = reader.CreateEntry() + vals = [] + for i in range(reader.GetNEntries()): + reader.LoadEntry(i, entry) + vals.append(entry["x"]) + assert vals == pytest.approx([1.0, 2.0, 3.0, 4.0, 5.0]) + + +def test_add_field_ntuple_duplicate(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="already exists"): + f["mytuple"].add_fields({"x": np.int32}) + + +def test_extend_ntuple_multiple_times(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend({"x": np.array([4, 5, 6], dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend({"x": np.array([7, 8, 9], dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert ak.all( + f["mytuple"]["x"].array() + == np.array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.float32) + ) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 9 + + +def test_add_multiple_fields_ntuple(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"y": np.int32, "z": np.float64}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert ak.all( + f["mytuple"]["x"].array() == np.array([1, 2, 3], dtype=np.float32) + ) + assert ak.all(f["mytuple"]["y"].array() == np.zeros(3, dtype=np.int32)) + assert ak.all(f["mytuple"]["z"].array() == np.zeros(3, dtype=np.float64)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 3 + + +def test_extend_ntuple_wrong_fields(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = { + "x": np.array([1, 2, 3], dtype=np.float32), + "y": np.array([4, 5, 6], dtype=np.int32), + } + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError): + f["mytuple"].extend( + {"x": np.array([7, 8, 9], dtype=np.float32)} + ) # missing y + + +def test_ntuple_dtypes(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = { + "x_f32": np.array([1, 2, 3], dtype=np.float32), + "x_f64": np.array([1, 2, 3], dtype=np.float64), + "x_i32": np.array([1, 2, 3], dtype=np.int32), + "x_i64": np.array([1, 2, 3], dtype=np.int64), + } + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields( + { + "z_f32": np.float32, + "z_f64": np.float64, + "z_i32": np.int32, + "z_i64": np.int64, + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all(nt["x_f32"].array() == np.array([1, 2, 3], dtype=np.float32)) + assert ak.all(nt["z_i64"].array() == np.zeros(3, dtype=np.int64)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 3 + + +def test_ntuple_variable_length(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = { + "x": ak.Array([[1, 2], [3, 4, 5], [6]]), + } + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend({"x": ak.Array([[7, 8, 9], [10]])}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert ak.all( + f["mytuple"]["x"].array() + == ak.Array([[1, 2], [3, 4, 5], [6], [7, 8, 9], [10]]) + ) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 5 + + +def test_ntuple_mixed_types_extend(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = { + "pt": np.array([10.0, 20.0, 30.0], dtype=np.float32), + "jets": ak.Array([[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]]), + } + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend( + { + "pt": np.array([40.0, 50.0], dtype=np.float32), + "jets": ak.Array([[7.0, 8.0, 9.0], [10.0]]), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all( + nt["pt"].array() == np.array([10, 20, 30, 40, 50], dtype=np.float32) + ) + assert ak.all( + nt["jets"].array() == ak.Array([[1, 2], [3], [4, 5, 6], [7, 8, 9], [10]]) + ) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 5 + entry = reader.CreateEntry() + vals = [] + for i in range(reader.GetNEntries()): + reader.LoadEntry(i, entry) + vals.append(entry["pt"]) + assert vals == pytest.approx([10.0, 20.0, 30.0, 40.0, 50.0]) + + +def test_ntuple_add_field_then_extend(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + # add new field (backfilled with zeros) + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"z": np.int32}) + + # now extend with both fields + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend( + { + "x": np.array([4, 5, 6], dtype=np.float32), + "z": np.array([40, 50, 60], dtype=np.int32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all(nt["x"].array() == np.array([1, 2, 3, 4, 5, 6], dtype=np.float32)) + assert ak.all( + nt["z"].array() == np.array([0, 0, 0, 40, 50, 60], dtype=np.int32) + ) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 6 + + +def test_ntuple_extend_empty(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f.mkrntuple("mytuple", {"x": np.dtype("float32")}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend({"x": np.array([1, 2, 3], dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert nt.num_entries == 3 + assert ak.all(nt["x"].array() == np.array([1, 2, 3], dtype=np.float32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 3 + + +def test_ntuple_multiple_in_file(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["tuple1"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + f["tuple2"] = {"y": np.array([4, 5, 6], dtype=np.int32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["tuple1"].extend({"x": np.array([4, 5], dtype=np.float32)}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert ak.all( + f["tuple1"]["x"].array() == np.array([1, 2, 3, 4, 5], dtype=np.float32) + ) + assert ak.all(f["tuple2"]["y"].array() == np.array([4, 5, 6], dtype=np.int32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader1 = ROOT.RNTupleReader.Open("tuple1", os.path.join(tmp_path, "test.root")) + reader2 = ROOT.RNTupleReader.Open("tuple2", os.path.join(tmp_path, "test.root")) + assert reader1.GetNEntries() == 5 + assert reader2.GetNEntries() == 3 + + +def test_ntuple_multiple_add_fields_then_extend(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"y": np.int32}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"z": np.float64}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend( + { + "x": np.array([4, 5, 6], dtype=np.float32), + "y": np.array([40, 50, 60], dtype=np.int32), + "z": np.array([400.0, 500.0, 600.0], dtype=np.float64), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all(nt["x"].array() == np.array([1, 2, 3, 4, 5, 6], dtype=np.float32)) + assert ak.all( + nt["y"].array() == np.array([0, 0, 0, 40, 50, 60], dtype=np.int32) + ) + assert ak.all( + nt["z"].array() == np.array([0, 0, 0, 400, 500, 600], dtype=np.float64) + ) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 6 + + +def test_ntuple_add_field_and_extend_same_session(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"y": np.int32}) + f["mytuple"].extend( + { + "x": np.array([4, 5], dtype=np.float32), + "y": np.array([40, 50], dtype=np.int32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all(nt["x"].array() == np.array([1, 2, 3, 4, 5], dtype=np.float32)) + assert ak.all(nt["y"].array() == np.array([0, 0, 0, 40, 50], dtype=np.int32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 5 + entry = reader.CreateEntry() + vals = [] + for i in range(reader.GetNEntries()): + reader.LoadEntry(i, entry) + vals.append(entry["x"]) + assert vals == pytest.approx([1.0, 2.0, 3.0, 4.0, 5.0]) + + +def test_ntuple_accept_new_fields(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + # should raise without accept_new_fields + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="not in this RNTuple"): + f["mytuple"].extend( + { + "x": np.array([4, 5], dtype=np.float32), + "z": np.array([40, 50], dtype=np.int32), + } + ) + + # with accept_new_fields=True - z backfilled with zeros, then user values + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].extend( + { + "x": np.array([4, 5], dtype=np.float32), + "z": np.array([40, 50], dtype=np.int32), + }, + accept_new_fields=True, + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert ak.all(nt["x"].array() == np.array([1, 2, 3, 4, 5], dtype=np.float32)) + assert ak.all(nt["z"].array() == np.array([0, 0, 0, 40, 50], dtype=np.int32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 5 + entry = reader.CreateEntry() + vals = [] + for i in range(reader.GetNEntries()): + reader.LoadEntry(i, entry) + vals.append(entry["x"]) + assert vals == pytest.approx([1.0, 2.0, 3.0, 4.0, 5.0]) + + +def test_ntuple_add_subfield(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = ak.Array( + [ + {"particle": {"pt": 1.0, "eta": 2.0}}, + {"particle": {"pt": 3.0, "eta": 4.0}}, + ] + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"particle.phi": np.float32}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert "particle.phi" in nt.keys() + assert ak.all(nt["particle.phi"].array() == np.zeros(2, dtype=np.float32)) + assert ak.all(nt["particle"].array().pt == np.array([1.0, 3.0])) + assert ak.all(nt["particle"].array().phi == np.zeros(2, dtype=np.float32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 2 + + +def test_ntuple_add_nested_subfield(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = ak.Array( + [ + {"particle": {"track": {"pt": 1.0}}}, + {"particle": {"track": {"pt": 3.0}}}, + ] + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"particle.track.phi": np.float32}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert "particle.track.phi" in nt.keys() + assert ak.all(nt["particle.track.phi"].array() == np.zeros(2, dtype=np.float32)) + + if has_root and hasattr(ROOT, "RNTupleReader"): + reader = ROOT.RNTupleReader.Open("mytuple", os.path.join(tmp_path, "test.root")) + assert reader.GetNEntries() == 2 + + +def test_ntuple_add_subfield_nonexistent_parent(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="not found"): + f["mytuple"].add_fields({"nonexistent.phi": np.float32}) + + +def test_ntuple_add_subfield_typed_parent(tmp_path): + # structs written by ROOT have C++ typenames and cannot have subfields added + src = skhep_testdata.data_path("test_nested_structs_rntuple_v1-0-0-0.root") + shutil.copy(src, os.path.join(tmp_path, "test.root")) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + # my_struct has typename TopStruct + with pytest.raises(ValueError, match="TopStruct"): + f["ntuple"].add_fields({"my_struct.new_field": np.float32}) + + # sub_struct has typename SubStruct + with pytest.raises(ValueError, match="SubStruct"): + f["ntuple"].add_fields({"sub_struct.new_field": np.float32}) + + # sub_sub_struct has typename SubSubSruct + with pytest.raises(ValueError, match="SubSubSruct"): + f["ntuple"].add_fields({"sub_sub_struct.new_field": np.float32}) + + +def test_ntuple_add_subfield_to_collection(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"jets": ak.Array([[1.0, 2.0], [3.0]])} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="type"): + f["mytuple"].add_fields({"jets.x": np.float32}) + + +def test_ntuple_add_subfield_to_collection_of_records(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = ak.Array( + {"jets": [[{"pt": 1.0, "eta": 0.0}], [{"pt": 2.0, "eta": 3.0}]]} + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="not a record"): + f["mytuple"].add_fields({"jets.phi": np.float32}) + + +def test_ntuple_add_subfield_to_variant(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = ak.Array( + {"variant": ak.Array([{"jet": {"pt": 1.0, "eta": 2.0}}, 2])} + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises((ValueError, AssertionError)): + f["mytuple"].add_fields({"variant.jet.eta": np.float32}) + + +def test_ntuple_num_entries(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + assert nt.num_entries == 3 + + +def test_ntuple_add_field_duplicate_after_extension(tmp_path): + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"z": np.int32}) + + # try to add z again - should fail even though it's an extension field + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="already exists"): + f["mytuple"].add_fields({"z": np.float32}) + + +def test_ntuple_extend_root_written_raises(tmp_path): + # ROOT-written RNTuples use split encoding which uproot cannot write + src = skhep_testdata.data_path("test_int_float_rntuple_v1-0-0-0.root") + shutil.copy(src, os.path.join(tmp_path, "test.root")) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="column encodings"): + f["ntuple"].extend( + { + "one_integers": np.array([100, 200], dtype=np.int32), + "two_floats": np.array([1.5, 2.5], dtype=np.float32), + } + ) + + +def test_ntuple_add_fields_multi_cluster_groups(tmp_path): + # multiple cluster groups (one cluster each) — should work + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["ntuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["ntuple"].extend({"x": np.array([4, 5, 6], dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["ntuple"].extend({"x": np.array([7, 8, 9], dtype=np.float32)}) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["ntuple"].add_fields({"y": np.int32}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert f["ntuple"].num_entries == 9 + assert np.all(f["ntuple"]["x"].array() == np.arange(1, 10, dtype=np.float32)) + assert np.all(f["ntuple"]["y"].array() == 0) + + +def test_ntuple_add_fields_sequential_same_session(tmp_path): + # test that holding onto a WritableNTuple object and calling add_fields twice works + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = ak.Array([{"x": 1.0}, {"x": 2.0}, {"x": 3.0}]) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + nt.add_fields({"y": np.int32}) + nt.add_fields({"z": np.float64}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert set(f["mytuple"].keys()) == {"x", "y", "z"} + assert np.all(f["mytuple"]["y"].array() == 0) + assert np.all(f["mytuple"]["z"].array() == 0.0) + + +def test_ntuple_add_fields_then_extend_same_object(tmp_path): + # test add_fields then extend using same nt object (not re-opening f["mytuple"]) + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] + nt.add_fields({"y": np.int32}) + nt.extend( + { + "x": np.array([4, 5], dtype=np.float32), + "y": np.array([40, 50], dtype=np.int32), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert ak.all( + f["mytuple"]["x"].array() == np.array([1, 2, 3, 4, 5], dtype=np.float32) + ) + assert ak.all( + f["mytuple"]["y"].array() == np.array([0, 0, 0, 40, 50], dtype=np.int32) + ) + + +def test_ntuple_add_subfield_correct_parent(tmp_path): + # verify p2.track.phi goes to p2.track not p1.track + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = ak.Array( + [{"p1": {"track": {"pt": 1.0}}, "p2": {"track": {"pt": 2.0}}}] * 2 + ) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"].add_fields({"p2.track.phi": np.float32}) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + keys = list(f["mytuple"].keys()) + assert "p2.track.phi" in keys + assert "p1.track.phi" not in keys + + +def test_ntuple_update_root_written_file_opens(tmp_path): + # ROOT-written files use UUID versions other than 1 — verify uproot.update can open them + # This test uses a ROOT-written file and checks it can be accessed in update mode + src = skhep_testdata.data_path("test_int_float_rntuple_v1-0-0-0.root") + shutil.copy(src, os.path.join(tmp_path, "test.root")) + + # should not raise — previously failed with assert uuid_version == 1 + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + # accessing the ntuple should work (even though we can't extend ROOT-written files) + with pytest.raises(ValueError, match="column encodings"): + f["ntuple"].extend( + { + "one_integers": np.array([1], dtype=np.int32), + "two_floats": np.array([1.0], dtype=np.float32), + } + ) + + +def test_ntuple_root_written_add_fields_raises(tmp_path): + # ROOT-written files use split encoding which uproot cannot write + # add_fields should raise a clear error not corrupt data + src = skhep_testdata.data_path("ntpl001_staff_rntuple_v1-0-0-0.root") + shutil.copy(src, os.path.join(tmp_path, "test.root")) + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + with pytest.raises(ValueError, match="column encodings"): + f["Staff"].add_fields({"new_field": np.float32}) + + +def test_ntuple_hold_object_across_operations(tmp_path): + # hold WritableNTuple object across add_fields and extend + with uproot.recreate(os.path.join(tmp_path, "test.root")) as f: + f["mytuple"] = {"x": np.array([1, 2, 3], dtype=np.float32)} + + with uproot.update(os.path.join(tmp_path, "test.root")) as f: + nt = f["mytuple"] # hold the object + nt.add_fields({"y": np.int32}) + nt.add_fields({"z": np.float64}) + nt.extend( + { + "x": np.array([4, 5], dtype=np.float32), + "y": np.array([10, 20], dtype=np.int32), + "z": np.array([1.1, 2.2], dtype=np.float64), + } + ) + + with uproot.open(os.path.join(tmp_path, "test.root")) as f: + assert set(f["mytuple"].keys()) == {"x", "y", "z"} + assert np.all( + f["mytuple"]["x"].array() == np.array([1, 2, 3, 4, 5], dtype=np.float32) + ) + assert np.all( + f["mytuple"]["y"].array() == np.array([0, 0, 0, 10, 20], dtype=np.int32) + )