first/last return 'true vectors' so dogroups.c knows not to recycle length-1 and pad with NA

Author: mattdowle

R/last.R

@@ -20,60 +20,64 @@ last = function(x, n=1L, na.rm=FALSE, ...) {
   stopifnot(isTRUEorFALSE(na.rm) || identical(na.rm,"row"))
   stopifnot(is.numeric(n), length(n)==1L, n>=0L)
   n = as.integer(n)
-  .headtail = if (first) utils::head else utils::tail
-  if (isFALSE(na.rm) || n==0L)
-    return(.headtail(x, n=n, ...))
   if (is.data.frame(x)) {
     if (!nrow(x)) return(x)
     if (identical(na.rm, "row")) {   # any NA on the row removes that row
       nna = which_(.Call(Cdt_na, x, seq_along(x)), bool=FALSE)
-      # from na.omit.data.table without calling na.omit which would subset all non-NA rows
-      # TODO: n and first/last could be passed to Cdt_na and it could stop after finding n
-      nna = .headtail(nna, n=n)
-      if (length(nna) < min(n,nrow(x))) {
-        # to match optimized na.rm=TRUE behavior; e.g. when .SD is one column
-        # TODO: extra argument all.na=NA|NULL (or pad.na=) could control this
-        pad = rep.int(NA, min(n,nrow(x))-length(nna))
-        # returning min(n,nrow(x)) is what optimized one-column does because GForce needs to be deterministic by group
-        # currently; i.e. number of items per group doesn't depend on how many NA there are
-        nna = if (first) c(nna, pad) else c(pad, nna)
-      }
-      ans = x[nna,,drop=FALSE]
-      # DT[NA] returns NULL for list columns. TODO: change [.data.table to return NA for list columns
-      # In the meantime, fix up the NULLs here in first/last
-      for (col in which(vapply_1b(ans, is.list))) {
-        for (i in which(vapply_1b(ans[[col]], is.null))) {
-          set(ans, i, col, NA)
+      # very similar to na.omit.data.table
+      # TODO: n and first/last could be passed to Cdt_na and it could stop after finding n (it already does that in gsumm.c when gforce optimized)
+      nna = .firstlastVector(nna, n=n, first=first, na.rm=FALSE)
+      ans = .Call(CsubsetDT, x, nna, seq_along(x))  # works on DF too
+    } else {
+      ans = lapply(x, .firstlastVector, n=n, first=first, na.rm=na.rm)
+      if (na.rm) {
+        l = vapply_1i(ans, length)
+        m = max(l)
+        for (i in which(l<m)) {
+          ans[[i]] = c(ans[[i]], rep(NA, m-l[i]))
         }
+        # any row.names won't align to the values now in the result so don't retain them
       }
-      return(ans)
-    }
-    # else na.rm==TRUE; select the first/last non-NA within each column
-    ans = lapply(x, .narmVector, n=n, first=first)
-    l = vapply_1i(ans, length)
-    m = max(l)
-    for (i in which(l<m)) {  # pad with NA
-      ans[[i]] = c(ans[[i]], rep(NA, m-l[i]))
     }
     if (is.data.table(x)) setDT(ans) else setDF(ans)
     setattr(ans, "class", class(x))
+    if (!isTRUE(na.rm) && length(rn<-attr(x,"row.names")))
+      setattr(ans, "row.names", if (isFALSE(na.rm)) .firstlastVector(rn, n=n, first=first, na.rm=FALSE)
+                                               else rn[nna])
     return(ans)
   }
   if (!length(x))
     return(x)
-  if (is.vector(x) && !isFALSE(na.rm))
-    return(.narmVector(x, n=n, first=first))
-  if (!isFALSE(na.rm))
-    stopf("na.rm=TRUE|'row' is not currently supported for '%s'", class(x)[1L])
-  .headtail(x, n=n, ...)  
-  # TODO when n=1, return(x[length(x)]) would save method dispatch overhead
-  # TODO and previous version had lx = length(x); if (!lx) x else x[[lx]].  So empty input returned empty
+  if (!is.vector(x)) {
+    if (!isFALSE(na.rm))
+      stopf("na.rm=TRUE|'row' is not currently supported for '%s'", class(x)[1L])
+    return((if (first) utils::head else utils::tail)(x, n=n, ...))  # e.g. matrix
+  }
+  return(.firstlastVector(x, n=n, first=first, na.rm=!isFALSE(na.rm)))  # !isFALSE to convert 'row' to TRUE
 }
 
-.narmVector = function(x, n, first) {
-  nna = which_(is.na(x) | (is.list(x) & vapply_1b(x,is.null)), bool=FALSE)   # TODO: again, n and first/last could be passed to C here
-  if (!length(nna)) x[0L]
-  else if (n==1L)   x[nna[if (first) 1L else length(nna)]]
-  else              x[(if (first) utils::head else utils::tail)(nna, n)]  # TODO: avoid dispatch here and do ourselves since just a vector
+.firstlastVector = function(x, n, first, na.rm) {
+  if (!length(x)) return(x)
+  if (n==0L) return(x[0L])
+  ans = if (na.rm) {
+    nna = which_(if (is.list(x)) vapply_1b(x,function(y){is.null(y)||(length(y)==1L&&is.na(y))})
+                            else is.na(x), bool=FALSE)   # TODO: again, n and first/last could be passed to C here
+    if (!length(nna)) x[0L]
+    else {y=min(n,length(nna)); x[nna[if (first) seq.int(1L,y) else seq.int(length(nna)-y+1L,length(nna))]]}
+  } else {
+    y=min(n,length(x)); x[if (first) seq.int(1L,y) else seq.int(length(x)-y+1L,length(x))] 
+  }
+  if (n>1L || na.rm)   # n!=length(ans)
+    .Call("Csettruelength", ans, length(ans))
+  # for dogroups.c to know that shorter results (including when na.rm results in a length-1) should be padded with NA to match the length of longer items
+  # head and tail with na.rm=TRUE are by their nature returning a vector and therefore shouldn't be recycled when length-1; test 2240.81
+  # TODO: new function pad() could be provided so user can do things like DT[, .(pad(na.omit(B)), pad(na.omit(C))), by=grp]
+  #         to avoid the error 'Supplied 2 items for column 1 of group 1 which has 3 rows ...'
+  #       and/or pad= could be added to [.data.table to allow padding all results
+  # Since gforce_dynamic optimizes head/tail it knows to pad and that's optimized. However, default last(x) and first(x) (i.e. n=1 na.rm=FALSE) are
+  # single-valued like mean,median etc and are recycled in the same way. This is consistent with n=1 na.rm=FALSE already not being treated as
+  # gforce_dynamic in gsumm.c either.
+  # ***** TODO *****: n=1 na.rm=TRUE is not vector result and should be recycled to be consistent with n=1 na.rm=FALSE. It's just n>1 that is true vector then.
+  ans
 }
 

inst/tests/tests.Rraw

@@ -18865,14 +18865,14 @@ for (opt in c(0, Inf)) {
                 output=out)
   test(test_no+.22, DT[, lapply(.SD, first, na.rm=TRUE), by=grp, verbose=TRUE], ans, output=out)
   test(test_no+.23, DT[, first(.SD, na.rm='row'), by=grp, verbose=TRUE],
-                data.table(grp=1:4, A=c(NA,2L,NA,3L), B=c(NA,3,NA,pi), C=c(NA,"b",NA,"d"), D=c(NA,3i,NA,4i), E=list(NA,c("a","b"),NA,3:4)),
+                data.table(grp=c(2L,4L), A=c(2L,3L), B=c(3,pi), C=c("b","d"), D=c(3i,4i), E=list(c("a","b"),3:4)),
                 notOutput="GForce optimized")   # TODO: could try to implement gforce optimized na.rm='row' in future
   test(test_no+.24, DT[, last(.SD, na.rm=TRUE), by=grp, verbose=TRUE],
                 ans<-data.table(grp=1:4, A=c(1L,2L,NA,3L), B=c(1,3,3,pi), C=c("a","b","c","d"), D=c(1i,3i,NA,4i), E=list(NA,c("a","b"),list(1:2),3:4)),
                 output=out)
   test(test_no+.25, DT[, lapply(.SD, last, na.rm=TRUE), by=grp, verbose=TRUE], ans, output=out)
   test(test_no+.26, DT[, last(.SD, na.rm='row'), by=grp, verbose=TRUE],
-                data.table(grp=1:4, A=c(NA,2L,NA,3L), B=c(NA,3,NA,pi), C=c(NA,"b",NA,"d"), D=c(NA,3i,NA,4i), E=list(NA,c("a","b"),NA,3:4)),
+                data.table(grp=c(2L,4L), A=c(2L,3L), B=c(3,pi), C=c("b","d"), D=c(3i,4i), E=list(c("a","b"),3:4)),
                 output="GForce FALSE")
   test(test_no+.27, DT[, .(last(A,na.rm=TRUE), first(B, na.rm=TRUE), last(C)), by=grp, verbose=TRUE],
                 data.table(grp=1:4, V1=c(1L, 2L, NA, 3L), V2=c(1,pi,3,pi), V3=c(NA,NA,"c",NA)),
@@ -18886,13 +18886,13 @@ for (opt in c(0, Inf)) {
                 data.table(grp=INT(1,2,2,3,4,4), A=INT(1,2,NA,NA,3,3), B=c(1,pi,3,3,pi,NA), C=c("a","b",NA,"c","d",NA), D=c(1i,2i,3i,NA,4i,NA), E=list(NA,c("a","b"),NA,list(1:2),3:4,NA)),
                 output=out)
   test(test_no+.42, DT[, last(.SD, na.rm="row", n=2), by=grp, verbose=TRUE],
-                data.table(grp=INT(1,1,2,2,3,4,4), A=c(NA,NA,NA,2L,NA,NA,3L), B=c(NA,NA,NA,3,NA,NA,pi), C=c(NA,NA,NA,"b",NA,NA,"d"), D=c(NA,NA,NA,3i,NA,NA,4i), E=list(NA,NA,NA,c("a","b"),NA,NA,3:4)),
+                data.table(grp=INT(2,4), A=c(2L,3L), B=c(3,pi), C=c("b","d"), D=c(3i,4i), E=list(c("a","b"),3:4)),
                 output="GForce FALSE")
   test(test_no+.43, DT[, first(.SD, na.rm=TRUE, n=2), by=grp, verbose=TRUE],
                 data.table(grp=INT(1,2,2,3,4,4), A=INT(1,2,NA,NA,3,3), B=c(1,pi,3,3,pi,NA), C=c("a","b",NA,"c","d",NA), D=c(1i,2i,3i,NA,4i,NA), E=list(NA,c("a","b"),NA,list(1:2),3:4,NA)),
                 output=out)
   test(test_no+.44, DT[, first(.SD, na.rm="row", n=2), by=grp, verbose=TRUE],
-                data.table(grp=INT(1,1,2,2,3,4,4), A=c(NA,NA,2L,NA,NA,3L,NA), B=c(NA,NA,3,NA,NA,pi,NA), C=c(NA,NA,"b",NA,NA,"d",NA), D=c(NA,NA,3i,NA,NA,4i,NA), E=list(NA,NA,c("a","b"),NA,NA,3:4,NA)),
+                data.table(grp=INT(2,4), A=c(2L,3L), B=c(3,pi), C=c("b","d"), D=c(3i,4i), E=list(c("a","b"),3:4)),
                 output="GForce FALSE")
   test(test_no+.51, DT[, first(.SD, na.rm=TRUE, n=0), by=grp, verbose=TRUE], DT[0], output="GForce FALSE")
   test(test_no+.52, DT[, last(.SD, na.rm=TRUE, n=0),  by=grp, verbose=TRUE], DT[0], output="GForce FALSE")
@@ -18913,9 +18913,8 @@ for (opt in c(0, Inf)) {
   test(test_no+.74, DT[,last(A, na.rm=NA),by=grp,verbose=TRUE], error="na.rm", output=out)
 
   # aligning two gforce dynamic columns the same between optimized and unoptimized
-  #   needs to be top aligned otherwise dogroups.c would need knowledge of R first/last (last was aligned at the bottom in the PR before this test added)
-  #   but outstanding is that dogroups.c recycles length-1 whereas gforce_dynamic currently NA fills with vector output from
-  #     first/last n>1 na.rm=TRUE in mind. So this test fails when not optimized (2240.81) currently.
+  # needs to be top aligned otherwise dogroups.c would need knowledge of whether first or last was called (last was aligned at the bottom in the PR before
+  # this test added). dogroups.c now just needs to know whether the vector is a true vector to know not to recycle length-1 and to pad with NA
   test(test_no+.81, DT[, .(first(B, n=2, na.rm=TRUE), last(C, na.rm=TRUE)), by=grp, verbose=TRUE],
                          data.table(grp=INT(1,2,2,3,4), V1=c(1,pi,3,3,pi), V2=c("a","b",NA,"c","d")),
                          output=out)

src/assign.c

@@ -277,6 +277,12 @@ SEXP truelength(SEXP x) {
   return ScalarInteger(isNull(x) ? 0 : TRUELENGTH(x));
 }
 
+SEXP settruelength(SEXP x, SEXP n) {
+  // currently just for first/last and dogroups.c; see comments at the end of last.R
+  SET_TRUELENGTH(x, INTEGER(n)[0]);
+  return R_NilValue;
+}
+
 SEXP selfrefokwrapper(SEXP x, SEXP verbose) {
   return ScalarInteger(_selfrefok(x,FALSE,LOGICAL(verbose)[0]));
 }

src/dogroups.c

@@ -414,15 +414,21 @@ SEXP dogroups(SEXP dt, SEXP dtcols, SEXP groups, SEXP grpcols, SEXP jiscols, SEX
         // thislen>0
         if (TYPEOF(source) != TYPEOF(target))
           error(_("Column %d of result for group %d is type '%s' but expecting type '%s'. Column types must be consistent for each group."), j+1, i+1, type2char(TYPEOF(source)), type2char(TYPEOF(target)));
-        if (thislen>1 && thislen!=maxn && grpn>0) {  // grpn>0 for grouping empty tables; test 1986
-          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, j+1, i+1, maxn);
-        }
         bool copied = false;
         if (isNewList(target) && anySpecialStatic(source)) {  // see comments in anySpecialStatic()
           source = PROTECT(copyAsPlain(source));
           copied = true;
         }
-        memrecycle(target, R_NilValue, thisansloc, maxn, source, 0, -1, 0, "");
+        if (TRUELENGTH(source)==LENGTH(source)) {
+          // first() and last() set truelength to mark that it is a true vector; see comments at the end of last.R and test 2240.81
+          // a true vector is not recycled when length-1 and is padded with NA to match the length of the longest result
+          memrecycle(target, R_NilValue, thisansloc, thislen, source, 0, -1, 0, "");  // just using memrecycle to copy contents
+          writeNA(target, thisansloc+thislen, maxn-thislen, true);                    // pad with NA
+        } else {
+          if (thislen>1 && thislen!=maxn && grpn>0)  // grpn>0 for grouping empty tables; test 1986
+            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, j+1, i+1, maxn);
+          memrecycle(target, R_NilValue, thisansloc, maxn, source, 0, -1, 0, "");
+        }
         if (copied) UNPROTECT(1);
       }
     }

src/init.c

@@ -54,6 +54,7 @@ SEXP shallowwrapper();
 SEXP alloccolwrapper();
 SEXP selfrefokwrapper();
 SEXP truelength();
+SEXP settruelength();
 SEXP setcharvec();
 SEXP setcolorder();
 SEXP chmatch_R();
@@ -148,6 +149,7 @@ R_CallMethodDef callMethods[] = {
 {"Calloccolwrapper", (DL_FUNC) &alloccolwrapper, -1},
 {"Cselfrefokwrapper", (DL_FUNC) &selfrefokwrapper, -1},
 {"Ctruelength", (DL_FUNC) &truelength, -1},
+{"Csettruelength", (DL_FUNC) &settruelength, -1},
 {"Csetcharvec", (DL_FUNC) &setcharvec, -1},
 {"Csetcolorder", (DL_FUNC) &setcolorder, -1},
 {"Cchmatch", (DL_FUNC) &chmatch_R, -1},