Performance issues with evaluation of custom tree data structure in Julia

Viewed 44

I am implementing a Binary Tree in Julia. The binary tree has nodes and leafs. The nodes point to left and right children, which are also nodes/leafs objects. The following code exemplifies the data structure:

using TimerOutputs

mutable struct NodeLeaf
    isleaf::Bool
    value::Union{Nothing,Float64}
    split::Union{Nothing,Float64}
    column::Union{Nothing,Int64}
    left::Union{Nothing,NodeLeaf}
    right::Union{Nothing,NodeLeaf}
end

function evaluate(node::NodeLeaf, x)::Float64
    while !node.isleaf
        if x[node.column] < node.split
            node = node.left
        else
            node = node.right
        end
    end
    return node.value
end

function build_random_tree(max_depth)
    if max_depth == 0
        return NodeLeaf(true, randn(), randn(), rand(1:10), nothing, nothing)
    else
        return NodeLeaf(false, randn(), randn(), rand(1:10), build_random_tree(max_depth - 1), build_random_tree(max_depth - 1))
    end
end

function main()
    my_random_tree = build_random_tree(4)
    @timeit to "evaluation" for i in 1:1000000
        evaluate(my_random_tree, randn(10))
    end
end

const to = TimerOutput()
main()
show(to)

I notice that a lot of allocations occur in the evaluate function, but I don't see the reason why this is the case:

julia mytree.jl
 ───────────────────────────────────────────────────────────────────────
                               Time                    Allocations      
                      ───────────────────────   ────────────────────────
   Tot / % measured:       476ms /  21.6%            219MiB /  62.7%    

 Section      ncalls     time    %tot     avg     alloc    %tot      avg
 ───────────────────────────────────────────────────────────────────────
 evaluation        1    103ms  100.0%   103ms    137MiB  100.0%   137MiB
 ───────────────────────────────────────────────────────────────────────  

As I increase the evaluation loop, the allocation continues to increase without bound. Can anybody explain why allocation grows so much and please suggest how to avoid this issue? Thanks.

EDIT I simplified too much the code for the example. The actual code is accessing DataFrames, so the main looks like this:

using DataFrames
function main()
    my_random_tree = build_random_tree(7)
    df = DataFrame(A=1:1000000)
    for i in 1:9
        df[!, string(i)] = collect(1:1000000)
    end

    @timeit to "evaluation" for i in 1:size(df, 1)
        evaluate(my_random_tree, @view df[i, :])
    end
end

I expect this to yield 0 allocations, but that isn't true:

julia mytree.jl
 ───────────────────────────────────────────────────────────────────────
                               Time                    Allocations      
                      ───────────────────────   ────────────────────────
   Tot / % measured:       551ms /  20.5%            305MiB /  45.0%    

 Section      ncalls     time    %tot     avg     alloc    %tot      avg
 ───────────────────────────────────────────────────────────────────────
 evaluation        1    113ms  100.0%   113ms    137MiB  100.0%   137MiB
 ───────────────────────────────────────────────────────────────────────%   

On the other hand, if I use a plain array I don't get allocations:

function main()
    my_random_tree = build_random_tree(7)
    df = randn(1000000, 10)

    @timeit to "evaluation" for i in 1:size(df, 1)
        evaluate(my_random_tree, @view df[i, :])
    end
end
julia mytree.jl
 ───────────────────────────────────────────────────────────────────────
                               Time                    Allocations      
                      ───────────────────────   ────────────────────────
   Tot / % measured:       465ms /   5.7%            171MiB /   0.0%    

 Section      ncalls     time    %tot     avg     alloc    %tot      avg
 ───────────────────────────────────────────────────────────────────────
 evaluation        1   26.4ms  100.0%  26.4ms     0.00B     - %    0.00B
 ───────────────────────────────────────────────────────────────────────%
1 Answers

The thing that allocates is randn not evaluation. Switch to randn!:

julia> using Random

julia> function main()
           my_random_tree = build_random_tree(4)
           x = randn(10)
           @allocated for i in 1:1000000
               evaluate(my_random_tree, randn!(x))
           end
       end
main (generic function with 1 method)

julia> main()
0

EDIT

Solution with DataFrames.jl:

function bar(mrt, nti)
    @timeit to "evaluation" for nt in nti
        evaluate(mrt, nt)
    end
end

function main()
    my_random_tree = build_random_tree(7)
    df = DataFrame(A=1:1000000)
    for i in 1:9
        df[!, string(i)] = collect(1:1000000)
    end
    bar(my_random_tree, Tables.namedtupleiterator(df))
end
Related