groupby() using two arrays in julia?

Question

I have two arrays with same dimension:

a1 = [1,1,3,4,6,6]
a2 = [1,2,3,4,5,6]

And I want to group both of them with respect to array a1 and get the mean of the array a2 for each group. My output is coming from array a2, as mentioned below:

result:
1.5
3.0
4.0
5.5

Please suggest an approach to achieve this task. Thanks!!

user_1269567 · Answer

Here is a solution using DataFrames.jl:

julia> using DataFrames, Statistics

julia> df = DataFrame(a1 = [1,1,3,4,6,6], a2 = [1,2,3,4,5,6]);

julia> combine(groupby(df, :a1), :a2 => mean)
4×2 DataFrame
 Row │ a1     a2_mean
     │ Int64  Float64
─────┼────────────────
   1 │     1      1.5
   2 │     3      3.0
   3 │     4      4.0
   4 │     6      5.5

EDIT:

Here are the timings (as usual in Julia you need to remember that the first time you run some function it has to be compiled which takes time):

julia> using DataFrames, Statistics

(@v1.6) pkg> st DataFrames # I am using main branch, as it should be released this week
      Status `D:\.julia\environments\v1.6\Project.toml`
  [a93c6f00] DataFrames v0.22.7 `https://github.com/JuliaData/DataFrames.jl.git#main`

julia> df = DataFrame(a1=rand(1:1000, 10^8), a2=rand(10^8)); # 10^8 rows in 1000 random groups

julia> @time combine(groupby(df, :a1), :a2 => mean); # first run includes compilation time
  3.781717 seconds (6.76 M allocations: 1.151 GiB, 6.73% gc time, 84.20% compilation time)

julia> @time combine(groupby(df, :a1), :a2 => mean); # second run is just execution time
  0.442082 seconds (294 allocations: 762.990 MiB)

Note that e.g. data.table (if this is your reference) on similar data is noticeably slower:

> library(data.table) # using 4 threads
> df = data.table(a1 = sample(1:1000, 10^8, replace=T), a2 = runif(10^8));
> system.time(df[, .(mean(a2)), by = a1])
   user  system elapsed 
   4.72    1.20    2.00

user_4658686 · Answer

In case you are interested in using Chain.jl in addition to DataFrames.jl, Bogumił Kamiński's answer might then look like this:

julia> using DataFrames, Statistics, Chain

julia> df = DataFrame(a1 = [1,1,3,4,6,6], a2 = [1,2,3,4,5,6]);

julia> @chain df begin
           groupby(:a1)
           combine(:a2 => mean)
       end

4×2 DataFrame
 Row │ a1     a2_mean 
     │ Int64  Float64 
─────┼────────────────
   1 │     1      1.5
   2 │     3      3.0
   3 │     4      4.0
   4 │     6      5.5

groupby() using two arrays in julia?

2 Answers