Julia 朱莉娅：数据帧的类型稳定性

如何以类型稳定的方式访问数据帧的列

假设我有以下数据：

df = DataFrame(x = fill(1.0, 1000000), y = fill(1, 1000000), z = fill("1", 1000000))

现在我想做一些递归计算（所以我不能使用

transform

）

这有着糟糕的表现：

julia> @time foo!(df)
  0.144921 seconds (6.00 M allocations: 91.529 MiB)

此简化示例中的快速修复方法如下所示：

function bar!(df::DataFrame)
    x::Vector{Float64} = df.x
    for i in length(x)
        if (i > 1) x[i] += x[i-1] end
    end
end

然而，我正在寻找一个可推广的解决方案，例如当递归计算只是指定为一个函数时

function foo2!(df::DataFrame, fn::Function)
    for i in 1:nrow(df)
        if (i > 1) fn(df, i) end
    end
end

function my_fn(df::DataFrame, i::Int64)
    x::Vector{Float64} = df.x
    x[i] += x[i-1]
end

虽然这（几乎）没有分配，但仍然非常缓慢

julia> @time foo2!(df, my_fn)
  0.050465 seconds (1 allocation: 16 bytes)

是否有一种方法是高性能的，并且允许这种灵活性/通用性

编辑：我还应该提到，在实践中，函数

fn

依赖于哪些列是未知的。Ie我正在寻找一种允许性能访问/更新

fn

中任意列的方法。所需的列可以与

fn

一起指定为

向量{Symbol}

，如有必要

编辑2：我试着按如下方式使用屏障函数，但没有效果

function foo3!(df::DataFrame, fn::Function, colnames::Vector{Symbol})
    cols = map(cname -> df[!,cname], colnames)
    for i in 1:nrow(df)
        if (i > 1) fn(cols..., i) end
    end
end

function my_fn1(x::Vector{Float64}, i::Int64)
    x[i] += x[i-1]
end

function my_fn2(x::Vector{Float64}, y::Vector{Int64}, i::Int64)
    x[i] += x[i-1] * y[i-1]
end

此问题旨在（避免对宽数据帧进行过度编译）中介绍了如何处理此问题

通常，您应该减少索引到数据帧的次数。因此，在这种情况下，请：

julia> function foo3!(x::AbstractVector, fn::Function)
           for i in 2:length(x)
               fn(x, i)
           end
       end
foo3! (generic function with 1 method)

julia> function my_fn(x::AbstractVector, i::Int64)
           x[i] += x[i-1]
       end
my_fn (generic function with 1 method)

julia> @time foo3!(df.x, my_fn)
  0.010746 seconds (16.60 k allocations: 926.036 KiB)

julia> @time foo3!(df.x, my_fn)
  0.002301 seconds

---

（我使用的是希望传递自定义函数的版本）

我当前的方法是将数据帧包装到结构中，并重载

getindex

/

setindex。为了能够按名称访问列，还需要使用生成的函数进行一些额外的技巧。虽然这是一个性能很好的解决方案，但它也是一个相当粗糙的解决方案，我希望有一个只使用数据帧的更优雅的解决方案

为简单起见，假设所有（相关）列都是
Float64
类型

struct DataFrameWrapper{colnames}
    cols::Vector{Vector{Float64}}
end

function df_to_vectors(df::AbstractDataFrame, colnames::Vector{Symbol})::Vector{Vector{Float64}}
    res = Vector{Vector{Float64}}(undef, length(colnames))
    for i in 1:length(colnames)
        res[i] = df[!,colnames[i]]
    end
    res
end

function DataFrameWrapper{colnames}(df::AbstractDataFrame) where colnames
    DataFrameWrapper{colnames}(df_to_vectors(df, collect(colnames)))
end

get_colnames(::Type{DataFrameWrapper{colnames}}) where colnames = colnames

@generated function get_col_index(x::DataFrameWrapper, ::Val{col})::Int64 where col
    id = findfirst(y -> y == col, get_colnames(x))
    :($id)
end

Base.@propagate_inbounds Base.getindex(x::DataFrameWrapper, col::Val)::Vector{Float64} = x.cols[get_col_index(x, col)]
Base.@propagate_inbounds Base.getindex(x::DataFrameWrapper, col::Symbol)::Vector{Float64} = getindex(x, Val(col))
Base.@propagate_inbounds Base.setindex!(x::DataFrameWrapper, value::Float64, row::Int64, col::Val) = setindex!(x.cols[get_col_index(x, col)], value, row)
Base.@propagate_inbounds Base.setindex!(x::DataFrameWrapper, value::Float64, row::Int64, col::Symbol) = setindex!(x, value, row, Val(col))

不幸的是，对于我的用例来说，这还不够通用，因为必须对
foo3
中传递给
fn
的列进行硬编码。本质上，我正在寻找一种方法，将递归逻辑（实际上还包含大量簿记、度量计算等）与状态更新逻辑（可能取决于数据框架中的任何列）分离开来。您可以轻松地将列名作为
符号传递
，然后执行
df[！，colname]
然后调用worker函数。因此，您不必硬编码任何内容。对于任意数量的列，我将如何做到这一点？例如我有
fn
和
向量{Symbol}
。只需在
向量{Symbol}
上迭代即可。如果你给我一个你需要的具体例子，我可以向你提出一个解决方案（在你的问题中它不存在）。我用另一个例子更新了这个问题，突出了我想做的事情。
@time foo3!(df, my_fn1, [:x])
@time foo3!(df, my_fn2, [:x, :y])

julia> function foo3!(x::AbstractVector, fn::Function)
           for i in 2:length(x)
               fn(x, i)
           end
       end
foo3! (generic function with 1 method)

julia> function my_fn(x::AbstractVector, i::Int64)
           x[i] += x[i-1]
       end
my_fn (generic function with 1 method)

julia> @time foo3!(df.x, my_fn)
  0.010746 seconds (16.60 k allocations: 926.036 KiB)

julia> @time foo3!(df.x, my_fn)
  0.002301 seconds

struct DataFrameWrapper{colnames}
    cols::Vector{Vector{Float64}}
end

function df_to_vectors(df::AbstractDataFrame, colnames::Vector{Symbol})::Vector{Vector{Float64}}
    res = Vector{Vector{Float64}}(undef, length(colnames))
    for i in 1:length(colnames)
        res[i] = df[!,colnames[i]]
    end
    res
end

function DataFrameWrapper{colnames}(df::AbstractDataFrame) where colnames
    DataFrameWrapper{colnames}(df_to_vectors(df, collect(colnames)))
end

get_colnames(::Type{DataFrameWrapper{colnames}}) where colnames = colnames

@generated function get_col_index(x::DataFrameWrapper, ::Val{col})::Int64 where col
    id = findfirst(y -> y == col, get_colnames(x))
    :($id)
end

Base.@propagate_inbounds Base.getindex(x::DataFrameWrapper, col::Val)::Vector{Float64} = x.cols[get_col_index(x, col)]
Base.@propagate_inbounds Base.getindex(x::DataFrameWrapper, col::Symbol)::Vector{Float64} = getindex(x, Val(col))
Base.@propagate_inbounds Base.setindex!(x::DataFrameWrapper, value::Float64, row::Int64, col::Val) = setindex!(x.cols[get_col_index(x, col)], value, row)
Base.@propagate_inbounds Base.setindex!(x::DataFrameWrapper, value::Float64, row::Int64, col::Symbol) = setindex!(x, value, row, Val(col))