feat: support duckplyr archive backend

[dshemetov]

This is an experimental PR to see how far I could get with a clear idea and an LLM tool - support duckplyr tibbles as a backend for epi archives. The end result seems worth sharing.

The commits log follow my process:

• create an abstraction layer between the internal storage $DT and the functions that operate on it - archive_accessors.R
• replace all call sites with this abstraction layer, creating thin wrappers for data.table operations
• introduce duckplyr alternatives and hook them in via S3 dispatch - archive_duck.R
• re-use the existing test suite by shadowing default data.table archive constructors with duckplyr constructors - helper-archive-backend.R

The full test suite passes with both backends! Very encouraging for duckplyr. It's also great that the original data.table backend is unaffected, so duckplyr can sit behind a feature flag.

Some commands (using just)

# Run test suite with old backend
just test-dt

# Run test suite with duckplyr backend
just test-duck

# Run both back to back
just test-backends

The shadowing is enabled via EPIPROCESS_TEST_ARCHIVE_BACKEND=duck, but currently it only shadows during tests. Some experiments with this backend:

# Working examples from https://cmu-delphi.github.io/epiprocess/reference/epix_slide.html?q=epix_slide#null
arch <- as_duckdb_archive(epidatasets::archive_cases_dv_subset_all_states)

requested_versions <-
  seq(as.Date("2020-09-02"), as.Date("2020-09-15"), by = "1 day")

arch %>%
  epix_slide(
    geowide_percent_cli_max_time = max(time_value[!is.na(percent_cli)]),
    geowide_percent_cli_rpt_lag = .version - geowide_percent_cli_max_time,
    .versions = requested_versions
  )

arch %>% 
  group_by(geo_value) %>%
    epix_slide(
      percent_cli_max_time = max(time_value[!is.na(percent_cli)]),
      percent_cli_rpt_lag = .version - percent_cli_max_time,
      .versions = requested_versions
    )

arch2 <- as_duckdb_archive(case_death_rate_archive)
arch2 %>%
  epix_slide(
    .versions = as.Date(c("2021-10-01", "2021-10-08")),
    function(x, g, v) {
      epipredict::arx_forecaster(
        x,
        outcome = "death_rate",
        predictors = c("death_rate_7d_av", "case_rate_7d_av")
      )$predictions
    }
  )

# From epix_merge docs
s1 <- tibble::tibble(
  geo_value = c("ca", "ca", "ca"),
  time_value = as.Date(c("2024-08-01", "2024-08-02", "2024-08-03")),
  version = as.Date(c("2024-08-01", "2024-08-02", "2024-08-03")),
  signal1 = c(14, 11, 9)
)
# The s2 signal at August 1st gets revised from 3 to 5 on August 3rd
s2 <- tibble::tibble(
  geo_value = c("ca", "ca", "ca"),
  time_value = as.Date(c("2024-08-01", "2024-08-01", "2024-08-02")),
  version = as.Date(c("2024-08-02", "2024-08-03", "2024-08-03")),
  signal2 = c(3, 5, 2),
)
s1 <- s1 %>% as_duckdb_archive()
s2 <- s2 %>% as_duckdb_archive()
merged <- epix_merge(s1, s2, sync = "locf")
merged

Known issues:

• autoplot needs to be written less eagerly to help duckplyr
• calls to archive_col are expensive for duckplyr (materializes a column), so any places where it's mostly used for metadata or schema or inexpensive checks, those should be replaced with a more specific and cheaper accessor
• various documentation boilerplate things are borked
• docstrings, examples, vignettes still contain references to $DT

[brookslogan]

Thanks, exciting to see the duckplyr backend being built. I've only made it through DUCK.md and jumped to a few random parts. Need to actually sanity-check the implementation. I have a bundle of questions, posting now & revisiting another day.

[brookslogan]

suggestion: rename to archive_schema or archive_tbl_ptype

Oh, I misunderstood, these are all logical columns. Can we collect(slice(archive_data(x), integer())) or collect(filter(archive_data(x), FALSE)) to make it give actual type info as well?

[brookslogan]

question: can we just use vapply is.factor on archive_colmask instead?
issue: as_duckdb_tibble preserves factor columns for me, duckplyr 1.2.1. So hardcoded return FALSE seems wrong.

[brookslogan]

suggestion:

• rename archive_data to archive_data_lazy
• add archive_data_general_shallow, which
  • uses old shallow copy tibble for DTs (avoid full table copy)
  • uses archive_data_lazy for duckplyr (avoid collecting ducks)

[brookslogan]

musing: really wish we could eliminate this bare data.table backend-related stuff. But it seems that it does support / supported some things that duckplyr doesn't (list columns and quantile_pred columns, though latter is actually currently broken for data.tables as well; maybe factors though we have mixed results on those). It'd be nice to transition the data.table backend to dtplyr underneath, but I guess that should be handled in a separate PR.

[brookslogan]

suggestion: replace remaining usage with archive_data_general_shallow, proposed in a comment below

[brookslogan]

musing: this function has some duplicate boilerplate, and as_duckdb_epi_archive is probably copying the data multiple times.

It feels like it might be nicer to try to get toward only as_epi_archive (with alias as_dtbl_epi_archive) and as_duck_epi_archive, with the same functionality provided, involving:

• reviving optional geo_type and time_type args in as_epi_archive
• making compactify default to choose FALSE if we're not backed by RAM, TRUE otherwise

but this seems like substantial work, and refactoring data.table to lazy_dt might make this easier, or data.table backend elimination altogether might obviate this. Maybe not worth the time?

[brookslogan]

question: do epipredict tests still pass? I thought there was some reason we couldn't clean this up earlier.

[brookslogan]

Suggested change

	as_epi_archive_like <- function(template, x, ...) UseMethod("as_epi_archive_like")
	reconstruct_epi_archive <- function(x, template, ...) UseMethod("reconstruct_epi_archive", template)

and update name & arg order in usage

[brookslogan]

suggestion: use vec_slice(.x, NA)... I'm hoping eventually we can allow tibble columns. Like ~ .x[NA] this is the same size as .x; for size 1 use vec_slice(.x, NA_integer_) or maybe vec_cast(NA, .x).

[brookslogan]

issue?: this seems like it will be significantly slower for lazy_dt compared to old data.table approach, maybe by 2--3 times? Doing a little bit of experimentation, I can't find a dplyr-verb-based way to speed it up. And actually data.table rolling join approach now seems faster than the unique fromLast approach (before it seemed like figuring out the epikeytimes to query made it more expensive, but current tests seem to disagree / suggest that something changed).

I thought we were actually materially dissatisfied with epix_as_of performance in the past, so this seems like it may still be a problem. Though I'm not sure if the current benchmarks are that problematic or not. Is this one of those cases where using Dates slows everything down and if we convert to double and back it will be much faster? Or does grouping/splitting by geo_value slow things down?

lzdts <- list(
  tidyr::expand_grid(g = 1:3000, v = 1:100, l = 1:60) %>%
    mutate(t = v - l) %>%
    select(g, t, v, l) %>%
    arrange(g, t, v) %>%
    as.data.table(key = c("g", "t", "v")) %>%
    `[`(seq_len(nrow(.)) %% 7 != 0,) %>% # hack to make it actually have v != request v sometimes
    lazy_dt(),
  tidyr::expand_grid(g = 1:50, v = 1:300) %>%
    reframe(.by = c(g,v), l = 1:v) %>%
    mutate(t = v - l) %>%
    select(g, t, v, l) %>%
    arrange(g, t, v) %>%
    as.data.table(key = c("g", "t", "v")) %>%
    `[`(seq_len(nrow(.)) %% 7 != 0,) %>% # hack to make it actually have v != request v sometimes
    `[`(v <= max(v) - g) %>% # make some gt's missing in our maxv query below
    lazy_dt()
)

lapply(lzdts, function(lzdt) {
  DT <- as.data.table(lzdt, key = c("g", "t", "v"))
  maxv <- max(DT$v) - 20
  bench::mark(
    a = lzdt %>% filter(v <= maxv) %>% arrange(desc(v)) %>% distinct(pick(all_of(c("g", "t"))), .keep_all = TRUE) %>% select(-v) %>% arrange(g, t) %>% collect(),
    a2 = lzdt %>% filter(v <= maxv) %>% arrange(across(all_of(c("g", "t"))),desc(v)) %>% distinct(pick(all_of(c("g", "t"))), .keep_all = TRUE) %>% select(-v) %>% collect(),
    a3 = lzdt %>% filter(v <= maxv) %>% arrange(pick(all_of(c("g", "t"))),desc(v)) %>% distinct(pick(all_of(c("g", "t"))), .keep_all = TRUE) %>% select(-v) %>% collect(),
    b = lzdt %>% filter(v <= maxv) %>% arrange(pick(all_of(c("g", "t"))), v) %>% slice_tail(n = 1, by = all_of(c("g", "t"))) %>% select(-v) %>% collect(),
    c = lzdt %>% filter(v <= maxv) %>% slice_tail(n = 1, by = all_of(c("g", "t"))) %>% select(-v) %>% collect(),
    d = {
      request_ektvs <- lzdt %>% filter(v <= maxv) %>% distinct(pick(all_of(c("g", "t")))) %>% mutate(v = maxv)
      left_join(request_ektvs %>% as_tibble(), lzdt %>% as_tibble(), by = join_by(g, t, closest(x$v >= y$v))) %>% select(-c(v.x, v.y))
    },
    e = {
      lzdt %>% filter(v <= maxv) %>% as_tibble() %>% slice_max(v, by = c(g,t)) %>% select(-v)
    },
    e2 = {
      lzdt %>% filter(v <= maxv) %>% as_tibble() %>% slice_tail(n = 1, by = c(g,t)) %>% select(-v)
    },
    e3 = {
      lzdt %>% filter(v <= maxv) %>% as_tibble() %>% arrange(g, t, desc(v)) %>% vctrs::vec_slice(vctrs::vec_duplicate_id(.[c("g", "t")]) == seq_len(nrow(.))) %>% select(-v)
    },
    f = {
      request_ektvs <- lzdt %>% filter(v <= maxv) %>% summarize(.by = all_of(c("g", "t")), v = max(v))
      left_join(request_ektvs, lzdt, by = c("g", "t", "v")) %>% select(-v) %>% collect()
    },
    g = unique(DT[v <= maxv,], by = c("g", "t"), fromLast = TRUE) %>% as.data.frame() %>% as_tibble() %>% select(-v),
    g2 = {
      DT <- as.data.table(lzdt, key = c("g", "t", "v"))
      unique(DT[v <= maxv,], by = c("g", "t"), fromLast = TRUE) %>% as.data.frame() %>% as_tibble() %>% select(-v)
    },
    h = {
      DT <- as.data.table(lzdt, key = c("g", "t", "v"))
      request_ektvs <- unique(DT, by = c("g", "t"))[, c("g", "t")][, v := maxv]
      DT[request_ektvs, on = c("g", "t", "v"), roll = TRUE, nomatch = NULL] %>% as.data.frame() %>% as_tibble() %>% select(-v)
    },
    min_time = 5
  )
})
#> [[1]]
#> # A tibble: 13 × 13
#>    expression      min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time result   memory
#>    <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm> <list>   <list>
#>  1 a          545.68ms 573.83ms    1.73          NA    1.16      9     6      5.19s <tibble> <NULL>
#>  2 a2         397.86ms 433.47ms    2.33          NA    1.17     12     6      5.14s <tibble> <NULL>
#>  3 a3         388.76ms 415.51ms    2.41          NA    0.926    13     5       5.4s <tibble> <NULL>
#>  4 b             2.09s    2.14s    0.462         NA    3.08      3    20      6.49s <tibble> <NULL>
#>  5 c             2.03s    2.08s    0.479         NA    3.36      3    21      6.26s <tibble> <NULL>
#>  6 d             1.43s    1.53s    0.642         NA    1.44      4     9      6.23s <tibble> <NULL>
#>  7 e             47.3s    47.3s    0.0211        NA    3.30      1   156      47.3s <tibble> <NULL>
#>  8 e2           24.56s   24.56s    0.0407        NA    3.50      1    86     24.56s <tibble> <NULL>
#>  9 e3         669.22ms 703.94ms    1.39          NA    1.39      7     7      5.05s <tibble> <NULL>
#> 10 f          671.86ms  691.6ms    1.38          NA    1.03      8     6      5.81s <tibble> <NULL>
#> 11 g          176.87ms 202.32ms    4.91          NA    1.77     25     9      5.09s <tibble> <NULL>
#> 12 g2         175.52ms 195.92ms    5.05          NA    1.36     26     7      5.15s <tibble> <NULL>
#> 13 h          175.39ms 184.59ms    5.35          NA    0.396    27     2      5.05s <tibble> <NULL>
#> # ℹ 2 more variables: time <list>, gc <list>
#> 
#> [[2]]
#> # A tibble: 13 × 13
#>    expression      min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time result   memory
#>    <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm> <list>   <list>
#>  1 a           60.14ms  65.28ms    15.1          NA    0.398    76     2      5.02s <tibble> <NULL>
#>  2 a2          79.67ms  86.45ms    11.5          NA    0.396    58     2      5.04s <tibble> <NULL>
#>  3 a3          77.62ms  85.21ms    11.6          NA    0.598    58     3      5.01s <tibble> <NULL>
#>  4 b          114.38ms 130.45ms     7.58         NA    1.20     38     6      5.01s <tibble> <NULL>
#>  5 c           98.15ms 113.44ms     8.71         NA    1.19     44     6      5.05s <tibble> <NULL>
#>  6 d          128.65ms 144.16ms     6.86         NA    0.980    35     5       5.1s <tibble> <NULL>
#>  7 e             1.51s    1.54s     0.643        NA    2.25      4    14      6.22s <tibble> <NULL>
#>  8 e2         723.54ms 810.93ms     1.25         NA    2.32      7    13      5.61s <tibble> <NULL>
#>  9 e3          74.82ms  91.74ms    10.9          NA    0.988    55     5      5.06s <tibble> <NULL>
#> 10 f           82.06ms  92.37ms    10.3          NA    0.395    52     2      5.06s <tibble> <NULL>
#> 11 g           21.45ms  24.85ms    38.7          NA    0.998   194     5      5.01s <tibble> <NULL>
#> 12 g2          22.03ms  25.88ms    37.9          NA    0.998   190     5      5.01s <tibble> <NULL>
#> 13 h           14.51ms  18.58ms    53.5          NA    0.399   268     2      5.01s <tibble> <NULL>
#> # ℹ 2 more variables: time <list>, gc <list>
#> 
#> Warning messages:
#> 1: Some expressions had a GC in every iteration; so filtering is disabled. 
#> 2: Some expressions had a GC in every iteration; so filtering is disabled.