Introduction

When learning R, we are always told to vectorise, vectorise, vectorise. This is because R is an interpreted language. Every time R runs a line of R code, it has to look up each name, match the arguments of each function call and set up a new environment for each function it enters. A vectorised function such as *, cumsum() or match() does that once and then loops over the elements in compiled code, where each step is very cheap. A for loop written in R does it on every iteration.

set.seed(42)
x <- runif(1e5)
loop <- function(x) {
  out <- numeric(length(x))
  for (i in seq_along(x)) out[i] <- x[i] * 2
  out
}
bench::mark(loop = loop(x), vectorised = x * 2)[, c("expression", "median", "mem_alloc")]
# # A tibble: 2 x 3
#   expression   median mem_alloc
#   <bch:expr> <bch:tm> <bch:byt>
# 1 loop         2.22ms     819KB
# 2 vectorised  59.02us     781KB

The loop gives the same answer and uses the same memory, but it is nearly 40 times slower. Bear in mind that lapply(), vapply() and friends are loops too. They tidy the loop away, but each iteration still calls an R function, so what matters is how many times we go round and how much work we do each time.

Grouped operations make this advice tricky to follow. mutate(df, y = x - mean(x)) on grouped data means “evaluate x - mean(x) for each group”, and since the expression can be any R code at all, there is no way around evaluating it once per group. What we can do is make everything around that loop vectorised, and keep the work inside it to a minimum.

{poorman} has not always done this. When I first wrote about grouped operations in {poorman}, they used the classic split-apply-combine approach: split the data.frame into one data.frame per group, run the verb on each piece and rbind() the results back together. It was simple, but it was slow once there were more than a few hundred groups. I have since rewritten the grouped verbs, and in this post I want to share the tricks that made the difference. All of them are plain base R, so you can use them in your own code too.

Split-Apply-Combine Versus Vectorised

To keep things concrete, let’s centre a column within groups, which is what mutate(x = x - mean(x)) does on grouped data. The split-apply-combine version is how most of us would first write it.

centre_split <- function(df) {
  df$row <- seq_len(nrow(df))
  pieces <- split(df, df$g)
  pieces <- lapply(pieces, function(piece) {
    piece$centred <- piece$x - mean(piece$x)
    piece
  })
  res <- do.call(rbind, pieces)
  res$centred[order(res$row)]
}

And here is the version {poorman} has ended up with, stripped down to the bare bones.

centre_fast <- function(df) {
  ids <- match(df$g, unique(df$g))
  f <- structure(ids, levels = as.character(seq_len(max(ids))), class = "factor")
  pieces <- split.default(df$x, f)
  res <- unlist(lapply(pieces, function(x) x - mean(x)), use.names = FALSE)
  rows <- unlist(split.default(seq_along(ids), f), use.names = FALSE)
  back <- integer(length(rows))
  back[rows] <- seq_along(rows)
  res[back]
}

Both still have an lapply() over the groups. The difference is what else happens. centre_split() builds a data.frame for every group, adds a column to each one and then binds them all back together. centre_fast() only splits the one column it needs, does nothing in the loop except the calculation itself, and combines the results with a single unlist().

To time them, I use bench::press(), which runs bench::mark() once for each value of n_groups, on 100,000 rows. bench::mark() also checks that both functions return the same result.

res <- bench::press(
  n_groups = c(10, 1e3, 1e4),
  {
    set.seed(42)
    df <- data.frame(g = sample(n_groups, 1e5, replace = TRUE), x = runif(1e5))
    bench::mark(split = centre_split(df), fast = centre_fast(df), iterations = 10L)
  },
  .quiet = TRUE
)
res[, c("expression", "n_groups", "median", "mem_alloc")]
# # A tibble: 6 x 4
#   expression n_groups   median mem_alloc
#   <bch:expr>    <dbl> <bch:tm> <bch:byt>
# 1 split            10  29.62ms    26.1MB
# 2 fast             10   3.76ms    7.57MB
# 3 split          1000  86.71ms   61.97MB
# 4 fast           1000   6.23ms    7.67MB
# 5 split         10000 522.85ms   55.61MB
# 6 fast          10000  26.27ms    6.18MB

centre_fast() is about 8 times faster with 10 groups, 14 times faster with 1,000 groups and 20 times faster with 10,000 groups, and it uses far less memory. Let’s go through the tricks one at a time. For the rest of the post, I’ll use this data for the timings, with 1,000 groups:

set.seed(42)
n <- 1e5
df <- data.frame(g = sample(1e3, n, replace = TRUE), x = runif(n))
g <- df$g

and this small dataset to show what each trick does.

sales <- data.frame(
  shop = c("b", "a", "b", "c", "a", "b"),
  day = c(1L, 1L, 2L, 1L, 1L, 2L),
  amount = c(2, 4, 6, 8, 10, 12)
)
sales
#   shop day amount
# 1    b   1      2
# 2    a   1      4
# 3    b   2      6
# 4    c   1      8
# 5    a   1     10
# 6    b   2     12

Tip 1: Number the Groups With match()

The first job is to work out which rows belong to which group. The obvious way is to compare every row with every group.

lapply(unique(sales$shop), function(shop) which(sales$shop == shop))
# [[1]]
# [1] 1 3 6
# 
# [[2]]
# [1] 2 5
# 
# [[3]]
# [1] 4

That is n * G comparisons for n rows and G groups, and each comparison allocates a new logical vector the length of the data. It is much faster to give each row a group number in one go with match().

match(sales$shop, unique(sales$shop))
# [1] 1 2 1 3 2 1

Each value gets the position where it first appears in unique(sales$shop), so the groups are numbered in the order they appear. Both unique() and match() hand their work to C code inside R itself (you can see the call to .Internal() if you print match), and that code uses a hash table. A hash table is a lookup table where each value is turned into a number, its hash, which says where in the table the value is stored. Looking a value up then means working out its hash and checking that one place, rather than searching the whole table, so it takes about the same time however many values are stored. unique() goes through the vector once and adds each value it hasn’t seen before to a table; match() puts the values of unique(sales$shop) in a table and looks up each element of sales$shop in it. Either way, each row is looked at once, whatever the number of groups.

bench::mark(
  compare = lapply(unique(g), function(value) which(g == value)),
  match = match(g, unique(g)),
  check = FALSE
)[, c("expression", "median", "mem_alloc")]
# # A tibble: 2 x 3
#   expression   median mem_alloc
#   <bch:expr> <bch:tm> <bch:byt>
# 1 compare       162ms  764.87MB
# 2 match         653us    2.16MB

(check = FALSE is needed here because the two give their answers in different forms; Tip 2 turns the group numbers into the rows of each group.)

With more than one grouping column, we need a number for each combination of values. The original version of {poorman} used base R’s interaction() for this, which builds a label for each row by pasting the values together with a ".". Unfortunately, pasting can give two different groups the same label.

interaction(c("a.b", "a"), c("c", "b.c"))
# [1] a.b.c a.b.c
# Levels: a.b.c a.b.b.c a.c

("a.b", "c") and ("a", "b.c") are different groups, but both become a.b.c. interaction() also makes a level for every possible combination of values, whether or not it appears in the data, which can use a lot of memory with several columns. Instead, I number each column with match() and combine the numbers with a little arithmetic, a bit like turning two digits into one number.

shop_id <- match(sales$shop, unique(sales$shop))
day_id <- match(sales$day, unique(sales$day))
combined <- (shop_id - 1) * max(day_id) + day_id
combined
# [1] 1 3 2 5 3 2
match(combined, unique(combined))
# [1] 1 2 3 4 2 3

Each (shop, day) pair gets its own number, so no two groups can collide. Calling match() again on the combined number brings the ids back down to 1, 2, ..., G, so they stay small when there are more columns to add. The - 1 is a double, so the arithmetic is done in doubles and can’t overflow R’s integers.

Tip 2: Build Factors by Hand

Given the group ids, split() gives us the rows of each group. But split() wants a factor, and calling factor(ids) would sort the unique values and match() every value against them again: work we have already done. The ids are already the integer codes 1:G, so I attach the levels and class attributes directly with structure() and have a valid factor at no cost.

ids <- match(combined, unique(combined))
f <- structure(ids, levels = as.character(seq_len(max(ids))), class = "factor")
f
# [1] 1 2 3 4 2 3
# Levels: 1 2 3 4
split.default(seq_along(ids), f)
# $`1`
# [1] 1
# 
# $`2`
# [1] 2 5
# 
# $`3`
# [1] 3 6
# 
# $`4`
# [1] 4

On 100,000 rows with 1,000 groups, this is 4 times faster than going through factor().

ids <- match(g, unique(g))
bench::mark(
  factor = split(seq_along(ids), factor(ids)),
  structure = split(seq_along(ids), structure(ids, levels = as.character(seq_len(max(ids))), class = "factor"))
)[, c("expression", "median", "mem_alloc")]
# # A tibble: 2 x 3
#   expression   median mem_alloc
#   <bch:expr> <bch:tm> <bch:byt>
# 1 factor       5.63ms     4.2MB
# 2 structure    1.28ms     1.2MB

This only works because the ids are whole numbers from 1 to G with no gaps, which match() guarantees.

Tip 3: Split Columns, Not data.frames

split() on a data.frame calls [.data.frame once for each group. [.data.frame is a long R function that has to handle every way of indexing a data.frame and checks the row names every time. Splitting each column on its own avoids all of that.

f <- structure(ids, levels = as.character(seq_len(max(ids))), class = "factor")
bench::mark(
  data_frame = split(df, f),
  columns = lapply(df, split.default, f),
  check = FALSE
)[, c("expression", "median", "mem_alloc")]
# # A tibble: 2 x 3
#   expression   median mem_alloc
#   <bch:expr> <bch:tm> <bch:byt>
# 1 data_frame  13.01ms     3.9MB
# 2 columns      2.04ms    2.02MB

There is a second choice to make here. If we already have the rows of each group, we can also split a column by indexing it once per group, lapply(rows, function(i) x[i]). Which is faster depends on the size of the groups.

res <- bench::press(
  n_groups = c(100, 1e3, 3e3, 1e4),
  {
    set.seed(42)
    x <- runif(1e5)
    ids <- match(sample(n_groups, 1e5, replace = TRUE), seq_len(n_groups))
    f <- structure(ids, levels = as.character(seq_len(n_groups)), class = "factor")
    rows <- split.default(seq_along(x), f)
    bench::mark(index = lapply(rows, function(i) x[i]), split = split.default(x, f), check = FALSE)
  },
  .quiet = TRUE
)
res[, c("expression", "n_groups", "median", "mem_alloc")]
# # A tibble: 8 x 4
#   expression n_groups   median mem_alloc
#   <bch:expr>    <dbl> <bch:tm> <bch:byt>
# 1 index           100 232.35us  786.77KB
# 2 split           100 964.07us    1.15MB
# 3 index          1000 526.28us  835.98KB
# 4 split          1000   1.02ms     1.2MB
# 5 index          3000    1.1ms  944.86KB
# 6 split          3000    1.1ms    1.32MB
# 7 index         10000    2.9ms  131.44KB
# 8 split         10000   1.18ms  561.22KB

split.default() costs about 1 ms whatever the number of groups, because it always goes through the whole column. Indexing costs one R function call per group, so it is cheap with a few large groups and expensive with many small ones. The two meet at about 3,000 groups, which is about 33 rows per group, so {poorman} indexes when the groups have 32 rows or more on average and uses split.default() otherwise.

Tip 4: Combine With One unlist()

do.call(rbind, pieces) has the same problem as splitting a data.frame, but in reverse. rbind.data.frame() has to cope with columns in a different order, factors with different levels, duplicated row names and more, and it does those checks for every piece. If each group gives a plain vector, unlist() combines them in one call.

pieces_df <- lapply(split(df, f), function(piece) data.frame(m = mean(piece$x)))
pieces_vec <- lapply(split.default(df$x, f), mean)
bench::mark(
  rbind = do.call(rbind, pieces_df)$m,
  unlist = unlist(pieces_vec, use.names = FALSE)
)[, c("expression", "median", "mem_alloc")]
# # A tibble: 2 x 3
#   expression   median mem_alloc
#   <bch:expr> <bch:tm> <bch:byt>
# 1 rbind        7.11ms  267.74KB
# 2 unlist       7.38us    7.86KB

With 1,000 one row pieces, unlist() is about 1,000 times faster. Of course unlist() drops attributes, so factors and dates need a little extra care, but for plain numbers and characters this is all it takes.

Tip 5: Put the Rows Back With One Permutation

After the combine, the results are in group order: all of group 1, then all of group 2 and so on. mutate() has to return them in the original row order. unlist(rows) tells us which row each result belongs to, and writing seq_along(rows) into those positions turns that around, giving the position of each row’s result.

rows <- unlist(split.default(seq_along(shop_id), shop_id), use.names = FALSE)
rows
# [1] 1 3 6 2 5 4
back <- integer(length(rows))
back[rows] <- seq_along(rows)
back
# [1] 1 4 2 6 5 3

match() numbered shop b first, so group order is the three b rows, then the two a rows, then c. Row 1 is shop b, so its result is 1st in group order; row 2 is shop a, so its result is 4th, after the three results for b; and so on. Now any result can be put back with a single subset.

totals <- unlist(lapply(split(sales$amount, shop_id), sum), use.names = FALSE)
by_group <- rep(totals, tabulate(shop_id))
by_group[back]
# [1] 20 14 20  8 14 20

back only has to be worked out once, and since it is used with [, it works for every new column whatever its type.

Tip 6: Remove the Loop When You Can

Some grouped calculations do not need a loop over the groups at all. If you know the function in advance, base R often has a grouped version already. The mean from our example can be done with rowsum(), which sums each group in C.

centre_rowsum <- function(df) {
  ids <- match(df$g, unique(df$g))
  df$x - (rowsum(df$x, ids, reorder = FALSE) / tabulate(ids))[ids]
}
res <- bench::press(
  n_groups = c(10, 1e3, 1e4),
  {
    set.seed(42)
    df <- data.frame(g = sample(n_groups, 1e5, replace = TRUE), x = runif(1e5))
    bench::mark(fast = centre_fast(df), rowsum = centre_rowsum(df), iterations = 10L)
  },
  .quiet = TRUE
)
res[, c("expression", "n_groups", "median", "mem_alloc")]
# # A tibble: 6 x 4
#   expression n_groups   median mem_alloc
#   <bch:expr>    <dbl> <bch:tm> <bch:byt>
# 1 fast             10   3.81ms    7.49MB
# 2 rowsum           10   1.64ms    4.68MB
# 3 fast           1000   5.55ms    7.67MB
# 4 rowsum         1000   1.64ms    4.71MB
# 5 fast          10000  22.65ms    6.18MB
# 6 rowsum        10000   6.59ms    5.15MB

{poorman} can’t do this for mutate(), because it doesn’t know what the expression will be, but it can for verbs that don’t take an expression to evaluate for each group.

My favourite example is filling missing values, as in fill(). Filling down means replacing each NA with the last value before it that isn’t missing. Take the position of each element, set it to 0 where the value is missing, and take the running maximum with cummax(). That gives the position of the last value that isn’t missing.

x <- c(NA, 3, NA, NA, 5, NA)
seq_along(x) * !is.na(x)
# [1] 0 2 0 0 5 0
cummax(seq_along(x) * !is.na(x))
# [1] 0 2 2 2 5 5

The first element has no value before it, so it gets 0. To do this within groups, put the rows in group order and always count the first element of each group, missing or not. That fixes the 0, and it also stops a value leaking into the next group.

first <- c(TRUE, FALSE, FALSE, TRUE, FALSE, FALSE)  # two groups: positions 1-3 and 4-6
from <- cummax(seq_along(x) * (!is.na(x) | first))
from
# [1] 1 2 2 4 5 5
x[from]
# [1] NA  3  3 NA  5  5

Position 4 stays NA rather than taking the 3 from the first group. Here is a version that splits the data and uses the trick inside each group, and a version that uses it once on the whole column.

fill_one <- function(x) {
  pos <- seq_along(x)
  x[cummax(pos * (!is.na(x) | pos == 1L))]
}
fill_split <- function(x, ids) {
  f <- structure(ids, levels = as.character(seq_len(max(ids))), class = "factor")
  unsplit(lapply(split(x, f), fill_one), f)
}
fill_whole <- function(x, ids) {
  idx <- order(ids)
  sorted <- ids[idx]
  first <- c(TRUE, sorted[-1L] != sorted[-length(sorted)])
  y <- x[idx]
  out <- x
  out[idx] <- y[cummax(seq_along(y) * (!is.na(y) | first))]
  out
}
res <- bench::press(
  n_groups = c(10, 1e3, 1e4),
  {
    set.seed(42)
    ids <- match(sample(n_groups, 1e5, replace = TRUE), seq_len(n_groups))
    z <- ifelse(runif(1e5) < 0.3, NA, runif(1e5))
    bench::mark(split = fill_split(z, ids), whole = fill_whole(z, ids), iterations = 10L)
  },
  .quiet = TRUE
)
res[, c("expression", "n_groups", "median", "mem_alloc")]
# # A tibble: 6 x 4
#   expression n_groups   median mem_alloc
#   <bch:expr>    <dbl> <bch:tm> <bch:byt>
# 1 split            10   3.64ms    7.69MB
# 2 whole            10   1.96ms    9.16MB
# 3 split          1000   4.91ms    8.13MB
# 4 whole          1000   2.07ms    9.16MB
# 5 split         10000  11.53ms    3.45MB
# 6 whole         10000   2.09ms    9.16MB

Both functions use the same cummax() trick, so the only difference is the loop over the groups. With 10 groups that loop costs little, but with 10,000 groups fill_whole() is 5 times faster. It uses a little more memory, for the reordered copy of the column.

How {poorman} Puts It Together

So far these have been small base R functions. Here is how {poorman} uses the same tricks in its verbs.

library(poorman)
gs <- group_by(sales, shop)
group_data(gs)
#   shop   .rows
# 1    a    2, 5
# 2    b 1, 3, 6
# 3    c       4

expr <- quote(amount / sum(amount))
rows <- group_data(gs)$.rows
pieces <- lapply(rows, function(i) sales$amount[i])
envs <- lapply(pieces, function(amount) list2env(list(amount = amount), envir = new.env(parent = globalenv())))
results <- lapply(envs, function(env) eval(expr, env))
results
# [[1]]
# [1] 0.2857143 0.7142857
# 
# [[2]]
# [1] 0.1 0.3 0.6
# 
# [[3]]
# [1] 1

idx <- unlist(rows)
back <- integer(length(idx))
back[idx] <- seq_along(idx)
unlist(results)[back]
# [1] 0.1000000 0.2857143 0.3000000 1.0000000 0.7142857 0.6000000
mutate(gs, share = amount / sum(amount))$share
# [1] 0.1000000 0.2857143 0.3000000 1.0000000 0.7142857 0.6000000

var <- "amount"
mutate(gs, double = .data[[var]] * 2)
#   shop day amount double
# 1    b   1      2      4
# 2    a   1      4      8
# 3    b   2      6     12
# 4    c   1      8     16
# 5    a   1     10     20
# 6    b   2     12     24

env <- list2env(list(amount = c(4, 10)), envir = new.env())
makeActiveBinding(".data", function() {
  cat("Building .data\n")
  data.frame(amount = env$amount)
}, env)
eval(quote(sum(amount)), env)
# [1] 14
eval(quote(sum(.data[["amount"]])), env)
# Building .data
# [1] 14

keep <- sales$amount > 4
keep
# [1] FALSE FALSE  TRUE  TRUE  TRUE  TRUE
cumsum(keep)
# [1] 0 0 1 2 3 4
group_data(filter(gs, amount > 4))
#   shop .rows
# 1    a     3
# 2    b  1, 4
# 3    c     2

which(!duplicated(ids))
# [1] 1 2 3 4
distinct(gs, day)
#   shop day
# 1    b   1
# 2    a   1
# 3    b   2
# 4    c   1

unlist(lapply(group_data(gs)$.rows, head, 2))
# [1] 2 5 1 3 4
slice_head(gs, n = 2)
#   shop day amount
# 2    a   1      4
# 5    a   1     10
# 1    b   1      2
# 3    b   2      6
# 4    c   1      8

The Results

Here is the code I used to compare {poorman} 0.3.0, the version with these changes, with the current CRAN release, 0.2.8. I ran it twice, in separate R sessions, once with each version installed (install.packages("poorman") for 0.2.8 and remotes::install_github("nathaneastwood/poorman") for 0.3.0). group_by() times making the grouped data; every other line times a grouped operation on gdf, the data grouped by g.

library(poorman)
res <- bench::press(
  n_groups = c(10, 1e3, 1e4),
  {
    set.seed(42)
    n <- 1e5
    df <- data.frame(
      g = sample(n_groups, n, replace = TRUE),
      x = runif(n),
      y = sample(5L, n, replace = TRUE)
    )
    df$z <- ifelse(runif(n) < 0.3, NA, df$x)
    gdf <- group_by(df, g)
    bench::mark(
      group_by = group_by(df, g),
      summarise = summarise(gdf, m = mean(x)),
      mutate = mutate(gdf, m = x - mean(x)),
      filter = filter(gdf, x > mean(x)),
      slice_head = slice_head(gdf, n = 2),
      distinct = distinct(gdf, y),
      fill = fill(gdf, z),
      iterations = 3L,
      check = FALSE
    )
  },
  .quiet = TRUE
)
res[, c("expression", "n_groups", "median", "mem_alloc")]

0.3.0 is faster for every verb and every number of groups: between 4 and 22 times faster with 10 groups, and between 21 and about 3,800 times faster with 10,000 groups. fill(), which no longer loops over groups, went from 12.4 seconds to 3.3 ms. Most of the old time and memory went on comparing every row with every group, as in Tip 1: 0.2.8 did this in group_by(), and every grouped verb did it again to find the groups of its result. That is why the old times grow so quickly with the number of groups, and why mutate() with 10,000 groups allocated 18.7 GB. It now allocates 11.6 MB.

None of this needed anything outside of base R. The loop over the groups is still there where it has to be, but it is now the only part of the code that runs once per group.

{poorman} 0.3.0 will be on CRAN soon, and then you will be able to install it with

install.packages("poorman")

Until then, you can install it from GitHub with

remotes::install_github("nathaneastwood/poorman")

As always, please report any bugs or differences from {dplyr} on GitHub.