使用 data.table 有条件地计算所需的聚合

Question

考虑以下代码段：

library("data.table")
dT <- data.table(keyCol=sample(x = c("A", "B", "C"), size = 20, replace = TRUE), 
                 valCol=rpois(n = 20, lambda=10))
# head(dT)

keyCol  valCol
<chr>   <int>
A   11
C   14
C   9
B   9
C   11
C   10

我想计算一些聚合（valCol 的唯一计数，行数）分组 keyCol 列：

res <- dT[, .(unique__=length(unique(valCol)), count__=.N), by="keyCol"]
# res

keyCol  unique__    count__
<chr>   <int>   <int>
A   4   4
C   5   8
B   6   8

但我想有条件地计算这些聚合，即有时我只想要唯一，有时我只想要计数，有时我想要两者。一种可能的解决方案是使用多个 if else 条件：

getCount <- TRUE
getUnique <- TRUE
aggColName <- "valCol"

if(getCount & getUnique){
    res <- dT[, .(unique__=length(unique(valCol)), count__=.N), by="keyCol"]
} else if (getCount){
    res <- dT[, .(count__=.N), by="keyCol"]
} else {
    res <- dT[, .(unique__=length(unique(valCol))), by="keyCol"]    
}
# res

keyCol  unique__    count__
<chr>   <int>   <int>
A   4   4
C   5   8
B   6   8

来自Python背景，if else阶梯上面好像写了一堆多余的代码。例如，在 pandas 上，我们可以传递包含列名和聚合函数的元组，或者我们可以解压包含列名和聚合函数的字典。

是否有使用 data.table 执行相同操作的简单方法，即使用某些变量传递列名和聚合函数，或有条件地传递它们？

Answer 1

这可能不是您要找的东西，但也许可以帮助您入门。

library("data.table")

dt <- data.table(keyCol=sample(x = c("A", "B", "C"), size = 20, replace = TRUE), 
                 valCol1=rpois(n = 20, lambda=10),
                 valCol2=sample(20))

chrVals <- rep(c("valCol1", "valCol2"), each = 3)
lFuns <- rep(list(length, uniqueN, sum), 2)
nms <- c("length1", "unique1", "sum1", "length2", "unique2", "sum2")

dt[, setNames(lapply(seq_along(chrVals), function(i) lFuns[[i]](.SD[[i]])), nms), .SDcols = chrVals, by = "keyCol"]
#>    keyCol length1 unique1 sum1 length2 unique2 sum2
#> 1:      B       4       3   46       4       4   39
#> 2:      C      11       9  104      11      11  129
#> 3:      A       5       5   39       5       5   42

idx <- c(1,3:5)
dt[, setNames(lapply(idx, function(i) lFuns[[i]](.SD[[i]])), nms[idx]), .SDcols = chrVals, by = "keyCol"]
#>    keyCol length1 sum1 length2 unique2
#> 1:      B       4   46       4       4
#> 2:      C      11  104      11      11
#> 3:      A       5   39       5       5

Answer 2

你的具体例子有点棘手，因为 .N 不是一个函数，这是一个象征。稍后我会回过头来，但我们首先假设您只有功能。您可以编写自己的助手来帮助您实现您想要的目标：

helper <- function(dt, by, colNames, funs) {
    dt[, by = by, .SDcols = colNames, Map(funs, .SD, f = function(fn, col) {
        fn(col)
    })]
}

helper(dT, "keyCol", "valCol", list(
    unique__ = function(x) { length(unique(x)) }
))
#    keyCol unique__
# 1:      B        8
# 2:      A        4
# 3:      C        2

helper(dT, "keyCol", "valCol", list(
    max__ = max
))
#    keyCol max__
# 1:      B    13
# 2:      A    11
# 3:      C    13

helper(dT, "keyCol", "valCol", list(
    unique__ = function(x) { length(unique(x)) },
    max__ = max
))
#    keyCol unique__ max__
# 1:      B        8    13
# 2:      A        4    11
# 3:      C        2    13

借助助手，您可以通过指定 .SDcols 来提供不同数量的列， Map 应用您提供的功能。

但是，你必须知道 data.table 经常可以根据它“看到”的内容来优化调用，在 helper 背后隐藏一些逻辑意味着你可能会失去它。举个例子：

dT[, max(valCol), by = keyCol, verbose = TRUE]
Detected that j uses these columns: valCol 
Finding groups using forderv ... forder.c received 20 rows and 1 columns
0.000s elapsed (0.000s cpu) 
Finding group sizes from the positions (can be avoided to save RAM) ... 0.000s elapsed (0.000s cpu) 
Getting back original order ... forder.c received a vector type 'integer' length 3
0.000s elapsed (0.000s cpu) 
lapply optimization is on, j unchanged as 'max(valCol)'
GForce optimized j to 'gmax(valCol)'
Making each group and running j (GForce TRUE) ... gforce initial population of grp took 0.000
gforce assign high and low took 0.000
gforce eval took 0.000
0.000s elapsed (0.001s cpu) 
   keyCol V1
1:      B 13
2:      A 11
3:      C 13

你看到那里GForce TRUE。如果您在 helper 内的调用中启用详细信息，你会注意到 GForce 总是 FALSE.

现在，关于 .N，您还可以为此编写一个帮助程序并使用函数“包装”符号，但老实说，这有点老套。我的版本只是在堆栈的环境中搜索 .N 直到找到它（所以它假设 data.table 会计算出在堆栈中的某个点，在这种情况下是正确的，YMMV)。

count <- function(ignored) {
    n <- 1L
    e <- parent.frame(n)
    ans <- get0(".N", e, inherits = FALSE)
    while (is.null(ans) && !identical(e, .GlobalEnv)) {
        n <- n + 1L
        e <- parent.frame(n)
        ans <- get0(".N", e, inherits = FALSE)
    }
    ans
}

helper(dT, "keyCol", "valCol", list(
    unique__ = function(x) { length(unique(x)) },
    count__ = count
))
#    keyCol unique__ count__
# 1:      B        8      10
# 2:      A        4       6
# 3:      C        2       4

使用 data.table 有条件地计算所需的聚合

Conditionally calculate the required aggregates using data.table

aggregate

r

data.table