diff --git a/Project.toml b/Project.toml index 5769720f..eb51ffde 100644 --- a/Project.toml +++ b/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" diff --git a/src/root.jl b/src/root.jl index 8f14a24c..7ca3fdd9 100644 --- a/src/root.jl +++ b/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" && return Vector{Float64}, _jaggtype + streamer.fTypeName == "vector" && return Vector{Int32}, _jaggtype + elseif _jaggtype == Offsetjaggjagg || _jaggtype == Offset6jaggjagg + streamer.fTypeName == "vector" && return Vector{Vector{Float64}}, _jaggtype + streamer.fTypeName == "vector" && return Vector{Vector{Int32}}, _jaggtype + end + + end + + # some standard cases m = match(r"vector<(.*)>", classname) if m!==nothing elname = m[1] diff --git a/src/utils.jl b/src/utils.jl index 71225063..7fde645f 100644 --- a/src/utils.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index ab2aae6b..78d971f3 100644 --- a/test/runtests.jl +++ b/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") diff --git a/test/samples/triply_jagged_via_custom_class.root b/test/samples/triply_jagged_via_custom_class.root new file mode 100644 index 00000000..6e1986bd Binary files /dev/null and b/test/samples/triply_jagged_via_custom_class.root differ