-
Notifications
You must be signed in to change notification settings - Fork 9
/
bm-write-csv.R
89 lines (83 loc) · 2.76 KB
/
bm-write-csv.R
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
#' Benchmark CSV writing
#'
#' @section Parameters:
#' * `source` A CSV file path to write to
#' * `writer` One of `c("arrow", "data.table", "vroom", "readr",)`
#' * `input` One of `c("arrow_table", "data_frame")`
#'
#' @export
write_csv <- Benchmark(
"write_csv",
setup = function(source = names(known_sources),
writer = "arrow",
compression = c("uncompressed", "gzip"),
input = c("arrow_table", "data_frame")) {
writer <- match.arg(writer, c("arrow", "data.table", "vroom", "readr", "base"))
compression <- match.arg(compression, c("uncompressed", "gzip"))
input <- match.arg(input)
# source defaults are retrieved from the function definition (all available
# known_sources) and then read the source in as a data.frame
source <- ensure_source(source)
df <- read_source(source, as_data_frame = match.arg(input) == "data_frame")
ext <- switch(
compression,
uncompressed = ".csv",
gzip = ".csv.gz",
paste0(".csv.", compression)
)
# Map string param name to functions
BenchEnvironment(
write_csv_func = get_csv_writer(writer),
source = source,
df = df,
ext = ext
)
},
# delete the results before each iteration
before_each = {
result_file <- tempfile(fileext = ext)
},
# the benchmark to run
run = {
write_csv_func(df, result_file)
},
# after each iteration, check the dimensions and delete the results
after_each = {
stopifnot(identical(dim(df), dim(arrow::open_dataset(result_file, format = "csv"))))
stopifnot("Output file does not exist" = file.exists(result_file))
unlink(result_file)
},
valid_params = function(params) {
## Only arrow fns will accept an arrow_table
drop <- ( params$input == "arrow_table" & params$writer != "arrow" )
params[!drop,]
},
packages_used = function(params) {
params$writer
}
)
#' Get a CSV writer
#'
#' @param writer the writer to use
#'
#' @return the csv writer
#' @export
get_csv_writer <- function(writer) {
library(writer, character.only = TRUE, warn.conflicts = FALSE)
if (writer == "arrow") {
return(function(...) arrow::write_csv_arrow(...))
} else if (writer == "readr") {
return(function(..., as_data_frame) readr::write_csv(...))
} else if (writer == "data.table") {
return(function(..., as_data_frame) data.table::fwrite(...))
} else if (writer == "vroom") {
return(function(..., as_data_frame) vroom::vroom_write(..., delim = ","))
} else if (writer == "base") {
return(function(df, result_file) {
if (tools::file_ext(result_file) == "gz") result_file <- gzfile(result_file)
utils::write.csv(df, result_file, row.names = FALSE)
})
} else {
stop("Unsupported writer: ", writer, call. = FALSE)
}
}