Julia: type-stability with DataFrames - julia

Question

How can I access the columns of a DataFrame in a type-stable way?
Let's assume I have the following data:
df = DataFrame(x = fill(1.0, 1000000), y = fill(1, 1000000), z = fill("1", 1000000))
And now I want to do some recursive computation (so I cannot use transform)
function foo!(df::DataFrame)
for i in 1:nrow(df)
if (i > 1) df.x[i] += df.x[i-1] end
end
end
This has terrible performance:
julia> #time foo!(df)
0.144921 seconds (6.00 M allocations: 91.529 MiB)
A quick fix in this simplified example would be the following:
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
julia> #time bar!(df)
0.000004 seconds
However, I'm looking for a solution that is generalisable, eg when the recursive computation is just specified as a function
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
While this (almost) doesn't allocate, it is still very slow.
julia> #time foo2!(df, my_fn)
0.050465 seconds (1 allocation: 16 bytes)
Is there an approach that is performant and allows this kind of flexibility / generalisability?
EDIT: I should also mention that in practice it is not known a priori on which columns the function fn depends on. Ie I'm looking for an approach that allows performant access to / updating of arbitrary columns inside fn. The needed columns could be specified together with fn as a Vector{Symbol} for example if necessary.
EDIT 2: I tried using barrier functions as follows, but it's not performant
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
#time foo3!(df, my_fn1, [:x])
#time foo3!(df, my_fn2, [:x, :y])

Answer 1

This issue is intended (to avoid excessive compilation for wide data frames) and the ways how to handle it are explained in https://github.com/bkamins/Julia-DataFrames-Tutorial/blob/master/11_performance.ipynb.
In general you should reduce the number of times you index into a data frame. So in this case do:
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
(I am using the version where you want to have a custom function passed)

Answer 2

My current approach involves wrapping the DataFrame in a struct and overloading getindex / setindex!. Some additional trickery using generated functions is needed to get the ability to access columns by name. While this is performant, it is also a quite hacky, and I was hoping there was a more elegant solution using only DataFrames.
For simplicity this assumes all (relevant) columns are of Float64 type.
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))