How to properly add two new columns (in dataframe) from a function that returns two arrays?

General Usage
,
1

Hello, I am a Julia newbie and I have a question.
The example below sort of simplified my current problem—

function myfun(id)
  a = [id .+ 1, id .+ 2, id .+ 2]
  b = [id .- 1, id .- 2, id .- 3]
  return [a, b]
end

mydf = DataFrame(id = [7, 8])
combine(groupby(mydf, :id), :id => myfun)

For each group (groupby id in my data mydf), I run a function, myfun, which generates two arrays based on a variable in mydf.
What I want to do is to add the two arrays (a and b) from the function to the original data as two columns.

However… the command below generates only a single column

combine(groupby(mydf, :id), :id => myfun)

4×2 DataFrame
 Row │ id     id_myfun          
     │ Int64  Array…            
─────┼──────────────────────────
   1 │     7  [[8], [9], [9]]
   2 │     7  [[6], [5], [4]]
   3 │     8  [[9], [10], [10]]
   4 │     8  [[7], [6], [5]]

Ideally, I want my result to be like below…

 Row │ id     │  a    │b 
     │ Int64  │ Int64 │ Int64
─────┼──────────────────────────
   1 │     7  | 8     │ 6 
   2 │     7  | 9     │ 5
   3 │     7  | 9     │ 4 
   4 │     7  | 9     │ 7 
   5 │     7  | 10    │ 6
   6 │     7  | 10    │ 5

Is there anyone who could help me?

Thanks!

1 Like
2

Here’s one way:

julia> df_ab = transform(mydf, :id => ByRow(myfun) => [:a, :b])
2×3 DataFrame
 Row │ id     a            b         
     │ Int64  Array…       Array…    
─────┼───────────────────────────────
   1 │     7  [8, 9, 9]    [6, 5, 4]
   2 │     8  [9, 10, 10]  [7, 6, 5]

julia> flatten(df_ab, [:a, :b])
6×3 DataFrame
 Row │ id     a      b     
     │ Int64  Int64  Int64 
─────┼─────────────────────
   1 │     7      8      6
   2 │     7      9      5
   3 │     7      9      4
   4 │     8      9      7
   5 │     8     10      6
   6 │     8     10      5
4 Likes
4

Another way

julia> function myfun(id)
           a = [id .+ 1, id .+ 2, id .+ 2]
           b = [id .- 1, id .- 2, id .- 3]
           id = fill(id, length(a))
           (; id, a, b)
       end;

julia> vcat(DataFrame.(myfun.([7, 8]))...)
6×3 DataFrame
 Row │ id     a      b
     │ Int64  Int64  Int64
─────┼─────────────────────
   1 │     7      8      6
   2 │     7      9      5
   3 │     7      9      4
   4 │     8      9      7
   5 │     8     10      6
   6 │     8     10      5
2 Likes
5

Many Thanks @sijo and @rikh!
These are perfect answers for me.

6

Some alternatives for reference :slight_smile:

combine(groupby(mydf, :id), :id => Base.splat(hcat)∘myfun∘only => [:a, :b])

combine(groupby(mydf, :id), :id => NamedTuple{(:a, :b)}∘myfun∘only => AsTable)

# If the `id` column is not required in the result:
combine(mydf, :id => Base.splat(vcat)∘ByRow(Base.splat(hcat)∘myfun) => [:a, :b])

There’s also an issue discussing the idea of making this easier.

1 Like
7

Thank you for listing all these excellent alternatives.

One quick follow-up question—
Is their performance comparable? I wonder if some of them are dominantly better than the other in terms of speed.

8

I don’t know… probably best to run benchmarks for your real use case.

9

Thanks! I will do some tests.