Conditionally calculate the required aggregates using data.table

user_9136348 0
Viewed 65

Consider following snippet:

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

I want to calculate some aggregates (unique count of valCol, number of rows) grouping keyCol column:

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

But I want to calculate these aggregates conditionally, i.e. sometimes I want unique only, sometimes I want count only, and sometimes I want both. One possible solution is to use multiple if else conditions:

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

Coming from Python background, above if else ladder seems like writing bunch of redundant codes. For example, on pandas, we can pass tuple containing the column names, and aggregate functions, or we can unpack a dictionary containing column name, and aggregate functions.

Is there simple way of doing the same using data.table, i.e. passing column names, and aggregates function using some variables, or passing them conditionally?

data.table aggregate r
2 Answers

This might not be quite what you're looking for, but maybe it will get you started.

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
user_94634890

Your specific example is a bit tricky because .N is not a function, it's a symbol. I'll come back to that later, but let's first assume you only have functions. You could write your own helper that helps you achieve what you want:

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

With the helper, you can provide different number of columns by specifing .SDcols, and Map applies the functions you supply.

However, you must know that data.table can often optimize calls based on what it "sees", and hiding some logic behind helper means you may lose that. Take this example:

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

You see there GForce TRUE. If you enable verbosity in the call within helper, you'll notice GForce is always FALSE.

Now, regarding .N, you could also write a helper for that and "wrap" the symbol with a function, but it's honestly kind of hacky. My version simply searches for .N in the stack's environments until it finds it (so it assumes data.table will have computed that at some point in the stack, which in this case is true, 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
user_57939050
Related