Skip to content

Instantly share code, notes, and snippets.

@xiaodaigh
Created August 30, 2020 05:57
Show Gist options
  • Save xiaodaigh/132bf69b3a54faad2535517b7e4b49a9 to your computer and use it in GitHub Desktop.
Save xiaodaigh/132bf69b3a54faad2535517b7e4b49a9 to your computer and use it in GitHub Desktop.
df1[B] = df2[B] where df1 and df2 are DataFrames B is boolean array
using DataFrames
df1 = DataFrame(a = repeat([1], 100), b = "a")
df2 = DataFrame(a = repeat([2], 100), b = "b")
B = Array{Bool, 2}(undef, 100, 2)
df1[B] # doesn't work
# Let's overload get index get index
function Base.getindex(df::AbstractDataFrame, B::AbstractArray{Bool, 2})
@assert size(B) == size(df)
res = []
for (colnumber, Bcol) in enumerate(eachcol(B))
res = vcat(res, df[Bcol, colnumber])
end
res
end
df1[B]
findall(B)
# How do I make assignment work? df1[B] = df2[B]
# Overload the Base.setindex! method
function Base.setindex!(df::AbstractDataFrame, vals_to_assign, B::AbstractArray{Bool, 2})
@assert size(B) == size(df)
idx_to_assign = findall(B)
@assert length(idx_to_assign) == length(vals_to_assign)
for (idx, val) in zip(idx_to_assign, vals_to_assign)
df[idx] = val
end
vals_to_assign
end
# check the vaules should be 1 and "a"
df1[B]
# check the vaules should be 2 and "b"
df2[B]
# assign
df1[B] = df2[B]
# viola!
df1[B]
df1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment