Julia DataFrame - Preferred way to update values in one column based on the values in another column

Viewed 1032

I am trying to learn the best way in Julia to update a DataFrame in one column based on the value in another column. I primarily use R, in which I can write this:

library(data.table)

dt <- data.table(A = 1:5, B = 6:10)
# A  B
# ----
# 1  6
# 2  7
# 3  8
# 4  9
# 5  10

dt[A < 3, B := B * 2]
# A  B
# ----
# 1  12
# 2  14
# 3  8
# 4  9
# 5  10

...which updates B in place where A < 3, and very succinctly too. The only Julia equivalent I've found is:

using DataFrames

df = DataFrame(A = 1:5, B = 6:10)

transform!(df, [:A, :B] => ByRow((a, b) -> a < 3 ? b * 2 : a))

Is this the only method? And is it preferred? Are there any packages that can make this more succinct?

2 Answers

Currently you can do this:

@. df.B = ifelse(df.A < 3, 2*df.B, df.B)

or this

df.B = @. ifelse(df.A < 3, 2*df.B, df.B)

or this

dfv = view(df, df.A .< 3, :);
@. dfv.B = dfv.B * 2

The third approach is something similar to what data.table provides, but you have to create a view separately first (something that data.table does in one command).

The difference between the first and the second approach is whether :B column is updated in-place or replaced.

While Bogumił Kamiński is the authority here and his answer shows a very nice, short alternative solution I believe there is room for more discussion and examples.

Firstly, your original solution is certainly idiomatic and I personally would prefer it in many cases. Mutation is an important tool for high-performance languages but it is also a source of many problems. The Julia convention (common for certain other languages as well) of adding an exclamation point (or "bang!") to functions that mutate their argument help greatly in keeping track of the places where mutation happens.

In-place assignments (.=) is also commonplace in Julia code but in my own code and in the code of many library authors these would typically be found in functions named in accordance with the convention so that the intent is clear at the call site. If you find that you often needed the kind of operation you're describing you could write a function:

function filter_transform!(df, pred, args...)
    fdf = filter(pred, df; view=true)
    fdf .= transform(copy(fdf), (col => f => col for (col,f) in args)...)
end

and to use it:

filter_transform!(df, :A => <(3), :B => b -> b*2)

which I think reads nicely. The function works for your original example but does not allow for assignment to another column or using several columns in the calculation or adding columns (like data.table allows, with NA filling the skipped rows of new columns). This slightly more complicated function achieves that:

applyif(cond, f, x, default) = cond ? f(x) : default
function filter_transform!(df, filt_cols, args...)
    colset = Set(keys(eachcol(df)))
    fcols, pred = filt_cols
    for (cols, (f, assign)) in args
        df[!, assign] .= in(assign, colset) ? df[!, assign] : missing
        @. df[!, assign] = applyif(pred(df[!, fcols]), f, df[!,cols], df[!, assign])
    end
end

using it:

filter_transform!(df, :A => <(3), :B => (b -> b*2) => :B)

Finally, you asked about packages there are indeed packages that you could check out:

Query.jl (this example reassigns instead of modifying in-place):

using Query
df = df |> @mutate(B = _.A < 3 ? _.B*2 : _.B)

DataFramesMeta (also reassigning):

using DataFramesMeta
df = @eachrow df :B = :A < 3 ? :B*2 : :B
Related