MatthewR
MatthewR

Reputation: 2780

Using indexing to subset rows in a data.table

I am looking for an efficient way to subset data.tables based on a vector.

This doesn't work:

library(data.table)
z <- data.table(mtcars)
ids <- z[wt > 3.5, "mpg"] 
z[mpg %in% ids, ] 

I would expect it to work - based on this behavior for data.frames.

ids2 <- mtcars[mtcars$wt > 3.5, "mpg"] 
mtcars[mtcars$mpg %in% ids2, ] 

If I just type out the numbers it works:

z[ mpg %in% c( 14.3 , 10.4) , ]

I have tried converting classes:

subset(z, mpg %in% ids)
subset(z, as.numeric (mpg) %in% as.numeric( ids))
subset(z, as.character(mpg) %in% as.character( ids ))

Upvotes: 3

Views: 187

Answers (2)

chinsoon12
chinsoon12

Reputation: 25208

With regards to efficiency mentioned in OP, here are a few methods for your considerations. Using if(any(... is the fastest.

library(data.table)
M <- 1e7
set.seed(0L)
z <- data.table(
    mpg=c(sample(LETTERS[1:13], M, replace=TRUE),
        sample(LETTERS[14:26], M, replace=TRUE)),
    wt=c(rnorm(M), rnorm(M, 3.5))
)

op <- function() {
    z[mpg %in% z[wt > 3.5, mpg]]
}

dtUniq <- function() {
    z[mpg %in% z[wt > 3.5, unique(mpg)] ]
}

useif <- function() {
    z[, if(any(wt > 3.5)) .SD, by=mpg]
}

equijoin <- function() {
    z[unique(z[wt > 3.5, "mpg"]), on=.(mpg)]
}

identical(setorder(op(), mpg), setorder(dtUniq(), mpg))
#[1] TRUE

identical(setorder(op(), mpg), setorder(useif(), mpg))
#[1] TRUE

identical(setorder(op(), mpg), setorder(equijoin(), mpg))
#[1] TRUE

library(microbenchmark)
microbenchmark(op(), dtUniq(), useif(), equijoin(), times=3L)

timings:

Unit: seconds
       expr      min       lq     mean   median       uq      max neval
       op() 4.386059 4.546217 4.676105 4.706374 4.821128 4.935883     3
   dtUniq() 4.340680 4.562036 4.658450 4.783393 4.817335 4.851276     3
    useif() 1.098922 1.197751 1.238395 1.296581 1.308131 1.319682     3
 equijoin() 2.345496 2.389649 2.424599 2.433802 2.464151 2.494500     3

The timings will ultimately depend on the dimensions of your dataset and other requirements.

Upvotes: 2

Jaap
Jaap

Reputation: 83275

z[wt > 3.5, "mpg"] will give you a data.table back, you need to unquote mpg to get a vector:

ids <- z[wt > 3.5, mpg] 
z[mpg %in% ids]

which gives:

     mpg cyl  disp  hp drat    wt  qsec vs am gear carb
 1: 14.3   8 360.0 245 3.21 3.570 15.84  0  0    3    4
 2: 19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4
 3: 16.4   8 275.8 180 3.07 4.070 17.40  0  0    3    3
 4: 17.3   8 275.8 180 3.07 3.730 17.60  0  0    3    3
 5: 15.2   8 275.8 180 3.07 3.780 18.00  0  0    3    3
 6: 10.4   8 472.0 205 2.93 5.250 17.98  0  0    3    4
 7: 10.4   8 460.0 215 3.00 5.424 17.82  0  0    3    4
 8: 14.7   8 440.0 230 3.23 5.345 17.42  0  0    3    4
 9: 15.5   8 318.0 150 2.76 3.520 16.87  0  0    3    2
10: 15.2   8 304.0 150 3.15 3.435 17.30  0  0    3    2
11: 13.3   8 350.0 245 3.73 3.840 15.41  0  0    3    4
12: 19.2   8 400.0 175 3.08 3.845 17.05  0  0    3    2
13: 15.0   8 301.0 335 3.54 3.570 14.60  0  1    5    8

Alternatively you could also make use of the ability of data.table to set keys:

z <- data.table(mtcars, key = 'mpg')
ids <- z[wt > 3.5, unique(mpg)] 
z[.(ids)]

Upvotes: 4

Related Questions