在单个 R data.table 中按组有效地定位

Posted 2023-03-11

【中文标题】在单个 R data.table 中按组有效地定位

【英文标题】：efficiently locf by groups in a single R data.table

【发布时间】：2016-08-31 19:24:44

【问题描述】：

我有一个大而宽的data.table（20m 行），由一个人 ID 键入，但有很多列（~150）有很多空值。每列都是我希望为每个人发扬光大的记录状态/属性。每个人可能有 10 到 10,000 个观察值，集合中大约有 500,000 人。一个人的价值观不能“渗透”到下一个人，所以我的解决方案必须适当地尊重人员 ID 列和组。

出于演示目的 - 这是一个非常小的示例输入：

DT = data.table(
  id=c(1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3),
  aa=c("A", NA, "B", "C", NA, NA, "D", "E", "F", NA, NA, NA),
  bb=c(NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA),
  cc=c(1, NA, NA, NA, NA, 4, NA, 5, 6, NA, 7, NA)
)

看起来像这样：

    id aa bb cc
 1:  1  A NA  1
 2:  1 NA NA NA
 3:  1  B NA NA
 4:  1  C NA NA
 5:  2 NA NA NA
 6:  2 NA NA  4
 7:  2  D NA NA
 8:  2  E NA  5
 9:  3  F NA  6
10:  3 NA NA NA
11:  3 NA NA  7
12:  3 NA NA NA

我的预期输出如下所示：

    id aa bb cc
 1:  1  A NA  1
 2:  1  A NA  1
 3:  1  B NA  1
 4:  1  C NA  1
 5:  2 NA NA NA
 6:  2 NA NA  4
 7:  2  D NA  4
 8:  2  E NA  5
 9:  3  F NA  6
10:  3  F NA  6
11:  3  F NA  7
12:  3  F NA  7

我找到了一个有效的data.table 解决方案，但它在我的大型数据集上非常慢：

DT[, na.locf(.SD, na.rm=FALSE), by=id]

我发现使用 dplyr 的等效解决方案同样慢。

GRP = DT %>% group_by(id)
data.table(GRP %>% mutate_each(funs(blah=na.locf(., na.rm=FALSE))))

我希望我可以使用data.table 功能提出一个滚动的“自我”加入，但我似乎无法做到正确（我怀疑我需要使用.N，但我只是还没弄明白）。

此时我想我必须在 Rcpp 中写一些东西来有效地应用分组的 locf。

我是 R 新手，但我不是 C++ 新手——所以我有信心我能做到。我只是觉得应该有一种有效的方法在 R 中使用data.table 来做到这一点。

【问题讨论】：

我很确定DT[, lapply(.SD, na.locf, F), by = id] 会更快

我实际上是从那个开始的，发现性能更差。

Rolling self join 在这里看起来很重要，我记得有些问题既有na.locf 又有rolling joins 的答案，所以我想你可以在当前的 SO 知识库中找到答案。

使用有序的“id”，也许你可以使用类似：tmp = c(TRUE, DT$id[-1] != DT$id[-nrow(DT)]); DT[, lapply(.SD, function(x) x[cummax(((!is.na(x)) | tmp) * seq_len(nrow(DT)))])]?

@alexis_laz - 哇！棒极了！它可以工作，并且比 data.table 解决方案快 2 个数量级。你能帮我理解代码在做什么吗？另外，您的评论应该成为答案，以便我可以将其标记为已解决。

【参考方案1】：

一个非常简单的na.locf 可以通过转发（cummax）非NA 索引（(!is.na(x)) * seq_along(x)）和相应的子集来构建：

x = c(1, NA, NA, 6, 4, 5, 4, NA, NA, 2)
x[cummax((!is.na(x)) * seq_along(x))]
# [1] 1 1 1 6 4 5 4 4 4 2

这使用na.rm = TRUE 参数复制na.locf，要获得na.rm = FALSE 行为，我们只需要确保cummax 中的第一个元素是TRUE：

x = c(NA, NA, 1, NA, 2)
x[cummax(c(TRUE, tail((!is.na(x)) * seq_along(x), -1)))]
#[1] NA NA  1  1  2

在这种情况下，我们不仅需要考虑非NA 索引，还需要考虑（已排序或待排序）“id”列更改值的索引：

id = c(10, 10, 11, 11, 11, 12, 12, 12, 13, 13)
c(TRUE, id[-1] != id[-length(id)])
# [1]  TRUE FALSE  TRUE FALSE FALSE  TRUE FALSE FALSE  TRUE FALSE

结合以上：

id = c(10, 10, 11, 11, 11, 12, 12, 12, 13, 13)
x =  c(1,  NA, NA, 6,  4,  5,  4,  NA, NA, 2)

x[cummax(((!is.na(x)) | c(TRUE, id[-1] != id[-length(id)])) * seq_along(x))]
# [1]  1  1 NA  6  4  5  4  4 NA  2

注意，这里我们将OR 的第一个元素与TRUE，即使其等于TRUE，从而得到na.rm = FALSE 行为。

对于这个例子：

id_change = DT[, c(TRUE, id[-1] != id[-.N])]
DT[, lapply(.SD, function(x) x[cummax(((!is.na(x)) | id_change) * .I)])]
#    id aa bb cc
# 1:  1  A NA  1
# 2:  1  A NA  1
# 3:  1  B NA  1
# 4:  1  C NA  1
# 5:  2 NA NA NA
# 6:  2 NA NA  4
# 7:  2  D NA  4
# 8:  2  E NA  5
# 9:  3  F NA  6
#10:  3  F NA  6
#11:  3  F NA  7
#12:  3  F NA  7

【讨论】：

反对票对我来说不是很明显，我们将不胜感激

很棒的答案 imo - 这不仅是常规 na.locf 的更快版本，而且还添加了修改以按组执行此操作（假设已排序组），没有 实际上做了一个by 循环（这将在每个组中引入额外的eval 并且会减慢它的速度）。除非我遗漏了什么——这应该是标准的na.locf 实现，而不是zoo 所做的rle。

@eddi ：感谢您的编辑。我猜zoo::na.locf 更灵活，不过，我相信对于简单的情况，cummax 版本的4-5 * length(x) 扫描应该非常简单。而且，确实证明在函数中传递每个列指针一次并虚拟地“按”组应用是很方便的。

我可以补充一点，在我最初的 20m 行测试集上，第一个建议的 lapply 解决方案需要 40 小时才能完成。新代码只需 4 分钟！我怀疑 Rcpp 可以做得比这更好。

@carl.anderson 我做了一个快速的信封测试，你会很容易地通过 Rcpp 获得 2-3 倍的改进