Plan cumulative functions of a logical tree

docs/adr/0123-plan-cumulative-functions-of-a-logical-tree.md

@@ -0,0 +1,33 @@
+---
+status: accepted
+---
+
+# Plan cumulative functions of a logical tree
+
+ADR 0085 lets a planned arithmetic tree take `cumsum()`, `cummax()` and
+`cummin()` of a Stata numeric operand, which dispatch to the package's
+`Math.dta_numeric()`. Of a comparison, such as the running count
+`cumsum(b5 == 1)`, the tree went to the loop. The comparison gives a bare
+logical vector, and each group's evaluation called `Ops.dta_numeric()` and
+base's `cumsum()`. Over the 42,739 women in the first 100,000 rows of the
+India 2021 DHS births file, `mutate(b, c = cumsum(b5 == 1), .by = caseid)`
+took about 0.30 seconds, where a tibble takes about 0.060, and
+`filter(b, cumsum(b5 == 0) <= 1, .by = caseid)` about 0.43, where a tibble
+takes about 0.059.
+
+The tree now takes `cumsum()`, `cummax()` and `cummin()` of a logical
+tree, which gives a bare integer vector for each group. The call must find
+base's primitive, which dispatches on no bare logical vector. As base's
+primitive computes it, the value is the running count, maximum or minimum
+of the logical values as integers, and missing from the first missing
+value on. A logical vector no longer than `INT_MAX` cannot overflow
+`cumsum()`'s integer.
+
+Such a vector may be a mutate()'s whole expression, or compare with a bare
+number, as the group's positions from `row_number()` do (ADR 0099), which
+gives a logical tree that mutate() and filter() take. Arithmetic of it, a
+summary of it, and abs() and the rounding functions of a logical vector,
+which give other types, still go to the loop.
+
+Results are unchanged. The calls above now take about 0.050 and 0.059
+seconds.

r-package/dtatools/NEWS.md

@@ -5,6 +5,13 @@
   Use `summ()` for Stata's `summarize` command, with the same arguments and
   return value.
 
+* A grouped `mutate()` or `filter()` now plans `cumsum()`, `cummax()` and
+  `cummin()` of a comparison or another logical tree, such as the running
+  count `cumsum(b5 == 1)`, and comparisons of their results with a number.
+  Over the 42,739 women in the first 100,000 rows of the India 2021 DHS
+  births file, `mutate(b, c = cumsum(b5 == 1), .by = caseid)` falls from
+  about 0.30 to 0.050 seconds (ADR 0123).
+
 * A grouped `slice()` of numbers, such as `slice(gb, 1)`, `slice(gb, -1)`
   or `slice(gb, 1:3)`, now takes the rows of every group after the first
   in one native call. Over the 42,739 women in the first 100,000 rows of

r-package/dtatools/R/dibble-summaries.R

@@ -1905,7 +1905,8 @@
     # A Math function of one unnamed argument dispatches to the package's
     # Math.dta_numeric() (ADR 0085). round() may also take a finite number
     # of digits, unnamed or named `digits`, which the method passes on to
-    # base's round() (ADR 0111).
+    # base's round() (ADR 0111). cumsum(), cummax() and cummin() of a
+    # logical tree call base's primitive (ADR 0123).
     # An empty argument gives round() its default digits, and binding it
     # here would make reading it an error.
     rounds <- identical(generic, "round") && length(expression) == 3L &&

r-package/dtatools/src/summarize.c

@@ -4179,7 +4179,8 @@ SEXP C_dtatools_compare_plan(SEXP chunks, SEXP generic, SEXP operand, SEXP left,
                             with digits a finite number;
    and abs(), floor(), ceiling(), trunc() and round() of a bare number
    from ADR 0098's operands, which give a bare number as base R computes
-   it.
+   it. ADR 0123 adds cumsum(), cummax() and cummin() of a logical node,
+   which give a bare integer vector as base R computes it.
    ADR 0098 adds operands that give a bare number, which count as
    summaries:
    list("lsummary", generic, node, frame)
@@ -4217,7 +4218,8 @@ SEXP C_dtatools_compare_plan(SEXP chunks, SEXP generic, SEXP operand, SEXP left,
                             numeric operand, with na.rm as frame binds it,
                             whose columns hold the group's rows. */
 enum { ARITH_STATA, ARITH_DOUBLE, ARITH_INTEGER, TREE_LOGICAL, TREE_SCALAR, TREE_ONE,
-       PLAIN_INTEGER, PLAIN_DOUBLE, TREE_PLAIN, PLAIN_RANKS, PLAIN_LOGICAL, TREE_FLAG };
+       PLAIN_INTEGER, PLAIN_DOUBLE, TREE_PLAIN, PLAIN_RANKS, PLAIN_LOGICAL, TREE_FLAG,
+       TREE_COUNTS };
 #define ARITH_SHIFTS 8
 
 /* The positions of the dependencies, the records and the records' list of
@@ -4866,6 +4868,31 @@ static SEXP tree_math_plain(int which, SEXP value, const double *digits) {
     return Rf_ScalarReal(y);
 }
 
+/* ADR 0123. cumsum(), cummax() or cummin() of a bare logical vector, as
+   base's primitive computes it for one: an integer vector, missing from
+   the first missing value on. A logical vector no longer than INT_MAX
+   cannot overflow cumsum()'s integer. R_NilValue for anything else. */
+static SEXP tree_math_counts(int which, SEXP value) {
+    if (TYPEOF(value) != LGLSXP || ALTREP(value) || ANY_ATTRIB(value) || which < 5 ||
+        which > 7 || XLENGTH(value) > INT_MAX) return R_NilValue;
+    R_xlen_t length = XLENGTH(value);
+    SEXP result = PROTECT(Rf_allocVector(INTSXP, length));
+    const int *in = LOGICAL(value);
+    int *out = INTEGER(result), running = 0, missing = 0;
+    for (R_xlen_t i = 0; i < length; i++) {
+        if (missing || in[i] == NA_LOGICAL) {
+            missing = 1;
+            out[i] = NA_INTEGER;
+            continue;
+        }
+        if (which == 5) running += in[i];
+        else if (i == 0 || (which == 6 ? in[i] > running : in[i] < running)) running = in[i];
+        out[i] = running;
+    }
+    UNPROTECT(1);
+    return result;
+}
+
 /* The primitives the closures on a Math function's way call, from the
    package's namespace for Math.dta_numeric(), vec_math.dta_numeric(),
    `.restore_missing()` and `.collapse_missing()`, and from base's for
@@ -5237,6 +5264,11 @@ static int arith_tree_admit(arith_tree *tree, SEXP node, int depth, R_xlen_t cou
         if (inner == PLAIN_INTEGER || inner == PLAIN_DOUBLE)
             return which <= 4 && arith_tree_math_plain(tree, which) ?
                 (which == 0 && inner == PLAIN_INTEGER ? PLAIN_INTEGER : PLAIN_DOUBLE) : -1;
+        /* cumsum(), cummax() or cummin() of a logical vector, on which
+           base's primitive dispatches to no method, gives a bare integer
+           vector, which compares as the group's positions do (ADR 0123). */
+        if (inner == TREE_LOGICAL)
+            return which >= 5 && arith_tree_math_plain(tree, which) ? TREE_COUNTS : -1;
         return inner == ARITH_STATA && arith_tree_math(tree, which) ? ARITH_STATA : -1;
     }
     int before = *columns;
@@ -5266,9 +5298,12 @@ static int arith_tree_admit(arith_tree *tree, SEXP node, int depth, R_xlen_t cou
     }
     if (left == TREE_LOGICAL || right == TREE_LOGICAL || left == PLAIN_LOGICAL ||
         right == PLAIN_LOGICAL) return -1;
-    int ranks = left == PLAIN_RANKS || right == PLAIN_RANKS;
-    int plain_left = left == PLAIN_INTEGER || left == PLAIN_DOUBLE || left == PLAIN_RANKS,
-        plain_right = right == PLAIN_INTEGER || right == PLAIN_DOUBLE || right == PLAIN_RANKS;
+    int ranks = left == PLAIN_RANKS || right == PLAIN_RANKS || left == TREE_COUNTS ||
+        right == TREE_COUNTS;
+    int plain_left = left == PLAIN_INTEGER || left == PLAIN_DOUBLE || left == PLAIN_RANKS ||
+        left == TREE_COUNTS,
+        plain_right = right == PLAIN_INTEGER || right == PLAIN_DOUBLE ||
+        right == PLAIN_RANKS || right == TREE_COUNTS;
     if (left != ARITH_STATA && right != ARITH_STATA) {
         if (!plain_left && !plain_right) return -1;
         /* A comparison of two bare numbers gives a logical scalar, and one
@@ -5835,6 +5870,12 @@ static SEXP arith_tree_value(arith_tree *tree, SEXP node, SEXP columns, R_xlen_t
         if (!tree_math_digits(node, which, &storage, &digits)) return R_NilValue;
         SEXP inner = arith_tree_value(tree, VECTOR_ELT(node, 2), columns, group, size, frame);
         if (inner == R_NilValue) return R_NilValue;
+        if (TYPEOF(inner) == LGLSXP) {
+            PROTECT(inner);
+            SEXP value = tree_math_counts(which, inner);
+            UNPROTECT(1);
+            return value;
+        }
         if (!Rf_isObject(inner)) return tree_math_plain(which, inner, digits);
         if (!arith_tree_stata(inner)) return R_NilValue;
         PROTECT(inner);
@@ -5904,7 +5945,7 @@ static int arith_tree_root(arith_tree *tree, SEXP node, R_xlen_t count) {
         return kind == ARITH_STATA && (tree->operators || tree->maths) &&
             tree->summaries + tree->picks > 0 ? TREE_ONE : -1;
     }
-    return columns > 0 && (kind == TREE_LOGICAL || kind == TREE_SCALAR ||
+    return columns > 0 && (kind == TREE_LOGICAL || kind == TREE_SCALAR || kind == TREE_COUNTS ||
                            (kind == ARITH_STATA && (tree->operators || tree->maths))) ?
         kind : -1;
 }

r-package/dtatools/tests/testthat/test-summary-plan.R

@@ -9150,3 +9150,77 @@ test_that("a value an old file wrote leaves the groups to the loop where a slice
                          info = info)
     }
 })
+
+# ADR 0123: cumsum(), cummax() and cummin() of a logical tree, which give a
+# bare integer vector for each group, at the root or in a comparison.
+
+test_that("cumulative functions of a logical tree give what the group loop gives", {
+    skip_if_not_installed("dplyr", "1.2.1")
+    data <- .plan_data()
+    grouped <- dplyr::group_by(data, g)
+    planned <- cases <- 0L
+    check <- function(code) {
+        result <- .arith_check(.between_global_call(code, data = data, grouped = grouped))
+        info <- paste(deparse(code), collapse = " ")
+        expect_true(result$same, info = info)
+        expect_false(result$error, info = info)
+        planned <<- planned + result$planned
+        cases <<- cases + 1L
+    }
+    # Each function of a comparison of each column, whose missing and tagged
+    # missing values make the rest of a group's results missing.
+    for (name in c("cumsum", "cummax", "cummin")) for (column in setdiff(names(data), "g")) {
+        test <- rlang::call2(">", rlang::sym(column), 1)
+        check(rlang::expr(dplyr::mutate(data, k = !!rlang::call2(name, test), .by = g)))
+        check(rlang::expr(dplyr::mutate(grouped, k = !!rlang::call2(name, test))))
+    }
+    expect_identical(planned, cases)
+    # Other logical trees, and comparisons of the counts with a number, in
+    # mutate() and filter().
+    planned <- cases <- 0L
+    for (code in list(
+        quote(dplyr::mutate(data, k = cumsum(is.na(int)), .by = g)),
+        quote(dplyr::mutate(data, k = cummax(!(byte == 2)), .by = g)),
+        quote(dplyr::mutate(data, k = cumsum(byte == 4 | int > 10), .by = g)),
+        quote(dplyr::mutate(data, k = cummin((float < 1)), .by = g)),
+        quote(dplyr::mutate(data, k = (cumsum(labelled == 1)), .by = g)),
+        quote(dplyr::mutate(data, k = cumsum(cumsum(long > 3) > 1), .by = g)),
+        quote(dplyr::mutate(data, k = cumsum(byte > 1) == 1L, .by = g)),
+        quote(dplyr::mutate(data, k = 2 <= cumsum(described > 0), .by = g)),
+        quote(dplyr::filter(data, cumsum(byte > 1) <= 1, .by = g)),
+        quote(dplyr::filter(grouped, cummax(int > 5) == 0)),
+        quote(dplyr::filter(grouped, cumsum(double > 0) > 1 & byte > 0))))
+        check(code)
+    expect_identical(planned, cases)
+    # between() as the logical tree, which needs ADR 0081's plan.
+    planned <- cases <- 0L
+    check(quote(dplyr::mutate(data, k = cumsum(between(int, 0, 9)), .by = g)))
+    expect_identical(planned, as.integer(.between_plan_expected()))
+})
+
+test_that("cumulative functions of a logical tree decline what the loop would see differently", {
+    skip_if_not_installed("dplyr", "1.2.1")
+    data <- .plan_data()
+    check <- function(code, ...) {
+        result <- .arith_check(.plan_global_call(code, data = data, ...))
+        info <- paste(deparse(code), collapse = " ")
+        expect_true(result$same, info = info)
+        expect_identical(result$planned, 0L, info = info)
+    }
+    # Another function of the name, where R finds it first, and a column
+    # named for it.
+    check(quote(dplyr::mutate(data, k = cumsum(byte > 1), .by = g)),
+          cumsum = function(x) base::cumsum(x) + 1L)
+    named <- dplyr::mutate(data, cumsum = 1)
+    check(quote(dplyr::mutate(named, k = cumsum(byte > 1), .by = g)), named = named)
+    # Arithmetic of the counts, the other Math functions and a logical
+    # scalar, which no column gives, a summary, and rowwise data.
+    for (code in list(
+        quote(dplyr::mutate(data, k = cumsum(byte > 1) + 1L, .by = g)),
+        quote(dplyr::mutate(data, k = cumsum(byte > 1) * int, .by = g)),
+        quote(dplyr::mutate(data, k = abs(byte > 1), .by = g)),
+        quote(dplyr::mutate(data, k = cumsum(dplyr::n() > 1), .by = g)),
+        quote(dplyr::summarise(data, k = cumsum(byte > 2), .by = g)),
+        quote(dplyr::mutate(dplyr::rowwise(data), k = cumsum(byte > 1)))))
+        check(code)
+})

r-package/dtatools/tools/native-test-manifest.json

@@ -18794,7 +18794,7 @@
         },
         {
           "path": "tests/testthat/test-summary-plan.R",
-          "sha256": "8a991f5efebbfb497820028d701560cdf5c7edb350414b301b42c2f7d0868faf"
+          "sha256": "294d2baa82a01d3e9d3af1849004b34c12ab6584a237a3caf842371ecac05fa8"
         }
       ],
       "blocks": [
@@ -21378,6 +21378,28 @@
           "min_pass_before_skip": 0,
           "reason": "Optional dplyr comparison retained in the present-dependency suite.",
           "skip_message": "Reason: {dplyr} is not installed"
+        },
+        {
+          "file": "test-summary-plan.R",
+          "test": "cumulative functions of a logical tree give what the group loop gives",
+          "occurrence": 1,
+          "skip": "require",
+          "min_pass": 123,
+          "warnings": 0,
+          "min_pass_before_skip": 0,
+          "reason": "Optional dplyr comparison retained in the present-dependency suite.",
+          "skip_message": "Reason: {dplyr} is not installed"
+        },
+        {
+          "file": "test-summary-plan.R",
+          "test": "cumulative functions of a logical tree decline what the loop would see differently",
+          "occurrence": 1,
+          "skip": "require",
+          "min_pass": 16,
+          "warnings": 0,
+          "min_pass_before_skip": 0,
+          "reason": "Optional dplyr comparison retained in the present-dependency suite.",
+          "skip_message": "Reason: {dplyr} is not installed"
         }
       ],
       "children": [