Allow setnafill() to accept logical vectors in cols (#4113)

NEWS.md

@@ -36,6 +36,8 @@
 
 7. Rows can now be deleted by reference using `DT[i, .ROW := NULL]`, avoiding a full copy of the table for large row-removal operations, [#635](https://github.com/Rdatatable/data.table/issues/635). This has been one of data.table's most requested features. Target rows must be selected with the `i` expression, `by`/`keyby` are not supported, and keys/indices are cleared after deletion. The new experimental helper `setallocrow()` prepares columns for by-reference row operations. Thanks @arunsrinivasan for the feature request, @ben-schwen for the implementation, and @aitap for review and assistance.
 
+8. `setnafill()` now accepts a logical vector for the `cols` argument, which must be the same length as the number of columns in `x`, [#4113](https://github.com/Rdatatable/data.table/issues/4113). Thanks to @MichaelChirico for the suggestion and @venom1204 for the PR.
+
 ### BUG FIXES
 
 1. `fread()` with `skip=0` and `(header=TRUE|FALSE)` no longer skips the first row when it has fewer fields than subsequent rows, [#7463](https://github.com/Rdatatable/data.table/issues/7463). Thanks @emayerhofer for the report and @ben-schwen for the fix.

R/shift.R

@@ -33,5 +33,7 @@ nafill = function(x, type=c("const","locf","nocb"), fill=NA, nan=NA) {
 
 setnafill = function(x, type=c("const","locf","nocb"), fill=NA, nan=NA, cols=seq_along(x)) {
   type = match.arg(type)
+  if (!is.list(x)) stop("in-place update is supported only for list", call. = FALSE)
+  cols = .Call(CcolnamesInt, x, cols, FALSE, FALSE)
   invisible(.Call(CnafillR, x, type, fill, nan_is_na(nan), TRUE, cols))
 }

inst/tests/tests.Rraw

@@ -21850,3 +21850,10 @@ a = c(a=1L)
 test(2379.12, as.IDate(d) + 1L, as.IDate(as.Date(d) + 1L))
 test(2379.13, as.IDate("2020-01-01") + a, as.IDate(as.Date("2020-01-01") + a))
 test(2379.14, as.IDate(d) - 1L, as.IDate(as.Date(d) - 1L))
+
+# #4113 setnafill could accept cols=logical(ncol(x))
+test(2380.01, {DT = data.table(a=c(1,NA,3), b=c(4,NA,6), c=c(7,NA,9)); DT1 = copy(DT); DT2 = copy(DT); setnafill(DT1, type="locf", cols=c(TRUE,FALSE,TRUE)); setnafill(DT2, type="locf", cols=c(1L,3L)); identical(DT1, DT2)}, TRUE)
+test(2380.02, {DT = data.table(a=c(1,NA), b=c(2,NA), c=c(3,NA)); DT1 = copy(DT); DT2 = copy(DT); setnafill(DT1, type="locf", cols=c(TRUE,FALSE,TRUE)); setnafill(DT2, type="locf", cols=c("a","c")); identical(DT1, DT2)}, TRUE)
+test(2380.03, {DT = data.table(a=c(1,NA), b=c(2,NA)); before = copy(DT); setnafill(DT, type="locf", cols=c(FALSE,FALSE)); identical(DT, before)}, TRUE)
+test(2380.04, {DT = data.table(a=c(1,NA), b=c("x",NA), c=c(3,NA)); cols = sapply(DT, is.numeric); setnafill(DT, type="locf", cols=cols); DT}, data.table(a=c(1,1), b=c("x",NA), c=c(3,3)))
+test(2380.05, setnafill(data.table(a=1,b=2,c=3), type="locf", cols=c(TRUE,NA,FALSE)), error="argument specifying columns is a logical vector containing NA at position 2")

man/nafill.Rd

@@ -18,7 +18,7 @@ setnafill(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA, cols=seq_along(x)
   \item{type}{ Character, one of \emph{"const"}, \emph{"locf"} or \emph{"nocb"}. Defaults to \code{"const"}. }
   \item{fill}{ Value to be used to replace missing observations. See examples. }
   \item{nan}{ Either \code{NaN} or \code{NA}; if the former, \code{NaN} is treated as distinct from \code{NA}, otherwise, they are treated the same during replacement. See Examples. }
-  \item{cols}{ Numeric or character vector specifying columns to be updated. }
+  \item{cols}{ Numeric, character or logical vector specifying columns to be updated. A logical vector must be the same length as the number of columns in \code{x}. }
 }
 \details{
   Supported types are \emph{logical}, \emph{integer}, \emph{double}, \emph{character}, and \emph{factor}, as well as classes built on top of these such as \code{Date}, \code{IDate}, and \code{POSIXct}.

src/utils.c

@@ -156,6 +156,20 @@ SEXP colnamesInt(SEXP x, SEXP cols, SEXP check_dups, SEXP skip_absent) {
       else if(bskip_absent && icols[i]>nx)
         icols[i] = 0L;
     }
+  } else if (isLogical(cols) && nc == nx) { 
+    int count = 0;
+    const int *bcols = LOGICAL_RO(cols);
+    for (int i=0; i<nc; i++) {
+      if (bcols[i] == 1) count++;
+      else if (bcols[i] == NA_LOGICAL)
+        error(_("argument specifying columns is a logical vector containing NA at position %d"), i+1);
+    }
+    ricols = PROTECT(allocVector(INTSXP, count)); protecti++;
+    int *icols = INTEGER(ricols);
+    int target = 0;
+    for (int i=0; i<nc; i++) {
+      if (bcols[i] == 1) icols[target++] = i + 1;
+    }
   } else if (isString(cols)) {
     SEXP xnames = PROTECT(getAttrib(x, R_NamesSymbol)); protecti++;
     if (isNull(xnames))