在一个函数调用中向 R data.table 添加多个列?

我有一个函数,它返回列表中的两个值。这两个值都需要在两个新列中添加到 data.table。函数的求值代价很高,因此我希望避免计算函数两次。下面是一个例子:

library(data.table)
example(data.table)
DT
x y  v
1: a 1 42
2: a 3 42
3: a 6 42
4: b 1  4
5: b 3  5
6: b 6  6
7: c 1  7
8: c 3  8
9: c 6  9

这是我函数的一个例子。请记住,我说过这是一个昂贵的计算,除此之外,没有办法从其他给定的值中推断出一个返回值(如下例所示) :

myfun <- function (y, v)
{
ret1 = y + v
ret2 = y - v
return(list(r1 = ret1, r2 = ret2))
}

下面是我在一个语句中添加两个栏目的方法,但是这个语句需要调用 myfun 两次:

DT[,new1:=myfun(y,v)$r1][,new2:=myfun(y,v)$r2]


x y  v new1 new2
1: a 1 42   43  -41
2: a 3 42   45  -39
3: a 6 42   48  -36
4: b 1  4    5   -3
5: b 3  5    8   -2
6: b 6  6   12    0
7: c 1  7    8   -6
8: c 3  8   11   -5
9: c 6  9   15   -3

对于如何做到这一点有什么建议吗?我可以保存在一个单独的环境中的 r2每次我调用我的乐趣,我只是需要一个方法来添加两个列的参考一次。

95029 次浏览

data.table v1.8.3开始,您可以这样做:

DT[, c("new1","new2") := myfun(y,v)]

另一个选项是存储函数的输出并逐个添加列:

z <- myfun(DT$y,DT$v)
head(DT[,new1:=z$r1][,new2:=z$r2])
#      x y  v new1 new2
# [1,] a 1 42   43  -41
# [2,] a 3 42   45  -39
# [3,] a 6 42   48  -36
# [4,] b 1  4    5   -3
# [5,] b 3  5    8   -2
# [6,] b 6  6   12    0

为什么不让你的函数接受一个数据帧并直接返回一个数据帧呢?

myfun <- function (DT)
{
DT$ret1 = with(DT, y + v)
DT$ret2 = with(DT, y - v)
return(DT)
}

要在前一个答案的基础上构建,可以将 lapply与输出多列的函数一起使用。这样就可以在 data.table 的更多列中使用该函数。

 myfun <- function(a,b){
res1 <- a+b
res2 <- a-b
list(res1,res2)
}


DT <- data.table(z=1:10,x=seq(3,30,3),t=seq(4,40,4))
DT


## DT
##     z  x  t
## 1:  1  3  4
## 2:  2  6  8
## 3:  3  9 12
## 4:  4 12 16
## 5:  5 15 20
## 6:  6 18 24
## 7:  7 21 28
## 8:  8 24 32
## 9:  9 27 36
## 10: 10 30 40


col <- colnames(DT)
DT[, paste0(c('r1','r2'),rep(col,each=2)):=unlist(lapply(.SD,myfun,z),
recursive=FALSE),.SDcols=col]
## > DT
##     z  x  t r1z r2z r1x r2x r1t r2t
## 1:  1  3  4   2   0   4   2   5   3
## 2:  2  6  8   4   0   8   4  10   6
## 3:  3  9 12   6   0  12   6  15   9
## 4:  4 12 16   8   0  16   8  20  12
## 5:  5 15 20  10   0  20  10  25  15
## 6:  6 18 24  12   0  24  12  30  18
## 7:  7 21 28  14   0  28  14  35  21
## 8:  8 24 32  16   0  32  16  40  24
## 9:  9 27 36  18   0  36  18  45  27
## 10: 10 30 40  20   0  40  20  50  30

答案无法使用,例如当函数未向量化时。

例如,在下列情况下,它不会按预期的方式工作:

myfun <- function (y, v, g)
{
ret1 = y + v + length(g)
ret2 = y - v + length(g)
return(list(r1 = ret1, r2 = ret2))
}
DT
#    v y                  g
# 1: 1 1                  1
# 2: 1 3                4,2
# 3: 1 6              9,8,6


DT[,c("new1","new2"):=myfun(y,v,g)]
DT
#    v y     g new1 new2
# 1: 1 1     1    5    3
# 2: 1 3   4,2    7    5
# 3: 1 6 9,8,6   10    8

它将始终添加列 g的大小,而不是 g中每个向量的大小

在这种情况下,解决办法是:

DT[, c("new1","new2") := data.table(t(mapply(myfun,y,v,g)))]
DT
#    v y     g new1 new2
# 1: 1 1     1    3    1
# 2: 1 3   4,2    6    4
# 3: 1 6 9,8,6   10    8

如果一个函数返回一个矩阵,你可以通过包装函数来实现相同的行为,首先将矩阵转换成列表。我想知道 data.table 是否应该自动处理它?

matrix2list <- function(mat){
unlist(apply(mat,2,function(x) list(x)),FALSE)
}


DT <- data.table(A=1:10)


myfun <- function(x) matrix2list(cbind(x+1,x-1))


DT[,c("c","d"):=myfun(A)]


##>DT
##      A  c d
##  1:  1  2 0
##  2:  2  3 1
##  3:  3  4 2
##  4:  4  5 3
##  5:  5  6 4
##  6:  6  7 5
##  7:  7  8 6
##  8:  8  9 7
##  9:  9 10 8
## 10: 10 11 9