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); diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 24ace98..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 @@ -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 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(*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(*other.summaries_); - } else { - summaries_.reset(); - } - } - return *this; -} + const InvertedListClusters& other) = default; auto InvertedListClusters::get_docs(idx_t idx) const -> std::span { return {docs_.data() + offsets_[idx], @@ -147,19 +132,113 @@ 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 = summarize_(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::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(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(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 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); + std::copy_n(values + static_cast(j) * esz, esz, + csc_value_.data() + static_cast(pos) * esz); + } + } +} + +template +void InvertedListClusters::score_summaries_typed( + const term_t* q_idx, const T* q_val, size_t q_len, + std::vector& out) const { + 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 + // 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 = 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 * static_cast(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& 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 { - summaries_ = std::make_unique( - std::move(summarize_(vectors, docs_, offsets_, alpha))); + score_summaries_typed(q_idx, q_val_bytes, q_len, out); } } @@ -175,10 +254,27 @@ void InvertedListClusters::serialize(IOWriter* writer) const { writer->write(const_cast(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_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(cluster_id_t), nnz); + writer->write(const_cast(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(); - 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 \ No newline at end of file +} // namespace nsparse diff --git a/nsparse/cluster/inverted_list_clusters.h b/nsparse/cluster/inverted_list_clusters.h index f65ebd3..9b6beb9 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,21 +34,54 @@ 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; + // 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_; + + // 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) }; } // namespace nsparse diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index 283c899..332cb8d 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,12 @@ 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. 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(); std::vector cluster_order = @@ -143,14 +144,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 +151,59 @@ 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,18 +218,26 @@ 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); + std::vector score_scratch; bool first_list = true; for (const auto& term : cuts) { if (term >= clustered_inverted_lists.size()) [[unlikely]] { @@ -220,11 +245,18 @@ auto SeismicIndex::single_query(const 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; } + // 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_; diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index 41b1111..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 @@ -40,6 +41,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 +56,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 = @@ -197,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( @@ -228,50 +237,69 @@ 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); - auto [distances, labels] = - single_query(dense, 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 std::vector& cuts, int k, + 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; for (const auto& term : cuts) { if (term >= clustered_inverted_lists.size()) [[unlikely]] { @@ -279,11 +307,19 @@ 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; } + // 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 023f99c..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,9 +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 std::vector& cuts, int k, float heap_factor, - const ScalarQuantizer& query_sq, + // `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, + 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_; diff --git a/tests/inverted_list_clusters_test.cpp b/tests/inverted_list_clusters_test.cpp index 8c5066e..85ac184 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 { @@ -65,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 @@ -148,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 @@ -189,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) { @@ -210,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) { @@ -230,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) { @@ -249,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) { @@ -270,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) { @@ -379,3 +400,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_reference) { + // 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); + + const size_t n_clusters = clusters.cluster_size(); + ASSERT_EQ(n_clusters, 2U); + + // 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 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 + // (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_scores(queries[q], query_values[q]); + + std::vector got; + clusters.score_summaries_transposed( + 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) { + ASSERT_FLOAT_EQ(got[c], expected[c]) + << "query " << q << ", cluster " << 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 ==============