Skip to content

Build a grouped mutate() of arithmetic from the columns' rows - #456

Merged
jbearak merged 2 commits into
mainfrom
claude/arith-rows
Oct 26, 2026
Merged

jbearak merged 2 commits into
mainfrom
claude/arith-rows

Conversation

@jbearak

@jbearak jbearak commented Oct 26, 2026 •

Copy link
Copy Markdown
Owner

Summary

ADR 0065's plan of arithmetic in a grouped mutate() evaluated the tree for each group after the first from the columns' chunks, and returned a result vector for each group, which the steps after the loop combined. #455 leaves a column that only the first group has read unsplit, but this plan asked for the chunks, so it split the column anyway. With many small groups, the splits, the groups' results and the combine took most of the time. Over the 42,739 women in the first 100,000 rows of the India 2021 DHS births file, mutate(b, m = (b2 - b4) / 12, .by = caseid) took about 0.24 seconds and mutate(b, k = b3 - min(b3), .by = caseid) about 0.13. A tibble takes about 0.07 for each.

This PR adds C_dtatools_arith_plan_rows(). When a tree holds only numbers, (, columns, mean(), min(), max(), sum() or median() of a column, and the four operators, and only the first group has read each of its columns, the call evaluates the tree for every group at once from the whole columns:

  • It gathers each column's rows once, in the groups' order, as the split gathers each chunk.
  • It computes each summary of a column for every group from those rows, with Split a column for later groups on their first read #455's helpers, now split into a gather step and a step that reads each group's span.
  • For each operator it computes every group's values as the native small arithmetic does, then types each group's values as .dta_computed() types a small result: the minimum storage the package's method passes, the first storage from it that holds every value, and each value as that storage reads it back.
  • The first group's result, which the loop built, must match what the call gives that group.
  • Assemble a mutate() of a summary of one column #454's dtatools_planned_summary_rows() then builds the column the combine would build from every group's values and storages.

.arith_planned() takes the first group's result as head and tries this call first in a mutate(). .arith_plan_columns() takes unsplit = TRUE to return the whole columns after the same admission checks. Any other tree, a column already split, or a decline goes to the plan of the chunks as before.

Behavior

Values, attributes, ALTREP forms, warnings and errors are unchanged. ADR 0118 records the decision. The columns must be compact Stata numerics that the native gather takes, the rows must be bare and cover each row once, and every group must be shorter than the size gate, under which each group's arithmetic reaches the native call. The R code checks the same closures for a debugger as the plan of the chunks does.

Timing

Best of three runs in fresh R sessions, in seconds, with the builds interleaved, over the first 100,000 rows of the India 2021 DHS births file. Before is #455's head.

Call Before After Tibble
mutate(b, m = (b2 - b4) / 12, .by = caseid) 0.24 0.011 0.070
mutate(b, k = b3 - min(b3), .by = caseid) 0.13 0.010 0.067
mutate(b, across(c(b2, b3, bord), ~ .x - mean(.x)), .by = caseid) 0.74 0.024 0.33

GC time fell from about 0.08, 0.012 and 0.37 seconds to none in each call.

Tests

  • A test runs 13 shapes of arithmetic of each of the eight columns of Assemble a mutate() of a summary of one column #454's data, such as x - mean(x), (x - median(x)) * 2 and x * 1e6, and each operator on every pair of columns, 360 cases. Each must give the same value, warnings and result form with the plans as without them. 287 build their column from the rows. The other 73 read the double column, which is not compact. A case that builds its column splits no column.
  • A test runs 80 random trees of columns, numbers and summaries over random data whose groups interleave and whose values sit at each storage's limits, so that the groups' results take different storages and join.
  • A test divides a column by min() and max() of a column declared double whose second group has no value. Those summaries are Inf and -Inf, which the plan of the chunks reads back as missing, so the second group's results must be missing.
  • A test runs cases the call leaves to the loop: a group at the size gate, a column the lag() plan has split, lag() in the tree, rowwise and ungrouped data, summarise() and filter(). It also checks that a column an earlier expression built, and a column the result replaces, build from the rows.
  • A test debugs each closure on the way while the plan decides, and checks that the plan declines.
  • Two tests in test-compact-combine.R count the combine's native calls, and their arithmetic now builds its column without the combine. One adds t = identity(i + 1L), whose groups the loop still combines, and expects two calls where it expected three. The other wraps weight - min(weight) in identity(), so that its first grouped mutate() in a fresh session still reaches the combine.
  • The full suite passes: 168 files, 2,263 tests. The native manifest and repository guard tests pass. Codex review is clean.

Arithmetic of columns, numbers and summaries of columns in a grouped
mutate() now evaluates every group at once from the columns' rows, when
only the first group has read each column, and builds the column the
combine would build (ADR 0118).
@coderabbitai

coderabbitai Bot commented Oct 26, 2026 •

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

# Conflicts:
#	r-package/dtatools/tools/native-test-manifest.json
@jbearak
jbearak changed the base branch from claude/lazy-first-split to main October 26, 2026 23:55
@jbearak
jbearak marked this pull request as ready for review October 26, 2026 23:55
@jbearak
jbearak merged commit 22dd077 into main Oct 26, 2026
4 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant