last/first get argument na.rm

NEWS.md

@@ -50,6 +50,8 @@
 
 14. `nafill()` and `setnafill()` gain a `limit` argument to restrict the maximum number of consecutive `NA` values filled, [#7677](https://github.com/Rdatatable/data.table/issues/7677). Thanks to @jaynewton for the suggestion and @venom1204 and @ben-schwen for the PR.
 
+15. `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.
+
 ### 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

@@ -2093,30 +2093,21 @@ replace_dot_alias = function(e) {
       }
     }
 
-    # adding ghead/gtail(n) support for n > 1 #5060 #523
-    q3 = 0L
-    if (!is.symbol(jsub)) {
-      headTail_arg = function(q) {
-        if (length(q)==3L && length(q3 <- q[[3L]])==1L && is.numeric(q3) &&
-         (q[[1L]]) %chin% c("ghead", "gtail") && q3!=1L) q3
-        else 0L
-      }
-      if (jsub %iscall% "list"){
-        q3 = max(sapply(jsub, headTail_arg))
-      } else if (length(jsub)==3L) {
-        q3 = headTail_arg(jsub)
-      }
-    }
-    if (q3 > 0L) {
-      grplens = pmin.int(q3, len__)
-      g = lapply(g, rep.int, times=grplens)
-    } else if (.is_nrows(jsub)) {
+    if (.is_nrows(jsub)) {
+      # multi-n shift() is not marked gforce_dynamic (its per-item results
+      # need unpacking into separate columns below), so it needs this branch
+      # specifically.
       g = lapply(g, rep.int, times=len__)
       # unpack list of lists for nrows functions
       zip_items = function(ll) do.call(mapply, c(list(FUN = c), ll, SIMPLIFY=FALSE, USE.NAMES=FALSE))
       if (all(vapply_1b(ans, is.list))) {
         ans = lapply(ans, zip_items)
       }
+    } else if (!is.null(dynlens <- attr(ans, "gforce_dynamic", exact=TRUE))) {
+      # ghead/gtail/gfirst/glast with n>1 and shift() with a single n:
+      # gforce() has already validated 'ans' is consistently 'dynlens' items
+      # per group in gsumm.c
+      g = lapply(g, rep.int, times=dynlens)
     }
     ans = c(g, ans)
   } else {
@@ -3346,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) .Call(Cgfirst, x)
-glast = function(x) .Call(Cglast, x)
+gfirst = function(x, n=1L) .Call(Cgfirst, x, as.integer(n))
+glast = function(x, n=1L) .Call(Cglast, x, as.integer(n))
 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) {
@@ -3399,6 +3390,12 @@ 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.
+.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
+}
 `.g[_ok` = function(q, x, envir=parent.frame(3L)) {
   length(q) == 3L &&
     is_constantish(q[[3L]], check_singleton = TRUE) &&
@@ -3458,6 +3455,7 @@ is_constantish = function(q, check_singleton=FALSE) {
       "shift" = .gshift_ok(q),
       "weighted.mean" = .gweighted.mean_ok(q, x),
       "tail" = , "head" = .ghead_ok(q),
+      "first" = , "last" = .gfirstlast_ok(q, envir),
       "[[" = , "[" = `.g[_ok`(q, x, envir),
       FALSE
     ))

inst/tests/optimize.Rraw

@@ -110,7 +110,7 @@ test(1184.2,optimize=1L, DT[, mean(v), by=x, verbose=TRUE], output="(GForce FALS
 test(1185.1,optimize=c(0L, 1L, 2L), DT[, list(sum(y), sum(v), sum(y,na.rm=TRUE), sum(v,na.rm=TRUE)), by=x],
            data.table(x=c("a","b","c","d"), V1=c(NA,10L,NA,NA), V2=c(6,NA,NA,NA), V3=c(4L,10L,7L,0L), V4=c(6,10,15,0)))
 # test 1185.1 subsumes 1185.1 and 1187.1 for testing all levels
-test(1185.2,optimize=c(0L,1L,Inf), DT[, list(mean(y), mean(v), mean(y,na.rm=TRUE), mean(v,na.rm=TRUE)), by=x, verbose=TRUE], 
+test(1185.2,optimize=c(0L,1L,Inf), DT[, list(mean(y), mean(v), mean(y,na.rm=TRUE), mean(v,na.rm=TRUE)), by=x, verbose=TRUE],
   output=c("All optimizations.*off", "Old mean.*changed j", "GForce optimized j to"))
 # test 1185.2 subsumes 1185.2, 1185.3, and 1185.4 for testing all levels
 MyVar = TRUE
@@ -334,7 +334,7 @@ test(1594.08,optimize=opt, DT[(idx), var(z1, na.rm=TRUE), by=x], ans)
 test(1594.09,optimize=opt,DT[, lapply(.SD, sd, na.rm=FALSE), by=x])
 test(1594.10,optimize=opt, DT[, lapply(.SD, sd, na.rm=TRUE), by=x], DT[, lapply(.SD, stats::sd, na.rm=TRUE), by=x])
 test(1594.11,optimize=opt, DT[, lapply(.SD, sd, na.rm=TRUE), by=x, verbose=TRUE], output=out)
- 
+
 test(1594.12,optimize=opt, DT[, lapply(.SD, prod, na.rm=FALSE), by=x])
 test(1594.13,optimize=opt, DT[, lapply(.SD, prod, na.rm=TRUE), by=x])
 test(1594.14,optimize=opt, DT[, lapply(.SD, prod, na.rm=TRUE), by=x, verbose=TRUE], output=out)
@@ -490,3 +490,36 @@ test(2283 + 0.19, optimize=opts,
        names(M[, c(mpg[1], list(mpg, b=hp), c(lapply(.SD, mean))), by="cyl", .SDcols=c("vs", "am")]),
        c("cyl", "V1", "V2", "b", "vs", "am"),
        context=sprintf("optimize=%s [XIX]", format(opt)))
+
+# first()/last() with n>1 are now GForce optimized via gfirst()/glast(), #4239
+opt = 0:2
+DT = data.table(g=c(1,1,1,2,2), x=1:5, y=11:15)
+out = c("GForce FALSE", "GForce FALSE", "GForce optimized j to")
+test(2284.01, optimize=opt, DT[, first(x, n=2), by=g, verbose=TRUE],
+              data.table(g=c(1,1,2,2), V1=c(1L,2L,4L,5L)), output=out)
+test(2284.02, optimize=opt, DT[, .(first(x,n=2), mean(y)), by=g, verbose=TRUE],
+              data.table(g=c(1,1,2,2), V1=c(1L,2L,4L,5L), V2=c(12,12,14.5,14.5)), output=out)
+test(2284.03, optimize=opt, DT[, .(first(x,n=2), first(y,n=3)), by=g], error="Supplied 2 items for column 1 of group 1 which has 3 rows")
+DT = data.table(a=1:3,b=(1:9)/10)
+test(2284.04, optimize=opt, copy(DT)[, v := first(b, n=3), a, verbose=TRUE], copy(DT)[, v := b], output=out)
+test(2284.05, DT[, v := first(b, n=2), a], error="Supplied 6 items to be assigned to 9 items of column 'v'.")
+DT = data.table(g=c(1,1,1,2,2), s=c("a","b","c","d","e"), l=list(1,2,3,4,5))
+test(2284.06, DT[, .(first(s,n=2), last(l,n=2)), by=g],
+              data.table(g=c(1,1,2,2), V1=c("a","b","d","e"), V2=list(2,3,4,5)))
+DT4 = data.table(g=c(1,1,1,2,2), lg=c(TRUE,FALSE,TRUE,FALSE,TRUE), db=c(1.5,2.5,3.5,4.5,5.5), cx=complex(1:5, 5:1))
+test(2284.07, DT4[, .(first(lg,n=2), first(db,n=2), first(cx,n=2)), by=g],
+              data.table(g=c(1,1,2,2), V1=c(TRUE,FALSE,FALSE,TRUE), V2=c(1.5,2.5,4.5,5.5),
+                         V3=complex(c(1,2,4,5), c(5,4,2,1))))
+# fall back for n<1L
+DT = data.table(g=c(1,1,1,2,2), x=1:5, y=11:15)
+test(2284.08, DT[, first(x, n=0L), by=g, verbose=TRUE],
+              data.table(g=numeric(), V1=integer()), output="GForce is on, but not activated")
+
+# ghead()/gtail() and gshift() now also go through gforce_dynamic too
+DT5 = data.table(g=c(1,1,1,2,2), x=1:5, y=1:5)
+test(2285.01, optimize=opt, DT5[, .(shift(x), mean(x)), by=g, verbose=TRUE],
+              data.table(g=c(1,1,1,2,2), V1=c(NA,1L,2L,NA,4L), V2=c(2,2,2,4.5,4.5)), output=out)
+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

inst/tests/tests.Rraw

@@ -17857,8 +17857,8 @@ test(2233.25, copy(DT)[a!=4, v := head(b, 3L), a, verbose=TRUE], copy(DT)[a!=4,
 DT = data.table(a=1:3,b=(1:9)/10)
 test(2233.26, DT[, c("v1","v2") := .(min(b), max(b)), a,         verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v1=(1:3)/10, v2=(7:9)/10), output="GForce optimized j to")
 test(2233.27, DT[, c("v1","v2") := .(head(b,3L), tail(b,3L)), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v1=(1:9)/10, v2=(1:9)/10), output="GForce optimized j to")
-test(2233.28, DT[, c("v1","v2") := .(head(b,3L), tail(b,2L)), a], error="Supplied 6 items to be assigned to 9 items of column 'v2'.")
-test(2233.29, DT[, c("v1","v2") := .(head(b,2L), tail(b,3L)), a], error="Supplied 6 items to be assigned to 9 items of column 'v1'.")
+test(2233.28, DT[, c("v1","v2") := .(head(b,3L), tail(b,2L)), a], error="Supplied 2 items for column 2 of group 1 which has 3 rows")
+test(2233.29, DT[, c("v1","v2") := .(head(b,2L), tail(b,3L)), a], error="Supplied 2 items for column 1 of group 1 which has 3 rows")
 test(2233.30, DT[, c("v1","v2") := .(head(b,2L), tail(b,2L)), a], error="Supplied 6 items to be assigned to 9 items of column 'v1'.")
 test(2233.31, DT[, c("v1","v2") := .(min(b), max(b)), a,         verbose=TRUE], DT[, c("v1","v2") := .(base::min(b), base::max(b)), a          ], output="GForce optimized j to")
 test(2233.32, DT[, c("v1","v2") := .(head(b,3L), tail(b,3L)), a, verbose=TRUE], DT[, c("v1","v2") := .(utils::head(b,3L), utils::tail(b,3L)), a], output="GForce optimized j to")

src/data.table.h

@@ -137,6 +137,7 @@ extern SEXP sym_index;
 extern SEXP sym_BY;
 extern SEXP sym_starts, char_starts;
 extern SEXP sym_maxgrpn;
+extern SEXP sym_gforce_dynamic;
 extern SEXP sym_anyna;
 extern SEXP sym_anyinfnan;
 extern SEXP sym_anynotascii;
@@ -454,8 +455,8 @@ SEXP rleid(SEXP, SEXP);
 SEXP gmedian(SEXP, SEXP);
 SEXP gtail(SEXP, SEXP);
 SEXP ghead(SEXP, SEXP);
-SEXP glast(SEXP);
-SEXP gfirst(SEXP);
+SEXP glast(SEXP, SEXP);
+SEXP gfirst(SEXP, SEXP);
 SEXP gnthvalue(SEXP, SEXP);
 SEXP dim(SEXP);
 SEXP warn_matrix_column_r(SEXP);

src/gsumm.c

@@ -37,6 +37,59 @@ static int nbit(int n)
   return nb;
 }
 
+/* Some GForce functions (currently gfirst/glast/ghead/gtail with n>1, and gshift) return, for at least
+   one group, a different number of items than 1 (the norm for e.g. gmean, gsum). Such a column carries
+   a 'gforce_dynamic' attribute (set by gfirstlast and gshift, below): a scalar integer w, meaning the
+   column has MIN(w, grpsize[g]) items for group g (gshift sets w=INT_MAX, since it always returns
+   exactly grpsize[g]). When a query combines such a column with an ordinary (fixed, 1-per-group) result
+   in the same by= (e.g. .(shift(x), mean(y))), this function replicates the fixed result out to match,
+   #1414. When a query combines two dynamic columns that don't actually agree in length for some group
+   (e.g. .(head(x,2), head(y,3)) where the group has >=3 rows), this errors rather than reconciling them
+   with NA. Otherwise optimize>=2 (GForce) would succeed at a query that optimize<2 (dogroups.c) can't
+   do at all, making the optimize level change correctness rather than just speed Returns gans unchanged
+   (no allocation) when no column carries the attribute, so the (very common) case of no dynamic-length
+   GForce result in the query costs nothing extra. */
+static SEXP gforce_align_dynamic(SEXP gans) {
+  const int nans = length(gans);
+  int max_w = 0;
+  for (int i=0; i<nans; ++i) {
+    SEXP att = getAttrib(VECTOR_ELT(gans, i), sym_gforce_dynamic);
+    if (isNull(att)) continue;
+    const int this_w = INTEGER(att)[0];
+    if (this_w>max_w) max_w=this_w;
+  }
+  if (!max_w) return gans;  // nothing in this j is gforce_dynamic; most common case, return untouched with no allocation
+  int nprotect = 0;
+  SEXP lens = PROTECT(allocVector(INTSXP, ngrp)); nprotect++;
+  int *lensp = INTEGER(lens);
+  for (int g=0; g<ngrp; ++g) lensp[g] = MIN(grpsize[g], max_w);
+
+  SEXP ans = PROTECT(allocVector(VECSXP, nans)); nprotect++;
+  for (int i=0; i<nans; ++i) {
+    SEXP tt = VECTOR_ELT(gans, i);
+    SEXP att = getAttrib(tt, sym_gforce_dynamic);
+    if (isNull(att)) {
+      // an ordinary (1-per-group) result combined with a dynamic one; rep it out to match
+      SET_VECTOR_ELT(ans, i, eval(PROTECT(lang3(install("rep.int"), tt, lens)), R_GlobalEnv));
+      UNPROTECT(1);
+      continue;
+    }
+    const int w = INTEGER(att)[0];
+    for (int g=0; g<ngrp; ++g) {
+      const int thislen = MIN(grpsize[g], w);
+      if (thislen != lensp[g])
+        error(_("Supplied %d items for column %d of group %d which has %d rows. The RHS length must either be 1 (single values are ok) or match the LHS length exactly. If you wish to 'recycle' the RHS please use rep() explicitly to make this intent clear to readers of your code."), thislen, i+1, g+1, lensp[g]);
+    }
+    // this column's own length, MIN(grpsize[g],w) summed over g, already equals sum(lensp) (just
+    // checked above) and is already in the right group-by-group order; use it as-is, no copy needed
+    setAttrib(tt, sym_gforce_dynamic, R_NilValue);
+    SET_VECTOR_ELT(ans, i, tt);
+  }
+  setAttrib(ans, sym_gforce_dynamic, lens);  // expose the final per-group lengths so [.data.table can align/replicate the group columns to match
+  UNPROTECT(nprotect);
+  return ans;
+}
+
 /*
   Functions with GForce optimization are internally parallelized to speed up
     grouped summaries over a large data.table. OpenMP is used here to
@@ -205,16 +258,17 @@ SEXP gforce(SEXP env, SEXP jsub, SEXP o, SEXP f, SEXP l, SEXP irowsArg) {
   oo = INTEGER(o);
   ff = INTEGER(f);
 
-  SEXP ans = PROTECT( eval(jsub, env) );
+  int nprotect = 0;
+  SEXP ans = PROTECT( eval(jsub, env) ); nprotect++;
   if (verbose) { Rprintf(_("gforce eval took %.3f\n"), wallclock()-started); started=wallclock(); }
   // if this eval() fails with R error, R will release grp for us. Which is why we use R_alloc above.
   if (isVectorAtomic(ans)) {
-    SEXP tt = PROTECT(allocVector(VECSXP, 1));
+    SEXP tt = PROTECT(allocVector(VECSXP, 1)); nprotect++;
     SET_VECTOR_ELT(tt, 0, ans);
-    UNPROTECT(2);
-    return tt;
+    ans = tt;
   }
-  UNPROTECT(1);
+  ans = PROTECT(gforce_align_dynamic(ans)); nprotect++;  // no-op (no extra allocation) unless a gfirst()/glast() with n>1 is present, #4446 #4239
+  UNPROTECT(nprotect);
   return ans;
 }
 
@@ -922,7 +976,10 @@ SEXP gmedian(SEXP x, SEXP narmArg) {
 
 static SEXP gfirstlast(SEXP x, const bool first, const int w, const bool headw) {
   // w: which item (1 other than for gnthvalue when could be >1)
-  // headw: select 1:w of each group when first=true, and (n-w+1):n when first=false (i.e. tail)
+  // headw: select 1:w of each group when first=true, and (n-w+1):n when first=false (i.e. tail).
+  //        When TRUE, the result is marked with a 'gforce_dynamic' attribute #4239
+  //        so that gforce() can correctly replicate ordinary (fixed, 1-per-group) results
+  //        against it, and validate it against other dynamic results, when combined in the same by=.
   const bool nosubset = irowslen == -1;
   const bool issorted = !isunsorted; // make a const-bool for use inside loops
   const int n = nosubset ? length(x) : irowslen;
@@ -988,29 +1045,34 @@ static SEXP gfirstlast(SEXP x, const bool first, const int w, const bool headw)
   default:
     error(_("Type '%s' is not supported by GForce head/tail/first/last/`[`. Either add the namespace prefix (e.g. utils::head(.)) or turn off GForce optimization using options(datatable.optimize=1)"), type2char(TYPEOF(x)));
   }
+  if (headw) setAttrib(ans, sym_gforce_dynamic, ScalarInteger(w));  // so gforce() can recompute MIN(w, grpsize[g]) per group
   copyMostAttrib(x, ans);
   UNPROTECT(1);
   return(ans);
 }
 
-SEXP glast(SEXP x) {
-  return gfirstlast(x, false, 1, false);
+SEXP glast(SEXP x, SEXP nArg) {
+  if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "glast is only implemented for n>0. This should have been caught before"); // # nocov
+  const int n=INTEGER(nArg)[0];
+  return gfirstlast(x, false, n, n>1);
 }
 
-SEXP gfirst(SEXP x) {
-  return gfirstlast(x, true, 1, false);
+SEXP gfirst(SEXP x, SEXP nArg) {
+  if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "gfirst is only implemented for n>0. This should have been caught before"); // # nocov
+  const int n=INTEGER(nArg)[0];
+  return gfirstlast(x, true, n, n>1);
 }
 
 SEXP gtail(SEXP x, SEXP nArg) {
   if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "gtail is only implemented for n>0. This should have been caught before"); // # nocov
   const int n=INTEGER(nArg)[0];
-  return n==1 ? glast(x) : gfirstlast(x, false, n, true);
+  return gfirstlast(x, false, n, n>1);
 }
 
 SEXP ghead(SEXP x, SEXP nArg) {
-  if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "gtail is only implemented for n>0. This should have been caught before"); // # nocov
+  if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "ghead is only implemented for n>0. This should have been caught before"); // # nocov
   const int n=INTEGER(nArg)[0];
-  return n==1 ? gfirst(x) : gfirstlast(x, true, n, true);
+  return gfirstlast(x, true, n, n>1);
 }
 
 SEXP gnthvalue(SEXP x, SEXP nArg) {
@@ -1279,7 +1341,14 @@ SEXP gshift(SEXP x, SEXP nArg, SEXP fillArg, SEXP typeArg) {
     }
     copyMostAttrib(x, tmp); // needed for integer64 because without not the correct class of int64 is assigned
   }
-  UNPROTECT(nprotect);
   // consistency with plain shift(): "strip" the list in the 1-input case, for convenience
-  return isVectorAtomic(x) && length(ans) == 1 ? VECTOR_ELT(ans, 0) : ans;
+  if (isVectorAtomic(x) && length(ans) == 1) {
+    SEXP tmp = VECTOR_ELT(ans, 0);
+    // mark single gshift with gforce_dynamic to use MIN(w, grpsize[g])
+    setAttrib(tmp, sym_gforce_dynamic, ScalarInteger(INT_MAX));
+    UNPROTECT(nprotect);
+    return tmp;
+  }
+  UNPROTECT(nprotect);
+  return ans;
 }

src/init.c

@@ -29,6 +29,7 @@ SEXP sym_index;
 SEXP sym_BY;
 SEXP sym_starts, char_starts;
 SEXP sym_maxgrpn;
+SEXP sym_gforce_dynamic;
 SEXP sym_anyna;
 SEXP sym_anyinfnan;
 SEXP sym_anynotascii;
@@ -336,6 +337,7 @@ void attribute_visible R_init_data_table(DllInfo *info)
   sym_index = install("index");
   sym_BY = install(".BY");
   sym_maxgrpn = install("maxgrpn");
+  sym_gforce_dynamic = install("gforce_dynamic");
   sym_anyna = install("anyna");
   sym_anyinfnan = install("anyinfnan");
   sym_anynotascii = install("anynotascii");