fix: TTree-writing bugs (sliced string arrays, nested recarrays, fLen, TGraph y-range, oversized compression blocks)

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]
 

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

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")

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)}"

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