Read the external id map alongside the mmapped CSR

nsparse/id_map_index.cpp

@@ -9,7 +9,15 @@
 
 #include "nsparse/id_map_index.h"
 
+#include <cstdint>
+#include <filesystem>
+#include <fstream>
+#include <limits>
 #include <memory>
+#include <stdexcept>
+#include <string>
+#include <utility>
+#include <vector>
 
 #include "nsparse/id_selector.h"
 #include "nsparse/io/index_io.h"
@@ -78,6 +86,96 @@ void IDMapIndex::add_with_ids(idx_t n, const idx_t* indptr,
         external_to_internal_[ids[i]] = old_size + i;
     }
 }
+
+std::vector<idx_t> IDMapIndex::read_id_file(const char* id_path) {
+    check_if_file_valid(id_path, "id map file");
+
+    std::ifstream in(id_path, std::ios::binary);
+    if (!in.is_open()) {
+        throw std::runtime_error(std::string("cannot open id map file: ") +
+                                 id_path);
+    }
+
+    int64_t count = 0;
+    in.read(reinterpret_cast<char*>(&count), sizeof(count));
+    if (!in) {
+        throw std::runtime_error(std::string("truncated id map file: ") +
+                                 id_path);
+    }
+    if (count < 0) {
+        throw std::invalid_argument(std::string("negative id map count in: ") +
+                                    id_path);
+    }
+    const auto map_size = static_cast<size_t>(count);
+
+    // Guard the byte-size arithmetic against wraparound before relying on it.
+    if (map_size > (std::numeric_limits<size_t>::max() - sizeof(int64_t)) /
+                       sizeof(idx_t)) {
+        throw std::invalid_argument(std::string("id map count is too large: ") +
+                                    id_path);
+    }
+
+    // Reject a truncated or oversized file up front.
+    const size_t expected_bytes = sizeof(int64_t) + map_size * sizeof(idx_t);
+    if (std::filesystem::file_size(id_path) != expected_bytes) {
+        throw std::invalid_argument(
+            std::string(
+                "id map file size does not match its count (expected ") +
+            std::to_string(expected_bytes) + " bytes): " + id_path);
+    }
+
+    std::vector<idx_t> internal_to_external(map_size);
+    if (map_size > 0) {
+        in.read(reinterpret_cast<char*>(internal_to_external.data()),
+                static_cast<std::streamsize>(map_size * sizeof(idx_t)));
+        if (!in) {
+            throw std::runtime_error(std::string("truncated id map file: ") +
+                                     id_path);
+        }
+    }
+    return internal_to_external;
+}
+
+void IDMapIndex::read_csr_and_ids(const char* csr_path, const char* id_path,
+                                  Residency residency) {
+    if (delegate_ == nullptr) {
+        throw std::logic_error("IDMapIndex has no delegate index");
+    }
+    check_if_file_valid(csr_path, "csr file");
+
+    // Fully validate and load the id file BEFORE ingesting the CSR, so a
+    // missing/malformed/truncated id file leaves this index untouched. Only the
+    // count-vs-CSR-row check below can fail once the delegate has ingested; on
+    // ANY throw from this method the half-built index must be discarded.
+    std::vector<idx_t> internal_to_external = read_id_file(id_path);
+
+    // The id file is known-good; now ingest the vectors (borrowed from the
+    // mapping when residency == kMmap). Afterward num_vectors() reflects the
+    // CSR rows.
+    delegate_->read_csr(csr_path, residency);
+
+    // The map is row-aligned with the CSR, so its count must equal the vectors
+    // the delegate just ingested.
+    const size_t delegate_size = delegate_->num_vectors();
+    if (internal_to_external.size() != delegate_size) {
+        throw std::invalid_argument(
+            "id map count (" + std::to_string(internal_to_external.size()) +
+            ") does not match the CSR vector count (" +
+            std::to_string(delegate_size) + "): " + id_path);
+    }
+
+    set_id_map(std::move(internal_to_external));
+}
+
+void IDMapIndex::set_id_map(std::vector<idx_t>&& internal_to_external) {
+    internal_to_external_ = std::move(internal_to_external);
+    external_to_internal_.clear();
+    external_to_internal_.reserve(internal_to_external_.size());
+    for (size_t i = 0; i < internal_to_external_.size(); ++i) {
+        external_to_internal_[internal_to_external_[i]] = static_cast<idx_t>(i);
+    }
+}
+
 void IDMapIndex::write_index(IOWriter* io_writer) {
     // Write internal_to_external_ vector
     size_t map_size = internal_to_external_.size();
@@ -89,23 +187,18 @@ void IDMapIndex::write_index(IOWriter* io_writer) {
     nsparse::detail::write_index(delegate_.get(), io_writer, true);
 }
 
-void IDMapIndex::read_index(IOReader* io_reader, const IndexHeader& header,
+void IDMapIndex::read_index(IOReader* io_reader, const IndexHeader& /*header*/,
                             int io_flags) {
-    // Read internal_to_external_ vector
+    // Read the id map into a local vector, then load the delegate.
     size_t map_size = 0;
     io_reader->read(&map_size, sizeof(size_t), 1);
-    internal_to_external_.resize(map_size);
+    std::vector<idx_t> internal_to_external(map_size);
     if (map_size > 0) {
-        io_reader->read(internal_to_external_.data(), sizeof(idx_t), map_size);
+        io_reader->read(internal_to_external.data(), sizeof(idx_t), map_size);
     }
 
     delegate_.reset(nsparse::detail::read_index(io_reader, true, io_flags));
 
-    // Rebuild external_to_internal_ from internal_to_external_
-    external_to_internal_.clear();
-    external_to_internal_.reserve(map_size);
-    for (size_t i = 0; i < map_size; ++i) {
-        external_to_internal_[internal_to_external_[i]] = static_cast<idx_t>(i);
-    }
+    set_id_map(std::move(internal_to_external));
 }
 }  // namespace nsparse

nsparse/id_map_index.h

@@ -101,6 +101,10 @@ class IDMapIndex : public Index, public IndexIO {
 
     void add_with_ids(idx_t n, const idx_t* indptr, const term_t* indices,
                       const float* values, const idx_t* ids) override;
+
+    void read_csr_and_ids(const char* csr_path, const char* id_path,
+                          Residency residency = Residency::kInMemory);
+
     [[nodiscard]] uint32_t format_version() const override {
         return kFormatVersion;
     }
@@ -109,6 +113,10 @@ class IDMapIndex : public Index, public IndexIO {
                     int io_flags = 0) override;
 
 private:
+    std::vector<idx_t> read_id_file(const char* id_path);
+
+    void set_id_map(std::vector<idx_t>&& internal_to_external);
+
     // Owns the wrapped delegate index. Using unique_ptr ensures the delegate is
     // freed when the IDMapIndex is destroyed (previously a raw pointer with a
     // defaulted destructor, which leaked the delegate and everything it owned).

nsparse/utils/checks.h

@@ -11,8 +11,10 @@
 #define COMMON_H
 
 #include <cstddef>
+#include <filesystem>
 #include <limits>
 #include <stdexcept>
+#include <string>
 
 namespace nsparse {
 
@@ -24,6 +26,16 @@ T* throw_if_null(T* ptr, const char* msg = "unexpected nullptr") {
     return ptr;
 }
 
+// Rejects a file path that is null or does not exist, so a caller gets a clear
+// error before trying to open it. `what` names the file in the message.
+inline void check_if_file_valid(const char* path, const char* what = "file") {
+    throw_if_null(path, "file path must not be null");
+    if (!std::filesystem::exists(path)) {
+        throw std::invalid_argument(std::string(what) +
+                                    " does not exist: " + path);
+    }
+}
+
 template <typename T>
 T throw_if_not_positive(T value, const char* msg = "value must be positive") {
     if (value <= 0) {

python_tests/test_seismic_index.py

@@ -67,6 +67,51 @@ def test_with_id_map(corpus, queries, oracle, doc_ids):
     assert recall_at_k(labels, want_external) >= RECALL_FLOOR
 
 
+def _write_interchange_csr(path, corpus):
+    """Corpus as an interchange CSR: int64 header {n, dim, nnz}, int64 indptr,
+    int32 indices, float32 values -- the layout nsparse.convert consumes."""
+    with open(path, "wb") as out:
+        np.array(
+            [corpus.n, corpus.dim, corpus.indices.size], dtype=np.int64
+        ).tofile(out)
+        corpus.indptr.astype(np.int64).tofile(out)
+        corpus.indices.astype(np.int32).tofile(out)
+        corpus.values.astype(np.float32).tofile(out)
+
+
+def _write_id_map(path, external_ids):
+    """The id-map file read_csr_and_ids reads: int64 count, then int32 ids,
+    row-aligned with the CSR."""
+    with open(path, "wb") as out:
+        np.array([external_ids.size], dtype=np.int64).tofile(out)
+        external_ids.astype(np.int32).tofile(out)
+
+
+def test_id_map_from_csr_and_id_files(corpus, queries, oracle, doc_ids, tmp_path):
+    """read_csr_and_ids builds an idmap from a native CSR (borrowed via mmap)
+    plus a separate id file -- the memory-saving build path -- and must return
+    the caller's external ids, matching the in-RAM add_with_ids path."""
+    interchange = tmp_path / "corpus.csr"
+    native = tmp_path / "corpus.mcsr"
+    id_file = tmp_path / "ids.bin"
+    _write_interchange_csr(interchange, corpus)
+    nsparse.convert(str(interchange), str(native))
+    _write_id_map(id_file, doc_ids)
+
+    index = nsparse.index_factory(corpus.dim, f"idmap,{SPEC}")
+    index.read_csr_and_ids(str(native), str(id_file), nsparse.Residency_kMmap)
+    index.build()
+
+    _, labels = search(index, queries)
+    returned = labels[labels >= 0]
+    assert returned.size > 0, "every query should return at least one hit"
+    assert np.isin(returned, doc_ids).all(), "returned ids must be caller ids"
+
+    want_labels, _ = oracle
+    want_external = np.where(want_labels >= 0, doc_ids[want_labels], -1)
+    assert recall_at_k(labels, want_external) >= RECALL_FLOOR
+
+
 def test_exact_match(index, queries, oracle):
     """An enumerable selector of size <= k switches search to the exact path.
 

tests/csr_interchange_test_util.h

@@ -18,15 +18,17 @@
 #include <system_error>
 #include <vector>
 
+#include "nsparse/types.h"
+
 // Shared helpers for the mmap-CSR build path, used by both the regular and the
 // disk-resident index suites: write a corpus as an interchange CSR (the layout
 // csr_layout::convert consumes) and manage the interchange + native temp files.
 namespace nsparse::csr_test {
 
 // Writes a corpus as an interchange CSR: int64 header {rows, num_cols, nnz},
 // int64 indptr[rows + 1], int32 indices[nnz], float values[nnz]. Templated on
-// the corpus struct (any type exposing .n / .indptr / .indices / .values), so it
-// serves any test corpus. The values are written verbatim, so a convert +
+// the corpus struct (any type exposing .n / .indptr / .indices / .values), so
+// it serves any test corpus. The values are written verbatim, so a convert +
 // read_csr(kMmap) build sees the exact same vectors as add().
 template <class Corpus>
 void write_interchange_csr(const std::string& path, const Corpus& c,
@@ -47,12 +49,26 @@ void write_interchange_csr(const std::string& path, const Corpus& c,
               static_cast<std::streamsize>(c.values.size() * sizeof(float)));
 }
 
+// Writes the id-map file that IDMapIndex::read_csr_and_ids reads:
+// [int64 count][idx_t external_id x count]. Row-aligned with the CSR, so
+// external_ids[i] is the external id of CSR row i.
+inline void write_id_map_file(const std::string& path,
+                              const std::vector<idx_t>& external_ids) {
+    std::ofstream out(path, std::ios::binary);
+    const int64_t count = static_cast<int64_t>(external_ids.size());
+    out.write(reinterpret_cast<const char*>(&count), sizeof(count));
+    out.write(
+        reinterpret_cast<const char*>(external_ids.data()),
+        static_cast<std::streamsize>(external_ids.size() * sizeof(idx_t)));
+}
+
 // An interchange CSR temp file and the native path convert writes it to, both
 // removed on destruction.
 class TempCsrFiles {
 public:
     explicit TempCsrFiles(const std::string& stem)
-        : interchange_(std::filesystem::temp_directory_path() / (stem + ".csr")),
+        : interchange_(std::filesystem::temp_directory_path() /
+                       (stem + ".csr")),
           native_(std::filesystem::temp_directory_path() / (stem + ".mcsr")) {
         std::error_code ignored;
         std::filesystem::remove(interchange_, ignored);