如何在 data.table 中的用户定义函数中获取分组列的值

Question

我想开始使用 data.table 而不是 dplyr，因为我必须加快计算速度。 现在我对组中每个元素的代码使用分组列的值。 但是 data.table 放弃了它。例如

g <- function(x) {
   browser()
}

DT = data.table(x=rep(c("a","b","c"),each=3), y=c(1,3,6), v=1:9)

DT[,list(a = g(.SD)), keyby="x"]

当从浏览器查询x 的值时给出：

   y v
1: 1 1
2: 3 2
3: 6 3

对于组的第一个元素。

有什么方法可以获取 g() 中每个组的 x 值吗？

更新： 我正在使用公式取决于组的函数，例如

g <- function(data) {
   if (x == "a") {
       return(y-v)
   } else {
       return(v-y)
   }
}

Answer 1

首先，假设这或多或少是您的功能（意思是：它很短和/或相似），我会直接在j 中执行此操作，如下所示：

DT[, .(a = (y-v) * (2L*(x=="a") - 1L)), by="x"]
# or if it's too cryptic
DT[, .(a = if (x=="a") y-v else v-y), by="x"]

现在假设您的函数更复杂，至少还有其他两种方法可以解决这个问题。

• 首先，来自分组列的默认值的长度为 1。因此，您可以编写一个带有附加参数的函数，该参数是组值，如下所示：

  foo <- function(dt, grp) {
      if (grp == "a") dt[, y-v]
      else dt[, v-y]
  }
  DT[, .(a = foo(.SD, x)), by="x"]
  

  这里grp 的长度为1（如上所述）。

• 使用.SDcols 也可以在.SD 中添加分组列。

  foo <- function(dt) {
      if (dt$x[1L] == "a") dt[, y-v]
      else dt[, v-y]
  }
  DT[, .(a = foo(.SD)), by="x", .SDcols = c("x", "y", "v")]
  

  虽然我更喜欢第一种方法，因为它不会不必要地创建一个只有一个唯一值的额外列。