Performance of fill() with labelled data after group_by()

[etiennebacher]

This problem is very similar to tidyverse/tidyr#520. When there is a large number of groups, fill() is much slower with labelled data than with only numeric data.

library(dplyr, warn.conflicts = FALSE)
library(tidyr)
library(haven)

set.seed(2)
n <- 1e4
my_sample <- sample(c(1:10, NA), n, replace = TRUE)
df <- tibble(
  group = sample(paste("id", 1:(n/4)), n, replace = TRUE),
  num = my_sample,
  lab = haven::labelled(my_sample)
) %>% 
  group_by(group)

bench::mark(
  num = fill(df, num, .direction = "updown"),
  lab = fill(df, lab, .direction = "updown"),
  check = FALSE
)[1:4]
#> # A tibble: 2 x 4
#>   expression      min   median `itr/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl>
#> 1 num         37.09ms  41.65ms    23.3  
#> 2 lab           1.99s    1.99s     0.502
#> Warning message:
#> Some expressions had a GC in every iteration; so filtering is disabled. 

Note that the timing is similar when the data is not grouped:

set.seed(2)
n <- 1e4
my_sample <- sample(c(1:10, NA), n, replace = TRUE)
df <- tibble(
  num = my_sample,
  lab = haven::labelled(my_sample)
) 

bench::mark(
  num = fill(df, num, .direction = "updown"),
  lab = fill(df, lab, .direction = "updown"),
  check = FALSE
)[1:4]
#> # A tibble: 2 x 4
#>   expression      min   median `itr/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl>
#> 1 num          22.4ms   24.4ms      39.4
#> 2 lab          26.1ms   32.9ms      29.7

---

Done with development versions of tidyr, dplyr and haven.

[DavisVaughan]

It is expected that custom types like the labelled type from haven will run slower in general vs the base R types like integer/double. That is unfortunately just a restriction that is hard to overcome (but hey at least fill() works here now!) because it has to work out if it can combine the result from each group together into a new labelled object, and that takes time.

This isn't limited to fill(). Any grouped operation with this column that returns another labelled object will be slower:

library(dplyr, warn.conflicts = FALSE)
library(tidyr)
library(haven)

set.seed(2)
n <- 1e4
my_sample <- sample(c(1:10, NA), n, replace = TRUE)

df <- tibble(
  group = sample(paste("id", 1:(n/4)), n, replace = TRUE),
  num = my_sample,
  lab = haven::labelled(my_sample)
) %>% 
  group_by(group)

# mine seems to run a little faster
bench::mark(
  num = fill(df, num, .direction = "updown"),
  lab = fill(df, lab, .direction = "updown"),
  check = FALSE
)
#> Warning: Some expressions had a GC in every iteration; so filtering is disabled.
#> # A tibble: 2 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 num         59.59ms  61.66ms    16.1      2.67MB     23.3
#> 2 lab           1.18s    1.18s     0.848  416.91KB     25.4

# but here is a mutate(), which also runs relatively slow
bench::mark(mutate(df, lab2 = haven::zap_label(lab)))
#> Warning: Some expressions had a GC in every iteration; so filtering is disabled.
#> # A tibble: 1 × 6
#>   expression                                    min   median `itr/sec` mem_alloc
#>   <bch:expr>                               <bch:tm> <bch:tm>     <dbl> <bch:byt>
#> 1 mutate(df, lab2 = haven::zap_label(lab))    918ms    918ms      1.09     350KB
#> # … with 1 more variable: `gc/sec` <dbl>

@hadley idk if the ptype2 / cast methods for labelled can be optimized at all, but they seem to do a lot of work and that seems to be where most of the time is spent (i.e. determining the common type of the group results, and then casting each group result to that type).

haven:::vec_ptype2.haven_labelled.haven_labelled
#> function (x, y, ..., x_arg = "", y_arg = "") 
#> {
#>     data_type <- vec_ptype2(vec_data(x), vec_data(y), ..., x_arg = x_arg, 
#>         y_arg = y_arg)
#>     x_labels <- vec_cast_named(attr(x, "labels"), data_type, 
#>         x_arg = x_arg)
#>     y_labels <- vec_cast_named(attr(y, "labels"), data_type, 
#>         x_arg = y_arg)
#>     labels <- c(x_labels, y_labels[!y_labels %in% x_labels])
#>     label <- attr(x, "label", exact = TRUE) %||% attr(y, "label", 
#>         exact = TRUE)
#>     new_labelled(data_type, labels = labels, label = label)
#> }
#> <bytecode: 0x7fa351ce48a8>
#> <environment: namespace:haven>

haven:::vec_cast.haven_labelled.haven_labelled
#> function (x, to, ..., x_arg = "", to_arg = "") 
#> {
#>     out_data <- vec_cast(vec_data(x), vec_data(to), ..., x_arg = x_arg, 
#>         to_arg = to_arg)
#>     x_labels <- attr(x, "labels")
#>     to_labels <- attr(to, "labels")
#>     out_labels <- to_labels %||% x_labels
#>     out <- labelled(out_data, labels = out_labels, label = attr(x, 
#>         "label", exact = TRUE))
#>     if (is.double(x) && !is.double(out)) {
#>         lossy <- is_tagged_na(x)
#>         maybe_lossy_cast(out, x, to, lossy, x_arg = x_arg, to_arg = to_arg, 
#>             details = "Only doubles can hold tagged na values.")
#>     }
#>     if (!is.null(to_labels)) {
#>         lossy <- x %in% x_labels[!x_labels %in% out_labels]
#>         maybe_lossy_cast(out, x, to, lossy, x_arg = x_arg, to_arg = to_arg, 
#>             details = paste0("Values are labelled in `", x_arg, 
#>                 "` but not in `", to_arg, "`."))
#>     }
#>     out
#> }
#> <bytecode: 0x7fa351d602f0>
#> <environment: namespace:haven>

[hadley]

Transferred to haven, since it seems like that where the fix will need to happen.

[gorcha]

We could avoid unnecessary casts, label processing etc. by checking if the input vectors are identical before performing any further processing:

vec_ptype2.haven_labelled.haven_labelled <- function(x, y, ..., x_arg = "", y_arg = "") {
  # Don't preform any further processing if the input vectors have matching
  # metadata
  if (identical(attributes(x), attributes(y)))
    return(vec_slice(x, 0L))

  data_type <- vec_ptype2(vec_data(x), vec_data(y), ..., x_arg = x_arg, y_arg = y_arg)

  # Prefer labels from LHS
  x_labels <- vec_cast_named(attr(x, "labels"), data_type, x_arg = x_arg)
  y_labels <- vec_cast_named(attr(y, "labels"), data_type, x_arg = y_arg)
  labels <- c(x_labels, y_labels[!y_labels %in% x_labels])

  # Prefer labels from LHS
  label <- attr(x, "label", exact = TRUE) %||% attr(y, "label", exact = TRUE)

  new_labelled(data_type, labels = labels, label = label)
}

This seems like the minimum amount of processing in these cases unless there are some vctrs tricks we can take advantage of. It decreases the time by about a third on my machine. What do you think @hadley?

[gorcha]

For reference, timings with the current implementation are:

bench::mark(
  num = fill(df, num, .direction = "updown"),
  lab = fill(df, lab, .direction = "updown"),
  check = FALSE
)
#> Warning: Some expressions had a GC in every iteration; so filtering is disabled.
#> # A tibble: 2 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 num            48ms   50.6ms    19.7      2.77MB     31.6
#> 2 lab              1s       1s     0.998  443.66KB     31.9

And with the identical check:

bench::mark(
  num = fill(df, num, .direction = "updown"),
  lab = fill(df, lab, .direction = "updown"),
  check = FALSE
)
#> Warning: Some expressions had a GC in every iteration; so filtering is disabled.
#> # A tibble: 2 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 num          47.4ms   50.6ms     19.7     2.77MB     29.5
#> 2 lab         606.1ms  606.1ms      1.65  448.85KB     36.3

[DavisVaughan]

I think you can avoid vec_slice(x, 0L) too, and just return x. vctrs should ensure that x is a zero slice ahead of time

[hadley]

@gorcha yeah, that seems like a good start. Maybe benchmark with this to make it more clear what is the haven overhead?

set.seed(2)
n <- 1e4
my_sample <- sample(c(1:10, NA), n, replace = TRUE)
lab <- haven::labelled(my_sample)
pieces <- unname(split(lab, (seq_len(n) - 1) %/% 4))

bench::mark(labelled = vctrs::vec_c(!!!pieces))

And can we take the same shortcut in vec_cast.haven_labelled.haven_labelled()?

[gorcha]

Oh great, thanks @DavisVaughan!

@hadley with the vec_slice removed it looks like this is cutting the time roughly in half.

Current version:

set.seed(2)
n <- 1e4
my_sample <- sample(c(1:10, NA), n, replace = TRUE)
lab <- haven::labelled(my_sample)
pieces <- unname(split(lab, (seq_len(n) - 1) %/% 4))

bench::mark(labelled = vctrs::vec_c(!!!pieces))
#> Warning: Some expressions had a GC in every iteration; so filtering is disabled.
#> # A tibble: 1 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 labelled      432ms    435ms      2.30     233KB     23.0

And with the identical check:

set.seed(2)
n <- 1e4
my_sample <- sample(c(1:10, NA), n, replace = TRUE)
lab <- haven::labelled(my_sample)
pieces <- unname(split(lab, (seq_len(n) - 1) %/% 4))

bench::mark(labelled = vctrs::vec_c(!!!pieces))
#> Warning: Some expressions had a GC in every iteration; so filtering is disabled.
#> # A tibble: 1 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 labelled      196ms    196ms      5.05     234KB     26.9

[hadley]

Looks great. You'll do a PR?

[gorcha]

Yup shall do!

[hadley]

@DavisVaughan do you happen to know why the attributes get copied here?

lab <- haven::labelled(letters[1:3])
lobstr::sxp(lab)
lobstr::sxp(lab[1])
lobstr::sxp(lab[2:3])

If they weren't copied, the comparison would be even faster. Presumably that happens somewhere deep in R, so maybe it's not worth exploring.

[DavisVaughan]

I'm fairly certain that comes from vec_restore_default(). i.e. [ -> vec_slice which does vec_proxy -> do the slice -> vec_restore

I think a shallow copy of the attribute pair list object is probably unavoidable. It happens here:
https://github.com/r-lib/vctrs/blob/48ea794800a4752e2f1419535b4d7cc1238e3d44/src/proxy-restore.c#L26

We do a lot of special tweaking to that attribute pair list to ensure it is restored correctly (i.e. making sure the OBJECT bit is set correctly, retaining names from x rather than from to) which makes a copy unavoidable.

It is only a shallow copy though, the individual attributes themselves are the same