Reduce memory usage of reads with small n_max

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.

src/reader.cpp

@@ -1,3 +1,4 @@
+#include <algorithm>
 #include <memory>
 #include <string>
 
@@ -89,12 +90,18 @@ struct ParserDeleter {
 };
 typedef std::unique_ptr<readstat_parser_t, ParserDeleter> 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<ReaderInput> 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();
 }
 

src/reader.h

@@ -1,6 +1,7 @@
 #ifndef HAVEN_READER_H
 #define HAVEN_READER_H
 
+#include <algorithm>
 #include <map>
 #include <set>
 #include <string>
@@ -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<int> 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<R_xlen_t>(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<int> 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_);