data.table SD returns as many rows as asked per group with NA fill instead of as many as exist

Question

I am trying to mimic some dplyr code in data.table, which I rarely use. Not getting the result I want.

For example, with dplyr, when I ask for up to 5 rows per group, using slice, I get the desired result - as many rows as I ask, or as many exist per group, whichever is lower:

library(dplyr)
mtcars %>%
  group_by(cyl, am) %>%
  slice(1:5)

    mpg cyl  disp  hp drat    wt  qsec vs am gear carb
1  24.4   4 146.7  62 3.69 3.190 20.00  1  0    4    2
2  22.8   4 140.8  95 3.92 3.150 22.90  1  0    4    2
3  21.5   4 120.1  97 3.70 2.465 20.01  1  0    3    1
4  22.8   4 108.0  93 3.85 2.320 18.61  1  1    4    1
5  32.4   4  78.7  66 4.08 2.200 19.47  1  1    4    1
6  30.4   4  75.7  52 4.93 1.615 18.52  1  1    4    2
7  33.9   4  71.1  65 4.22 1.835 19.90  1  1    4    1
8  27.3   4  79.0  66 4.08 1.935 18.90  1  1    4    1
9  21.4   6 258.0 110 3.08 3.215 19.44  1  0    3    1
10 18.1   6 225.0 105 2.76 3.460 20.22  1  0    3    1
11 19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4
12 17.8   6 167.6 123 3.92 3.440 18.90  1  0    4    4
13 21.0   6 160.0 110 3.90 2.620 16.46  0  1    4    4
14 21.0   6 160.0 110 3.90 2.875 17.02  0  1    4    4
15 19.7   6 145.0 175 3.62 2.770 15.50  0  1    5    6
16 18.7   8 360.0 175 3.15 3.440 17.02  0  0    3    2
17 14.3   8 360.0 245 3.21 3.570 15.84  0  0    3    4
18 16.4   8 275.8 180 3.07 4.070 17.40  0  0    3    3
19 17.3   8 275.8 180 3.07 3.730 17.60  0  0    3    3
20 15.2   8 275.8 180 3.07 3.780 18.00  0  0    3    3
21 15.8   8 351.0 264 4.22 3.170 14.50  0  1    5    4
22 15.0   8 301.0 335 3.54 3.570 14.60  0  1    5    8

But, with data.table and .SD, I get as many rows as I ask, no matter if fewer rows exist per group, and the extra rows are filled with NA in non-grouped columns. Any idea how to get what I get from above dplyr code with data.table?

library(data.table)
as.data.table(mtcars)[, .SD[1:5], by = .(cyl, am)]
    cyl am  mpg  disp  hp drat    wt  qsec vs gear carb
 1:   6  1 21.0 160.0 110 3.90 2.620 16.46  0    4    4
 2:   6  1 21.0 160.0 110 3.90 2.875 17.02  0    4    4
 3:   6  1 19.7 145.0 175 3.62 2.770 15.50  0    5    6
 4:   6  1   NA    NA  NA   NA    NA    NA NA   NA   NA
 5:   6  1   NA    NA  NA   NA    NA    NA NA   NA   NA
 6:   4  1 22.8 108.0  93 3.85 2.320 18.61  1    4    1
 7:   4  1 32.4  78.7  66 4.08 2.200 19.47  1    4    1
 8:   4  1 30.4  75.7  52 4.93 1.615 18.52  1    4    2
 9:   4  1 33.9  71.1  65 4.22 1.835 19.90  1    4    1
10:   4  1 27.3  79.0  66 4.08 1.935 18.90  1    4    1
11:   6  0 21.4 258.0 110 3.08 3.215 19.44  1    3    1
12:   6  0 18.1 225.0 105 2.76 3.460 20.22  1    3    1
13:   6  0 19.2 167.6 123 3.92 3.440 18.30  1    4    4
14:   6  0 17.8 167.6 123 3.92 3.440 18.90  1    4    4
15:   6  0   NA    NA  NA   NA    NA    NA NA   NA   NA
16:   8  0 18.7 360.0 175 3.15 3.440 17.02  0    3    2
17:   8  0 14.3 360.0 245 3.21 3.570 15.84  0    3    4
18:   8  0 16.4 275.8 180 3.07 4.070 17.40  0    3    3
19:   8  0 17.3 275.8 180 3.07 3.730 17.60  0    3    3
20:   8  0 15.2 275.8 180 3.07 3.780 18.00  0    3    3
21:   4  0 24.4 146.7  62 3.69 3.190 20.00  1    4    2
22:   4  0 22.8 140.8  95 3.92 3.150 22.90  1    4    2
23:   4  0 21.5 120.1  97 3.70 2.465 20.01  1    3    1
24:   4  0   NA    NA  NA   NA    NA    NA NA   NA   NA
25:   4  0   NA    NA  NA   NA    NA    NA NA   NA   NA
26:   8  1 15.8 351.0 264 4.22 3.170 14.50  0    5    4
27:   8  1 15.0 301.0 335 3.54 3.570 14.60  0    5    8
28:   8  1   NA    NA  NA   NA    NA    NA NA   NA   NA
29:   8  1   NA    NA  NA   NA    NA    NA NA   NA   NA
30:   8  1   NA    NA  NA   NA    NA    NA NA   NA   NA
    cyl am  mpg  disp  hp drat    wt  qsec vs gear carb

user_3732271 · Answer

slice automatically readjusts based on the number of rows per group, here we can use head. With data.table, indexing if we provide index and the index is greater than the .N (i.e. the total number of rows per group), it will create an NA row for each index that is outside the range

library(data.table)
as.data.table(mtcars)[, head(.SD, 5), by = .(cyl, am)]
#    cyl am  mpg  disp  hp drat    wt  qsec vs gear carb
# 1:   6  1 21.0 160.0 110 3.90 2.620 16.46  0    4    4
# 2:   6  1 21.0 160.0 110 3.90 2.875 17.02  0    4    4
# 3:   6  1 19.7 145.0 175 3.62 2.770 15.50  0    5    6
# 4:   4  1 22.8 108.0  93 3.85 2.320 18.61  1    4    1
# 5:   4  1 32.4  78.7  66 4.08 2.200 19.47  1    4    1
# 6:   4  1 30.4  75.7  52 4.93 1.615 18.52  1    4    2
# 7:   4  1 33.9  71.1  65 4.22 1.835 19.90  1    4    1
# 8:   4  1 27.3  79.0  66 4.08 1.935 18.90  1    4    1
# 9:   6  0 21.4 258.0 110 3.08 3.215 19.44  1    3    1
#10:   6  0 18.1 225.0 105 2.76 3.460 20.22  1    3    1
#11:   6  0 19.2 167.6 123 3.92 3.440 18.30  1    4    4
#12:   6  0 17.8 167.6 123 3.92 3.440 18.90  1    4    4
#13:   8  0 18.7 360.0 175 3.15 3.440 17.02  0    3    2
#14:   8  0 14.3 360.0 245 3.21 3.570 15.84  0    3    4
#15:   8  0 16.4 275.8 180 3.07 4.070 17.40  0    3    3
#16:   8  0 17.3 275.8 180 3.07 3.730 17.60  0    3    3
#17:   8  0 15.2 275.8 180 3.07 3.780 18.00  0    3    3
#18:   4  0 24.4 146.7  62 3.69 3.190 20.00  1    4    2
#19:   4  0 22.8 140.8  95 3.92 3.150 22.90  1    4    2
#20:   4  0 21.5 120.1  97 3.70 2.465 20.01  1    3    1
#21:   8  1 15.8 351.0 264 4.22 3.170 14.50  0    5    4
#22:   8  1 15.0 301.0 335 3.54 3.570 14.60  0    5    8
#    cyl am  mpg  disp  hp drat    wt  qsec vs gear carb

Also, slice should work with data.table

library(dplyr)
as.data.table(mtcars)[, slice(.SD, 1:5), by = .(cyl, am)]

data.table SD returns as many rows as asked per group with NA fill instead of as many as exist

1 Answers