edit loadGenotypeRegion so it filter variants before reading dosages

R/genotypeIo.R

@@ -1447,6 +1447,9 @@ loadGenotypeRegion <- function(
     )
     variants <- .rangesSnpInfo(SummarizedExperiment::rowRanges(panel))
     snpIdx <- .loadGenoSnpIdx(variants, which)
+    # Apply the indel and whitelist filters to the variant index before any
+    # dosage is read, so only the requested variants are loaded.
+    snpIdx <- .loadGenoKeepIdx(variants, snpIdx, keepIndel, keepVariantsPath)
     # Samples x variants matrix (pecotmr convention), its columns named by
     # the file's own variant ids; callers handle missing.
     extracted <- list(
@@ -1459,9 +1462,8 @@ loadGenotypeRegion <- function(
             .snpInfoToVariantInfo(slice(variants, snpIdx))
         )
     )
-    filtered <- .loadGenoPostFilter(extracted, keepIndel, keepVariantsPath)
     result <- .loadGenoInvertStochastic(
-        filtered,
+        extracted,
         genotype,
         stochasticMetaPath,
         stochasticMetaFormat
@@ -1506,35 +1508,23 @@ loadGenotypeRegion <- function(
     )
 }
 
-# Apply the indel-drop and variant-whitelist post-filters to (X, variant_info).
+# Restrict the SNP indices to the variants kept by the indel and whitelist
+# filters. Runs before the dosage read, so filtered-out variants are never
+# loaded.
 # @noRd
-.loadGenoPostFilter <- function(result, keepIndel, keepVariantsPath) {
-    snpOnly <- if (keepIndel) {
-        result
-    } else {
-        .loadGenoSubsetVariants(
-            result,
-            isSnpAlleles(result$variant_info$A1, result$variant_info$A2)
-        )
+.loadGenoKeepIdx <- function(variants, snpIdx, keepIndel, keepVariantsPath) {
+    if (keepIndel && is.null(keepVariantsPath)) {
+        return(snpIdx)
     }
-    if (is.null(keepVariantsPath)) {
-        return(snpOnly)
+    info <- .snpInfoToVariantInfo(slice(variants, snpIdx))
+    keep <- rep(TRUE, length(snpIdx))
+    if (!keepIndel) {
+        keep <- keep & isSnpAlleles(info$A1, info$A2)
     }
-    .loadGenoSubsetVariants(
-        snpOnly,
-        matchVariantsToKeep(snpOnly$variant_info, keepVariantsPath)
-    )
-}
-
-# Subset a loaded (X, variant_info) pair along the variant axis, keeping the
-# dosage columns and the metadata rows in step.
-# @noRd
-.loadGenoSubsetVariants <- function(result, keep) {
-    list_assign(
-        result,
-        X = result$X[, keep, drop = FALSE],
-        variant_info = result$variant_info[keep, , drop = FALSE]
-    )
+    if (!is.null(keepVariantsPath)) {
+        keep <- keep & matchVariantsToKeep(info, keepVariantsPath)
+    }
+    snpIdx[keep]
 }
 
 # Restore a stochastic sketch's original scale from its metadata sidecar.

tests/testthat/test_genotypeIo.R

@@ -622,6 +622,27 @@ test_that("loadGenotypeRegion filters by keep_variants_path for plink2", {
     expect_equal(ncol(result), 3L)
 })
 
+test_that("loadGenotypeRegion with keepVariantsPath reads the same dosages as a full load", {
+    skip_if_not_installed("pgenlibr")
+    td <- test_path("test_data")
+    full <- loadGenotypeRegion(
+        file.path(td, "test_variants"),
+        returnVariantInfo = TRUE
+    )
+    keep_file <- tempfile(fileext = ".tsv")
+    on.exit(unlink(keep_file), add = TRUE)
+    keep_df <- full$variant_info[c(2, 5, 9), c("chrom", "pos", "A2", "A1")]
+    vroom::vroom_write(keep_df, keep_file, delim = "\t")
+    kept <- loadGenotypeRegion(
+        file.path(td, "test_variants"),
+        keepVariantsPath = keep_file,
+        returnVariantInfo = TRUE
+    )
+    expect_equal(kept$X, full$X[, c(2, 5, 9), drop = FALSE])
+    expect_equal(kept$variant_info, full$variant_info[c(2, 5, 9), , drop = FALSE],
+        ignore_attr = TRUE)
+})
+
 test_that("loadGenotypeRegion attaches afreq info for plink2", {
     skip_if_not_installed("pgenlibr")
     td <- test_path("test_data")