Seismic search optimizations

nsparse/cluster/inverted_list_clusters.cpp

@@ -14,6 +14,7 @@
 #include <cstdint>
 #include <limits>
 #include <span>
+#include <stdexcept>
 #include <unordered_map>
 #include <vector>
 
@@ -118,6 +119,32 @@ InvertedListClusters::InvertedListClusters(
         docs_.insert(docs_.end(), doc_ids.begin(), doc_ids.end());
         offsets_.push_back(docs_.size());
     }
+    sort_cluster_docs();
+}
+
+// Sort the doc ids within each cluster ascending. Within-cluster iteration
+// order does not affect search results (the top-k heap and the visited dedup
+// are order-independent), but it dictates the memory access pattern of the
+// per-doc gather: indptr[doc_id] and the forward-index rows it points at are
+// spread across the multi-GB corpus. Ascending doc ids turn that per-doc random
+// gather into a monotonic sweep the hardware prefetcher and TLB can follow.
+void InvertedListClusters::sort_cluster_docs() {
+    if (offsets_.size() <= 1) return;
+    for (size_t c = 0; c + 1 < offsets_.size(); ++c) {
+        std::sort(docs_.begin() + offsets_[c], docs_.begin() + offsets_[c + 1]);
+    }
+}
+
+void InvertedListClusters::validate_doc_ids(size_t num_docs) const {
+    // docs_ is sorted per cluster but not globally, so scan the flat array.
+    // One O(nnz) pass at load; negligible against deserialize + the sort above.
+    for (const idx_t doc_id : docs_) {
+        if (static_cast<size_t>(doc_id) >= num_docs) {
+            throw std::out_of_range(
+                "InvertedListClusters: doc id out of range for corpus size; "
+                "index is corrupt or does not match its vectors");
+        }
+    }
 }
 
 InvertedListClusters::InvertedListClusters(const InvertedListClusters& other) =
@@ -291,6 +318,10 @@ void InvertedListClusters::deserialize(IOReader* reader) {
         offsets_.resize(n_offsets);
         reader->read(offsets_.data(), sizeof(idx_t), n_offsets);
     }
+    // Order within each cluster does not affect results but makes the per-doc
+    // gather monotonic (see sort_cluster_docs). Applied on load so existing
+    // serialized indexes benefit without a rebuild.
+    sort_cluster_docs();
 
     reader->read(&n_clusters_, sizeof(size_t), 1);
     reader->read(&element_size_, sizeof(size_t), 1);

nsparse/cluster/inverted_list_clusters.h

@@ -41,6 +41,13 @@ class InvertedListClusters : public Serializable {
     void serialize(IOWriter* writer) const override;
     void deserialize(IOReader* reader) override;
 
+    // Validate that every stored doc id is within [0, num_docs). The search
+    // path indexes a VisitedSet sized to num_docs by doc id with no per-
+    // candidate bounds check (that would cost a branch on the hot path), so an
+    // out-of-domain id from a corrupt or mismatched serialized index would be
+    // an out-of-bounds write. Called once at load to fail loudly instead.
+    void validate_doc_ids(size_t num_docs) const;
+
     // 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
@@ -56,6 +63,9 @@ class InvertedListClusters : public Serializable {
 private:
     // Build the term-major (CSC) transpose from a per-cluster CSR summary.
     void build_transpose(const SparseVectors& summaries);
+    // Sort doc ids ascending within each cluster (result-order invariant) so
+    // the per-doc gather over the forward index is monotonic, not random.
+    void sort_cluster_docs();
     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;

nsparse/seismic_index.cpp

@@ -16,7 +16,6 @@
 #include <numeric>
 #include <vector>
 
-#include "absl/container/flat_hash_set.h"
 #include "nsparse/cluster/inverted_list_clusters.h"
 #include "nsparse/cluster/random_kmeans.h"
 #include "nsparse/exact_matcher.h"
@@ -32,6 +31,7 @@
 #include "nsparse/utils/prefetch.h"
 #include "nsparse/utils/ranker.h"
 #include "nsparse/utils/vector_process.h"
+#include "nsparse/utils/visited_set.h"
 
 namespace nsparse {
 namespace {
@@ -41,9 +41,10 @@ constexpr int kElementSize = U32;
 void query_single_inverted_list(
     const SparseVectors* vectors, const InvertedListClusters& cluster_invlist,
     const std::vector<float>& dense, const term_t* q_idx, const float* q_val,
-    size_t q_len, std::vector<float>& score_scratch, const float heap_factor,
+    size_t q_len, std::vector<float>& score_scratch,
+    std::vector<uint32_t>& cluster_order, const float heap_factor,
     const bool first_list, const SearchParameters* search_parameters,
-    detail::TopKHolder<idx_t>& heap, absl::flat_hash_set<idx_t>& visited) {
+    detail::TopKHolder<idx_t>& heap, detail::VisitedSet& visited) {
     // Skip empty clusters
     size_t csize = cluster_invlist.cluster_size();
     if (csize == 0) {
@@ -58,39 +59,39 @@ void query_single_inverted_list(
     cluster_invlist.score_summaries_transposed(
         q_idx, reinterpret_cast<const uint8_t*>(q_val), q_len, score_scratch);
     const std::vector<float>& summary_scores = score_scratch;
-    size_t num_vectors = vectors->num_vectors();
-
-    std::vector<size_t> cluster_order =
-        detail::reorder_clusters(summary_scores, first_list);
+    const size_t n_clusters = summary_scores.size();
 
     const auto& [indptr, indices, values] = vectors->get_all_data();
 
-    for (const size_t& cluster_id : cluster_order) {
-        const auto& cluster_score = summary_scores[cluster_id];
+    // Process one cluster: prune by summary score, then score its docs.
+    // Returns false when the early-out fires on the first (sorted) list, which
+    // means no later cluster can qualify either — the caller stops iterating.
+    auto process_cluster = [&](size_t cluster_id) -> bool {
+        const float cluster_score = summary_scores[cluster_id];
         if (heap.full() && (cluster_score * heap_factor < heap.peek_score())) {
-            if (first_list) {
-                break;
-            }
-            continue;
+            // On the first list clusters are visited in descending score order,
+            // so once one falls below the threshold every later one does too.
+            return !first_list;
         }
         const auto& docs = cluster_invlist.get_docs(cluster_id);
         const size_t n_docs = docs.size();
-        static constexpr size_t kPrefetchDist1 = 2;  // vector data prefetch
-        static constexpr size_t kPrefetchDist2 = 4;  // indptr prefetch
         for (size_t i = 0; i < n_docs; ++i) {
-            const auto& doc_id = docs[i];
-            if (i + kPrefetchDist2 < n_docs) {
-                detail::prefetch_indptr(indptr, docs[i + kPrefetchDist2]);
-            }
+            const idx_t doc_id = docs[i];
+            // Prefetch one doc ahead, only the leading lines of the upcoming
+            // row; the row is contiguous so the hardware streamer pulls the
+            // tail, while bounding outstanding software prefetches keeps the
+            // line-fill buffers from saturating (measured optimum ~4 lines).
+            static constexpr size_t kPrefetchDist1 = 1;
             if (i + kPrefetchDist1 < n_docs) {
-                const idx_t next_doc = docs[i + kPrefetchDist1];
-                const idx_t next_start = indptr[next_doc];
-                const size_t next_len = indptr[next_doc + 1] - next_start;
-                detail::prefetch_vector(indices + next_start,
-                                        values + next_start, next_len);
+                const idx_t nd = docs[i + kPrefetchDist1];
+                const idx_t next_start = indptr[nd];
+                const size_t next_len = indptr[nd + 1] - next_start;
+                static constexpr size_t kPrefetchHeadLines = 4;
+                detail::prefetch_vector_head(indices + next_start,
+                                             values + next_start, next_len,
+                                             kPrefetchHeadLines);
             }
-            auto [_, inserted] = visited.insert(doc_id);
-            if (!inserted) {
+            if (!visited.insert(static_cast<size_t>(doc_id))) {
                 continue;
             }
             if (id_selector != nullptr && !id_selector->is_member(doc_id)) {
@@ -102,6 +103,29 @@ void query_single_inverted_list(
                 indices + start, values + start, len, dense.data());
             heap.add(score, doc_id);
         }
+        return true;
+    };
+
+    if (first_list) {
+        // Only the first list is score-ordered; sort a per-thread scratch of
+        // narrow (u32) cluster ids instead of allocating a size_t vector per
+        // list. Clusters per list stay well under 2^32.
+        cluster_order.resize(n_clusters);
+        std::iota(cluster_order.begin(), cluster_order.end(), 0U);
+        std::ranges::sort(cluster_order, [&](uint32_t a, uint32_t b) {
+            return summary_scores[a] > summary_scores[b];
+        });
+        for (const uint32_t cluster_id : cluster_order) {
+            if (!process_cluster(cluster_id)) {
+                break;
+            }
+        }
+    } else {
+        // Later lists are visited in natural cluster order, so iterate directly
+        // — no order array, iota, or sort needed.
+        for (size_t cluster_id = 0; cluster_id < n_clusters; ++cluster_id) {
+            process_cluster(cluster_id);
+        }
     }
 }
 }  // namespace
@@ -187,8 +211,14 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices,
 #pragma omp parallel
     {
         std::vector<float> dense(dim, 0.0F);
-        absl::flat_hash_set<idx_t> visited;
-        visited.reserve(static_cast<size_t>(std::max(k, 1)) * 4096);
+        // Generation-stamped visited set over the doc-id domain: O(1) reset per
+        // query and a single indexed load per candidate instead of a hashed
+        // random probe (the doc loop is memory-bound on random gathers).
+        detail::VisitedSet visited(vectors_->num_vectors());
+        // Per-thread scratch reused across queries: the per-cluster summary
+        // score buffer and the sorted cluster-order buffer (first list only).
+        std::vector<float> score_scratch;
+        std::vector<uint32_t> cluster_order;
 
 #pragma omp for schedule(dynamic, 64)
         for (idx_t query_idx = 0; query_idx < n; ++query_idx) {
@@ -200,7 +230,8 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices,
                 detail::top_k_tokens(q_indices, q_values, len, parameters->cut);
             auto [distances, labels] =
                 single_query(dense, visited, q_indices, q_values, len, cuts, k,
-                             parameters->heap_factor, search_parameters);
+                             parameters->heap_factor, search_parameters,
+                             score_scratch, cluster_order);
             result_distances[query_idx] = std::move(distances);
             result_labels[query_idx] = std::move(labels);
         }
@@ -219,11 +250,13 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices,
  * @return std::pair<std::vector<float>, std::vector<idx_t>>
  */
 auto SeismicIndex::single_query(std::vector<float>& dense,
-                                absl::flat_hash_set<idx_t>& visited,
+                                detail::VisitedSet& visited,
                                 const term_t* q_indices, const float* q_values,
                                 size_t q_len, const std::vector<term_t>& cuts,
                                 int k, float heap_factor,
-                                SearchParameters* search_parameters)
+                                SearchParameters* search_parameters,
+                                std::vector<float>& score_scratch,
+                                std::vector<uint32_t>& cluster_order)
     -> pair_of_score_id_vector_t {
     size_t num_docs = vectors_->num_vectors();
     if (num_docs == 0) {
@@ -234,10 +267,9 @@ auto SeismicIndex::single_query(std::vector<float>& dense,
     for (size_t i = 0; i < q_len; ++i) {
         dense[q_indices[i]] = q_values[i];
     }
-    visited.clear();
+    visited.new_query();
 
     detail::TopKHolder<idx_t> holder(k);
-    std::vector<float> score_scratch;
     bool first_list = true;
     for (const auto& term : cuts) {
         if (term >= clustered_inverted_lists.size()) [[unlikely]] {
@@ -246,8 +278,8 @@ auto SeismicIndex::single_query(std::vector<float>& dense,
         const auto& cluster_invlist = clustered_inverted_lists[term];
         query_single_inverted_list(vectors_.get(), cluster_invlist, dense,
                                    q_indices, q_values, q_len, score_scratch,
-                                   heap_factor, first_list, search_parameters,
-                                   holder, visited);
+                                   cluster_order, heap_factor, first_list,
+                                   search_parameters, holder, visited);
         first_list = false;
     }
 
@@ -287,5 +319,14 @@ void SeismicIndex::read_index(IOReader* io_reader) {
     SeismicInvertedListsWriter inv_list_writer({});
     inv_list_writer.deserialize(io_reader);
     clustered_inverted_lists = std::move(inv_list_writer.release());
+    // The search path indexes a VisitedSet by doc id without a per-candidate
+    // bounds check; validate the loaded ids fall in the corpus domain once here
+    // so a corrupt/mismatched index fails loudly instead of writing OOB.
+    if (vectors_ != nullptr) {
+        const size_t num_docs = vectors_->num_vectors();
+        for (const auto& cluster_invlist : clustered_inverted_lists) {
+            cluster_invlist.validate_doc_ids(num_docs);
+        }
+    }
 }
 }  // namespace nsparse

nsparse/seismic_index.h

@@ -12,13 +12,13 @@
 #include <array>
 #include <vector>
 
-#include "absl/container/flat_hash_set.h"
 #include "nsparse/cluster/inverted_list_clusters.h"
 #include "nsparse/index.h"
 #include "nsparse/io/io.h"
 #include "nsparse/seismic_common.h"
 #include "nsparse/sparse_vectors.h"
 #include "nsparse/types.h"
+#include "nsparse/utils/visited_set.h"
 
 namespace nsparse {
 
@@ -67,12 +67,15 @@ class SeismicIndex : public Index, public IndexIO {
     // `dense` and `visited` are per-thread scratch reused across the queries a
     // thread handles (see search()). `dense` must be all-zero on entry and is
     // restored to all-zero on exit via a sparse clear over the query's own
-    // dims (q_indices/q_len); `visited` is cleared on entry.
-    auto single_query(std::vector<float>& dense,
-                      absl::flat_hash_set<idx_t>& visited,
+    // dims (q_indices/q_len); `visited` starts a new generation on entry.
+    // `score_scratch` and `cluster_order` are per-thread scratch reused across
+    // queries (resized in place), avoiding a per-query/per-list allocation.
+    auto single_query(std::vector<float>& dense, detail::VisitedSet& visited,
                       const term_t* q_indices, const float* q_values,
                       size_t q_len, const std::vector<term_t>& cuts, int k,
-                      float heap_factor, SearchParameters* search_parameters)
+                      float heap_factor, SearchParameters* search_parameters,
+                      std::vector<float>& score_scratch,
+                      std::vector<uint32_t>& cluster_order)
         -> pair_of_score_id_vector_t;
     std::unique_ptr<SparseVectors> vectors_;
     SeismicClusterParameters cluster_parameter_;