R Faster way to create lag by entire group using data.table

user_11939840 0
Viewed 263

I have a simple data.table as follows-

  ID = c(rep("A", 1000), rep("B", 1000), rep("C", 1000), rep("D", 1000))
  val = c("a", "a", "a", "b", "b", "c", "c","d","d","d","d","e","e","f","f","g","g","g","g","g")

  dt = data.table(ID, val)

I want to add a new column to this data.table which will have the lag of val by group ID.

Here is the expected output

> head(dt, 20)
     ID val val_lag
 1:  A   a    <NA>
 2:  A   a    <NA>
 3:  A   a    <NA>
 4:  A   b       a
 5:  A   b       a
 6:  A   c       b
 7:  A   c       b
 8:  A   d       c
 9:  A   d       c
10:  A   d       c
11:  A   d       c
12:  A   e       d
13:  A   e       d
14:  A   f       e
15:  A   f       e
16:  A   g       f
17:  A   g       f
18:  A   g       f
19:  A   g       f
20:  A   g       f

The current solution I am using is -

dt[, val_lag := with(rle(val), rep(c(NA, head(values, -1)), lengths)), by = ID]

However, this solution is super slow on the actual dataset, which is very large and has millions of rows. Is there any faster way to solve this problem?

Following is the performance result of all methods discussed in this post -

  microbenchmark::microbenchmark(rles = dt[, val_lag1 := with(rle(val), rep(c(NA, head(values, -1)), lengths)), by = ID],
                                 chinsoon = dt[, val_lag := shift(val)[nafill(replace(seq.int(.N), rowid(rleid(val)) > 1L, NA_integer_), "locf")], by = ID],
                                 TiC = dt[, val_lag3 := c(NA,rle(val)$values)[cumsum(c(0,head(val,-1)!=tail(val,-1)))+1], by = ID],
                                 times = 1000
  )

Unit: milliseconds
     expr      min       lq     mean   median       uq      max neval cld
     rles 1.549548 1.781014 2.750187 2.096805 2.743668 46.65326  1000  a 
 chinsoon 1.766827 2.060233 3.059109 2.379477 3.077080 67.16040  1000  a 
      TiC 1.986808 2.226933 3.472451 2.624236 3.397165 60.67802  1000   b

Thanks!

data.table r
3 Answers

Here is another option:

dt[, val_lag := shift(val)[nafill(replace(seq.int(.N), rowid(rleid(val)) > 1L, NA_integer_), "locf")]]

timing code:

library(data.table)
set.seed(0L)
nr <- 1e6
ng <- 1e5
dt = data.table(ID=sample(ng, nr, TRUE), val=as.character(sample(nr, nr, TRUE)))
setorder(dt, ID, val)

microbenchmark::microbenchmark(times = 3L,
    opt = dt[, val_lag := shift(val)[nafill(replace(seq.int(.N), rowid(rleid(val)) > 1L, NA_integer_), "locf")]],
    rle = dt[, val_lag := with(rle(val), rep(c(NA, head(values, -1)), lengths)), by = ID]
)

timings:

Unit: milliseconds
 expr       min        lq      mean    median        uq       max neval
  opt  133.8857  159.8922  265.2029  185.8987  330.8614  475.8242     3
  rle 3097.6005 3123.5422 3193.2654 3149.4839 3241.0978 3332.7117     3

edit: added an example of what is happening:

index         |    1    2    3    4    5    6    7    8    9   10
value         |    a    a    a    b    b    c    c    c    d    d

shifted (s)   |   NA    a    a    a    b    b    c    c    c    d
rowid+rleid   |    1    2    3    1    2    1    2    3    1    2
replace       |    1   NA   NA    4   NA    6   NA   NA    9   NA <In ?nafill, Only double and integer data types are currently supported. Hence, nafill the indices before accessing>
nafill        |    1    1    1    4    4    6    6    6    9    9
using s above | s[1] s[1] s[1] s[4] s[4] s[6] s[6] s[6] s[9] s[9]
user_19894800

I guess you can try the code like below

dt[,val_tag := c(NA,rle(val)$values)[cumsum(c(0,head(val,-1)!=tail(val,-1)))+1],ID]

and you will see

      ID val val_tag
   1:  A   a    <NA>
   2:  A   a    <NA>
   3:  A   a    <NA>
   4:  A   b       a
   5:  A   b       a
  ---
3996:  D   g       f
3997:  D   g       f
3998:  D   g       f
3999:  D   g       f
4000:  D   g       f

and

> tail(dt,30)
    ID val val_tag
 1:  D   d       c
 2:  D   e       d
 3:  D   e       d
 4:  D   f       e
 5:  D   f       e
 6:  D   g       f
 7:  D   g       f
 8:  D   g       f
 9:  D   g       f
10:  D   g       f
11:  D   a       g
12:  D   a       g
13:  D   a       g
14:  D   b       a
15:  D   b       a
16:  D   c       b
17:  D   c       b
18:  D   d       c
19:  D   d       c
20:  D   d       c
21:  D   d       c
22:  D   e       d
23:  D   e       d
24:  D   f       e
25:  D   f       e
26:  D   g       f
27:  D   g       f
28:  D   g       f
29:  D   g       f
30:  D   g       f
    ID val val_tag
user_121587570

I am not sure whether grouping by ID is necessary in your example. You could basically look up a shifted named vector, which seems faster:

library(data.table)
library(microbenchmark)
ID = c(rep("A", 1000), rep("B", 1000), rep("C", 1000), rep("D", 1000))
val = c("a", "a", "a", "b", "b", "c", "c","d","d","d","d","e","e","f","f","g","g","g","g","g")
dt = data.table(ID, val)
lt <- setNames(c(NaN, seq_along(unique(val))), c(NA_character_, unique(val)))

microbenchmark(
    rle = dt[, val_lag := with(rle(val), rep(c(NA, head(values, -1)), lengths)), by = ID],
    TiC = dt[, val_lag := shift(unique(val))[as.integer(factor(paste(ID, val)))], ID], 
    me = dt[, val_lag := names(lt)[lt[val]]], 
    control = list(warmup=10)
)
#> Unit: microseconds
#>  expr      min        lq      mean    median       uq      max neval cld
#>   rle  614.544  653.2165  772.9975  775.7005  844.245 1391.390   100  b 
#>   TiC 1249.129 1286.2355 1578.1695 1412.4135 1553.035 6148.756   100   c
#>    me  330.570  346.1440  414.7982  386.9125  440.422  910.842   100 a

identical(dt[, val_lag:=names(lt)[lt[val]]],  
          dt[, val_lag := with(rle(val), rep(c(NA, head(values, -1)), lengths)), by = ID])
#> [1] TRUE

Created on 2021-01-31 by the reprex package (v1.0.0)

user_127287480
Related