将特定于组的函数应用于 Julia 数据帧

Question

我想根据在另一列中指定的类别(每个函数一个)将两个可能的函数之一应用于DataFrame的列。

我对此的用例是将来自两个不同位置的观测列表的 JD 时间转换为 BJD(重心儒略日期)，这取决于该位置。

例如，如果我有一个这样的表：

julia> using DataFrames
julia> df = DataFrame(:t => 1:5, :location => rand(["locA", "locB"], 5))
5×2 DataFrame
Row │ t      location 
│ Int64  String   
─────┼─────────────────
1 │     1  locB
2 │     2  locA
3 │     3  locA
4 │     4  locB
5 │     5  locA

其中 BJD 时间在locA= t^2 和时间在locB= -t^2(只是为了清楚起见)，我该如何添加这个新列？

我尝试的一种方法是只应用三元函数：

julia> loc_to_BJD(t, loc) = loc == "locA" ? t^2 : -t^2
loc_to_BJD (generic function with 1 method)
julia> df.loc_to_BJD = loc_to_BJD.(df.t, df.location); df
5×3 DataFrame
Row │ t      location  loc_to_BJD 
│ Int64  String    Int64      
─────┼─────────────────────────────
1 │     1  locB              -1
2 │     2  locA               4
3 │     3  locA               9
4 │     4  locB             -16
5 │     5  locA              25

但是每次像这样对类别进行检查似乎效率非常低(除非编译器在幕后做一些聪明的事情？

在这种情况下，使用组是否更有意义？不过，在尝试此操作时，我遇到了尝试向SubDataFrame添加列的问题，因此我认为我没有从正确的方向接近这一点。

Answer 1

我不确定你为什么认为这会是低效的，因为在我看来，有人必须去检查这些location，看看它们是否是 A，而是 IANACS。想到的三种方法的快速基准：

julia> using BenchmarkTools, DataFrames
julia> function add1(d)
loc_to_BJD(t, loc) = loc == "locA" ? t^2 : -t^2
d[!, "loc_to_BJD"] = loc_to_BJD.(d.t, d.location)
return d
end
add1 (generic function with 1 method)
julia> function add2(d)
d[!, "loc_to_BJD"] = ifelse.(d.location .== "locA", d.t .^ 2, (-1) .* d.t .^ 2)
end
add2 (generic function with 1 method)
julia> function add3(d)
transform!(d, [:t, :location] => ByRow((t, l) -> t == "locA" ? t^2 : -t^2) => :loc_to_BJD)
end
add3 (generic function with 1 method)
julia> @btime add1($df);
10.627 ms (6 allocations: 7.63 MiB)
julia> @btime add2($df);
11.291 ms (19 allocations: 7.63 MiB)
julia> @btime add3($df);
1.965 ms (163 allocations: 7.64 MiB)

我很惊讶地看到transform!在这里比简单的ifelse/三元运算符版本做得好得多，我觉得我在基准测试中做错了什么，但看不到它是什么。也许检查您是否在现实世界的用例中找到相同的内容，但我相信Bogumil很快就会

解释发生了什么:)编辑：感谢Bogumil在下面的评论，这是正确的add3：

julia> function add3(d)
transform!(d, [:t, :location] => ByRow((t, l) -> l == "locA" ? t^2 : -t^2) => :loc_to_BJD)
end
add3 (generic function with 1 method)
julia> @btime add3($df);
11.177 ms (163 allocations: 7.64 MiB)

所以毕竟没有神秘感！