我如何在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"]

现在假设你的功能更复杂，至少有两种方法可以解决这个问题。

• 首先，默认分组列的值为length-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")]
  

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