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

* incorporate review comments

---------

Co-authored-by: Benjamin Schwendinger <benjaminschwe@gmail.com>

Author: venom1204

NEWS.md

@@ -46,6 +46,8 @@
 
 12. `frank()` gains an `order` argument (matching `frankv()`) and now intercepts the unary minus symbol (e.g., `frank(-dates)`) to support reverse ranking even for types where unary `-` is not defined in R, such as `Date` or `character` vectors, [#5489](https://github.com/Rdatatable/data.table/issues/5489). Thanks @hope-data-science for the request and @venom1204 for the implementation.
 
+13. `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,10 @@ 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.logical(cols)) {
+    if (length(cols) != length(x)) stopf("'cols' is a logical vector of length %d but there are %d columns", length(cols), length(x))
+    if (anyNA(cols)) stopf("'cols' contains NA at position %d", which(is.na(cols))[1L])
+    cols = which(cols)
+  }
   invisible(.Call(CnafillR, x, type, fill, nan_is_na(nan), TRUE, cols))
 }

inst/tests/tests.Rraw

@@ -21974,3 +21974,14 @@ test(2386.10, frank(-dates, order=-1L), frankv(dates, order=-1L), warning=warn)
 test(2386.11, frank(+dates), frankv(dates, order=1L))
 chars = c("b", "a", "c", "a")
 test(2386.12, frank(-chars), frankv(chars, order=-1L))
+
+# setnafill accepts logical cols #4113
+DT = data.table(a=c(1,NA,3), b=c(4,NA,6), c=c(7,NA,9))
+test(2387.01, setnafill(copy(DT), type="locf", cols=c(TRUE,FALSE,TRUE)), setnafill(copy(DT), type="locf", cols=c(1L,3L)))
+test(2387.02, setnafill(copy(DT), type="locf", cols=c(TRUE,FALSE,TRUE)), setnafill(copy(DT), type="locf", cols=c("a","c")))
+DT2 = data.table(a=c(1,NA), b=c(2,NA))
+test(2387.03, setnafill(copy(DT2), type="locf", cols=integer()), setnafill(copy(DT2), type="locf", cols=c(FALSE,FALSE)))
+DT3 = data.table(a=c(1,NA), b=c("x",NA), c=c(3,NA))
+test(2387.04, setnafill(copy(DT3), type="locf", cols=sapply(DT3, is.numeric)), data.table(a=c(1,1), b=c("x",NA), c=c(3,3)))
+test(2387.05, setnafill(copy(DT3), type="locf", cols=c(TRUE,NA,FALSE)), error="'cols' contains NA at position 2")
+test(2387.06, setnafill(copy(DT3), type="locf", cols=c(TRUE,FALSE)), error="'cols' is a logical vector of length 2 but there are 3 columns")

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}.