add na.rm capabilites

Author: ben-schwen

NEWS.md

@@ -50,6 +50,8 @@
 
 14. `first()` and `last()` with `n>1` are now GForce optimized (e.g. `DT[, first(x, n=3), by=grp]`), [#4239](https://github.com/Rdatatable/data.table/issues/4239). Also adds a new internal `gforce_dynamic` mechanism to track any GForce result which returns other than exactly 1 row per group so that results are correctly replicated. Thanks to @nbenn for the report and @ben-schwen and @mattdowle for the implementation.
 
+15. `first()` and `last()` gain an `na.rm` argument to skip missing values (e.g. `DT[, first(x, na.rm=TRUE), by=grp]`), [#4239](https://github.com/Rdatatable/data.table/issues/4239) and [#4446](https://github.com/Rdatatable/data.table/issues/4446). A group with no non-missing values returns `NA` for `n=1` (matching `median()`/`var()`), or zero rows for `n>1`. GForce optimized. Thanks to @nbenn and @MichaelChirico for the reports and @ben-schwen and @mattdowle for the implementation.
+
 ### 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/data.table.R

@@ -3337,8 +3337,8 @@ gfuns = c(gdtfuns,
 `g[` = `g[[` = function(x, n) .Call(Cgnthvalue, x, as.integer(n)) # n is of length=1 here.
 ghead = function(x, n) .Call(Cghead, x, as.integer(n))
 gtail = function(x, n) .Call(Cgtail, x, as.integer(n))
-gfirst = function(x, n=1L) .Call(Cgfirst, x, as.integer(n))
-glast = function(x, n=1L) .Call(Cglast, x, as.integer(n))
+gfirst = function(x, n=1L, na.rm=FALSE) .Call(Cgfirst, x, as.integer(n), na.rm)
+glast = function(x, n=1L, na.rm=FALSE) .Call(Cglast, x, as.integer(n), na.rm)
 gsum = function(x, na.rm=FALSE) .Call(Cgsum, x, na.rm)
 gmean = function(x, na.rm=FALSE) .Call(Cgmean, x, na.rm)
 gweighted.mean = function(x, w, ..., na.rm=FALSE) {
@@ -3390,11 +3390,14 @@ is_constantish = function(q, check_singleton=FALSE) {
   length(q) == 3L &&
     is_constantish(q[[3L]], check_singleton = TRUE)
 }
-# first(x, n) / last(x, n) with n>1, #4446 #4239.
+# first(x, n=, na.rm=) / last(x, n=, na.rm=), #4446 #4239. Called only when length(q)>=3
+# (the length(q)==2L case, i.e. first(x)/last(x), is already handled earlier in .gforce_ok).
 .gfirstlast_ok = function(q, envir) {
-  length(q) == 3L &&
-    is_constantish(q[[3L]], check_singleton = TRUE) &&
-    is.numeric(n <- eval(q[[3L]], envir)) && length(n)==1L && !is.na(n) && n>=1
+  q = match.call(first, q)  # first's signature is the same as last's: first(x, n=1L, na.rm=FALSE, ...)
+  (is.null(q[["n"]]) || (is_constantish(q[["n"]], check_singleton=TRUE) &&
+    is.numeric(n <- eval(q[["n"]], envir)) && length(n)==1L && !is.na(n) && n>=1)) &&
+  (is.null(q[["na.rm"]]) || (is_constantish(q[["na.rm"]], check_singleton=TRUE) &&
+    isTRUEorFALSE(eval(q[["na.rm"]], envir))))
 }
 `.g[_ok` = function(q, x, envir=parent.frame(3L)) {
   length(q) == 3L &&
@@ -3449,7 +3452,7 @@ is_constantish = function(q, check_singleton=FALSE) {
   if (!is.null(q1)) {
     q2 = .unwrap_conversions(q[[2L]])
     if (!is.symbol(q2) || (!q2 %chin% names(x) && q2 != ".I")) return(FALSE)
-    if (length(q)==2L || (.arg_is_narm(q) && is_constantish(q[[3L]]) &&
+    if (length(q)==2L || (length(q)==3L && .arg_is_narm(q) && is_constantish(q[[3L]]) &&
         !(is.symbol(q[[3L]]) && q[[3L]] %chin% names(x)))) return(TRUE)
     return(switch(as.character(q1),
       "shift" = .gshift_ok(q),

R/last.R

@@ -2,15 +2,20 @@
 # for xts class objects it will dispatch to xts::last
 # reworked to avoid loading xts namespace (#3857) then again to fix dispatching of xts class (#4053)
 # nocov start. Tests 19.* in other.Rraw, not in the main suite.
-last = function(x, n=1L, ...) {
+last = function(x, n=1L, na.rm=FALSE, ...) {
   verbose = isTRUE(getOption("datatable.verbose", FALSE))
   if (!inherits(x, "xts")) {
     if (nargs()>1L) {
       if ("package:xts" %chin% search()) {
         if (verbose)
           catf("%s: using %s: %s\n", "last", "xts::last", "!is.xts(x) & nargs>1 & 'package:xts'%in%search()")
-        xts::last(x, n=n, ...)
+        xts::last(x, n=n, na.rm=na.rm, ...)
+      } else if (is.null(dim(x)) && !is.data.frame(x)) {
+        if (verbose)
+          catf("%s: using %s: %s\n", "last", "'.firstlast'", "!is.xts(x) & nargs>1 & is.null(dim(x)) & !is.data.frame(x)")
+        .firstlast(x, n=n, first=FALSE, na.rm=na.rm)
       } else {
+        if (!isFALSE(na.rm)) stopf("na.rm=TRUE is not currently supported for '%s'", class(x)[1L])
         # nocov start
         if (verbose)
           catf("%s: using %s: %s\n", "last", "utils::tail", "!is.xts(x) & nargs>1 & !'package:xts'%in%search()")
@@ -39,19 +44,24 @@ last = function(x, n=1L, ...) {
       stopf("'xts' class passed to %s function but 'xts' is not available, you should have 'xts' installed already", "data.table::last") # nocov
     if (verbose)
       catf("%s: using %s: %s\n", "last", "xts::last", "is.xts(x)")
-    xts::last(x, n=n, ...)
+    xts::last(x, n=n, na.rm=na.rm, ...)
   }
 }
 
-first = function(x, n=1L, ...) {
+first = function(x, n=1L, na.rm=FALSE, ...) {
   verbose = isTRUE(getOption("datatable.verbose", FALSE))
   if (!inherits(x, "xts")) {
     if (nargs()>1L) {
       if ("package:xts" %chin% search()) {
         if (verbose)
           catf("%s: using %s: %s\n", "first", "xts::first", "!is.xts(x) & nargs>1 & 'package:xts'%in%search()")
-        xts::first(x, n=n, ...)
+        xts::first(x, n=n, na.rm=na.rm, ...)
+      } else if (is.null(dim(x)) && !is.data.frame(x)) {
+        if (verbose)
+          catf("%s: using %s: %s\n", "first", "'.firstlast'", "!is.xts(x) & nargs>1 & is.null(dim(x)) & !is.data.frame(x)")
+        .firstlast(x, n=n, first=TRUE, na.rm=na.rm)
       } else {
+        if (!isFALSE(na.rm)) stopf("na.rm=TRUE is not currently supported for '%s'", class(x)[1L])
         # nocov start
         if (verbose)
           catf("%s: using %s: %s\n", "first", "utils::head", "!is.xts(x) & nargs>1 & !'package:xts'%in%search()")
@@ -80,7 +90,19 @@ first = function(x, n=1L, ...) {
       stopf("'xts' class passed to %s function but 'xts' is not available, you should have 'xts' installed already", "data.table::first") # nocov
     if (verbose)
       catf("%s: using %s: %s\n", "first", "xts::first", "is.xts(x)")
-    xts::first(x, n=n, ...)
+    xts::first(x, n=n, na.rm=na.rm, ...)
   }
 }
 # nocov end
+
+.firstlast = function(x, n, first, na.rm) {
+  if (!isTRUEorFALSE(na.rm)) stopf("'%s' must be TRUE or FALSE", "na.rm")
+  if (!na.rm) return(if (first) utils::head(x, n=n) else utils::tail(x, n=n))
+  if (!length(x)) return(x)
+  # matches 'missing' used for GForce first()/last()
+  # for a list, an element is missing when it is NULL or a length-1 logical NA, not (only) when is.na()
+  isna = if (is.list(x)) vapply(x, function(el) is.null(el) || (is.logical(el) && length(el)==1L && is.na(el)), FALSE) else is.na(x)
+  nna = which(!isna)
+  if (!length(nna)) return(if (n==1L) x[NA_integer_] else x[0L])
+  x[if (first) utils::head(nna, n) else utils::tail(nna, n)]
+}

inst/tests/optimize.Rraw

@@ -522,4 +522,55 @@ test(2285.01, optimize=opt, DT5[, .(shift(x), mean(x)), by=g, verbose=TRUE],
 test(2285.02, copy(DT5)[, c("s","m") := .(shift(x), mean(x)), by=g],
               data.table(g=c(1,1,1,2,2), x=1:5, y=1:5, s=c(NA,1L,2L,NA,4L), m=c(2,2,2,4.5,4.5)))
 test(2285.03, optimize=opt, DT5[, .(head(x,2), head(y,3)), by=g], error="Supplied 2 items for column 1 of group 1 which has 3 rows")
-# head/tail with := and mismatched n: previously errored (Supplied N items...), now aligns/pads instead, see tests 2233.28/2233.29 in tests.Rraw
+
+# first()/last() gain na.rm=, #4239 #4446. GForce optimized (gfirst()/glast()) for both the simple
+# (n=1, always exactly 1 row per group) and n>1 (data-dependent rows per group, using the same
+# gforce_dynamic mechanism as n>1 without na.rm) cases.
+opt = 0:2
+DT = data.table(g=c(1,1,1,2,2), x=c(1,NA,3,NA,5), y=11:15)
+out = c("GForce FALSE", "GForce FALSE", "GForce optimized j")
+test(2286.01, optimize=opt, DT[, first(x, na.rm=TRUE), by=g, verbose=TRUE],
+              data.table(g=c(1,2), V1=c(1,5)), output=out)
+test(2286.02, optimize=opt, DT[, .(first(x, na.rm=TRUE), mean(y)), by=g, verbose=TRUE],
+              data.table(g=c(1,2), V1=c(1,5), V2=c(12,14.5)), output=out)
+test(2286.03, optimize=opt, DT[, last(x, na.rm=TRUE), by=g, verbose=TRUE],
+              data.table(g=c(1,2), V1=c(3,5)), output=out)
+test(2286.04, optimize=opt, DT[, first(x, n=2, na.rm=TRUE), by=g, verbose=TRUE],
+              data.table(g=c(1,1,2), V1=c(1,3,5)), output=out)
+test(2286.05, optimize=opt, DT[, last(x, n=2, na.rm=TRUE), by=g],
+              data.table(g=c(1,1,2), V1=c(1,3,5)))
+test(2286.06, optimize=opt, DT[, .(first(x, n=2, na.rm=TRUE), mean(y)), by=g],
+              data.table(g=c(1,1,2), V1=c(1,3,5), V2=c(12,12,14.5)))
+test(2286.07, optimize=opt, copy(DT)[, v := first(x, na.rm=TRUE), by=g, verbose=TRUE],
+              data.table(g=c(1,1,1,2,2), x=c(1,NA,3,NA,5), y=11:15, v=c(1,1,1,5,5)), output=out)
+test(2286.08, DT[, v := first(x, n=2, na.rm=TRUE), by=g],
+              error="Supplied 3 items to be assigned to 5 items of column 'v'.")
+
+# all-NA group returns NA, like gmedian/gvar do for insufficient data, not an empty/dropped group
+DT = data.table(g=c(1,1,2,2), x=c(NA,NA,3,4))
+test(2286.11, optimize=opt, DT[, first(x, na.rm=TRUE), by=g],
+              data.table(g=c(1,2), V1=c(NA,3)))
+test(2286.12, optimize=opt, DT[, first(x, n=2, na.rm=TRUE), by=g],
+              data.table(g=2, V1=c(3,4)))
+DT = data.table(g=c(1,1,1,1), x=c(1,2,3,4), y=c(NA,6,7,8))
+test(2286.13, optimize=opt, DT[, .(first(x, n=2), first(y, n=3, na.rm=TRUE)), by=g],
+              error="Supplied 2 items for column 1 of group 1 which has 3 rows")
+DT = data.table(g=c(1,1,1,1), x=c(1,NA,3,4), y=c(NA,NA,NA,9))
+test(2286.14, optimize=opt, DT[, .(first(x, n=2, na.rm=TRUE), first(y, n=2, na.rm=TRUE)), by=g],
+              data.table(g=c(1,1), V1=c(1,3), V2=c(9,9)))
+DT = data.table(g=c(1,1,1,2,2), s=c("a",NA,"c",NA,"e"), l=list(1,NA,3,NA,5),
+                  lg=c(TRUE,NA,FALSE,NA,TRUE), cx=c(1+1i,NA,3+3i,NA,5+5i))
+test(2286.21, DT[, .(first(s,na.rm=TRUE)), by=g], data.table(g=c(1,2), V1=c("a","e")))
+test(2286.22, DT[, .(first(l,na.rm=TRUE)), by=g], data.table(g=c(1,2), V1=list(1,5)))
+test(2286.23, DT[, .(first(lg,na.rm=TRUE)), by=g], data.table(g=c(1,2), V1=c(TRUE,TRUE)))
+test(2286.24, DT[, .(first(cx,na.rm=TRUE)), by=g], data.table(g=c(1,2), V1=c(1+1i,5+5i)))
+if (test_bit64) {
+  DT[, i64 := as.integer64(c(1,NA,3,NA,5))]
+  test(2286.25, DT[, .(first(i64,na.rm=TRUE)), by=g], data.table(g=c(1,2), V1=bit64::as.integer64(c(1,5))))
+}
+test(2286.31, first(c(1,NA,3,NA,5), na.rm=TRUE), 1)
+test(2286.32, first(c(1,NA,3,NA,5), n=2, na.rm=TRUE), c(1,3))
+test(2286.33, last(c(1,NA,3,NA,5), n=2, na.rm=TRUE), c(3,5))
+test(2286.34, first(c(NA_real_,NA_real_), na.rm=TRUE), NA_real_)
+test(2286.35, first(list(1,NULL,3,NA,5), n=3, na.rm=TRUE), list(1,3,5))
+test(2286.36, first(x <- data.frame(a=1:2), na.rm=TRUE), error="na.rm=TRUE is not currently supported for 'data.frame'")

man/last.Rd

@@ -8,13 +8,18 @@ or data.table. The main difference to head/tail is that the default for \code{n}
 rather than 6.
 }
 \usage{
-first(x, n=1L, \dots)
-last(x, n=1L, \dots)
+first(x, n=1L, na.rm=FALSE, \dots)
+last(x, n=1L, na.rm=FALSE, \dots)
 }
 \arguments{
 \item{x}{ A vector, list, data.frame or data.table. Otherwise the S3 method
 of \code{xts::first} is deployed. }
 \item{n}{ A numeric vector length 1. How many items to select. }
+\item{na.rm}{ \code{TRUE} or \code{FALSE} (default). When \code{TRUE}, missing values in \code{x}
+(\code{NA}, or \code{NULL} list elements) are skipped, so that the first/last \code{n}
+\emph{non-missing} items are returned; fewer than \code{n} are returned if fewer than \code{n}
+non-missing items are found. Only supported when \code{x} is a plain vector or list (not
+\code{data.frame}, \code{data.table}, matrix/array, or \code{xts}). }
 \item{\dots}{ Not applicable for \code{data.table} first/last. Any arguments here
 are passed through to \code{xts}'s first/last. }
 }
@@ -37,5 +42,12 @@ first(x) # same as head(x, 1)
 last(1:5) # [1] 5
 x = data.table(x=1:5, y=6:10)
 last(x) # same as tail(x, 1)
+
+first(c(NA, 2, NA, 4), na.rm=TRUE) # [1] 2
+last(c(1, NA, 3, NA), n=2, na.rm=TRUE) # [1] 1 3
+
+DT = data.table(g=c(1,1,1,2,2), x=c(1,NA,3,NA,5))
+DT[, first(x, na.rm=TRUE), by=g] # first non-missing x per group
+DT[, first(x, n=2, na.rm=TRUE), by=g] # up to 2 non-missing x per group; fewer rows if fewer are non-missing
 }
 \keyword{ data }

src/data.table.h

@@ -453,8 +453,8 @@ SEXP rleid(SEXP, SEXP);
 SEXP gmedian(SEXP, SEXP);
 SEXP gtail(SEXP, SEXP);
 SEXP ghead(SEXP, SEXP);
-SEXP glast(SEXP, SEXP);
-SEXP gfirst(SEXP, SEXP);
+SEXP glast(SEXP, SEXP, SEXP);
+SEXP gfirst(SEXP, SEXP, SEXP);
 SEXP gnthvalue(SEXP, SEXP);
 SEXP dim(SEXP);
 SEXP warn_matrix_column_r(SEXP);