Julia: type-stability with DataFrames

user_7756364 0
Viewed 130

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])
julia
2 Answers

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)

user_12695670

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))
user_77563640
Related