Support optional DuckDB-backed inputs

DESCRIPTION

@@ -50,7 +50,9 @@ Suggests:
     tidyseurat,
     tidySingleCellExperiment,
     bayesplot,
-    remotes
+    remotes,
+    duckdb,
+    dbplyr
 Additional_repositories:
     https://mc-stan.org/r-packages/
 SystemRequirements: CmdStan (https://mc-stan.org/users/interfaces/cmdstan), C++14

NAMESPACE

@@ -6,6 +6,7 @@ S3method(sccomp_calculate_residuals,sccomp_tbl)
 S3method(sccomp_estimate,DFrame)
 S3method(sccomp_estimate,Seurat)
 S3method(sccomp_estimate,SingleCellExperiment)
+S3method(sccomp_estimate,tbl_duckdb_connection)
 S3method(sccomp_estimate,data.frame)
 S3method(sccomp_predict,sccomp_tbl)
 S3method(sccomp_proportional_fold_change,sccomp_tbl)

R/sccomp_estimate.R

@@ -532,6 +532,13 @@ sccomp_estimate.DFrame <- function(.data,
     )
 }
 
+#' @export
+sccomp_estimate.tbl_duckdb_connection <- function(.data, ...) {
+  check_and_install_packages(c("duckdb", "dbplyr"))
+
+  sccomp_estimate.data.frame(.data, .collect_after_count = TRUE, ...)
+}
+
 
 #' @importFrom purrr when
 #' @importFrom rlang is_symbolic
@@ -908,8 +915,9 @@ sccomp_glm_data_frame_counts = function(.data,
                                         mcmc_seed = sample_seed(),
                                         max_sampling_iterations = 20000,
                                                                                  pass_fit = TRUE,
-                                         sig_figs = 9,
-                                         cache_stan_model = sccomp_stan_models_cache_dir,
+                                                                                 sig_figs = 9,
+                                                                                 cache_stan_model = sccomp_stan_models_cache_dir,
+                                                                                 .collect_after_count = FALSE,
                                         ...) {
 
   # Prepare column same enquo
@@ -948,7 +956,13 @@ sccomp_glm_data_frame_counts = function(.data,
 
 
   # Make rectangular data
-  .data = .data |> make_rectangular_data(!!.sample, !!.cell_group, !!.count, formula_composition)
+  .data = .data |> make_rectangular_data(
+    !!.sample,
+    !!.cell_group,
+    !!.count,
+    formula_composition,
+    collect_after_count = .collect_after_count
+  )
 
   # Check if test_composition_above_logit_fold_change is 0, as the Bayesian FDR does not allow it
   if(test_composition_above_logit_fold_change <= 0)
@@ -1151,16 +1165,17 @@ sccomp_glm_data_frame_counts = function(.data,
 #' @return A rectangular data frame with zeros added for missing combinations
 #' @keywords internal
 #' @noRd
-make_rectangular_data = function(.data, .sample, .cell_group, .count, formula_composition) {
+make_rectangular_data = function(.data, .sample, .cell_group, .count, formula_composition, collect_after_count = FALSE) {
 
   .sample = enquo(.sample)
   .cell_group = enquo(.cell_group)
   .count = enquo(.count)
 
-  if(
-    .data |> count(!!.sample) |> distinct(n) |> nrow() > 1 || 
-    .data |> count(!!.cell_group) |> distinct(n) |> nrow() > 1 
-  ){
+  sample_counts <- .data |> count(!!.sample) |> distinct(n)
+  cell_group_counts <- .data |> count(!!.cell_group) |> distinct(n)
+  if (collect_after_count) .data <- dplyr::collect(.data)
+
+  if (sample_counts |> nrow() > 1 || cell_group_counts |> nrow() > 1) {
     warning(sprintf("sccomp says: the input data frame does not have the same number of `%s`, for all `%s`. We have made it so, adding 0s for the missing sample/feature pairs.", quo_name(.cell_group), quo_name(.sample)))
 
     .data |> 

tests/testthat/test-duckdb.R

@@ -0,0 +1,46 @@
+test_that("sccomp_estimate has a DuckDB table method", {
+  expect_true(is.function(getS3method("sccomp_estimate", "tbl_duckdb_connection")))
+})
+
+test_that("sccomp_estimate gives the same results for DuckDB tables", {
+  skip_if_not_installed("duckdb")
+  skip_if_not_installed("dbplyr")
+  skip_cmdstan()
+
+  data("counts_obj", package = "sccomp", envir = environment())
+  con <- DBI::dbConnect(duckdb::duckdb(), dbdir = ":memory:")
+  on.exit(DBI::dbDisconnect(con, shutdown = TRUE), add = TRUE)
+  DBI::dbWriteTable(con, "counts_obj", counts_obj)
+
+  estimate_args <- list(
+    formula_composition = ~type,
+    sample = "sample",
+    cell_group = "cell_group",
+    abundance = "count",
+    inference_method = "pathfinder",
+    cores = 1,
+    mcmc_seed = 12345,
+    max_sampling_iterations = 1000,
+    verbose = FALSE
+  )
+
+  data_frame_estimate <- do.call(
+    sccomp_estimate,
+    c(list(counts_obj), estimate_args)
+  )
+  duckdb_estimate <- do.call(
+    sccomp_estimate,
+    c(list(dplyr::tbl(con, "counts_obj")), estimate_args)
+  )
+
+  data_frame_results <- data_frame_estimate |>
+      dplyr::arrange(cell_group, parameter) |>
+      dplyr::select(cell_group, parameter, c_effect, c_lower, c_upper) |>
+      as.data.frame()
+  duckdb_results <- duckdb_estimate |>
+      dplyr::arrange(cell_group, parameter) |>
+      dplyr::select(cell_group, parameter, c_effect, c_lower, c_upper) |>
+      as.data.frame()
+
+  expect_equal(data_frame_results, duckdb_results, tolerance = 1e-2)
+})