add arithmetic

Author: ben-schwen

R/data.table.R

@@ -3344,31 +3344,59 @@ is_constantish = function(q, check_singleton=FALSE) {
 
 .gforce_ok = function(q, x, envir=parent.frame(2L)) {
   if (is.N(q)) return(TRUE) # For #334
+  if (!is.call(q)) return(FALSE)  # plain columns are not gforce-able since they might not aggregate (see test 104.1)
+
   q1 = .get_gcall(q)
-  if (is.null(q1)) return(FALSE)
-  q2 = if (.is_type_conversion(q[[2L]]) && is.symbol(q[[2L]][[2L]])) q[[2L]][[2L]] else q[[2L]]
-  if (!q2 %chin% names(x) && q2 != ".I") return(FALSE)  # 875
-  if (length(q)==2L || (.arg_is_narm(q) && is_constantish(q[[3L]]) &&
-      !(is.symbol(q[[3L]]) && q[[3L]] %chin% names(x)))) return(TRUE)
-  switch(as.character(q1),
-    "shift" = .gshift_ok(q),
-    "weighted.mean" = .gweighted.mean_ok(q, x),
-    "tail" = , "head" = .ghead_ok(q),
-    "[[" = , "[" = `.g[_ok`(q, x, envir),
-    FALSE
-  )
+  if (!is.null(q1)) {
+    q2 = if (.is_type_conversion(q[[2L]]) && is.symbol(q[[2L]][[2L]])) q[[2L]][[2L]] else q[[2L]]
+    if (!q2 %chin% names(x) && q2 != ".I") return(FALSE)  # 875
+    if (length(q)==2L || (.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),
+      "weighted.mean" = .gweighted.mean_ok(q, x),
+      "tail" = , "head" = .ghead_ok(q),
+      "[[" = , "[" = `.g[_ok`(q, x, envir),
+      FALSE
+    ))
+  }
+
+  # check if arithmetic operator -> recursively validate ALL branches (like in AST)
+  if (is.symbol(q[[1L]]) && q[[1L]] %chin% c("+", "-", "*", "/", "^", "%%", "%/%")) {
+    for (i in 2:length(q)) {
+      if (!.gforce_ok(q[[i]], x, envir)) return(FALSE)
+    }
+    return(TRUE)
+  }
+
+  FALSE
 }
 
 .gforce_jsub = function(q, names_x, envir=parent.frame(2L)) {
-  call_name = if (is.symbol(q[[1L]])) q[[1L]] else q[[1L]][[3L]] # latter is like data.table::shift, #5942. .gshift_ok checked this will work.
-  q[[1L]] = as.name(paste0("g", call_name))
-  # gforce needs to evaluate arguments before calling C part TODO: move the evaluation into gforce_ok
-  # do not evaluate vars present as columns in x
-  if (length(q) >= 3L) {
-    for (i in 3:length(q)) {
-      if (is.symbol(q[[i]]) && !(q[[i]] %chin% names_x)) q[[i]] = eval(q[[i]], envir) # tests 1187.2 & 1187.4
+  if (!is.call(q)) return(q)
+
+  q1 = .get_gcall(q)
+  if (!is.null(q1)) {
+    call_name = if (is.symbol(q[[1L]])) q[[1L]] else q[[1L]][[3L]] # latter is like data.table::shift, #5942. .gshift_ok checked this will work.
+    q[[1L]] = as.name(paste0("g", call_name))
+    # gforce needs to evaluate arguments before calling C part TODO: move the evaluation into gforce_ok
+    # do not evaluate vars present as columns in x
+    if (length(q) >= 3L) {
+      for (i in 3:length(q)) {
+        if (is.symbol(q[[i]]) && !(q[[i]] %chin% names_x)) q[[i]] = eval(q[[i]], envir) # tests 1187.2 & 1187.4
+      }
+    }
+    return(q)
+  }
+
+  # if arithmetic operator, recursively substitute its operands. we know what branches are valid from .gforce_ok
+  if (is.symbol(q[[1L]]) && q[[1L]] %chin% c("+", "-", "*", "/", "^", "%%", "%/%")) {
+    for (i in 2:length(q)) {
+      q[[i]] = .gforce_jsub(q[[i]], names_x, envir)
     }
+    return(q)
   }
+  # should not reach here since .gforce_ok
   q
 }
 

inst/tests/tests.Rraw

@@ -21218,6 +21218,7 @@ test(2344.04, key(DT[, .(V4 = c("b", "a"), V2, V5 = c("y", "x"), V1)]), c("V1",
 # gforce should also work with Map in j #5336
 # conversions should not turn gforce off #2934
 # lapply gforce should also work without .SD #5032
+# support arithmetic in j with gforce #3815
 dt = data.table(a=1:4, b=1:2)
 out = c("GForce FALSE", "GForce FALSE", "GForce TRUE")
 test(2345.01, levels=0:2, dt[, max(as.character(a)), by=b, verbose=TRUE], data.table(b=1:2, V1=c("3","4")), output=out)
@@ -21231,3 +21232,7 @@ out = c("GForce FALSE", "lapply optimization changed j", "GForce TRUE")
 test(2345.06, levels=0:2, dt[, lapply(list(b, c), sum), by=a, verbose=TRUE], output=out)
 test(2345.07, levels=0:2, dt[, c(list(sum(b), sum(c))), by=a, verbose=TRUE], output=out)
 test(2345.08, levels=0:2, names(dt[, lapply(list(b, c), sum), by=a]))
+dt = data.table(a=1:4, b=1:2)
+out = c("GForce FALSE", "GForce FALSE", "GForce TRUE")
+test(2345.09, levels=0:2, dt[, .(max(a)-min(a)), by=b, verbose=TRUE], output=out)
+dt = data.table(a=1:4, b=1:2)