WIP: Avoid dense n x n allocations in upper_gather()

R/util.R

@@ -219,18 +219,57 @@ upper_gather <- function(mat, names=NULL, diag=NULL, cnames = c("Var1", "Var2",
       }
     }
 
-    # create a triangler matrix to melt
-    # use NA_real_ for performance
-    trimat <- matrix(data=NA_real_, nrow=length(names), ncol=length(names))
-    # fill only lower half of the matrix (transpose later to keep the order)
-    trimat[row(trimat)>col(trimat)] <- as.numeric(mat)
-    colnames(trimat) <- names
-    rownames(trimat) <- names
+    # Build the long form directly, without ever allocating an n x n matrix.
+    # The previous implementation created a dense n x n triangular matrix, its
+    # row()/col() index matrices, a logical mask, a transpose, and a data frame
+    # copy inside mat_to_df. That is 6 n x n allocations (about 60 * n^2 bytes),
+    # which is what made this the dominant cost of the pairwise functions.
+    #
+    # Output contract that is reproduced here exactly:
+    #   - one row per cell of the full n x n grid, in row major order, so the
+    #     first name varies slowly and the second name varies fast
+    #   - the value is filled only on the strict upper half (i < j); the rest
+    #     stays NA unless diag is given
+    #   - 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 of
+    #     mat are placed in their original order
+    n <- length(names)
+    names_chr <- as.character(names)
+
+    value <- rep(NA_real_, n * n)
+    if (n >= 2L) {
+      # i is the slow index, j the fast one, over the strict upper half
+      i <- rep.int(seq_len(n - 1L), times = (n - 1L):1L)
+      j <- sequence(nvec = (n - 1L):1L, from = 2:n)
+      value[(i - 1L) * n + j] <- as.numeric(mat)
+    }
     if(!is.null(diag)){
       # fill diagonal elements
-      trimat[row(trimat)==col(trimat)] = rep(diag, length(names))
+      value[seq.int(from = 1L, by = n + 1L, length.out = n)] <- rep(diag, length.out = n)
+    }
+
+    var1 <- rep(names_chr, each = n)
+    var2 <- rep.int(names_chr, times = n)
+
+    # mat_to_df() applies na.rm first and zero.rm after it; both keep the row
+    # order, so a single combined mask gives the same result.
+    keep <- NULL
+    if(na.rm){
+      keep <- !is.na(value)
+    }
+    if(zero.rm){
+      not_zero <- is.na(value) | value != 0
+      keep <- if(is.null(keep)) not_zero else keep & not_zero
     }
-    mat_to_df(t(trimat), na.rm=na.rm, cnames=cnames, zero.rm = zero.rm)
+    if(!is.null(keep)){
+      var1 <- var1[keep]
+      var2 <- var2[keep]
+      value <- value[keep]
+    }
+
+    df <- tibble::tibble(Var1 = var1, Var2 = var2, value = value)
+    colnames(df) <- cnames
+    df
   } else {
     # diag can be NULL or FALSE
     if(is.null(diag)){
@@ -247,34 +286,53 @@ upper_gather <- function(mat, names=NULL, diag=NULL, cnames = c("Var1", "Var2",
       r_names <- seq(nrow(tmat))
     }
 
-    # remove 0 if zero.rm is TRUE
-    ind_mat <- if(zero.rm){
-      tmat != 0
-    } else {
-      is.na(tmat) | !is.na(tmat) # Just return matrix of same shape with TRUE for all the values.
-    }
-    # preserve NA if na.rm is FALSE
-    if(!na.rm){
-      ind_mat <- is.na(ind_mat) | ind_mat
-    }
-    # get indice of matrix
-    ind <- Matrix::which(ind_mat, arr.ind = TRUE)
+    if(zero.rm){
+      # remove 0. tmat != 0 keeps a sparse matrix sparse, so this stays cheap.
+      ind_mat <- tmat != 0
+      # preserve NA if na.rm is FALSE
+      if(!na.rm){
+        ind_mat <- is.na(ind_mat) | ind_mat
+      }
+      # get indice of matrix
+      ind <- Matrix::which(ind_mat, arr.ind = TRUE)
 
-    # remove duplicated pairs
-    # by comparing indice
-    filtered <- if(diag) {
-      ind[ind[,2] <= ind[,1], ]
-    } else {
-      ind[ind[,2] < ind[,1], ]
-    }
+      # remove duplicated pairs
+      # by comparing indice
+      filtered <- if(diag) {
+        ind[ind[,2] <= ind[,1], ]
+      } else {
+        ind[ind[,2] < ind[,1], ]
+      }
 
-    # when there is only one index pairs,
-    # filtered becomes a vector, not matrix
-    # but matrix is expected later
-    # so should be converted to matrix with
-    # one row
-    if(is.vector(filtered)){
-      filtered <- t(as.matrix(filtered))
+      # when there is only one index pairs,
+      # filtered becomes a vector, not matrix
+      # but matrix is expected later
+      # so should be converted to matrix with
+      # one row
+      if(is.vector(filtered)){
+        filtered <- t(as.matrix(filtered))
+      }
+    } else {
+      # zero.rm FALSE means "keep the zeros too", so every cell is selected and
+      # na.rm cannot remove anything either. The previous implementation spelled
+      # that out as is.na(tmat) | !is.na(tmat), an all TRUE matrix of the same
+      # shape, and then called which() on it, which materialises every one of
+      # the n^2 cells (for a sparse input it even forces the sparse class to
+      # store all of them) and builds an n^2 x 2 index matrix on top.
+      # The selected indices are known in closed form, so build only the half
+      # that survives the row/column comparison below.
+      # Matrix::which(, arr.ind = TRUE) walks in column major order, so the
+      # pairs are ordered by column and then by row, which is reproduced here.
+      n_row <- nrow(tmat)
+      n_col <- ncol(tmat)
+      col_seq <- seq_len(n_col)
+      first_row <- if(diag) col_seq else col_seq + 1L
+      counts <- pmax(n_row - first_row + 1L, 0L)
+      has_any <- counts > 0L
+      filtered <- cbind(
+        row = sequence(nvec = counts[has_any], from = first_row[has_any]),
+        col = rep.int(col_seq[has_any], times = counts[has_any])
+      )
     }
 
     # this creates pairs of row and column indices

bench/micro/upper_gather_bench.R

@@ -0,0 +1,104 @@
+# Micro benchmark for the upper_gather() dist vector path.
+#
+# Compares the implementation on this branch against the verbatim previous one
+# on the argument combination the pairwise distance functions actually use
+# (distinct = TRUE, so na.rm = FALSE and zero.rm = FALSE), for a range of n.
+#
+# Usage:
+#   Rscript bench/micro/upper_gather_bench.R [path/to/old_impl.R] [n1,n2,...]
+#
+# See upper_gather_equivalence.R for how to produce the old implementation file.
+# Peak memory is gc()'s "max used" across the call, in MB, which counts R heap
+# only. Timings are elapsed seconds from system.time().
+
+args <- commandArgs(trailingOnly = TRUE)
+OLD_IMPL <- if (length(args) >= 1) args[1] else "/tmp/upper_gather_old_impl.R"
+NS <- if (length(args) >= 2) as.integer(strsplit(args[2], ",")[[1]]) else c(500L, 1000L, 2000L, 4000L)
+
+suppressMessages(library(Matrix))
+suppressMessages(library(magrittr))
+source(OLD_IMPL)
+source("R/util.R", local = FALSE)
+
+measure <- function(fn, ...) {
+  invisible(gc(reset = TRUE, full = TRUE))
+  t <- system.time(res <- fn(...))[["elapsed"]]
+  g <- gc(full = TRUE)
+  # "max used (Mb)" is always the LAST column of gc(), but its INDEX is not
+  # fixed: gc() adds a "limit (Mb)" column when a memory limit is set (see
+  # mem.maxVSize / R_MAX_VSIZE), which shifts column 6 from "max used (Mb)" to
+  # "max used" in cells. Index from the end so both shapes report Mb.
+  peak <- sum(g[, ncol(g)])
+  list(sec = t, peak_mb = peak, rows = nrow(res))
+}
+
+cat(sprintf("%6s %6s %10s %10s %10s %10s %12s %12s %7s\n",
+            "n", "dist_s", "old_s", "new_s", "old_peakMB", "new_peakMB",
+            "old_rows", "new_rows", "speedup"))
+
+for (n in NS) {
+  set.seed(1)
+  m <- matrix(stats::runif(n * 5), nrow = n)
+  rownames(m) <- paste0("r", seq_len(n))
+
+  invisible(gc(reset = TRUE, full = TRUE))
+  t_dist <- system.time(d <- as.vector(stats::dist(m)))[["elapsed"]]
+
+  old <- try(measure(upper_gather_old, d, rownames(m), diag = NULL,
+                     cnames = c("Var1", "Var2", "value"), na.rm = FALSE, zero.rm = FALSE),
+             silent = TRUE)
+  if (inherits(old, "try-error")) {
+    cat(sprintf("%6d %6.3f  OLD FAILED: %s\n", n, t_dist,
+                conditionMessage(attr(old, "condition"))))
+    old <- list(sec = NA_real_, peak_mb = NA_real_, rows = NA_integer_)
+  }
+
+  new <- try(measure(upper_gather, d, rownames(m), diag = NULL,
+                     cnames = c("Var1", "Var2", "value"), na.rm = FALSE, zero.rm = FALSE),
+             silent = TRUE)
+  if (inherits(new, "try-error")) {
+    cat(sprintf("%6d %6.3f  NEW FAILED: %s\n", n, t_dist,
+                conditionMessage(attr(new, "condition"))))
+    next
+  }
+
+  cat(sprintf("%6d %6.3f %10.3f %10.3f %10.1f %10.1f %12s %12s %6.1fx\n",
+              n, t_dist, old$sec, new$sec, old$peak_mb, new$peak_mb,
+              format(old$rows, big.mark = ","), format(new$rows, big.mark = ","),
+              old$sec / new$sec))
+}
+
+# The matrix path with zero.rm = FALSE, which is what do_cosine_sim.kv uses.
+# A sparse input is the interesting case: the previous code forced the sparse
+# class to store all n^2 cells before calling which() on them.
+cat("\nmatrix path, sparse input, na.rm = FALSE / zero.rm = FALSE (do_cosine_sim.kv)\n")
+cat(sprintf("%6s %10s %10s %10s %10s %12s %12s %7s\n",
+            "n", "old_s", "new_s", "old_peakMB", "new_peakMB", "old_rows", "new_rows", "speedup"))
+for (n in NS) {
+  set.seed(2)
+  nnz <- max(1L, as.integer(n * n * 0.01))
+  sm <- Matrix::sparseMatrix(
+    i = sample.int(n, nnz, replace = TRUE),
+    j = sample.int(n, nnz, replace = TRUE),
+    x = stats::runif(nnz), dims = c(n, n),
+    dimnames = list(paste0("r", seq_len(n)), paste0("r", seq_len(n))))
+
+  old <- try(measure(upper_gather_old, sm, diag = FALSE,
+                     cnames = c("Var1", "Var2", "value"), na.rm = FALSE, zero.rm = FALSE),
+             silent = TRUE)
+  if (inherits(old, "try-error")) {
+    cat(sprintf("%6d  OLD FAILED: %s\n", n, conditionMessage(attr(old, "condition"))))
+    old <- list(sec = NA_real_, peak_mb = NA_real_, rows = NA_integer_)
+  }
+  new <- try(measure(upper_gather, sm, diag = FALSE,
+                     cnames = c("Var1", "Var2", "value"), na.rm = FALSE, zero.rm = FALSE),
+             silent = TRUE)
+  if (inherits(new, "try-error")) {
+    cat(sprintf("%6d  NEW FAILED: %s\n", n, conditionMessage(attr(new, "condition"))))
+    next
+  }
+  cat(sprintf("%6d %10.3f %10.3f %10.1f %10.1f %12s %12s %6.1fx\n",
+              n, old$sec, new$sec, old$peak_mb, new$peak_mb,
+              format(old$rows, big.mark = ","), format(new$rows, big.mark = ","),
+              old$sec / new$sec))
+}