perf(index): optional uncompressed index

src/index_single_mode.cpp

@@ -1535,6 +1535,7 @@ std::array<uint8_t, kIndexHeaderSize> encodeIndexHeader(const IndexParamsHeader&
     put32(20, static_cast<uint32_t>(p.l));
     h[24] = p.hpc ? 1 : 0;
     h[25] = p.open ? 1 : 0;
+    h[26] = p.uncompressed ? 1 : 0;
     return h;
 }
 
@@ -1550,11 +1551,13 @@ bool readIndexHeader(const std::string& path, IndexParamsHeader& out) {
     out.l = static_cast<int32_t>(get32(20));
     out.hpc = h[24] != 0;
     out.open = h[25] != 0;
+    out.uncompressed = h[26] != 0;
     return true;
 }
 }  // namespace index_single_mode
 
-void index_single_mode::IndexBuilder::writeIndex(const std::string& path, int numThreads, int zstdLevel) {
+void index_single_mode::IndexBuilder::writeIndex(const std::string& path, int numThreads, int zstdLevel,
+                                                 bool uncompressed) {
     output::step("Serializing index...");
 
     kj::Array<capnp::word> flatArray = capnp::messageToFlatArray(outMessage);
@@ -1572,14 +1575,27 @@ void index_single_mode::IndexBuilder::writeIndex(const std::string& path, int nu
     ph.l = indexBuilder.getL();
     ph.hpc = indexBuilder.getHpc();
     ph.open = indexBuilder.getOpen();
+    ph.uncompressed = uncompressed;
     const auto header = encodeIndexHeader(ph);
 
-    // Many independent 64 MB frames so the index inflates in parallel on load.
-    constexpr size_t kIndexFrameSize = 64ull * 1024 * 1024;
-    if (!panmap_zstd::compressToFile(data, dataSize, path, zstdLevel, numThreads, kIndexFrameSize,
-                                     header.data(), header.size())) {
-        output::error("failed to write compressed index to {}", path);
-        std::exit(1);
+    if (uncompressed) {
+        // Store the raw capnp bytes so load() can mmap the file and hand them straight to
+        // capnp (zero-copy, no decompression pass) -- bigger on disk, faster to load.
+        std::ofstream out(path, std::ios::binary | std::ios::trunc);
+        out.write(reinterpret_cast<const char*>(header.data()), header.size());
+        out.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(dataSize));
+        if (!out) {
+            output::error("failed to write uncompressed index to {}", path);
+            std::exit(1);
+        }
+    } else {
+        // Many independent 64 MB frames so the index inflates in parallel on load.
+        constexpr size_t kIndexFrameSize = 64ull * 1024 * 1024;
+        if (!panmap_zstd::compressToFile(data, dataSize, path, zstdLevel, numThreads, kIndexFrameSize,
+                                         header.data(), header.size())) {
+            output::error("failed to write compressed index to {}", path);
+            std::exit(1);
+        }
     }
 
     output::done("Index written to " + path);

src/index_single_mode.hpp

@@ -31,6 +31,10 @@ constexpr size_t kIndexHeaderSize = 32;
 struct IndexParamsHeader {
     int32_t k = 0, s = 0, t = 0, l = 0;
     bool hpc = false, open = false;
+    // When set, the payload after the header is raw (unframed) capnp bytes rather than
+    // ZSTD frames, so load can mmap it and hand the bytes straight to capnp (zero-copy,
+    // no decompression pass). Bigger on disk, faster to load.
+    bool uncompressed = false;
 };
 std::array<uint8_t, kIndexHeaderSize> encodeIndexHeader(const IndexParamsHeader& p);
 // Reads the header from the start of `path`. Returns false if absent / not this format.
@@ -211,7 +215,8 @@ class IndexBuilder {
     // Compute 4x4 substitution spectrum from tree mutations and store in index
     void computeSubstitutionSpectrum();
 
-    void writeIndex(const std::string& path, int numThreads = 0, int zstdLevel = 7);
+    void writeIndex(const std::string& path, int numThreads = 0, int zstdLevel = 7,
+                    bool uncompressed = false);
 
     // Parameter getters (for testing)
     int getK() const { return k_; }

src/main.cpp

@@ -7,6 +7,7 @@
 
 #include <boost/program_options.hpp>
 #include <boost/filesystem.hpp>
+#include <boost/iostreams/device/mapped_file.hpp>
 #include <boost/iostreams/filtering_streambuf.hpp>
 #include <boost/iostreams/filter/lzma.hpp>
 #include <boost/algorithm/string.hpp>
@@ -126,6 +127,7 @@ struct Config {
     // Metagenomic options
     bool indexPacked = false;
     bool readPacked = false;
+    bool indexUncompressed = false;
     bool noProgress = false;
     size_t topOc = 1000;
     uint32_t maskReads = 0;
@@ -196,22 +198,40 @@ struct Config {
 class IndexReader : public ::capnp::MessageReader {
    public:
     std::vector<uint8_t> data;
+    // Held for the uncompressed path; the mmap must outlive `reader` (declared after it,
+    // so it is destroyed first), which points directly into the mapped bytes.
+    boost::iostreams::mapped_file_source mmapFile;
     std::unique_ptr<::capnp::FlatArrayMessageReader> reader;
 
     explicit IndexReader(const std::string& path, int numThreads = 0) : ::capnp::MessageReader(makeOptions()) {
-        // New-format indexes carry a small uncompressed param header before the frames;
+        // New-format indexes carry a small uncompressed param header before the payload;
         // skip it. Old-format indexes (no header) decompress from offset 0.
         index_single_mode::IndexParamsHeader ph;
-        const size_t dataOffset =
-            index_single_mode::readIndexHeader(path, ph) ? index_single_mode::kIndexHeaderSize : 0;
-        if (!panmap_zstd::decompressFromFile(path, data, numThreads, dataOffset)) {
-            throw std::runtime_error("Failed to decompress index: " + path);
+        const bool hasHeader = index_single_mode::readIndexHeader(path, ph);
+        const size_t dataOffset = hasHeader ? index_single_mode::kIndexHeaderSize : 0;
+
+        const capnp::word* words = nullptr;
+        size_t numWords = 0;
+        if (hasHeader && ph.uncompressed) {
+            // Uncompressed index: mmap and hand the bytes straight to capnp (zero-copy,
+            // no decompression pass). dataOffset (32) is 8-aligned within the page-aligned
+            // mapping, so the word cast is properly aligned.
+            mmapFile.open(path);
+            if (!mmapFile.is_open()) {
+                throw std::runtime_error("Failed to mmap index: " + path);
+            }
+            words = reinterpret_cast<const capnp::word*>(mmapFile.data() + dataOffset);
+            numWords = (mmapFile.size() - dataOffset) / sizeof(capnp::word);
+        } else {
+            if (!panmap_zstd::decompressFromFile(path, data, numThreads, dataOffset)) {
+                throw std::runtime_error("Failed to decompress index: " + path);
+            }
+            words = reinterpret_cast<const capnp::word*>(data.data());
+            numWords = data.size() / sizeof(capnp::word);
         }
 
         reader = std::make_unique<::capnp::FlatArrayMessageReader>(
-            kj::ArrayPtr<const capnp::word>(reinterpret_cast<const capnp::word*>(data.data()),
-                                            data.size() / sizeof(capnp::word)),
-            makeOptions());
+            kj::ArrayPtr<const capnp::word>(words, numWords), makeOptions());
     }
 
     kj::ArrayPtr<const capnp::word> getSegment(uint id) override { return reader->getSegment(id); }
@@ -417,7 +437,7 @@ bool buildIndex(const Config& cfg) {
                                             cfg.extentGuard);
     builder.buildIndexParallel(cfg.threads);
     builder.computeSubstitutionSpectrum();
-    builder.writeIndex(cfg.index, cfg.threads, cfg.zstdLevel);
+    builder.writeIndex(cfg.index, cfg.threads, cfg.zstdLevel, cfg.indexUncompressed);
 
     auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - t0).count();
     output::done("index", cfg.index, fmt::format("{} nodes", output::fmt_count(tg->trees[0].allNodes.size())), ms);
@@ -1865,6 +1885,9 @@ int main(int argc, char** argv) {
     metagenomic.add_options()(
         "index-packed", po::bool_switch(&cfg.indexPacked), "Build packed capnp message (default false)")(
         "read-packed", po::bool_switch(&cfg.readPacked), "Read packed capnp message (default false)")(
+        "index-uncompressed", po::bool_switch(&cfg.indexUncompressed),
+        "Store the index uncompressed so it is mmap'd on load (no decompression pass; "
+        "larger on disk, faster to load; default false)")(
         "no-progress", po::bool_switch(&cfg.noProgress), "Disable progress bars");
 
     po::options_description em("Metagenomic: EM");