diff --git a/NEWS.md b/NEWS.md index 73555477..fe4a7123 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,6 +1,7 @@ # haven (development version) * haven no longer depends on readr (#834). As a consequence, the `read_*()` functions no longer read compressed (`.gz`, `.bz2`, `.xz`, `.zip`) files; decompress them before reading. +* All `read_*()` functions with small `n_max` no longer pre-allocate more rows than needed, greatly reducing memory usage of header-only reads (#845). * Updated to ReadStat dev (0dc2aee). * `write_sav()` no longer produces files that `read_sav()` reads back incorrectly when a long string variable's internal segment name collides with another variable name (#798). * SAS catalog reader now correctly reads big-endian files. diff --git a/src/reader.cpp b/src/reader.cpp index 7b516a78..f566ab67 100644 --- a/src/reader.cpp +++ b/src/reader.cpp @@ -1,3 +1,4 @@ +#include #include #include @@ -89,12 +90,18 @@ struct ParserDeleter { }; typedef std::unique_ptr Parser; -Parser haven_init_parser() { +Parser haven_init_parser(long n_max) { Parser parser(readstat_parser_init()); readstat_set_metadata_handler(parser.get(), reader_metadata); readstat_set_note_handler(parser.get(), reader_note); readstat_set_variable_handler(parser.get(), reader_variable); - readstat_set_value_handler(parser.get(), reader_value); + // With n_max = 0 we only need the header. n_max = -1 means read all rows + if (n_max != 0) { + readstat_set_value_handler(parser.get(), reader_value); + } + if (n_max > 0) { + readstat_set_row_limit(parser.get(), n_max); + } readstat_set_value_label_handler(parser.get(), reader_value_label); readstat_set_error_handler(parser.get(), print_error); return parser; @@ -114,11 +121,6 @@ void haven_init_io(readstat_parser_t* parser, ReaderInput& builder_input) { builder_input.encoding == "" ? NULL : builder_input.encoding.c_str()); } -void haven_set_row_limit(readstat_parser_t* parser, long n) { - // readstat uses 0 to specify "all rows" but what we want is "minimal rows" - readstat_set_row_limit(parser, n == 0 ? 1 : n); -} - void haven_parse( FileExt ext, readstat_parser_t* parser, @@ -269,9 +271,9 @@ cpp11::list df_parse(FileExt ext, SEXP spec, const ReaderOptions& opts) { Reader builder(ext, user_na); builder.skipCols(opts.cols_skip); builder.setQuiet(opts.n_max == 0); + builder.setMaxRows(opts.n_max); - Parser parser = haven_init_parser(); - haven_set_row_limit(parser.get(), opts.n_max); + Parser parser = haven_init_parser(opts.n_max); readstat_set_row_offset(parser.get(), opts.rows_skip); std::unique_ptr input = make_input(spec, encoding); @@ -287,10 +289,6 @@ cpp11::list df_parse(FileExt ext, SEXP spec, const ReaderOptions& opts) { haven_parse(ext, parser.get(), *input, &builder); - if (opts.n_max >= 0) { - builder.limitRows(opts.n_max); // must enforce n_max = 0 - } - return builder.output(); } diff --git a/src/reader.h b/src/reader.h index d0d7d0bb..355f0a41 100644 --- a/src/reader.h +++ b/src/reader.h @@ -1,6 +1,7 @@ #ifndef HAVEN_READER_H #define HAVEN_READER_H +#include #include #include #include @@ -49,7 +50,8 @@ inline double haven_double_value(readstat_value_t value) { class Reader { FileVendor vendor_; - int nrows_, nrowsAlloc_; + int nrows_; // number of rows read so far + int nrowsAlloc_; // number of rows currently allocated in output vectors int ncols_; cpp11::writable::list output_; cpp11::writable::strings names_; @@ -62,32 +64,35 @@ class Reader { std::set colsSkip_; bool quiet_; + long n_max_; public: Reader(FileExt ext, bool user_na = false) : vendor_(extVendor(ext)), nrows_(0), + nrowsAlloc_(0), ncols_(0), output_(static_cast(0)), user_na_(user_na), - quiet_(false) {} + quiet_(false), + n_max_(-1) {} // Suppress metadata warnings, e.g. when only reading the header void setQuiet(bool quiet) { quiet_ = quiet; } + void setMaxRows(long n) { n_max_ = n; } + void skipCols(cpp11::integers cols) { std::set cols_set(cols.begin(), cols.end()); colsSkip_ = cols_set; } void setInfo(int obs_count, int var_count) { - if (obs_count < 0) { - // If unknown, start with 1e5, and use doubling strategy - nrowsAlloc_ = 1e5; - nrows_ = 0; - } else { - nrowsAlloc_ = nrows_ = obs_count; - } + // If unknown, start with 1e5, and use doubling strategy + nrowsAlloc_ = obs_count < 0 ? 1e5 : obs_count; + // But never allocate more rows than what we request + if (n_max_ >= 0 && n_max_ < nrowsAlloc_) + nrowsAlloc_ = n_max_; if (var_count < 1) { return; // sas7bcat has var_count = 0 @@ -158,8 +163,9 @@ class Reader { VarType var_type = var_types_[var_index]; + // Grow by doubling, always ensuring progress if (obs_index >= nrowsAlloc_) - resizeCols(nrowsAlloc_ * 2); + resizeCols((obs_index + 1) * 2); if (obs_index >= nrows_) nrows_ = obs_index + 1; @@ -237,12 +243,6 @@ class Reader { } } - void limitRows(long n) { - if (nrows_ > n) { - nrows_ = n; - } - } - cpp11::list output() { if (nrows_ != nrowsAlloc_) resizeCols(nrows_);