Filling a StructArray in a type-stable manner

General Usage Performance
1

I am trying to create StructArray’s reading from a ROOT file with the UnROOT.jl package, which provides a NamedTuple interface, in a type-stable manner. The following code illustrates the problem:

using StructArrays
using BenchmarkTools

struct Vector3f
    x::Float32
    y::Float32
    z::Float32
end

struct Foo
    a::Int64
    b::Int64
    c::Int64  # not in data
    vector::Vector3f
end

function StructArray{T,M}(dat::NamedTuple) where {T <: Number, M}  # number types
    getproperty(dat, M)
end
function StructArray{T,M}(dat::NamedTuple) where {T,M}             # generic composite types
    len = length(dat[1])
    tup = Tuple( map(zip(fieldnames(T),fieldtypes(T))) do (fn,ft)
                    if fn == :c
                        fill(0,len)      # missing in data
                    else
                        StructArray{ft, Symbol(M, :_ ,fn)}(dat)
                    end
                end
        )
    StructArray{T}(tup)
end

function StructArray{Foo,M}(dat::NamedTuple, ::Bool) where {M}    # very concrete for Foo type
    tup = (
        getproperty(dat, Symbol(M,:_a)),
        getproperty(dat, Symbol(M,:_b)),
        fill(0,100),
        StructArray{Vector3f}(( getproperty(dat, Symbol(M,:_vector_x)), 
                                getproperty(dat, Symbol(M,:_vector_y)),
                                getproperty(dat, Symbol(M,:_vector_z)) ))
    )
    StructArray{Foo}(tup)
end

data = (Foos_a=fill(1,100), 
        Foos_b=fill(2,100), 
        Foos_vector_x=fill(1.,100),
        Foos_vector_y=fill(2.,100),
        Foos_vector_z=fill(3.,100))

I get

julia> @btime foos1 = StructArray{Foo, :Foos}(data);
  10.000 μs (101 allocations: 5.67 KiB)

julia> @btime foos2 = StructArray{Foo, :Foos}(data, true);
  152.556 ns (2 allocations: 960 bytes)

The first uses a generic implementation, which is what I would want for my application, but it is factors slower that the second implementation, which is written explicitly for the concrete type Foo.
Does anybody has an idea how I would write StructArray{T,M}(dat::NamedTuple) where {T,M} in a type-stable manner? Thanks.

2

is that really a bottleneck? I think ideally we should do type-space computation “once” near the beginning of the whole workload, and then never worry about it.

In other words, as long as this doesn’t scale with the amount of data we process, it’s not an issue I think

3

At least this is a reproducible bottleneck. I need to do each event the operation of combining a set of vectors into a StructArray. I do not know how to do something at the beginning and reuse it over an over later for each event.

1 Like
4

I just don’t see how this can avoid allocation.

I think if you work with RNTuple, we can come up a way to expose the storage “column” directly, so you don’t have to work with NamedTuples maybe. As it is right now, only way I can think of is if you do @generated