在R中的data.table中进行快速子集筛选

Question

给定一个 data.table，我想快速地对其中的项目进行子集筛选。例如：

dt = data.table(a=1:10, key="a")
dt[a > 3 & a <= 7]

这还是比较慢的。我知道我可以使用连接获取单独的行，但是有没有一种方法可以利用已排序的 data.table 快速获取这种子集？

这是我正在做的：

dt1 = data.table(id = 1, ym = c(199001, 199006, 199009, 199012), last_ym = c(NA, 199001, 199006, 199009), v = 1:4, key=c("id", "ym"))
dt2 = data.table(id = 1, ym = c(199001, 199002, 199003, 199004, 199005, 199006, 199007, 199008, 199009, 199010, 199011, 199012), v2 = 1:12, key=c("id","ym"))

对于每个id，这里只有一个，在dt1中的ym，我想要将当前ym和dt1中的上一个ym之间v2的值求和。也就是说，对于dt1中的ym == 199006，我想要返回list(v2 = 2 + 3 + 4 + 5 + 6)。这些是dt2中与当前ym相等或小于当前ym的v2的值（不包括前一个ym）。代码如下：

expr = expression({ #browser();
 cur_id = id; 
 cur_ym = ym; 
 cur_dtb = dt2[J(cur_id)][ym <= cur_ym & ym > last_ym]; 
 setkey(cur_dtb , ym);
 list(r = sum(cur_dtb$v2))
})

dt1[,eval(expr ),by=list(id, ym)]

- Alex

---

我不认为这有多慢？使用system.time在一个有1000万行的data.table上只花了0.28秒。 - Señor O

你正在使用 for 循环吗？ - Señor O

不，我不是。让我举个例子。 - Alex

你的例子在我的电脑上根本无法运行。 :/ - joran

我对data.table不是很熟练，但也许dt1[dt2,roll = -Inf]是一个可行的方案？ - joran

Answer 1

为了避免逻辑条件，需要对 dt1 和 dt2 进行滚动连接。然后在 id 内将 ym 向前移动一个位置。最后按照 id 和 ym 对 v2 进行求和。

setkey(dt1, id, last_ym)
setkey(dt2, id, ym)
dt1[dt2,, roll = TRUE][
       , list(v2 = v2, ym = c(last_ym[1], head(ym, -1))), by = id][
       , list(v2 = sum(v2)), by = list(id, ym)]

请注意，我们希望从last_ym开始对所有内容求和，因此dt1上的键必须是last_ym而不是ym。

结果如下：

   id     ym v2
1:  1 199001  1
2:  1 199006 20
3:  1 199009 24
4:  1 199012 33

更新：更正

- G. Grothendieck

---

抱歉，我有点困惑所有的 setkey 表达式... 它们看起来不正确...？ - Alex

我已经修正了它们以对应我实际使用的内容，并添加了一些解释。 - G. Grothendieck

Answer 2

无论 data.table 是否已排序，您都受限于首先评估 a > 3 & a <= 7 所需的时间：

> dt = data.table(a=1:10000000, key="a")
> system.time(dt$a > 3 & dt$a <= 7)
   user  system elapsed 
   0.18    0.01    0.20 
> system.time(dt[,a > 3 & a <= 7])
   user  system elapsed 
   0.18    0.05    0.24 
> system.time(dt[a > 3 & a <= 7])
   user  system elapsed 
   0.25    0.07    0.31

另一种方法：

> system.time({Indices = dt$a > 3 & dt$a <= 7 ; dt[Indices]})
user  system elapsed 
0.28    0.03    0.31 

多个子集

如果您在临时基础上分解因素而不是首先全部完成，则可能会出现速度问题：

> dt <- data.table(A=sample(letters, 10000, replace=T))
> system.time(for(i in unique(dt$A)) dt[A==i])
   user  system elapsed 
   5.16    0.42    5.59 
> system.time(dt[,.SD,by=A])
   user  system elapsed 
   0.32    0.03    0.36

- Señor O

---

刚刚发布了一个示例，也许你能想到更好的方法来完成这个任务。 - Alex