diff --git a/src/uproot/compression.py b/src/uproot/compression.py index 9ca041792..c007e056c 100644 --- a/src/uproot/compression.py +++ b/src/uproot/compression.py @@ -557,6 +557,16 @@ def compress(data: bytes, compression: Compression) -> bytes: else: checksum = b"" + # An incompressible block can compress to *more* than the input (zlib + # adds ~5 bytes/16 KiB, lz4 adds a header plus an 8-byte checksum). If + # the compressed size doesn't fit in the 3-byte header field, the + # high byte would be silently truncated, corrupting the block. ROOT + # decides whether to compress on a whole-buffer basis, so we fall back + # to storing ``data`` uncompressed (as the final guard below does) when + # any block fails to shrink or would overflow the 3-byte size field. + if len_compressed >= len(block) or len_compressed > _3BYTE_MAX: + return data + uncompressed_size = _4byte.pack(len(block))[:-1] compressed_size = _4byte.pack(len_compressed)[:-1] diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index 6b80763d3..18f0fddd8 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -573,7 +573,7 @@ def extend(self, file, sink, data): if k in provided and not numpy.array_equal(v, provided[k]): raise ValueError( - f"branch {kk!r} provided both as an explicit array and generated as a counter, and they disagree" + f"branch {k!r} provided both as an explicit array and generated as a counter, and they disagree" ) provided[k] = v @@ -646,9 +646,17 @@ def extend(self, file, sink, data): layout = awkward.to_layout(branch_array) if isinstance( layout, - (awkward.contents.ListArray, awkward.contents.RegularArray), + ( + awkward.contents.ListArray, + awkward.contents.RegularArray, + awkward.contents.ListOffsetArray, + ), ): - layout = layout.to_ListOffsetArray64() + # start_at_zero=True normalizes offsets[0] to 0 and + # trims the content, which is required for a sliced + # array (e.g. ak.Array([...])[2:]) to produce a basket + # that reads back correctly. + layout = layout.to_ListOffsetArray64(True) lengths = numpy.asarray(awkward.num(layout)) which_big = lengths >= 255 @@ -693,7 +701,7 @@ def extend(self, file, sink, data): ) else: big_endian = uproot._util.ensure_numpy(branch_array).astype( - datum["dtype"] + datum["dtype"], copy=False ) if big_endian.shape != (len(branch_array),) + datum["shape"]: raise ValueError( @@ -776,9 +784,12 @@ def extend(self, file, sink, data): datum = self._branch_data[self._branch_lookup[branch_name]] if datum["dtype"] == ">U0": - totbytes, zipbytes, location = self.write_string_basket( + totbytes, zipbytes, location, fLen = self.write_string_basket( sink, branch_name, compression, big_endian, big_endian_offsets ) + # fLen is the size of the biggest string; accumulate the max + # across extends and keep it per-branch (not shared tree-wide). + datum["fLen"] = max(datum.get("fLen", 0), fLen) datum["fEntryOffsetLen"] = 4 * (len(big_endian_offsets) - 1) elif big_endian_offsets is None: @@ -1310,7 +1321,7 @@ def write_updates(self, sink): sink.write( position, uproot.models.TLeaf._tleaf2_format0.pack( - self._metadata["fLen"], + datum.get("fLen", 0), datum["dtype"].itemsize, 0, datum["kind"] == "counter", @@ -1519,12 +1530,8 @@ def write_string_basket(self, sink, branch_name, compression, array, offsets): ) compressed_data = uproot.compression.compress(uncompressed_data, compression) - # get size of biggest string - self._metadata["fLen"] = ( - 0 - if len(offsets) == 1 - else max([offsets[i + 1] - offsets[i] for i in range(len(offsets) - 1)]) - ) + # get size of biggest string in this basket + fLen = 0 if len(offsets) <= 1 else int(numpy.diff(offsets).max()) fLast = offsets[-1] offsets[-1] = 0 @@ -1578,7 +1585,7 @@ def write_string_basket(self, sink, branch_name, compression, array, offsets): sink.set_file_length(self._freesegments.fileheader.end) sink.flush() - return fKeylen + fObjlen, fNbytes, location + return fKeylen + fObjlen, fNbytes, location, fLen _tbasket_offsets_length = struct.Struct(">I") @@ -1602,7 +1609,7 @@ def recarray_to_dict(array): for field_name in array.dtype.fields: field = array[field_name] if field.dtype.fields is not None: - for subfield_name, subfield in recarray_to_dict(field): + for subfield_name, subfield in recarray_to_dict(field).items(): out[field_name + "." + subfield_name] = subfield else: out[field_name] = field diff --git a/src/uproot/writing/identify.py b/src/uproot/writing/identify.py index 745b6bc70..83ecdd62b 100644 --- a/src/uproot/writing/identify.py +++ b/src/uproot/writing/identify.py @@ -502,9 +502,7 @@ def to_writable(obj): fNbins = len(edges) - 1 fXmin, fXmax = edges[0], edges[-1] - if numpy.allclose( - edges, numpy.linspace(fXmin, fXmax, len(edges), edges.dtype) - ): + if numpy.allclose(edges, numpy.linspace(fXmin, fXmax, len(edges))): edges = numpy.array([], dtype=">f8") else: edges = edges.astype(">f8") @@ -551,18 +549,14 @@ def to_writable(obj): fXaxis_fNbins = len(xedges) - 1 fXmin, fXmax = xedges[0], xedges[-1] - if numpy.allclose( - xedges, numpy.linspace(fXmin, fXmax, len(xedges), xedges.dtype) - ): + if numpy.allclose(xedges, numpy.linspace(fXmin, fXmax, len(xedges))): xedges = numpy.array([], dtype=">f8") else: xedges = xedges.astype(">f8") fYaxis_fNbins = len(yedges) - 1 fYmin, fYmax = yedges[0], yedges[-1] - if numpy.allclose( - yedges, numpy.linspace(fYmin, fYmax, len(yedges), yedges.dtype) - ): + if numpy.allclose(yedges, numpy.linspace(fYmin, fYmax, len(yedges))): yedges = numpy.array([], dtype=">f8") else: yedges = yedges.astype(">f8") @@ -625,27 +619,21 @@ def to_writable(obj): fXaxis_fNbins = len(xedges) - 1 fXmin, fXmax = xedges[0], xedges[-1] - if numpy.allclose( - xedges, numpy.linspace(fXmin, fXmax, len(xedges), xedges.dtype) - ): + if numpy.allclose(xedges, numpy.linspace(fXmin, fXmax, len(xedges))): xedges = numpy.array([], dtype=">f8") else: xedges = xedges.astype(">f8") fYaxis_fNbins = len(yedges) - 1 fYmin, fYmax = yedges[0], yedges[-1] - if numpy.allclose( - yedges, numpy.linspace(fYmin, fYmax, len(yedges), yedges.dtype) - ): + if numpy.allclose(yedges, numpy.linspace(fYmin, fYmax, len(yedges))): yedges = numpy.array([], dtype=">f8") else: yedges = yedges.astype(">f8") fZaxis_fNbins = len(zedges) - 1 fZmin, fZmax = zedges[0], zedges[-1] - if numpy.allclose( - zedges, numpy.linspace(fZmin, fZmax, len(zedges), zedges.dtype) - ): + if numpy.allclose(zedges, numpy.linspace(fZmin, fZmax, len(zedges))): zedges = numpy.array([], dtype=">f8") else: zedges = zedges.astype(">f8") @@ -703,7 +691,7 @@ def _fXbins_maybe_regular(axis, boost_histogram): if boost_histogram is None: edges = axis.edges fXmin, fXmax = edges[0], edges[-1] - if numpy.allclose(edges, numpy.linspace(fXmin, fXmax, len(edges), edges.dtype)): + if numpy.allclose(edges, numpy.linspace(fXmin, fXmax, len(edges))): return numpy.array([], dtype=">f8") else: return edges.astype(">f8") diff --git a/src/uproot/writing/interpret.py b/src/uproot/writing/interpret.py index ad40e68dd..38ea2d2f7 100644 --- a/src/uproot/writing/interpret.py +++ b/src/uproot/writing/interpret.py @@ -139,7 +139,7 @@ def _as_TGraph( raise ValueError(f"y has to be 1D, but is {len(y.shape)}D!") if minY is None: - new_minY = np.min(x) + new_minY = np.min(y) elif not isinstance(minY, numbers.Real): raise ValueError( f"uproot.as_TGraph minY has to be None or a number, not {type(minY)}" @@ -148,7 +148,7 @@ def _as_TGraph( new_minY = minY if maxY is None: - new_maxY = np.max(x) + new_maxY = np.max(y) elif not isinstance(maxY, numbers.Real): raise ValueError( f"uproot.as_TGraph minY has to be None or a number, not {type(maxY)}" diff --git a/tests/test_1657_fix_ttree_writing_bugs.py b/tests/test_1657_fix_ttree_writing_bugs.py new file mode 100644 index 000000000..7b039bddf --- /dev/null +++ b/tests/test_1657_fix_ttree_writing_bugs.py @@ -0,0 +1,173 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE + +import os + +import numpy as np +import pytest + +import uproot +import uproot.compression + +awkward = pytest.importorskip("awkward") + + +def test_sliced_string_branch_roundtrip(tmp_path): + # A sliced string array has offsets[0] != 0; the string-branch path of + # Tree.extend must normalize this or the basket is unreadable. + newfile = os.path.join(tmp_path, "sliced_string.root") + + big = "x" * 300 # also exercise the >= 255 length-extension path + cases = [ + awkward.Array(["aaa", "bb", "cccc", "dd", "eeeee"])[2:], + awkward.Array(["aaa", big, "bb", "cccc", "dd"])[2:], + awkward.Array([big, "aaa", "bb", "cccc"])[1:], + awkward.Array(["aaa", "bb", "cccc", "dd", "eeeee"]), # unsliced still ok + ] + + for i, arr in enumerate(cases): + fn = os.path.join(tmp_path, f"sliced_string_{i}.root") + with uproot.recreate(fn) as f: + tree = f.mktree("t", {"s": "string"}) + tree.extend({"s": arr}) + with uproot.open(fn) as f: + assert f["t/s"].array().tolist() == arr.tolist() + + +def test_string_branch_flen_per_branch(tmp_path): + # TLeafC fLen must be tracked per-branch and accumulated as a max across + # extends, not shared at the tree level or overwritten each extend. + newfile = os.path.join(tmp_path, "flen.root") + + a1 = awkward.Array(["x" * 24, "yy", "z"]) # max length 24 + b1 = awkward.Array(["aaa", "b", "cc"]) # max length 3 + a2 = awkward.Array(["p", "q", "r"]) + b2 = awkward.Array(["x" * 40, "y", "z"]) # bigger on second extend + + with uproot.recreate(newfile) as f: + tree = f.mktree("t", {"s1": "string", "s2": "string"}) + tree.extend({"s1": a1, "s2": b1}) + tree.extend({"s1": a2, "s2": b2}) + + with uproot.open(newfile) as f: + # fLen includes the 1-byte length prefix + assert f["t/s1"].member("fLeaves")[0].member("fLen") == 25 + assert f["t/s2"].member("fLeaves")[0].member("fLen") == 41 + assert f["t/s1"].array().tolist() == a1.tolist() + a2.tolist() + assert f["t/s2"].array().tolist() == b1.tolist() + b2.tolist() + + +def test_recarray_to_dict_nested_struct(tmp_path): + from uproot.writing._cascadetree import recarray_to_dict + + dt = np.dtype([("a", np.int32), ("nested", [("x", np.float64), ("y", np.int16)])]) + arr = np.zeros(5, dtype=dt) + arr["a"] = np.arange(5) + arr["nested"]["x"] = np.arange(5) * 1.5 + arr["nested"]["y"] = np.arange(5) + + out = recarray_to_dict(arr) + assert set(out) == {"a", "nested.x", "nested.y"} + assert out["a"].tolist() == [0, 1, 2, 3, 4] + assert out["nested.x"].tolist() == [0.0, 1.5, 3.0, 4.5, 6.0] + assert out["nested.y"].tolist() == [0, 1, 2, 3, 4] + + +def test_tgraph_default_y_range_uses_y(): + x = np.array([0.0, 1.0, 2.0, 3.0]) + y = np.array([100.0, 200.0, 150.0, 175.0]) + g = uproot.as_TGraph({"x": x, "y": y}) + + span = 200.0 - 100.0 + assert g.member("fMinimum") == pytest.approx(100.0 - 0.1 * span) + assert g.member("fMaximum") == pytest.approx(200.0 + 0.1 * span) + + +def test_histogram_regular_edge_detection(tmp_path): + # The numpy.linspace regularity check must still collapse regular edges to + # an empty fXbins and keep irregular edges. + newfile = os.path.join(tmp_path, "hist.root") + + entries = np.ones(10) + regular = np.linspace(0, 10, 11) + irregular = np.array([0, 1, 2, 4, 8, 16, 17, 18, 19, 20, 21], dtype=float) + + with uproot.recreate(newfile) as f: + f["hreg"] = (entries, regular) + f["hirr"] = (entries, irregular) + + with uproot.open(newfile) as f: + assert len(f["hreg"].member("fXaxis").member("fXbins")) == 0 + assert len(f["hirr"].member("fXaxis").member("fXbins")) == len(irregular) + + +def test_compress_incompressible_block_no_corruption(): + # An incompressible block can compress to more than 2**24-1 bytes; the + # 3-byte header field must not be silently truncated. Mixing a compressible + # block with an incompressible one used to produce a buffer that was + # smaller than the input (so it was kept) but had a corrupt block header. + # The fix falls back to storing the data uncompressed in that case. + from uproot.compression import ( + ZLIB, + LZ4, + ZSTD, + compress, + decompress, + _3BYTE_MAX, + ) + from uproot.source.chunk import Chunk + from uproot.source.cursor import Cursor + + rng = np.random.default_rng(12345) + block1 = np.zeros(_3BYTE_MAX, dtype=np.uint8).tobytes() + block2 = rng.integers(0, 256, size=_3BYTE_MAX, dtype=np.uint8).tobytes() + data = block1 + block2 + + for comp in [ZLIB(1), LZ4(1), ZSTD(1)]: + out = compress(data, comp) + # The fix returns the input unchanged (stored uncompressed) instead of + # emitting a buffer with a truncated 3-byte block header. + assert out is data + + # Any compressed buffer that compress() *does* return must read back. + compressible = (np.arange(_3BYTE_MAX + 1000, dtype=np.uint8) % 4).tobytes() + for comp in [ZLIB(1), LZ4(1), ZSTD(1)]: + out = compress(compressible, comp) + chunk = Chunk.wrap(None, np.frombuffer(out, dtype=np.uint8)) + result = decompress(chunk, Cursor(0), {}, len(out), len(compressible)) + assert np.asarray(result.raw_data).tobytes() == compressible + + +def test_write_incompressible_branch_roundtrip(tmp_path): + # End-to-end: writing ~17 MB of incompressible random bytes (more than one + # compression block) into a compressed TTree branch must read back exactly. + newfile = os.path.join(tmp_path, "incompressible.root") + + rng = np.random.default_rng(20240610) + values = rng.integers( + np.iinfo(np.int64).min, + np.iinfo(np.int64).max, + size=17_000_000 // 8, + dtype=np.int64, + ) + + with uproot.recreate(newfile, compression=uproot.ZLIB(1)) as f: + f.mktree("t", {"x": np.int64}) + f["t"].extend({"x": values}) + + with uproot.open(newfile) as f: + np.testing.assert_array_equal(f["t/x"].array(library="np"), values) + + +def test_compress_still_compresses_compressible_data(): + from uproot.compression import ZLIB, LZ4, ZSTD, compress, decompress + from uproot.source.chunk import Chunk + from uproot.source.cursor import Cursor + + data = (np.arange(1_000_000, dtype=np.int64) % 7).tobytes() + + for comp in [ZLIB(1), LZ4(1), ZSTD(1)]: + out = compress(data, comp) + assert len(out) < len(data) + chunk = Chunk.wrap(None, np.frombuffer(out, dtype=np.uint8)) + result = decompress(chunk, Cursor(0), {}, len(out), len(data)) + assert np.asarray(result.raw_data).tobytes() == data