Add Zarr to benchmarks (#446)

lazappi · keller-mark · Artur-man · web-flow · commit 897914bb5aba · 2026-05-12T20:03:28.000+02:00
* WIP * More tests passing * Fix df read bug * More tests passing after fixing zero-dimensional get bug in pizzarr * WIP: writing * Fix more tests * Zarr df writing * WIP: ZarrAnnData class * Tests passing * Tests that compare h5ad to zarr * Use Rarr to read full numeric arrays * Fix bugs. Add test for from_SingleCellExperiment with Zarr * Add a to_dense param to ZarrAnnData constructor. Add overwrite params internally. * Update * Backwards dense/sparse * Simplify how obs and var names handled in ZarrAnnData (similar to #171) * update extdata and documentation * fix set/get zarr _index, update text example.zarr and update tests similar to HDF5AnnData * Fix test * Revert unnecessary changes * Formatting * Add comments * remove unnecessary example zarr store * lintr and R check for zarr related utilities and functions, updated some documentation * add pizzarr to Suggests and README * proj * add keller-mark/pizzarr to Remotes * zip example.zarr * adapt read_zarr to Rarr * adapt write_zarr to Rarr * remove old scripts * update write_zarr * initial update to ZarrAnnData * update ZarrAnnData, documentation, and implement read_zarr_rec_array * review read zarr helpers, and update tests * update read_zarr, read tests pass * some updates for writing zarr * update write_empty_zarr * remove pizzarr, update documentation * remove pizzarr from tests * fix test-ZarrAnnData * update ZarrAnnData to imitate HDF5AnnData * check redundant files, correct lines * update example_h5ad.py, add zarr and change to example_files.py * add new test example * some linting changes * remove read/write_zattrs since implemented in Rarr * access read/write_zarr_attr * add some missing tests * update readers, update tests * correct nullable string zarr array write/read, introduce ordering in categorical zarr array * do some linting, fix commented out code * update some zarr writers and classes * fix documentation * fix compression interface for zarr * full lint check * fix examples * check, biocheck and lintr * fix development status * air format * air format test * update example.zarr.zip, skip some test (waiting for Rarr) * update example.zarr, fix some read_zarr_ * fix examples * remove overwrite * R code styling * fixes from @lazappi * air format * update some documentation * fix some tests * more fixes on anndata-zarr integration * update ZarrAnnData$initialize * update zarr compression * fix column-order here, C based ordering for arrays * implement roundtrip tests for anndata-zarr * add zarr to vignettes * update README and software_design.rmd * update AnnData-usage * update write_zarr documentation * update write_zarr_null * fix rec_array, update tests and example datasets * fix duplicate chunks in Rmd * add write_zarr_null * update write string array (zarr), air and lint * implement writing empty zarr elements * update tests for rec_array conformance of h5ad and zarr * update mapping conformance test for h5ad and zarr * implement H5_ITER like ordering and fix h5ad vs zarr testing * air and lint * fix test bug * do not call expect_equal outside of test * implement examples, test and datasets for zarr v3 * fix issues, lint and update example datasets to new anndata version * lint and merge * revert some lines * small changes * revert small changes * air format some tests * Set v2 in write_zarr_* helpers * Fix stop message in write_zarr_element() * Fix roxygen comment in write_zarr_element() * Expand compression list in as_ZarrAnnData() * Fix H5_ITER_INC_ORDERING docs * Fix as_ZarrAnnData() compression docs * Fix Zarr varm roundtrip test * Review duplicate entry in README * Fix typo in AnnData-usage docs * Fix comma in software design vignette * Adjust class descriptions in software design vignette * Roxygenise * Minor text fixes * Document .get_compressor * Minor fixes to function docs * Comment logic in create_zarr_group() * Fix indentation in read_zarr_sparse_array() * Add construct sparse matrix helper * Add ZARR_METADATA_FILES vector * Eval Zarr chunks in vignettes * Merge test-Zarrv3-read.R into test-Zarr-read.R * Combine roundtrip tests * Add roundtrip test helpers * Refactor test-h5ad-zarr.R to use helper * Refactor example files script - Use uv comments for dependencies - Update dependencies to latest stable versions - Add progress messages - Use Ruff for formatting * Remove H5_ITER_INC_ORDERING() * Remove Zarr compression comment * Fix factor creation in read_zarr_categorical() * Pin Rarr version * Delete existing Zarr path before writing * Add helper functions for accessing Zarr keys * Update read_zarr_element() error message * Add dimname warnings to ZarrAnnData * Add Zarr writeability checks/tests * Roxygenise, lint, style * Use setup-bioc for all GHA * Update WORDLIST * Add .venv to .Rbuildignore * Clean up test output * Add Zarr to benchmarks * run air format --------- Co-authored-by: Mark Keller <7525285+keller-mark@users.noreply.github.com> Co-authored-by: Artur-man <artur-man@hotmail.com> Co-authored-by: Robrecht Cannoodt <rcannood@gmail.com>
diff --git a/benchmarks/lib/helpers.R b/benchmarks/lib/helpers.R
@@ -83,6 +83,23 @@ generate_bench_h5ad <- function(x_type, n_obs, n_vars, cache_dir) {
   path
 }
 
+#' Convert an H5AD bench file to a Zarr store and cache it
+#'
+#' @param x_type Matrix type key (matches h5ad_paths names)
+#' @param h5ad_path Path to the corresponding H5AD file
+#' @param cache_dir Directory to cache generated stores
+#' @return Path to the generated Zarr store directory
+generate_bench_zarr <- function(x_type, h5ad_path, cache_dir) {
+  path <- file.path(cache_dir, paste0("bench_", x_type, ".zarr"))
+  if (dir.exists(path)) {
+    return(path)
+  }
+  ad <- reticulate::import("anndata", convert = FALSE)
+  adata_py <- ad$read_h5ad(h5ad_path)
+  adata_py$write_zarr(path)
+  path
+}
+
 # ---------------------------------------------------------------------------
 # bench::mark → BMF JSON conversion
 # ---------------------------------------------------------------------------
diff --git a/benchmarks/run_benchmarks.R b/benchmarks/run_benchmarks.R
@@ -73,6 +73,22 @@ h5ad_paths <- setNames(
 )
 cat("\n")
 
+cat("Generating Zarr test data (converting from H5AD)...\n")
+zarr_paths <- setNames(
+  vapply(
+    x_types,
+    function(xt) {
+      cat(sprintf("  %s... ", xt))
+      path <- generate_bench_zarr(xt, h5ad_paths[[xt]], cache_dir)
+      cat("done\n")
+      path
+    },
+    character(1)
+  ),
+  x_types
+)
+cat("\n")
+
 # ---------------------------------------------------------------------------
 # Run selected suites
 # ---------------------------------------------------------------------------
@@ -88,12 +104,12 @@ for (suite in suites_to_run) {
 
   suite_results <- switch(
     suite,
-    read = bench_read(h5ad_paths, opts$iterations, x_types),
-    write = bench_write(h5ad_paths, opts$iterations, x_types),
-    get = bench_get(h5ad_paths, opts$iterations),
-    set = bench_set(h5ad_paths, opts$iterations),
-    convert = bench_convert(h5ad_paths, opts$iterations, x_types),
-    subset = bench_subset(h5ad_paths, opts$iterations),
+    read = bench_read(h5ad_paths, opts$iterations, x_types, zarr_paths),
+    write = bench_write(h5ad_paths, opts$iterations, x_types, zarr_paths),
+    get = bench_get(h5ad_paths, opts$iterations, zarr_paths),
+    set = bench_set(h5ad_paths, opts$iterations, zarr_paths),
+    convert = bench_convert(h5ad_paths, opts$iterations, x_types, zarr_paths),
+    subset = bench_subset(h5ad_paths, opts$iterations, zarr_paths),
     {
       warning("Unknown suite: ", suite)
       list()
diff --git a/benchmarks/suites/bench_convert.R b/benchmarks/suites/bench_convert.R
@@ -5,7 +5,7 @@
 # format conversions (InMemory↔SCE, InMemory↔Seurat).
 # =============================================================================
 
-bench_convert <- function(h5ad_paths, iterations, x_types) {
+bench_convert <- function(h5ad_paths, iterations, x_types, zarr_paths) {
   results <- list()
 
   # --- Backend conversions (per X type) ---
@@ -47,6 +47,43 @@ bench_convert <- function(h5ad_paths, iterations, x_types) {
     )
   }
 
+  # --- Zarr ↔ InMemory conversions (per X type) ---
+  for (xt in x_types) {
+    zarr_path <- zarr_paths[[xt]]
+
+    # Zarr → InMemory
+    env <- new.env(parent = globalenv())
+    env$.ad <- read_zarr(zarr_path, as = "ZarrAnnData")
+
+    results <- c(
+      results,
+      run_one_benchmark(
+        name = paste0("convert_Zarr_to_InMemory_", xt),
+        expr = quote(.ad$as_InMemoryAnnData()),
+        iterations = iterations,
+        env = env
+      )
+    )
+
+    # InMemory → Zarr
+    env2 <- new.env(parent = globalenv())
+    env2$.ad <- read_zarr(zarr_path, as = "InMemoryAnnData")
+
+    results <- c(
+      results,
+      run_one_benchmark(
+        name = paste0("convert_InMemory_to_Zarr_", xt),
+        expr = quote({
+          .tmp <- tempfile()
+          .result <- .ad$as_ZarrAnnData(.tmp)
+          unlink(.tmp, recursive = TRUE)
+        }),
+        iterations = iterations,
+        env = env2
+      )
+    )
+  }
+
   # --- Format conversions (using float_csparse as representative) ---
   path <- h5ad_paths[["float_csparse"]]
   ad <- read_h5ad(path, as = "InMemoryAnnData")
diff --git a/benchmarks/suites/bench_get.R b/benchmarks/suites/bench_get.R
@@ -42,15 +42,24 @@
   colnames = quote(colnames(.ad))
 )
 
-bench_get <- function(h5ad_paths, iterations) {
+bench_get <- function(h5ad_paths, iterations, zarr_paths) {
   results <- list()
   path <- h5ad_paths[["float_csparse"]]
 
-  for (backend in c("InMemoryAnnData", "HDF5AnnData")) {
-    short <- if (backend == "InMemoryAnnData") "InMemory" else "HDF5"
+  for (backend in c("InMemoryAnnData", "HDF5AnnData", "ZarrAnnData")) {
+    short <- switch(
+      backend,
+      InMemoryAnnData = "InMemory",
+      HDF5AnnData = "HDF5",
+      ZarrAnnData = "Zarr"
+    )
 
     # Open the AnnData
-    ad <- read_h5ad(path, as = backend)
+    ad <- if (backend == "ZarrAnnData") {
+      read_zarr(zarr_paths[["float_csparse"]], as = "ZarrAnnData")
+    } else {
+      read_h5ad(path, as = backend)
+    }
 
     # --- Slot getters ---
     for (slot in .bench_slots) {
diff --git a/benchmarks/suites/bench_read.R b/benchmarks/suites/bench_read.R
@@ -5,7 +5,7 @@
 # across different X matrix types.
 # =============================================================================
 
-bench_read <- function(h5ad_paths, iterations, x_types) {
+bench_read <- function(h5ad_paths, iterations, x_types, zarr_paths) {
   results <- list()
 
   for (xt in x_types) {
@@ -37,5 +37,32 @@ bench_read <- function(h5ad_paths, iterations, x_types) {
     )
   }
 
+  # Read from Zarr store
+  for (xt in x_types) {
+    path <- zarr_paths[[xt]]
+
+    # Read Zarr → InMemoryAnnData
+    results <- c(
+      results,
+      run_one_benchmark(
+        name = paste0("read_zarr_InMemory_", xt),
+        expr = quote(read_zarr(.path, as = "InMemoryAnnData")),
+        setup = bquote(.path <- .(path)),
+        iterations = iterations
+      )
+    )
+
+    # Open Zarr lazily → ZarrAnnData
+    results <- c(
+      results,
+      run_one_benchmark(
+        name = paste0("read_zarr_Zarr_", xt),
+        expr = quote(read_zarr(.path, as = "ZarrAnnData")),
+        setup = bquote(.path <- .(path)),
+        iterations = iterations
+      )
+    )
+  }
+
   results
 }
diff --git a/benchmarks/suites/bench_set.R b/benchmarks/suites/bench_set.R
@@ -4,7 +4,7 @@
 # Benchmarks setting every AnnData slot on both InMemory and HDF5 backends.
 # =============================================================================
 
-bench_set <- function(h5ad_paths, iterations) {
+bench_set <- function(h5ad_paths, iterations, zarr_paths) {
   results <- list()
   path <- h5ad_paths[["float_csparse"]]
 
@@ -22,12 +22,25 @@ bench_set <- function(h5ad_paths, iterations) {
     "uns"
   )
 
-  for (backend in c("InMemoryAnnData", "HDF5AnnData")) {
-    short <- if (backend == "InMemoryAnnData") "InMemory" else "HDF5"
+  for (backend in c("InMemoryAnnData", "HDF5AnnData", "ZarrAnnData")) {
+    short <- switch(
+      backend,
+      InMemoryAnnData = "InMemory",
+      HDF5AnnData = "HDF5",
+      ZarrAnnData = "Zarr"
+    )
 
     for (slot in slots) {
-      # For HDF5, we need a fresh writable copy for each slot
-      if (backend == "HDF5AnnData") {
+      # Each backend needs a fresh writable instance per slot
+      if (backend == "ZarrAnnData") {
+        # Copy Zarr store directory so each slot gets a fresh writable copy
+        zarr_path <- zarr_paths[["float_csparse"]]
+        tmp_parent <- tempfile()
+        dir.create(tmp_parent, recursive = TRUE)
+        file.copy(zarr_path, tmp_parent, recursive = TRUE)
+        tmp <- file.path(tmp_parent, basename(zarr_path))
+        ad <- read_zarr(tmp, as = "ZarrAnnData", mode = "r+")
+      } else if (backend == "HDF5AnnData") {
         tmp <- tempfile(fileext = ".h5ad")
         file.copy(path, tmp)
         ad <- suppressWarnings(
@@ -55,7 +68,9 @@ bench_set <- function(h5ad_paths, iterations) {
         )
       )
 
-      if (backend == "HDF5AnnData") {
+      if (backend == "ZarrAnnData") {
+        unlink(tmp, recursive = TRUE)
+      } else if (backend == "HDF5AnnData") {
         ad$close()
         unlink(tmp)
       }
diff --git a/benchmarks/suites/bench_subset.R b/benchmarks/suites/bench_subset.R
@@ -5,13 +5,22 @@
 # materialization back to concrete implementations.
 # =============================================================================
 
-bench_subset <- function(h5ad_paths, iterations) {
+bench_subset <- function(h5ad_paths, iterations, zarr_paths) {
   results <- list()
   path <- h5ad_paths[["float_csparse"]]
 
-  for (backend in c("InMemoryAnnData", "HDF5AnnData")) {
-    short <- if (backend == "InMemoryAnnData") "InMemory" else "HDF5"
-    ad <- read_h5ad(path, as = backend)
+  for (backend in c("InMemoryAnnData", "HDF5AnnData", "ZarrAnnData")) {
+    short <- switch(
+      backend,
+      InMemoryAnnData = "InMemory",
+      HDF5AnnData = "HDF5",
+      ZarrAnnData = "Zarr"
+    )
+    ad <- if (backend == "ZarrAnnData") {
+      read_zarr(zarr_paths[["float_csparse"]], as = "ZarrAnnData")
+    } else {
+      read_h5ad(path, as = backend)
+    }
 
     n_obs <- ad$n_obs()
     n_vars <- ad$n_vars()
@@ -123,7 +132,22 @@ bench_subset <- function(h5ad_paths, iterations) {
       )
     )
 
-    # Clean up
+    # --- Materialize view → Zarr ---
+    results <- c(
+      results,
+      run_one_benchmark(
+        name = paste0("materialize_to_Zarr_", short),
+        expr = quote({
+          .tmp <- tempfile()
+          .result <- .view$as_ZarrAnnData(.tmp)
+          unlink(.tmp, recursive = TRUE)
+        }),
+        iterations = iterations,
+        env = env4
+      )
+    )
+
+    # Clean up (ZarrAnnData holds no persistent file handles)
     if (backend == "HDF5AnnData") {
       ad$close()
     }
diff --git a/benchmarks/suites/bench_write.R b/benchmarks/suites/bench_write.R
@@ -5,7 +5,7 @@
 # with different compression settings and X matrix types.
 # =============================================================================
 
-bench_write <- function(h5ad_paths, iterations, x_types) {
+bench_write <- function(h5ad_paths, iterations, x_types, zarr_paths) {
   results <- list()
 
   compressions <- c("none", "gzip")
@@ -57,5 +57,50 @@ bench_write <- function(h5ad_paths, iterations, x_types) {
     }
   }
 
+  # Write to Zarr store
+  for (xt in x_types) {
+    path <- zarr_paths[[xt]]
+
+    for (compression in compressions) {
+      # Write from InMemoryAnnData → Zarr
+      env <- new.env(parent = globalenv())
+      env$.ad <- read_zarr(path, as = "InMemoryAnnData")
+      env$.compression <- compression
+
+      results <- c(
+        results,
+        run_one_benchmark(
+          name = paste0("write_zarr_InMemory_", xt, "_", compression),
+          expr = quote({
+            .tmp <- tempfile()
+            .ad$as_ZarrAnnData(.tmp, compression = .compression)
+            unlink(.tmp, recursive = TRUE)
+          }),
+          iterations = iterations,
+          env = env
+        )
+      )
+
+      # Write from ZarrAnnData → Zarr
+      env2 <- new.env(parent = globalenv())
+      env2$.ad <- read_zarr(path, as = "ZarrAnnData")
+      env2$.compression <- compression
+
+      results <- c(
+        results,
+        run_one_benchmark(
+          name = paste0("write_zarr_Zarr_", xt, "_", compression),
+          expr = quote({
+            .tmp <- tempfile()
+            .ad$as_ZarrAnnData(.tmp, compression = .compression)
+            unlink(.tmp, recursive = TRUE)
+          }),
+          iterations = iterations,
+          env = env2
+        )
+      )
+    }
+  }
+
   results
 }
