Seismic search optimization

benchmarks/index_search_benchmark.cpp

@@ -235,10 +235,16 @@ BENCHMARK(BM_Seismic_Search)
 // Seismic Scalar Quantized Index
 // ---------------------------------------------------------------------------
 static void BM_SeismicSQ_Search(benchmark::State& state) {
-    auto& fix = IndexFixture::get(
-        "seismic_sq,quantizer=8bit|vmin=0.0|vmax=3.0|lambda=6000|beta=400|"
-        "alpha=0.4",
-        ".seismic_sq.dat");
+    // Quantizer width is configurable so int16 vs int8 can be compared
+    // (NSPARSE_SQ_BITS = "8bit" or "16bit"; default "8bit"). Each width caches
+    // to its own .dat suffix.
+    const char* bits_env = std::getenv("NSPARSE_SQ_BITS");
+    std::string bits =
+        (bits_env != nullptr && bits_env[0] != '\0') ? bits_env : "8bit";
+    std::string descriptor =
+        "seismic_sq,quantizer=" + bits +
+        "|vmin=0.0|vmax=3.0|lambda=6000|beta=400|alpha=0.4";
+    auto& fix = IndexFixture::get(descriptor, ".seismic_sq_" + bits + ".dat");
     const int k = static_cast<int>(state.range(0));
     const int n_queries = static_cast<int>(fix.query.nrow);
 

nsparse/cluster/inverted_list_clusters.cpp

@@ -10,7 +10,9 @@
 #include "nsparse/cluster/inverted_list_clusters.h"
 
 #include <algorithm>
+#include <cassert>
 #include <cstdint>
+#include <limits>
 #include <span>
 #include <unordered_map>
 #include <vector>
@@ -22,12 +24,14 @@ namespace nsparse {
 namespace {
 
 /**
- * @brief Generage summary sparse vector for posting lists
+ * @brief Generate the per-cluster summary sparse vector (CSR) for a posting
+ *        list's clusters.
  *
  * @param vectors inverted index
- * @param group_of_doc_ids a list of posting list
+ * @param group_of_doc_ids flattened doc ids of all clusters
+ * @param offsets cluster boundaries into group_of_doc_ids
  * @param alpha prune ratio
- * @return SparseVectors
+ * @return SparseVectors  one summary vector per cluster
  */
 template <class T>
 SparseVectors summarize_(const SparseVectors* vectors,
@@ -39,7 +43,6 @@ SparseVectors summarize_(const SparseVectors* vectors,
     if (offsets.size() <= 1) {
         return summarized_vectors;
     }
-    const auto element_size = vectors->get_element_size();
     const auto& indptr_data = vectors->indptr_data();
     const auto& indices_data = vectors->indices_data();
     const auto& values_data = vectors->values_data();
@@ -117,28 +120,10 @@ InvertedListClusters::InvertedListClusters(
     }
 }
 
-InvertedListClusters::InvertedListClusters(const InvertedListClusters& other) {
-    docs_ = other.docs_;
-    offsets_ = other.offsets_;
-    if (other.summaries_ != nullptr) {
-        summaries_ = std::make_unique<SparseVectors>(*other.summaries_);
-    } else {
-        summaries_.reset();
-    }
-}
+InvertedListClusters::InvertedListClusters(const InvertedListClusters& other) =
+    default;
 InvertedListClusters& InvertedListClusters::operator=(
-    const InvertedListClusters& other) {
-    if (this != &other) {
-        docs_ = other.docs_;
-        offsets_ = other.offsets_;
-        if (other.summaries_ != nullptr) {
-            summaries_ = std::make_unique<SparseVectors>(*other.summaries_);
-        } else {
-            summaries_.reset();
-        }
-    }
-    return *this;
-}
+    const InvertedListClusters& other) = default;
 
 auto InvertedListClusters::get_docs(idx_t idx) const -> std::span<const idx_t> {
     return {docs_.data() + offsets_[idx],
@@ -147,19 +132,113 @@ auto InvertedListClusters::get_docs(idx_t idx) const -> std::span<const idx_t> {
 
 void InvertedListClusters::summarize(const SparseVectors* vectors,
                                      float alpha) {
-    if (summaries_ != nullptr) {
-        summaries_.reset();
-    }
     const auto element_size = vectors->get_element_size();
+    SparseVectors summaries;
     if (element_size == U32) {
-        summaries_ = std::make_unique<SparseVectors>(
-            std::move(summarize_<float>(vectors, docs_, offsets_, alpha)));
+        summaries = summarize_<float>(vectors, docs_, offsets_, alpha);
     } else if (element_size == U16) {
-        summaries_ = std::make_unique<SparseVectors>(
-            std::move(summarize_<uint16_t>(vectors, docs_, offsets_, alpha)));
+        summaries = summarize_<uint16_t>(vectors, docs_, offsets_, alpha);
+    } else {
+        summaries = summarize_<uint8_t>(vectors, docs_, offsets_, alpha);
+    }
+    build_transpose(summaries);
+}
+
+// Build the term-major (CSC) transpose directly from a per-cluster CSR summary.
+// The CSR summary is transient; only the transpose is retained.
+void InvertedListClusters::build_transpose(const SparseVectors& summaries) {
+    n_clusters_ = summaries.num_vectors();
+    element_size_ = summaries.get_element_size();
+    term_ids_.clear();
+    term_ptr_.clear();
+    csc_cluster_.clear();
+    csc_value_.clear();
+    if (n_clusters_ == 0) {
+        return;
+    }
+    // cluster_id_t (uint16) must be able to represent every cluster index in
+    // this list; beta keeps this well under 2^16 for any workable config.
+    assert(n_clusters_ <= std::numeric_limits<cluster_id_t>::max() + size_t{1});
+
+    const auto* indptr = summaries.indptr_data();
+    const auto* indices = summaries.indices_data();
+    const auto* values = summaries.values_data();  // raw bytes
+    const size_t nnz = static_cast<size_t>(indptr[n_clusters_]);
+    const size_t esz = element_size_;
+
+    // Distinct summary terms, ascending. Collected by sort+unique over the
+    // summary indices so no dimension-sized scratch is needed (the working set
+    // stays proportional to nnz).
+    term_ids_.assign(indices, indices + nnz);
+    std::ranges::sort(term_ids_);
+    term_ids_.erase(std::ranges::unique(term_ids_).begin(), term_ids_.end());
+    const size_t n_terms = term_ids_.size();
+
+    // Maps a term to its index in term_ids_ (its CSC column).
+    auto term_column = [this](term_t term) -> size_t {
+        return static_cast<size_t>(std::ranges::lower_bound(term_ids_, term) -
+                                   term_ids_.begin());
+    };
+
+    // Count entries per term, then prefix-sum into CSC offsets term_ptr_.
+    term_ptr_.assign(n_terms + 1, 0);
+    for (size_t j = 0; j < nnz; ++j) {
+        term_ptr_[term_column(indices[j]) + 1]++;
+    }
+    for (size_t t = 0; t < n_terms; ++t) term_ptr_[t + 1] += term_ptr_[t];
+
+    csc_cluster_.resize(nnz);
+    csc_value_.resize(nnz * esz);
+    std::vector<idx_t> cursor(term_ptr_.begin(), term_ptr_.end() - 1);
+    for (size_t cluster = 0; cluster < n_clusters_; ++cluster) {
+        const idx_t start = indptr[cluster];
+        const idx_t end = indptr[cluster + 1];
+        for (idx_t j = start; j < end; ++j) {
+            const size_t col = term_column(indices[j]);
+            const idx_t pos = cursor[col]++;
+            csc_cluster_[pos] = static_cast<cluster_id_t>(cluster);
+            std::copy_n(values + static_cast<size_t>(j) * esz, esz,
+                        csc_value_.data() + static_cast<size_t>(pos) * esz);
+        }
+    }
+}
+
+template <class T>
+void InvertedListClusters::score_summaries_typed(
+    const term_t* q_idx, const T* q_val, size_t q_len,
+    std::vector<float>& out) const {
+    const T* csc_values = reinterpret_cast<const T*>(csc_value_.data());
+    for (size_t i = 0; i < q_len; ++i) {
+        const term_t term = q_idx[i];
+        // Locate the query term among the summaries' distinct terms. Absent
+        // terms (including any out of the summaries' range) contribute nothing.
+        auto it = std::ranges::lower_bound(term_ids_, term);
+        if (it == term_ids_.end() || *it != term) {
+            continue;
+        }
+        const size_t col = static_cast<size_t>(it - term_ids_.begin());
+        const float qv = static_cast<float>(q_val[i]);
+        const idx_t start = term_ptr_[col];
+        const idx_t end = term_ptr_[col + 1];
+        for (idx_t j = start; j < end; ++j) {
+            out[csc_cluster_[j]] += qv * static_cast<float>(csc_values[j]);
+        }
+    }
+}
+
+void InvertedListClusters::score_summaries_transposed(
+    const term_t* q_idx, const uint8_t* q_val_bytes, size_t q_len,
+    std::vector<float>& out) const {
+    out.assign(n_clusters_, 0.0F);
+    if (n_clusters_ == 0 || term_ids_.empty()) return;
+    if (element_size_ == U32) {
+        score_summaries_typed<float>(
+            q_idx, reinterpret_cast<const float*>(q_val_bytes), q_len, out);
+    } else if (element_size_ == U16) {
+        score_summaries_typed<uint16_t>(
+            q_idx, reinterpret_cast<const uint16_t*>(q_val_bytes), q_len, out);
     } else {
-        summaries_ = std::make_unique<SparseVectors>(
-            std::move(summarize_<uint8_t>(vectors, docs_, offsets_, alpha)));
+        score_summaries_typed<uint8_t>(q_idx, q_val_bytes, q_len, out);
     }
 }
 
@@ -175,10 +254,27 @@ void InvertedListClusters::serialize(IOWriter* writer) const {
         writer->write(const_cast<idx_t*>(offsets_.data()), sizeof(idx_t),
                       n_offsets);
     }
-    if (summaries_ == nullptr) {
-        empty_sparse_vectors.serialize(writer);
-    } else {
-        summaries_->serialize(writer);
+
+    // Transposed (CSC) summary store.
+    size_t n_clusters = n_clusters_;
+    writer->write(&n_clusters, sizeof(size_t), 1);
+    size_t element_size = element_size_;
+    writer->write(&element_size, sizeof(size_t), 1);
+    size_t n_terms = term_ids_.size();
+    writer->write(&n_terms, sizeof(size_t), 1);
+    if (n_terms > 0) {
+        writer->write(const_cast<term_t*>(term_ids_.data()), sizeof(term_t),
+                      n_terms);
+        writer->write(const_cast<idx_t*>(term_ptr_.data()), sizeof(idx_t),
+                      n_terms + 1);
+    }
+    size_t nnz = csc_cluster_.size();
+    writer->write(&nnz, sizeof(size_t), 1);
+    if (nnz > 0) {
+        writer->write(const_cast<cluster_id_t*>(csc_cluster_.data()),
+                      sizeof(cluster_id_t), nnz);
+        writer->write(const_cast<uint8_t*>(csc_value_.data()), sizeof(uint8_t),
+                      nnz * element_size_);
     }
 }
 
@@ -195,11 +291,25 @@ void InvertedListClusters::deserialize(IOReader* reader) {
         offsets_.resize(n_offsets);
         reader->read(offsets_.data(), sizeof(idx_t), n_offsets);
     }
-    summaries_ = std::make_unique<SparseVectors>();
-    summaries_->deserialize(reader);
-    if (summaries_->num_vectors() == 0) {
-        summaries_.reset();
+
+    reader->read(&n_clusters_, sizeof(size_t), 1);
+    reader->read(&element_size_, sizeof(size_t), 1);
+    size_t n_terms = 0;
+    reader->read(&n_terms, sizeof(size_t), 1);
+    if (n_terms > 0) {
+        term_ids_.resize(n_terms);
+        reader->read(term_ids_.data(), sizeof(term_t), n_terms);
+        term_ptr_.resize(n_terms + 1);
+        reader->read(term_ptr_.data(), sizeof(idx_t), n_terms + 1);
+    }
+    size_t nnz = 0;
+    reader->read(&nnz, sizeof(size_t), 1);
+    if (nnz > 0) {
+        csc_cluster_.resize(nnz);
+        reader->read(csc_cluster_.data(), sizeof(cluster_id_t), nnz);
+        csc_value_.resize(nnz * element_size_);
+        reader->read(csc_value_.data(), sizeof(uint8_t), nnz * element_size_);
     }
 }
 
-}  // namespace nsparse
+}  // namespace nsparse

nsparse/cluster/inverted_list_clusters.h

@@ -16,7 +16,6 @@
 #include "nsparse/io/io.h"
 #include "nsparse/sparse_vectors.h"
 #include "nsparse/types.h"
-#include "nsparse/utils/dense_vector_matrix.h"
 
 namespace nsparse {
 
@@ -35,21 +34,54 @@ class InvertedListClusters : public Serializable {
 
     auto get_docs(idx_t idx) const -> std::span<const idx_t>;
 
-    auto summaries() const -> const SparseVectors& { return *summaries_; }
-
     void summarize(const SparseVectors* vectors, float alpha);
 
-    size_t cluster_size() const {
-        return summaries_ == nullptr ? 0 : summaries_->num_vectors();
-    }
+    size_t cluster_size() const { return n_clusters_; }
 
     void serialize(IOWriter* writer) const override;
     void deserialize(IOReader* reader) override;
 
+    // Accumulate per-cluster summary scores for a query into `out` (resized to
+    // the cluster count) using the term-major transpose. The query is given as
+    // its sparse (term, value) pairs; `q_val_bytes` points at the query values
+    // in the same element width as the stored summaries (float / uint16 /
+    // uint8), matching how the dense path reinterprets the query buffer. A
+    // query term is located by binary-searching term_ids_ (the distinct summary
+    // terms, ascending), so no dimension-sized lookup table is retained and the
+    // per-list footprint is proportional to the summaries' nnz.
+    void score_summaries_transposed(const term_t* q_idx,
+                                    const uint8_t* q_val_bytes, size_t q_len,
+                                    std::vector<float>& out) const;
+
 private:
+    // Build the term-major (CSC) transpose from a per-cluster CSR summary.
+    void build_transpose(const SparseVectors& summaries);
+    template <class T>
+    void score_summaries_typed(const term_t* q_idx, const T* q_val,
+                               size_t q_len, std::vector<float>& out) const;
+
     std::vector<idx_t> docs_;
     std::vector<idx_t> offsets_;
-    std::unique_ptr<SparseVectors> summaries_;
+
+    // Cluster ids within a posting list are bounded by beta (clusters per
+    // list), a small fraction of lambda (docs kept per list) that stays far
+    // below 2^16 for any workable configuration; the canonical Rust seismic
+    // stores this same field and likewise caps the summary count at 2^16. A
+    // 16-bit cluster id therefore halves this array vs a 32-bit one at no
+    // recall cost (build_transpose asserts the bound holds).
+    using cluster_id_t = uint16_t;
+
+    // Term-major transpose of the cluster summaries (replaces a CSR store). For
+    // each distinct summary term term_ids_[i], entries
+    // [term_ptr_[i], term_ptr_[i + 1]) in csc_cluster_/csc_value_ hold the
+    // (cluster id, summary value) pairs for that term. Scoring iterates the
+    // query's terms and scatter-adds q_val * summary_value into out[cluster].
+    size_t n_clusters_ = 0;
+    size_t element_size_ = U32;              // width of each csc_value_ entry
+    std::vector<term_t> term_ids_;           // distinct summary terms, ascending
+    std::vector<idx_t> term_ptr_;            // CSC offsets, size term_ids_+1
+    std::vector<cluster_id_t> csc_cluster_;  // cluster id per entry
+    std::vector<uint8_t> csc_value_;         // summary value per entry (bytes)
 };
 
 }  // namespace nsparse