Skip to content

WIP: Avoid dense n x n allocations in upper_gather() - #1635

Open
kei51e wants to merge 4 commits into
masterfrom
perf/upper-gather-avoid-dense-matrix
Open

WIP: Avoid dense n x n allocations in upper_gather()#1635
kei51e wants to merge 4 commits into
masterfrom
perf/upper-gather-avoid-dense-matrix

Conversation

@kei51e

@kei51e kei51e commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Problem

upper_gather() is the shared long-form builder behind every pairwise function
(do_dist.kv / do_dist.cols / do_kl_dist.kv / do_cosine_sim.kv / pair_count).
It was the dominant cost of all of them, by a wide margin.

Measured on a 4000 x 5 input: stats::dist() itself takes 0.07s, while the
upper_gather() step took 8.3s and peaked at 1.1GB. That is a 117x ratio, so the
distance computation was never the bottleneck -- the reshape after it was.

The dist vector path allocated six dense n x n objects: the triangular matrix, its
row() and col() index matrices, a logical mask, a transpose, and one more copy inside
mat_to_df(). That is roughly 60 * n^2 bytes, i.e. ~6GB at n = 10000.

The matrix path had a second copy of the same problem. With zero.rm = FALSE it built
is.na(tmat) | !is.na(tmat) -- an all-TRUE matrix of the same shape -- purely so that
which() could be handed every cell. For a sparse input that forces the sparse class to
store all n^2 cells, which is heavier than a plain dense logical, and which() then
builds an n^2 x 2 index matrix on top of it.

What changed

R/util.R, upper_gather() only. No caller was changed, and mat_to_df() is untouched
(the vector path simply no longer calls it). The signature and the defaults are the same.

  • Vector path (dist): six n x n allocations -> zero. The long form is built directly
    from the dist vector. stats::dist stores the lower half in column-major order, which
    is the same sequence as the upper half in row-major order, so the values land in their
    original order with no matrix involved.
  • Matrix path, zero.rm = FALSE: no more all-cell materialisation. The surviving
    indices are known in closed form, so only the half that passes the row/column comparison
    is generated, in the same column-major order Matrix::which(arr.ind = TRUE) produced.
  • zero.rm = TRUE (the pair_count path) is untouched -- tmat != 0 keeps a sparse
    matrix sparse, so it never had the problem.

Equivalence

bench/micro/upper_gather_equivalence.R (added) compares this implementation against the
verbatim previous one -- extracted from the base revision, not retyped -- with
identical(), so row order, row count, column names, column types, NA positions and class
all have to match.

49 cases: the six argument combinations the callers actually use, every caller's exact
arguments, degenerate cases (n = 1, n = 2, all-NA, all-zero, single surviving pair, 1x1),
non-square and wide matrices, sparse dgCMatrix inputs with and without dimnames, and
multibyte / punctuation-heavy / near-colliding / duplicated dimension names.

47 of 49 are identical(), one is an error whose message matches exactly, and two are
the intentional divergence below. The harness names those two explicitly and exits non-zero
on any other difference.

One intentional behaviour change -- please review

Duplicated dimension names on the vector path. The previous implementation raised
The `.data` argument of `add_column()` must have unique names as of tibble 3.0.0 from
inside mat_to_df(). This one returns the natural result with the duplicated names kept.

This also removes an inconsistency: the matrix path already accepted duplicated names
and returned a result (case D24 in the harness is identical() between old and new, both
succeeding). Only the vector path errored, and only because of an internal tibble detail
leaking out, not because of anything the function guarantees. Making the two paths agree
seems right, but it is a behaviour change, so flagging it rather than burying it.

Benchmark

bench/micro/upper_gather_bench.R (added). Same arguments the callers use
(na.rm = FALSE, zero.rm = FALSE), gc() peak, row count checked equal every time.
Measured locally on R 4.6.1 (Apple silicon):

Vector path (dist):

n dist() before after before peak after peak rows speedup
500 0.000s 0.038s 0.017s 178 MB 143 MB 250,000 2.2x
1000 0.002s 0.074s 0.011s 172 MB 144 MB 1,000,000 6.7x
2000 0.006s 0.273s 0.043s 329 MB 242 MB 4,000,000 6.3x
4000 0.069s 0.806s 0.170s 975 MB 631 MB 16,000,000 4.7x

Matrix path, sparse input:

n before after before peak after peak rows speedup
500 0.026s 0.004s 256 MB 187 MB 124,750 6.5x
1000 0.081s 0.013s 381 MB 226 MB 499,500 6.2x
2000 0.473s 0.055s 503 MB 350 MB 1,999,000 8.6x
4000 2.228s 0.218s 1337 MB 819 MB 7,998,000 10.2x

On a Linux box with more headroom the vector path reached 10.3x at n = 4000.

What this does NOT fix

Worth being explicit, so nobody reads more into the numbers than is there:

  • Memory drops by about a third, not by an order of magnitude. The n x n intermediates
    are gone, but the output itself is n^2 rows x 3 columns whenever
    na.rm = FALSE, zero.rm = FALSE (which is 4 of the 5 callers), i.e. ~24 * n^2 bytes.
    Reducing that means changing what the function returns, which is a separate decision and
    is deliberately not in this PR.
  • do_cmdscale is not fixed by this. It rebuilds an n x n matrix of its own via
    simple_cast() + t() + as.dist(). It does get faster input, but its own allocation
    is untouched.

Testing

  • test_util.R, test_pairwise.R, test_pair_count.R, test_stats_wrapper.R all pass
    locally under devtools::load_all() on R 4.6.1 -- zero failures. Those are the four
    files that reach upper_gather (30 test_that blocks).
  • Two test_that blocks added to test_util.R pinning the output contract the callers
    depend on but nothing tested: the full n x n grid, the row-major order, which half
    carries values, column types, and verbatim dimension names.
  • Non-ASCII characters appear only in test data (dimension names), as intended.

Two follow-ups noticed while doing this, not addressed here:

  • do_kl_dist has no test coverage at all -- it is the only one of the five callers
    with none.
  • The vector path returns a tibble while the matrix path returns a data.frame. That
    asymmetry is pre-existing; it was preserved deliberately, but it does not look intended.

Note on R version

The new code uses sequence(nvec =, from =), which needs R >= 4.0.0. DESCRIPTION
still nominally says R (>= 2.10), which was already inaccurate, and the product ships
R 4.x -- so this is not a practical constraint, but it is a real one.

claude and others added 4 commits August 25, 2026 06:50
upper_gather() dominated the cost of every pairwise function that calls it
(distance, cosine similarity, KL distance, pair counting). Measured on a
4000 x 5 input, stats::dist() itself took 0.071s while the upper_gather()
step took 8.3s and peaked at 1.1GB.

The dist vector path built a dense n x n triangular matrix, the row() and
col() index matrices for it, a logical mask, a transpose, and one more n x n
copy inside mat_to_df(): six n x n allocations, about 60 * n^2 bytes, which is
6GB at n = 10000. It now builds the long form directly from the dist vector.
stats::dist stores the lower half in column major order, which is the same
sequence as the upper half in row major order, so the values keep their
original order and no n x n object is allocated at all.

The matrix path had a second copy of the same problem. With zero.rm = FALSE it
evaluated is.na(tmat) | !is.na(tmat), an all TRUE matrix of the same shape,
purely to hand every cell to which(). For a sparse input that forces the sparse
class to store all n^2 cells, which is heavier than a plain dense logical. The
selected indices are known in closed form there, so only the surviving half is
generated.

The output contract is unchanged: same row order, row count, column names,
column types, and NA positions. bench/micro/upper_gather_equivalence.R compares
this implementation against the previous one with identical() over the six
argument combinations the callers use plus degenerate, non square, multibyte
and duplicated name cases: 47 of 49 identical. The two remaining cases are
duplicated dimension names on the vector path, where the previous
implementation raised an error from tibble inside mat_to_df() and this one
returns the natural result.

mat_to_df() itself is deliberately left untouched; the vector path simply no
longer calls it.
The five callers of upper_gather() depend on details that were not covered by
any test: the full n x n grid returned when na.rm and zero.rm are FALSE, the
row major order, which half of the grid carries values, the column types, and
the fact that dimension names are copied into the output verbatim.

Adds two test_that blocks to test_util.R covering the six argument
combinations the callers use, plus the degenerate cases (n = 1, n = 2, all NA)
and dimension names that are multibyte, punctuation heavy, near colliding or
duplicated.

Adds bench/micro/upper_gather_bench.R next to the equivalence harness. It
compares both paths against the previous implementation over a range of n and
reports elapsed time, gc peak and row count.
gc() inserts a "limit (Mb)" column when a memory limit is set, which shifts
column 6 from "max used (Mb)" to "max used" in cells. The benchmark indexed
column 6 directly, so on a machine with a limit set it reported cell counts
instead of megabytes (about 9.5e6 where the value is 178 Mb). "max used (Mb)"
is the last column in both shapes, so index from the end.
@kei51e kei51e changed the title Avoid dense n x n allocations in upper_gather() WIP: Avoid dense n x n allocations in upper_gather() Aug 28, 2026
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.

2 participants