I am running some simulations at work, and I wanted to experiment with different approaches for performing a groupby and then running the simulation. The idea is to run a simulation of each group (identified by an id column) using multiple datasets all keyed by the same id column. There are about 5 DataFrames, and one of the only has on row per id. So running 5 joins seemed wasteful.
Using some advice from the Slack #Data channel, I came up with the following four approaches and their timings on a mock dataset. Here’s the code:
using DataFrames
using Query
using Tables
using BenchmarkTools
randparams(n) = join.([rand('a':'z', 100) for _ in 1:n])
function sim(params, data)
total = 0.0
for p in params, x in data
total += Int(p)*x
end
return total
end
function datan(n_ids, n_data, n_df)
ids = collect(1:n_ids)
data_ids = collect(1:n_data)
a = DataFrame(id=ids, p=randparams(n_ids))
b = [
DataFrame(
:id => vcat([ids for _ in 1:n_data]...),
:data_id => [i for i in data_ids for _ in 1:n_data],
Symbol("x$j") => rand(n_ids*n_data),
)
for j in 1:(n_df-1)
]
a,b
end
function simgroups_naivejoin(a, b)
ab = join(a, b[1]; on=:id)
for df in b[2:end]
ab = join(ab, df; on=[:id, :data_id])
end
result = Dict()
for g in groupby(ab, :id)
params = g.p[1]
result[g.id[1]] = 0.0
for i in 1:length(b)
data = g[!, Symbol("x$i")]
result[g.id[1]] += sim(params, data)
end
end
return result
end
function simgroups_groupbykeys(a, b)
ag = groupby(a, :id)
bg = [groupby(df, :id) for df in b]
a_gk_by_pk = Dict(tuple(k...) => k for k in keys(ag))
b_gk_by_pk = [Dict(tuple(k...) => k for k in keys(dfg)) for dfg in bg]
result = Dict()
for ka in keys(a_gk_by_pk)
if all(haskey(gk_by_pk, ka) for gk_by_pk in b_gk_by_pk)
a_gk= a_gk_by_pk[ka]
params = ag[a_gk].p[1]
result[ka[1]] = 0.0
for (i, gk_by_pk, dfg) in zip(1:length(bg), b_gk_by_pk, bg)
bgkey = gk_by_pk[ka]
data = dfg[bgkey][!, Symbol("x$i")]
result[ka[1]] += sim(params, data)
end
end
end
return result
end
function simgroups_groupjoinselect(a, b)
ab = a |> @groupjoin(
b[1], _.id, _.id,
(x,y) -> (id=x.id, params=x.p, data=(y,)))
for (i, df) in zip(2:length(b), b[2:end])
data_col = Symbol("data$i")
ab = ab |> @groupjoin(
df, _.id, _.id,
(x,y) -> (id=x.id, params=x.params, data=(x.data..., y)))
end
result = Dict()
for g in ab
result[g.id] = 0.0
for (i, gdata) in zip(1:length(b), g.data)
x_col = Symbol("x$i")
data = vec(Tables.matrix(Tables.select(gdata, x_col)))
result[g.id] += sim(g.params, data)
end
end
return result
end
function simgroups_indexview(a, b)
result = Dict()
mask = [Array{Bool}(undef, size(df)[1]) for df in b]
for g in groupby(a, :id)
id = g.id[1]
params = g.p[1]
result[g.id[1]] = 0.0
for i in 1:length(b)
df = b[i]
mask[i] .= df.id .== id
data = view(df, mask[i], Symbol("x$i"))
result[g.id[1]] += sim(params, data)
end
end
return result
end
function benchmark(n, m)
@assert m >= 2
a, b = datan(n, n, m)
println("naive join")
@btime simgroups_naivejoin($a, $b)
println("groupby keys")
@btime simgroups_groupbykeys($a, $b)
println("groupjoin select")
@btime simgroups_groupjoinselect($a, $b)
println("index view")
@btime simgroups_indexview($a, $b)
println("DONE")
end
And here’s the results:
julia> benchmark(1000, 5)
naive join
4.281 s (6077351 allocations: 855.46 MiB)
groupby keys
1.983 s (141394 allocations: 159.02 MiB)
groupjoin select
1.323 s (171862 allocations: 255.23 MiB)
index view
9.774 s (99897 allocations: 37.66 MiB)
DONE
The Query plus Tables case (groupjoinselect) also provides a pretty flexible approach for handling the data, so it’s pretty encouraging for other use cases too. Obvious the results will vary a lot based on number of columns and how they are used, but I thought other mind find this interesting too.
In Python, I’ve mostly done approach one or two, but that’s largely because of performance considerations, since Pandas is usually the fastest options, and their indexing API is not bad if you’re used to it.
This example is part of my process of learning Julia so I might use it on the job for something more substantial.