add broadcasting fusion (#126)

* add broadcasting fusion

* fix allocation and let getindex be lazy

* better lazyevent display

* fix CI

* add tree displaying tests

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.6.3"
+version = "0.7.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/displays.jl

@@ -55,6 +55,7 @@ printnode(io::IO, k::TKeyNode) = print(io, "$(k.name) ($(k.classname))")
 
 Base.show(tree::LazyTree; kwargs...) = _show(stdout, tree; crop=:both, kwargs...)
 Base.show(io::IO, tree::LazyTree; kwargs...) = _show(io, tree; kwargs...)
+Base.show(io::IO, ::MIME"text/plain", tree::LazyTree) = _show(io, tree)
 function _show(io::IO, tree::LazyTree; kwargs...)
     _hs = _make_header(tree)
     _ds = displaysize(io)

src/iteration.jl

@@ -179,14 +179,18 @@ function _localindex_newbasket!(ba::LazyBranch{T,J,B}, idx::Integer, tid::Int) w
     return idx - br.start + 1
 end
 
-Base.IndexStyle(::Type{<:LazyBranch}) = IndexCartesian()
+Base.IndexStyle(::Type{<:LazyBranch}) = IndexLinear()
 
 function Base.iterate(ba::LazyBranch{T,J,B}, idx=1) where {T,J,B}
     idx > ba.L && return nothing
     return (ba[idx], idx + 1)
 end
 
-struct LazyTree{T}
+struct LazyEvent{T<:TypedTables.Table}
+    tree::T
+    idx::Int64
+end
+struct LazyTree{T} <: AbstractVector{LazyEvent{T}}
     treetable::T
 end
 
@@ -195,15 +199,18 @@ end
 Base.propertynames(lt::LazyTree) = propertynames(innertable(lt))
 Base.getproperty(lt::LazyTree, s::Symbol) = getproperty(innertable(lt), s)
 
-# a specific branch
-Base.getindex(lt::LazyTree, row::Int) = innertable(lt)[row]
+Base.broadcastable(lt::LazyTree) = lt
+Base.IndexStyle(::Type{<:LazyTree}) = IndexLinear()
+Base.getindex(lt::LazyTree, row::Int) = LazyEvent(innertable(lt), row)
+# kept lazy for broadcasting purpose
+Base.getindex(lt::LazyTree, row::CartesianIndex{1}) = LazyEvent(innertable(lt), row[1])
 function Base.getindex(lt::LazyTree, rang::UnitRange)
     return LazyTree(innertable(lt)[rang])
 end
-Base.getindex(lt::LazyTree, ::typeof(!), s::Symbol) = lt[:, s]
-Base.getindex(lt::LazyTree, ::Colon, s::Symbol) = getproperty(innertable(lt), s) # the real deal
 
 # a specific event
+Base.getindex(lt::LazyTree, ::typeof(!), s::Symbol) = lt[:, s]
+Base.getindex(lt::LazyTree, ::Colon, s::Symbol) = getproperty(innertable(lt), s) # the real deal
 Base.getindex(lt::LazyTree, row::Int, col::Symbol) = lt[:, col][row]
 Base.getindex(lt::LazyTree, rows::UnitRange, col::Symbol) = lt[:, col][rows]
 Base.getindex(lt::LazyTree, ::Colon) = lt[1:end]
@@ -220,6 +227,8 @@ Base.getindex(e::Iterators.Enumerate{LazyTree{T}}, row::Int) where T = (row, fir
 # interfacing Table
 Base.names(lt::LazyTree) = collect(String.(propertynames(innertable(lt))))
 Base.length(lt::LazyTree) = length(innertable(lt))
+Base.ndims(::Type{<:LazyTree}) = 1
+Base.size(lt::LazyTree) = size(innertable(lt))
 
 function getbranchnamesrecursive(obj)
     out = Vector{String}()
@@ -282,17 +291,15 @@ function LazyTree(f::ROOTFile, s::AbstractString, branch::Union{AbstractString,R
     return LazyTree(f, s, [branch])
 end
 
-struct LazyEvent{T<:TypedTables.Table}
-    tree::T
-    idx::Int64
-end
 function Base.show(io::IO, evt::LazyEvent)
     idx = Core.getfield(evt, :idx)
     fields = propertynames(Core.getfield(evt, :tree))
     nfields = length(fields)
     sfields = nfields < 20 ? ": $(fields)" : ""
-    show(io, "LazyEvent $(idx) with $(nfields) fields$(sfields)")
+    println(io, "UnROOT.LazyEvent at index $(idx) with $(nfields) columns:")
+    show(io, collect(evt))
 end
+
 function Base.getproperty(evt::LazyEvent, s::Symbol)
     @inbounds getproperty(Core.getfield(evt, :tree), s)[Core.getfield(evt, :idx)]
 end
@@ -335,5 +342,5 @@ function _clusterbytes(lbs::AbstractVector{<:LazyBranch}; compressed=false)
     return bytes
 end
 
-Tables.columns(t::LazyTree) = NamedTuple((p, getproperty(t, p)) for p in propertynames(t))
+Tables.columns(t::LazyTree) = Tables.columns(innertable(t))
 Tables.partitions(t::LazyTree) = (t[r] for r in _clusterranges(t))

test/runtests.jl

@@ -368,7 +368,10 @@ end
     @test length(collect(eachmatch(r"Float32\[", s))) == 0
     _io = IOBuffer()
     show(_io, t)
-    @test length(split(String(take!(_io)),'\n')) > length(t)
+    _iostring = String(take!(_io))
+    @test length(split(_iostring,'\n')) > length(t)
+    @test occursin("───────", _iostring)
+    @test !occursin("NamedTuple", _iostring)
     show(_io, t; crop=:both)
     @test length(split(String(take!(_io)),'\n')) <= Base.displaysize()[1]
     close(f)
@@ -689,7 +692,7 @@ end
 
 @testset "Basic C++ types" begin
     f = UnROOT.samplefile("tree_basictypes.root")
-    onesrow = LazyTree(f,"t")[2] |> values .|> first .|> Int
+    onesrow = LazyTree(f,"t")[2] |> collect |> values .|> first .|> Int
     @test all(onesrow .== 1)
 end
 
@@ -712,3 +715,21 @@ end
     @test sum(UnROOT._clusterbytes([t.b1]; compressed=true)) == 33493.0 # same as uproot4
     @test sum(UnROOT._clusterbytes([t.b2]; compressed=true)) == 23710.0 # same as uproot4
 end
+
+
+@static if VERSION > v"1.5.0"
+    @testset "Broadcast fusion" begin
+        rootfile = ROOTFile(joinpath(SAMPLES_DIR, "NanoAODv5_sample.root"))
+        t = LazyTree(rootfile, "Events", "nMuon")
+        testf(evt) = evt.nMuon == 4
+        testf2(evt) = evt.nMuon == 4
+        alloc1 = @allocated a1 = testf.(t)
+        alloc1 += @allocated a2 = testf2.(t)
+        alloc1 += @allocated idx1 = findall(a1 .& a2)
+        alloc2 = @allocated idx2 = findall(@. testf(t) & testf2(t))
+        @assert !isempty(idx1)
+        @test idx1 == idx2
+        # compiler optimization is good on 1.8
+        @test alloc1 > 1.4*alloc2
+    end
+end