data.table update by reference for different columns

Viewed 46

I tried to sets columns trtxxp and trtxxa as NA, when the sdx is NA.

dummy data:

library(data.table)
dt <- data.table(sd1 = c(1:3, NA, 4:5, NA, 6:10, NA, NA),
                 sd2 = c(1:5, NA, 6:7, NA, NA, 10:13),
                 trt01p = "p", trt01a = "a",
                 trt02p = "p", trt02a = "a")
dt
    sd1 sd2 trt01p trt01a trt02p trt02a
 1:   1   1      p      a      p      a
 2:   2   2      p      a      p      a
 3:   3   3      p      a      p      a
 4:  NA   4      p      a      p      a
 5:   4   5      p      a      p      a
 6:   5  NA      p      a      p      a
 7:  NA   6      p      a      p      a
 8:   6   7      p      a      p      a
 9:   7  NA      p      a      p      a
10:   8  NA      p      a      p      a
11:   9  10      p      a      p      a
12:  10  11      p      a      p      a
13:  NA  12      p      a      p      a
14:  NA  13      p      a      p      a

I know I could achieved by below lines:

dt[is.na(sd1), `:=`(trt01p = NA,
                    trt01a = NA)]

dt[is.na(sd2), `:=`(trt02p = NA,
                    trt02a = NA)]
dt
    sd1 sd2 trt01p trt01a trt02p trt02a
 1:   1   1      p      a      p      a
 2:   2   2      p      a      p      a
 3:   3   3      p      a      p      a
 4:  NA   4   <NA>   <NA>      p      a
 5:   4   5      p      a      p      a
 6:   5  NA      p      a   <NA>   <NA>
 7:  NA   6   <NA>   <NA>      p      a
 8:   6   7      p      a      p      a
 9:   7  NA      p      a   <NA>   <NA>
10:   8  NA      p      a   <NA>   <NA>
11:   9  10      p      a      p      a
12:  10  11      p      a      p      a
13:  NA  12   <NA>   <NA>      p      a
14:  NA  13   <NA>   <NA>      p      a

But I since a got a lots of columns, so I tried with .SD, lapply and .SDcols but failed (only trt01p is update correctly)

trt.col <- c("trt01p", "trt01a", "trt02a", "trt02p")
sd.col <- c("sd1", "sd2")

dt[, (trt.col) := lapply(.SD, function(x) ifelse(is.na(x), NA, get(trt.col))),
     .SDcols = sort(c(sd.col, sd.col))][]
dt
    sd1 sd2 trt01p trt01a trt02p trt02a
 1:   1   1      p      p      p      p
 2:   2   2      p      p      p      p
 3:   3   3      p      p      p      p
 4:  NA   4   <NA>   <NA>      p      p
 5:   4   5      p      p      p      p
 6:   5  NA      p      p   <NA>   <NA>
 7:  NA   6   <NA>   <NA>      p      p
 8:   6   7      p      p      p      p
 9:   7  NA      p      p   <NA>   <NA>
10:   8  NA      p      p   <NA>   <NA>
11:   9  10      p      p      p      p
12:  10  11      p      p      p      p
13:  NA  12   <NA>   <NA>      p      p
14:  NA  13   <NA>   <NA>      p      p

Any suggestions on how to do this? Thank you.

1 Answers

I think MichaelChirico's suggest of a for loop may look like this:

cols <- list(sd1=c("trt01p", "trt01a"), sd2=c("trt02a", "trt02p"))
for (col in names(cols)) set(dt, which(is.na(dt[[col]])), cols[[col]], value = NA)
dt
#       sd1   sd2 trt01p trt01a trt02p trt02a
#     <int> <int> <char> <char> <char> <char>
#  1:     1     1      p      a      p      a
#  2:     2     2      p      a      p      a
#  3:     3     3      p      a      p      a
#  4:    NA     4   <NA>   <NA>      p      a
#  5:     4     5      p      a      p      a
#  6:     5    NA      p      a   <NA>   <NA>
#  7:    NA     6   <NA>   <NA>      p      a
#  8:     6     7      p      a      p      a
#  9:     7    NA      p      a   <NA>   <NA>
# 10:     8    NA      p      a   <NA>   <NA>
# 11:     9    10      p      a      p      a
# 12:    10    11      p      a      p      a
# 13:    NA    12   <NA>   <NA>      p      a
# 14:    NA    13   <NA>   <NA>      p      a

(Though I think I'm missing a level of data.table-elegance somewhere.)

The named list is providing a sort of dependency: the names represent the columns that you want to test for NA values, and the contents for each of them are the columns that need to be updated in the presence of the condition.

Related