Julia自定义二叉树评估时异常内存分配与性能问题求解

基础实现与初始问题 #

在Julia中实现包含内部节点、叶节点的二叉树结构时，节点持有左右子节点（均为NodeLeaf类型）的引用，初始实现代码如下：

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)

初始测试中观察到evaluate执行段存在大量内存分配，分配量随评估循环次数线性增长，对应性能统计输出如下：

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
 ───────────────────────────────────────────────────────────────────────  

实际场景测试与对照 #

初始示例做了逻辑简化，实际使用场景需要访问DataFrame格式存储的数据，对应main函数实现如下：

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

该场景下测试仍存在大量内存分配，性能统计如下：

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
 ───────────────────────────────────────────────────────────────────────%   

作为对照，使用普通二维数组存储数据，采用完全相同的视图访问方式传参时，评估过程无内存分配，对应代码与测试结果如下：

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
 ───────────────────────────────────────────────────────────────────────%

内存分配成因 #

• 初始测试版本的分配与evaluate核心逻辑无关：循环内每次调用randn(10)都会在堆上生成新的10元素Float64数组，统计到的分配全部来自测试输入构造，并非树遍历逻辑产生。
• DataFrame行视图的类型不稳定问题：@view df[i, :]返回的DataFrame行对象是抽象类型容器，字段类型无法在编译期确定，每次索引取值都会触发动态派发和小对象堆分配；而普通二维数组的行视图是类型确定的SubArray，编译器可以完成全流程静态推断，因此可以实现零分配。
• 结构体设计存在优化空间：NodeLeaf所有字段都使用Union{Nothing, T}类型定义，即使叶节点不需要split/column/left/right字段、内部节点不需要value字段，统一的Union类型会让每次字段访问都增加运行时类型检查开销，一定程度上阻碍编译器优化。

零分配优化方案 #

• 重构二叉树类型体系，区分内部节点与叶节点，消除不必要的Union类型：

  abstract type AbstractNode end
  
  mutable struct LeafNode <: AbstractNode
      value::Float64
  end
  
  mutable struct InternalNode <: AbstractNode
      split::Float64
      column::Int
      left::AbstractNode
      right::AbstractNode
  end
  

  拆分后每个节点的字段类型完全确定，无需Nothing类型兜底，可大幅减少运行时类型检查开销。
• 优化DataFrame访问逻辑：避免逐行取整行视图传入评估函数，改为直接按列索引取值，使用类型完全确定的元组或静态数组存储单行特征，从根源消除动态派发。例如评估单行时可直接构造固定长度的Float64元组或SVector存储特征，这类结构可被编译器完全优化到栈上，无堆分配。
• 为evaluate函数增加具体类型标注，约束输入特征x的类型为确定的容器类型（如AbstractVector{Float64}、固定长度的Float64元组、静态数组），辅助编译器完成类型推断，消除无意义的动态分发。
• 性能测试时注意拆分开销统计维度：将测试数据构造逻辑移到评估循环外，避免将输入构造的内存分配误算到核心遍历逻辑的开销中。

内容的提问来源于stack exchange，提问作者Vladimir Vargas