diff --git a/nsparse/CMakeLists.txt b/nsparse/CMakeLists.txt index 6e31237..3aa6bd1 100644 --- a/nsparse/CMakeLists.txt +++ b/nsparse/CMakeLists.txt @@ -10,6 +10,7 @@ set(NSPARSE_SRC brutal_index.cpp seismic_index.cpp seismic_scalar_quantized_index.cpp + disk_seismic_index.cpp sparse_vectors.cpp index_factory.cpp id_map_index.cpp diff --git a/nsparse/disk_seismic_index.cpp b/nsparse/disk_seismic_index.cpp new file mode 100644 index 0000000..369ef8a --- /dev/null +++ b/nsparse/disk_seismic_index.cpp @@ -0,0 +1,346 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#include "nsparse/disk_seismic_index.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#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/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/distance_simd.h" +#include "nsparse/utils/mmap_cursor.h" +#include "nsparse/utils/mmap_file.h" +#include "nsparse/utils/ranker.h" +#include "nsparse/utils/vector_process.h" + +namespace nsparse { +namespace { + +constexpr int kElementSize = U32; + +// Score one doc's vector against the dense query and push it, honoring +// visited-dedup and the id selector. Shared by both vector sources. +inline void score_doc(idx_t doc_id, const term_t* comps, const float* vals, + size_t nnz, const std::vector& dense, + const IDSelector* id_selector, + detail::TopKHolder& heap, + absl::flat_hash_set& visited) { + if (!visited.insert(doc_id).second) { + return; + } + if (id_selector != nullptr && !id_selector->is_member(doc_id)) { + return; + } + const float score = + detail::dot_product_float_dense(comps, vals, nnz, dense.data()); + heap.add(score, doc_id); +} + +// A candidate block: its summary score and its (posting list, cluster) address. +struct BlockCandidate { + float score; + uint32_t pl; + uint32_t cid; +}; + +// Score every document of one block against the dense query, dedup via +// `visited` and honor the id selector. Vectors come from the inline forward +// index (fwd) when loaded, else from the in-RAM CSR of a fresh build. +void score_block(const detail::InlineForwardIndex* fwd, + const SparseVectors* vectors, + const std::vector& clusters, uint32_t pl, + uint32_t cid, const std::vector& dense, + const IDSelector* id_selector, size_t element_size, + detail::TopKHolder& heap, + absl::flat_hash_set& visited) { + if (fwd != nullptr) { + const detail::BlockView bv = fwd->block(pl, cid); + if (bv.absent()) { + return; + } + for (uint32_t i = 0; i < bv.n_docs; ++i) { + score_doc( + static_cast(bv.doc_ids[i]), bv.doc_comps(i), + reinterpret_cast(bv.doc_vals(i, element_size)), + bv.nnz(i), dense, id_selector, heap, visited); + } + } else if (vectors != nullptr) { + const auto data = vectors->get_all_data(); + const idx_t* const indptr = data.indptr_data; + const term_t* const indices = data.indices_data; + const float* const values = data.values_data; + for (const idx_t doc_id : clusters[pl].get_docs(cid)) { + const idx_t start = indptr[doc_id]; + const size_t len = indptr[doc_id + 1] - start; + score_doc(doc_id, indices + start, values + start, len, dense, + id_selector, heap, visited); + } + } +} +} // namespace + +DiskSeismicIndex::DiskSeismicIndex(int dim) + : MmapIndex(dim), + cluster_parameter_(detail::kDefaultSeismicClusterParams) {} + +DiskSeismicIndex::DiskSeismicIndex(int dim, SeismicClusterParameters parameter) + : MmapIndex(dim), cluster_parameter_(parameter) {} + +void DiskSeismicIndex::add(idx_t n, const idx_t* indptr, const term_t* indices, + const float* values) { + throw_if_not_positive(n); + throw_if_any_null(indptr, indices, values); + const size_t indptr_size = n + 1; + const size_t nnz = indptr[n]; + if (vectors_ == nullptr) { + // Fresh container: start the count at 0 so a stale num_vectors_ (e.g. + // left by a prior mmap load, which has no vectors_) cannot accumulate. + num_vectors_ = 0; + vectors_ = std::make_unique( + SparseVectorsConfig{.element_size = kElementSize, + .dimension = static_cast(dimension_)}); + } + vectors_->add_vectors(indptr, indptr_size, indices, nnz, + reinterpret_cast(values), + nnz * kElementSize); + num_vectors_ += n; +} + +void DiskSeismicIndex::build() { + clustered_inverted_lists = detail::build_inverted_lists_clusters( + get_vectors(), + {.element_size = kElementSize, + .dimension = static_cast(get_dimension())}, + cluster_parameter_); +} + +auto DiskSeismicIndex::search(idx_t n, const idx_t* indptr, + const term_t* indices, const float* values, int k, + SearchParameters* search_parameters) + -> pair_of_score_id_vectors_t { + // Quit early when there is nothing to score: no vectors, no queries, or no + // forward-vector source (fwd_ empty and vectors_ null — a corrupt or + // uninitialized index). + if (num_vectors_ == 0 || n == 0 || + (fwd_.num_blocks() == 0 && vectors_ == nullptr)) { + return { + std::vector>(n, std::vector(k, -1.0F)), + std::vector>( + n, std::vector(k, detail::INVALID_IDX))}; + } + + // Resolve `cut` and `k_prime`. A DiskSeismicSearchParameters carries + // k_prime; a plain SeismicSearchParameters (or null) uses the default + // budget. The id-selector exact-match fast path is omitted (it needs an + // in-RAM SparseVectors a mapped index lacks); the selector is still honored + // in score_block. + const DiskSeismicSearchParameters defaults; + int cut = defaults.cut; + int k_prime = defaults.k_prime; + if (const auto* disk_parameters = + dynamic_cast( + search_parameters)) { + cut = disk_parameters->cut; + k_prime = disk_parameters->k_prime; + } else if (const auto* seismic_parameters = + dynamic_cast( + search_parameters)) { + cut = seismic_parameters->cut; + } + // else: search_parameters is null or an unrelated SearchParameters subtype, + // so cut and k_prime keep the defaults initialized above. + + // k_prime is a block budget, not a document count: a block holds many docs, + // so k_prime < k is valid (a few blocks can still fill k, and a short-fall + // is padded like any under-budget search). Only a non-positive budget is + // rejected. + if (k_prime <= 0) { + throw std::invalid_argument( + "DiskSeismicIndex: k_prime (block budget) must be positive"); + } + + std::vector> result_distances(n, std::vector(k, -1.0F)); + std::vector> result_labels( + n, std::vector(k, detail::INVALID_IDX)); + const size_t dim = static_cast(dimension_); + +#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, cut); + auto [distances, labels] = + single_query(dense, visited, q_indices, q_values, len, cuts, k, + k_prime, search_parameters); + result_distances[query_idx] = std::move(distances); + result_labels[query_idx] = std::move(labels); + } + } + return {result_distances, result_labels}; +} + +auto DiskSeismicIndex::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, int k_prime, + SearchParameters* search_parameters) -> pair_of_score_id_vector_t { + if (num_vectors_ == 0) { + return {{}, {}}; + } + // Scatter the query into the dense lookup table. + for (size_t i = 0; i < q_len; ++i) { + dense[q_indices[i]] = q_values[i]; + } + visited.clear(); + + // Prefer the mapped inline forward index; fall back to the in-RAM vectors + // of a fresh build. + const detail::InlineForwardIndex* fwd = + fwd_.num_blocks() > 0 ? &fwd_ : nullptr; + const SparseVectors* vectors = fwd == nullptr ? vectors_.get() : nullptr; + const size_t element_size = fwd != nullptr + ? fwd->element_size() + : static_cast(kElementSize); + const IDSelector* id_selector = search_parameters == nullptr + ? nullptr + : search_parameters->get_id_selector(); + + // Collect every candidate block across the cut lists with its summary + // score. score_summaries_transposed dots the full query with each cluster + // summary, so scores are comparable across posting lists for the global + // ranking. + std::vector candidates; + std::vector score_scratch; + for (const term_t term : cuts) { + if (term >= clustered_inverted_lists.size()) [[unlikely]] { + continue; + } + const InvertedListClusters& cluster_invlist = + clustered_inverted_lists[term]; + const size_t n_clusters = cluster_invlist.cluster_size(); + if (n_clusters == 0) { + continue; + } + cluster_invlist.score_summaries_transposed( + q_indices, reinterpret_cast(q_values), q_len, + score_scratch); + for (size_t cid = 0; cid < n_clusters; ++cid) { + candidates.push_back( + {score_scratch[cid], term, static_cast(cid)}); + } + } + + // Select the global top-k_prime blocks by summary score. The order among + // them does not affect the result (dedup + exact scoring are + // order-independent). + const size_t budget = + std::min(static_cast(k_prime), candidates.size()); + if (budget < candidates.size()) { + std::nth_element(candidates.begin(), candidates.begin() + budget, + candidates.end(), + [](const BlockCandidate& a, const BlockCandidate& b) { + return a.score > b.score; + }); + candidates.resize(budget); + } + + // Score the selected blocks; visited dedups docs shared across them. + detail::TopKHolder holder(k); + for (const BlockCandidate& candidate : candidates) { + score_block(fwd, vectors, clustered_inverted_lists, candidate.pl, + candidate.cid, dense, id_selector, element_size, holder, + visited); + } + + // Reset only the query's own positions (q_len is small) instead of + // memset-ing the whole dim-sized table; mirrors the scatter at entry. The + // scattered writes are bounded by q_len and negligible next to the block + // reads/scoring that dominate the query. + 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, detail::INVALID_IDX); + return {scores, ids}; +} + +void DiskSeismicIndex::write_index(IOWriter* io_writer) { + const uint64_t nv = num_vectors_; + io_writer->write(const_cast(&nv), sizeof(uint64_t), 1); + SeismicInvertedListsWriter inv_list_writer(clustered_inverted_lists); + inv_list_writer.serialize(io_writer); + // Inline forward index, built from the same clusters + vectors. An empty + // corpus uses a correctly-typed empty SparseVectors (element_size must be a + // valid width even with zero vectors) so the section still round-trips. + SparseVectors empty_vectors({.element_size = kElementSize, + .dimension = static_cast(dimension_)}); + const SparseVectors& v = vectors_ != nullptr ? *vectors_ : empty_vectors; + detail::InlineForwardIndex forward(clustered_inverted_lists, v); + forward.serialize(io_writer); +} + +void DiskSeismicIndex::read_index(IOReader* /*io_reader*/, int /*io_flags*/) { + // The inline forward index is borrowed from a mapping, never copied onto + // the heap, so this index has no copying read path. + throw std::runtime_error( + "DiskSeismicIndex is mmap-only; load with read_index(file, " + "IndexIoFlag::kUseMmap)"); +} + +DiskSeismicIndex* DiskSeismicIndex::mmap_index(int dimension, + const char* index_file, + size_t pos) { + throw_if_null(index_file, "index_file must not be null"); + auto index = std::make_unique(dimension); + + MmapFile mmap_file(std::string{index_file}); + MmapCursor cursor(mmap_file.data(), mmap_file.size()); + cursor.skip(pos); + + // Same order write_index wrote them: doc count, summaries, inline forward. + const uint64_t nv = cursor.read_scalar(); + SeismicInvertedListsWriter inv_list_writer; + inv_list_writer.mmap_deserialize(&cursor); + detail::InlineForwardIndex forward; + forward.mmap_deserialize(&cursor); + + // Committed only once everything parsed. mapped_file_ last: the summaries + // and the forward index borrow from it, and moving it does not move the + // mapping. + index->clustered_inverted_lists = std::move(inv_list_writer.release()); + index->fwd_ = std::move(forward); + index->num_vectors_ = nv; + index->mapped_file_ = std::move(mmap_file); + return index.release(); +} +} // namespace nsparse diff --git a/nsparse/disk_seismic_index.h b/nsparse/disk_seismic_index.h new file mode 100644 index 0000000..60b18b4 --- /dev/null +++ b/nsparse/disk_seismic_index.h @@ -0,0 +1,108 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#ifndef DISK_SEISMIC_INDEX_H +#define DISK_SEISMIC_INDEX_H +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/io/inline_forward_index_io.h" +#include "nsparse/io/io.h" +#include "nsparse/mmap_index.h" +#include "nsparse/seismic_common.h" +#include "nsparse/seismic_index.h" // SeismicSearchParameters +#include "nsparse/types.h" + +namespace nsparse { + +// Default block budget K': the primary recall/latency knob. +inline constexpr int kDefaultBlockBudget = 50; + +// `cut` (inherited) bounds which posting lists' summaries are scored; `k_prime` +// caps how many blocks are read (the k_prime highest-scoring ones). Replaces +// the inherited heap_factor, which DiskSeismicIndex ignores. +struct DiskSeismicSearchParameters : public SeismicSearchParameters { + int k_prime = kDefaultBlockBudget; + DiskSeismicSearchParameters() = default; + DiskSeismicSearchParameters(int cut, int k_prime) + : SeismicSearchParameters(cut, /*heap_factor=*/1.0F), + k_prime(k_prime) {} +}; + +// A seismic index whose per-document forward vectors live on disk in the +// block-contiguous (inline) layout of io/inline_forward_index_io.h, borrowed +// via mmap at search time; the cluster summaries stay in RAM. Search scores +// every candidate block's summary across the `cut` posting lists, selects the +// global top-k_prime blocks, and reads and scores only those. Pass a +// DiskSeismicSearchParameters to set k_prime, else kDefaultBlockBudget is used. +// +// A block's vectors come from fwd_.block() (the mapping) once loaded, else from +// the in-RAM SparseVectors of a fresh build. +// +// mmap-only: load with read_index(file, kUseMmap); the copying read throws. +class DiskSeismicIndex : public MmapIndex, public IndexIO { +public: + static constexpr std::array name = {'D', 'S', 'E', 'I'}; + + explicit DiskSeismicIndex(int dim); + DiskSeismicIndex(int dim, SeismicClusterParameters parameter); + ~DiskSeismicIndex() override = default; + std::array id() const override { return name; } + + DiskSeismicIndex(const DiskSeismicIndex&) = delete; + DiskSeismicIndex& operator=(const DiskSeismicIndex&) = delete; + + void add(idx_t n, const idx_t* indptr, const term_t* indices, + const float* values) override; + void build() override; + + // Persisted, since a mapped index has no in-RAM vectors_ to derive it from. + size_t num_vectors() const override { return num_vectors_; } + + // Borrows a serialized index from a file mapping. `pos` is where the + // payload begins. + static DiskSeismicIndex* mmap_index(int dimension, const char* index_file, + size_t pos); + +protected: + std::vector clustered_inverted_lists; + +private: + void write_index(IOWriter* io_writer) override; + // Unsupported: the inline forward index is mmap-only. Throws. + void read_index(IOReader* io_reader, int io_flags = 0) override; + + auto search(idx_t n, const idx_t* indptr, const term_t* indices, + const float* values, int k, + SearchParameters* search_parameters = nullptr) + -> pair_of_score_id_vectors_t override; + + // Selects the global top-k_prime blocks across `cuts` by summary score and + // scores their docs. `dense` is per-thread scratch, zero on entry and + // restored on exit; `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, + int k_prime, SearchParameters* search_parameters) + -> pair_of_score_id_vector_t; + + SeismicClusterParameters cluster_parameter_; + // Borrows from the base's mapped_file_, so it is declared here (destroyed + // before the mapping). Empty after a fresh build (search uses vectors_ + // then). + detail::InlineForwardIndex fwd_; + size_t num_vectors_ = 0; +}; +} // namespace nsparse + +#endif // DISK_SEISMIC_INDEX_H diff --git a/nsparse/index_factory.cpp b/nsparse/index_factory.cpp index 49eb758..7fcf35d 100644 --- a/nsparse/index_factory.cpp +++ b/nsparse/index_factory.cpp @@ -16,6 +16,7 @@ #include #include "nsparse/brutal_index.h" +#include "nsparse/disk_seismic_index.h" #include "nsparse/id_map_index.h" #include "nsparse/index.h" #include "nsparse/inverted_index.h" @@ -111,6 +112,16 @@ Index* index_factory(int dimension, const char* description) { .seed = seed}); } + if (index_type == "disk_seismic") { + int lambda = std::stoi(get_param("lambda", "10")); + int beta = std::stoi(get_param("beta", "5")); + float alpha = std::stof(get_param("alpha", "0.5")); + int seed = std::stoi(get_param("seed", std::to_string(kRandomSeed))); + return new DiskSeismicIndex( + dimension, + {.lambda = lambda, .beta = beta, .alpha = alpha, .seed = seed}); + } + if (index_type == "seismic_sq") { std::string quantizer_str = get_param("quantizer", "8bit"); QuantizerType quantizer_type = QuantizerType::QT_8bit; diff --git a/nsparse/io/index_io.cpp b/nsparse/io/index_io.cpp index 135e7c1..3e41500 100644 --- a/nsparse/io/index_io.cpp +++ b/nsparse/io/index_io.cpp @@ -14,6 +14,7 @@ #include #include "nsparse/brutal_index.h" +#include "nsparse/disk_seismic_index.h" #include "nsparse/id_map_index.h" #include "nsparse/inverted_index.h" #include "nsparse/io/file_io.h" @@ -28,6 +29,7 @@ constexpr uint32_t SEIS = fourcc(SeismicIndex::name); constexpr uint32_t SESQ = fourcc(SeismicScalarQuantizedIndex::name); constexpr uint32_t IDMP = fourcc(IDMapIndex::name); constexpr uint32_t INVT = fourcc(InvertedIndex::name); +constexpr uint32_t DSEI = fourcc(DiskSeismicIndex::name); // Closes a stream once, on whichever path leaves the scope. // @@ -82,6 +84,8 @@ Index* mmap_index_payload(uint32_t id, int dimension, const char* file_name, pos); case INVT: return InvertedIndex::mmap_index(dimension, file_name, pos); + case DSEI: + return DiskSeismicIndex::mmap_index(dimension, file_name, pos); default: return nullptr; } @@ -108,6 +112,8 @@ Index* read_header(IOReader* io_reader) { return new SeismicIndex(dimension); case SESQ: return new SeismicScalarQuantizedIndex(dimension); + case DSEI: + return new DiskSeismicIndex(dimension); case IDMP: return new IDMapIndex(); case INVT: diff --git a/nsparse/python/swignsparse.swig b/nsparse/python/swignsparse.swig index 344c699..b5bd011 100644 --- a/nsparse/python/swignsparse.swig +++ b/nsparse/python/swignsparse.swig @@ -25,6 +25,7 @@ #include "nsparse/seismic_common.h" #include "nsparse/seismic_index.h" #include "nsparse/seismic_scalar_quantized_index.h" +#include "nsparse/disk_seismic_index.h" #include "nsparse/id_map_index.h" #include "nsparse/utils/scalar_quantizer.h" #include "nsparse/cluster/random_kmeans.h" @@ -138,6 +139,7 @@ import_array(); %ignore nsparse::SeismicIndex::mmap_index; %ignore nsparse::SeismicScalarQuantizedIndex::mmap_index; %ignore nsparse::InvertedIndex::mmap_index; +%ignore nsparse::DiskSeismicIndex::mmap_index; %ignore nsparse::InvertedListClusters::serialize; %ignore nsparse::InvertedListClusters::deserialize; @@ -188,6 +190,7 @@ import_array(); %include "nsparse/seismic_common.h" %include "nsparse/seismic_index.h" %include "nsparse/seismic_scalar_quantized_index.h" +%include "nsparse/disk_seismic_index.h" %include "nsparse/inverted_index.h" %include "nsparse/id_map_index.h" %include "nsparse/cluster/random_kmeans.h" diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5a80709..d404ee3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -30,6 +30,7 @@ set(NSPARSE_TEST_SRC ranker_test.cpp scalar_quantizer_test.cpp seismic_common_test.cpp + disk_seismic_index_test.cpp seismic_index_test.cpp seismic_invlists_writer_test.cpp seismic_scalar_quantized_index_test.cpp diff --git a/tests/disk_seismic_index_test.cpp b/tests/disk_seismic_index_test.cpp new file mode 100644 index 0000000..7b6f9e9 --- /dev/null +++ b/tests/disk_seismic_index_test.cpp @@ -0,0 +1,279 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#include "nsparse/disk_seismic_index.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "nsparse/index_factory.h" +#include "nsparse/io/buffered_io.h" +#include "nsparse/io/index_io.h" +#include "nsparse/seismic_index.h" +#include "nsparse/types.h" + +namespace nsparse { +namespace { + +// Fixed cluster params + seed so SeismicIndex and DiskSeismicIndex build +// bit-identical clusters and their results can be compared exactly. +constexpr int kLambda = 32; +constexpr int kBeta = 8; +constexpr float kAlpha = 0.4F; +constexpr int kSeed = 42; +constexpr int kDim = 400; + +SeismicClusterParameters cluster_params() { + return {.lambda = kLambda, .beta = kBeta, .alpha = kAlpha, .seed = kSeed}; +} + +struct CSR { + std::vector indptr; + std::vector indices; + std::vector values; + idx_t n = 0; +}; + +// A reproducible random sparse corpus: each row has a random number of distinct +// terms (sorted ascending, CSR convention) with values in (0, 1]. +CSR make_corpus(idx_t rows, unsigned seed) { + std::mt19937 rng(seed); + std::uniform_int_distribution nnz_dist(5, 25); + std::uniform_int_distribution term_dist(0, kDim - 1); + std::uniform_real_distribution val_dist(0.05F, 1.0F); + CSR c; + c.n = rows; + c.indptr.push_back(0); + for (idx_t r = 0; r < rows; ++r) { + std::set terms; + const int nnz = nnz_dist(rng); + while (static_cast(terms.size()) < nnz) + terms.insert(term_dist(rng)); + for (const int t : terms) { + c.indices.push_back(static_cast(t)); + c.values.push_back(val_dist(rng)); + } + c.indptr.push_back(static_cast(c.indices.size())); + } + return c; +} + +void add_corpus(Index& index, const CSR& c) { + index.add(c.n, c.indptr.data(), c.indices.data(), c.values.data()); +} + +// Index file removed on destruction. write_index/read_index take char*. +class TempIndexFile { +public: + explicit TempIndexFile(const std::string& name) + : path_(std::filesystem::temp_directory_path() / name) { + std::filesystem::remove(path_); + } + ~TempIndexFile() { std::filesystem::remove(path_); } + TempIndexFile(const TempIndexFile&) = delete; + TempIndexFile& operator=(const TempIndexFile&) = delete; + char* c_str() { return path_str_.data(); } + +private: + std::filesystem::path path_{}; + std::string path_str_ = path_.string(); +}; + +using ScoreIds = + std::pair>, std::vector>>; + +// K' large enough to select every candidate block (saturates the budget). +constexpr int kAllBlocks = 1'000'000; + +ScoreIds search_all(Index& index, const CSR& queries, int k, + SearchParameters* params) { + std::vector distances(static_cast(queries.n) * k); + std::vector labels(static_cast(queries.n) * k); + index.search(queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, distances.data(), labels.data(), + params); + ScoreIds out; + out.first.resize(queries.n); + out.second.resize(queries.n); + for (idx_t q = 0; q < queries.n; ++q) { + for (int j = 0; j < k; ++j) { + out.first[q].push_back(distances[static_cast(q) * k + j]); + out.second[q].push_back(labels[static_cast(q) * k + j]); + } + } + return out; +} + +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) { + EXPECT_EQ(a.second[q], b.second[q]) << "labels differ at query " << q; + ASSERT_EQ(a.first[q].size(), b.first[q].size()); + for (size_t j = 0; j < a.first[q].size(); ++j) { + EXPECT_FLOAT_EQ(a.first[q][j], b.first[q][j]) + << "score differs at query " << q << " rank " << j; + } + } +} + +} // namespace + +// Bit-exact anchor: reading a selected block's vectors from the inline mapping +// (fwd_) gives the same result as reading them from the in-RAM CSR (vectors_) +// of a fresh build, at the same K' -- same clusters, same selected blocks, same +// docs and dots. +TEST(DiskSeismicIndex, MappedReloadMatchesFreshBuild) { + const CSR corpus = make_corpus(1500, /*seed=*/1); + const CSR queries = make_corpus(40, /*seed=*/2); + + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + DiskSeismicSearchParameters params(/*cut=*/10, /*k_prime=*/32); + const ScoreIds fresh = search_all(disk, queries, 10, ¶ms); // vectors_ + + TempIndexFile file("nsparse_disk_seismic_parity.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + EXPECT_EQ(mapped->num_vectors(), static_cast(corpus.n)); + expect_same_results(search_all(*mapped, queries, 10, ¶ms), + fresh); // fwd_ +} + +// K' is a monotonic depth budget: the global top-K' blocks are nested by +// summary score, so a larger budget scores a superset of docs and every rank's +// score can only improve. This also proves the parameter is actually consumed. +TEST(DiskSeismicIndex, BlockBudgetIsMonotone) { + const CSR corpus = make_corpus(1500, /*seed=*/7); + const CSR queries = make_corpus(40, /*seed=*/8); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + TempIndexFile file("nsparse_disk_seismic_monotone.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + + ScoreIds prev; + bool have_prev = false; + bool any_improved = false; + for (const int kp : {1, 4, 16, 64, kAllBlocks}) { + DiskSeismicSearchParameters params(/*cut=*/10, kp); + const ScoreIds cur = search_all(*mapped, queries, 10, ¶ms); + if (have_prev) { + for (size_t q = 0; q < cur.first.size(); ++q) { + for (size_t j = 0; j < cur.first[q].size(); ++j) { + EXPECT_GE(cur.first[q][j], prev.first[q][j] - 1e-6F) + << "rank " << j << " regressed as K' grew, query " << q; + if (cur.first[q][j] > prev.first[q][j] + 1e-6F) { + any_improved = true; + } + } + } + } + prev = cur; + have_prev = true; + } + EXPECT_TRUE(any_improved) << "varying K' had no effect on any result"; +} + +// Once K' reaches the candidate-block count, a larger budget changes nothing. +TEST(DiskSeismicIndex, BlockBudgetSaturates) { + const CSR corpus = make_corpus(1200, /*seed=*/3); + const CSR queries = make_corpus(30, /*seed=*/4); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + TempIndexFile file("nsparse_disk_seismic_saturate.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + DiskSeismicSearchParameters big(/*cut=*/10, kAllBlocks); + DiskSeismicSearchParameters bigger(/*cut=*/10, kAllBlocks * 2); + expect_same_results(search_all(*mapped, queries, 10, &big), + search_all(*mapped, queries, 10, &bigger)); +} + +// K' (block budget) must be positive. +TEST(DiskSeismicIndex, RejectsNonPositiveBlockBudget) { + const CSR corpus = make_corpus(300, /*seed=*/5); + const CSR queries = make_corpus(2, /*seed=*/6); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + std::vector distances(2 * 10); + std::vector labels(2 * 10); + Index& idx = disk; // search() is a public method of the Index base + for (const int bad : {0, -1}) { + DiskSeismicSearchParameters params(/*cut=*/10, bad); + EXPECT_THROW(idx.search(queries.n, queries.indptr.data(), + queries.indices.data(), queries.values.data(), + 10, distances.data(), labels.data(), ¶ms), + std::invalid_argument); + } +} + +// mmap-only: the copying read path is unsupported. +TEST(DiskSeismicIndex, CopyingReadThrows) { + const CSR corpus = make_corpus(300, /*seed=*/5); + DiskSeismicIndex disk(kDim, cluster_params()); + add_corpus(disk, corpus); + disk.build(); + TempIndexFile file("nsparse_disk_seismic_copying.idx"); + write_index(&disk, file.c_str()); + // No kUseMmap -> falls to the copying reader, which throws. + EXPECT_THROW(read_index(file.c_str(), /*io_flags=*/0), std::runtime_error); + // A stream reader cannot be mapped either, so the flag still throws. + BufferedIOWriter writer; + write_index(&disk, &writer); + BufferedIOReader reader(writer.data()); + EXPECT_THROW(read_index(&reader, IndexIoFlag::kUseMmap), + std::runtime_error); +} + +// An empty index (no docs) round-trips: count 0, no results, no crash. +TEST(DiskSeismicIndex, EmptyIndexRoundTrip) { + DiskSeismicIndex disk(kDim, cluster_params()); // no add(), no build() + TempIndexFile file("nsparse_disk_seismic_empty.idx"); + write_index(&disk, file.c_str()); + std::unique_ptr mapped( + read_index(file.c_str(), IndexIoFlag::kUseMmap)); + ASSERT_NE(mapped, nullptr); + EXPECT_EQ(mapped->num_vectors(), 0U); + + const CSR queries = make_corpus(3, /*seed=*/6); + std::vector distances(3 * 5, -1.0F); + std::vector labels(3 * 5, detail::INVALID_IDX); + SeismicSearchParameters params(10, 1.0F); + mapped->search(queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), 5, distances.data(), labels.data(), + ¶ms); // must not crash; nothing found +} + +// index_factory understands the "disk_seismic" descriptor. +TEST(DiskSeismicIndex, FactoryCreatesIt) { + std::unique_ptr index( + index_factory(kDim, "disk_seismic,lambda=32|beta=8|alpha=0.4|seed=42")); + ASSERT_NE(index, nullptr); + EXPECT_EQ(index->id(), DiskSeismicIndex::name); +} + +} // namespace nsparse