catch first/last n>1 := by group with helpful error

mattdowle · mattdowle · commit beded9565d62 · 2022-08-10T23:22:04.000-06:00
diff --git a/R/data.table.R b/R/data.table.R
@@ -1864,7 +1864,8 @@ replace_dot_alias = function(e) {
     #fix for #1683
     if (use.I) assign(".I", seq_len(nrow(x)), thisEnv)
     ans = gforce(thisEnv, jsub, o__, f__, len__, irows,  # irows needed for #971
-                 .Call(CsubsetVector, groups, grpcols))  # just a list() subset to make C level neater; doesn't copy column contents
+                 .Call(CsubsetVector, groups, grpcols),  # just a list() subset to make C level neater; doesn't copy column contents
+                 lhs)  # for now this just prevents := with new feature first/last n>1; in future see TODO below
   } else {
     ans = .Call(Cdogroups, x, xcols, groups, grpcols, jiscols, xjiscols, grporder, o__, f__, len__, jsub, SDenv, cols, newnames, !missing(on), verbose)
   }
@@ -1886,7 +1887,7 @@ replace_dot_alias = function(e) {
   # Grouping by by: i is by val, icols NULL, o__ may be subset of x, f__ points to o__ (or x if !length o__)
   # TO DO: setkey could mark the key whether it is unique or not.
   if (!is.null(lhs)) {
-    if (GForce) { # GForce should work with := #1414
+    if (GForce) { # GForce should work with := #1414. TODO: move down into gforce at C level to save creating/rep'ing ans and grpcols wastefully
       vlen = length(ans[[1L]])   # TODO: this might be ngrp when na.rm=TRUE and one group has 2 and another 0, so needs enhancing here (by passing all-1 back from gans?)
       # replicate vals if GForce returns 1 value per group
       jvals = if (vlen==length(len__)) lapply(tail(ans, -length(grpcols)), rep, times=len__) else tail(ans, -length(grpcols))  # see comment in #4245 for why rep instead of rep.int
@@ -3004,7 +3005,7 @@ gshift = function(x, n=1L, fill=NA, type=c("lag", "lead", "shift", "cyclic")) {
   stopifnot(is.numeric(n))
   .Call(Cgshift, x, as.integer(n), fill, type)
 }
-gforce = function(env, jsub, o, f, l, rows, grpcols) .Call(Cgforce, env, jsub, o, f, l, rows, grpcols)
+gforce = function(env, jsub, o, f, l, rows, grpcols, lhs) .Call(Cgforce, env, jsub, o, f, l, rows, grpcols, lhs)
 
 .prepareFastSubset = function(isub, x, enclos, notjoin, verbose = FALSE){
   ## helper that decides, whether a fast binary search can be performed, if i is a call
diff --git a/inst/tests/tests.Rraw b/inst/tests/tests.Rraw
@@ -18623,60 +18623,60 @@ test(2232.4, unique(DT, by='g', cols='v3'), error="non-existing column(s)")
 # support := with GForce #1414
 options(datatable.optimize = 2L)
 DT = data.table(a=1:3,b=(1:9)/10)
-test(2233.01, DT[, v := min(b), a, verbose=TRUE],      data.table(a=1:3, b=(1:9)/10, v=(1:3)/10), output="GForce optimized j to")
+test(2233.01, DT[, v := min(b), a, verbose=TRUE],   data.table(a=1:3, b=(1:9)/10, v=(1:3)/10), output="GForce optimized j to")
 # GForce returning full length
-test(2233.02, DT[, v := head(b, 3L), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v=(1:9)/10), output="GForce optimized j to")
+test(2233.02, DT[, v := shift(b), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v=c(NA,NA,NA,(1:6)/10)), output="GForce optimized j to")
 # GForce neither returning 1 per group nor full length
-test(2233.03, DT[, v := head(b, 2L), a], error="Supplied 6 items to be assigned to 9 items of column 'v'.")
-# compare to non GForce version
+test(2233.03, DT[, v := head(b, 2L), a], error="could be implemented")
+# ensure base:: and utils:: return the same result
 DT = data.table(a=1:3,b=(1:9)/10)
-test(2233.04, copy(DT)[, v := min(b), a,      verbose=TRUE], copy(DT)[, v := base::min(b), a,     ], output="GForce optimized j to")
-test(2233.05, copy(DT)[, v := head(b, 3L), a, verbose=TRUE], copy(DT)[, v := utils::head(b, 3L), a], output="GForce optimized j to")
+test(2233.04, copy(DT)[, v := min(b), a, verbose=TRUE], copy(DT)[, v := base::min(b), a,     ], output="GForce optimized j to")
+test(2233.05, copy(DT)[, v := b[1:3], a], copy(DT)[, v := utils::head(b, 3L), a])
 
 # with key and grouping by key
 DT = data.table(a=1:3,b=(1:9)/10, key="a")
-test(2233.06, DT[, v := min(b), a, verbose=TRUE],      data.table(a=1:3, b=(1:9)/10, v=(1:3)/10, key="a"), output="GForce optimized j to")
-test(2233.07, DT[, v := head(b, 3L), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v=(1:9)/10, key="a"), output="GForce optimized j to")
-test(2233.08, DT[, v := head(b, 2L), a], error="Supplied 6 items to be assigned to 9 items of column 'v'.")
+test(2233.06, DT[, v := min(b), a, verbose=TRUE],   data.table(a=1:3, b=(1:9)/10, v=(1:3)/10, key="a"), output="GForce optimized j to")
+test(2233.07, DT[, v := shift(b), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v=c(NA,NA,NA,(1:6)/10), key="a"), output="GForce optimized j to")
+test(2233.08, DT[, v := head(b, 2L), a], error="could be implemented")
 DT = data.table(a=1:3,b=(1:9)/10, key="a")
-test(2233.09, copy(DT)[, v := min(b), a,      verbose=TRUE], copy(DT)[, v := base::min(b), a,     ], output="GForce optimized j to")
-test(2233.10, copy(DT)[, v := head(b, 3L), a, verbose=TRUE], copy(DT)[, v := utils::head(b, 3L), a], output="GForce optimized j to")
+test(2233.09, copy(DT)[, v := min(b), a, verbose=TRUE], copy(DT)[, v := base::min(b), a,     ], output="GForce optimized j to")
+test(2233.10, copy(DT)[, v := b[1:3], a], copy(DT)[, v := utils::head(b, 3L), a])
 
 # with key and grouping by nonkey
 DT = data.table(a=1:3,b=(1:9)/10,c=(3:1),key="c")
-test(2233.11, DT[, v := min(b), a, verbose=TRUE],      data.table(a=1:3, b=(1:9)/10, c=(3:1), v=(1:3)/10, key="c"), output="GForce optimized j to")
-test(2233.12, DT[, v := head(b, 3L), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, c=(3:1), v=(1:9)/10, key="c"), output="GForce optimized j to")
-test(2233.13, DT[, v := head(b, 2L), a], error="Supplied 6 items to be assigned to 9 items of column 'v'.")
+test(2233.11, DT[, v := min(b), a, verbose=TRUE],   data.table(a=1:3, b=(1:9)/10, c=(3:1), v=(1:3)/10, key="c"), output="GForce optimized j to")
+test(2233.12, DT[, v := shift(b), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, c=(3:1), v=c(NA,NA,NA,1:6)/10, key="c"), output="GForce optimized j to")
+test(2233.13, DT[, v := head(b, 2L), a], error="could be implemented")
 DT = data.table(a=1:3,b=(1:9)/10,c=(3:1),key="c")
-test(2233.14, copy(DT)[, v := min(b), a,      verbose=TRUE], copy(DT)[, v := base::min(b), a,     ], output="GForce optimized j to")
-test(2233.15, copy(DT)[, v := head(b, 3L), a, verbose=TRUE], copy(DT)[, v := utils::head(b, 3L), a], output="GForce optimized j to")
+test(2233.14, copy(DT)[, v := min(b), a, verbose=TRUE], copy(DT)[, v := base::min(b), a,     ], output="GForce optimized j to")
+test(2233.15, copy(DT)[, v := b[1:3], a], copy(DT)[, v := utils::head(b, 3L), a])
 
 # with key and keyby by nonkey
 DT = data.table(a=1:3,b=(1:9)/10,c=(3:1),key="c")
-test(2233.16, copy(DT)[, v := min(b),      keyby=a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, c=(3:1), v=(1:3)/10, key="a"), output="GForce optimized j to")
-test(2233.17, copy(DT)[, v := head(b, 3L), keyby=a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, c=(3:1), v=(1:9)/10, key="a"), output="GForce optimized j to")
-test(2233.18, copy(DT)[, v := head(b, 2L), keyby=a], error="Supplied 6 items to be assigned to 9 items of column 'v'.")
+test(2233.16, copy(DT)[, v := min(b),   keyby=a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, c=(3:1), v=(1:3)/10, key="a"), output="GForce optimized j to")
+test(2233.17, copy(DT)[, v := shift(b), keyby=a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, c=(3:1), v=c(NA,NA,NA,1:6)/10, key="a"), output="GForce optimized j to")
+test(2233.18, copy(DT)[, v := head(b, 2L), keyby=a], error="could be implemented")
 DT = data.table(a=1:3,b=(1:9)/10,c=(3:1),key="c")
-test(2233.19, copy(DT)[, v := min(b),      keyby=a, verbose=TRUE], copy(DT)[, v := base::min(b),       keyby=a], output="GForce optimized j to")
-test(2233.20, copy(DT)[, v := head(b, 3L), keyby=a, verbose=TRUE], copy(DT)[, v := utils::head(b, 3L), keyby=a], output="GForce optimized j to")
+test(2233.19, copy(DT)[, v := min(b), keyby=a, verbose=TRUE], copy(DT)[, v := base::min(b),       keyby=a], output="GForce optimized j to")
+test(2233.20, copy(DT)[, v := b[1:3], keyby=a], copy(DT)[, v := utils::head(b, 3L), keyby=a])
 # with irows
 DT = data.table(a=1:3,b=(1:9)/10)
-test(2233.21, DT[a==2, v := min(b), a, verbose=TRUE],      data.table(a=1:3, b=(1:9)/10, v=c(NA,0.2,NA)), output="GForce optimized j to")
-test(2233.22, DT[a!=4, v := head(b, 3L), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v=(1:9)/10),     output="GForce optimized j to")
-test(2233.23, DT[a!=4, v := head(b, 2L), a], error="Supplied 6 items to be assigned to 9 items of column 'v'.")
+test(2233.21, DT[a==2, v := min(b), a, verbose=TRUE],   data.table(a=1:3, b=(1:9)/10, v=c(NA,0.2,NA)), output="GForce optimized j to")
+test(2233.22, DT[a!=4, v := shift(b), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v=c(NA,NA,NA,1:6)/10), output="GForce optimized j to")
+test(2233.23, DT[a!=4, v := head(b, 2L), a], error="could be implemented")
 DT = data.table(a=1:3,b=(1:9)/10)
-test(2233.24, copy(DT)[a==2, v := min(b), a,      verbose=TRUE], copy(DT)[a==2, v := base::min(b), a,     ], output="GForce optimized j to")
-test(2233.25, copy(DT)[a!=4, v := head(b, 3L), a, verbose=TRUE], copy(DT)[a!=4, v := utils::head(b, 3L), a], output="GForce optimized j to")
+test(2233.24, copy(DT)[a==2, v := min(b), a, verbose=TRUE], copy(DT)[a==2, v := base::min(b), a,     ], output="GForce optimized j to")
+test(2233.25, copy(DT)[a!=4, v := b[1:3], a], copy(DT)[a!=4, v := utils::head(b, 3L), a])
 
 # multiple assignments
 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.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")
+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") := .(shift(b), shift(b,2L)), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v1=c(NA,NA,NA,1:6)/10, v2=c(NA,NA,NA,NA,NA,NA,1:3)/10), output="GForce optimized j to")
+test(2233.28, DT[, c("v1","v2") := .(shift(b), head(b,3L)), a],                error="head.*:=.*could be implemented")
+test(2233.29, DT[, c("v1","v2") := .(min(b), shift(b)), a, verbose=TRUE],      data.table(a=1:3, b=(1:9)/10, v1=(1:3)/10, v2=c(NA,NA,NA,1:6)/10), output="GForce optimized j to")
+test(2233.30, DT[, c("v1","v2") := .(head(b,2L), tail(b,2L)), a],              error="head.*:=.*could be implemented")
+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],              error="head.*:=.*could be implemented")
 
 # gforce needs to evaluate variable arguments before calling C part (part of test 101.17 in programming.Rraw)
 set.seed(108)
diff --git a/src/gsumm.c b/src/gsumm.c
@@ -21,6 +21,9 @@ static int *oo = NULL;
 static int *ff = NULL;
 static int isunsorted = 0;
 
+// for first/last with n>1 to error when used with :=, until implemented
+static bool assignByRef = false;
+
 // from R's src/cov.c (for variance / sd)
 #ifdef HAVE_LONG_DOUBLE
 # define SQRTL sqrtl
@@ -37,10 +40,11 @@ static int nbit(int n)
   return nb;
 }
 
-SEXP gforce(SEXP env, SEXP jsub, SEXP o, SEXP f, SEXP l, SEXP irowsArg, SEXP grpcols) {
+SEXP gforce(SEXP env, SEXP jsub, SEXP o, SEXP f, SEXP l, SEXP irowsArg, SEXP grpcols, SEXP lhs) {
   int nprotect=0;
   double started = wallclock();
   const bool verbose = GetVerbose();
+  assignByRef = !isNull(lhs);
   if (TYPEOF(env) != ENVSXP) error(_("env is not an environment"));
   // The type of jsub is pretty flexible in R, so leave checking to eval() below.
   if (!isInteger(o)) error(_("%s is not an integer vector"), "o");
@@ -1113,6 +1117,8 @@ static SEXP gfirstlast(const SEXP x, const bool first, const SEXP nArg, const bo
   if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<0)
     error(_("Internal error, gfirstlast is not implemented for n<0. This should have been caught before. Please report to data.table issue tracker.")); // # nocov
   const int w = INTEGER(nArg)[0];
+  if (w>1 && assignByRef)
+    error(_("Is first/last/head/tail with n>1 and := by group intentional? Please provide a use case to the GitHub issue tracker. It could be implemented."));
   // select 1:w when first=TRUE, and (n-w+1):n when first=FALSE
   // or select w'th item when nthvalue=TRUE; e.g. the n=0 case in test 280
   const bool nosubset = irowslen == -1;
