[Fix] [DiskSeismic]: take the exact-match path when the ID selector is smaller than k

nsparse/disk_seismic_index_base.cpp

@@ -13,6 +13,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <span>
 #include <stdexcept>
 #include <string>
 #include <utility>
@@ -21,13 +22,17 @@
 #include "absl/container/flat_hash_set.h"
 #include "nsparse/cluster/inverted_list_clusters.h"
 #include "nsparse/disk_seismic_search.h"
+#include "nsparse/exact_matcher.h"
 #include "nsparse/id_selector.h"
 #include "nsparse/index.h"
+#include "nsparse/io/align.h"
 #include "nsparse/io/inline_forward_index_io.h"
 #include "nsparse/io/seismic_invlists_writer.h"
+#include "nsparse/seismic_common.h"
 #include "nsparse/sparse_vectors.h"
 #include "nsparse/types.h"
 #include "nsparse/utils/checks.h"
+#include "nsparse/utils/ranker.h"
 #include "nsparse/utils/mmap_cursor.h"
 #include "nsparse/utils/mmap_file.h"
 
@@ -77,9 +82,51 @@ auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr,
         return detail::initialize_padded_results(n, k);
     }
 
-    // The id-selector exact-match fast path is omitted (it needs an in-RAM
-    // SparseVectors a mapped index lacks); the selector is still honored per
-    // candidate doc inside the shared search core.
+    // An enumerable selector of size <= k must return every member, so score
+    // exactly the selected docs rather than only those in the top-k' blocks.
+    // get_vectors() is doc-id-addressable for a fresh or mmap-CSR build; a
+    // mapped serialized index resolves docs through doc_locators_ instead.
+    if (search_parameters != nullptr &&
+        detail::should_run_exact_match(search_parameters->get_id_selector(), k,
+                                       nullptr)) {
+        const auto& selector = *dynamic_cast<const IDSelectorEnumerable*>(
+            search_parameters->get_id_selector());
+        const SparseVectors* vectors = get_vectors();
+        if (vectors != nullptr) {
+            const size_t element_size = code_element_size();
+            const size_t nnz = indptr[n];
+            std::vector<uint8_t> query_scratch;
+            const uint8_t* query_codes =
+                encode_query(values, nnz, search_parameters, query_scratch);
+            SparseVectors query_vectors(
+                {.element_size = element_size,
+                 .dimension = static_cast<size_t>(get_dimension())});
+            query_vectors.add_vectors(indptr, static_cast<size_t>(n) + 1,
+                                      indices, nnz, query_codes,
+                                      nnz * element_size);
+            auto [distances, labels] = detail::ExactMatcher::search(
+                vectors, &selector, &query_vectors, element_size, k);
+            // Decode only the filled prefix; ExactMatcher pads the tail with
+            // (-1.0F, INVALID_IDX), which must not be decoded.
+            for (size_t q = 0; q < distances.size(); ++q) {
+                size_t filled = 0;
+                while (filled < labels[q].size() &&
+                       labels[q][filled] != detail::INVALID_IDX) {
+                    ++filled;
+                }
+                std::vector<float> head(distances[q].begin(),
+                                        distances[q].begin() + filled);
+                decode_scores(head, search_parameters);
+                std::copy(head.begin(), head.end(), distances[q].begin());
+            }
+            return {distances, labels};
+        }
+        if (doc_locators_ != nullptr) {
+            return exact_match_directory(n, indptr, indices, values, k,
+                                         selector, search_parameters);
+        }
+    }
+
     const detail::DiskSeismicCutBudget budget =
         detail::resolve_cut_and_budget(search_parameters);
     const int cut = budget.cut;
@@ -165,6 +212,72 @@ void DiskSeismicIndexBase::write_index(IOWriter* io_writer) {
     const SparseVectors& v = vectors_ != nullptr ? *vectors_ : empty_vectors;
     detail::InlineForwardIndex forward(clustered_inverted_lists, v);
     forward.serialize(io_writer);
+    write_doc_directory(io_writer, v);
+}
+
+void DiskSeismicIndexBase::write_doc_directory(
+    IOWriter* io_writer, const SparseVectors& vectors) const {
+    const size_t num_docs = vectors.num_vectors();
+    const size_t element_size = vectors.get_element_size();
+
+    // Default to remainder; the loop below overrides docs it finds in a block.
+    std::vector<detail::DocLocator> locators(
+        num_docs, {detail::DocLocator::kRemainder, 0, 0});
+    std::vector<bool> covered(num_docs, false);
+    // Must mirror InlineForwardIndex::write_body's block/slot iteration so the
+    // recorded (posting_list, block, slot) match the blocks it writes.
+    // First occurrence wins; every copy of a doc's vector is identical.
+    for (size_t pl = 0; pl < clustered_inverted_lists.size(); ++pl) {
+        const InvertedListClusters& list = clustered_inverted_lists[pl];
+        const size_t n_clusters = list.cluster_size();
+        for (size_t block = 0; block < n_clusters; ++block) {
+            const std::span<const idx_t> docs = list.get_docs(block);
+            for (size_t slot = 0; slot < docs.size(); ++slot) {
+                const idx_t doc_id = docs[slot];
+                if (doc_id < 0 || static_cast<size_t>(doc_id) >= num_docs ||
+                    covered[doc_id]) {
+                    continue;
+                }
+                covered[doc_id] = true;
+                locators[doc_id] = {static_cast<uint32_t>(pl),
+                                    static_cast<uint32_t>(block),
+                                    static_cast<uint32_t>(slot)};
+            }
+        }
+    }
+
+    // Docs in no block, in doc-id order: keep their full vectors here and point
+    // the locator at the row.
+    SparseVectors remainder(
+        {.element_size = element_size,
+         .dimension = static_cast<size_t>(get_dimension())});
+    const idx_t* indptr = vectors.indptr_data();
+    const term_t* indices = vectors.indices_data();
+    const uint8_t* values = vectors.values_data();
+    uint32_t remainder_row = 0;
+    for (size_t doc_id = 0; doc_id < num_docs; ++doc_id) {
+        if (covered[doc_id]) {
+            continue;
+        }
+        const idx_t start = indptr[doc_id];
+        const size_t nnz = static_cast<size_t>(indptr[doc_id + 1] - start);
+        const idx_t row_indptr[2] = {0, static_cast<idx_t>(nnz)};
+        remainder.add_vectors(
+            row_indptr, 2, indices + start, nnz,
+            values + static_cast<size_t>(start) * element_size,
+            nnz * element_size);
+        locators[doc_id] = {detail::DocLocator::kRemainder, remainder_row, 0};
+        ++remainder_row;
+    }
+
+    // Aligned u64 doc count, then the locator array, then the remainder
+    // vectors; load_mapped_payload reads them back in this order.
+    io_align::pad_to(io_writer, detail::kMinBlockAlign);
+    uint64_t count = num_docs;
+    io_writer->write(&count, sizeof(uint64_t), 1);
+    io_align::write_padded(io_writer, locators.data(), locators.size(),
+                           alignof(detail::DocLocator));
+    remainder.serialize(io_writer);
 }
 
 void DiskSeismicIndexBase::read_index(IOReader* /*io_reader*/,
@@ -189,6 +302,20 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor,
 
     clustered_inverted_lists = std::move(inv_list_writer.release());
     fwd_ = std::move(forward);
+
+    // Doc-locator directory (borrowed in place) then remainder vectors, in the
+    // order write_doc_directory wrote them.
+    io_align::skip_padding(cursor, detail::kMinBlockAlign);
+    const uint64_t num_locators = cursor->read_scalar<uint64_t>();
+    if (num_locators != num_vectors_) {
+        throw std::runtime_error(
+            "DiskSeismic index: doc-locator count disagrees with vector count");
+    }
+    io_align::skip_padding(cursor, alignof(detail::DocLocator));
+    doc_locators_ = cursor->read_array<detail::DocLocator>(num_locators);
+    num_locators_ = num_locators;
+    remainder_.mmap_deserialize(cursor);
+
     // Now that the summaries and forward index are populated (still borrowing
     // from `mapped`, which is alive here), let the concrete index reject a
     // width mismatch before we commit.
@@ -199,4 +326,100 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor,
     index_mapping_ = std::move(mapped);
 }
 
+auto DiskSeismicIndexBase::get_doc(idx_t doc_id, size_t element_size) const
+    -> DocSlice {
+    const detail::DocLocator loc = doc_locators_[doc_id];
+    if (loc.posting_list == detail::DocLocator::kRemainder) {
+        if (loc.block >= remainder_.num_vectors()) {
+            throw std::runtime_error(
+                "DiskSeismic exact match: remainder locator out of range");
+        }
+        const idx_t* r_indptr = remainder_.indptr_data();
+        const idx_t r_start = r_indptr[loc.block];
+        return {remainder_.indices_data() + r_start,
+                remainder_.values_data() +
+                    static_cast<size_t>(r_start) * element_size,
+                static_cast<size_t>(r_indptr[loc.block + 1] - r_start)};
+    }
+    const detail::BlockView bv = fwd_.block(loc.posting_list, loc.block);
+    if (bv.absent() || loc.slot >= bv.n_docs ||
+        bv.doc_ids[loc.slot] != static_cast<uint32_t>(doc_id)) {
+        throw std::runtime_error(
+            "DiskSeismic exact match: doc locator does not resolve to its doc");
+    }
+    return {bv.doc_comps(loc.slot), bv.doc_vals(loc.slot, element_size),
+            bv.nnz(loc.slot)};
+}
+
+auto DiskSeismicIndexBase::exact_match_directory(
+    idx_t n, const idx_t* indptr, const term_t* indices, const float* values,
+    int k, const IDSelectorEnumerable& selector,
+    const SearchParameters* search_parameters) const
+    -> pair_of_score_id_vectors_t {
+    const size_t element_size = code_element_size();
+    const size_t total_nnz = indptr[n];
+    std::vector<uint8_t> query_scratch;
+    const uint8_t* query_codes =
+        encode_query(values, total_nnz, search_parameters, query_scratch);
+    const std::vector<idx_t> ids = selector.ordered_ids();
+
+    std::vector<std::vector<float>> result_distances(n);
+    std::vector<std::vector<idx_t>> result_labels(n);
+    const size_t dense_bytes =
+        static_cast<size_t>(get_dimension()) * element_size;
+
+#pragma omp parallel
+    {
+        // Per-thread dense query buffer, cleared per query below.
+        std::vector<uint8_t> dense(dense_bytes, 0);
+#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 =
+                static_cast<size_t>(indptr[query_idx + 1] - start);
+            const term_t* q_indices = indices + start;
+            const uint8_t* q_codes =
+                query_codes + static_cast<size_t>(start) * element_size;
+            for (size_t i = 0; i < len; ++i) {
+                std::copy_n(
+                    q_codes + i * element_size, element_size,
+                    dense.data() +
+                        static_cast<size_t>(q_indices[i]) * element_size);
+            }
+
+            detail::TopKHolder<idx_t> holder(k);
+            for (const idx_t doc_id : ids) {
+                if (doc_id < 0 ||
+                    static_cast<uint64_t>(doc_id) >= num_locators_) {
+                    continue;  // out-of-range member: nothing to score
+                }
+                const DocSlice doc = get_doc(doc_id, element_size);
+                // Dot the doc's slice against the dense query via a 2-entry
+                // indptr.
+                const idx_t slice_indptr[2] = {0,
+                                               static_cast<idx_t>(doc.nnz)};
+                const float score = detail::compute_similarity(
+                    0, slice_indptr, doc.comps, doc.vals, dense.data(),
+                    element_size);
+                holder.add(score, doc_id);
+            }
+
+            // Decode before padding so the -1 pad is not decoded.
+            auto [scores, labels] = holder.top_k_items_descending();
+            decode_scores(scores, search_parameters);
+            scores.resize(k, -1.0F);
+            labels.resize(k, detail::INVALID_IDX);
+
+            for (size_t i = 0; i < len; ++i) {
+                std::fill_n(dense.data() + static_cast<size_t>(q_indices[i]) *
+                                               element_size,
+                            element_size, uint8_t{0});
+            }
+            result_distances[query_idx] = std::move(scores);
+            result_labels[query_idx] = std::move(labels);
+        }
+    }
+    return {result_distances, result_labels};
+}
+
 }  // namespace nsparse

nsparse/disk_seismic_index_base.h

@@ -15,16 +15,32 @@
 #include <vector>
 
 #include "nsparse/cluster/inverted_list_clusters.h"
+#include "nsparse/id_selector.h"
 #include "nsparse/index.h"
 #include "nsparse/io/inline_forward_index_io.h"
 #include "nsparse/io/io.h"
 #include "nsparse/mmap_index.h"
+#include "nsparse/sparse_vectors.h"
 #include "nsparse/types.h"
 #include "nsparse/utils/mmap_cursor.h"
 #include "nsparse/utils/mmap_file.h"
 
 namespace nsparse {
 
+namespace detail {
+// Locates one full copy of a doc's vector. When posting_list == kRemainder the
+// vector is row `block` of the remainder store; otherwise it is slot `slot` of
+// inline-forward block (posting_list, block).
+struct DocLocator {
+    uint32_t posting_list;
+    uint32_t block;
+    uint32_t slot;
+    static constexpr uint32_t kRemainder = UINT32_MAX;
+};
+static_assert(sizeof(DocLocator) == 12,
+              "DocLocator is borrowed from the mapping");
+}  // namespace detail
+
 // Shared implementation of the two disk-resident SEISMIC indexes: the cluster
 // summaries live in RAM, the per-document forward vectors live on disk in the
 // block-contiguous (inline) layout and are borrowed via mmap at search time,
@@ -116,8 +132,38 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO {
     void read_index(IOReader* io_reader, const IndexHeader& header,
                     int io_flags = 0) override;
 
+    // Appends the doc-locator directory and remainder vectors, built from the
+    // same clusters and vectors as the inline forward.
+    void write_doc_directory(IOWriter* io_writer,
+                             const SparseVectors& vectors) const;
+
+    // One doc's within-doc slice: component ids, element_size-wide codes, and
+    // the count. Borrowed from the live mapping or remainder_ (both outlive the
+    // call), so valid only within one search().
+    struct DocSlice {
+        const term_t* comps = nullptr;
+        const uint8_t* vals = nullptr;
+        size_t nnz = 0;
+    };
+    // Resolves one selected doc's full vector through the doc-locator
+    // directory: an inline-forward block slot, or a row of remainder_ when the
+    // doc was pruned from every block. Throws on a corrupt locator.
+    [[nodiscard]] DocSlice get_doc(idx_t doc_id, size_t element_size) const;
+    // Scores every selected doc directly through the doc-locator directory, for
+    // a mapped index. Requires doc_locators_ populated.
+    [[nodiscard]] auto exact_match_directory(
+        idx_t n, const idx_t* indptr, const term_t* indices,
+        const float* values, int k, const IDSelectorEnumerable& selector,
+        const SearchParameters* search_parameters) const
+        -> pair_of_score_id_vectors_t;
+
     SeismicClusterParameters cluster_parameter_;
     size_t num_vectors_ = 0;
+    // A borrowed per-doc locator table plus the full vectors of the docs that
+    // are missing from every block, used by the mapped-index exact match.
+    const detail::DocLocator* doc_locators_ = nullptr;
+    uint64_t num_locators_ = 0;
+    SparseVectors remainder_;
 };
 
 }  // namespace nsparse

nsparse/exact_matcher.h

@@ -29,7 +29,13 @@ class ExactMatcher {
         const auto* values = vectors->values_data();
         detail::TopKHolder<idx_t> holder(k);
         auto ids = id_selector->ordered_ids();
+        const auto num_vectors = vectors->num_vectors();
         for (auto doc_id : ids) {
+            // A selector may name ids outside this index; skip them rather than
+            // index indptr out of bounds.
+            if (doc_id < 0 || static_cast<size_t>(doc_id) >= num_vectors) {
+                continue;
+            }
             auto score = detail::compute_similarity(
                 doc_id, indptr, indices, values, dense, element_size);
             holder.add(score, doc_id);