Simple and efficient way to subset a data frame using values and names in a vector

Viewed 124

Given a dataset (let say stored as data frame) in the form:

> n <- 10   
> set.seed(123)
> ds.df <- data.frame(col1 = round(rnorm(n,2,4), digit = 1),
                    col2 = sample.int(2, n, replace = TRUE),
                    col3 = sample.int(n*10, n),
                    col4 = sample(letters, n, replace = TRUE))

is there a simple and efficient way to subset it, using a vector of value that defines multiple equalities the subset should respect? Something like:

> subset_v <- c(col1 = -0.2, col4 = "i")
> ds.subset <- subset(ds.df, subset_v)
> ds.subset
  col1 col2 col3 col4
1 -0.2    1    9    i

where the function subset(ds.df,subset_v) should return the subset that respect:

ds.df[ ds.df$col1 == subset_v["col1"] & ds.df$col2 == subset_v["col2"] & ds.df$col4 == subset_v["col4"], ]

But this last expression isn't very handy and I would like to be able to have any column without knowing them into advance.

I did something that works:

subset <- function(ds.df,subset_v){
    sub = rep(TRUE, nrow(ds.df))
    for(cn in names(subset_v)){
       sub=sub & (ds.df[,cn] == subset_v[[cn]])
    }
    ds.df[sub,]
}

But I feel like there is a much better and more efficient way to do it (maybe removing the for loop somehow).

3 Answers

Personally, I wonder whether it is a good idea to use a named vector to subset a dataframe, since it can only be used for equality =, while larger than and smaller than cannot be expressed this way. I would recommend using a quoted expression instead of a named vector (see approach below).

However, I figured out a tidyverse way to write a function with said functionality:

library(tidyverse)

set.seed(123)
n <- 10 

ds.df <- data.frame(col1 = round(rnorm(n,2,4), digit=1),
                   col2 = sample.int(2, n, replace=T),
                   col3 = sample.int(n*10, n),
                   col4 = sample(letters, n, replace=T))

new_filter <- function (data, expr) {
  exprs_ls <- purrr::imap(expr, ~ rlang::exprs(!! rlang::sym(.y) == !!.x))
  filter(data, !!! unname(unlist(exprs_ls)))
}

new_filter(ds.df, c(col1 = -0.2, col4 = "i"))
#>   col1 col2 col3 col4
#> 1 -0.2    1    9    i

Created on 2020-06-17 by the reprex package (v0.3.0)

Below is my alternative approach. In base R you can use quote to quote the subset expression (instead of creating a vector) and then you can use eval to evaluate it inside subset.

n <- 10   

ds.df=data.frame(col1=round(rnorm(n,2,4),digit=1),
                 col2=sample.int(2,n,replace=T),
                 col3=sample.int(n*10,n),
                 col4=sample(letters,n,replace=T))


subset_v = quote(col1 > 2 & col3 > 40)

subset(ds.df, eval(subset_v))
#>    col1 col2 col3 col4
#> 1   6.6    1   93    m
#> 2   7.0    2   62    j
#> 4   3.9    1   94    t
#> 7   4.5    1   46    r
#> 8   2.8    2   98    h
#> 10  4.9    1   78    p

Created on 2020-06-17 by the reprex package (v0.3.0)


Same approach but using dplyr filter

library(dplyr)

n <- 10 

ds.df = data.frame(col1 = round(rnorm(n,2,4), digit=1),
                   col2 = sample.int(2, n, replace=T),
                   col3 = sample.int(n*10, n),
                   col4 = sample(letters, n, replace=T))

filter_v = expr(col1 > 2 & col3 > 40)

filter(ds.df, !! filter_v)

#>   col1 col2 col3 col4
#> 1  3.3    1   70    a
#> 2  2.5    2   82    q
#> 3  3.6    1   51    z

Created on 2020-06-17 by the reprex package (v0.3.0)

In data.table you could do this:

setDT(ds.df)
subset_v = list(col1=-3.3, col2=1, col4="e")
ds.df[as.list(subset_v), on = names(subset_v)]

#    col1 col2 col3 col4
# 1: -3.3    1   29    e

Reproducible data:

set.seed(20)
n <- 10   
ds.df <- data.frame(
  col1 = round(rnorm(n, 2, 4), digit = 1),
  col2 = sample.int(2, n, replace = TRUE),
  col3 = sample.int(n*10, n),
  col4 = sample(letters, n, replace = TRUE)
)

I think you're looking for merge:

subset <- function(ds.df,subset_v){
    filter = data.frame(as.list(subset_v))
    merge(ds.df,filter,by=names(filter),all=F)
}

This works on both data.frame and data.table, and with data.table, should be the same underneath as @sindri_baldur's answer, so, if you're already using data.table, the primary difference is whether you prefer typing merge(x,y,by=z,all=F) or x[y,on=z].

Related