r - Choose groups to keep/drop in data.table

Question

How can I drop/keep groups according to a condition in data.table? Is there a better method than adding a new column, then filtering on that column and removing it?

set.seed(0)
dt <- data.table(a = rep(1:3, rep(3, 3)), b = sample(1:5, 9, T))

#    a b
# 1: 1 4
# 2: 1 1
# 3: 1 2
# 4: 2 1
# 5: 2 4
# 6: 2 2
# 7: 3 4
# 8: 3 3
# 9: 3 4

#data.table
dt[, keep := 2 %in% b, by = a][keep == T][, keep := NULL][]
#    a b
# 1: 1 5
# 2: 1 2
# 3: 1 2
# 4: 2 3
# 5: 2 5
# 6: 2 2

# dplyr
dt %>% 
  group_by(a) %>% 
  filter(2 %in% b)

# # A tibble: 6 x 2
# # Groups:   a [2]
#       a     b
#   <int> <int>
# 1     1     5
# 2     1     2
# 3     1     2
# 4     2     3
# 5     2     5
# 6     2     2

Benchmark to see if .I is faster. 2015 Macbook Pro

bench <- 
  map(10^(4:7)
      , ~ {
        df <- data.table(name = sample(1:.x, 3*.x, T)
                         , a = runif(3*.x)
                         , b = runif(3*.x)
                         , c = runif(3*.x))

        dt <- data.table(a = rep(1:.x, rep(10, .x)), b = sample(1:10, 10*.x, T))

        microbenchmark(dt[, if(2 %in% b) .SD, a]
                      , dt[dt[, .I[2 %in% b], a]$V1] )
      })


bench

[[1]]
Unit: milliseconds
                         expr      min       lq     mean   median       uq       max neval
   dt[, if (2 %in% b) .SD, a] 13.04827 17.36046 21.15155 19.19119 22.94641  43.04519   100
 dt[dt[, .I[2 %in% b], a]$V1] 17.32547 22.92023 27.09775 24.87586 28.39789 108.47604   100

[[2]]
Unit: milliseconds
                         expr      min       lq     mean   median       uq      max neval
   dt[, if (2 %in% b) .SD, a] 123.9118 143.7802 162.6719 154.4713 173.2986 428.4141   100
 dt[dt[, .I[2 %in% b], a]$V1] 158.2975 177.3303 206.3611 193.4460 224.5091 435.3982   100

[[3]]
Unit: seconds
                         expr     min       lq     mean   median       uq      max neval
   dt[, if (2 %in% b) .SD, a] 1.23310 1.351067 1.448680 1.402827 1.517017 1.852797   100
 dt[dt[, .I[2 %in% b], a]$V1] 1.58702 1.704344 1.826468 1.778590 1.947943 2.243176   100

[[4]]
Unit: seconds
                         expr      min       lq     mean   median       uq      max neval
   dt[, if (2 %in% b) .SD, a] 14.44317 14.65889 14.79806 14.78217 14.91571 15.29134   100
 dt[dt[, .I[2 %in% b], a]$V1] 18.04774 18.36764 18.48804 18.45732 18.53073 20.73805   100

Answer 1

You can use an if condition with .SD after grouping dt by column a:

dt[, if(2 %in% b) .SD, a]

#   a b
#1: 1 5
#2: 1 2
#3: 1 2
#4: 2 3
#5: 2 5
#6: 2 2

From ?.SD, .SD is a data.table containing the Subset of x's Data for each group. Combined with if condition, we return nothing if 2 is not in column b and drop the corresponding group.