Passing Symbol argument to @where for selecting rows of a DataFrame with Julia

user_2741402 0
Viewed 301

I want to create a subsample of a dataframe based on the values of one of the columns.

As discussed in a previous question (see /questions/58220143/julia-dataframe-select-rows-based-values-of-one-column-belonging-to-a-set), I can use the following syntax:

using DataFrames, DataFramesMeta
DT = DataFrame(ID = [1, 1, 2,2,3,3, 4,4], x1 = rand(8))

@where(DT, in([1,4]).(:ID))

Now if I develop a function that, lets call it subsampleDF, where the argument varID is the symbol corresponding to the variable upon which selection is made. 

function subsampleDF( data::DataFrame, varID::Symbol,subs::Vector{Int})
    d = @where(data, in(subs).(varID))
end

subsampleDF(DT, :ID,  [1,4])

Then I get the following error. I am not sure of why it's wrong as the varID is actually equal to :ID ?

> ArgumentError: invalid row index of type Bool

Stacktrace:
 [1] Type at C:\Users\djourdain\.julia\packages\DataFrames\yH0f6\src\dataframerow\dataframerow.jl:60 [inlined]
 [2] getindex at C:\Users\djourdain\.julia\packages\DataFrames\yH0f6\src\dataframerow\dataframerow.jl:93 [inlined]
 [3] where(::DataFrame, ::getfield(Main, Symbol("##10#11")){Symbol,Array{Int64,1}}) at C:\Users\djourdain\.julia\packages\DataFramesMeta\c88dH\src\DataFramesMeta.jl:186
 [4] subsampleDF(::DataFrame, ::Symbol, ::Array{Int64,1}) at .\In[8]:2
 [5] top-level scope at In[9]:1

Thank you.

Damien

julia dataframe
2 Answers

In DataFramesMeta.jl use cols to reference a column using an expression. So a correct way to define your function is:

function subsampleDF( data::DataFrame, varID::Symbol,subs::Vector{Int})
    @where(data, in(subs).(cols(varID)))
end
user_12695670

Not a direct answer to your question (not sure what's going on with the @where macro), but would you be okay with just using plain DataFrames?

using DataFrames
DT = DataFrame(ID = [1, 1, 2,2,3,3, 4,4], x1 = rand(8))

function subsampleDF( data::DataFrame, varID::Symbol,subs::Vector{Int})
    data[in(subs).(data[!, varID]), :]
end

subsampleDF(DT, :ID,  [1,4])

should work as expected.

user_24998920
Related