Build the quantized indexes from an mmapped codes CSR

nsparse/disk_seismic_index_base.h

@@ -64,8 +64,10 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO {
 
     // --- Hooks the concrete indexes implement. ---
 
-    // Stored value width in bytes: 4 (float) or 1/2 (quantized codes).
-    [[nodiscard]] virtual size_t code_element_size() const = 0;
+    // Stored value width in bytes: 4 (float) or 1/2 (quantized codes). Kept
+    // pure (overriding MmapIndex's default) so each concrete disk type must
+    // state its width.
+    [[nodiscard]] size_t code_element_size() const override = 0;
 
     // Encode nnz float values to the stored width, returning a pointer to
     // code_element_size()-byte-per-value data. `scratch` backs the result when

nsparse/disk_seismic_scalar_quantized_index.cpp

@@ -77,16 +77,6 @@ DiskSeismicScalarQuantizedIndex::DiskSeismicScalarQuantizedIndex(
     SeismicClusterParameters parameter, int dim)
     : DiskSeismicIndexBase(dim, parameter), sq_(quantizer_type, vmin, vmax) {}
 
-void DiskSeismicScalarQuantizedIndex::read_csr(const char* file_path,
-                                               Residency residency) {
-    if (residency == Residency::kMmap) {
-        throw std::invalid_argument(
-            "mmap residency is not available for a quantized index: a mapped "
-            "CSR is borrowed as float, and this index searches over codes");
-    }
-    MmapIndex::read_csr(file_path, residency);
-}
-
 size_t DiskSeismicScalarQuantizedIndex::code_element_size() const {
     return sq_.bytes_per_value();
 }

nsparse/disk_seismic_scalar_quantized_index.h

@@ -68,12 +68,6 @@ class DiskSeismicScalarQuantizedIndex : public DiskSeismicIndexBase {
                                                        const char* index_file,
                                                        size_t pos);
 
-    // Only the copying residency. A mapped CSR is borrowed at the width it was
-    // written in, which is float, whereas this index searches over codes: the
-    // values have to pass through add() to be quantized.
-    void read_csr(const char* file_path,
-                  Residency residency = Residency::kInMemory) override;
-
 private:
     [[nodiscard]] uint32_t format_version() const override {
         return kFormatVersion;

nsparse/mmap_index.h

@@ -72,10 +72,15 @@ class MmapIndex : public Index {
     // so they are always gone before the spill is released.
     detail::ClusteredListsSpill batch_spill_;
 
+    // The stored value width, in bytes, which is also what read_mcsr borrows a
+    // native CSR's values at: float (the default) for the unquantized types, or
+    // the quantizer's code width for a quantizing index, so the same mapped-CSR
+    // path yields codes borrowed in place rather than floats. Non-pure because
+    // MmapIndex is instantiated directly (e.g. TestMmapIndex); the disk family
+    // re-declares it pure so each concrete type must state its width.
+    [[nodiscard]] virtual size_t code_element_size() const { return U32; }
+
 private:
-    // Values are borrowed at their stored width, so a quantizing index cannot
-    // use this path.
-    static constexpr size_t kMmapElementSize = U32;
     bool is_mmap_index_ = false;
 
     // Points at the native layout (csr_layout.h) rather than copying it. Terms
@@ -114,10 +119,17 @@ class MmapIndex : public Index {
                 file_path);
         }
 
+        // The value width the borrow reinterprets in place: float for the
+        // unquantized types, a quantizer's code width for a quantizing one. A
+        // file written at a different width fails the size check below, so a
+        // codes CSR handed to a float index (or vice versa) is rejected rather
+        // than misread.
+        const size_t element_size = code_element_size();
+
         const size_t indptr_size = static_cast<size_t>(num_rows) + 1;
         const auto nnz_size = static_cast<size_t>(nnz);
         if (file.size() !=
-            csr_layout::native_file_size(indptr_size, nnz_size)) {
+            csr_layout::native_file_size(indptr_size, nnz_size, element_size)) {
             throw std::invalid_argument(
                 std::string("CSR file is not in the native layout (convert it "
                             "with csr_layout::convert): ") +
@@ -128,17 +140,19 @@ class MmapIndex : public Index {
         const auto* indices = cursor.read_array<term_t>(nnz_size);
         cursor.skip(csr_layout::native_values_offset(indptr_size, nnz_size) -
                     cursor.pos());
-        const auto* values = cursor.read_array<float>(nnz_size);
+        // Borrowed as raw bytes at the value width. The values offset is padded
+        // to alignof(float), which satisfies any code width, so the in-place
+        // reinterpret downstream stays aligned.
+        const auto* values = cursor.read_array<uint8_t>(nnz_size * element_size);
 
         // Validates before borrowing, so a corrupt file throws here rather
         // than faulting during search.
         auto vectors =
             std::make_unique<SparseVectors>(SparseVectors::map_vectors(
-                {.element_size = kMmapElementSize,
+                {.element_size = element_size,
                  .dimension = static_cast<size_t>(dimension_)},
-                indptr, indptr_size, indices, nnz_size,
-                reinterpret_cast<const uint8_t*>(values),
-                nnz_size * kMmapElementSize));
+                indptr, indptr_size, indices, nnz_size, values,
+                nnz_size * element_size));
 
         // Committed last, so a rejected file leaves the index untouched.
         mapped_file_ = std::move(file);

nsparse/seismic_scalar_quantized_index.cpp

@@ -146,16 +146,6 @@ SeismicScalarQuantizedIndex::SeismicScalarQuantizedIndex(
       sq_(quantizer_type, vmin, vmax),
       cluster_parameter_(parameter) {}
 
-void SeismicScalarQuantizedIndex::read_csr(const char* file_path,
-                                           Residency residency) {
-    if (residency == Residency::kMmap) {
-        throw std::invalid_argument(
-            "mmap residency is not available for a quantized index: a mapped "
-            "CSR is borrowed as float, and this index searches over codes");
-    }
-    MmapIndex::read_csr(file_path, residency);
-}
-
 void SeismicScalarQuantizedIndex::add(idx_t n, const idx_t* indptr,
                                       const term_t* indices,
                                       const float* values) {

nsparse/seismic_scalar_quantized_index.h

@@ -63,13 +63,13 @@ class SeismicScalarQuantizedIndex : public MmapIndex, public IndexIO {
                                                    const char* index_file,
                                                    size_t pos);
 
-    // Only the copying residency. A mapped CSR is borrowed at the width it was
-    // written in, which is float, whereas this index searches over codes: the
-    // values have to pass through add() to be quantized.
-    void read_csr(const char* file_path,
-                  Residency residency = Residency::kInMemory) override;
-
 private:
+    // Stored value width: this index's code width, so read_mcsr borrows a codes
+    // CSR in place (rather than a float one) via MmapIndex::code_element_size.
+    [[nodiscard]] size_t code_element_size() const override {
+        return sq_.bytes_per_value();
+    }
+
     // interfaces of IndexIO
     [[nodiscard]] uint32_t format_version() const override {
         return kFormatVersion;

nsparse/utils/csr_layout.h

@@ -22,12 +22,16 @@
 //   int64  indptr[r + 1]  idx_t  indptr[r + 1]
 //   int32  indices[nnz]   term_t indices[nnz]
 //                         <padding to 4 bytes>
-//   float  values[nnz]    float  values[nnz]
+//   float  values[nnz]    value  values[nnz]
 //
 // Interchange is what scipy writes and what Index::read_csr narrows while
 // copying. Native stores in-memory widths so MmapIndex can borrow the arrays in
 // place; values is padded because term_t is 2 bytes, and an odd nnz would leave
-// the floats misaligned for the mapped reader's in-place reinterpret.
+// the floats misaligned for the mapped reader's in-place reinterpret. The
+// native value width is the index's own: 4-byte float for the unquantized
+// types, or 1-/2-byte scalar codes for a quantizing one (see native_file_size's
+// element_size). The pad is to alignof(float), which also satisfies the smaller
+// code widths.
 //
 // File size cannot tell the layouts apart (rows=1, nnz=4 is 64 bytes either
 // way), so conversion is explicit and the results are kept apart by suffix.
@@ -53,8 +57,12 @@ constexpr size_t native_values_offset(size_t indptr_size, size_t nnz) {
     return unaligned + padding(unaligned);
 }
 
-constexpr size_t native_file_size(size_t indptr_size, size_t nnz) {
-    return native_values_offset(indptr_size, nnz) + nnz * sizeof(float);
+// `element_size` is the native value width: sizeof(float) for the unquantized
+// types (the default), or the quantizer's byte width (1 or 2) for a quantizing
+// index whose CSR holds codes rather than floats.
+constexpr size_t native_file_size(size_t indptr_size, size_t nnz,
+                                  size_t element_size = sizeof(float)) {
+    return native_values_offset(indptr_size, nnz) + nnz * element_size;
 }
 
 std::string native_path(const std::string& path);

tests/csr_interchange_test_util.h

@@ -14,11 +14,13 @@
 #include <cstdint>
 #include <filesystem>
 #include <fstream>
+#include <stdexcept>
 #include <string>
 #include <system_error>
 #include <vector>
 
 #include "nsparse/types.h"
+#include "nsparse/utils/csr_layout.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
@@ -49,6 +51,47 @@ void write_interchange_csr(const std::string& path, const Corpus& c,
               static_cast<std::streamsize>(c.values.size() * sizeof(float)));
 }
 
+// Writes a corpus of pre-quantized codes as a NATIVE CSR -- the layout read_mcsr
+// borrows when a quantizing index maps it: int64 header {rows, num_cols, nnz},
+// idx_t indptr[rows+1], term_t indices[nnz], pad to alignof(float), then
+// `element_size`-byte codes[nnz]. `codes` holds nnz*element_size bytes,
+// row-aligned with `indices`. This is the code-width analog of
+// csr_layout::convert's output, written directly since the codes path has no
+// interchange form.
+inline void write_native_codes_csr(const std::string& path,
+                                   const std::vector<idx_t>& indptr,
+                                   const std::vector<term_t>& indices,
+                                   const std::vector<uint8_t>& codes,
+                                   int64_t num_cols, size_t element_size) {
+    // The declared width must describe the buffer, or the file would not match
+    // the layout read_mcsr validates -- a mismatch here is a test-authoring bug.
+    if (codes.size() != indices.size() * element_size) {
+        throw std::invalid_argument(
+            "write_native_codes_csr: codes size does not match indices * "
+            "element_size");
+    }
+    std::ofstream out(path, std::ios::binary);
+    const std::array<int64_t, 3> header = {
+        static_cast<int64_t>(indptr.size()) - 1, num_cols,
+        static_cast<int64_t>(indices.size())};
+    out.write(reinterpret_cast<const char*>(header.data()),
+              header.size() * sizeof(int64_t));
+    out.write(reinterpret_cast<const char*>(indptr.data()),
+              static_cast<std::streamsize>(indptr.size() * sizeof(idx_t)));
+    out.write(reinterpret_cast<const char*>(indices.data()),
+              static_cast<std::streamsize>(indices.size() * sizeof(term_t)));
+    const size_t values_pos = csr_layout::kHeaderBytes +
+                              indptr.size() * sizeof(idx_t) +
+                              indices.size() * sizeof(term_t);
+    const std::array<char, alignof(float)> pad{};
+    if (const size_t pad_bytes = csr_layout::padding(values_pos);
+        pad_bytes > 0) {
+        out.write(pad.data(), static_cast<std::streamsize>(pad_bytes));
+    }
+    out.write(reinterpret_cast<const char*>(codes.data()),
+              static_cast<std::streamsize>(codes.size()));
+}
+
 // 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.

tests/disk_seismic_scalar_quantized_index_test.cpp

@@ -331,4 +331,62 @@ TEST(DiskSeismicSQIndex, FactoryCreatesIt) {
               QuantizerType::QT_16bit);
 }
 
+// Building the disk quantized index from a native codes CSR borrowed via mmap
+// must match the add()-fed build: the upstream writes codes at the quantizer's
+// width, read_csr(kMmap) borrows them, and build -> persist -> mmap-reload ->
+// search is bit-exact to feeding the same corpus (quantized by add()) then
+// building. Both see identical codes because add() runs the same
+// ScalarQuantizer::encode the test wrote. Also asserts a width-mismatched codes
+// file is rejected rather than misread.
+TEST(DiskSeismicSQIndex, MmapCodesCsrBuildMatchesAddBuild) {
+    const CSR corpus = make_corpus(1500, /*seed=*/1);
+    const CSR queries = make_corpus(40, /*seed=*/2);
+    DiskSeismicSearchParameters params(/*cut=*/25, /*k_prime=*/32);
+
+    // Reference: float corpus fed through add() (which quantizes), reloaded.
+    DiskSeismicScalarQuantizedIndex added(QuantizerType::QT_8bit, 0.0F, 1.0F,
+                                          cluster_params(), kDim);
+    add_corpus(added, corpus);
+    added.build();
+    TempIndexFile added_file("nsparse_dssq_addbuild.idx");
+    write_index(&added, added_file.c_str());
+    std::unique_ptr<Index> added_mapped(
+        read_index(added_file.c_str(), IndexIoFlag::kUseMmap));
+    ASSERT_NE(added_mapped, nullptr);
+    const ScoreIds fresh = search_all(*added_mapped, queries, 10, &params);
+
+    // Under test: the same corpus pre-quantized to codes, borrowed via mmap.
+    const ScalarQuantizer sq(QuantizerType::QT_8bit, 0.0F, 1.0F);
+    const size_t element_size = sq.bytes_per_value();
+    std::vector<uint8_t> codes(corpus.indices.size() * element_size);
+    sq.encode(corpus.values.data(), codes.data(), corpus.indices.size());
+    TempCsrFiles csr("nsparse_dssq_codes");
+    write_native_codes_csr(csr.native(), corpus.indptr, corpus.indices, codes,
+                           kDim, element_size);
+
+    DiskSeismicScalarQuantizedIndex mapped(QuantizerType::QT_8bit, 0.0F, 1.0F,
+                                           cluster_params(), kDim);
+    mapped.read_csr(csr.native().c_str(), Residency::kMmap);
+    ASSERT_EQ(mapped.num_vectors(), static_cast<size_t>(corpus.n));
+    mapped.build();
+    TempIndexFile built_file("nsparse_dssq_mmapbuild.idx");
+    write_index(&mapped, built_file.c_str());
+    std::unique_ptr<Index> built_mapped(
+        read_index(built_file.c_str(), IndexIoFlag::kUseMmap));
+    ASSERT_NE(built_mapped, nullptr);
+    EXPECT_EQ(built_mapped->num_vectors(), static_cast<size_t>(corpus.n));
+    expect_same_results(search_all(*built_mapped, queries, 10, &params), fresh);
+
+    // A codes file whose width does not match the index's quantizer is rejected
+    // by the native-layout size check: 16-bit-wide values fed to an 8-bit index.
+    std::vector<uint8_t> wide_codes(corpus.indices.size() * 2);
+    TempCsrFiles wrong("nsparse_dssq_wrongwidth");
+    write_native_codes_csr(wrong.native(), corpus.indptr, corpus.indices,
+                           wide_codes, kDim, /*element_size=*/2);
+    DiskSeismicScalarQuantizedIndex eight(QuantizerType::QT_8bit, 0.0F, 1.0F,
+                                          cluster_params(), kDim);
+    EXPECT_THROW(eight.read_csr(wrong.native().c_str(), Residency::kMmap),
+                 std::invalid_argument);
+}
+
 }  // namespace nsparse

tests/disk_seismic_test_util.h

@@ -87,6 +87,7 @@ inline void add_corpus(Index& index, const CSR& c) {
 // re-exported here so the disk suite reaches them via `disk_seismic_test`.
 using csr_test::TempCsrFiles;
 using csr_test::write_interchange_csr;
+using csr_test::write_native_codes_csr;
 
 // Index file removed on destruction. write_index/read_index take char*.
 class TempIndexFile {

tests/id_map_index_test.cpp

@@ -25,8 +25,10 @@
 #include "nsparse/io/buffered_io.h"
 #include "nsparse/io/index_io.h"
 #include "nsparse/seismic_index.h"
+#include "nsparse/seismic_scalar_quantized_index.h"
 #include "nsparse/types.h"
 #include "nsparse/utils/csr_layout.h"
+#include "nsparse/utils/scalar_quantizer.h"
 #include "tests/csr_interchange_test_util.h"
 
 namespace {
@@ -514,6 +516,55 @@ TEST(IDMapReadCsrAndId, MatchesAddWithIdsBuild) {
     EXPECT_TRUE(saw_external) << "labels should be translated to external ids";
 }
 
+// read_csr_and_ids over a QUANTIZED delegate: since CR-2 the quantized types
+// accept a mmapped (codes) CSR, so the id-map wrapper works over them too --
+// this used to throw. The upstream writes codes at the quantizer's width; the
+// build is bit-exact to add_with_ids feeding the same corpus (which add()
+// quantizes with the same ScalarQuantizer the test wrote).
+TEST(IDMapReadCsrAndId, MatchesAddWithIdsBuildQuantized) {
+    const Corpus corpus = make_corpus(300, kDim, /*seed=*/1);
+    const Corpus queries = make_corpus(20, kDim, /*seed=*/2);
+    const std::vector<idx_t> ids = make_external_ids(corpus.n);
+    constexpr int k = 10;
+    auto delegate = [] {
+        return new nsparse::SeismicScalarQuantizedIndex(
+            nsparse::QuantizerType::QT_8bit, 0.0F, 1.0F, kClusterParams, kDim);
+    };
+
+    // Reference: add_with_ids() feeds floats, which the delegate quantizes.
+    nsparse::IDMapIndex added(delegate());
+    added.add_with_ids(corpus.n, corpus.indptr.data(), corpus.indices.data(),
+                       corpus.values.data(), ids.data());
+    added.build();
+    const auto expected = search_corpus(added, queries, k);
+
+    // Under test: the same corpus pre-quantized to codes, borrowed via mmap.
+    const nsparse::ScalarQuantizer sq(nsparse::QuantizerType::QT_8bit, 0.0F,
+                                      1.0F);
+    const size_t element_size = sq.bytes_per_value();
+    std::vector<uint8_t> codes(corpus.indices.size() * element_size);
+    sq.encode(corpus.values.data(), codes.data(), corpus.indices.size());
+    nsparse::csr_test::TempCsrFiles csr("nsparse_idmap_sq_src");
+    nsparse::csr_test::write_native_codes_csr(csr.native(), corpus.indptr,
+                                             corpus.indices, codes, kDim,
+                                             element_size);
+    TempIdFile idfile("nsparse_idmap_sq_src.ids");
+    nsparse::csr_test::write_id_map_file(idfile.path(), ids);
+
+    nsparse::IDMapIndex mapped(delegate());
+    mapped.read_csr_and_ids(csr.native().c_str(), idfile.path().c_str(),
+                            nsparse::Residency::kMmap);
+    ASSERT_EQ(mapped.num_vectors(), static_cast<size_t>(corpus.n));
+    mapped.build();
+    const auto got = search_corpus(mapped, queries, k);
+
+    EXPECT_EQ(got.first, expected.first) << "external-id labels differ";
+    ASSERT_EQ(got.second.size(), expected.second.size());
+    for (size_t i = 0; i < got.second.size(); ++i) {
+        EXPECT_FLOAT_EQ(got.second[i], expected.second[i]) << "score at " << i;
+    }
+}
+
 // The id map is row-aligned with the CSR, so a count that disagrees with the
 // delegate's vector count is rejected.
 TEST(IDMapReadCsrAndId, CountMismatchThrows) {