34

As simple as it is, I don't know a data.table solution to select the first n rows in groups in a data table. Can you please help me out?

paljenczy
  • 4,779
  • 8
  • 33
  • 46

2 Answers2

70

As an alternative:

dt[, .SD[1:3], cyl]

When you look at speed on the example dataset, the head method is on par with the .I method of @eddi. Comparing with the microbenchmark package:

microbenchmark(head = dt[, head(.SD, 3), cyl],
               SD = dt[, .SD[1:3], cyl], 
               I = dt[dt[, .I[1:3], cyl]$V1],
               times = 10, unit = "relative")

results in:

Unit: relative
 expr      min       lq     mean   median       uq       max neval cld
 head 1.000000 1.000000 1.000000 1.000000 1.000000 1.0000000    10  a 
   SD 2.156562 2.319538 2.306065 2.365190 2.318540 2.1908401    10   b
    I 1.001810 1.029511 1.007371 1.018514 1.016583 0.9442973    10  a

However, data.table is specifically designed for large datasets. So, running this comparison again:

# creating a 30 million dataset
largeDT <- dt[,.SD[sample(.N, 1e7, replace = TRUE)], cyl]
# running the benchmark on the large dataset
microbenchmark(head = largeDT[, head(.SD, 3), cyl],
               SD = largeDT[, .SD[1:3], cyl], 
               I = largeDT[largeDT[, .I[1:3], cyl]$V1],
               times = 10, unit = "relative")

results in:

Unit: relative
 expr      min       lq     mean   median       uq     max neval cld
 head 2.279753 2.194702 2.221330 2.177774 2.276986 2.33876    10   b
   SD 2.060959 2.187486 2.312009 2.236548 2.568240 2.55462    10   b
    I 1.000000 1.000000 1.000000 1.000000 1.000000 1.00000    10  a

Now the .I method is clearly the fastest one.

Update 2016-02-12:

With the most recent development version of the data.table package, the .I method still wins. Whether the .SD method or the head() method is faster seems to depend on the size of the dataset. Now the benchmark gives:

Unit: relative
 expr      min       lq     mean   median       uq      max neval cld
 head 2.093240 3.166974 3.473216 3.771612 4.136458 3.052213    10   b
   SD 1.840916 1.939864 2.658159 2.786055 3.112038 3.411113    10   b
    I 1.000000 1.000000 1.000000 1.000000 1.000000 1.000000    10  a

However with a somewhat smaller dataset (but still quite big), the odds change:

largeDT2 <- dt[,.SD[sample(.N, 1e6, replace = TRUE)], cyl]

the benchmark is now slightly in favor of the head method over the .SD method:

Unit: relative
 expr      min       lq     mean   median       uq      max neval cld
 head 1.808732 1.917790 2.087754 1.902117 2.340030 2.441812    10   b
   SD 1.923151 1.937828 2.150168 2.040428 2.413649 2.436297    10   b
    I 1.000000 1.000000 1.000000 1.000000 1.000000 1.000000    10  a
Community
  • 1
  • 1
Jaap
  • 81,064
  • 34
  • 182
  • 193
  • Was this done with 1.9.7? I know some operations with `.SD` have recently been optimized... – MichaelChirico Feb 11 '16 at 19:49
  • @MichaelChirico Yes, but I'm using a relatively old version of 1.9.7 (shortly after the introduction of `rowid`). So, it depends on how recent recently is ;-) – Jaap Feb 12 '16 at 09:20
  • I think the relevant commits were around the beginning of November – MichaelChirico Feb 12 '16 at 13:37
  • 1
    performance may depends on cardinality of `cyl` – jangorecki Feb 23 '16 at 19:13
  • I don't find the same results anymore, I find `.SD` faster than `head` faster than `.I`. @Jaap do you consider updating your answer ? :) – Samuel Allain Jun 22 '21 at 15:01
  • @SamuelAllain I still get more or less the same result. However this might be due to the number of threads `data.table` is using (this benchmark is run with 1 thread). Its on my todo-list to enable `openmp` on my machine so I can test with more threads. – Jaap Jun 22 '21 at 16:21
  • 1
    @Jaap thanks for testing again. I tested with setDTthreads(1) but still didn't find your results. I think @jangorecki is right, the results highly depend on the the cardinality of the grouping variable. When cardinality is under 6, `.I` wins for me, but otherwise, .SD and head are faster. – Samuel Allain Jun 28 '21 at 10:14
  • I made a simulation with moving cardinality, It indeed affects the ranking of methods. I cannot post my graph on this topic, unfortunately. – Samuel Allain Jun 28 '21 at 13:48
19

We can use head with .SD

library(data.table)

dt <- data.table(mtcars)

> dt[, head(.SD, 3), by = "cyl"]

   cyl  mpg  disp  hp drat    wt  qsec vs am gear carb
1:   6 21.0 160.0 110 3.90 2.620 16.46  0  1    4    4
2:   6 21.0 160.0 110 3.90 2.875 17.02  0  1    4    4
3:   6 21.4 258.0 110 3.08 3.215 19.44  1  0    3    1
4:   4 22.8 108.0  93 3.85 2.320 18.61  1  1    4    1
5:   4 24.4 146.7  62 3.69 3.190 20.00  1  0    4    2
6:   4 22.8 140.8  95 3.92 3.150 22.90  1  0    4    2
7:   8 18.7 360.0 175 3.15 3.440 17.02  0  0    3    2
8:   8 14.3 360.0 245 3.21 3.570 15.84  0  0    3    4
9:   8 16.4 275.8 180 3.07 4.070 17.40  0  0    3    3
paljenczy
  • 4,779
  • 8
  • 33
  • 46
  • 1
    note that actually `head.data.table` (`getAnywhere("head.data.table")`) is just calling a (slightly more robust) version of @Jaap's answer. – MichaelChirico Jan 12 '16 at 20:39