From 7f3f2c8b826f7349bcc62638dad2492c9f93f3b8 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Mon, 13 Jul 2026 08:29:33 +0000 Subject: [PATCH 1/8] Optimize SeismicIndex search: per-thread scratch, drop query copy, dynamic scheduling Signed-off-by: Liyun Xiu --- nsparse/seismic_index.cpp | 105 ++++++++++++++++++++++++-------------- nsparse/seismic_index.h | 13 +++-- 2 files changed, 78 insertions(+), 40 deletions(-) diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index 283c899..f89dd12 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -143,14 +143,6 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, return {std::vector>(n), std::vector>(n)}; } - size_t indptr_size = n + 1; - size_t nnz = indptr[n]; // Total non-zeros - // Create query vectors from input - SparseVectors query_vectors({.element_size = kElementSize, - .dimension = static_cast(dimension_)}); - query_vectors.add_vectors(indptr, indptr_size, indices, nnz, - reinterpret_cast(values), - nnz * kElementSize); SeismicSearchParameters default_params; const auto* parameters = @@ -158,35 +150,61 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, ? dynamic_cast(search_parameters) : &default_params; - // if filter ids size is <= k, just run exact match - if (search_parameters != nullptr && - detail::should_run_exact_match(search_parameters->get_id_selector(), k, - &query_vectors)) { - return detail::ExactMatcher::search( - vectors_.get(), - dynamic_cast( - search_parameters->get_id_selector()), - &query_vectors, kElementSize, k); + // The input batch CSR (indptr/indices/values) is already absolute-indexed, + // so we index it directly in the common (no-filter) path — no need to + // re-materialize a SparseVectors or copy its arrays. Only the exact-match + // id-selector fast path still needs a SparseVectors view, so build one + // lazily just for that case. + if (search_parameters != nullptr) { + const IDSelector* sel = search_parameters->get_id_selector(); + // should_run_exact_match ignores its queries arg (only inspects the + // selector + k), so nullptr is fine here. + if (sel != nullptr && + detail::should_run_exact_match(sel, k, nullptr)) { + size_t indptr_size = n + 1; + size_t nnz = indptr[n]; + SparseVectors query_vectors( + {.element_size = kElementSize, + .dimension = static_cast(dimension_)}); + query_vectors.add_vectors( + indptr, indptr_size, indices, nnz, + reinterpret_cast(values), nnz * kElementSize); + return detail::ExactMatcher::search( + vectors_.get(), + dynamic_cast(sel), &query_vectors, + kElementSize, k); + } } std::vector> result_distances(n); std::vector> result_labels(n); - // For each query vector - const auto* query_indptr = query_vectors.indptr_data(); - const auto* query_indices = query_vectors.indices_data(); - const auto* query_values = query_vectors.values_data_float(); + const size_t dim = static_cast(dimension_); -#pragma omp parallel for - for (idx_t query_idx = 0; query_idx < n; ++query_idx) { - const auto& dense = query_vectors.get_dense_vector_float(query_idx); - const idx_t start = query_indptr[query_idx]; - const size_t len = query_indptr[query_idx + 1] - start; - const auto& cuts = detail::top_k_tokens( - query_indices + start, query_values + start, len, parameters->cut); - auto [distances, labels] = single_query( - dense, cuts, k, parameters->heap_factor, search_parameters); - result_distances[query_idx] = std::move(distances); - result_labels[query_idx] = std::move(labels); + // Per-thread scratch reused across all queries a thread handles: a + // dimension-sized dense query buffer (kept all-zero between queries via a + // sparse clear inside single_query) and the visited-doc set. This replaces + // the previous per-query allocation of both. schedule(dynamic, 64) matches + // the coarse-chunk scheduling used elsewhere in the codebase. +#pragma omp parallel + { + std::vector dense(dim, 0.0F); + absl::flat_hash_set visited; + visited.reserve(static_cast(std::max(k, 1)) * 4096); + +#pragma omp for schedule(dynamic, 64) + for (idx_t query_idx = 0; query_idx < n; ++query_idx) { + const idx_t start = indptr[query_idx]; + const size_t len = indptr[query_idx + 1] - start; + const term_t* q_indices = indices + start; + const float* q_values = values + start; + const auto& cuts = 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); + result_distances[query_idx] = std::move(distances); + result_labels[query_idx] = std::move(labels); + } } return {result_distances, result_labels}; @@ -201,17 +219,24 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, * @param heap_factor * @return std::pair, std::vector> */ -auto SeismicIndex::single_query(const std::vector& dense, - const std::vector& cuts, int k, - float heap_factor, +auto SeismicIndex::single_query(std::vector& dense, + absl::flat_hash_set& visited, + const term_t* q_indices, const float* q_values, + size_t q_len, const std::vector& cuts, + int k, float heap_factor, SearchParameters* search_parameters) -> pair_of_score_id_vector_t { size_t num_docs = vectors_->num_vectors(); if (num_docs == 0) { return {{}, {}}; } - absl::flat_hash_set visited; - visited.reserve(cuts.size() * 5000); + + // Scatter the query into the reused dense buffer (all-zero on entry). + for (size_t i = 0; i < q_len; ++i) { + dense[q_indices[i]] = q_values[i]; + } + visited.clear(); + detail::TopKHolder holder(k); bool first_list = true; for (const auto& term : cuts) { @@ -225,6 +250,12 @@ auto SeismicIndex::single_query(const std::vector& dense, first_list = false; } + // Restore the dense buffer to all-zero for the next query on this thread + // (sparse clear over only the dims this query touched). + for (size_t i = 0; i < q_len; ++i) { + dense[q_indices[i]] = 0.0F; + } + auto [scores, ids] = holder.top_k_items_descending(); scores.resize(k, -1.0F); ids.resize(k, INVALID_IDX); diff --git a/nsparse/seismic_index.h b/nsparse/seismic_index.h index 1050380..c9ffc8b 100644 --- a/nsparse/seismic_index.h +++ b/nsparse/seismic_index.h @@ -12,6 +12,7 @@ #include #include +#include "absl/container/flat_hash_set.h" #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/index.h" #include "nsparse/io/io.h" @@ -63,9 +64,15 @@ class SeismicIndex : public Index, public IndexIO { SearchParameters* search_parameters = nullptr) -> pair_of_score_id_vectors_t override; - auto single_query(const std::vector& dense, - const std::vector& cuts, int k, float heap_factor, - SearchParameters* search_parameters) + // `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& dense, + absl::flat_hash_set& visited, + const term_t* q_indices, const float* q_values, + size_t q_len, const std::vector& cuts, int k, + float heap_factor, SearchParameters* search_parameters) -> pair_of_score_id_vector_t; std::unique_ptr vectors_; SeismicClusterParameters cluster_parameter_; From 44fa2d6546a2e0538f9776c59e234529a1cc8fd5 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 15 Jul 2026 03:37:59 +0000 Subject: [PATCH 2/8] Add term-major summary transpose for query-driven cluster scoring Build a per-posting-list CSC transpose of each cluster's summary vectors at index load, then score summaries by iterating the query's non-zero terms (scatter-add into a small per-list scores array) instead of gathering a dimension-sized dense query buffer once per summary. The old path did an AVX-512 gather of dense[term] for every term of every summary in each scanned list -- a scattered read over a ~120 KB buffer that overflows L1/L2 and dominated search time. The transpose keeps the working set to the query's ~49 active terms plus a few-KB L1-resident scores array. Search throughput +72% at k=10 / +27% at k=100 (base_full, 8.84M docs, 8 threads, AVX-512), with byte-identical top-10 results vs the prior version (recall unchanged). Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 65 ++++++++++++++++++++++ nsparse/cluster/inverted_list_clusters.h | 25 +++++++++ nsparse/seismic_index.cpp | 48 ++++++++-------- 3 files changed, 114 insertions(+), 24 deletions(-) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 24ace98..ee40565 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -200,6 +200,71 @@ void InvertedListClusters::deserialize(IOReader* reader) { if (summaries_->num_vectors() == 0) { summaries_.reset(); } + build_transpose(); +} + +void InvertedListClusters::build_transpose() const { + if (transpose_built_) return; + transpose_built_ = true; + if (summaries_ == nullptr || summaries_->num_vectors() == 0) return; + const size_t n_clusters = summaries_->num_vectors(); + const auto* indptr = summaries_->indptr_data(); + const auto* indices = summaries_->indices_data(); + const auto* values = summaries_->values_data_float(); + const size_t nnz = static_cast(indptr[n_clusters]); + + // Distinct terms present across all summaries, ascending. + size_t dim = summaries_->get_dimension(); + term_lookup_.assign(dim, 0); + for (size_t j = 0; j < nnz; ++j) term_lookup_[indices[j]] = 1; + term_ids_.clear(); + for (size_t t = 0; t < dim; ++t) { + if (term_lookup_[t]) { + term_lookup_[t] = static_cast(term_ids_.size()) + 1; + term_ids_.push_back(static_cast(t)); + } + } + const size_t n_terms = term_ids_.size(); + + // Count entries per term, build term_ptr_ (CSC offsets). + term_ptr_.assign(n_terms + 1, 0); + for (size_t j = 0; j < nnz; ++j) { + term_ptr_[term_lookup_[indices[j]]]++; // note: lookup is +1 based + } + for (size_t t = 0; t < n_terms; ++t) term_ptr_[t + 1] += term_ptr_[t]; + + csc_cluster_.resize(nnz); + csc_value_.resize(nnz); + std::vector cursor(term_ptr_.begin(), term_ptr_.end() - 1); + for (size_t c = 0; c < n_clusters; ++c) { + const idx_t s = indptr[c]; + const idx_t e = indptr[c + 1]; + for (idx_t j = s; j < e; ++j) { + const int32_t ti = term_lookup_[indices[j]] - 1; + const idx_t pos = cursor[ti]++; + csc_cluster_[pos] = static_cast(c); + csc_value_[pos] = values[j]; + } + } +} + +void InvertedListClusters::score_summaries_transposed( + const term_t* q_idx, const float* q_val, size_t q_len, + std::vector& out) const { + const size_t n_clusters = cluster_size(); + out.assign(n_clusters, 0.0F); + if (!transpose_built_ || term_lookup_.empty()) return; + for (size_t i = 0; i < q_len; ++i) { + const term_t t = q_idx[i]; + const int32_t li = term_lookup_[t]; + if (li == 0) continue; // term absent from all summaries + const float qv = q_val[i]; + const idx_t s = term_ptr_[li - 1]; + const idx_t e = term_ptr_[li]; + for (idx_t j = s; j < e; ++j) { + out[csc_cluster_[j]] += qv * csc_value_[j]; + } + } } } // namespace nsparse \ No newline at end of file diff --git a/nsparse/cluster/inverted_list_clusters.h b/nsparse/cluster/inverted_list_clusters.h index f65ebd3..2c252ab 100644 --- a/nsparse/cluster/inverted_list_clusters.h +++ b/nsparse/cluster/inverted_list_clusters.h @@ -46,10 +46,35 @@ class InvertedListClusters : public Serializable { void serialize(IOWriter* writer) const override; void deserialize(IOReader* reader) override; + // Build (once) a term-major transpose of the summaries so per-query summary + // scoring can be driven by the sparse query's active terms instead of + // gathering dense[term] for every summary term. For each distinct summary + // term t, it stores the list of (cluster_id, value) pairs. Scoring: for + // each query term (t, qv), add qv*value to scores[cluster_id]. Avoids the + // dense-buffer gather entirely. + void build_transpose() const; + bool has_transpose() const { return transpose_built_; } + // term -> [start,end) into csc_cluster_/csc_value_ (size = dimension+1 via + // map) + void score_summaries_transposed(const term_t* q_idx, const float* q_val, + size_t q_len, + std::vector& out) const; + private: std::vector docs_; std::vector offsets_; std::unique_ptr summaries_; + + // Transposed (term-major) summary storage, built lazily. + mutable bool transpose_built_ = false; + mutable std::vector + term_ptr_; // per-term offset into csc arrays (sorted terms) + mutable std::vector + term_ids_; // distinct terms present (ascending) + mutable std::vector + term_lookup_; // dim -> index into term_ids_ (+1), 0 = absent + mutable std::vector csc_cluster_; // cluster id per entry + mutable std::vector csc_value_; // summary value per entry }; } // namespace nsparse diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index f89dd12..adcf8e3 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -38,13 +38,12 @@ namespace { constexpr int kElementSize = U32; -void query_single_inverted_list(const SparseVectors* vectors, - const InvertedListClusters& cluster_invlist, - const std::vector& dense, - const float heap_factor, const bool first_list, - const SearchParameters* search_parameters, - detail::TopKHolder& heap, - absl::flat_hash_set& visited) { +void query_single_inverted_list( + const SparseVectors* vectors, const InvertedListClusters& cluster_invlist, + const std::vector& dense, const term_t* q_idx, const float* q_val, + size_t q_len, std::vector& score_scratch, const float heap_factor, + const bool first_list, const SearchParameters* search_parameters, + detail::TopKHolder& heap, absl::flat_hash_set& visited) { // Skip empty clusters size_t csize = cluster_invlist.cluster_size(); if (csize == 0) { @@ -53,10 +52,11 @@ void query_single_inverted_list(const SparseVectors* vectors, const IDSelector* id_selector = search_parameters == nullptr ? nullptr : search_parameters->get_id_selector(); - const auto& summaries = cluster_invlist.summaries(); - // compute dp with all summaries - auto summary_scores = - detail::dot_product_float_vectors_dense(&summaries, dense.data()); + // Query-driven summary scoring via the term-major transpose, avoiding the + // per-summary gather into the dimension-sized dense buffer. + cluster_invlist.score_summaries_transposed(q_idx, q_val, q_len, + score_scratch); + const std::vector& summary_scores = score_scratch; size_t num_vectors = vectors->num_vectors(); std::vector cluster_order = @@ -159,20 +159,18 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, const IDSelector* sel = search_parameters->get_id_selector(); // should_run_exact_match ignores its queries arg (only inspects the // selector + k), so nullptr is fine here. - if (sel != nullptr && - detail::should_run_exact_match(sel, k, nullptr)) { + if (sel != nullptr && detail::should_run_exact_match(sel, k, nullptr)) { size_t indptr_size = n + 1; size_t nnz = indptr[n]; SparseVectors query_vectors( {.element_size = kElementSize, .dimension = static_cast(dimension_)}); - query_vectors.add_vectors( - indptr, indptr_size, indices, nnz, - reinterpret_cast(values), nnz * kElementSize); + query_vectors.add_vectors(indptr, indptr_size, indices, nnz, + reinterpret_cast(values), + nnz * kElementSize); return detail::ExactMatcher::search( - vectors_.get(), - dynamic_cast(sel), &query_vectors, - kElementSize, k); + vectors_.get(), dynamic_cast(sel), + &query_vectors, kElementSize, k); } } @@ -197,11 +195,11 @@ auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, const size_t len = indptr[query_idx + 1] - start; const term_t* q_indices = indices + start; const float* q_values = values + start; - const auto& cuts = 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); + const auto& cuts = + 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); result_distances[query_idx] = std::move(distances); result_labels[query_idx] = std::move(labels); } @@ -238,6 +236,7 @@ auto SeismicIndex::single_query(std::vector& dense, visited.clear(); detail::TopKHolder holder(k); + std::vector score_scratch; bool first_list = true; for (const auto& term : cuts) { if (term >= clustered_inverted_lists.size()) [[unlikely]] { @@ -245,6 +244,7 @@ auto SeismicIndex::single_query(std::vector& 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); first_list = false; From 014a30f98d9a4020421e4403d8f97d60018cbde2 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 15 Jul 2026 06:47:29 +0000 Subject: [PATCH 3/8] Address PR review: binary-search transpose lookup, add consistency test - Replace the retained dimension-sized term_lookup_ table with a binary search over the sorted distinct summary terms (term_ids_). This removes the per-posting-list dim-sized allocation (~120 KB each), so the transpose's footprint is now proportional to the summaries' nnz rather than the vocabulary dimension. - Query terms outside the summaries' term range are now skipped safely via the lower_bound miss (no unchecked dim-sized index). - Remove the unused has_transpose() accessor. - Rename local s/e to start/end in the scoring/build loops. - Add a unit test asserting the transposed (CSC) summary scores exactly match the reference dense-gather over the CSR summaries, across present, absent, and out-of-range query terms. Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 60 +++++++++++--------- nsparse/cluster/inverted_list_clusters.h | 21 +++---- tests/inverted_list_clusters_test.cpp | 65 ++++++++++++++++++++++ 3 files changed, 109 insertions(+), 37 deletions(-) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index ee40565..119d783 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -213,36 +213,37 @@ void InvertedListClusters::build_transpose() const { const auto* values = summaries_->values_data_float(); const size_t nnz = static_cast(indptr[n_clusters]); - // Distinct terms present across all summaries, ascending. - size_t dim = summaries_->get_dimension(); - term_lookup_.assign(dim, 0); - for (size_t j = 0; j < nnz; ++j) term_lookup_[indices[j]] = 1; - term_ids_.clear(); - for (size_t t = 0; t < dim; ++t) { - if (term_lookup_[t]) { - term_lookup_[t] = static_cast(term_ids_.size()) + 1; - term_ids_.push_back(static_cast(t)); - } - } + // Distinct terms present across all summaries, 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(); - // Count entries per term, build term_ptr_ (CSC offsets). + // Maps a term to its index in term_ids_ (its CSC column). + auto term_column = [this](term_t term) -> size_t { + return static_cast(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_lookup_[indices[j]]]++; // note: lookup is +1 based + 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); std::vector cursor(term_ptr_.begin(), term_ptr_.end() - 1); - for (size_t c = 0; c < n_clusters; ++c) { - const idx_t s = indptr[c]; - const idx_t e = indptr[c + 1]; - for (idx_t j = s; j < e; ++j) { - const int32_t ti = term_lookup_[indices[j]] - 1; - const idx_t pos = cursor[ti]++; - csc_cluster_[pos] = static_cast(c); + 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); csc_value_[pos] = values[j]; } } @@ -253,15 +254,20 @@ void InvertedListClusters::score_summaries_transposed( std::vector& out) const { const size_t n_clusters = cluster_size(); out.assign(n_clusters, 0.0F); - if (!transpose_built_ || term_lookup_.empty()) return; + if (!transpose_built_ || term_ids_.empty()) return; for (size_t i = 0; i < q_len; ++i) { - const term_t t = q_idx[i]; - const int32_t li = term_lookup_[t]; - if (li == 0) continue; // term absent from all summaries + 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(it - term_ids_.begin()); const float qv = q_val[i]; - const idx_t s = term_ptr_[li - 1]; - const idx_t e = term_ptr_[li]; - for (idx_t j = s; j < e; ++j) { + 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 * csc_value_[j]; } } diff --git a/nsparse/cluster/inverted_list_clusters.h b/nsparse/cluster/inverted_list_clusters.h index 2c252ab..ca10aee 100644 --- a/nsparse/cluster/inverted_list_clusters.h +++ b/nsparse/cluster/inverted_list_clusters.h @@ -53,9 +53,8 @@ class InvertedListClusters : public Serializable { // each query term (t, qv), add qv*value to scores[cluster_id]. Avoids the // dense-buffer gather entirely. void build_transpose() const; - bool has_transpose() const { return transpose_built_; } - // term -> [start,end) into csc_cluster_/csc_value_ (size = dimension+1 via - // map) + // Accumulate per-cluster summary scores for a query into `out` (sized to + // the cluster count) using the transpose built by build_transpose(). void score_summaries_transposed(const term_t* q_idx, const float* q_val, size_t q_len, std::vector& out) const; @@ -65,14 +64,16 @@ class InvertedListClusters : public Serializable { std::vector offsets_; std::unique_ptr summaries_; - // Transposed (term-major) summary storage, built lazily. + // Transposed (term-major) summary storage, built lazily. A query term is + // located by binary-searching term_ids_ (the distinct summary terms, in + // ascending order); no dimension-sized lookup table is retained, so the + // per-list footprint is proportional to the summaries' nnz, not to the + // vocabulary dimension. mutable bool transpose_built_ = false; - mutable std::vector - term_ptr_; // per-term offset into csc arrays (sorted terms) - mutable std::vector - term_ids_; // distinct terms present (ascending) - mutable std::vector - term_lookup_; // dim -> index into term_ids_ (+1), 0 = absent + // Distinct terms present across all summaries, ascending. term_ids_[i] owns + // csc entries [term_ptr_[i], term_ptr_[i + 1]). + mutable std::vector term_ids_; + mutable std::vector term_ptr_; mutable std::vector csc_cluster_; // cluster id per entry mutable std::vector csc_value_; // summary value per entry }; diff --git a/tests/inverted_list_clusters_test.cpp b/tests/inverted_list_clusters_test.cpp index 8c5066e..94374d5 100644 --- a/tests/inverted_list_clusters_test.cpp +++ b/tests/inverted_list_clusters_test.cpp @@ -16,6 +16,7 @@ #include "nsparse/io/buffered_io.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" +#include "nsparse/utils/distance.h" namespace { @@ -379,3 +380,67 @@ TEST(InvertedListClusters, move_assignment) { ASSERT_EQ(doc_span[0], 2); ASSERT_EQ(doc_span[1], 3); } + +// Verifies that the query-driven transposed summary scorer +// (score_summaries_transposed, CSC) produces exactly the same per-cluster +// scores as the reference dense-gather over the CSR summaries. This is the +// data-consistency check between the CSR and CSC representations. +TEST(InvertedListClusters, transposed_scoring_matches_dense_gather) { + // Two clusters with overlapping and distinct summary terms. + std::vector> docs = {{0, 1}, {2}}; + nsparse::InvertedListClusters clusters(docs); + auto vectors = create_float_vectors( + {{0, 3}, {1, 3}, {2, 5}}, {{1.0F, 2.0F}, {4.0F, 1.0F}, {3.0F, 2.0F}}, + /*dimension=*/10); + clusters.summarize(&vectors, 1.0F); + clusters.build_transpose(); + + const size_t n_clusters = clusters.cluster_size(); + ASSERT_EQ(n_clusters, 2U); + + // Reference: dense-gather each summary against a dense query vector. + const auto& summaries = clusters.summaries(); + const size_t dim = 10; + auto reference_scores = [&](const std::vector& q_idx, + const std::vector& q_val) { + std::vector dense(dim, 0.0F); + for (size_t i = 0; i < q_idx.size(); ++i) dense[q_idx[i]] = q_val[i]; + return nsparse::detail::dot_product_float_vectors_dense(&summaries, + dense.data()); + }; + + // Query terms exercising: a shared term (3), a term unique to one summary + // (0), a term absent from all summaries (7), and an out-of-range term id + // (9999 > dim) that must be safely ignored. + std::vector> queries = { + {3}, {0, 5}, {0, 3, 5}, {7}, {9999}, {3, 7, 9999}}; + std::vector> query_values = { + {2.0F}, {1.5F, 0.5F}, {1.0F, 1.0F, 1.0F}, + {3.0F}, {2.0F}, {2.0F, 1.0F, 1.0F}}; + + for (size_t q = 0; q < queries.size(); ++q) { + std::vector expected; + // Reference cannot index an out-of-range dim, so score it only over + // in-range terms; the transposed path must match that (it ignores the + // out-of-range term the same way). + std::vector in_range_idx; + std::vector in_range_val; + for (size_t i = 0; i < queries[q].size(); ++i) { + if (queries[q][i] < dim) { + in_range_idx.push_back(queries[q][i]); + in_range_val.push_back(query_values[q][i]); + } + } + expected = reference_scores(in_range_idx, in_range_val); + + std::vector got; + clusters.score_summaries_transposed( + queries[q].data(), query_values[q].data(), queries[q].size(), got); + + ASSERT_EQ(got.size(), expected.size()); + for (size_t c = 0; c < got.size(); ++c) { + ASSERT_FLOAT_EQ(got[c], expected[c]) + << "query " << q << ", cluster " << c; + } + } +} From 20b036b273424353a331a0df0e8f0883383dc99a Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 15 Jul 2026 07:03:42 +0000 Subject: [PATCH 4/8] Build summary transpose in build() and copy paths, not just deserialize() The term-major transpose was only built in deserialize(), so an index built in-process (via build()/summarize()) or copied left transpose_built_ == false. score_summaries_transposed() then returned all-zero summary scores, so cluster pruning/ordering degenerated and depended on the random k-means assignment -- observable as a flaky SeismicIndexSearch.search_respects_k_limit (doc2 sometimes returned over the higher-scoring doc1). Rebuild the transpose at the end of summarize() (resetting the built flag first) and copy the transpose members in the copy constructor / assignment operator, so every construction path has a consistent transpose. Fixes the flaky test (100/100 now; whole seismic + cluster suite green over repeated runs). Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 119d783..a8f9204 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -125,6 +125,11 @@ InvertedListClusters::InvertedListClusters(const InvertedListClusters& other) { } else { summaries_.reset(); } + transpose_built_ = other.transpose_built_; + term_ids_ = other.term_ids_; + term_ptr_ = other.term_ptr_; + csc_cluster_ = other.csc_cluster_; + csc_value_ = other.csc_value_; } InvertedListClusters& InvertedListClusters::operator=( const InvertedListClusters& other) { @@ -136,6 +141,11 @@ InvertedListClusters& InvertedListClusters::operator=( } else { summaries_.reset(); } + transpose_built_ = other.transpose_built_; + term_ids_ = other.term_ids_; + term_ptr_ = other.term_ptr_; + csc_cluster_ = other.csc_cluster_; + csc_value_ = other.csc_value_; } return *this; } @@ -161,6 +171,10 @@ void InvertedListClusters::summarize(const SparseVectors* vectors, summaries_ = std::make_unique( std::move(summarize_(vectors, docs_, offsets_, alpha))); } + // Rebuild the term-major transpose from the freshly-computed summaries so + // the build() path (not just deserialize()) has it available at query time. + transpose_built_ = false; + build_transpose(); } void InvertedListClusters::serialize(IOWriter* writer) const { From 6bc4cd15766b38477af09a7240269a424a55404d Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 15 Jul 2026 07:51:35 +0000 Subject: [PATCH 5/8] Make summary transpose the sole store; use it for SQ index; drop summaries_ Extends the term-major (CSC) summary transpose to be element-size-aware (float / uint16 / uint8 with matching accumulation) and routes the scalar-quantized index's summary scoring through it as well, so both index types share the query-driven scoring path. With the transpose now covering every query path, the CSR summaries_ member is removed entirely: InvertedListClusters stores only the transpose, summarize() builds it directly from the transient per-cluster summary, and serialization reads/writes the CSC form. This addresses the reviewer's question about what still reads the CSR after the transpose is built (nothing, now). Note: this changes the on-disk index format (CSC instead of CSR summaries); existing saved .dat files must be rebuilt. Tests that inspected summaries() are rewritten to assert via the public score_summaries_transposed() (a one-hot query recovers a summary value). Full suite green (362 tests), SQ + seismic + cluster suites stable over repeated runs. Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 238 +++++++++++---------- nsparse/cluster/inverted_list_clusters.h | 58 +++-- nsparse/seismic_index.cpp | 7 +- nsparse/seismic_scalar_quantized_index.cpp | 25 ++- nsparse/seismic_scalar_quantized_index.h | 3 +- tests/inverted_list_clusters_test.cpp | 110 ++++++---- tests/seismic_index_test.cpp | 4 +- 7 files changed, 245 insertions(+), 200 deletions(-) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index a8f9204..9013522 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -22,12 +22,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 SparseVectors summarize_(const SparseVectors* vectors, @@ -39,7 +41,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,38 +118,10 @@ InvertedListClusters::InvertedListClusters( } } -InvertedListClusters::InvertedListClusters(const InvertedListClusters& other) { - docs_ = other.docs_; - offsets_ = other.offsets_; - if (other.summaries_ != nullptr) { - summaries_ = std::make_unique(*other.summaries_); - } else { - summaries_.reset(); - } - transpose_built_ = other.transpose_built_; - term_ids_ = other.term_ids_; - term_ptr_ = other.term_ptr_; - csc_cluster_ = other.csc_cluster_; - csc_value_ = other.csc_value_; -} +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(*other.summaries_); - } else { - summaries_.reset(); - } - transpose_built_ = other.transpose_built_; - term_ids_ = other.term_ids_; - term_ptr_ = other.term_ptr_; - csc_cluster_ = other.csc_cluster_; - csc_value_ = other.csc_value_; - } - return *this; -} + const InvertedListClusters& other) = default; auto InvertedListClusters::get_docs(idx_t idx) const -> std::span { return {docs_.data() + offsets_[idx], @@ -157,79 +130,40 @@ auto InvertedListClusters::get_docs(idx_t idx) const -> std::span { 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( - std::move(summarize_(vectors, docs_, offsets_, alpha))); + summaries = summarize_(vectors, docs_, offsets_, alpha); } else if (element_size == U16) { - summaries_ = std::make_unique( - std::move(summarize_(vectors, docs_, offsets_, alpha))); + summaries = summarize_(vectors, docs_, offsets_, alpha); } else { - summaries_ = std::make_unique( - std::move(summarize_(vectors, docs_, offsets_, alpha))); + summaries = summarize_(vectors, docs_, offsets_, alpha); } - // Rebuild the term-major transpose from the freshly-computed summaries so - // the build() path (not just deserialize()) has it available at query time. - transpose_built_ = false; - build_transpose(); + build_transpose(summaries); } -void InvertedListClusters::serialize(IOWriter* writer) const { - size_t n_docs = docs_.size(); - writer->write(&n_docs, sizeof(size_t), 1); - if (n_docs > 0) { - writer->write(const_cast(docs_.data()), sizeof(idx_t), n_docs); - } - size_t n_offsets = offsets_.size(); - writer->write(&n_offsets, sizeof(size_t), 1); - if (n_offsets > 0) { - writer->write(const_cast(offsets_.data()), sizeof(idx_t), - n_offsets); - } - if (summaries_ == nullptr) { - empty_sparse_vectors.serialize(writer); - } else { - summaries_->serialize(writer); +// 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; } -} - -void InvertedListClusters::deserialize(IOReader* reader) { - size_t n_docs = 0; - reader->read(&n_docs, sizeof(size_t), 1); - if (n_docs > 0) { - docs_.resize(n_docs); - reader->read(docs_.data(), sizeof(idx_t), n_docs); - } - size_t n_offsets = 0; - reader->read(&n_offsets, sizeof(size_t), 1); - if (n_offsets > 0) { - offsets_.resize(n_offsets); - reader->read(offsets_.data(), sizeof(idx_t), n_offsets); - } - summaries_ = std::make_unique(); - summaries_->deserialize(reader); - if (summaries_->num_vectors() == 0) { - summaries_.reset(); - } - build_transpose(); -} -void InvertedListClusters::build_transpose() const { - if (transpose_built_) return; - transpose_built_ = true; - if (summaries_ == nullptr || summaries_->num_vectors() == 0) return; - const size_t n_clusters = summaries_->num_vectors(); - const auto* indptr = summaries_->indptr_data(); - const auto* indices = summaries_->indices_data(); - const auto* values = summaries_->values_data_float(); - const size_t nnz = static_cast(indptr[n_clusters]); + 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(indptr[n_clusters_]); + const size_t esz = element_size_; - // Distinct terms present across all summaries, ascending. Collected by - // sort+unique over the summary indices so no dimension-sized scratch is - // needed (the working set stays proportional to nnz). + // 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()); @@ -249,26 +183,26 @@ void InvertedListClusters::build_transpose() const { for (size_t t = 0; t < n_terms; ++t) term_ptr_[t + 1] += term_ptr_[t]; csc_cluster_.resize(nnz); - csc_value_.resize(nnz); + csc_value_.resize(nnz * esz); std::vector cursor(term_ptr_.begin(), term_ptr_.end() - 1); - for (size_t cluster = 0; cluster < n_clusters; ++cluster) { + 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); - csc_value_[pos] = values[j]; + std::copy_n(values + static_cast(j) * esz, esz, + csc_value_.data() + static_cast(pos) * esz); } } } -void InvertedListClusters::score_summaries_transposed( - const term_t* q_idx, const float* q_val, size_t q_len, +template +void InvertedListClusters::score_summaries_typed( + const term_t* q_idx, const T* q_val, size_t q_len, std::vector& out) const { - const size_t n_clusters = cluster_size(); - out.assign(n_clusters, 0.0F); - if (!transpose_built_ || term_ids_.empty()) return; + const T* csc_values = reinterpret_cast(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 @@ -278,13 +212,99 @@ void InvertedListClusters::score_summaries_transposed( continue; } const size_t col = static_cast(it - term_ids_.begin()); - const float qv = q_val[i]; + const float qv = static_cast(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 * csc_value_[j]; + out[csc_cluster_[j]] += qv * static_cast(csc_values[j]); } } } -} // namespace nsparse \ No newline at end of file +void InvertedListClusters::score_summaries_transposed( + const term_t* q_idx, const uint8_t* q_val_bytes, size_t q_len, + std::vector& out) const { + out.assign(n_clusters_, 0.0F); + if (n_clusters_ == 0 || term_ids_.empty()) return; + if (element_size_ == U32) { + score_summaries_typed( + q_idx, reinterpret_cast(q_val_bytes), q_len, out); + } else if (element_size_ == U16) { + score_summaries_typed( + q_idx, reinterpret_cast(q_val_bytes), q_len, out); + } else { + score_summaries_typed(q_idx, q_val_bytes, q_len, out); + } +} + +void InvertedListClusters::serialize(IOWriter* writer) const { + size_t n_docs = docs_.size(); + writer->write(&n_docs, sizeof(size_t), 1); + if (n_docs > 0) { + writer->write(const_cast(docs_.data()), sizeof(idx_t), n_docs); + } + size_t n_offsets = offsets_.size(); + writer->write(&n_offsets, sizeof(size_t), 1); + if (n_offsets > 0) { + writer->write(const_cast(offsets_.data()), sizeof(idx_t), + n_offsets); + } + + // 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_ids_.data()), sizeof(term_t), + n_terms); + writer->write(const_cast(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(csc_cluster_.data()), + sizeof(int32_t), nnz); + writer->write(const_cast(csc_value_.data()), sizeof(uint8_t), + nnz * element_size_); + } +} + +void InvertedListClusters::deserialize(IOReader* reader) { + size_t n_docs = 0; + reader->read(&n_docs, sizeof(size_t), 1); + if (n_docs > 0) { + docs_.resize(n_docs); + reader->read(docs_.data(), sizeof(idx_t), n_docs); + } + size_t n_offsets = 0; + reader->read(&n_offsets, sizeof(size_t), 1); + if (n_offsets > 0) { + offsets_.resize(n_offsets); + reader->read(offsets_.data(), sizeof(idx_t), n_offsets); + } + + 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(int32_t), nnz); + csc_value_.resize(nnz * element_size_); + reader->read(csc_value_.data(), sizeof(uint8_t), nnz * element_size_); + } +} + +} // namespace nsparse diff --git a/nsparse/cluster/inverted_list_clusters.h b/nsparse/cluster/inverted_list_clusters.h index ca10aee..7fa9e5c 100644 --- a/nsparse/cluster/inverted_list_clusters.h +++ b/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,47 +34,46 @@ class InvertedListClusters : public Serializable { auto get_docs(idx_t idx) const -> std::span; - 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; - // Build (once) a term-major transpose of the summaries so per-query summary - // scoring can be driven by the sparse query's active terms instead of - // gathering dense[term] for every summary term. For each distinct summary - // term t, it stores the list of (cluster_id, value) pairs. Scoring: for - // each query term (t, qv), add qv*value to scores[cluster_id]. Avoids the - // dense-buffer gather entirely. - void build_transpose() const; - // Accumulate per-cluster summary scores for a query into `out` (sized to - // the cluster count) using the transpose built by build_transpose(). - void score_summaries_transposed(const term_t* q_idx, const float* q_val, - size_t q_len, + // 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& out) const; private: + // Build the term-major (CSC) transpose from a per-cluster CSR summary. + void build_transpose(const SparseVectors& summaries); + template + void score_summaries_typed(const term_t* q_idx, const T* q_val, + size_t q_len, std::vector& out) const; + std::vector docs_; std::vector offsets_; - std::unique_ptr summaries_; - // Transposed (term-major) summary storage, built lazily. A query term is - // located by binary-searching term_ids_ (the distinct summary terms, in - // ascending order); no dimension-sized lookup table is retained, so the - // per-list footprint is proportional to the summaries' nnz, not to the - // vocabulary dimension. - mutable bool transpose_built_ = false; - // Distinct terms present across all summaries, ascending. term_ids_[i] owns - // csc entries [term_ptr_[i], term_ptr_[i + 1]). - mutable std::vector term_ids_; - mutable std::vector term_ptr_; - mutable std::vector csc_cluster_; // cluster id per entry - mutable std::vector csc_value_; // summary value per entry + // 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_ids_; // distinct summary terms, ascending + std::vector term_ptr_; // CSC offsets, size term_ids_+1 + std::vector csc_cluster_; // cluster id per entry + std::vector csc_value_; // summary value per entry (bytes) }; } // namespace nsparse diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index adcf8e3..332cb8d 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -53,9 +53,10 @@ void query_single_inverted_list( ? nullptr : search_parameters->get_id_selector(); // Query-driven summary scoring via the term-major transpose, avoiding the - // per-summary gather into the dimension-sized dense buffer. - cluster_invlist.score_summaries_transposed(q_idx, q_val, q_len, - score_scratch); + // per-summary gather into the dimension-sized dense buffer. The summaries + // hold float values, so the query values are passed as their raw bytes. + cluster_invlist.score_summaries_transposed( + q_idx, reinterpret_cast(q_val), q_len, score_scratch); const std::vector& summary_scores = score_scratch; size_t num_vectors = vectors->num_vectors(); diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index 41b1111..d9e8c2e 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -40,6 +40,8 @@ namespace { void query_single_inverted_list(const SparseVectors* vectors, const InvertedListClusters& cluster_invlist, const std::vector& dense, + const term_t* q_idx, const uint8_t* q_val_bytes, + size_t q_len, std::vector& score_scratch, float heap_factor, bool first_list, const SearchParameters* search_parameters, detail::TopKHolder& heap, @@ -53,10 +55,11 @@ void query_single_inverted_list(const SparseVectors* vectors, ? nullptr : search_parameters->get_id_selector(); const auto element_size = vectors->get_element_size(); - const auto& summaries = cluster_invlist.summaries(); - // compute dp with all summaries - std::vector summary_scores = - detail::calculate_summary_scores(element_size, &summaries, dense); + // Query-driven summary scoring via the term-major transpose. The query + // values are the quantized codes in the summaries' element width. + cluster_invlist.score_summaries_transposed(q_idx, q_val_bytes, q_len, + score_scratch); + const std::vector& summary_scores = score_scratch; size_t num_vectors = vectors->num_vectors(); std::vector cluster_order = @@ -251,9 +254,10 @@ auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, parameters->cut); } - auto [distances, labels] = - single_query(dense, cuts, k, parameters->heap_factor, query_sq, - search_parameters); + const uint8_t* q_val_bytes = query_values + start * element_size; + auto [distances, labels] = single_query( + dense, query_indices + start, q_val_bytes, len, cuts, k, + parameters->heap_factor, query_sq, search_parameters); result_distances[query_idx] = std::move(distances); result_labels[query_idx] = std::move(labels); } @@ -262,8 +266,9 @@ auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, } auto SeismicScalarQuantizedIndex::single_query( - const std::vector& dense, const std::vector& cuts, int k, - float heap_factor, const ScalarQuantizer& query_sq, + const std::vector& dense, const term_t* q_idx, + const uint8_t* q_val_bytes, size_t q_len, const std::vector& cuts, + int k, float heap_factor, const ScalarQuantizer& query_sq, SearchParameters* search_parameters) -> pair_of_score_id_vector_t { size_t num_docs = vectors_->num_vectors(); if (num_docs == 0) { @@ -272,6 +277,7 @@ auto SeismicScalarQuantizedIndex::single_query( absl::flat_hash_set visited; visited.reserve(cuts.size() * 5000); detail::TopKHolder holder(k); + std::vector score_scratch; bool first_list = true; for (const auto& term : cuts) { if (term >= clustered_inverted_lists.size()) [[unlikely]] { @@ -279,6 +285,7 @@ auto SeismicScalarQuantizedIndex::single_query( } const auto& cluster_invlist = clustered_inverted_lists[term]; query_single_inverted_list(vectors_.get(), cluster_invlist, dense, + q_idx, q_val_bytes, q_len, score_scratch, heap_factor, first_list, search_parameters, holder, visited); first_list = false; diff --git a/nsparse/seismic_scalar_quantized_index.h b/nsparse/seismic_scalar_quantized_index.h index 023f99c..96cc44b 100644 --- a/nsparse/seismic_scalar_quantized_index.h +++ b/nsparse/seismic_scalar_quantized_index.h @@ -65,7 +65,8 @@ class SeismicScalarQuantizedIndex : public Index, public IndexIO { -> pair_of_score_id_vectors_t override; auto encode(const float* values, size_t nnz, SearchParameters* search_parameters) -> std::vector; - auto single_query(const std::vector& dense, + auto single_query(const std::vector& dense, const term_t* q_idx, + const uint8_t* q_val_bytes, size_t q_len, const std::vector& cuts, int k, float heap_factor, const ScalarQuantizer& query_sq, SearchParameters* search_parameters) diff --git a/tests/inverted_list_clusters_test.cpp b/tests/inverted_list_clusters_test.cpp index 94374d5..85ac184 100644 --- a/tests/inverted_list_clusters_test.cpp +++ b/tests/inverted_list_clusters_test.cpp @@ -66,6 +66,24 @@ nsparse::SparseVectors create_uint16_vectors( return vectors; } +// Recovers summary value[cluster][term] via the public scoring API: a one-hot +// query {term: 1} makes score_summaries_transposed return, per cluster, exactly +// that cluster's summary value at `term` (0 if absent). Replaces the old +// summaries().get_dense_vector*() inspection now that only the transpose is +// stored. `q_val_t` is the query value type matching the summaries' element +// width (float / uint16_t / uint8_t). +template +std::vector summary_values_at_term( + const nsparse::InvertedListClusters& clusters, nsparse::term_t term) { + std::vector q_idx = {term}; + std::vector q_val = {static_cast(1)}; + std::vector out; + clusters.score_summaries_transposed( + q_idx.data(), reinterpret_cast(q_val.data()), + q_idx.size(), out); + return out; +} + } // namespace // Constructor tests @@ -149,7 +167,11 @@ TEST(InvertedListClusters, copy_constructor_with_summaries) { nsparse::InvertedListClusters copy(original); ASSERT_EQ(copy.cluster_size(), 2); - ASSERT_EQ(copy.summaries().num_vectors(), 2); + // The copy must carry a working transpose: cluster 0 has term 0, cluster 1 + // has term 2 (both with max value 1.5 across their docs' overlapping term). + auto term0 = summary_values_at_term(copy, 0); + ASSERT_EQ(term0.size(), 2U); + ASSERT_GT(term0[0], 0.0F); } // Copy assignment tests @@ -190,12 +212,10 @@ TEST(InvertedListClusters, summarize_float_single_cluster) { clusters.summarize(&vectors, 1.0F); ASSERT_EQ(clusters.cluster_size(), 1); - const auto& summaries = clusters.summaries(); - ASSERT_EQ(summaries.num_vectors(), 1); - - auto dense = summaries.get_dense_vector_float(0); - ASSERT_FLOAT_EQ(dense[0], 3.0F); // max(1.0, 3.0) - ASSERT_FLOAT_EQ(dense[1], 2.0F); // max(2.0, 1.0) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 0)[0], + 3.0F); // max(1.0, 3.0) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 1)[0], + 2.0F); // max(2.0, 1.0) } TEST(InvertedListClusters, summarize_float_multiple_clusters) { @@ -211,11 +231,10 @@ TEST(InvertedListClusters, summarize_float_multiple_clusters) { ASSERT_EQ(clusters.cluster_size(), 2); - auto dense0 = clusters.summaries().get_dense_vector_float(0); - ASSERT_FLOAT_EQ(dense0[0], 2.0F); // max(1.0, 2.0) - - auto dense1 = clusters.summaries().get_dense_vector_float(1); - ASSERT_FLOAT_EQ(dense1[1], 3.0F); + // Cluster 0 has term 0 -> 2.0; cluster 1 has term 1 -> 3.0. + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 0)[0], + 2.0F); // max(1.0, 2.0) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 1)[1], 3.0F); } TEST(InvertedListClusters, summarize_uint8) { @@ -231,9 +250,10 @@ TEST(InvertedListClusters, summarize_uint8) { ASSERT_EQ(clusters.cluster_size(), 1); - auto dense = clusters.summaries().get_dense_vector(0); - ASSERT_EQ(dense[0], 30); // max(10, 30) - ASSERT_EQ(dense[1], 20); // max(20, 10) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 0)[0], + 30.0F); // max(10, 30) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 1)[0], + 20.0F); // max(20, 10) } TEST(InvertedListClusters, summarize_uint16) { @@ -250,10 +270,10 @@ TEST(InvertedListClusters, summarize_uint16) { ASSERT_EQ(clusters.cluster_size(), 1); - auto dense = clusters.summaries().get_dense_vector(0); - auto* values = reinterpret_cast(dense.data()); - ASSERT_EQ(values[0], 300); // max(100, 300) - ASSERT_EQ(values[1], 200); // max(200, 100) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 0)[0], + 300.0F); // max(100, 300) + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 1)[0], + 200.0F); // max(200, 100) } TEST(InvertedListClusters, summarize_with_alpha_pruning) { @@ -271,10 +291,10 @@ TEST(InvertedListClusters, summarize_with_alpha_pruning) { clusters.summarize(&vectors, 0.5F); ASSERT_EQ(clusters.cluster_size(), 1); - const auto& summaries = clusters.summaries(); - auto dense = summaries.get_dense_vector_float(0); - // Only term 0 should be non-zero - ASSERT_GT(dense[0], 0.0F); + // Only term 0 should survive alpha pruning (non-zero); the pruned terms + // contribute nothing. + ASSERT_GT(summary_values_at_term(clusters, 0)[0], 0.0F); + ASSERT_FLOAT_EQ(summary_values_at_term(clusters, 1)[0], 0.0F); } TEST(InvertedListClusters, summarize_replaces_existing) { @@ -385,7 +405,7 @@ TEST(InvertedListClusters, move_assignment) { // (score_summaries_transposed, CSC) produces exactly the same per-cluster // scores as the reference dense-gather over the CSR summaries. This is the // data-consistency check between the CSR and CSC representations. -TEST(InvertedListClusters, transposed_scoring_matches_dense_gather) { +TEST(InvertedListClusters, transposed_scoring_matches_reference) { // Two clusters with overlapping and distinct summary terms. std::vector> docs = {{0, 1}, {2}}; nsparse::InvertedListClusters clusters(docs); @@ -393,20 +413,29 @@ TEST(InvertedListClusters, transposed_scoring_matches_dense_gather) { {{0, 3}, {1, 3}, {2, 5}}, {{1.0F, 2.0F}, {4.0F, 1.0F}, {3.0F, 2.0F}}, /*dimension=*/10); clusters.summarize(&vectors, 1.0F); - clusters.build_transpose(); const size_t n_clusters = clusters.cluster_size(); ASSERT_EQ(n_clusters, 2U); - // Reference: dense-gather each summary against a dense query vector. - const auto& summaries = clusters.summaries(); + // Reference summary matrix: recover value[cluster][term] via the one-hot + // helper, so we can independently compute the expected query dot product + // per cluster without inspecting the internal storage. const size_t dim = 10; + std::vector> summary(n_clusters, + std::vector(dim, 0.0F)); + for (nsparse::term_t t = 0; t < dim; ++t) { + auto vals = summary_values_at_term(clusters, t); + for (size_t c = 0; c < n_clusters; ++c) summary[c][t] = vals[c]; + } auto reference_scores = [&](const std::vector& q_idx, const std::vector& q_val) { - std::vector dense(dim, 0.0F); - for (size_t i = 0; i < q_idx.size(); ++i) dense[q_idx[i]] = q_val[i]; - return nsparse::detail::dot_product_float_vectors_dense(&summaries, - dense.data()); + std::vector out(n_clusters, 0.0F); + for (size_t c = 0; c < n_clusters; ++c) { + for (size_t i = 0; i < q_idx.size(); ++i) { + if (q_idx[i] < dim) out[c] += q_val[i] * summary[c][q_idx[i]]; + } + } + return out; }; // Query terms exercising: a shared term (3), a term unique to one summary @@ -419,23 +448,14 @@ TEST(InvertedListClusters, transposed_scoring_matches_dense_gather) { {3.0F}, {2.0F}, {2.0F, 1.0F, 1.0F}}; for (size_t q = 0; q < queries.size(); ++q) { - std::vector expected; - // Reference cannot index an out-of-range dim, so score it only over - // in-range terms; the transposed path must match that (it ignores the - // out-of-range term the same way). - std::vector in_range_idx; - std::vector in_range_val; - for (size_t i = 0; i < queries[q].size(); ++i) { - if (queries[q][i] < dim) { - in_range_idx.push_back(queries[q][i]); - in_range_val.push_back(query_values[q][i]); - } - } - expected = reference_scores(in_range_idx, in_range_val); + std::vector expected = + reference_scores(queries[q], query_values[q]); std::vector got; clusters.score_summaries_transposed( - queries[q].data(), query_values[q].data(), queries[q].size(), got); + queries[q].data(), + reinterpret_cast(query_values[q].data()), + queries[q].size(), got); ASSERT_EQ(got.size(), expected.size()); for (size_t c = 0; c < got.size(); ++c) { diff --git a/tests/seismic_index_test.cpp b/tests/seismic_index_test.cpp index d95ea72..4b39671 100644 --- a/tests/seismic_index_test.cpp +++ b/tests/seismic_index_test.cpp @@ -141,10 +141,8 @@ TEST(SeismicIndexBuild, build_creates_summaries_for_clusters) { index.build(); auto& inv_lists = index.get_clustered_inverted_lists(); - // Term 0 has 2 docs -> clusters with summaries + // Term 0 has 2 docs -> 2 clusters, each with a summary. EXPECT_EQ(inv_lists[0].cluster_size(), 2); - const auto& summaries = inv_lists[0].summaries(); - EXPECT_EQ(summaries.num_vectors(), 2); } // ============== Constructor tests ============== From a2567cb2a5c3a420bce1f07885f7bf062e4d0c8d Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 15 Jul 2026 08:41:39 +0000 Subject: [PATCH 6/8] Make SQ search benchmark quantizer width configurable Read the quantizer width from NSPARSE_SQ_BITS ("8bit" or "16bit", default "8bit") in BM_SeismicSQ_Search and cache each width to its own .dat suffix, so float / int16 / int8 search throughput can be compared without editing the benchmark. Default behavior is unchanged when the variable is unset. Signed-off-by: Liyun Xiu --- benchmarks/index_search_benchmark.cpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/benchmarks/index_search_benchmark.cpp b/benchmarks/index_search_benchmark.cpp index c276d93..5f6bda6 100644 --- a/benchmarks/index_search_benchmark.cpp +++ b/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(state.range(0)); const int n_queries = static_cast(fix.query.nrow); From c4cc9a8567e6366053fd259c216f05831bbd156b Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 16 Jul 2026 02:28:06 +0000 Subject: [PATCH 7/8] Narrow summary-transpose cluster id from int32 to uint16 The term-major (CSC) summary transpose stored cluster ids as int32, but clusters per posting list are bounded by beta (a small fraction of lambda) and stay far below 2^16 for any workable config. Narrowing csc_cluster_ to uint16 halves that array at no recall cost. On base_full (8.84M SPLADE docs) this shrinks the summary store ~20.8% (float) to ~30.3% (SQ-8bit) and the on-disk index ~12-19%, with search performance retained (SQ-8bit slightly faster from reduced memory traffic) and recall unchanged. This changes the serialized index format, so cached indexes must be rebuilt. The canonical Rust seismic stores this same field and likewise caps the summary count at 2^16; build_transpose asserts the bound holds. Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 13 +++++++++---- nsparse/cluster/inverted_list_clusters.h | 18 +++++++++++++----- 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 9013522..73ff26b 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -10,7 +10,9 @@ #include "nsparse/cluster/inverted_list_clusters.h" #include +#include #include +#include #include #include #include @@ -154,6 +156,9 @@ void InvertedListClusters::build_transpose(const SparseVectors& summaries) { 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::max() + size_t{1}); const auto* indptr = summaries.indptr_data(); const auto* indices = summaries.indices_data(); @@ -191,7 +196,7 @@ void InvertedListClusters::build_transpose(const SparseVectors& summaries) { 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); + csc_cluster_[pos] = static_cast(cluster); std::copy_n(values + static_cast(j) * esz, esz, csc_value_.data() + static_cast(pos) * esz); } @@ -266,8 +271,8 @@ void InvertedListClusters::serialize(IOWriter* writer) const { size_t nnz = csc_cluster_.size(); writer->write(&nnz, sizeof(size_t), 1); if (nnz > 0) { - writer->write(const_cast(csc_cluster_.data()), - sizeof(int32_t), nnz); + writer->write(const_cast(csc_cluster_.data()), + sizeof(cluster_id_t), nnz); writer->write(const_cast(csc_value_.data()), sizeof(uint8_t), nnz * element_size_); } @@ -301,7 +306,7 @@ void InvertedListClusters::deserialize(IOReader* reader) { reader->read(&nnz, sizeof(size_t), 1); if (nnz > 0) { csc_cluster_.resize(nnz); - reader->read(csc_cluster_.data(), sizeof(int32_t), 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_); } diff --git a/nsparse/cluster/inverted_list_clusters.h b/nsparse/cluster/inverted_list_clusters.h index 7fa9e5c..9b6beb9 100644 --- a/nsparse/cluster/inverted_list_clusters.h +++ b/nsparse/cluster/inverted_list_clusters.h @@ -63,17 +63,25 @@ class InvertedListClusters : public Serializable { std::vector docs_; std::vector offsets_; + // 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_ids_; // distinct summary terms, ascending - std::vector term_ptr_; // CSC offsets, size term_ids_+1 - std::vector csc_cluster_; // cluster id per entry - std::vector csc_value_; // summary value per entry (bytes) + size_t element_size_ = U32; // width of each csc_value_ entry + std::vector term_ids_; // distinct summary terms, ascending + std::vector term_ptr_; // CSC offsets, size term_ids_+1 + std::vector csc_cluster_; // cluster id per entry + std::vector csc_value_; // summary value per entry (bytes) }; } // namespace nsparse From ec8b600897b8848c2ea2a9b8c2233fd775172ee2 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 16 Jul 2026 02:54:53 +0000 Subject: [PATCH 8/8] Apply per-thread scratch to SQ search, matching SeismicIndex The scalar-quantized search path never received the allocation optimizations that SeismicIndex::search got. It allocated a fresh dimension-sized dense buffer and a visited-doc set per query, and re-materialized the query batch into a SparseVectors on every call. Mirror the float index: reuse a per-thread dense buffer (restored to all-zero via a sparse clear over each query's own dims) and visited set across the queries a thread handles, index the pre-quantized batch codes directly instead of building a SparseVectors (kept only for the id-selector exact-match path), and use schedule(dynamic, 64). Behavior-preserving: all 362 tests pass and recall@10 on base_full SQ-8bit is unchanged (0.8327). Search throughput improves ~6% on base_small (k=10 160.0k -> 169.8k QPS, k=100 66.9k -> 70.7k at 16 threads), where per-query allocation is a meaningful fraction of the work; the gain washes out on the full 8.84M corpus where per-doc scoring dominates. Signed-off-by: Liyun Xiu --- nsparse/seismic_scalar_quantized_index.cpp | 109 +++++++++++++-------- nsparse/seismic_scalar_quantized_index.h | 13 ++- 2 files changed, 79 insertions(+), 43 deletions(-) diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index d9e8c2e..63599fe 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -11,6 +11,7 @@ #include +#include #include #include #include @@ -200,17 +201,22 @@ auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, sq_params->vmax); } - // construct query vector + // Quantize the whole query batch once. `codes` holds the quantized values + // in the same CSR order as `indices`, so the batch (indptr/indices/codes) + // can be indexed directly below — no per-query SparseVectors needed. const size_t element_size = sq_.bytes_per_value(); - SparseVectors query_vectors({.element_size = element_size, - .dimension = static_cast(dimension_)}); std::vector codes = encode(values, nnz, search_parameters); - query_vectors.add_vectors(indptr, indptr_size, indices, nnz, codes.data(), - nnz * element_size); + const uint8_t* query_values = codes.data(); - // if filter ids size is <= k, just run exact match + // if filter ids size is <= k, just run exact match. Only this path needs a + // SparseVectors view of the query, so build one lazily just for it. if (detail::should_run_exact_match(search_parameters->get_id_selector(), k, - &query_vectors)) { + nullptr)) { + SparseVectors query_vectors( + {.element_size = element_size, + .dimension = static_cast(dimension_)}); + query_vectors.add_vectors(indptr, indptr_size, indices, nnz, + codes.data(), nnz * element_size); auto [distances, labels] = detail::ExactMatcher::search( vectors_.get(), dynamic_cast( @@ -231,51 +237,67 @@ auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, // query const auto* parameters = dynamic_cast(search_parameters); - const auto* query_indptr = query_vectors.indptr_data(); - const auto* query_indices = query_vectors.indices_data(); - const auto* query_values = query_vectors.values_data(); - -#pragma omp parallel for - for (idx_t query_idx = 0; query_idx < n; ++query_idx) { - const auto& dense = query_vectors.get_dense_vector(query_idx); - const idx_t start = query_indptr[query_idx]; - const size_t len = query_indptr[query_idx + 1] - start; - std::vector cuts; - if (element_size == U16) { - // start is element index, need byte offset for uint16_t access - cuts = detail::top_k_tokens( - query_indices + start, - reinterpret_cast(query_values + - start * sizeof(uint16_t)), - len, parameters->cut); - } else { - cuts = detail::top_k_tokens(query_indices + start, - query_values + start, len, - parameters->cut); - } + const size_t dense_bytes = + static_cast(dimension_) * element_size; + + // Per-thread scratch reused across all queries a thread handles: a + // dimension-sized quantized-code dense buffer (kept all-zero between queries + // via a sparse clear inside single_query) and the visited-doc set. This + // replaces the previous per-query allocation of both. schedule(dynamic, 64) + // matches the coarse-chunk scheduling used by SeismicIndex::search. +#pragma omp parallel + { + std::vector dense(dense_bytes, 0); + absl::flat_hash_set visited; + visited.reserve(static_cast(std::max(k, 1)) * 4096); - const uint8_t* q_val_bytes = query_values + start * element_size; - auto [distances, labels] = single_query( - dense, query_indices + start, q_val_bytes, len, cuts, k, - parameters->heap_factor, query_sq, search_parameters); - result_distances[query_idx] = std::move(distances); - result_labels[query_idx] = std::move(labels); +#pragma omp for schedule(dynamic, 64) + for (idx_t query_idx = 0; query_idx < n; ++query_idx) { + const idx_t start = indptr[query_idx]; + const size_t len = indptr[query_idx + 1] - start; + const term_t* q_indices = indices + start; + const uint8_t* q_val_bytes = query_values + start * element_size; + std::vector cuts; + if (element_size == U16) { + cuts = detail::top_k_tokens( + q_indices, + reinterpret_cast(q_val_bytes), len, + parameters->cut); + } else { + cuts = detail::top_k_tokens(q_indices, q_val_bytes, + len, parameters->cut); + } + + auto [distances, labels] = single_query( + dense, visited, q_indices, q_val_bytes, len, element_size, cuts, + k, parameters->heap_factor, query_sq, search_parameters); + result_distances[query_idx] = std::move(distances); + result_labels[query_idx] = std::move(labels); + } } return {result_distances, result_labels}; } auto SeismicScalarQuantizedIndex::single_query( - const std::vector& dense, const term_t* q_idx, - const uint8_t* q_val_bytes, size_t q_len, const std::vector& cuts, - int k, float heap_factor, const ScalarQuantizer& query_sq, + std::vector& dense, absl::flat_hash_set& visited, + const term_t* q_idx, const uint8_t* q_val_bytes, size_t q_len, + size_t element_size, const std::vector& cuts, int k, + float heap_factor, const ScalarQuantizer& query_sq, SearchParameters* search_parameters) -> pair_of_score_id_vector_t { size_t num_docs = vectors_->num_vectors(); if (num_docs == 0) { return {{}, {}}; } - absl::flat_hash_set visited; - visited.reserve(cuts.size() * 5000); + + // Scatter the query's quantized codes into the reused dense buffer (all-zero + // on entry): element_size contiguous bytes per non-zero dim. + for (size_t i = 0; i < q_len; ++i) { + std::copy_n(q_val_bytes + i * element_size, element_size, + dense.data() + static_cast(q_idx[i]) * element_size); + } + visited.clear(); + detail::TopKHolder holder(k); std::vector score_scratch; bool first_list = true; @@ -291,6 +313,13 @@ auto SeismicScalarQuantizedIndex::single_query( first_list = false; } + // Restore the dense buffer to all-zero for the next query on this thread + // (sparse clear over only the dims this query touched). + for (size_t i = 0; i < q_len; ++i) { + std::fill_n(dense.data() + static_cast(q_idx[i]) * element_size, + element_size, uint8_t{0}); + } + auto [distances, labels] = holder.top_k_items_descending(); // Decode quantized dot product scores diff --git a/nsparse/seismic_scalar_quantized_index.h b/nsparse/seismic_scalar_quantized_index.h index 96cc44b..179b226 100644 --- a/nsparse/seismic_scalar_quantized_index.h +++ b/nsparse/seismic_scalar_quantized_index.h @@ -14,6 +14,7 @@ #include #include +#include "absl/container/flat_hash_set.h" #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/index.h" #include "nsparse/seismic_index.h" @@ -65,10 +66,16 @@ class SeismicScalarQuantizedIndex : public Index, public IndexIO { -> pair_of_score_id_vectors_t override; auto encode(const float* values, size_t nnz, SearchParameters* search_parameters) -> std::vector; - auto single_query(const std::vector& dense, const term_t* q_idx, + // `dense` and `visited` are per-thread scratch reused across the queries a + // thread handles (see search()). `dense` (a dimension-sized quantized-code + // buffer, element_size bytes per dim) 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_idx/q_len); `visited` is cleared on entry. + auto single_query(std::vector& dense, + absl::flat_hash_set& visited, const term_t* q_idx, const uint8_t* q_val_bytes, size_t q_len, - const std::vector& cuts, int k, float heap_factor, - const ScalarQuantizer& query_sq, + size_t element_size, const std::vector& cuts, + int k, float heap_factor, const ScalarQuantizer& query_sq, SearchParameters* search_parameters) -> pair_of_score_id_vector_t; ScalarQuantizer sq_;