Improves streamer logic for multiple jaggednes (#231)

* Improves streamer logic

* Support for doubly/triply jagged int/double

* Bump version

* Fix jaggedness for BASE

Project.toml

@@ -1,7 +1,7 @@
 name = "UnROOT"
 uuid = "3cd96dde-e98d-4713-81e9-a4a1b0235ce9"
 authors = ["Tamas Gal", "Jerry Ling", "Johannes Schumann", "Nick Amin"]
-version = "0.10.0"
+version = "0.10.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/root.jl

@@ -130,7 +130,23 @@ function streamerfor(f::ROOTFile, name::AbstractString)
             return e
         end
     end
-    error("No streamer found for $name.")
+    missing
+end
+
+streamerfor(f::ROOTFile, branch::TBranch) = missing
+function streamerfor(f::ROOTFile, branch::TBranchElement)
+    fID = branch.fID
+    # According to ChatGPt: When fID is equal to -1, it means that the
+    # TBranch object has not been registered yet in the TTree's list of
+    # branches. This can happen, for example, when a TBranch object has been
+    # created, but has not been added to a TTree with the TTree::Branch()
+    # method.
+    #
+    # TODO: For now, we force it to be 0 in this case, until someone complains.
+    if fID == -1
+        fID = 0
+    end
+    streamerfor(f, branch.fClassName).streamer.fElements.elements[fID + 1]  # one-based indexing in Julia
 end
 
 
@@ -226,8 +242,12 @@ function interped_data(rawdata, rawoffsets, ::Type{T}, ::Type{J}) where {T, J<:J
     # we want the fundamental type as `reinterpret` will create vector
     if J === Nojagg
         return ntoh.(reinterpret(T, rawdata))
-    elseif J === Offsetjaggjagg # the branch is doubly jagged
-        jagg_offset = 10
+    elseif J === Offsetjaggjagg || J === Offset6jaggjagg # the branch is doubly jagged
+        if J === Offset6jaggjagg
+            jagg_offset = 6
+        else
+            jagg_offset = 10
+        end
         subT = eltype(eltype(T))
         out = VectorOfVectors(T(), Int32[1])
         @views for i in 1:(length(rawoffsets)-1)
@@ -346,6 +366,22 @@ function auto_T_JaggT(f::ROOTFile, branch; customstructs::Dict{String, Type})
             return _custom, Nojagg
         catch
         end
+
+        # check if we have an actual streamer
+        streamer = streamerfor(f, branch)
+        if !ismissing(streamer)
+            # TODO unify this with the "switch" block below and expand for more types!
+            if _jaggtype == Offsetjagg
+                streamer.fTypeName == "vector<double>" && return Vector{Float64}, _jaggtype
+                streamer.fTypeName == "vector<int>" && return Vector{Int32}, _jaggtype
+            elseif _jaggtype == Offsetjaggjagg || _jaggtype == Offset6jaggjagg
+                streamer.fTypeName == "vector<double>" && return Vector{Vector{Float64}}, _jaggtype
+                streamer.fTypeName == "vector<int>" && return Vector{Vector{Int32}}, _jaggtype
+            end
+
+        end
+
+        # some standard cases
         m = match(r"vector<(.*)>", classname)
         if m!==nothing
             elname = m[1]

src/utils.jl

@@ -55,32 +55,33 @@ struct Nojagg      <:JaggType  end
 struct Nooffsetjagg<:JaggType  end
 struct Offsetjagg  <:JaggType  end
 struct Offsetjaggjagg  <:JaggType  end
+# this is a preliminary workaround for 6 byte offset jaggedness
+struct Offset6jaggjagg  <:JaggType  end
 
 function JaggType(f, branch, leaf)
     # https://github.com/scikit-hep/uproot3/blob/54f5151fb7c686c3a161fbe44b9f299e482f346b/uproot3/interp/auto.py#L144
 
-    try
-        fID = branch.fID
-        # According to ChatGPt: When fID is equal to -1, it means that the
-        # TBranch object has not been registered yet in the TTree's list of
-        # branches. This can happen, for example, when a TBranch object has been
-        # created, but has not been added to a TTree with the TTree::Branch()
-        # method.
-        #
-        # TODO: For now, we force it to be 0 in this case, until someone complains.
-        if fID == -1
-            fID = 0
+    streamer = streamerfor(f, branch)
+    if !ismissing(streamer)
+        if typeof(streamer) <: TStreamerBasicType
+            (match(r"\[.*\]", leaf.fTitle) !== nothing) && return Nooffsetjagg
+            return Nojagg
+        end
+        if typeof(streamer) <: TStreamerBase
+            leaf isa TLeafElement && leaf.fLenType==0 && return Offsetjagg
+            return Nojagg
+        end
+        if streamer.fSTLtype == Const.kSTLvector
+            (match(r"\[.*\]", leaf.fTitle) !== nothing) && return Offset6jaggjagg
+            return Offsetjagg
         end
-        streamer = streamerfor(f, branch.fClassName).streamer.fElements.elements[fID + 1]  # one-based indexing in Julia
-        (streamer.fSTLtype == Const.kSTLvector) && return Offsetjagg
-    catch
     end
 
+    # TODO this might be redundant but for now it works
     (match(r"\[.*\]", leaf.fTitle) !== nothing) && return Nooffsetjagg
     leaf isa TLeafElement && leaf.fLenType==0 && return Offsetjagg
     !hasproperty(branch, :fClassName) && return Nojagg
 
-
     return Nojagg
 end
 

test/runtests.jl

@@ -374,6 +374,36 @@ end
                        ]
 end
 
+@testset "Doubly jagged via custom class" begin
+    f = UnROOT.samplefile("triply_jagged_via_custom_class.root")
+    b = LazyBranch(f, "E/Evt/w")
+    @test 3 == length(b)
+    @test [3.3329158e6, 3.6047424e22, 3.1181261e8] ≈ b[1]
+    @test [3.3329158e6, 2.4039883e23, 4.5337845e6] ≈ b[2]
+    @test [3.3329158e6, 4.9458148e24, 42805.383] ≈ b[3]
+    close(f)
+end
+
+@testset "Triply jagged stuff via custom class" begin
+    f = UnROOT.samplefile("triply_jagged_via_custom_class.root")
+    b = LazyBranch(f, "E/Evt/trks/trks.rec_stages")
+    @test 3 == length(b)
+    @test 38 == length(b[3])
+    @test 5 == length(b[3][1])
+    @test [1, 2, 3, 4, 5] == b[2][1]
+    b = LazyBranch(f, "E/Evt/trks/trks.fitinf")
+    @test 3 == length(b)
+    @test 38 == length(b[3])
+    @test 17 == length(b[3][1])
+    @test b[3][1] ≈ [
+        0.02175229704756278, 0.013640169301181978, -27.59387722710992, 45.0,
+        90072.00780343144, 52.67760965729974, 3.9927948910593525, 10.0,
+        0.3693832125931179, 0.3693832125931179, 72.02530060292972, 0.0, 0.0,
+        13461.937901498717, -Inf, 1518.0, 56.0,
+    ]
+    close(f)
+end
+
 @testset "NanoAOD" begin
     rootfile = ROOTFile(joinpath(SAMPLES_DIR, "NanoAODv5_sample.root"))
     event = UnROOT.array(rootfile, "Events/event")