From 78c4cec24c92aeee26e94642d6c007a942b9eac5 Mon Sep 17 00:00:00 2001 From: Zirui Song Date: Tue, 8 Sep 2026 09:19:01 +0000 Subject: [PATCH 1/2] Fix a bug: Disk-resident indexes drop selected docs when the ID selector is smaller than k Signed-off-by: Zirui Song --- nsparse/disk_seismic_index_base.cpp | 231 +++++++++++++++++- nsparse/disk_seismic_index_base.h | 33 +++ nsparse/exact_matcher.h | 6 + tests/disk_seismic_index_test.cpp | 93 +++++++ ...sk_seismic_scalar_quantized_index_test.cpp | 122 +++++++++ tests/disk_seismic_test_util.h | 39 +++ 6 files changed, 521 insertions(+), 3 deletions(-) diff --git a/nsparse/disk_seismic_index_base.cpp b/nsparse/disk_seismic_index_base.cpp index d01e24f..482fa64 100644 --- a/nsparse/disk_seismic_index_base.cpp +++ b/nsparse/disk_seismic_index_base.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -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( + 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 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(get_dimension())}); + query_vectors.add_vectors(indptr, static_cast(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 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_mapped(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 locators( + num_docs, {detail::DocLocator::kRemainder, 0, 0}); + std::vector 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 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(doc_id) >= num_docs || + covered[doc_id]) { + continue; + } + covered[doc_id] = true; + locators[doc_id] = {static_cast(pl), + static_cast(block), + static_cast(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(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(indptr[doc_id + 1] - start); + const idx_t row_indptr[2] = {0, static_cast(nnz)}; + remainder.add_vectors( + row_indptr, 2, indices + start, nnz, + values + static_cast(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(); + 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(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,102 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, index_mapping_ = std::move(mapped); } +auto DiskSeismicIndexBase::exact_match_mapped( + 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 query_scratch; + const uint8_t* query_codes = + encode_query(values, total_nnz, search_parameters, query_scratch); + const std::vector ids = selector.ordered_ids(); + + std::vector> result_distances(n); + std::vector> result_labels(n); + const size_t dense_bytes = + static_cast(get_dimension()) * element_size; + +#pragma omp parallel + { + // Per-thread dense query buffer, cleared per query below. + std::vector 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(indptr[query_idx + 1] - start); + const term_t* q_indices = indices + start; + const uint8_t* q_codes = + query_codes + static_cast(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(q_indices[i]) * element_size); + } + + detail::TopKHolder holder(k); + for (const idx_t doc_id : ids) { + if (doc_id < 0 || + static_cast(doc_id) >= num_locators_) { + continue; // out-of-range member: nothing to score + } + const detail::DocLocator loc = doc_locators_[doc_id]; + const term_t* comps = nullptr; + const uint8_t* vals = nullptr; + size_t doc_nnz = 0; + 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]; + doc_nnz = + static_cast(r_indptr[loc.block + 1] - r_start); + comps = remainder_.indices_data() + r_start; + vals = remainder_.values_data() + + static_cast(r_start) * element_size; + } else { + 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(doc_id)) { + throw std::runtime_error( + "DiskSeismic exact match: doc locator does not " + "resolve to its doc"); + } + comps = bv.doc_comps(loc.slot); + vals = bv.doc_vals(loc.slot, element_size); + doc_nnz = bv.nnz(loc.slot); + } + // Dot the doc's slice against the dense query via a 2-entry + // indptr. + const idx_t slice_indptr[2] = {0, static_cast(doc_nnz)}; + const float score = detail::compute_similarity( + 0, slice_indptr, comps, 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(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 diff --git a/nsparse/disk_seismic_index_base.h b/nsparse/disk_seismic_index_base.h index 25d1b79..aeeb6f5 100644 --- a/nsparse/disk_seismic_index_base.h +++ b/nsparse/disk_seismic_index_base.h @@ -15,16 +15,32 @@ #include #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,25 @@ 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; + // Scores every selected doc directly through the directory, for a mapped + // index. Requires doc_locators_ populated. + auto exact_match_mapped(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 diff --git a/nsparse/exact_matcher.h b/nsparse/exact_matcher.h index 05dfad1..02048d8 100644 --- a/nsparse/exact_matcher.h +++ b/nsparse/exact_matcher.h @@ -29,7 +29,13 @@ class ExactMatcher { const auto* values = vectors->values_data(); detail::TopKHolder 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(doc_id) >= num_vectors) { + continue; + } auto score = detail::compute_similarity( doc_id, indptr, indices, values, dense, element_size); holder.add(score, doc_id); diff --git a/tests/disk_seismic_index_test.cpp b/tests/disk_seismic_index_test.cpp index d2833b9..c2d10d6 100644 --- a/tests/disk_seismic_index_test.cpp +++ b/tests/disk_seismic_index_test.cpp @@ -50,6 +50,99 @@ TEST(DiskSeismicIndex, MappedReloadMatchesFreshBuild) { fresh); // fwd_ } +// An enumerable selector of size <= k must return every member, not just those +// the block budget scores. Many-block corpus + tiny k' so the budget cannot +// incidentally cover the scattered members. In-RAM path. +TEST(DiskSeismicIndex, SmallSelectorReturnsAllMembersInMemory) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + + std::vector members = {5, 100, 250, 500, 900, 1200, 1600, 1999}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + expect_all_members_returned(search_all(disk, queries, 10, ¶ms), + members); +} + +// Same contract on the mmap-loaded serialized index (fwd_ path). +TEST(DiskSeismicIndex, SmallSelectorReturnsAllMembersMmap) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + TempIndexFile file("nsparse_disk_seismic_exact_match.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + + std::vector members = {5, 100, 250, 500, 900, 1200, 1600, 1999}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + expect_all_members_returned(search_all(*mapped, queries, 10, ¶ms), + members); +} + +// The mapped path (directory + remainder) must return the same members and +// scores as the in-RAM path. +TEST(DiskSeismicIndex, ExactMatchMappedMatchesInMemory) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + std::vector members = {5, 100, 250, 500, 900, 1200, 1600, 1999}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + const ScoreIds in_memory = search_all(disk, queries, 10, ¶ms); + TempIndexFile file("nsparse_disk_seismic_exact_parity.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + expect_same_results(search_all(*mapped, queries, 10, ¶ms), in_memory); +} + +// Docs pruned from every block must still be returned and scored, through the +// remainder store. The victims are provably fully pruned (see the helper), so +// with k_prime=1 they can only surface via the mapped exact-match path. +TEST(DiskSeismicIndex, RemainderPathReturnsFullyPrunedMembers) { + const int n_fillers = 40; + const int n_victims = 3; + const CSR corpus = make_corpus_with_remainder(n_fillers, n_victims); + const CSR queries = make_corpus(5, /*seed=*/3); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + std::vector members; + for (int v = 0; v < n_victims; ++v) { + members.push_back(n_fillers + v); + } + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + const ScoreIds in_memory = search_all(disk, queries, 10, ¶ms); + TempIndexFile file("nsparse_disk_seismic_remainder.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + const ScoreIds got = search_all(*mapped, queries, 10, ¶ms); + expect_all_members_returned(got, members); + expect_same_results(got, in_memory); +} + // Building from a native CSR borrowed via mmap must match building the same // corpus fed through add(): convert -> read_csr(kMmap) -> build -> persist -> // mmap-reload -> search is bit-exact to the add()-fed build. This is also the diff --git a/tests/disk_seismic_scalar_quantized_index_test.cpp b/tests/disk_seismic_scalar_quantized_index_test.cpp index e7d94d9..b8e69e1 100644 --- a/tests/disk_seismic_scalar_quantized_index_test.cpp +++ b/tests/disk_seismic_scalar_quantized_index_test.cpp @@ -121,6 +121,128 @@ TEST(DiskSeismicSQIndex, MappedReloadMatchesFreshBuild8bit) { fresh); // fwd_ } +// An enumerable selector of size <= k must return every member, not just those +// the block budget scores. Many-block corpus + tiny k' so the budget cannot +// incidentally cover the scattered members. In-RAM path. +TEST(DiskSeismicSQIndex, SmallSelectorReturnsAllMembersInMemory) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicScalarQuantizedIndex disk(QuantizerType::QT_8bit, 0.0F, 1.0F, + cluster_params(), kDim); + add_corpus(disk, corpus); + disk.build(); + + std::vector members = {5, 100, 250, 500, 900, 1200, 1600, 1999}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + expect_all_members_returned(search_all(disk, queries, 10, ¶ms), + members); +} + +// Same contract on the mmap-loaded serialized index. +TEST(DiskSeismicSQIndex, SmallSelectorReturnsAllMembersMmap) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicScalarQuantizedIndex disk(QuantizerType::QT_8bit, 0.0F, 1.0F, + cluster_params(), kDim); + add_corpus(disk, corpus); + disk.build(); + TempIndexFile file("nsparse_dssq_exact_match.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + + std::vector members = {5, 100, 250, 500, 900, 1200, 1600, 1999}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + expect_all_members_returned(search_all(*mapped, queries, 10, ¶ms), + members); +} + +// The mapped path (directory + remainder) must return the same members and +// scores as the in-RAM path. +TEST(DiskSeismicSQIndex, ExactMatchMappedMatchesInMemory) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicScalarQuantizedIndex disk(QuantizerType::QT_8bit, 0.0F, 1.0F, + cluster_params(), kDim); + add_corpus(disk, corpus); + disk.build(); + std::vector members = {5, 100, 250, 500, 900, 1200, 1600, 1999}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + const ScoreIds in_memory = search_all(disk, queries, 10, ¶ms); + TempIndexFile file("nsparse_dssq_exact_parity.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + expect_same_results(search_all(*mapped, queries, 10, ¶ms), in_memory); +} + +// Docs pruned from every block must still be returned and scored, through the +// remainder store. The victims are provably fully pruned (see the helper), so +// with k_prime=1 they can only surface via the mapped exact-match path. +TEST(DiskSeismicSQIndex, RemainderPathReturnsFullyPrunedMembers) { + const int n_fillers = 40; + const int n_victims = 3; + const CSR corpus = make_corpus_with_remainder(n_fillers, n_victims); + const CSR queries = make_corpus(5, /*seed=*/3); + DiskSeismicScalarQuantizedIndex disk(QuantizerType::QT_8bit, 0.0F, 1.0F, + cluster_params(), kDim); + add_corpus(disk, corpus); + disk.build(); + std::vector members; + for (int v = 0; v < n_victims; ++v) { + members.push_back(n_fillers + v); + } + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + const ScoreIds in_memory = search_all(disk, queries, 10, ¶ms); + TempIndexFile file("nsparse_dssq_remainder.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + const ScoreIds got = search_all(*mapped, queries, 10, ¶ms); + expect_all_members_returned(got, members); + expect_same_results(got, in_memory); +} + +// A selector member outside the index (id >= num_vectors) is skipped, not read +// out of bounds, on both the in-RAM and the mapped exact-match paths. +TEST(DiskSeismicSQIndex, ExactMatchSkipsOutOfRangeSelectorMember) { + const CSR corpus = make_corpus(2000, /*seed=*/1); + const CSR queries = make_corpus(20, /*seed=*/2); + DiskSeismicScalarQuantizedIndex disk(QuantizerType::QT_8bit, 0.0F, 1.0F, + cluster_params(), kDim); + add_corpus(disk, corpus); + disk.build(); + std::vector members = {5, 100, static_cast(corpus.n) + 50}; + const std::vector valid = {5, 100}; + SetIDSelector selector(members.size(), members.data()); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/1); + params.set_id_selector(&selector); + + expect_all_members_returned(search_all(disk, queries, 10, ¶ms), valid); + TempIndexFile file("nsparse_dssq_oob.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + expect_all_members_returned(search_all(*mapped, queries, 10, ¶ms), + valid); +} + // Same fresh-vs-mapped parity at 16-bit (the other quantizer width). TEST(DiskSeismicSQIndex, MappedReloadMatchesFreshBuild16bit) { const CSR corpus = make_corpus(1500, /*seed=*/1); diff --git a/tests/disk_seismic_test_util.h b/tests/disk_seismic_test_util.h index 4d15fd1..f34cde6 100644 --- a/tests/disk_seismic_test_util.h +++ b/tests/disk_seismic_test_util.h @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -79,6 +80,31 @@ inline CSR make_corpus(idx_t rows, unsigned seed) { return c; } +// A corpus whose last `n_victims` docs are pruned from every posting list, so +// they end up in the remainder store rather than an inline block: each victim +// carries only term 0 at a weight below every filler, and term 0 has far more +// than kLambda higher-weight fillers (each filler also owns a unique term, so +// it always survives via its own single-doc list). Victim ids are +// [n_fillers, n_fillers + n_victims). Requires n_fillers > kLambda. +inline CSR make_corpus_with_remainder(int n_fillers, int n_victims) { + CSR c; + c.n = n_fillers + n_victims; + c.indptr.push_back(0); + for (int r = 0; r < n_fillers; ++r) { + c.indices.push_back(0); + c.values.push_back(1.0F); + c.indices.push_back(static_cast(r + 1)); + c.values.push_back(1.0F); + c.indptr.push_back(static_cast(c.indices.size())); + } + for (int v = 0; v < n_victims; ++v) { + c.indices.push_back(0); + c.values.push_back(0.01F); + c.indptr.push_back(static_cast(c.indices.size())); + } + return c; +} + inline void add_corpus(Index& index, const CSR& c) { index.add(c.n, c.indptr.data(), c.indices.data(), c.values.data()); } @@ -129,6 +155,19 @@ inline ScoreIds search_all(Index& index, const CSR& queries, int k, return out; } +// Asserts every selector member appears in each query's results. +inline void expect_all_members_returned(const ScoreIds& got, + const std::vector& members) { + for (size_t q = 0; q < got.second.size(); ++q) { + const std::unordered_set returned(got.second[q].begin(), + got.second[q].end()); + for (const idx_t member : members) { + EXPECT_GT(returned.count(member), 0U) + << "query " << q << " dropped selector member " << member; + } + } +} + inline void expect_same_results(const ScoreIds& a, const ScoreIds& b) { ASSERT_EQ(a.second.size(), b.second.size()); for (size_t q = 0; q < a.second.size(); ++q) { From 1abc63e6983555613a867eef6644b4d109f3f17b Mon Sep 17 00:00:00 2001 From: Zirui Song Date: Wed, 9 Sep 2026 03:38:22 +0000 Subject: [PATCH 2/2] Address Liyun's comments Signed-off-by: Zirui Song --- nsparse/disk_seismic_index_base.cpp | 68 ++++++++++++++--------------- nsparse/disk_seismic_index_base.h | 25 ++++++++--- 2 files changed, 52 insertions(+), 41 deletions(-) diff --git a/nsparse/disk_seismic_index_base.cpp b/nsparse/disk_seismic_index_base.cpp index 482fa64..8e70603 100644 --- a/nsparse/disk_seismic_index_base.cpp +++ b/nsparse/disk_seismic_index_base.cpp @@ -122,8 +122,8 @@ auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr, return {distances, labels}; } if (doc_locators_ != nullptr) { - return exact_match_mapped(n, indptr, indices, values, k, selector, - search_parameters); + return exact_match_directory(n, indptr, indices, values, k, + selector, search_parameters); } } @@ -326,7 +326,32 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, index_mapping_ = std::move(mapped); } -auto DiskSeismicIndexBase::exact_match_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(r_start) * element_size, + static_cast(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(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 @@ -368,41 +393,14 @@ auto DiskSeismicIndexBase::exact_match_mapped( static_cast(doc_id) >= num_locators_) { continue; // out-of-range member: nothing to score } - const detail::DocLocator loc = doc_locators_[doc_id]; - const term_t* comps = nullptr; - const uint8_t* vals = nullptr; - size_t doc_nnz = 0; - 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]; - doc_nnz = - static_cast(r_indptr[loc.block + 1] - r_start); - comps = remainder_.indices_data() + r_start; - vals = remainder_.values_data() + - static_cast(r_start) * element_size; - } else { - 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(doc_id)) { - throw std::runtime_error( - "DiskSeismic exact match: doc locator does not " - "resolve to its doc"); - } - comps = bv.doc_comps(loc.slot); - vals = bv.doc_vals(loc.slot, element_size); - doc_nnz = bv.nnz(loc.slot); - } + 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(doc_nnz)}; + const idx_t slice_indptr[2] = {0, + static_cast(doc.nnz)}; const float score = detail::compute_similarity( - 0, slice_indptr, comps, vals, dense.data(), element_size); + 0, slice_indptr, doc.comps, doc.vals, dense.data(), + element_size); holder.add(score, doc_id); } diff --git a/nsparse/disk_seismic_index_base.h b/nsparse/disk_seismic_index_base.h index aeeb6f5..74353a7 100644 --- a/nsparse/disk_seismic_index_base.h +++ b/nsparse/disk_seismic_index_base.h @@ -136,12 +136,25 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO { // same clusters and vectors as the inline forward. void write_doc_directory(IOWriter* io_writer, const SparseVectors& vectors) const; - // Scores every selected doc directly through the directory, for a mapped - // index. Requires doc_locators_ populated. - auto exact_match_mapped(idx_t n, const idx_t* indptr, const term_t* indices, - const float* values, int k, - const IDSelectorEnumerable& selector, - const SearchParameters* search_parameters) 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_;