diff --git a/R/csv.R b/R/csv.R index d199a7939..d641eb627 100644 --- a/R/csv.R +++ b/R/csv.R @@ -171,7 +171,7 @@ read_cmdstan_csv <- function(files, files <- file.path(temp_storage, basename(files)) } format <- assert_valid_draws_format(format) - assert_file_exists(files, access = "r", extension = "csv") + assert_file_exists(files, access = "r", extension = c("csv", "gz", "bz2")) metadata <- NULL warmup_draws <- list() draws <- list() @@ -282,23 +282,7 @@ read_cmdstan_csv <- function(files, num_warmup_draws <- ceiling(metadata$iter_warmup / metadata$thin) num_post_warmup_draws <- ceiling(metadata$iter_sampling / metadata$thin) for (output_file in files) { - if (os_is_windows()) { - grep_path_repaired <- withr::with_path( - c( - toolchain_PATH_env_var() - ), - repair_path(Sys.which("grep.exe")) - ) - grep_path_quotes <- paste0('"', grep_path_repaired, '"') - fread_cmd <- paste0( - grep_path_quotes, - " -v \"^#\" --color=never \"", - wsl_safe_path(output_file, revert = TRUE), - "\"" - ) - } else { - fread_cmd <- paste0("grep -v '^#' --color=never '", path.expand(output_file), "'") - } + fread_cmd <- build_fread_cmd(output_file, "^#", invert = TRUE) if (length(sampler_diagnostics) > 0) { post_warmup_sd_id <- length(post_warmup_sampler_diagnostics) + 1 warmup_sd_id <- length(warmup_sampler_diagnostics) + 1 @@ -706,6 +690,43 @@ parse_generated_quantities_time <- function(line) { time } +# Build a command for data.table::fread() that filters CSV lines via grep. +# Handles gzip (.gz) and bzip2 (.bz2) compressed files. +# @param file Path to the CSV file (possibly compressed) +# @param pattern Regex pattern to filter by +# @param invert If TRUE, exclude matching lines (grep -v) +build_fread_cmd <- function(file, pattern, invert = FALSE) { + v_flag <- if (invert) "-v " else "" + compression <- if (grepl("\\.gz$", file, ignore.case = TRUE)) "gz" + else if (grepl("\\.bz2$", file, ignore.case = TRUE)) "bz2" + else NULL + + if (os_is_windows()) { + toolchain <- withr::with_path( + c(toolchain_PATH_env_var()), + list( + grep = repair_path(Sys.which("grep.exe")), + gz = repair_path(Sys.which("gzip.exe")), + bz2 = repair_path(Sys.which("bzip2.exe")) + ) + ) + file_path <- wsl_safe_path(file, revert = TRUE) + grep_args <- paste0("--color=never ", v_flag, '"', pattern, '"') + if (is.null(compression)) { + paste0('"', toolchain$grep, '" ', grep_args, ' "', file_path, '"') + } else { + paste0('"', toolchain[[compression]], '" -dc "', file_path, '" | "', + toolchain$grep, '" ', grep_args) + } + } else { + grep_cmd <- if (is.null(compression)) "grep" + else if (compression == "gz") "zgrep" + else "bzgrep" + paste0(grep_cmd, " --color=never ", v_flag, "'", pattern, "' '", + path.expand(file), "'") + } +} + #' Reads the sampling arguments and the diagonal of the #' inverse mass matrix from the comments in a CSV file. #' @@ -715,7 +736,7 @@ parse_generated_quantities_time <- function(line) { #' mass matrix (or its diagonal depending on the metric). #' read_csv_metadata <- function(csv_file) { - assert_file_exists(csv_file, access = "r", extension = "csv") + assert_file_exists(csv_file, access = "r", extension = c("csv", "gz", "bz2")) inv_metric_next <- FALSE csv_file_info <- list() csv_file_info$inv_metric <- NULL @@ -727,23 +748,7 @@ read_csv_metadata <- function(csv_file) { warmup_time <- 0 sampling_time <- 0 total_time <- 0 - if (os_is_windows()) { - grep_path_repaired <- withr::with_path( - c( - toolchain_PATH_env_var() - ), - repair_path(Sys.which("grep.exe")) - ) - grep_path_quotes <- paste0('"', grep_path_repaired, '"') - fread_cmd <- paste0( - grep_path_quotes, - " \"^[#a-zA-Z]\" --color=never \"", - wsl_safe_path(csv_file, revert = TRUE), - "\"" - ) - } else { - fread_cmd <- paste0("grep '^[#a-zA-Z]' --color=never '", path.expand(csv_file), "'") - } + fread_cmd <- build_fread_cmd(csv_file, "^[#a-zA-Z]", invert = FALSE) suppressWarnings( metadata <- data.table::fread( cmd = fread_cmd, diff --git a/tests/testthat/test-csv.R b/tests/testthat/test-csv.R index ada4f60f4..b53072da6 100644 --- a/tests/testthat/test-csv.R +++ b/tests/testthat/test-csv.R @@ -1024,3 +1024,28 @@ test_that("as_cmdstan_fit filters variables across methods", { expect_equal(pathfinder$summary()$variable, pathfinder_vars) expect_equal(pathfinder$metadata()$variables, pathfinder_vars) }) + +compress_csv <- function(src, ext) { + dest <- tempfile(fileext = paste0(".", ext)) + con <- if (identical(ext, "csv.gz")) gzfile(dest, "wt") else bzfile(dest, "wt") + writeLines(readLines(src), con) + close(con) + dest +} + +test_that("read_cmdstan_csv() reads compressed CSV files", { + csv_files <- c( + test_path("resources", "csv", "model1-1-warmup.csv"), + test_path("resources", "csv", "model1-2-warmup.csv") + ) + expected <- read_cmdstan_csv(csv_files) + compressed_files <- lapply( + c("csv.gz", "csv.bz2"), + function(ext) unname(vapply(csv_files, compress_csv, ext = ext, character(1))) + ) + on.exit(unlink(unlist(compressed_files)), add = TRUE) + + for (files in compressed_files) { + expect_equal(read_cmdstan_csv(files), expected) + } +})