Fix all_rows = TRUE to work with move to vctrs

Document behavior when `all_rows = TRUE` and `as_list_col = TRUE`; this behavior
might be different from before move to use `vctrs` though.

Add tests to cover this case + others, which was caught only via a vignette
build failing.

Author: lcbrooks

R/slide.R

@@ -62,8 +62,15 @@
 #'   when `as_list_col = FALSE`. Default is "_". Using `NULL` drops the prefix
 #'   from `new_col_name` entirely.
 #' @param all_rows If `all_rows = TRUE`, then all rows of `x` will be kept in
-#'   the output; otherwise, there will be one row for each time value in `x`
-#'   that acts as a reference time value. Default is `FALSE`.
+#'   the output even with `ref_time_values` provided, with some type of missing
+#'   value marker for the slide computation output column(s) for `time_value`s
+#'   outside `ref_time_values`; otherwise, there will be one row for each row in
+#'   `x` that had a `time_value` in `ref_time_values`. Default is `FALSE`. The
+#'   missing value marker is the result of `vctrs::vec_cast`ing `NA` to the type
+#'   of the slide computation output. If using `as_list_col = TRUE`, note that
+#'   the missing marker is a `NULL` entry in the list column; for certain
+#'   operations, you might want to replace these `NULL` entries with a different
+#'   `NA` marker.
 #' @return An `epi_df` object given by appending a new column to `x`, named
 #'   according to the `new_col_name` argument.
 #' 
@@ -306,10 +313,13 @@ epi_slide = function(x, f, ..., before, after, ref_time_values,
     # If all rows, then pad slide values with NAs, else filter down data group
     if (all_rows) {
       orig_values = slide_values
-      slide_values = rep(NA, nrow(.data_group))
-      slide_values[o] = orig_values
+      slide_values = vctrs::vec_rep(vctrs::vec_cast(NA, orig_values), nrow(.data_group))
+      # ^ using vctrs::vec_init would be shorter but docs don't guarantee it
+      # fills with NA equivalent.
+      vctrs::vec_slice(slide_values, o) = orig_values
+    } else {
+      .data_group = filter(.data_group, o)
     }
-    else .data_group = filter(.data_group, o)
     return(mutate(.data_group, !!new_col := slide_values))
   }
 

man/epi_slide.Rd

@@ -2,15 +2,23 @@
 
 d <- as.Date("2020-01-01")
 
-grouped = dplyr::bind_rows(
+ungrouped = dplyr::bind_rows(
   dplyr::tibble(geo_value = "ak", time_value = d + 1:200, value=1:200),
   dplyr::tibble(geo_value = "al", time_value = d + 1:5, value=-(1:5))
 ) %>%
-  as_epi_df() %>%
+  as_epi_df()
+grouped = ungrouped %>%
   group_by(geo_value)
-
 f = function(x, g) dplyr::tibble(value=mean(x$value), count=length(x$value))
 
+toy_edf = tibble::tribble(
+  ~geo_value, ~time_value, ~value    ,
+  "a"       , 1:10       , 2L^( 1:10),
+  "b"       , 1:10       , 2L^(11:20),
+  ) %>%
+  tidyr::unchop(c(time_value, value)) %>%
+  as_epi_df(as_of = 100)
+
 ## --- These cases generate errors (or not): ---
 test_that("`before` and `after` are both vectors of length 1", {
   expect_error(epi_slide(grouped, f, before = c(0,1), after = 0, ref_time_values = d+3),
@@ -88,13 +96,7 @@ test_that("these doesn't produce an error; the error appears only if the ref tim
 })
 
 test_that("computation output formats x as_list_col", {
-  toy_edf = tibble::tribble(
-    ~geo_value, ~time_value, ~value    ,
-    "a"       , 1:10       , 2L^( 1:10),
-    "b"       , 1:10       , 2L^(11:20),
-    ) %>%
-    tidyr::unchop(c(time_value, value)) %>%
-    as_epi_df(as_of = 100)
+  # See `toy_edf` definition at top of file.
   # We'll try 7d sum with a few formats.
   basic_result_from_size1 = tibble::tribble(
     ~geo_value, ~time_value, ~value    , ~slide_value                                    ,
@@ -170,3 +172,123 @@ test_that("epi_slide alerts if the provided f doesn't take enough args", {
   expect_warning(epi_slide(grouped, f_x_dots, before = 1L, ref_time_values = d+1),
     class = "epiprocess__assert_sufficient_f_args__mandatory_f_args_passed_to_f_dots")
 })
+
+test_that("`ref_time_values` + `all_rows = TRUE` works", {
+  # See `toy_edf` definition at top of file. We'll do variants of a slide
+  # returning the following:
+  basic_full_result = tibble::tribble(
+    ~geo_value, ~time_value, ~value    , ~slide_value                                    ,
+    "a"       , 1:10       , 2L^( 1:10), data.table::frollsum(2L^(1:10) + 2L^(11:20), c(1:7,rep(7L, 3L)), adaptive=TRUE, na.rm=TRUE),
+    "b"       , 1:10       , 2L^(11:20), data.table::frollsum(2L^(1:10) + 2L^(11:20), c(1:7,rep(7L, 3L)), adaptive=TRUE, na.rm=TRUE),
+    ) %>%
+    tidyr::unchop(c(time_value, value, slide_value)) %>%
+    dplyr::arrange(time_value) %>%
+    as_epi_df(as_of = 100)
+  # slide computations returning atomic vecs:
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ sum(.x$value)),
+    basic_full_result
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ sum(.x$value),
+                          ref_time_values = c(2L, 8L)),
+    basic_full_result %>% dplyr::filter(time_value %in% c(2L, 8L))
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ sum(.x$value),
+                          ref_time_values = c(2L, 8L), all_rows = TRUE),
+    basic_full_result %>%
+      dplyr::mutate(slide_value = dplyr::if_else(time_value %in% c(2L, 8L),
+                                                 slide_value, NA_integer_))
+  )
+  # slide computations returning data frames:
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value))),
+    basic_full_result %>% dplyr::rename(slide_value_value = slide_value)
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L)),
+    basic_full_result %>%
+      dplyr::filter(time_value %in% c(2L, 8L)) %>%
+      dplyr::rename(slide_value_value = slide_value)
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L), all_rows = TRUE),
+    basic_full_result %>%
+      dplyr::mutate(slide_value = dplyr::if_else(time_value %in% c(2L, 8L),
+                                                 slide_value, NA_integer_)) %>%
+      dplyr::rename(slide_value_value = slide_value)
+  )
+  # slide computations returning data frames with `as_list_col=TRUE`:
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          as_list_col = TRUE),
+    basic_full_result %>%
+      dplyr::mutate(slide_value = purrr::map(slide_value, ~ data.frame(value = .x)))
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L),
+                          as_list_col = TRUE),
+    basic_full_result %>%
+      dplyr::mutate(slide_value = purrr::map(slide_value, ~ data.frame(value = .x))) %>%
+      dplyr::filter(time_value %in% c(2L, 8L))
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L), all_rows = TRUE,
+                          as_list_col = TRUE),
+    basic_full_result %>%
+      dplyr::mutate(slide_value = purrr::map(slide_value, ~ data.frame(value = .x))) %>%
+      dplyr::mutate(slide_value = dplyr::if_else(time_value %in% c(2L, 8L),
+                                                 slide_value, list(NULL)))
+  )
+  # slide computations returning data frames, `as_list_col = TRUE`, `unnest`:
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          as_list_col = TRUE) %>%
+      unnest(slide_value, names_sep = "_"),
+    basic_full_result %>% dplyr::rename(slide_value_value = slide_value)
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L),
+                          as_list_col = TRUE) %>%
+      unnest(slide_value, names_sep = "_"),
+    basic_full_result %>%
+      dplyr::filter(time_value %in% c(2L, 8L)) %>%
+      dplyr::rename(slide_value_value = slide_value)
+  )
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L), all_rows = TRUE,
+                          as_list_col = TRUE) %>%
+      unnest(slide_value, names_sep = "_"),
+    basic_full_result %>%
+      # XXX unclear exactly what we want in this case. Current approach is
+      # compatible with `vctrs::vec_detect_missing` but breaks `tidyr::unnest`
+      # compatibility
+      dplyr::filter(time_value %in% c(2L, 8L)) %>%
+      dplyr::rename(slide_value_value = slide_value)
+  )
+  rework_nulls = function(slide_values_list) {
+    vctrs::vec_assign(
+      slide_values_list,
+      vctrs::vec_detect_missing(slide_values_list),
+      list(vctrs::vec_cast(NA, vctrs::vec_ptype_common(!!!slide_values_list)))
+    )
+  }
+  expect_identical(
+    toy_edf %>% epi_slide(before = 6L, ~ data.frame(value = sum(.x$value)),
+                          ref_time_values = c(2L, 8L), all_rows = TRUE,
+                          as_list_col = TRUE) %>%
+      mutate(slide_value = rework_nulls(slide_value)) %>%
+      unnest(slide_value, names_sep = "_"),
+    basic_full_result %>%
+      dplyr::mutate(slide_value = dplyr::if_else(time_value %in% c(2L, 8L),
+                                                 slide_value, NA_integer_)) %>%
+      dplyr::rename(slide_value_value = slide_value)
+  )
+})