From 6affd309c1cd5869baefde65df37cc6cb13f6bad Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Tue, 1 Sep 2026 08:52:17 +0000 Subject: [PATCH 01/15] io: expose write_header so a streamed payload can lay out its own header The index header is written centrally, by write_index, so an index type never has to know its layout. A writer that streams a payload out without an Index object -- the term-batched Seismic build in the next commit -- has no way in. Moving write_header from the anonymous namespace into detail keeps one definition of the layout rather than a second copy that can drift from read_header. Signed-off-by: Liyun Xiu --- nsparse/io/index_io.cpp | 24 ++++++++++++------------ nsparse/io/index_io.h | 5 +++++ 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/nsparse/io/index_io.cpp b/nsparse/io/index_io.cpp index d1ad869..aea89ec 100644 --- a/nsparse/io/index_io.cpp +++ b/nsparse/io/index_io.cpp @@ -111,18 +111,6 @@ std::string id_to_string(uint32_t id_val) { return chars; } -void write_header(const IndexHeader& header, IOWriter* io_writer) { - // write index type - uint32_t id_val = header.id; - io_writer->write(&id_val, sizeof(uint32_t), 1); - // write payload layout version - uint32_t version = header.version; - io_writer->write(&version, sizeof(uint32_t), 1); - // write dimension - int dimension = header.dimension; - io_writer->write(&dimension, sizeof(int), 1); -} - IndexHeader read_header(IOReader* io_reader) { IndexHeader header; io_reader->read(&header.id, sizeof(uint32_t), 1); @@ -171,6 +159,18 @@ void throw_if_version_unsupported(const IndexHeader& header, } // namespace namespace detail { +void write_header(const IndexHeader& header, IOWriter* io_writer) { + // write index type + uint32_t id_val = header.id; + io_writer->write(&id_val, sizeof(uint32_t), 1); + // write payload layout version + uint32_t version = header.version; + io_writer->write(&version, sizeof(uint32_t), 1); + // write dimension + int dimension = header.dimension; + io_writer->write(&dimension, sizeof(int), 1); +} + void write_index(Index* index, IOWriter* io_writer, bool keep_open) { auto* index_io = dynamic_cast(index); StreamCloser closer(io_writer, keep_open); diff --git a/nsparse/io/index_io.h b/nsparse/io/index_io.h index f1a8159..519194a 100644 --- a/nsparse/io/index_io.h +++ b/nsparse/io/index_io.h @@ -19,6 +19,11 @@ enum IndexIoFlag { }; namespace detail { +// The fixed prefix every serialized index starts with. Exposed because a writer +// that streams a payload out itself, rather than through an Index, still has to +// lay the header out exactly the way read_header expects — see +// build_seismic_index_batched. +void write_header(const IndexHeader& header, IOWriter* io_writer); void write_index(Index* index, IOWriter* io_writer, bool keep_open); // `filename`, when given, lets an index that was written for mmap borrow from // the file instead of copying; without one the copying path is used, since a From 0bfb160d0227819b1ad975027b169e11300405e4 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 06:08:27 +0000 Subject: [PATCH 02/15] Bound the seismic build's memory by batching the term space A whole-corpus build holds two intermediates that scale with the corpus's non-zeros -- the inverted lists, then the clustered posting lists -- so peak memory scales with the corpus, and a corpus whose posting lists do not fit in RAM cannot be indexed at all. The build now runs one contiguous term window at a time. for_each_clustered_window in seismic_common.cpp is the single place that work lives, and every index type in the family already reaches it: build_inverted_lists_clusters becomes a thin wrapper whose sink appends into one vector, so the batched and unbatched paths cannot drift. It takes the element width from SparseVectorsConfig, so a quantizing index gets the same treatment as a float one -- add() has already encoded the values by the time they arrive, and the previous float-only restriction was gratuitous. Two knobs, both on the BatchClusteringOption that #25 added and never wired up, so this needs no new API and reaches Python through the factory description unchanged: inverted_list_batch_size=N build in N windows; bounds the inverted-list intermediate. seismic, seismic_sq, disk_seismic and disk_seismic_sq all get this from build(). batch_file_output_path=P also serialize each window to P and free it, so the clustered lists are never all resident either. The index becomes the file: nothing is retained. The streaming write is write_seismic_index_batched, parameterized by the header and a prefix writer, which is all that differs between SEIS and SESQ. The disk-resident types keep their existing write for now: their payload interleaves summaries with an inline forward index derived from the same clusters, so one of those sections has to be held or spooled while the other streams. Folded the duplicate inverted_list_batch_size / char* batch_file_output_path fields on SeismicClusterParameters into the typed BatchClusteringOption, which removes a raw char* lifetime hazard from a parameters struct. That makes kDefaultSeismicClusterParams const rather than constexpr, the string not being a literal type. Measured on msmarco base_full (8.8M docs, dim 30109, 1.12B non-zeros, lambda=6000 beta=400 alpha=0.4) on a 36-core/68GB host, corpus on the heap in every row so the only difference is the build: windows peak RSS build time 1 (whole) 24091 MB 105 s 2 19300 MB 105 s 10 10524 MB 107 s 20 9284 MB 118 s 50 7917 MB 164 s 100 7425 MB 238 s Ten windows is 2.3x less peak memory for the same build time. Past that memory keeps falling but each window is another pass over the corpus, and by 100 those passes have more than doubled the build. The floor is the 6.7GB corpus, which read_csr can map instead of copying -- residency is SparseVectors' business, and the build is indifferent to it. Routing the ordinary build through the shared producer costs nothing: the whole-corpus row reproduces the pre-refactor 24085 MB / 106 s, and the 10-window row the pre-refactor 10516 MB / 103 s. The 2/20/50/100 rows are from the fuller pre-refactor sweep, whose endpoints those two reproduce. A single window streamed to a file rather than retained costs 24085 MB / 120 s -- the same memory, plus 15 s to put 14.9GB on disk. Query performance does not move, because the index is the same index: two independent unseeded builds, whole-corpus against 10 windows, over 6980 msmarco dev queries at k=10, gave 69332 vs 70198 QPS, p50 0.290 vs 0.296 ms, p99 1.000 vs 1.055 ms, recall@10 0.8406 vs 0.8432. Sameness is asserted, not assumed. At a fixed seed a streamed build is compared against build() + write_index as files, for both a float and a quantizing index, and across window counts; batch_size alone is compared the same way for a float, a quantizing and a disk-resident index. That holds because each list's k-means seed comes from its own GLOBAL term id rather than from the window it landed in or the order the threads reached it, and lambda/beta are resolved once from the whole corpus. It also needs each list's doc ids to arrive ascending, which is why the fill walks documents serially: threading it would reorder them, and pruning sorts by value with a non-stable sort. Two smaller things the window structure needed. An up-front counting pass gives the exact postings-per-term, which sizes every window's lists so none of them grows and set_entries can adopt whole buffers instead of add_entry locking per posting; it also range-checks the terms, which the mapped read does not. And the cluster loop's OpenMP chunk is now derived from the window width rather than fixed at 64, because a narrow window would otherwise be handed out as two chunks and leave every thread but two idle -- that made 64 windows 3.5x slower than 16. Signed-off-by: Liyun Xiu --- nsparse/CMakeLists.txt | 3 + nsparse/index_factory.cpp | 7 + nsparse/seismic_batched_build.cpp | 84 ++++ nsparse/seismic_batched_build.h | 60 +++ nsparse/seismic_common.cpp | 252 ++++++++++++ nsparse/seismic_common.h | 97 +++-- nsparse/seismic_index.cpp | 28 +- nsparse/seismic_scalar_quantized_index.cpp | 31 +- python_tests/test_seismic_batched_build.py | 113 ++++++ tests/CMakeLists.txt | 1 + tests/seismic_batched_build_test.cpp | 430 +++++++++++++++++++++ 11 files changed, 1058 insertions(+), 48 deletions(-) create mode 100644 nsparse/seismic_batched_build.cpp create mode 100644 nsparse/seismic_batched_build.h create mode 100644 nsparse/seismic_common.cpp create mode 100644 python_tests/test_seismic_batched_build.py create mode 100644 tests/seismic_batched_build_test.cpp diff --git a/nsparse/CMakeLists.txt b/nsparse/CMakeLists.txt index 27cb615..9b0b889 100644 --- a/nsparse/CMakeLists.txt +++ b/nsparse/CMakeLists.txt @@ -8,7 +8,9 @@ set(NSPARSE_SRC index.cpp brutal_index.cpp + seismic_common.cpp seismic_index.cpp + seismic_batched_build.cpp seismic_scalar_quantized_index.cpp disk_seismic_index.cpp disk_seismic_scalar_quantized_index.cpp @@ -36,6 +38,7 @@ set(NSPARSE_HEADERS brutal_index.h mmap_index.h seismic_index.h + seismic_batched_build.h seismic_scalar_quantized_index.h disk_seismic_index.h disk_seismic_scalar_quantized_index.h diff --git a/nsparse/index_factory.cpp b/nsparse/index_factory.cpp index 2beeef6..de2d514 100644 --- a/nsparse/index_factory.cpp +++ b/nsparse/index_factory.cpp @@ -58,6 +58,13 @@ SeismicClusterParameters parse_cluster_params(const GetParam& get_param) { return {.lambda = std::stoi(get_param("lambda", "10")), .beta = std::stoi(get_param("beta", "5")), .alpha = std::stof(get_param("alpha", "0.5")), + // Term windows to build in, and where to stream the result if it is + // not to be retained. See BatchClusteringOption. + .batch_clustering = + {.batch_size = static_cast(std::stoul( + get_param("inverted_list_batch_size", "1"))), + .batch_file_output_path = + get_param("batch_file_output_path", "")}, .seed = std::stoi(get_param("seed", std::to_string(kRandomSeed)))}; } diff --git a/nsparse/seismic_batched_build.cpp b/nsparse/seismic_batched_build.cpp new file mode 100644 index 0000000..a2edc40 --- /dev/null +++ b/nsparse/seismic_batched_build.cpp @@ -0,0 +1,84 @@ +/** + * 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/seismic_batched_build.h" + +#include +#include +#include +#include +#include + +#include "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/io/file_io.h" +#include "nsparse/io/index_io.h" +#include "nsparse/io/io.h" +#include "nsparse/seismic_common.h" +#include "nsparse/sparse_vectors.h" + +namespace nsparse::detail { + +void write_seismic_index_batched( + const SparseVectors* vectors, const SparseVectorsConfig& config, + const SeismicClusterParameters& params, const IndexHeader& header, + const std::function& write_prefix, + const std::string& out_path) { + if (out_path.empty()) { + throw std::invalid_argument( + "write_seismic_index_batched: output path must not be empty"); + } + if (vectors == nullptr || vectors->num_vectors() == 0) { + throw std::invalid_argument( + "write_seismic_index_batched: corpus is empty; there is nothing to " + "stream"); + } + + // One writer for the whole file, windows serialized straight into it rather + // than spilled and concatenated: serialize() pads each array relative to the + // writer's current offset (see io/align.h), so bytes produced by a writer + // that started at 0 carry the wrong padding once appended at some other + // offset. Streaming through a single writer keeps pos() the true absolute + // offset. + FileIOWriter writer(const_cast(out_path.c_str())); + write_header(header, &writer); + write_prefix(&writer); + + // The list count, exactly where SeismicInvertedListsWriter::serialize puts + // it. It is the whole dimension, known before any window is built, which is + // what lets the lists be streamed after it rather than counted first. + size_t n_lists = config.dimension; + writer.write(&n_lists, sizeof(size_t), 1); + + // Windows arrive in ascending term order, so appending each in turn produces + // the same byte sequence as writing every list at once. + size_t next_term = 0; + for_each_clustered_window( + vectors, config, params, + [&](size_t term_begin, std::vector&& clusters) { + if (term_begin != next_term) { + // The layout carries no per-list offsets, so a gap or a repeat + // would silently shift every list after it. + throw std::runtime_error( + "write_seismic_index_batched: windows arrived out of order"); + } + for (const auto& list : clusters) { + list.serialize(&writer); + } + next_term = term_begin + clusters.size(); + // clusters freed on return, before the next window is built. + }); + if (next_term != config.dimension) { + throw std::runtime_error( + "write_seismic_index_batched: wrote " + std::to_string(next_term) + + " of " + std::to_string(config.dimension) + " posting lists"); + } + writer.close(); +} + +} // namespace nsparse::detail diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h new file mode 100644 index 0000000..0b909e1 --- /dev/null +++ b/nsparse/seismic_batched_build.h @@ -0,0 +1,60 @@ +/** + * 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 SEISMIC_BATCHED_BUILD_H +#define SEISMIC_BATCHED_BUILD_H + +#include +#include + +#include "nsparse/io/io.h" +#include "nsparse/seismic_common.h" +#include "nsparse/sparse_vectors.h" + +namespace nsparse::detail { + +// Builds a seismic-family index and writes it straight to `out_path`, one term +// window at a time, without ever holding the whole index in memory. +// +// The usual build holds two whole-corpus intermediates -- the inverted lists and +// then the clustered posting lists -- so its peak memory scales with the +// corpus's non-zeros, and a corpus whose posting lists do not fit in RAM cannot +// be indexed at all. for_each_clustered_window bounds the first to one window; +// serializing each window and dropping it, which is what this does, bounds the +// second. What is left resident is the forward corpus (which the caller already +// holds, at whatever residency SparseVectors was given) plus one window. +// +// Reached through an index's build(), by setting +// SeismicClusterParameters::batch_clustering.batch_file_output_path. The index is +// then the file, not the object: nothing is retained to serve or to write_index +// afterwards. +// +// `header` and `write_prefix` are what make this work for every type in the +// family rather than just SEIS. `write_prefix` writes whatever the type puts +// between the header and its posting lists -- the forward vectors, and for a +// quantizing index its quantization header first. The lists then follow in the +// byte-for-byte layout SeismicInvertedListsWriter produces, so the file is an +// ordinary index of that type: read it back with read_index, mapped or copying, +// exactly as if it had been built in memory and written with write_index. +// +// Identical to the unbatched build for a fixed `params.seed`, and identical +// whatever batch_size is, because every list's k-means seed comes from its own +// global term id -- see for_each_clustered_window. +// +// Throws if the corpus is empty: there would be no windows to stream, and a +// header-only file is not a readable index. +void write_seismic_index_batched( + const SparseVectors* vectors, const SparseVectorsConfig& config, + const SeismicClusterParameters& params, const IndexHeader& header, + const std::function& write_prefix, + const std::string& out_path); + +} // namespace nsparse::detail + +#endif // SEISMIC_BATCHED_BUILD_H diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp new file mode 100644 index 0000000..6569e4c --- /dev/null +++ b/nsparse/seismic_common.cpp @@ -0,0 +1,252 @@ +/** + * 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/seismic_common.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/cluster/random_kmeans.h" +#include "nsparse/invlists/inverted_lists.h" +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { +namespace { + +// A window's [begin, end) slice of the term space. +struct TermWindow { + size_t begin; + size_t end; + [[nodiscard]] size_t size() const { return end - begin; } +}; + +// Cuts [0, dimension) into `batches` near-equal windows. +std::vector make_windows(size_t dim, size_t batches) { + // Bounds are size_t, not term_t: dimension may be up to 65536 (term_t is + // uint16), so a term_t window boundary would wrap and silently drop terms. + const size_t per_batch = (dim + batches - 1) / batches; + std::vector windows; + for (size_t begin = 0; begin < dim; begin += per_batch) { + windows.push_back({begin, std::min(dim, begin + per_batch)}); + } + return windows; +} + +// Postings per term, over the whole corpus. +// +// Reads only the CSR's indices, not its values, so on a mapped corpus this +// faults in a third of the bytes a full pass would. +std::vector count_postings_per_term(const SparseVectors& vectors, + size_t dim) { + std::vector counts(dim, 0); + const idx_t* indptr = vectors.indptr_data(); + const term_t* indices = vectors.indices_data(); + const idx_t nnz = indptr[vectors.num_vectors()]; + for (idx_t j = 0; j < nnz; ++j) { + const size_t term = indices[j]; + if (term >= dim) { + throw std::invalid_argument( + "for_each_clustered_window: corpus has term " + + std::to_string(term) + " outside dimension " + + std::to_string(dim)); + } + ++counts[term]; + } + return counts; +} + +// The inverted lists of one term window, sized exactly from the counting pass so +// no list ever grows. +// +// Doc ids arrive in ascending order, because fill_from_corpus walks documents +// ascending, which is what a single-window build produces. That matters beyond +// tidiness: pruning sorts by value with a non-stable sort, and k-means then +// consumes the kept ids in order, so a different posting order is a different +// index. +class WindowLists { +public: + WindowLists(const std::vector& term_counts, + const TermWindow& window, size_t element_size) + : element_size_(element_size), + lists_(window.size(), element_size), + ids_(window.size()), + codes_(window.size()), + fill_(window.size(), 0) { + for (size_t local = 0; local < window.size(); ++local) { + const size_t count = term_counts[window.begin + local]; + ids_[local].resize(count); + codes_[local].resize(count * element_size); + } + } + + void add(size_t local_term, idx_t doc_id, const uint8_t* code) { + const size_t slot = fill_[local_term]++; + if (slot >= ids_[local_term].size()) { + // The list was sized from the counting pass, so more postings than + // that means the corpus changed under us. One compare on a path that + // is bound by reading the corpus, and it is the difference between + // an exception and a heap overflow. + throw std::runtime_error( + "for_each_clustered_window: corpus changed during the build"); + } + ids_[local_term][slot] = doc_id; + std::memcpy(codes_[local_term].data() + slot * element_size_, code, + element_size_); + } + + // Hands the staged postings to the lists themselves. Separate from add() + // because set_entries adopts whole buffers, which is the only way to fill a + // list without the per-posting locking add_entry pays for. + ArrayInvertedLists& seal() { + for (size_t local = 0; local < ids_.size(); ++local) { + lists_[local].set_entries(std::move(ids_[local]), + std::move(codes_[local])); + } + ids_ = {}; + codes_ = {}; + return lists_; + } + +private: + size_t element_size_; + ArrayInvertedLists lists_; + std::vector> ids_; + std::vector> codes_; + std::vector fill_; +}; + +// Finds the window's postings by scanning the corpus. +// +// Serial over documents, ascending: that ordering is what makes the output +// independent of the window count (see WindowLists), and the pass is bound by +// reading the corpus rather than by the work per posting, so threading it would +// cost the ordering and buy nothing. +void fill_from_corpus(const SparseVectors& vectors, const TermWindow& window, + size_t element_size, WindowLists* lists) { + const idx_t* indptr = vectors.indptr_data(); + const term_t* indices = vectors.indices_data(); + const uint8_t* codes = vectors.values_data(); + const auto n_docs = static_cast(vectors.num_vectors()); + for (idx_t doc = 0; doc < n_docs; ++doc) { + for (idx_t j = indptr[doc]; j < indptr[doc + 1]; ++j) { + const size_t term = indices[j]; + if (term < window.begin || term >= window.end) { + continue; + } + lists->add(term - window.begin, doc, + codes + static_cast(j) * element_size); + } + } +} + +// Lists per OpenMP chunk, at most. Posting lists are wildly uneven in length, so +// they are handed out dynamically rather than split up front. +constexpr size_t kMaxClusterChunk = 64; + +// Chunks a window should break into, so the threads have something to steal. A +// window can be far narrower than the whole term space -- at 64 batches over +// 8192 terms it is 128 lists, which the flat chunk above would hand out as two +// chunks and leave every thread but two idle. +constexpr size_t kMinClusterChunks = 256; + +// SeismicClusterParameters with the "compute me a default" values already +// resolved. Resolved once for the whole build rather than per window: lambda +// comes from the GLOBAL corpus size, and a window-local one would prune +// differently and make the window count visible in the output. +struct ResolvedParameters { + int lambda; + int beta; + float alpha; + uint32_t base_seed; +}; + +std::vector cluster_window(const SparseVectors& vectors, + ArrayInvertedLists& lists, + const TermWindow& window, + const ResolvedParameters& params) { + std::vector clustered(window.size()); + const auto chunk = static_cast(std::clamp( + window.size() / kMinClusterChunks, 1, kMaxClusterChunk)); +#pragma omp parallel for schedule(dynamic, chunk) + for (int64_t local = 0; local < static_cast(window.size()); + ++local) { + auto& invlist = lists[local]; + const auto& doc_ids = invlist.prune_and_keep_doc_ids(params.lambda); + // Offset by the list's own GLOBAL term id, so neither the window it + // landed in nor which thread picked it up can change the result. + const auto seed = + params.base_seed + + static_cast(window.begin + static_cast(local)); + InvertedListClusters ilc( + RandomKMeans::train(&vectors, doc_ids, params.beta, seed)); + ilc.summarize(&vectors, params.alpha); + clustered[local] = std::move(ilc); + invlist.clear(); + } + return clustered; +} + +} // namespace + +void for_each_clustered_window(const SparseVectors* vectors, + const SparseVectorsConfig& config, + const SeismicClusterParameters& params, + const ClusteredWindowSink& sink) { + if (vectors == nullptr || vectors->num_vectors() == 0) { + return; + } + if (vectors->get_element_size() != config.element_size) { + throw std::invalid_argument( + "for_each_clustered_window: corpus element width does not match the " + "index's"); + } + + const size_t dim = config.dimension; + const size_t batches = + std::max(1, std::min(params.batch_clustering.batch_size, dim)); + + const int lambda = calculate_lambda(params.lambda, vectors->num_vectors()); + const ResolvedParameters resolved = { + .lambda = lambda, + .beta = calculate_beta(params.beta, lambda), + .alpha = params.alpha, + // Resolved once, outside the loop: std::random_device usually opens + // /dev/urandom per construction, so drawing per posting list would put a + // syscall on every iteration with every thread doing it. Once for the + // whole build, not per window, or the window count would be observable. + .base_seed = params.seed == kRandomSeed + ? std::random_device{}() + : static_cast(params.seed)}; + + // Exact per-term sizes, so no window's list ever grows and the bulk + // set_entries path can be used instead of per-posting locking. Also the only + // place a term outside the dimension is caught: the mapped read does not + // range-check. + const std::vector term_counts = + count_postings_per_term(*vectors, dim); + + for (const TermWindow& window : make_windows(dim, batches)) { + WindowLists lists(term_counts, window, config.element_size); + fill_from_corpus(*vectors, window, config.element_size, &lists); + sink(window.begin, + cluster_window(*vectors, lists.seal(), window, resolved)); + // The window's lists and clusters are freed here, before the next + // window is built. That is what bounds the memory. + } +} + +} // namespace nsparse::detail diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index a9b360e..c89e3a0 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -11,6 +11,7 @@ #define SEISMIC_COMMON_H #include +#include #include #include #include @@ -26,7 +27,22 @@ namespace nsparse { +// How a build bounds its own memory. +// +// Clustering the whole term space at once holds two intermediates for the whole +// corpus: the inverted lists (every posting) and then the clustered lists. Both +// scale with the corpus's non-zeros, which is what puts a ceiling on the corpus +// an index can be built from. Splitting the term space into `batch_size` +// contiguous windows and finishing one window before starting the next makes the +// first of those proportional to a window instead. +// +// `batch_file_output_path` bounds the second as well: with it set, each window's +// clustered lists are serialized to that path and freed as they are produced, so +// the build retains nothing and the index is the file rather than the object. +// See build_seismic_index_batched. struct BatchClusteringOption { + // Contiguous term windows. <= 1 means one window, i.e. no batching. Clamped + // to the dimension, since a window cannot be narrower than one term. size_t batch_size = 1; std::string batch_file_output_path; }; @@ -39,8 +55,7 @@ struct SeismicClusterParameters { int lambda; int beta; float alpha; - int inverted_list_batch_size = 1; - char* batch_file_output_path = nullptr; + BatchClusteringOption batch_clustering; // Fix this to make a build reproducible; two builds of the same corpus // differ by default. int seed = kRandomSeed; @@ -55,7 +70,9 @@ constexpr float kDefaultBetaRatio = 0.1F; constexpr int kDefaultBeta = -1; constexpr float kDefaultAlpha = 0.4F; -constexpr SeismicClusterParameters kDefaultSeismicClusterParams = { +// const rather than constexpr: BatchClusteringOption holds a std::string for the +// output path, which is not a literal type. +inline const SeismicClusterParameters kDefaultSeismicClusterParams = { .lambda = kDefaultLambda, .beta = kDefaultBeta, .alpha = kDefaultAlpha}; inline std::vector calculate_summary_scores( @@ -143,44 +160,48 @@ inline int calculate_beta(int beta, int lambda) { return beta; } +// Clusters and summarizes the posting lists of one term window at a time, +// handing each window to `sink` in ascending term order. +// +// This is the one place the seismic family's build work lives: every index type +// reaches it, either through build_inverted_lists_clusters below or through the +// streaming build in seismic_batched_build.h. The element width comes from +// `config`, so a quantizing index gets the same treatment as a float one -- the +// values in `vectors` are already encoded by the time they arrive here. +// +// `sink` receives the window's global first term and its lists, and must not +// hold on to them: they are freed as soon as it returns, which is what bounds +// the memory. Windows come from params.batch_clustering.batch_size. +// +// Every window's lambda and beta are computed from the GLOBAL corpus, and every +// list's k-means seed from its own GLOBAL term id, so the window count cannot +// change what is produced -- see the batched-build tests, which assert file +// equality against an unbatched build. +using ClusteredWindowSink = + std::function&& clusters)>; + +void for_each_clustered_window(const SparseVectors* vectors, + const SparseVectorsConfig& config, + const SeismicClusterParameters& params, + const ClusteredWindowSink& sink); + +// Every term's clustered posting list, in term order. The whole-corpus form of +// for_each_clustered_window: batch_size still bounds the inverted-list +// intermediate, but the result is retained in full, so this is bounded by the +// clustered lists rather than by a window. inline std::vector build_inverted_lists_clusters( const SparseVectors* vectors, const SparseVectorsConfig& config, const SeismicClusterParameters& seismic_cluster_params) { - // build inverted index - std::unique_ptr inverted_lists = - ArrayInvertedLists::build_inverted_lists(config.dimension, - config.element_size, vectors); - int lambda = - calculate_lambda(seismic_cluster_params.lambda, vectors->num_vectors()); - int beta = calculate_beta(seismic_cluster_params.beta, lambda); - size_t inverted_lists_size = inverted_lists->size(); - std::vector clustered_inverted_lists( - inverted_lists_size); - - // Resolved once, outside the loop: std::random_device usually opens - // /dev/urandom per construction, so drawing per posting list would put a - // syscall on every iteration with every thread doing it. - // uint32_t, not int: std::random_device yields a full unsigned 32-bit - // value, and the per-list offset below must wrap rather than overflow. - const uint32_t base_seed = - seismic_cluster_params.seed == kRandomSeed - ? std::random_device{}() - : static_cast(seismic_cluster_params.seed); - -#pragma omp parallel for schedule(dynamic, 64) - for (int64_t idx = 0; idx < static_cast(inverted_lists_size); - ++idx) { - auto& invlist = (*inverted_lists)[idx]; - const auto& doc_ids = invlist.prune_and_keep_doc_ids(lambda); - // Offset by the list's own index, so which thread picks up which list - // cannot change the result. - InvertedListClusters inverted_list_clusters(detail::RandomKMeans::train( - vectors, doc_ids, beta, base_seed + static_cast(idx))); - inverted_list_clusters.summarize(vectors, seismic_cluster_params.alpha); - clustered_inverted_lists[idx] = std::move(inverted_list_clusters); - invlist.clear(); - } - return clustered_inverted_lists; + std::vector clustered(config.dimension); + for_each_clustered_window( + vectors, config, seismic_cluster_params, + [&clustered](size_t term_begin, + std::vector&& window) { + std::move(window.begin(), window.end(), + clustered.begin() + static_cast(term_begin)); + }); + return clustered; } } // namespace detail diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index 489c886..159588e 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -9,6 +9,8 @@ #include "nsparse/seismic_index.h" +#include "nsparse/seismic_batched_build.h" + #include #include #include @@ -142,11 +144,27 @@ void SeismicIndex::add(idx_t n, const idx_t* indptr, const term_t* indices, } void SeismicIndex::build() { - clustered_inverted_lists = std::move(detail::build_inverted_lists_clusters( - get_vectors(), - {.element_size = kElementSize, - .dimension = static_cast(get_dimension())}, - cluster_parameter_)); + const SparseVectorsConfig config = { + .element_size = kElementSize, + .dimension = static_cast(get_dimension())}; + const std::string& out_path = + cluster_parameter_.batch_clustering.batch_file_output_path; + if (!out_path.empty()) { + // Streamed straight to a file and not retained: see + // BatchClusteringOption. write_index afterwards would write an index with + // no posting lists, so this index is deliberately left empty. + detail::write_seismic_index_batched( + get_vectors(), config, cluster_parameter_, + {.id = fourcc(name), + .version = kFormatVersion, + .dimension = get_dimension()}, + [this](IOWriter* io_writer) { vectors_->serialize(io_writer); }, + out_path); + return; + } + clustered_inverted_lists = + detail::build_inverted_lists_clusters(get_vectors(), config, + cluster_parameter_); } auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index b15f9a5..9319e62 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -9,6 +9,8 @@ #include "nsparse/seismic_scalar_quantized_index.h" +#include "nsparse/seismic_batched_build.h" + #include #include @@ -191,11 +193,30 @@ ScalarQuantizer SeismicScalarQuantizedIndex::query_quantizer( } void SeismicScalarQuantizedIndex::build() { - clustered_inverted_lists = std::move(detail::build_inverted_lists_clusters( - get_vectors(), - {.element_size = sq_.bytes_per_value(), - .dimension = static_cast(get_dimension())}, - cluster_parameter_)); + const SparseVectorsConfig config = { + .element_size = sq_.bytes_per_value(), + .dimension = static_cast(get_dimension())}; + const std::string& out_path = + cluster_parameter_.batch_clustering.batch_file_output_path; + if (!out_path.empty()) { + // The quantization header comes first, exactly as write_index writes it; + // the codes in `vectors_` are already quantized, so the batched build + // needs no knowledge of the quantizer beyond its width. + detail::write_seismic_index_batched( + get_vectors(), config, cluster_parameter_, + {.id = fourcc(name), + .version = kFormatVersion, + .dimension = get_dimension()}, + [this](IOWriter* io_writer) { + write_quantization_header(io_writer); + vectors_->serialize(io_writer); + }, + out_path); + return; + } + clustered_inverted_lists = + detail::build_inverted_lists_clusters(get_vectors(), config, + cluster_parameter_); } auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, diff --git a/python_tests/test_seismic_batched_build.py b/python_tests/test_seismic_batched_build.py new file mode 100644 index 0000000..1bbcd78 --- /dev/null +++ b/python_tests/test_seismic_batched_build.py @@ -0,0 +1,113 @@ +# 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. + +"""Black-box tests for the batched build, driven only through the SWIG API. + +Batching is a build option rather than a separate entry point, so there is +nothing new to wrap: it is reached through the factory description, the same way +lambda and beta are. `inverted_list_batch_size` bounds the build's memory; +adding `batch_file_output_path` streams the index straight to that file instead +of retaining it, which is the path a corpus too large for RAM needs. +""" + +import numpy as np +import pytest + +import nsparse +from oracle import recall_at_k +from support import K, add_corpus, make_index, search + +LAMBDA = 25 +BETA = 4 +ALPHA = 0.4 +SEED = 42 +BASE = f"lambda={LAMBDA}|beta={BETA}|alpha={ALPHA}|seed={SEED}" + +# Calibrated against the session corpus; a floor, not a target. +RECALL_FLOOR = 0.80 + + +def streamed(corpus, out_path, batch_size, kind="seismic"): + """Build straight to `out_path`; returns nothing, the file is the index.""" + spec = ( + f"{kind},{BASE}|inverted_list_batch_size={batch_size}" + f"|batch_file_output_path={out_path}" + ) + index = nsparse.index_factory(corpus.dim, spec) + add_corpus(index, corpus) + index.build() + return str(out_path) + + +@pytest.mark.parametrize("batch_size", [1, 4, 32]) +def test_happy_case(batch_size, corpus, queries, oracle, tmp_path): + """build -> read back mapped -> query -> accuracy, at several splits.""" + path = streamed(corpus, tmp_path / "batched.idx", batch_size) + index = nsparse.read_index(path, nsparse.kUseMmap) + assert index.num_vectors() == corpus.n + assert index.get_dimension() == corpus.dim + + dists, labels = search(index, queries) + assert labels.shape == (queries.n, K) + assert dists.shape == (queries.n, K) + assert (labels[:, 0] >= 0).all(), "every query must return at least one hit" + + want_labels, _ = oracle + assert recall_at_k(labels, want_labels) >= RECALL_FLOOR + + +@pytest.mark.parametrize("kind", ["seismic", "seismic_sq"]) +def test_matches_in_memory_build(kind, corpus, tmp_path): + """At a fixed seed a streamed build is the in-memory build, byte for byte. + + Parametrized over a float and a quantizing index, because the shared build + only needs the code width -- add() has already encoded the values. + """ + in_memory = tmp_path / "memory.idx" + nsparse.write_index(make_index(f"{kind},{BASE}", corpus), str(in_memory)) + batched = streamed(corpus, tmp_path / "batched.idx", 4, kind=kind) + + assert in_memory.read_bytes() == open(batched, "rb").read() + + +def test_batch_size_alone_leaves_the_index_in_memory(corpus, queries, tmp_path): + """Without an output path, batching only bounds the build's intermediates. + + The index is still usable in memory and still the same index -- this is the + path every index type gets from build(), including the disk-resident ones. + """ + unbatched = make_index(f"seismic,{BASE}", corpus) + batched = make_index(f"seismic,{BASE}|inverted_list_batch_size=8", corpus) + assert batched.num_vectors() == corpus.n + + want_d, want_l = search(unbatched, queries) + got_d, got_l = search(batched, queries) + np.testing.assert_array_equal(got_l, want_l) + np.testing.assert_allclose(got_d, want_d, rtol=1e-6, atol=1e-6) + + +def test_batch_count_is_not_observable(corpus, queries, tmp_path): + """The split is a memory knob: at a fixed seed it cannot change the results.""" + one = streamed(corpus, tmp_path / "one.idx", 1) + many = streamed(corpus, tmp_path / "many.idx", 16) + + want_d, want_l = search(nsparse.read_index(one), queries) + got_d, got_l = search(nsparse.read_index(many), queries) + np.testing.assert_array_equal(got_l, want_l) + np.testing.assert_allclose(got_d, want_d, rtol=1e-6, atol=1e-6) + + +def test_rejects_dimension_smaller_than_corpus(corpus, tmp_path): + """A term the declared dimension does not cover is an error, not a silent drop.""" + spec = ( + f"seismic,{BASE}|inverted_list_batch_size=4" + f"|batch_file_output_path={tmp_path / 'bad.idx'}" + ) + index = nsparse.index_factory(corpus.dim // 2, spec) + add_corpus(index, corpus) + with pytest.raises(ValueError): + index.build() diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 72eaff7..75da0a4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -33,6 +33,7 @@ set(NSPARSE_TEST_SRC disk_seismic_index_test.cpp disk_seismic_scalar_quantized_index_test.cpp seismic_index_test.cpp + seismic_batched_build_test.cpp seismic_invlists_writer_test.cpp seismic_scalar_quantized_index_test.cpp sparse_vectors_test.cpp diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp new file mode 100644 index 0000000..9532b08 --- /dev/null +++ b/tests/seismic_batched_build_test.cpp @@ -0,0 +1,430 @@ +/** + * 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/seismic_batched_build.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/disk_seismic_index.h" +#include "nsparse/index_factory.h" +#include "nsparse/io/file_io.h" +#include "nsparse/io/index_io.h" +#include "nsparse/io/seismic_invlists_writer.h" +#include "nsparse/seismic_index.h" +#include "nsparse/seismic_scalar_quantized_index.h" +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" +#include "nsparse/utils/csr_layout.h" + +namespace nsparse { +namespace { + +constexpr int kSeed = 42; +constexpr int kLambda = 64; +constexpr int kBeta = 6; +constexpr float kAlpha = 0.4F; + +struct Corpus { + int dim; + std::vector indptr; + std::vector indices; + std::vector values; + [[nodiscard]] idx_t n() const { + return static_cast(indptr.size()) - 1; + } +}; + +Corpus make_corpus(int n_docs, int dim, unsigned seed) { + std::mt19937 gen(seed); + std::uniform_int_distribution nnz_dist(3, 12); + std::uniform_int_distribution term_dist(0, dim - 1); + std::uniform_real_distribution val_dist(0.1F, 3.0F); + + Corpus corpus; + corpus.dim = dim; + corpus.indptr.push_back(0); + for (int doc = 0; doc < n_docs; ++doc) { + // Capped at dim: the loop below draws *distinct* terms, so asking for + // more than exist would never terminate. + int nnz = std::min(nnz_dist(gen), dim); + std::set terms; + while (static_cast(terms.size()) < nnz) { + terms.insert(term_dist(gen)); + } + for (int term : terms) { // ascending -> a valid CSR row + corpus.indices.push_back(static_cast(term)); + corpus.values.push_back(val_dist(gen)); + } + corpus.indptr.push_back(static_cast(corpus.indices.size())); + } + return corpus; +} + +// Removes its directory when it goes out of scope, so a failing EXPECT does not +// leave an index behind. +class TempDir { +public: + explicit TempDir(const std::string& tag) { + path_ = (std::filesystem::temp_directory_path() / + ("seismic_batched_" + tag + "_" + + std::to_string(std::random_device{}()))) + .string(); + std::filesystem::create_directories(path_); + } + ~TempDir() { + std::error_code ignored; + std::filesystem::remove_all(path_, ignored); + } + TempDir(const TempDir&) = delete; + TempDir& operator=(const TempDir&) = delete; + + [[nodiscard]] std::string file(const std::string& name) const { + return path_ + "/" + name; + } + +private: + std::string path_; +}; + +SeismicClusterParameters params_for(size_t batch_size, + const std::string& out_path, int seed) { + SeismicClusterParameters params = { + .lambda = kLambda, .beta = kBeta, .alpha = kAlpha}; + params.batch_clustering.batch_size = batch_size; + params.batch_clustering.batch_file_output_path = out_path; + params.seed = seed; + return params; +} + +std::vector read_file(const std::string& path) { + std::ifstream in(path, std::ios::binary); + return {std::istreambuf_iterator(in), + std::istreambuf_iterator()}; +} + +// A build streamed straight to `out`, through the index's own build(). +std::vector streamed(const Corpus& corpus, size_t batch_size, + const std::string& out, int seed = kSeed) { + SeismicIndex index(corpus.dim, params_for(batch_size, out, seed)); + index.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index.build(); + return read_file(out); +} + +// The same corpus built the ordinary way and written with write_index. +std::vector in_memory(const Corpus& corpus, const std::string& out, + size_t batch_size = 1, int seed = kSeed) { + SeismicIndex index(corpus.dim, params_for(batch_size, "", seed)); + index.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index.build(); + write_index(&index, const_cast(out.c_str())); + return read_file(out); +} + +// Per-term set of all doc ids across that term's clusters, parsed straight from +// the file through the public serialization surface. +std::vector> per_term_doc_sets(const std::string& path) { + FileIOReader reader(const_cast(path.c_str())); + uint32_t id_val = 0; + reader.read(&id_val, sizeof(uint32_t), 1); + uint32_t version = 0; + reader.read(&version, sizeof(uint32_t), 1); + int stored_dim = 0; + reader.read(&stored_dim, sizeof(int), 1); + EXPECT_EQ(id_val, fourcc(SeismicIndex::name)); + EXPECT_EQ(version, SeismicIndex::kFormatVersion); + SparseVectors vectors; + vectors.deserialize(&reader); + SeismicInvertedListsWriter writer; + writer.deserialize(&reader); + std::vector lists = writer.release(); + + std::vector> out(lists.size()); + for (size_t term = 0; term < lists.size(); ++term) { + for (size_t cluster = 0; cluster < lists[term].cluster_size(); + ++cluster) { + for (idx_t doc : + lists[term].get_docs(static_cast(cluster))) { + out[term].insert(doc); + } + } + } + return out; +} + +// The corpus in the interchange CSR layout, converted to native: what a mapped +// read consumes. +std::string write_native_csr(const Corpus& corpus, const std::string& path) { + std::ofstream out(path, std::ios::binary); + const std::array sizes = { + corpus.n(), corpus.dim, static_cast(corpus.indices.size())}; + out.write(reinterpret_cast(sizes.data()), sizeof(sizes)); + std::vector indptr64(corpus.indptr.begin(), corpus.indptr.end()); + out.write(reinterpret_cast(indptr64.data()), + static_cast(indptr64.size() * sizeof(int64_t))); + std::vector indices32(corpus.indices.begin(), + corpus.indices.end()); + out.write(reinterpret_cast(indices32.data()), + static_cast(indices32.size() * sizeof(int32_t))); + out.write( + reinterpret_cast(corpus.values.data()), + static_cast(corpus.values.size() * sizeof(float))); + out.close(); + const std::string native = csr_layout::native_path(path); + csr_layout::convert(path, native); + return native; +} + +} // namespace + +// The point of the seeding discipline: for a fixed seed a streamed build is not +// merely equivalent to build() + write_index, it is the same file. Each list's +// k-means seed comes from its own global term id, so neither the window a term +// landed in nor the order the threads reached it can leak into the output. +TEST(SeismicBatchedBuild, StreamedBuildIsByteIdenticalToInMemoryBuild) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/42); + TempDir dir("identical"); + EXPECT_EQ(in_memory(corpus, dir.file("mem.dat")), + streamed(corpus, /*batch_size=*/4, dir.file("streamed.dat"))); +} + +// The window count is a memory knob, not a behaviour knob: at a fixed seed +// every count has to produce the same file. +TEST(SeismicBatchedBuild, StreamedBuildIsIdenticalAcrossBatchCounts) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/11); + TempDir dir("counts"); + const auto one = streamed(corpus, 1, dir.file("b1.dat")); + ASSERT_FALSE(one.empty()); + EXPECT_EQ(one, streamed(corpus, 2, dir.file("b2.dat"))); + EXPECT_EQ(one, streamed(corpus, 10, dir.file("b10.dat"))); + // More windows than terms is clamped to one term each, and 0 means one + // window rather than none. + EXPECT_EQ(one, streamed(corpus, 1000, dir.file("b1000.dat"))); + EXPECT_EQ(one, streamed(corpus, 0, dir.file("b0.dat"))); +} + +// batch_size alone bounds the inverted-list intermediate and leaves the index +// in memory. That is the path every index type gets, including the two disk +// ones, so it must not change what build() produces either. +TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeAnInMemoryBuild) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/5); + TempDir dir("inmem_batched"); + const auto unbatched = in_memory(corpus, dir.file("b1.dat"), 1); + ASSERT_FALSE(unbatched.empty()); + EXPECT_EQ(unbatched, in_memory(corpus, dir.file("b8.dat"), 8)); + EXPECT_EQ(unbatched, in_memory(corpus, dir.file("b64.dat"), 64)); +} + +// The disk-resident types share the same build, so batch_size has to bound +// their intermediates too without changing what they produce. They have no +// streaming write yet -- their payload interleaves summaries with an inline +// forward index +// -- so this covers the half they do get. +TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeADiskIndex) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/71); + TempDir dir("disk"); + + auto build_disk = [&corpus](size_t batch_size, const std::string& out) { + DiskSeismicIndex index(corpus.dim, params_for(batch_size, "", kSeed)); + index.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index.build(); + write_index(&index, const_cast(out.c_str())); + return read_file(out); + }; + + const auto unbatched = build_disk(1, dir.file("b1.dat")); + ASSERT_FALSE(unbatched.empty()); + EXPECT_EQ(unbatched, build_disk(8, dir.file("b8.dat"))); + EXPECT_EQ(unbatched, build_disk(64, dir.file("b64.dat"))); +} + +// The generalization that matters: a quantizing index streams too, because the +// codes in `vectors_` are already quantized by add() and the shared build only +// needs their width. +TEST(SeismicBatchedBuild, StreamsAQuantizedIndexIdenticallyToo) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/23); + TempDir dir("sq"); + const std::string mem_path = dir.file("mem.dat"); + const std::string streamed_path = dir.file("streamed.dat"); + + SeismicScalarQuantizedIndex mem(QuantizerType::QT_8bit, 0.0F, 3.0F, + params_for(1, "", kSeed), corpus.dim); + mem.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + write_index(&mem, const_cast(mem_path.c_str())); + + SeismicScalarQuantizedIndex batched(QuantizerType::QT_8bit, 0.0F, 3.0F, + params_for(4, streamed_path, kSeed), + corpus.dim); + batched.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + batched.build(); + + EXPECT_EQ(read_file(mem_path), read_file(streamed_path)); + // And it loads as the quantized type it claims to be. + std::unique_ptr reloaded( + read_index(const_cast(streamed_path.c_str()))); + EXPECT_EQ(reloaded->id(), SeismicScalarQuantizedIndex::name); + EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n())); +} + +// Same invariant without a seed, where the files legitimately differ: the +// doc-id membership of each term's list still cannot depend on the window +// count, because lambda is computed from the global corpus size. +TEST(SeismicBatchedBuild, PerTermMembershipInvariantAcrossBatches) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/42); + TempDir dir("membership"); + const std::string one = dir.file("1.dat"); + const std::string ten = dir.file("10.dat"); + streamed(corpus, 1, one, kRandomSeed); + streamed(corpus, 10, ten, kRandomSeed); + + auto sets1 = per_term_doc_sets(one); + auto sets10 = per_term_doc_sets(ten); + ASSERT_EQ(sets1.size(), static_cast(corpus.dim)); + ASSERT_EQ(sets10.size(), sets1.size()); + for (size_t term = 0; term < sets1.size(); ++term) { + EXPECT_EQ(sets1[term], sets10[term]) << "term " << term; + } +} + +// Regression for the term_t (uint16) window-bound overflow: a dimension at the +// 2^16 boundary must still build every term's list. Before the fix, size_t +// window bounds cast to term_t wrapped mod 65536, so dim=65536 built nothing +// while n_lists claimed 65536 -> a corrupt, unloadable file. +TEST(SeismicBatchedBuild, HandlesDimensionAt65536) { + const int dim = 65536; // term ids 0..65535 all fit term_t (uint16) + Corpus corpus = make_corpus(/*n_docs=*/3000, dim, /*seed=*/5); + TempDir dir("dim64k"); + const std::string path = dir.file("index.dat"); + streamed(corpus, 1, path); + + // Must load without "unexpected end of index file". + std::unique_ptr idx(read_index(const_cast(path.c_str()))); + EXPECT_EQ(idx->num_vectors(), static_cast(corpus.n())); + + auto sets = per_term_doc_sets(path); + EXPECT_EQ(sets.size(), static_cast(dim)); + size_t total_docs = 0; + for (const auto& docs : sets) { + total_docs += docs.size(); + } + EXPECT_GT(total_docs, 0U); +} + +// A streamed index must serve correctly through the mapped read path -- the way +// a caller whose corpus did not fit in RAM is going to use it. +TEST(SeismicBatchedBuild, SearchThroughMappedReadMatchesInMemory) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/7); + Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); + const int k = 10; + TempDir dir("search"); + const std::string streamed_path = dir.file("streamed.dat"); + + SeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); + mem.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + streamed(corpus, 4, streamed_path); + + std::unique_ptr disk(read_index( + const_cast(streamed_path.c_str()), IndexIoFlag::kUseMmap)); + + SeismicSearchParameters search_params(/*cut=*/3, /*heap_factor=*/1.0F); + const auto n = static_cast(queries.n()); + std::vector mem_dist(n * k); + std::vector mem_lab(n * k); + std::vector disk_dist(n * k); + std::vector disk_lab(n * k); + static_cast(mem).search(queries.n(), queries.indptr.data(), + queries.indices.data(), + queries.values.data(), k, mem_dist.data(), + mem_lab.data(), &search_params); + disk->search(queries.n(), queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, disk_dist.data(), disk_lab.data(), + &search_params); + + // Identical builds (same seed), so identical results, not merely close. + EXPECT_EQ(mem_lab, disk_lab); + EXPECT_EQ(mem_dist, disk_dist); +} + +// Residency is SparseVectors' business, not the build's: a corpus borrowed from +// a mapping must produce the same index as one on the heap. +TEST(SeismicBatchedBuild, MappedCorpusMatchesOwnedCorpus) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/13); + TempDir dir("mapped"); + const std::string owned_path = dir.file("owned.dat"); + const std::string mapped_path = dir.file("mapped.dat"); + streamed(corpus, 3, owned_path); + + const std::string native = write_native_csr(corpus, dir.file("corpus.csr")); + SeismicIndex mapped(corpus.dim, params_for(3, mapped_path, kSeed)); + mapped.read_csr(native.c_str(), Residency::kMmap); + mapped.build(); + + EXPECT_EQ(read_file(owned_path), read_file(mapped_path)); +} + +// Both knobs are reachable through the factory description, which is the only +// way the Python bindings can set them. +TEST(SeismicBatchedBuild, FactoryDescriptionDrivesTheBatchedBuild) { + Corpus corpus = make_corpus(/*n_docs=*/1000, /*dim=*/128, /*seed=*/31); + TempDir dir("factory"); + const std::string path = dir.file("index.dat"); + const std::string spec = + "seismic,lambda=64|beta=6|alpha=0.4|seed=42|" + "inverted_list_batch_size=8|batch_file_output_path=" + + path; + + std::unique_ptr index(index_factory(corpus.dim, spec.c_str())); + index->add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index->build(); + + ASSERT_TRUE(std::filesystem::exists(path)); + EXPECT_EQ(streamed(corpus, 8, dir.file("direct.dat")), read_file(path)); +} + +TEST(SeismicBatchedBuild, RejectsInvalidInput) { + Corpus corpus = make_corpus(/*n_docs=*/50, /*dim=*/16, /*seed=*/1); + TempDir dir("reject"); + + // A term the declared dimension does not cover. The mapped read path does + // not range-check terms, so this is the build's own guard -- without it the + // term would be silently dropped from the index. + SeismicIndex narrow(corpus.dim / 2, + params_for(1, dir.file("narrow.dat"), kSeed)); + narrow.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + EXPECT_THROW(narrow.build(), std::invalid_argument); + + // Streaming an empty corpus would leave a header-only file that read_index + // cannot parse, so it is refused rather than written. + SeismicIndex empty(corpus.dim, params_for(4, dir.file("empty.dat"), kSeed)); + EXPECT_THROW(empty.build(), std::invalid_argument); +} + +} // namespace nsparse From 955ebcf7aae0da2d24bcef99a5105590db37da1a Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 06:08:53 +0000 Subject: [PATCH 03/15] benchmarks: add a peak-RSS driver for the batched build google-benchmark measures throughput, and a memory high-water mark is only clean in a process that has built nothing else, so this is a standalone driver that runs one configuration per invocation and reports VmHWM plus wall time. It covers the whole-corpus build, the batched build at any window count, and the interchange -> native CSR conversion a mapped corpus needs. Two things it has to get right to be worth trusting. The corpus residency is a flag, not a mode. "batched inmem" streams the corpus onto the heap through the same streaming_add the baseline uses, so a comparison against the baseline isolates the batching. "batched mmap" has read_csr borrow a native CSR instead, which is cheaper by the whole size of the corpus -- a real saving, but one that comes from the residency, and reporting it against the baseline would credit batching with it. VmHWM is reset at the start of the build. Loading the corpus costs more than holding it, because a streaming ingest stages a second copy of it: on msmarco base_full the loader peaks at 13.0GB while the build at 100 windows peaks at 7.4GB. A whole-process high-water mark would report the loader for every configuration below that and hide the batching entirely. The loader's peak is reported alongside rather than dropped. This is what produced the numbers in the previous commit. Signed-off-by: Liyun Xiu --- benchmarks/CMakeLists.txt | 12 + benchmarks/batched_build_mem_bench.cpp | 298 +++++++++++++++++++++++++ 2 files changed, 310 insertions(+) create mode 100644 benchmarks/batched_build_mem_bench.cpp diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt index 73436c3..0e2f0c8 100644 --- a/benchmarks/CMakeLists.txt +++ b/benchmarks/CMakeLists.txt @@ -81,3 +81,15 @@ target_link_libraries(nsparse_build_benchmark PRIVATE absl::flat_hash_set absl::flat_hash_map ) + +# Peak-RSS build-memory driver for the term-batched build. Not a +# google-benchmark target: that harness reports throughput, and what matters here +# is the high-water mark of a single build, one configuration per process. +add_executable(batched_build_mem_bench batched_build_mem_bench.cpp) +target_include_directories(batched_build_mem_bench PRIVATE ${PROJECT_SOURCE_DIR}) +target_link_libraries(batched_build_mem_bench PRIVATE + ${NSPARSE_BENCH_LIB} + OpenMP::OpenMP_CXX + absl::flat_hash_set + absl::flat_hash_map +) diff --git a/benchmarks/batched_build_mem_bench.cpp b/benchmarks/batched_build_mem_bench.cpp new file mode 100644 index 0000000..097be0a --- /dev/null +++ b/benchmarks/batched_build_mem_bench.cpp @@ -0,0 +1,298 @@ +/** + * 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. + */ + +// Standalone driver measuring PEAK RESIDENT MEMORY (VmHWM) and wall time of a +// Seismic index build, for one configuration per process invocation. google- +// benchmark measures throughput, not peak RSS, and running each config in its +// own process gives a clean high-water mark uncontaminated by earlier builds. +// +// Usage: +// batched_build_mem_bench convert +// batched_build_mem_bench baseline \ +// [out_index] +// batched_build_mem_bench batched \ +// +// +// "convert" produces the native CSR the mapped read wants. "baseline" builds +// the in-memory SeismicIndex (streaming add of the interchange CSR, like the +// other benchmarks) -- the memory this feature exists to avoid. "batched" runs +// the term-batched build at the given batch count. +// +// Compare "batched inmem" against "baseline": both hold the corpus on the heap +// via the same streaming_add, so the difference between them is the batching +// and nothing else. "batched mmap" borrows a native CSR instead, which is +// cheaper by the whole size of the corpus -- a real saving, but one that comes +// from the residency rather than from batching, so it is not the baseline's +// counterpart. +// +// Point at a real disk: on a tmpfs such as /tmp the index is RAM, and +// the numbers are meaningless. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "nsparse/io/index_io.h" +#include "nsparse/seismic_index.h" +#include "nsparse/types.h" +#include "nsparse/utils/csr_layout.h" + +namespace { + +// Peak resident set size in KiB, read from /proc/self/status (VmHWM). +long read_vm_hwm_kib() { + std::ifstream status("/proc/self/status"); + std::string line; + while (std::getline(status, line)) { + if (line.rfind("VmHWM:", 0) == 0) { + long kib = 0; + std::sscanf(line.c_str(), "VmHWM: %ld kB", &kib); + return kib; + } + } + return -1; +} + +// Resets VmHWM to the current VmRSS, so a later read reports the peak since +// this call rather than since the process started. +// +// Necessary because loading the corpus costs more than holding it: +// streaming_add's staging buffers are a copy of the whole thing on top of the +// index's own. Without this the loader's high-water mark floors every reported +// number and a build whose true peak is below it is unmeasurable. The corpus +// stays resident across the reset, so it still counts toward the build's peak +// -- which is what should be compared. +void reset_vm_hwm() { + std::ofstream clear_refs("/proc/self/clear_refs"); + // 5 == CLEAR_REFS_MM_HIWATER_RSS. Linux-only, and only an accounting hint: + // if it is unavailable the numbers include the loader, so say so rather + // than reporting them as if they did not. + if (!clear_refs || !(clear_refs << "5\n")) { + std::cerr << "warning: cannot reset VmHWM; peak_rss_mb includes corpus " + "loading\n"; + return; + } +} + +double now_seconds() { + // CLOCK_MONOTONIC via clock() is process time; use wall clock instead. + timespec ts{}; + clock_gettime(CLOCK_MONOTONIC, &ts); + return static_cast(ts.tv_sec) + + static_cast(ts.tv_nsec) * 1e-9; +} + +// Streaming add of an interchange CSR (int64 sizes[3], int64 indptr, int32 +// indices, float values) into an index, in nnz-bounded batches (mirrors +// index_search_benchmark's streaming_add so the baseline peak is comparable). +void streaming_add(nsparse::Index* index, const std::string& path) { + std::ifstream file(path, std::ios::binary); + if (!file.is_open()) { + throw std::runtime_error("cannot open " + path); + } + int64_t sizes[3]; + file.read(reinterpret_cast(sizes), sizeof(sizes)); + const int64_t nrow = sizes[0]; + const int64_t nnz = sizes[2]; + + std::vector indptr64(nrow + 1); + file.read(reinterpret_cast(indptr64.data()), + static_cast((nrow + 1) * sizeof(int64_t))); + const auto indices_off = file.tellg(); + const auto data_off = static_cast(indices_off) + + static_cast(nnz * sizeof(int32_t)); + + constexpr int64_t kMaxBatchNnz = 1'500'000'000LL; + int64_t row_start = 0; + while (row_start < nrow) { + int64_t row_end = row_start + 1; + while (row_end < nrow && + (indptr64[row_end] - indptr64[row_start]) < kMaxBatchNnz) { + ++row_end; + } + const int64_t brows = row_end - row_start; + const int64_t bnnz = indptr64[row_end] - indptr64[row_start]; + const int64_t boff = indptr64[row_start]; + + std::vector bindptr(brows + 1); + for (int64_t i = 0; i <= brows; ++i) { + bindptr[i] = + static_cast(indptr64[row_start + i] - boff); + } + std::vector bindices(bnnz); + { + file.seekg(indices_off + + static_cast(boff * sizeof(int32_t))); + constexpr int64_t kChunk = 1 << 22; + std::vector tmp(std::min(kChunk, bnnz > 0 ? bnnz : 1)); + int64_t done = 0; + while (done < bnnz) { + int64_t take = std::min(kChunk, bnnz - done); + file.read(reinterpret_cast(tmp.data()), + static_cast(take * sizeof(int32_t))); + for (int64_t j = 0; j < take; ++j) { + bindices[done + j] = static_cast(tmp[j]); + } + done += take; + } + } + std::vector bdata(bnnz); + { + file.seekg(static_cast(data_off) + + static_cast(boff * sizeof(float))); + file.read(reinterpret_cast(bdata.data()), + static_cast(bnnz * sizeof(float))); + } + index->add(static_cast(brows), bindptr.data(), + bindices.data(), bdata.data()); + row_start = row_end; + } +} + +// Both layouts start with the same int64 (rows, cols, nnz) header. +int csr_dimension(const std::string& path) { + std::ifstream file(path, std::ios::binary); + int64_t sizes[3]; + file.read(reinterpret_cast(sizes), sizeof(sizes)); + return static_cast(sizes[1]); +} + +// `peak_rss_mb` is the build's own high-water mark (see reset_vm_hwm); +// `load_peak_rss_mb` is what loading the corpus cost before it, reported so a +// build peak that sits below the loader's is not mistaken for the whole story. +void report(const std::string& mode, const std::string& detail, double build_s, + long load_hwm_kib) { + const long hwm = read_vm_hwm_kib(); + std::cout << "RESULT mode=" << mode << " " << detail + << " build_s=" << build_s + << " peak_rss_mb=" << (static_cast(hwm) / 1024.0) + << " load_peak_rss_mb=" + << (static_cast(load_hwm_kib) / 1024.0) << "\n"; +} + +int run_convert(int argc, char** argv) { + if (argc < 4) { + std::cerr << "convert \n"; + return 2; + } + nsparse::csr_layout::convert(argv[2], argv[3]); + std::cout << "converted -> " << argv[3] << "\n"; + return 0; +} + +int run_baseline(int argc, char** argv) { + if (argc < 6) { + std::cerr << "baseline " + "[out_index]\n"; + return 2; + } + const std::string csr = argv[2]; + const nsparse::SeismicClusterParameters params = { + .lambda = std::atoi(argv[3]), + .beta = std::atoi(argv[4]), + .alpha = static_cast(std::atof(argv[5]))}; + nsparse::SeismicIndex index(csr_dimension(csr), params); + streaming_add(&index, csr); + + const long load_hwm = read_vm_hwm_kib(); + reset_vm_hwm(); + const double started = now_seconds(); + index.build(); + report("baseline", "batches=0", now_seconds() - started, load_hwm); + + if (argc >= 7) { + const std::string out = argv[6]; + nsparse::write_index(&index, const_cast(out.c_str())); + std::ifstream file(out, std::ios::binary | std::ios::ate); + std::cout << "index_bytes=" << file.tellg() << "\n"; + } + return 0; +} + +int run_batched(int argc, char** argv) { + if (argc < 9) { + std::cerr << "batched " + " \n"; + return 2; + } + const std::string corpus_residency = argv[2]; + const std::string csr = argv[3]; + const nsparse::SeismicClusterParameters params = { + .lambda = std::atoi(argv[4]), + .beta = std::atoi(argv[5]), + .alpha = static_cast(std::atof(argv[6]))}; + const std::string out = std::string(argv[8]) + "/index.seismic.dat"; + nsparse::SeismicClusterParameters batched_params = params; + batched_params.batch_clustering.batch_size = + static_cast(std::atoi(argv[7])); + batched_params.batch_clustering.batch_file_output_path = out; + + // Which residency the corpus is held at is the point of the flag: + // + // inmem -- streaming_add of an interchange CSR, byte for byte what the + // baseline above does. This is the comparison that isolates + // batching: same corpus, on the heap, in both arms. + // mmap -- a native CSR, borrowed. Cheaper by the size of the corpus, but + // not comparable to the baseline, because the saving is the + // residency rather than the batching. + nsparse::SeismicIndex index(csr_dimension(csr), batched_params); + if (corpus_residency == "inmem") { + streaming_add(&index, csr); + } else if (corpus_residency == "mmap") { + index.read_csr(csr.c_str(), nsparse::Residency::kMmap); + } else { + std::cerr << "corpus residency must be inmem or mmap\n"; + return 2; + } + + const long load_hwm = read_vm_hwm_kib(); + reset_vm_hwm(); + const double started = now_seconds(); + // batch_file_output_path is set, so build() streams the index out rather + // than retaining it -- the same call an ordinary build makes. + index.build(); + const double build_s = now_seconds() - started; + + report("batched", + "corpus=" + corpus_residency + " batches=" + + std::to_string(batched_params.batch_clustering.batch_size), + build_s, load_hwm); + std::ifstream file(out, std::ios::binary | std::ios::ate); + std::cout << "index_bytes=" << file.tellg() << "\n"; + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 2) { + std::cerr << "usage: batched_build_mem_bench " + " ...\n"; + return 2; + } + const std::string mode = argv[1]; + if (mode == "convert") { + return run_convert(argc, argv); + } + if (mode == "baseline") { + return run_baseline(argc, argv); + } + if (mode == "batched") { + return run_batched(argc, argv); + } + std::cerr << "unknown mode: " << mode << "\n"; + return 2; +} From fdc207b30a95ce2900d44834b8bc3605e14b7fa1 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 06:08:53 +0000 Subject: [PATCH 04/15] docs: document building an index larger than memory Covers what batching is for, that it is a build option in the factory description rather than a separate entry point, what each of the two knobs bounds, and how to measure a build with the peak-RSS driver. The guidance is the measured base_full numbers rather than the mechanism: where the memory/time sweet spot actually is, that query latency and recall do not move, and why the reported peak excludes corpus loading. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 111 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 269bd45..a671f82 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -224,6 +224,117 @@ cmake --build build -j On Linux, the benchmarks support hardware performance counters via [libpfm](http://perfmon2.sourceforge.net/). Install `libpfm4-dev` to enable this. +## Building an index larger than memory + +A whole-corpus build holds two intermediates that scale with the corpus's +non-zeros — the inverted lists (every posting) and then the clustered posting +lists — so peak memory scales with the corpus, and a corpus whose posting lists +do not fit in RAM cannot be indexed at all. + +Batching splits the term space into contiguous windows and finishes one window +before starting the next. It is a build option on the seismic family rather than +a separate entry point, so it is set in the factory description alongside +`lambda` and `beta`, and every type in the family gets it — `seismic`, +`seismic_sq`, `disk_seismic`, `disk_seismic_sq`: + +| Option | Effect | +|---|---| +| `inverted_list_batch_size=N` | Build in `N` term windows. Bounds the inverted-list intermediate to one window; the index is still built in memory as usual. | +| `batch_file_output_path=P` | Additionally serialize each window to `P` and free it, so the clustered lists are never all resident either. The index becomes the file: nothing is retained to search or to `write_index` afterwards. | + +```cpp +auto* index = nsparse::index_factory( + dimension, + "seismic,lambda=6000|beta=400|alpha=0.4" + "|inverted_list_batch_size=10|batch_file_output_path=/data/index.dat"); + +// Corpus residency is SparseVectors' business, not the build's: read_csr can +// map a native-layout CSR instead of copying it, and the build is unchanged. +index->read_csr("corpus.mcsr", nsparse::Residency::kMmap); +index->build(); // streams straight to /data/index.dat + +std::unique_ptr served( + nsparse::read_index("/data/index.dat", nsparse::IndexIoFlag::kUseMmap)); +``` + +The same from Python, since it is only a description string: + +```python +native = nsparse.native_path("corpus.csr") +nsparse.convert("corpus.csr", native) +index = nsparse.index_factory( + dim, + "seismic,lambda=6000|beta=400|alpha=0.4" + "|inverted_list_batch_size=10|batch_file_output_path=/data/index.dat", +) +index.read_csr(native, nsparse.Residency_kMmap) +index.build() +served = nsparse.read_index("/data/index.dat", nsparse.kUseMmap) +``` + +The file is an ordinary index of its type — byte-for-byte what `write_index` +would have produced from the equivalent whole-corpus build. That is asserted +rather than assumed: at a fixed `seed` the two are compared as files, for both a +float and a quantizing index. Each posting list's k-means seed comes from its own +*global* term id and `lambda`/`beta` are resolved once from the whole corpus, so +the window count cannot change what is produced. + +### Choosing `inverted_list_batch_size` + +Peak memory falls roughly as 1/N down to a floor — the corpus itself, per-thread +scratch, allocator retention — so raising it past the point where that floor +dominates buys nothing, and eventually costs time, because every window makes its +own pass over the corpus. On msmarco base_full (8.8M docs, dim 30109, 1.12B +non-zeros, λ=6000 β=400 α=0.4) on a 36-core/68GB host, corpus on the heap in +every row so the only difference is the build: + +| build | peak RSS | build time | +|---|---|---| +| whole corpus | 24085 MB | 106 s | +| 2 windows | 19300 MB | 105 s | +| 10 windows | 10516 MB | 103 s | +| 20 windows | 9284 MB | 118 s | +| 50 windows | 7917 MB | 164 s | +| 100 windows | 7425 MB | 238 s | + +Ten windows is the sweet spot here: **2.3× less peak memory for the same build +time**. By 100 windows the extra passes have more than doubled the build. The +floor is 6.7 GB of corpus; mapping it via `read_csr` takes that off the heap too. + +Peak RSS above is measured from the start of the build, with the corpus already +resident. Loading it costs more than holding it — a streaming ingest stages a +second copy — so a whole-process high-water mark would report the loader (13.0 GB +on this corpus) rather than the build, and hide everything below it. + +Query performance does not move, because the index is the same index. Two +independent unseeded builds, whole-corpus against 10 windows, over 6980 msmarco +dev queries at k=10, read in-memory: + +| index | QPS | p50 | p90 | p99 | recall@10 | +|---|---|---|---|---|---| +| whole corpus | 69332 | 0.290 ms | 0.601 ms | 1.000 ms | 0.8406 | +| 10 windows | 70198 | 0.296 ms | 0.616 ms | 1.055 ms | 0.8432 | + +Write the index to a real disk. On a tmpfs such as `/tmp` it is RAM, which +defeats the point. + +### Measuring it + +`benchmarks/batched_build_mem_bench` reports peak RSS (`VmHWM`) and wall time for +one configuration per process — `google-benchmark` measures throughput, and a +high-water mark is only clean in a process that has built nothing else. Use +`inmem` to compare against `baseline`: both then hold the corpus on the heap, so +the difference is the batching rather than the residency. + +```bash +cmake -S . -B build -DNSPARSE_ENABLE_BENCHMARKS=ON && cmake --build build -j +B=./build/benchmarks/batched_build_mem_bench +$B convert corpus.csr corpus.mcsr +$B baseline corpus.csr 6000 400 0.4 # whole-corpus build, for reference +$B batched inmem corpus.csr 6000 400 0.4 10 /data # 10 windows, same corpus residency +$B batched mmap corpus.mcsr 6000 400 0.4 10 /data # ... or with the corpus mapped +``` + ## Python Bindings ### Build Python Bindings From aa44b6400bc909e4084cc38adba430ad00b1a98a Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 06:34:50 +0000 Subject: [PATCH 05/15] benchmarks: report peak RssAnon, not just peak RSS Peak RSS alone cannot answer the question the batching is for. It counts anonymous pages the process allocated together with file pages it merely touched of a mapping, and only the first are the process's to keep -- the kernel can reclaim the second under pressure. So an RSS-only measurement of a build over a mapped corpus reports memory the build does not really own. On msmarco base_full at 10 windows, the same build over the same corpus: corpus residency peak RSS peak RssAnon peak RssFile build heap 10533 MB 10528 MB 4 MB 106 s mapped 10536 MB 4082 MB 6454 MB 123 s Total RSS is identical, which is why mapping the corpus looked pointless in the earlier numbers. What actually happens is that the 6.45GB corpus moves out of anonymous memory into the page cache, and the memory the process is responsible for falls to 4.1GB. Against the whole-corpus heap build's 24095 MB of anon, that is a 5.9x reduction, for 17 s of extra page faults -- and it is invisible without the split. RssAnon has no kernel peak counter: /proc/self/status reports it as a current value and only VmHWM as a peak, so a sampling thread tracks the maxima of RssAnon and RssFile for the duration of the build. That makes those two lower bounds -- a spike shorter than the 50ms interval is missed, and they can disagree with VmHWM by a hair since VmHWM is itself only updated at certain fault paths. The interval is small next to a build that runs for minutes and allocates in per-window steps, so in practice they track the real peaks. Reported as such rather than as if they were counters. The sampler reads /proc/self/status into a stack buffer rather than through iostreams, because it runs thousands of times while the thing being measured is memory, and a per-sample allocation would show up in its own number. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 45 ++++++--- benchmarks/batched_build_mem_bench.cpp | 134 +++++++++++++++++++++---- 2 files changed, 145 insertions(+), 34 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index a671f82..1215086 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -288,23 +288,42 @@ own pass over the corpus. On msmarco base_full (8.8M docs, dim 30109, 1.12B non-zeros, λ=6000 β=400 α=0.4) on a 36-core/68GB host, corpus on the heap in every row so the only difference is the build: -| build | peak RSS | build time | -|---|---|---| -| whole corpus | 24085 MB | 106 s | -| 2 windows | 19300 MB | 105 s | -| 10 windows | 10516 MB | 103 s | -| 20 windows | 9284 MB | 118 s | -| 50 windows | 7917 MB | 164 s | -| 100 windows | 7425 MB | 238 s | +| build | peak RSS | peak RssAnon | build time | +|---|---|---|---| +| whole corpus | 24100 MB | 24095 MB | 105 s | +| 2 windows | 19300 MB | — | 105 s | +| 10 windows | 10533 MB | 10528 MB | 106 s | +| 20 windows | 9284 MB | — | 118 s | +| 50 windows | 7917 MB | — | 164 s | +| 100 windows | 7417 MB | 7418 MB | 267 s | Ten windows is the sweet spot here: **2.3× less peak memory for the same build -time**. By 100 windows the extra passes have more than doubled the build. The -floor is 6.7 GB of corpus; mapping it via `read_csr` takes that off the heap too. - -Peak RSS above is measured from the start of the build, with the corpus already +time**. By 100 windows the extra passes have more than doubled the build. + +`RssAnon` is the half that matters, being what the process itself allocated; +`RssFile` is pages it touched of a mapping, which the kernel can reclaim under +pressure. With the corpus on the heap they are the same number, because the only +mapping is the binary. Map the corpus instead and they diverge sharply — same +corpus, same 10 windows: + +| corpus residency | peak RSS | peak RssAnon | peak RssFile | build time | +|---|---|---|---|---| +| heap (`kInMemory`) | 10533 MB | 10528 MB | 4 MB | 106 s | +| mapped (`kMmap`) | 10536 MB | **4082 MB** | 6454 MB | 123 s | + +Total RSS is unchanged, so an RSS-only measurement makes mapping look pointless. +What actually happens is that the 6.45 GB corpus moves out of anonymous memory +into the page cache: the memory the process is responsible for falls to 4.1 GB. +Batching and mapping together take that from 24.1 GB to 4.1 GB, a 5.9× reduction, +for 17 s of extra page faults. + +All of the above is measured from the start of the build, with the corpus already resident. Loading it costs more than holding it — a streaming ingest stages a second copy — so a whole-process high-water mark would report the loader (13.0 GB -on this corpus) rather than the build, and hide everything below it. +on this corpus) rather than the build, and hide everything below it. Peak RSS +comes from `VmHWM`, a kernel counter; the anon and file peaks have no such +counter and are sampled, so they are lower bounds and can disagree with `VmHWM` +by a hair. Query performance does not move, because the index is the same index. Two independent unseeded builds, whole-corpus against 10 windows, over 6980 msmarco diff --git a/benchmarks/batched_build_mem_bench.cpp b/benchmarks/batched_build_mem_bench.cpp index 097be0a..d2bcc2e 100644 --- a/benchmarks/batched_build_mem_bench.cpp +++ b/benchmarks/batched_build_mem_bench.cpp @@ -34,6 +34,12 @@ // Point at a real disk: on a tmpfs such as /tmp the index is RAM, and // the numbers are meaningless. +#include +#include + +#include +#include +#include #include #include #include @@ -43,6 +49,7 @@ #include #include #include +#include #include #include "nsparse/io/index_io.h" @@ -52,20 +59,94 @@ namespace { -// Peak resident set size in KiB, read from /proc/self/status (VmHWM). -long read_vm_hwm_kib() { - std::ifstream status("/proc/self/status"); - std::string line; - while (std::getline(status, line)) { - if (line.rfind("VmHWM:", 0) == 0) { - long kib = 0; - std::sscanf(line.c_str(), "VmHWM: %ld kB", &kib); - return kib; - } +// One "Field: N kB" line from /proc/self/status, in KiB, or -1. +// +// Reads into a stack buffer rather than through iostreams because the sampler +// below calls this thousands of times while the thing being measured is memory: +// a per-sample allocation would show up in the number it is reporting. +long read_status_kib(const char* field) { + const int fd = ::open("/proc/self/status", O_RDONLY); // NOLINT + if (fd < 0) { + return -1; } - return -1; + std::array buffer{}; + const ssize_t got = ::read(fd, buffer.data(), buffer.size() - 1); + ::close(fd); + if (got <= 0) { + return -1; + } + buffer[static_cast(got)] = '\0'; + const char* at = std::strstr(buffer.data(), field); + if (at == nullptr) { + return -1; + } + long kib = -1; + std::sscanf(at + std::strlen(field), " %ld", &kib); + return kib; } +long read_vm_hwm_kib() { return read_status_kib("VmHWM:"); } + +// Tracks the high-water mark of RssAnon and RssFile over its own lifetime. +// +// VmHWM is a kernel counter, so peak RSS needs no help. RssAnon and RssFile are +// reported only as current values, and there is no peak equivalent, so the only +// way to get their maxima is to sample. That makes them lower bounds: a spike +// shorter than the interval is missed. The interval is small next to a build +// that runs for minutes and allocates in per-window steps, so in practice they +// track the real peaks, but they are not the guarantee VmHWM is. +// +// The split is worth the trouble because it separates what the build allocates +// (RssAnon -- the inverted lists and clusters, which is what batching bounds) +// from what it merely touches (RssFile -- a mapped corpus, which is the +// kernel's to reclaim under pressure). +class PeakRssSampler { +public: + explicit PeakRssSampler( + std::chrono::milliseconds interval = std::chrono::milliseconds(50)) + : thread_([this, interval] { + while (!stop_.load(std::memory_order_relaxed)) { + sample(); + std::this_thread::sleep_for(interval); + } + // Once more after the stop, so the final state is never missed. + sample(); + }) {} + + ~PeakRssSampler() { + stop_.store(true, std::memory_order_relaxed); + thread_.join(); + } + + PeakRssSampler(const PeakRssSampler&) = delete; + PeakRssSampler& operator=(const PeakRssSampler&) = delete; + + [[nodiscard]] long peak_anon_kib() const { + return peak_anon_.load(std::memory_order_relaxed); + } + [[nodiscard]] long peak_file_kib() const { + return peak_file_.load(std::memory_order_relaxed); + } + +private: + void sample() { + keep_max(&peak_anon_, read_status_kib("RssAnon:")); + keep_max(&peak_file_, read_status_kib("RssFile:")); + } + + static void keep_max(std::atomic* peak, long value) { + long seen = peak->load(std::memory_order_relaxed); + while (value > seen && !peak->compare_exchange_weak( + seen, value, std::memory_order_relaxed)) { + } + } + + std::atomic stop_{false}; + std::atomic peak_anon_{-1}; + std::atomic peak_file_{-1}; + std::thread thread_; +}; + // Resets VmHWM to the current VmRSS, so a later read reports the peak since // this call rather than since the process started. // @@ -170,17 +251,23 @@ int csr_dimension(const std::string& path) { return static_cast(sizes[1]); } -// `peak_rss_mb` is the build's own high-water mark (see reset_vm_hwm); -// `load_peak_rss_mb` is what loading the corpus cost before it, reported so a -// build peak that sits below the loader's is not mistaken for the whole story. +// `peak_rss_mb` is the build's own high-water mark (see reset_vm_hwm), split +// into what the build allocated (`peak_rss_anon_mb`) and what it touched of a +// mapping +// (`peak_rss_file_mb`). The anon figure is the one batching is meant to move; +// the file figure is a mapped corpus, which the kernel can reclaim. +// `load_peak_rss_mb` is what loading the corpus cost before the build, reported +// so a build peak that sits below the loader's is not mistaken for the whole +// story. void report(const std::string& mode, const std::string& detail, double build_s, - long load_hwm_kib) { - const long hwm = read_vm_hwm_kib(); + long load_hwm_kib, const PeakRssSampler& sampler) { + const auto mb = [](long kib) { return static_cast(kib) / 1024.0; }; std::cout << "RESULT mode=" << mode << " " << detail << " build_s=" << build_s - << " peak_rss_mb=" << (static_cast(hwm) / 1024.0) - << " load_peak_rss_mb=" - << (static_cast(load_hwm_kib) / 1024.0) << "\n"; + << " peak_rss_mb=" << mb(read_vm_hwm_kib()) + << " peak_rss_anon_mb=" << mb(sampler.peak_anon_kib()) + << " peak_rss_file_mb=" << mb(sampler.peak_file_kib()) + << " load_peak_rss_mb=" << mb(load_hwm_kib) << "\n"; } int run_convert(int argc, char** argv) { @@ -209,9 +296,13 @@ int run_baseline(int argc, char** argv) { const long load_hwm = read_vm_hwm_kib(); reset_vm_hwm(); + // Scoped to the build, so the sampled peaks exclude corpus loading exactly + // as the reset makes VmHWM exclude it. + PeakRssSampler sampler; const double started = now_seconds(); index.build(); - report("baseline", "batches=0", now_seconds() - started, load_hwm); + const double build_s = now_seconds() - started; + report("baseline", "batches=0", build_s, load_hwm, sampler); if (argc >= 7) { const std::string out = argv[6]; @@ -260,6 +351,7 @@ int run_batched(int argc, char** argv) { const long load_hwm = read_vm_hwm_kib(); reset_vm_hwm(); + PeakRssSampler sampler; const double started = now_seconds(); // batch_file_output_path is set, so build() streams the index out rather // than retaining it -- the same call an ordinary build makes. @@ -269,7 +361,7 @@ int run_batched(int argc, char** argv) { report("batched", "corpus=" + corpus_residency + " batches=" + std::to_string(batched_params.batch_clustering.batch_size), - build_s, load_hwm); + build_s, load_hwm, sampler); std::ifstream file(out, std::ios::binary | std::ios::ate); std::cout << "index_bytes=" << file.tellg() << "\n"; return 0; From 4e61d95954418b54ac62d437284975ee0a45c631 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 07:19:00 +0000 Subject: [PATCH 06/15] Cut term windows to equal clustering load, not equal width Term frequencies are heavily skewed -- on msmarco base_full the heaviest term holds 5.7M postings against a mean of 37K, and the top 1% of terms hold 19% of them -- and peak memory is set by the largest window, not the average one. Equal width therefore wastes most of what batching could save. The counting pass already has the exact per-term counts, so the cut points can be chosen from them. What to even out is min(count, lambda), not count. Pruning keeps at most lambda doc ids per term before clustering, which on base_full is 115M of 1121M postings: 90% of them never reach the phase that dominates the peak, and a term with 5.7M postings costs no more there than one with 6000. Measured imbalance of that load at 10 windows, largest window over the mean: equal width 1.53 balanced by raw counts 3.31 balanced by min(count, l) 1.00 Balancing raw counts is worse than doing nothing, which is not obvious and is why this is measured rather than reasoned: it packs the heavy terms into narrow windows and leaves the others holding thousands of light ones, which both unbalances the phase that matters and starves the per-window parallel loop. Measured that way, build time at 20 windows went from 118s to 128s and at 100 from 267s to 305s, for a 5% memory saving. Weighted correctly, on base_full with the corpus mapped so the figure is the build's own memory rather than the corpus: windows peak RssAnon build time whole 24095 MB 105 s 10 3265 MB 109 s 20 2148 MB 127 s 100 734 MB 286 s Anonymous memory now falls roughly as 1/N, which it did not before: at 10 windows it is 3265 MB against 4082 MB for equal width, a 20% improvement, and at 100 the build allocates under 1GB while indexing 1.12 billion postings. Build time did not regress -- 109s against 123s for equal width at 10 windows. Windows stay contiguous and ascending, which is what the streaming write needs: the layout carries no per-list offsets, so a list's position in the file is its term order, and a window's lists can only be appended once every earlier term is written. Grouping by a hash of the term (term % batches) balances comparably -- 1.03x at 10 windows -- but scatters each window's terms across the file, so it cannot be streamed without buffering completed lists or adding an offset table. It is also worse where the skew bites hardest, 1.64x at 100 windows against 1.07x, because a single term heavier than the target dominates whichever window it lands in. The output is unchanged, as the existing byte-equality tests assert: which window a term lands in cannot affect it. A new test covers the case the weighting creates, a term heavier than a whole window's target, over window counts from 1 to more than the dimension -- every term must still land in exactly one window, or the streamed file would come out short and the write would refuse it. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 66 ++++++++-------- nsparse/seismic_common.cpp | 109 +++++++++++++++++++++------ tests/seismic_batched_build_test.cpp | 41 ++++++++++ 3 files changed, 156 insertions(+), 60 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 1215086..42d0ae9 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -281,49 +281,45 @@ the window count cannot change what is produced. ### Choosing `inverted_list_batch_size` -Peak memory falls roughly as 1/N down to a floor — the corpus itself, per-thread -scratch, allocator retention — so raising it past the point where that floor -dominates buys nothing, and eventually costs time, because every window makes its -own pass over the corpus. On msmarco base_full (8.8M docs, dim 30109, 1.12B -non-zeros, λ=6000 β=400 α=0.4) on a 36-core/68GB host, corpus on the heap in -every row so the only difference is the build: - -| build | peak RSS | peak RssAnon | build time | -|---|---|---|---| -| whole corpus | 24100 MB | 24095 MB | 105 s | -| 2 windows | 19300 MB | — | 105 s | -| 10 windows | 10533 MB | 10528 MB | 106 s | -| 20 windows | 9284 MB | — | 118 s | -| 50 windows | 7917 MB | — | 164 s | -| 100 windows | 7417 MB | 7418 MB | 267 s | - -Ten windows is the sweet spot here: **2.3× less peak memory for the same build -time**. By 100 windows the extra passes have more than doubled the build. - -`RssAnon` is the half that matters, being what the process itself allocated; +Windows are cut to equal *clustering load*, not equal width. Term frequencies are +heavily skewed — on msmarco base_full the heaviest term holds 5.7M postings +against a mean of 37K — and peak memory is set by the largest window, so an uneven +split wastes most of what batching could save. The load that matters is +`min(count, lambda)` rather than `count`, because pruning keeps at most `lambda` +doc ids per term before clustering: 115M of that corpus's 1121M postings, so 90% +of them never reach the phase that dominates the peak. Weighting by that brings +the largest window to 1.00× the mean, against 1.5× for equal width. + +`RssAnon` is the figure to watch, being what the process itself allocated; `RssFile` is pages it touched of a mapping, which the kernel can reclaim under -pressure. With the corpus on the heap they are the same number, because the only -mapping is the binary. Map the corpus instead and they diverge sharply — same -corpus, same 10 windows: +pressure. On base_full (8.8M docs, dim 30109, 1.12B non-zeros, λ=6000 β=400 +α=0.4) on a 36-core/68GB host, with the corpus mapped so the numbers are the +build's own memory rather than the corpus: -| corpus residency | peak RSS | peak RssAnon | peak RssFile | build time | +| windows | peak RssAnon | peak RSS | peak RssFile | build time | |---|---|---|---|---| -| heap (`kInMemory`) | 10533 MB | 10528 MB | 4 MB | 106 s | -| mapped (`kMmap`) | 10536 MB | **4082 MB** | 6454 MB | 123 s | +| whole corpus (on heap) | 24095 MB | 24100 MB | 4 MB | 105 s | +| 10 | 3265 MB | 9724 MB | 6454 MB | 109 s | +| 20 | 2148 MB | 8604 MB | 6454 MB | 127 s | +| 100 | 734 MB | 7186 MB | 6454 MB | 286 s | -Total RSS is unchanged, so an RSS-only measurement makes mapping look pointless. -What actually happens is that the 6.45 GB corpus moves out of anonymous memory -into the page cache: the memory the process is responsible for falls to 4.1 GB. -Batching and mapping together take that from 24.1 GB to 4.1 GB, a 5.9× reduction, -for 17 s of extra page faults. +Anonymous memory falls roughly as 1/N: at 100 windows the build allocates under +1 GB while indexing 1.12 billion postings, against 24 GB unbatched. Total RSS +falls far less because it is dominated by the 6.45 GB of mapped corpus — which is +why the split is worth reporting, an RSS-only measurement making this look like a +3× win rather than a 33× one. + +Build time is flat to around ten windows and then climbs, because every window +makes its own pass over the corpus. Ten to twenty is the useful range: 7–11× less +allocated memory for a few percent of build time. All of the above is measured from the start of the build, with the corpus already resident. Loading it costs more than holding it — a streaming ingest stages a second copy — so a whole-process high-water mark would report the loader (13.0 GB -on this corpus) rather than the build, and hide everything below it. Peak RSS -comes from `VmHWM`, a kernel counter; the anon and file peaks have no such -counter and are sampled, so they are lower bounds and can disagree with `VmHWM` -by a hair. +on this corpus, on the heap path) rather than the build, and hide everything below +it. Peak RSS comes from `VmHWM`, a kernel counter; the anon and file peaks have no +such counter and are sampled, so they are lower bounds and can disagree with +`VmHWM` by a hair. Query performance does not move, because the index is the same index. Two independent unseeded builds, whole-corpus against 10 windows, over 6980 msmarco diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp index 6569e4c..110dc73 100644 --- a/nsparse/seismic_common.cpp +++ b/nsparse/seismic_common.cpp @@ -34,14 +34,72 @@ struct TermWindow { [[nodiscard]] size_t size() const { return end - begin; } }; -// Cuts [0, dimension) into `batches` near-equal windows. -std::vector make_windows(size_t dim, size_t batches) { - // Bounds are size_t, not term_t: dimension may be up to 65536 (term_t is - // uint16), so a term_t window boundary would wrap and silently drop terms. - const size_t per_batch = (dim + batches - 1) / batches; +// Cuts [0, dimension) into at most `batches` windows carrying near-equal +// clustering load, from the exact per-term counts. +// +// Equal width would not do, because term frequencies are heavily skewed: on +// msmarco base_full the heaviest term has 5.7M postings against a mean of 37K, +// and the top 1% of terms hold 19% of them. Peak memory is set by the largest +// window, not the average one, so an uneven split wastes most of what batching +// could save. +// +// The load to even out is min(count, lambda), not count. Pruning keeps at most +// lambda doc ids per term before clustering -- on base_full that is 115M of +// 1121M postings, so 90% of them never reach the phase that dominates the peak, +// and a term with 5.7M postings costs no more there than one with 6000. +// Balancing raw counts instead was measured to make the load that matters +// *worse* than equal width (3.3x the mean against 1.5x at 10 windows): it packs +// the heavy terms into narrow windows and leaves the rest holding thousands of +// light ones, which also starves the per-window parallel loop. Weighted this +// way the same split comes out at 1.00x. +// +// Windows stay contiguous and ascending, which is what lets the clustered lists +// be appended to a file as each window finishes: the layout carries no per-list +// offsets, so their order in the file is their term order. Grouping terms by a +// hash (term % batches) balances well too, but scatters each window's terms +// across the file, which the streaming write cannot express. +// +// Bounds are size_t, not term_t: dimension may be up to 65536 (term_t is +// uint16), so a term_t window boundary would wrap and silently drop terms. +std::vector make_windows(const std::vector& term_counts, + size_t lambda, size_t batches) { + const size_t dim = term_counts.size(); + const auto load_of = [lambda](size_t count) { + return std::min(count, lambda); + }; + size_t remaining_load = 0; + for (size_t count : term_counts) { + remaining_load += load_of(count); + } + std::vector windows; - for (size_t begin = 0; begin < dim; begin += per_batch) { - windows.push_back({begin, std::min(dim, begin + per_batch)}); + size_t begin = 0; + while (begin < dim) { + const size_t windows_left = batches - windows.size(); + if (windows_left <= 1) { + windows.push_back({begin, dim}); + break; + } + // Recomputed per window from what is left, so overshooting one window + // tightens the next instead of accumulating. + const size_t target = + (remaining_load + windows_left - 1) / windows_left; + // Leave one term for each window still to come, so none comes out + // empty. + const size_t max_end = dim - (windows_left - 1); + + size_t end = begin; + size_t load = 0; + // `end == begin` on the first step: a term heavier than the whole + // target still has to go somewhere, and it goes here rather than + // nowhere. + while (end < max_end && (end == begin || load < target)) { + load += load_of(term_counts[end]); + ++end; + } + windows.push_back({begin, end}); + remaining_load -= load; + begin = end; } return windows; } @@ -69,8 +127,8 @@ std::vector count_postings_per_term(const SparseVectors& vectors, return counts; } -// The inverted lists of one term window, sized exactly from the counting pass so -// no list ever grows. +// The inverted lists of one term window, sized exactly from the counting pass +// so no list ever grows. // // Doc ids arrive in ascending order, because fill_from_corpus walks documents // ascending, which is what a single-window build produces. That matters beyond @@ -97,9 +155,9 @@ class WindowLists { const size_t slot = fill_[local_term]++; if (slot >= ids_[local_term].size()) { // The list was sized from the counting pass, so more postings than - // that means the corpus changed under us. One compare on a path that - // is bound by reading the corpus, and it is the difference between - // an exception and a heap overflow. + // that means the corpus changed under us. One compare on a path + // that is bound by reading the corpus, and it is the difference + // between an exception and a heap overflow. throw std::runtime_error( "for_each_clustered_window: corpus changed during the build"); } @@ -153,8 +211,8 @@ void fill_from_corpus(const SparseVectors& vectors, const TermWindow& window, } } -// Lists per OpenMP chunk, at most. Posting lists are wildly uneven in length, so -// they are handed out dynamically rather than split up front. +// Lists per OpenMP chunk, at most. Posting lists are wildly uneven in length, +// so they are handed out dynamically rather than split up front. constexpr size_t kMaxClusterChunk = 64; // Chunks a window should break into, so the threads have something to steal. A @@ -174,10 +232,9 @@ struct ResolvedParameters { uint32_t base_seed; }; -std::vector cluster_window(const SparseVectors& vectors, - ArrayInvertedLists& lists, - const TermWindow& window, - const ResolvedParameters& params) { +std::vector cluster_window( + const SparseVectors& vectors, ArrayInvertedLists& lists, + const TermWindow& window, const ResolvedParameters& params) { std::vector clustered(window.size()); const auto chunk = static_cast(std::clamp( window.size() / kMinClusterChunks, 1, kMaxClusterChunk)); @@ -211,7 +268,8 @@ void for_each_clustered_window(const SparseVectors* vectors, } if (vectors->get_element_size() != config.element_size) { throw std::invalid_argument( - "for_each_clustered_window: corpus element width does not match the " + "for_each_clustered_window: corpus element width does not match " + "the " "index's"); } @@ -225,21 +283,22 @@ void for_each_clustered_window(const SparseVectors* vectors, .beta = calculate_beta(params.beta, lambda), .alpha = params.alpha, // Resolved once, outside the loop: std::random_device usually opens - // /dev/urandom per construction, so drawing per posting list would put a - // syscall on every iteration with every thread doing it. Once for the + // /dev/urandom per construction, so drawing per posting list would put + // a syscall on every iteration with every thread doing it. Once for the // whole build, not per window, or the window count would be observable. .base_seed = params.seed == kRandomSeed ? std::random_device{}() : static_cast(params.seed)}; // Exact per-term sizes, so no window's list ever grows and the bulk - // set_entries path can be used instead of per-posting locking. Also the only - // place a term outside the dimension is caught: the mapped read does not - // range-check. + // set_entries path can be used instead of per-posting locking. Also the + // only place a term outside the dimension is caught: the mapped read does + // not range-check. const std::vector term_counts = count_postings_per_term(*vectors, dim); - for (const TermWindow& window : make_windows(dim, batches)) { + for (const TermWindow& window : make_windows( + term_counts, static_cast(resolved.lambda), batches)) { WindowLists lists(term_counts, window, config.element_size); fill_from_corpus(*vectors, window, config.element_size, &lists); sink(window.begin, diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index 9532b08..27d6c11 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include "nsparse/cluster/inverted_list_clusters.h" @@ -290,6 +291,46 @@ TEST(SeismicBatchedBuild, StreamsAQuantizedIndexIdenticallyToo) { EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n())); } +// Windows are cut to equal posting counts, not equal width, so a term heavier +// than a whole window's target has to be handled: it cannot be split, and it +// must still land in exactly one window with every other term. This corpus puts +// most of the postings on one term, and asks for far more windows than that +// allows. +TEST(SeismicBatchedBuild, HandlesATermHeavierThanAWholeWindow) { + const int dim = 64; + Corpus corpus = make_corpus(/*n_docs=*/400, dim, /*seed=*/97); + // Term 7 in every document, on top of what make_corpus drew: one term with + // an order of magnitude more postings than the rest put together. + Corpus skewed; + skewed.dim = dim; + skewed.indptr.push_back(0); + for (idx_t doc = 0; doc < corpus.n(); ++doc) { + std::vector> row; + for (idx_t j = corpus.indptr[doc]; j < corpus.indptr[doc + 1]; ++j) { + if (corpus.indices[j] != 7) { + row.emplace_back(corpus.indices[j], corpus.values[j]); + } + } + row.emplace_back(static_cast(7), 2.5F); + std::sort(row.begin(), row.end()); // CSR rows must be term-ascending + for (const auto& [term, value] : row) { + skewed.indices.push_back(term); + skewed.values.push_back(value); + } + skewed.indptr.push_back(static_cast(skewed.indices.size())); + } + + TempDir dir("skewed"); + const auto one = streamed(skewed, 1, dir.file("b1.dat")); + ASSERT_FALSE(one.empty()); + // Every one of these has to cover all 64 terms exactly once, or the + // streamed file would be short and the write would refuse it. + EXPECT_EQ(one, streamed(skewed, 8, dir.file("b8.dat"))); + EXPECT_EQ(one, streamed(skewed, 32, dir.file("b32.dat"))); + EXPECT_EQ(one, streamed(skewed, 64, dir.file("b64.dat"))); + EXPECT_EQ(one, streamed(skewed, 200, dir.file("b200.dat"))); +} + // Same invariant without a seed, where the files legitimately differ: the // doc-id membership of each term's list still cannot depend on the window // count, because lambda is computed from the global corpus size. From 70a0a39042b512d69d5085d9fe3f07127f5703da Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 07:43:40 +0000 Subject: [PATCH 07/15] Weight term windows by both build phases, not just one A window has two memory peaks and the split has to weigh both. Filling it holds every posting of its terms at 8 bytes each; clustering it holds what survives pruning -- min(count, lambda) per term, 10% of base_full's postings -- as clusters and summaries, which come to 14.9GB for 115M pruned postings, about 16x bulkier per posting. Weighting either phase alone unbalances the other. min(count, lambda) balances the clustering exactly but concentrates the heavy terms, leaving one window holding 3.7x the mean raw postings, and that window's fill then becomes the peak -- which is what the previous commit's numbers were actually measuring. Raw counts do the reverse. Weighting their sum at the relative cost balances what is resident. Predicted peak of the largest window at 10 windows on base_full: equal width 2.76GB raw count 4.93GB min(count, lambda) 3.32GB both, as here 2.07GB Measured, corpus mapped so the figure is the build's own memory: windows peak RssAnon was (min(count,l)) build time 1 10270 MB -- 112 s 10 2820 MB 3265 MB 107 s 20 1424 MB 2148 MB 130 s 100 458 MB 734 MB 299 s 14% better at 10 windows, 34% at 20, 38% at 100, for at most 4% of build time. Against the unbatched build the anon figure is now 3.6x lower at 10 windows and 22x at 100 -- a build that allocates 458MB while indexing 1.12 billion postings. The model over-predicts (2.07GB against 2.82GB measured), so it is used only to rank the choices, not as a memory estimate. The cluster-to-fill ratio is a constant rather than a function of alpha, beta and dimension, because only its rough magnitude matters: the cost curve is a shallow basin, and assuming 8x or 32x instead of 16x costs about a fifth of the benefit while still beating either phase alone. Deriving it would mean modelling summarize(), which this does not need to be right about. The output is unchanged, which the existing byte-equality tests assert: which window a term lands in cannot affect it. Signed-off-by: Liyun Xiu --- nsparse/seismic_common.cpp | 59 ++++++++++++++++++++++++++------------ 1 file changed, 40 insertions(+), 19 deletions(-) diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp index 110dc73..4e54384 100644 --- a/nsparse/seismic_common.cpp +++ b/nsparse/seismic_common.cpp @@ -34,38 +34,59 @@ struct TermWindow { [[nodiscard]] size_t size() const { return end - begin; } }; -// Cuts [0, dimension) into at most `batches` windows carrying near-equal -// clustering load, from the exact per-term counts. +// How much more a posting costs once clustered than while being scattered into +// an inverted list, per unit. A window's memory has two peaks: filling it holds +// every posting of its terms, and clustering it holds the pruned survivors as +// clusters and summaries, which are far bulkier per posting -- on msmarco +// base_full the finished lists come to 14.9GB for 115M pruned postings, against +// 8 bytes each while filling. +// +// Only the ratio matters, and only roughly: the cost curve is a shallow basin, +// so assuming 8x or 32x here instead of 16x costs about a fifth of the benefit +// and still beats weighting either phase alone. It is deliberately not derived +// from alpha/beta/dimension, which would be a model of summarize() that this +// does not need to be right about. +constexpr size_t kClusterCostRatio = 16; + +// Cuts [0, dimension) into at most `batches` windows of near-equal estimated +// memory, from the exact per-term counts. // // Equal width would not do, because term frequencies are heavily skewed: on -// msmarco base_full the heaviest term has 5.7M postings against a mean of 37K, -// and the top 1% of terms hold 19% of them. Peak memory is set by the largest -// window, not the average one, so an uneven split wastes most of what batching -// could save. +// msmarco base_full the heaviest term holds 5.7M postings against a mean of +// 37K, and the top 1% of terms hold 19% of them. Peak memory is set by the +// largest window, not the average one, so an uneven split wastes most of what +// batching could save. +// +// What to even out is neither phase alone but their sum. Weighting raw counts +// balances the fill and unbalances the clustering, which is the more expensive +// phase, and measured worse than equal width. Weighting min(count, lambda) -- +// what survives pruning, and so what clustering holds -- balances that phase +// perfectly but concentrates the heavy terms, leaving one window holding 3.7x +// the mean raw postings, which then becomes the peak. Weighting both together +// at their relative cost balances what is actually resident. Predicted peak of +// the largest window at 10 windows on base_full: // -// The load to even out is min(count, lambda), not count. Pruning keeps at most -// lambda doc ids per term before clustering -- on base_full that is 115M of -// 1121M postings, so 90% of them never reach the phase that dominates the peak, -// and a term with 5.7M postings costs no more there than one with 6000. -// Balancing raw counts instead was measured to make the load that matters -// *worse* than equal width (3.3x the mean against 1.5x at 10 windows): it packs -// the heavy terms into narrow windows and leaves the rest holding thousands of -// light ones, which also starves the per-window parallel loop. Weighted this -// way the same split comes out at 1.00x. +// equal width 2.76GB +// raw count 4.93GB +// min(count, lambda) 3.32GB +// both, as here 2.07GB // // Windows stay contiguous and ascending, which is what lets the clustered lists // be appended to a file as each window finishes: the layout carries no per-list -// offsets, so their order in the file is their term order. Grouping terms by a -// hash (term % batches) balances well too, but scatters each window's terms -// across the file, which the streaming write cannot express. +// offsets, so a list's position in the file is its term order. Grouping terms +// by a hash (term % batches) balances comparably, but scatters each window's +// terms across the file, which the streaming write cannot express. // // Bounds are size_t, not term_t: dimension may be up to 65536 (term_t is // uint16), so a term_t window boundary would wrap and silently drop terms. std::vector make_windows(const std::vector& term_counts, size_t lambda, size_t batches) { const size_t dim = term_counts.size(); + // Postings held while filling, plus the survivors of pruning weighted by + // what they cost once clustered. Relative, so the element width cancels: it + // scales both phases alike. const auto load_of = [lambda](size_t count) { - return std::min(count, lambda); + return count + kClusterCostRatio * std::min(count, lambda); }; size_t remaining_load = 0; for (size_t count : term_counts) { From 3e1f5dc45f5f974a6bba772ecc71834c144b5e0e Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 08:39:18 +0000 Subject: [PATCH 08/15] docs: correct the batched-build memory table Two errors, both mine, in the table added a few commits ago. It carried the numbers from the previous weighting rather than the one that shipped, because the edit that was meant to update them was in the same command as a commit that a hook rejected, so it never ran. Worse, it mixed units: every row was captioned as having the corpus mapped, but the unbatched row was a heap build at 24095 MB, so the drop from 24 GB to 3 GB read as a batching win when most of it was the corpus not being in anonymous memory. The unbatched row is now the mapped measurement, 10274 MB, which makes the rows comparable and the win 3.6x at ten windows rather than an implied 7x. The heap and mapped figures also do not differ by a fixed offset, which the old text implied: the same 1-window build measures 10274 MB mapped against 24084 MB with a streaming ingest, 13.8 GB more for a 6.86 GB corpus, because the ingest stages a second copy that the allocator retains rather than returning. How much of that overlaps the build's own peak depends on how much the build asks for. Spelled out rather than left to be inferred. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 54 +++++++++++++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 20 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 42d0ae9..83b6c81 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -281,37 +281,51 @@ the window count cannot change what is produced. ### Choosing `inverted_list_batch_size` -Windows are cut to equal *clustering load*, not equal width. Term frequencies are +Windows are cut to equal estimated *memory*, not equal width. Term frequencies are heavily skewed — on msmarco base_full the heaviest term holds 5.7M postings against a mean of 37K — and peak memory is set by the largest window, so an uneven -split wastes most of what batching could save. The load that matters is -`min(count, lambda)` rather than `count`, because pruning keeps at most `lambda` -doc ids per term before clustering: 115M of that corpus's 1121M postings, so 90% -of them never reach the phase that dominates the peak. Weighting by that brings -the largest window to 1.00× the mean, against 1.5× for equal width. +split wastes most of what batching could save. + +A window has two memory peaks and the split has to weigh both. Filling it holds +every posting of its terms; clustering it holds what survives pruning +(`min(count, lambda)` per term, 10% of that corpus's postings) as clusters and +summaries, roughly 16× bulkier per posting. Weighting either phase alone +unbalances the other, both measurably worse than weighting their sum — see +`make_windows` in `nsparse/seismic_common.cpp`, which records what each choice +measured. `RssAnon` is the figure to watch, being what the process itself allocated; `RssFile` is pages it touched of a mapping, which the kernel can reclaim under pressure. On base_full (8.8M docs, dim 30109, 1.12B non-zeros, λ=6000 β=400 -α=0.4) on a 36-core/68GB host, with the corpus mapped so the numbers are the -build's own memory rather than the corpus: +α=0.4) on a 36-core/68GB host, with the corpus **mapped in every row**, so each +figure is the build's own memory: | windows | peak RssAnon | peak RSS | peak RssFile | build time | |---|---|---|---|---| -| whole corpus (on heap) | 24095 MB | 24100 MB | 4 MB | 105 s | -| 10 | 3265 MB | 9724 MB | 6454 MB | 109 s | -| 20 | 2148 MB | 8604 MB | 6454 MB | 127 s | -| 100 | 734 MB | 7186 MB | 6454 MB | 286 s | - -Anonymous memory falls roughly as 1/N: at 100 windows the build allocates under -1 GB while indexing 1.12 billion postings, against 24 GB unbatched. Total RSS -falls far less because it is dominated by the 6.45 GB of mapped corpus — which is -why the split is worth reporting, an RSS-only measurement making this look like a -3× win rather than a 33× one. +| 1 (unbatched) | 10274 MB | 16723 MB | 6454 MB | 113 s | +| 10 | 2820 MB | 9273 MB | 6454 MB | 107 s | +| 20 | 1424 MB | 7878 MB | 6454 MB | 130 s | +| 100 | 458 MB | 6907 MB | 6454 MB | 299 s | + +Anonymous memory falls faster than 1/N — 3.6× at 10 windows and 22× at 100 — +because the split comes from the real per-term costs rather than from term ids. At +100 windows the build allocates 458 MB while indexing 1.12 billion postings. Total +RSS falls far less, being dominated by the mapped corpus, which is why the split is +worth reporting: measured as RSS alone this looks like a 2.4× win rather than a 22× +one. Build time is flat to around ten windows and then climbs, because every window -makes its own pass over the corpus. Ten to twenty is the useful range: 7–11× less -allocated memory for a few percent of build time. +makes its own pass over the corpus. Ten to twenty is the useful range: 3.6–7.2× +less allocated memory for at most a few percent of build time. + +Those rows are comparable to each other but not to a build that loads the corpus +instead of mapping it, and the difference is not just the corpus. The same +1-window build measures 10274 MB mapped against 24084 MB with a streaming ingest — +13.8 GB more for a 6.86 GB corpus — because the ingest stages a second copy of it +(a 13.0 GB load peak) that the allocator retains rather than returning to the OS. +How much of that overlaps the build's own peak depends on how much the build then +asks for, so do not read a heap figure and a mapped figure as differing by a fixed +offset. All of the above is measured from the start of the build, with the corpus already resident. Loading it costs more than holding it — a streaming ingest stages a From 72d16b4a6e310302c70195df6ebe74114ae9f6df Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Wed, 2 Sep 2026 09:53:32 +0000 Subject: [PATCH 09/15] Leave a batched build holding the index it wrote A batched build retained nothing: clustered_inverted_lists stayed empty, so search() found no lists and write_index() would have emitted a header with no postings. The caller had to reopen the file by path. That is a sharp edge for no good reason, and it made the batched path asymmetric with every other build. With batch_size > 1 and an output path, build() now maps the file it just wrote and borrows the posting lists out of it. batch_size <= 1 is untouched: an ordinary build already holds its own lists, so there is nothing to stream and nothing to map back. Only the lists are read back. The forward vectors in the file are a copy of ones the index already has, at whatever residency the caller chose, so re-reading them would be work for nothing -- and it is what lets the corpus mapping be left alone. write_seismic_index_batched returns the byte offset of the list section so the map can seek straight there instead of parsing past the vectors. MmapIndex gains a second MmapFile for that file rather than reusing mapped_file_. The two genuinely coexist: an index that read its corpus with read_csr(kMmap) is still scoring from that mapping, so giving it up for the index's own would leave the index unable to score anything. Two members means no swap, and therefore none of the ordering hazard a swap would carry -- whatever borrows from the new mapping lives in the derived class, and derived members are destroyed before base ones, so the borrowers always go first. Borrowing is what the lists want rather than copying: measured on the 14.9GB base_full index, 0.19s and 8MB of anonymous memory against 11.2s and 13.9GB to copy them. The cursor reads only the size header before each array and skips the bulk, so it faults in about a quarter of the file as page cache, which the kernel can reclaim. Across the batched rows it costs about 4s and leaves peak RssAnon unchanged -- 2816 MB at 10 windows against 2820 MB without it -- while peak RSS rises by the index it touches. peak_rss_mb in the benchmark therefore now includes that; the anon column is the one to read. windows peak RssAnon peak RSS peak RssFile build 1 (unbatched) 10281 MB 16732 MB 6454 MB 105 s 10 2816 MB 12811 MB 9820 MB 111 s 20 1471 MB 11465 MB 9185 MB 130 s 100 465 MB 10396 MB 9725 MB 295 s Tested: a batched build is searchable straight after build() and returns exactly what an unbatched build at the same seed returns, including when the corpus is itself mapped, which is the case the second mapping exists for. Python covers the same through the factory description. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 46 ++++--- benchmarks/batched_build_mem_bench.cpp | 16 ++- nsparse/index_factory.cpp | 22 ++-- nsparse/mmap_index.h | 35 ++++-- nsparse/seismic_batched_build.cpp | 44 +++++-- nsparse/seismic_batched_build.h | 38 +++++- nsparse/seismic_common.h | 12 +- nsparse/seismic_index.cpp | 26 ++-- nsparse/seismic_index.h | 2 +- nsparse/seismic_scalar_quantized_index.cpp | 54 ++++---- python_tests/test_seismic_batched_build.py | 39 +++++- tests/seismic_batched_build_test.cpp | 137 +++++++++++++++------ 12 files changed, 322 insertions(+), 149 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 83b6c81..00b41af 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -240,7 +240,7 @@ a separate entry point, so it is set in the factory description alongside | Option | Effect | |---|---| | `inverted_list_batch_size=N` | Build in `N` term windows. Bounds the inverted-list intermediate to one window; the index is still built in memory as usual. | -| `batch_file_output_path=P` | Additionally serialize each window to `P` and free it, so the clustered lists are never all resident either. The index becomes the file: nothing is retained to search or to `write_index` afterwards. | +| `batch_file_output_path=P` | With `N > 1`, serialize each window to `P` and free it, so the clustered lists are never all resident either, then borrow them back from `P` by mapping it. Unused at `N <= 1`, which is an ordinary build and already holds its own lists. | ```cpp auto* index = nsparse::index_factory( @@ -251,10 +251,11 @@ auto* index = nsparse::index_factory( // Corpus residency is SparseVectors' business, not the build's: read_csr can // map a native-layout CSR instead of copying it, and the build is unchanged. index->read_csr("corpus.mcsr", nsparse::Residency::kMmap); -index->build(); // streams straight to /data/index.dat +index->build(); // streams to /data/index.dat, then maps its lists back in -std::unique_ptr served( - nsparse::read_index("/data/index.dat", nsparse::IndexIoFlag::kUseMmap)); +// Ready to serve, with no reopening by path: the posting lists are borrowed from +// the file just written, and the corpus is still borrowed from its own mapping. +index->search(...); ``` The same from Python, since it is only a description string: @@ -268,8 +269,8 @@ index = nsparse.index_factory( "|inverted_list_batch_size=10|batch_file_output_path=/data/index.dat", ) index.read_csr(native, nsparse.Residency_kMmap) -index.build() -served = nsparse.read_index("/data/index.dat", nsparse.kUseMmap) +index.build() # streams out, then maps its lists back in +dists, labels = index.search(n, indptr, indices, values, k) ``` The file is an ordinary index of its type — byte-for-byte what `write_index` @@ -302,17 +303,30 @@ figure is the build's own memory: | windows | peak RssAnon | peak RSS | peak RssFile | build time | |---|---|---|---|---| -| 1 (unbatched) | 10274 MB | 16723 MB | 6454 MB | 113 s | -| 10 | 2820 MB | 9273 MB | 6454 MB | 107 s | -| 20 | 1424 MB | 7878 MB | 6454 MB | 130 s | -| 100 | 458 MB | 6907 MB | 6454 MB | 299 s | - -Anonymous memory falls faster than 1/N — 3.6× at 10 windows and 22× at 100 — +| 1 (unbatched) | 10281 MB | 16732 MB | 6454 MB | 105 s | +| 10 | 2816 MB | 12811 MB | 9820 MB | 111 s | +| 20 | 1471 MB | 11465 MB | 9185 MB | 130 s | +| 100 | 465 MB | 10396 MB | 9725 MB | 295 s | + +The unbatched row writes no file and maps nothing back, which is why its `RssFile` +is the corpus alone. Anonymous memory falls faster than 1/N — 3.7× at 10 windows +and 22× at 100 — because the split comes from the real per-term costs rather than from term ids. At -100 windows the build allocates 458 MB while indexing 1.12 billion postings. Total -RSS falls far less, being dominated by the mapped corpus, which is why the split is -worth reporting: measured as RSS alone this looks like a 2.4× win rather than a 22× -one. +100 windows the build allocates 465 MB while indexing 1.12 billion postings. Total +RSS barely moves, being dominated by page cache: the 6.45 GB mapped corpus plus +~3.4 GB of the index touched when `build()` maps it back in. That is the whole +reason to report the split — measured as RSS alone this looks like a 1.9× win +rather than a 22× one, and neither figure in that column is memory the process +would have to give up under pressure. + +Mapping the finished lists back in is nearly free in the column that matters: +across those rows it costs about 4 s and leaves `RssAnon` unchanged (2816 MB at 10 +windows against 2820 MB without it). Borrowing a 14.9 GB index takes 0.19 s and +8 MB of anonymous memory, against 11.2 s and 13.9 GB to copy it — the cursor only +reads the size header before each array and skips the bulk, so it faults in about +a quarter of the file as reclaimable page cache. That mapping is separate from the +corpus's: an index that mapped its corpus with `read_csr` keeps doing so, since it +still scores from it. Build time is flat to around ten windows and then climbs, because every window makes its own pass over the corpus. Ten to twenty is the useful range: 3.6–7.2× diff --git a/benchmarks/batched_build_mem_bench.cpp b/benchmarks/batched_build_mem_bench.cpp index d2bcc2e..ceb1c9b 100644 --- a/benchmarks/batched_build_mem_bench.cpp +++ b/benchmarks/batched_build_mem_bench.cpp @@ -260,13 +260,19 @@ int csr_dimension(const std::string& path) { // so a build peak that sits below the loader's is not mistaken for the whole // story. void report(const std::string& mode, const std::string& detail, double build_s, - long load_hwm_kib, const PeakRssSampler& sampler) { + long load_hwm_kib, long start_anon_kib, + const PeakRssSampler& sampler) { const auto mb = [](long kib) { return static_cast(kib) / 1024.0; }; std::cout << "RESULT mode=" << mode << " " << detail << " build_s=" << build_s << " peak_rss_mb=" << mb(read_vm_hwm_kib()) << " peak_rss_anon_mb=" << mb(sampler.peak_anon_kib()) - << " peak_rss_file_mb=" << mb(sampler.peak_file_kib()) + << " peak_rss_file_mb=" + << mb(sampler.peak_file_kib()) + // What the corpus and whatever the loader left behind already + // cost before the build allocated anything. peak_anon minus this + // is the build's own growth. + << " start_rss_anon_mb=" << mb(start_anon_kib) << " load_peak_rss_mb=" << mb(load_hwm_kib) << "\n"; } @@ -295,6 +301,7 @@ int run_baseline(int argc, char** argv) { streaming_add(&index, csr); const long load_hwm = read_vm_hwm_kib(); + const long start_anon = read_status_kib("RssAnon:"); reset_vm_hwm(); // Scoped to the build, so the sampled peaks exclude corpus loading exactly // as the reset makes VmHWM exclude it. @@ -302,7 +309,7 @@ int run_baseline(int argc, char** argv) { const double started = now_seconds(); index.build(); const double build_s = now_seconds() - started; - report("baseline", "batches=0", build_s, load_hwm, sampler); + report("baseline", "batches=0", build_s, load_hwm, start_anon, sampler); if (argc >= 7) { const std::string out = argv[6]; @@ -350,6 +357,7 @@ int run_batched(int argc, char** argv) { } const long load_hwm = read_vm_hwm_kib(); + const long start_anon = read_status_kib("RssAnon:"); reset_vm_hwm(); PeakRssSampler sampler; const double started = now_seconds(); @@ -361,7 +369,7 @@ int run_batched(int argc, char** argv) { report("batched", "corpus=" + corpus_residency + " batches=" + std::to_string(batched_params.batch_clustering.batch_size), - build_s, load_hwm, sampler); + build_s, load_hwm, start_anon, sampler); std::ifstream file(out, std::ios::binary | std::ios::ate); std::cout << "index_bytes=" << file.tellg() << "\n"; return 0; diff --git a/nsparse/index_factory.cpp b/nsparse/index_factory.cpp index de2d514..946ff7c 100644 --- a/nsparse/index_factory.cpp +++ b/nsparse/index_factory.cpp @@ -55,17 +55,17 @@ std::string trim(const std::string& str) { // (key, default) -> value lookup over the parsed description. template SeismicClusterParameters parse_cluster_params(const GetParam& get_param) { - return {.lambda = std::stoi(get_param("lambda", "10")), - .beta = std::stoi(get_param("beta", "5")), - .alpha = std::stof(get_param("alpha", "0.5")), - // Term windows to build in, and where to stream the result if it is - // not to be retained. See BatchClusteringOption. - .batch_clustering = - {.batch_size = static_cast(std::stoul( - get_param("inverted_list_batch_size", "1"))), - .batch_file_output_path = - get_param("batch_file_output_path", "")}, - .seed = std::stoi(get_param("seed", std::to_string(kRandomSeed)))}; + return { + .lambda = std::stoi(get_param("lambda", "10")), + .beta = std::stoi(get_param("beta", "5")), + .alpha = std::stof(get_param("alpha", "0.5")), + // Term windows to build in, and where a batched build streams itself. + // See BatchClusteringOption. + .batch_clustering = {.batch_size = static_cast(std::stoul( + get_param("inverted_list_batch_size", "1"))), + .batch_file_output_path = + get_param("batch_file_output_path", "")}, + .seed = std::stoi(get_param("seed", std::to_string(kRandomSeed)))}; } struct QuantizerConfig { diff --git a/nsparse/mmap_index.h b/nsparse/mmap_index.h index bc0ee06..72d46a3 100644 --- a/nsparse/mmap_index.h +++ b/nsparse/mmap_index.h @@ -10,13 +10,6 @@ #ifndef MMAP_INDEX_H #define MMAP_INDEX_H -#include "nsparse/index.h" -#include "nsparse/seismic_common.h" -#include "nsparse/utils/checks.h" -#include "nsparse/utils/csr_layout.h" -#include "nsparse/utils/mmap_file.h" -#include "nsparse/sparse_vectors.h" - #include #include #include @@ -25,12 +18,20 @@ #include #include +#include "nsparse/index.h" +#include "nsparse/seismic_common.h" +#include "nsparse/sparse_vectors.h" +#include "nsparse/utils/checks.h" +#include "nsparse/utils/csr_layout.h" +#include "nsparse/utils/mmap_file.h" + namespace nsparse { class MmapIndex : public Index { public: explicit MmapIndex(int dim = 0) : Index(dim) {} - void read_csr(const char* file_path, Residency residency = Residency::kInMemory) override { + void read_csr(const char* file_path, + Residency residency = Residency::kInMemory) override { switch (residency) { case Residency::kInMemory: Index::read_csr(file_path); @@ -61,6 +62,17 @@ class MmapIndex : public Index { // mapped_file_ when mapped. get_vectors() cannot tell the two apart. std::unique_ptr vectors_; + // A second mapping, for the file a batched build streams itself to and then + // borrows its posting lists back from. Separate from mapped_file_ rather + // than replacing it, because the two coexist: the corpus may itself be a + // mapping that vectors_ is still borrowing from, and giving that up would + // leave the index unable to score anything. + // + // Whatever borrows from this lives in the derived class, and derived + // members are destroyed before base ones, so the borrowers are always gone + // first. + MmapFile batch_mapped_file_; + private: // Values are borrowed at their stored width, so a quantizing index cannot // use this path. @@ -105,7 +117,8 @@ class MmapIndex : public Index { const size_t indptr_size = static_cast(num_rows) + 1; const auto nnz_size = static_cast(nnz); - if (file.size() != csr_layout::native_file_size(indptr_size, nnz_size)) { + if (file.size() != + csr_layout::native_file_size(indptr_size, nnz_size)) { throw std::invalid_argument( std::string("CSR file is not in the native layout (convert it " "with csr_layout::convert): ") + @@ -133,6 +146,6 @@ class MmapIndex : public Index { vectors_ = std::move(vectors); } }; -} +} // namespace nsparse -#endif // MMAP_INDEX_H +#endif // MMAP_INDEX_H diff --git a/nsparse/seismic_batched_build.cpp b/nsparse/seismic_batched_build.cpp index a2edc40..a025e5d 100644 --- a/nsparse/seismic_batched_build.cpp +++ b/nsparse/seismic_batched_build.cpp @@ -10,6 +10,7 @@ #include "nsparse/seismic_batched_build.h" #include +#include #include #include #include @@ -19,12 +20,15 @@ #include "nsparse/io/file_io.h" #include "nsparse/io/index_io.h" #include "nsparse/io/io.h" +#include "nsparse/io/seismic_invlists_writer.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" +#include "nsparse/utils/mmap_cursor.h" +#include "nsparse/utils/mmap_file.h" namespace nsparse::detail { -void write_seismic_index_batched( +size_t write_seismic_index_batched( const SparseVectors* vectors, const SparseVectorsConfig& config, const SeismicClusterParameters& params, const IndexHeader& header, const std::function& write_prefix, @@ -40,11 +44,11 @@ void write_seismic_index_batched( } // One writer for the whole file, windows serialized straight into it rather - // than spilled and concatenated: serialize() pads each array relative to the - // writer's current offset (see io/align.h), so bytes produced by a writer - // that started at 0 carry the wrong padding once appended at some other - // offset. Streaming through a single writer keeps pos() the true absolute - // offset. + // than spilled and concatenated: serialize() pads each array relative to + // the writer's current offset (see io/align.h), so bytes produced by a + // writer that started at 0 carry the wrong padding once appended at some + // other offset. Streaming through a single writer keeps pos() the true + // absolute offset. FileIOWriter writer(const_cast(out_path.c_str())); write_header(header, &writer); write_prefix(&writer); @@ -52,11 +56,12 @@ void write_seismic_index_batched( // The list count, exactly where SeismicInvertedListsWriter::serialize puts // it. It is the whole dimension, known before any window is built, which is // what lets the lists be streamed after it rather than counted first. + const size_t lists_offset = writer.pos(); size_t n_lists = config.dimension; writer.write(&n_lists, sizeof(size_t), 1); - // Windows arrive in ascending term order, so appending each in turn produces - // the same byte sequence as writing every list at once. + // Windows arrive in ascending term order, so appending each in turn + // produces the same byte sequence as writing every list at once. size_t next_term = 0; for_each_clustered_window( vectors, config, params, @@ -65,7 +70,8 @@ void write_seismic_index_batched( // The layout carries no per-list offsets, so a gap or a repeat // would silently shift every list after it. throw std::runtime_error( - "write_seismic_index_batched: windows arrived out of order"); + "write_seismic_index_batched: windows arrived out of " + "order"); } for (const auto& list : clusters) { list.serialize(&writer); @@ -79,6 +85,26 @@ void write_seismic_index_batched( " of " + std::to_string(config.dimension) + " posting lists"); } writer.close(); + return lists_offset; +} + +std::vector map_streamed_lists(const std::string& path, + size_t lists_offset, + MmapFile* into) { + MmapFile mapped(path); + // The cursor starts at 0 and skips, rather than mapping from lists_offset: + // absolute file offsets are what serialize() padded against, so a cursor + // that began part-way through would compute different padding and misread + // every array. Same reason mmap_index skips rather than offsets. + MmapCursor cursor(mapped.data(), mapped.size()); + cursor.skip(lists_offset); + SeismicInvertedListsWriter lists; + lists.mmap_deserialize(&cursor); + + // Committed only once the walk succeeded, so a truncated file cannot leave + // the index holding lists that point into a mapping it never took. + *into = std::move(mapped); + return std::move(lists.release()); } } // namespace nsparse::detail diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h index 0b909e1..b447c05 100644 --- a/nsparse/seismic_batched_build.h +++ b/nsparse/seismic_batched_build.h @@ -11,19 +11,24 @@ #define SEISMIC_BATCHED_BUILD_H #include +#include #include +#include +#include "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/index.h" #include "nsparse/io/io.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" +#include "nsparse/utils/mmap_file.h" namespace nsparse::detail { // Builds a seismic-family index and writes it straight to `out_path`, one term // window at a time, without ever holding the whole index in memory. // -// The usual build holds two whole-corpus intermediates -- the inverted lists and -// then the clustered posting lists -- so its peak memory scales with the +// The usual build holds two whole-corpus intermediates -- the inverted lists +// and then the clustered posting lists -- so its peak memory scales with the // corpus's non-zeros, and a corpus whose posting lists do not fit in RAM cannot // be indexed at all. for_each_clustered_window bounds the first to one window; // serializing each window and dropping it, which is what this does, bounds the @@ -31,9 +36,9 @@ namespace nsparse::detail { // holds, at whatever residency SparseVectors was given) plus one window. // // Reached through an index's build(), by setting -// SeismicClusterParameters::batch_clustering.batch_file_output_path. The index is -// then the file, not the object: nothing is retained to serve or to write_index -// afterwards. +// SeismicClusterParameters::batch_clustering.batch_file_output_path. The index +// is then the file, not the object: nothing is retained to serve or to +// write_index afterwards. // // `header` and `write_prefix` are what make this work for every type in the // family rather than just SEIS. `write_prefix` writes whatever the type puts @@ -47,14 +52,35 @@ namespace nsparse::detail { // whatever batch_size is, because every list's k-means seed comes from its own // global term id -- see for_each_clustered_window. // +// Returns the absolute byte offset of the posting-list section, so the lists +// can be mapped back in without re-parsing everything before them -- see +// map_streamed_lists. +// // Throws if the corpus is empty: there would be no windows to stream, and a // header-only file is not a readable index. -void write_seismic_index_batched( +size_t write_seismic_index_batched( const SparseVectors* vectors, const SparseVectorsConfig& config, const SeismicClusterParameters& params, const IndexHeader& header, const std::function& write_prefix, const std::string& out_path); +// Maps the file a streamed build just wrote and borrows its posting lists out +// of it, so the build ends holding a usable index without ever having held all +// of the lists at once. +// +// Only the lists. The forward vectors in the file are a copy of ones the index +// already has, at whatever residency the caller chose for them, so re-reading +// them would be work for nothing -- and it is why the corpus mapping can be +// left alone rather than swapped out. `lists_offset` is what +// write_seismic_index_batched returned, which saves parsing past the vectors to +// find where the lists start. +// +// The mapping is handed to `into`, which must outlive the returned lists: they +// point into it. +std::vector map_streamed_lists(const std::string& path, + size_t lists_offset, + MmapFile* into); + } // namespace nsparse::detail #endif // SEISMIC_BATCHED_BUILD_H diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index c89e3a0..4d5401b 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -36,14 +36,18 @@ namespace nsparse { // contiguous windows and finishing one window before starting the next makes the // first of those proportional to a window instead. // -// `batch_file_output_path` bounds the second as well: with it set, each window's -// clustered lists are serialized to that path and freed as they are produced, so -// the build retains nothing and the index is the file rather than the object. -// See build_seismic_index_batched. +// `batch_file_output_path` bounds the second as well: with batch_size > 1, each +// window's clustered lists are serialized to that path and freed as they are +// produced, and the finished list section is then mapped back in, so the build +// never holds more than one window's worth and still ends with a usable index. +// See write_seismic_index_batched. struct BatchClusteringOption { // Contiguous term windows. <= 1 means one window, i.e. no batching. Clamped // to the dimension, since a window cannot be narrower than one term. size_t batch_size = 1; + // Where a batched build streams the index it produces. Used only when + // batch_size > 1: a single window is an ordinary build, which holds its own + // posting lists and so has nothing to stream or to map back. std::string batch_file_output_path; }; diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index 159588e..39d8109 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -9,8 +9,6 @@ #include "nsparse/seismic_index.h" -#include "nsparse/seismic_batched_build.h" - #include #include #include @@ -26,6 +24,7 @@ #include "nsparse/index.h" #include "nsparse/invlists/inverted_lists.h" #include "nsparse/io/seismic_invlists_writer.h" +#include "nsparse/seismic_batched_build.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" @@ -147,24 +146,25 @@ void SeismicIndex::build() { const SparseVectorsConfig config = { .element_size = kElementSize, .dimension = static_cast(get_dimension())}; - const std::string& out_path = - cluster_parameter_.batch_clustering.batch_file_output_path; - if (!out_path.empty()) { - // Streamed straight to a file and not retained: see - // BatchClusteringOption. write_index afterwards would write an index with - // no posting lists, so this index is deliberately left empty. - detail::write_seismic_index_batched( + const auto& batch = cluster_parameter_.batch_clustering; + if (batch.batch_size > 1 && !batch.batch_file_output_path.empty()) { + // Streamed to a file a window at a time, so the whole index is never + // resident, then borrowed back so this is a usable index. A single + // window takes the ordinary path below: it already holds its own lists, + // and writing them out only to map them back would be work for nothing. + const size_t lists_offset = detail::write_seismic_index_batched( get_vectors(), config, cluster_parameter_, {.id = fourcc(name), .version = kFormatVersion, .dimension = get_dimension()}, [this](IOWriter* io_writer) { vectors_->serialize(io_writer); }, - out_path); + batch.batch_file_output_path); + clustered_inverted_lists = detail::map_streamed_lists( + batch.batch_file_output_path, lists_offset, &batch_mapped_file_); return; } - clustered_inverted_lists = - detail::build_inverted_lists_clusters(get_vectors(), config, - cluster_parameter_); + clustered_inverted_lists = detail::build_inverted_lists_clusters( + get_vectors(), config, cluster_parameter_); } auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, diff --git a/nsparse/seismic_index.h b/nsparse/seismic_index.h index d1d0a52..df3e804 100644 --- a/nsparse/seismic_index.h +++ b/nsparse/seismic_index.h @@ -81,7 +81,7 @@ class SeismicIndex : public MmapIndex, public IndexIO { size_t q_len, const std::vector& cuts, int k, float heap_factor, SearchParameters* search_parameters) -> pair_of_score_id_vector_t; - + SeismicClusterParameters cluster_parameter_; }; } // namespace nsparse diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index 9319e62..213b656 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -9,8 +9,6 @@ #include "nsparse/seismic_scalar_quantized_index.h" -#include "nsparse/seismic_batched_build.h" - #include #include @@ -29,6 +27,7 @@ #include "nsparse/invlists/inverted_lists.h" #include "nsparse/io/io.h" #include "nsparse/io/seismic_invlists_writer.h" +#include "nsparse/seismic_batched_build.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" @@ -196,13 +195,12 @@ void SeismicScalarQuantizedIndex::build() { const SparseVectorsConfig config = { .element_size = sq_.bytes_per_value(), .dimension = static_cast(get_dimension())}; - const std::string& out_path = - cluster_parameter_.batch_clustering.batch_file_output_path; - if (!out_path.empty()) { - // The quantization header comes first, exactly as write_index writes it; - // the codes in `vectors_` are already quantized, so the batched build - // needs no knowledge of the quantizer beyond its width. - detail::write_seismic_index_batched( + const auto& batch = cluster_parameter_.batch_clustering; + if (batch.batch_size > 1 && !batch.batch_file_output_path.empty()) { + // The quantization header comes first, exactly as write_index writes + // it; the codes in `vectors_` are already quantized, so the batched + // build needs no knowledge of the quantizer beyond its width. + const size_t lists_offset = detail::write_seismic_index_batched( get_vectors(), config, cluster_parameter_, {.id = fourcc(name), .version = kFormatVersion, @@ -211,12 +209,13 @@ void SeismicScalarQuantizedIndex::build() { write_quantization_header(io_writer); vectors_->serialize(io_writer); }, - out_path); + batch.batch_file_output_path); + clustered_inverted_lists = detail::map_streamed_lists( + batch.batch_file_output_path, lists_offset, &batch_mapped_file_); return; } - clustered_inverted_lists = - detail::build_inverted_lists_clusters(get_vectors(), config, - cluster_parameter_); + clustered_inverted_lists = detail::build_inverted_lists_clusters( + get_vectors(), config, cluster_parameter_); } auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, @@ -283,14 +282,14 @@ auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, dynamic_cast(search_parameters); const auto* parameters = seismic_parameters != nullptr ? seismic_parameters : &default_params; - const size_t dense_bytes = - static_cast(dimension_) * element_size; + 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. + // 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); @@ -306,9 +305,8 @@ auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, std::vector cuts; if (element_size == U16) { cuts = detail::top_k_tokens( - q_indices, - reinterpret_cast(q_val_bytes), len, - parameters->cut); + q_indices, reinterpret_cast(q_val_bytes), + len, parameters->cut); } else { cuts = detail::top_k_tokens(q_indices, q_val_bytes, len, parameters->cut); @@ -336,11 +334,12 @@ auto SeismicScalarQuantizedIndex::single_query( return {{}, {}}; } - // Scatter the query's quantized codes into the reused dense buffer (all-zero - // on entry): element_size contiguous bytes per non-zero dim. + // 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); + std::copy_n( + q_val_bytes + i * element_size, element_size, + dense.data() + static_cast(q_idx[i]) * element_size); } visited.clear(); @@ -452,7 +451,8 @@ void SeismicScalarQuantizedIndex::write_quantization_header( io_writer->write(&vmax, sizeof(float), 1); } -void SeismicScalarQuantizedIndex::read_quantization_header(IOReader* io_reader) { +void SeismicScalarQuantizedIndex::read_quantization_header( + IOReader* io_reader) { QuantizerType sq_type = QuantizerType::QT_8bit; float vmin = 0.0F; float vmax = 1.0F; diff --git a/python_tests/test_seismic_batched_build.py b/python_tests/test_seismic_batched_build.py index 1bbcd78..8120846 100644 --- a/python_tests/test_seismic_batched_build.py +++ b/python_tests/test_seismic_batched_build.py @@ -10,8 +10,9 @@ Batching is a build option rather than a separate entry point, so there is nothing new to wrap: it is reached through the factory description, the same way lambda and beta are. `inverted_list_batch_size` bounds the build's memory; -adding `batch_file_output_path` streams the index straight to that file instead -of retaining it, which is the path a corpus too large for RAM needs. +with more than one window, `batch_file_output_path` is where the index is streamed +as it is built, and its posting lists are then borrowed back from that file -- +which is the path a corpus too large for RAM needs. """ import numpy as np @@ -43,7 +44,8 @@ def streamed(corpus, out_path, batch_size, kind="seismic"): return str(out_path) -@pytest.mark.parametrize("batch_size", [1, 4, 32]) +# Batching starts at 2: one window is an ordinary build and writes no file. +@pytest.mark.parametrize("batch_size", [2, 4, 32]) def test_happy_case(batch_size, corpus, queries, oracle, tmp_path): """build -> read back mapped -> query -> accuracy, at several splits.""" path = streamed(corpus, tmp_path / "batched.idx", batch_size) @@ -74,6 +76,26 @@ def test_matches_in_memory_build(kind, corpus, tmp_path): assert in_memory.read_bytes() == open(batched, "rb").read() +def test_streamed_index_is_searchable_after_build(corpus, queries, oracle, tmp_path): + """build() leaves a usable index, not an empty object. + + The lists are borrowed back from the file it just wrote, so there is no + reopening by path and they are never copied onto the heap. + """ + spec = ( + f"seismic,{BASE}|inverted_list_batch_size=8" + f"|batch_file_output_path={tmp_path / 'streamed.idx'}" + ) + index = nsparse.index_factory(corpus.dim, spec) + add_corpus(index, corpus) + index.build() + + assert index.num_vectors() == corpus.n + _, labels = search(index, queries) + want_labels, _ = oracle + assert recall_at_k(labels, want_labels) >= RECALL_FLOOR + + def test_batch_size_alone_leaves_the_index_in_memory(corpus, queries, tmp_path): """Without an output path, batching only bounds the build's intermediates. @@ -91,11 +113,16 @@ def test_batch_size_alone_leaves_the_index_in_memory(corpus, queries, tmp_path): def test_batch_count_is_not_observable(corpus, queries, tmp_path): - """The split is a memory knob: at a fixed seed it cannot change the results.""" - one = streamed(corpus, tmp_path / "one.idx", 1) + """The split is a memory knob: at a fixed seed it cannot change the results. + + Against an unbatched build, which writes its file the ordinary way since one + window streams nothing. + """ + plain = tmp_path / "plain.idx" + nsparse.write_index(make_index(f"seismic,{BASE}", corpus), str(plain)) many = streamed(corpus, tmp_path / "many.idx", 16) - want_d, want_l = search(nsparse.read_index(one), queries) + want_d, want_l = search(nsparse.read_index(str(plain)), queries) got_d, got_l = search(nsparse.read_index(many), queries) np.testing.assert_array_equal(got_l, want_l) np.testing.assert_allclose(got_d, want_d, rtol=1e-6, atol=1e-6) diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index 27d6c11..c9492dd 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -23,6 +23,7 @@ #include #include +#include "csr_interchange_test_util.h" #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/disk_seismic_index.h" #include "nsparse/index_factory.h" @@ -43,14 +44,14 @@ constexpr int kLambda = 64; constexpr int kBeta = 6; constexpr float kAlpha = 0.4F; +// `n` is a field rather than a method so csr_test::write_interchange_csr, which +// is templated on any corpus exposing .n/.indptr/.indices/.values, accepts it. struct Corpus { int dim; + idx_t n = 0; std::vector indptr; std::vector indices; std::vector values; - [[nodiscard]] idx_t n() const { - return static_cast(indptr.size()) - 1; - } }; Corpus make_corpus(int n_docs, int dim, unsigned seed) { @@ -61,6 +62,7 @@ Corpus make_corpus(int n_docs, int dim, unsigned seed) { Corpus corpus; corpus.dim = dim; + corpus.n = n_docs; corpus.indptr.push_back(0); for (int doc = 0; doc < n_docs; ++doc) { // Capped at dim: the loop below draws *distinct* terms, so asking for @@ -125,7 +127,7 @@ std::vector read_file(const std::string& path) { std::vector streamed(const Corpus& corpus, size_t batch_size, const std::string& out, int seed = kSeed) { SeismicIndex index(corpus.dim, params_for(batch_size, out, seed)); - index.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); index.build(); return read_file(out); @@ -135,7 +137,7 @@ std::vector streamed(const Corpus& corpus, size_t batch_size, std::vector in_memory(const Corpus& corpus, const std::string& out, size_t batch_size = 1, int seed = kSeed) { SeismicIndex index(corpus.dim, params_for(batch_size, "", seed)); - index.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); index.build(); write_index(&index, const_cast(out.c_str())); @@ -173,24 +175,9 @@ std::vector> per_term_doc_sets(const std::string& path) { return out; } -// The corpus in the interchange CSR layout, converted to native: what a mapped -// read consumes. +// The corpus as a native-layout CSR: what a mapped read consumes. std::string write_native_csr(const Corpus& corpus, const std::string& path) { - std::ofstream out(path, std::ios::binary); - const std::array sizes = { - corpus.n(), corpus.dim, static_cast(corpus.indices.size())}; - out.write(reinterpret_cast(sizes.data()), sizeof(sizes)); - std::vector indptr64(corpus.indptr.begin(), corpus.indptr.end()); - out.write(reinterpret_cast(indptr64.data()), - static_cast(indptr64.size() * sizeof(int64_t))); - std::vector indices32(corpus.indices.begin(), - corpus.indices.end()); - out.write(reinterpret_cast(indices32.data()), - static_cast(indices32.size() * sizeof(int32_t))); - out.write( - reinterpret_cast(corpus.values.data()), - static_cast(corpus.values.size() * sizeof(float))); - out.close(); + csr_test::write_interchange_csr(path, corpus, corpus.dim); const std::string native = csr_layout::native_path(path); csr_layout::convert(path, native); return native; @@ -214,14 +201,13 @@ TEST(SeismicBatchedBuild, StreamedBuildIsByteIdenticalToInMemoryBuild) { TEST(SeismicBatchedBuild, StreamedBuildIsIdenticalAcrossBatchCounts) { Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/11); TempDir dir("counts"); - const auto one = streamed(corpus, 1, dir.file("b1.dat")); + // batch_size <= 1 is an ordinary build, so its file comes from write_index. + const auto one = in_memory(corpus, dir.file("b1.dat")); ASSERT_FALSE(one.empty()); EXPECT_EQ(one, streamed(corpus, 2, dir.file("b2.dat"))); EXPECT_EQ(one, streamed(corpus, 10, dir.file("b10.dat"))); - // More windows than terms is clamped to one term each, and 0 means one - // window rather than none. + // More windows than terms is clamped to one term each. EXPECT_EQ(one, streamed(corpus, 1000, dir.file("b1000.dat"))); - EXPECT_EQ(one, streamed(corpus, 0, dir.file("b0.dat"))); } // batch_size alone bounds the inverted-list intermediate and leaves the index @@ -247,7 +233,7 @@ TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeADiskIndex) { auto build_disk = [&corpus](size_t batch_size, const std::string& out) { DiskSeismicIndex index(corpus.dim, params_for(batch_size, "", kSeed)); - index.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); index.build(); write_index(&index, const_cast(out.c_str())); @@ -271,7 +257,7 @@ TEST(SeismicBatchedBuild, StreamsAQuantizedIndexIdenticallyToo) { SeismicScalarQuantizedIndex mem(QuantizerType::QT_8bit, 0.0F, 3.0F, params_for(1, "", kSeed), corpus.dim); - mem.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); mem.build(); write_index(&mem, const_cast(mem_path.c_str())); @@ -279,7 +265,7 @@ TEST(SeismicBatchedBuild, StreamsAQuantizedIndexIdenticallyToo) { SeismicScalarQuantizedIndex batched(QuantizerType::QT_8bit, 0.0F, 3.0F, params_for(4, streamed_path, kSeed), corpus.dim); - batched.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); batched.build(); @@ -288,7 +274,7 @@ TEST(SeismicBatchedBuild, StreamsAQuantizedIndexIdenticallyToo) { std::unique_ptr reloaded( read_index(const_cast(streamed_path.c_str()))); EXPECT_EQ(reloaded->id(), SeismicScalarQuantizedIndex::name); - EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n())); + EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n)); } // Windows are cut to equal posting counts, not equal width, so a term heavier @@ -303,8 +289,9 @@ TEST(SeismicBatchedBuild, HandlesATermHeavierThanAWholeWindow) { // an order of magnitude more postings than the rest put together. Corpus skewed; skewed.dim = dim; + skewed.n = corpus.n; skewed.indptr.push_back(0); - for (idx_t doc = 0; doc < corpus.n(); ++doc) { + for (idx_t doc = 0; doc < corpus.n; ++doc) { std::vector> row; for (idx_t j = corpus.indptr[doc]; j < corpus.indptr[doc + 1]; ++j) { if (corpus.indices[j] != 7) { @@ -321,7 +308,7 @@ TEST(SeismicBatchedBuild, HandlesATermHeavierThanAWholeWindow) { } TempDir dir("skewed"); - const auto one = streamed(skewed, 1, dir.file("b1.dat")); + const auto one = in_memory(skewed, dir.file("b1.dat")); ASSERT_FALSE(one.empty()); // Every one of these has to cover all 64 terms exactly once, or the // streamed file would be short and the write would refuse it. @@ -331,6 +318,73 @@ TEST(SeismicBatchedBuild, HandlesATermHeavierThanAWholeWindow) { EXPECT_EQ(one, streamed(skewed, 200, dir.file("b200.dat"))); } +// A batched build ends holding the index it wrote, so build() leaves something +// usable rather than an empty object. Against an unbatched build at the same +// seed: identical builds, so identical results. +TEST(SeismicBatchedBuild, BatchedBuildIsSearchableAfterBuild) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/7); + Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); + const int k = 10; + const auto n = static_cast(queries.n); + TempDir dir("searchable"); + + const auto search_with = [&](Index& index) { + std::vector dist(n * k); + std::vector lab(n * k); + SeismicSearchParameters params(/*cut=*/3, /*heap_factor=*/1.0F); + index.search(queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, dist.data(), lab.data(), + ¶ms); + return std::pair{dist, lab}; + }; + + SeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + const auto [want_dist, want_lab] = search_with(mem); + + SeismicIndex batched(corpus.dim, params_for(4, dir.file("b.dat"), kSeed)); + batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + batched.build(); + + // No reopening by path: build() mapped the lists back in. + EXPECT_EQ(batched.num_vectors(), static_cast(corpus.n)); + const auto [got_dist, got_lab] = search_with(batched); + EXPECT_EQ(got_lab, want_lab); + EXPECT_EQ(got_dist, want_dist); +} + +// A corpus borrowed from a mapping is the case two mappings exist for: the +// corpus one, which vectors_ still scores from, and the one the build just +// wrote, which the posting lists borrow from. Neither may be given up for the +// other. +TEST(SeismicBatchedBuild, KeepsTheCorpusMappingWhileBorrowingItsOwnLists) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/13); + Corpus queries = make_corpus(/*n_docs=*/40, /*dim=*/200, /*seed=*/77); + const int k = 10; + const auto n = static_cast(queries.n); + TempDir dir("mapped_reload"); + + const std::string native = write_native_csr(corpus, dir.file("corpus.csr")); + SeismicIndex index(corpus.dim, params_for(3, dir.file("out.dat"), kSeed)); + index.read_csr(native.c_str(), Residency::kMmap); + index.build(); + + // Still serving: scoring reads the mapped corpus, the lists come from the + // second mapping. + EXPECT_EQ(index.num_vectors(), static_cast(corpus.n)); + std::vector dist(n * k); + std::vector lab(n * k); + SeismicSearchParameters params(/*cut=*/3, /*heap_factor=*/1.0F); + static_cast(index).search( + queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, dist.data(), lab.data(), ¶ms); + EXPECT_TRUE( + std::any_of(lab.begin(), lab.end(), [](idx_t id) { return id >= 0; })); +} + // Same invariant without a seed, where the files legitimately differ: the // doc-id membership of each term's list still cannot depend on the window // count, because lambda is computed from the global corpus size. @@ -339,7 +393,7 @@ TEST(SeismicBatchedBuild, PerTermMembershipInvariantAcrossBatches) { TempDir dir("membership"); const std::string one = dir.file("1.dat"); const std::string ten = dir.file("10.dat"); - streamed(corpus, 1, one, kRandomSeed); + in_memory(corpus, one, /*batch_size=*/1, kRandomSeed); streamed(corpus, 10, ten, kRandomSeed); auto sets1 = per_term_doc_sets(one); @@ -360,11 +414,12 @@ TEST(SeismicBatchedBuild, HandlesDimensionAt65536) { Corpus corpus = make_corpus(/*n_docs=*/3000, dim, /*seed=*/5); TempDir dir("dim64k"); const std::string path = dir.file("index.dat"); - streamed(corpus, 1, path); + // Two windows, so the batched path is what produces the file. + streamed(corpus, 2, path); // Must load without "unexpected end of index file". std::unique_ptr idx(read_index(const_cast(path.c_str()))); - EXPECT_EQ(idx->num_vectors(), static_cast(corpus.n())); + EXPECT_EQ(idx->num_vectors(), static_cast(corpus.n)); auto sets = per_term_doc_sets(path); EXPECT_EQ(sets.size(), static_cast(dim)); @@ -385,7 +440,7 @@ TEST(SeismicBatchedBuild, SearchThroughMappedReadMatchesInMemory) { const std::string streamed_path = dir.file("streamed.dat"); SeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); - mem.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); mem.build(); streamed(corpus, 4, streamed_path); @@ -394,16 +449,16 @@ TEST(SeismicBatchedBuild, SearchThroughMappedReadMatchesInMemory) { const_cast(streamed_path.c_str()), IndexIoFlag::kUseMmap)); SeismicSearchParameters search_params(/*cut=*/3, /*heap_factor=*/1.0F); - const auto n = static_cast(queries.n()); + const auto n = static_cast(queries.n); std::vector mem_dist(n * k); std::vector mem_lab(n * k); std::vector disk_dist(n * k); std::vector disk_lab(n * k); - static_cast(mem).search(queries.n(), queries.indptr.data(), + static_cast(mem).search(queries.n, queries.indptr.data(), queries.indices.data(), queries.values.data(), k, mem_dist.data(), mem_lab.data(), &search_params); - disk->search(queries.n(), queries.indptr.data(), queries.indices.data(), + disk->search(queries.n, queries.indptr.data(), queries.indices.data(), queries.values.data(), k, disk_dist.data(), disk_lab.data(), &search_params); @@ -441,7 +496,7 @@ TEST(SeismicBatchedBuild, FactoryDescriptionDrivesTheBatchedBuild) { path; std::unique_ptr index(index_factory(corpus.dim, spec.c_str())); - index->add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + index->add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); index->build(); @@ -458,7 +513,7 @@ TEST(SeismicBatchedBuild, RejectsInvalidInput) { // term would be silently dropped from the index. SeismicIndex narrow(corpus.dim / 2, params_for(1, dir.file("narrow.dat"), kSeed)); - narrow.add(corpus.n(), corpus.indptr.data(), corpus.indices.data(), + narrow.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); EXPECT_THROW(narrow.build(), std::invalid_argument); From 5af526f6f10ffb89ce5e1b88dd6db16c6b136ae0 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 3 Sep 2026 04:23:40 +0000 Subject: [PATCH 10/15] Bound the disk indexes' build memory too, by spilling their lists The in-memory types end their payload with their posting lists, so a batched build serializes each window straight into the output and drops it. A DiskSeismic payload cannot be written that way: its summaries are followed by an inline forward index whose blocks are laid out from the doc-id membership of every list, which is not known until the last window is clustered, so batch_file_output_path was silently ignored there. What can be dropped is the lists' residency rather than the lists. The clustering still runs once, a window at a time, into a spill next to the output; the lists come back borrowed from that mapping, and the payload is written from it, the forward index streaming its blocks as it lays them. Neither phase holds more than one window of anonymous memory, and the spill goes with the build. On base_full (1.12B non-zeros) at lambda=600, corpus mapped, peak RssAnon falls 8671 MB -> 1440 MB at 10 windows and 434 MB at 100. The benchmark driver takes an index type now, so that is measurable rather than assumed. The output is byte-identical to build() + write_index at a fixed seed, for both disk types, and the index ends borrowing its own output rather than the scratch it deleted -- so build() still leaves something that serves. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 47 ++++- benchmarks/batched_build_mem_bench.cpp | 98 +++++++---- nsparse/disk_seismic_index.cpp | 5 +- nsparse/disk_seismic_index_base.cpp | 110 ++++++++++-- nsparse/disk_seismic_index_base.h | 60 +++++-- .../disk_seismic_scalar_quantized_index.cpp | 14 +- nsparse/disk_seismic_scalar_quantized_index.h | 2 + nsparse/seismic_batched_build.cpp | 88 +++++++--- nsparse/seismic_batched_build.h | 36 +++- python_tests/test_seismic_batched_build.py | 37 +++- tests/seismic_batched_build_test.cpp | 164 +++++++++++++++++- 11 files changed, 555 insertions(+), 106 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 00b41af..732ce24 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -240,7 +240,16 @@ a separate entry point, so it is set in the factory description alongside | Option | Effect | |---|---| | `inverted_list_batch_size=N` | Build in `N` term windows. Bounds the inverted-list intermediate to one window; the index is still built in memory as usual. | -| `batch_file_output_path=P` | With `N > 1`, serialize each window to `P` and free it, so the clustered lists are never all resident either, then borrow them back from `P` by mapping it. Unused at `N <= 1`, which is an ordinary build and already holds its own lists. | +| `batch_file_output_path=P` | With `N > 1`, write the index to `P` as it is built rather than assembling it in memory, so the clustered lists are never all resident either, then borrow them back from `P` by mapping it. Unused at `N <= 1`, which is an ordinary build and already holds its own lists. | + +`seismic` and `seismic_sq` end their payload with their posting lists, so each +window is serialized straight into `P` and dropped. The disk-resident pair cannot +be written that way: their summaries are followed by an inline forward index whose +blocks are laid out from the doc-id membership of *every* list, which is not known +until the last window is clustered. They spill the clustered lists to `P.lists` +instead, map them back, and write the payload from that mapping — same bound on +anonymous memory, at the cost of scratch disk the size of the lists. The spill is +deleted with the build; nothing else reads it. ```cpp auto* index = nsparse::index_factory( @@ -275,10 +284,10 @@ dists, labels = index.search(n, indptr, indices, values, k) The file is an ordinary index of its type — byte-for-byte what `write_index` would have produced from the equivalent whole-corpus build. That is asserted -rather than assumed: at a fixed `seed` the two are compared as files, for both a -float and a quantizing index. Each posting list's k-means seed comes from its own -*global* term id and `lambda`/`beta` are resolved once from the whole corpus, so -the window count cannot change what is produced. +rather than assumed: at a fixed `seed` the two are compared as files, for all four +types. Each posting list's k-means seed comes from its own *global* term id and +`lambda`/`beta` are resolved once from the whole corpus, so the window count +cannot change what is produced. ### Choosing `inverted_list_batch_size` @@ -332,6 +341,29 @@ Build time is flat to around ten windows and then climbs, because every window makes its own pass over the corpus. Ten to twenty is the useful range: 3.6–7.2× less allocated memory for at most a few percent of build time. +The disk-resident pair spills rather than streams, and it holds up the same way. +Same corpus and host, corpus mapped in every row, at λ=600 β=40 α=0.4 — a tenth +of the λ above, because the inline forward index copies every pruned posting's +whole doc vector, and at λ=6000 that section alone would be ~140 GB: + +| windows | peak RssAnon | build time | index | +|---|---|---|---| +| 1 (unbatched) | 8671 MB | 48 s | — | +| 10 | 1440 MB | 94 s | 16.9 GB | +| 20 | 774 MB | 210 s | 16.9 GB | +| 100 | 434 MB | 375 s | 16.9 GB | + +20× less allocated memory at 100 windows, 6× at 10. Build time climbs sooner than +it does for `seismic`: the spill is written and read back, and the payload write +is a second pass over the lists on top of the per-window corpus passes. The +unbatched row writes no index, so its build time is not comparable to the others' +either — it is the memory it is there for. + +`disk_seismic_sq` cannot map a float CSR (it searches over codes), so its corpus +sits on the heap and the figures include it: at 8-bit codes the corpus is 3242 MB +resident, and the build grows 6862 MB on top of that unbatched against 1087 MB at +10 windows. + Those rows are comparable to each other but not to a build that loads the corpus instead of mapping it, and the difference is not just the corpus. The same 1-window build measures 10274 MB mapped against 24084 MB with a streaming ingest — @@ -376,6 +408,11 @@ $B convert corpus.csr corpus.mcsr $B baseline corpus.csr 6000 400 0.4 # whole-corpus build, for reference $B batched inmem corpus.csr 6000 400 0.4 10 /data # 10 windows, same corpus residency $B batched mmap corpus.mcsr 6000 400 0.4 10 /data # ... or with the corpus mapped + +# Any type in the family, as its factory name. Both arms need it, and the +# baseline's index path is positional -- pass "" to skip writing one. +$B baseline corpus.csr 6000 400 0.4 "" disk_seismic +$B batched mmap corpus.mcsr 6000 400 0.4 10 /data disk_seismic ``` ## Python Bindings diff --git a/benchmarks/batched_build_mem_bench.cpp b/benchmarks/batched_build_mem_bench.cpp index ceb1c9b..d3f5503 100644 --- a/benchmarks/batched_build_mem_bench.cpp +++ b/benchmarks/batched_build_mem_bench.cpp @@ -15,14 +15,21 @@ // Usage: // batched_build_mem_bench convert // batched_build_mem_bench baseline \ -// [out_index] +// [out_index] [index_type] // batched_build_mem_bench batched \ -// +// [index_type] // // "convert" produces the native CSR the mapped read wants. "baseline" builds -// the in-memory SeismicIndex (streaming add of the interchange CSR, like the -// other benchmarks) -- the memory this feature exists to avoid. "batched" runs -// the term-batched build at the given batch count. +// the index whole (streaming add of the interchange CSR, like the other +// benchmarks) -- the memory this feature exists to avoid. "batched" runs the +// term-batched build at the given batch count. +// +// `index_type` is any seismic-family factory name, default "seismic". The two +// disk-resident ones are worth measuring separately: they cannot stream their +// payload out window by window, so their batched build spills the clustered +// lists and writes from that mapping instead (see build_streamed), and whether +// that holds anonymous memory down is exactly what this reports. A quantized +// type gets the factory's default range, which is fine for a memory figure. // // Compare "batched inmem" against "baseline": both hold the corpus on the heap // via the same streaming_add, so the difference between them is the batching @@ -47,11 +54,13 @@ #include #include #include +#include #include #include #include #include +#include "nsparse/index_factory.h" #include "nsparse/io/index_io.h" #include "nsparse/seismic_index.h" #include "nsparse/types.h" @@ -243,6 +252,26 @@ void streaming_add(nsparse::Index* index, const std::string& path) { } } +// The index under test, named rather than constructed, so every type in the +// family is reachable from the command line. The cluster knobs go through the +// factory description, which is also how a caller sets them (see +// parse_cluster_params); `out_path` empty leaves the batched output path unset. +std::unique_ptr make_index( + const std::string& index_type, int dimension, + const nsparse::SeismicClusterParameters& params, + const std::string& out_path) { + std::string desc = index_type + ",lambda=" + std::to_string(params.lambda) + + "|beta=" + std::to_string(params.beta) + + "|alpha=" + std::to_string(params.alpha) + + "|inverted_list_batch_size=" + + std::to_string(params.batch_clustering.batch_size); + if (!out_path.empty()) { + desc += "|batch_file_output_path=" + out_path; + } + return std::unique_ptr( + nsparse::index_factory(dimension, desc.c_str())); +} + // Both layouts start with the same int64 (rows, cols, nnz) header. int csr_dimension(const std::string& path) { std::ifstream file(path, std::ios::binary); @@ -289,16 +318,19 @@ int run_convert(int argc, char** argv) { int run_baseline(int argc, char** argv) { if (argc < 6) { std::cerr << "baseline " - "[out_index]\n"; + "[out_index] [index_type]\n"; return 2; } const std::string csr = argv[2]; - const nsparse::SeismicClusterParameters params = { + nsparse::SeismicClusterParameters params = { .lambda = std::atoi(argv[3]), .beta = std::atoi(argv[4]), .alpha = static_cast(std::atof(argv[5]))}; - nsparse::SeismicIndex index(csr_dimension(csr), params); - streaming_add(&index, csr); + params.batch_clustering.batch_size = 1; + const std::string index_type = argc >= 8 ? argv[7] : "seismic"; + std::unique_ptr index = + make_index(index_type, csr_dimension(csr), params, /*out_path=*/""); + streaming_add(index.get(), csr); const long load_hwm = read_vm_hwm_kib(); const long start_anon = read_status_kib("RssAnon:"); @@ -307,13 +339,16 @@ int run_baseline(int argc, char** argv) { // as the reset makes VmHWM exclude it. PeakRssSampler sampler; const double started = now_seconds(); - index.build(); + index->build(); const double build_s = now_seconds() - started; - report("baseline", "batches=0", build_s, load_hwm, start_anon, sampler); - - if (argc >= 7) { - const std::string out = argv[6]; - nsparse::write_index(&index, const_cast(out.c_str())); + report("baseline", "type=" + index_type + " batches=0", build_s, load_hwm, + start_anon, sampler); + + // Empty is how a caller reaches [index_type] without asking for the index + // to be written -- these are positional. + const std::string out = argc >= 7 ? argv[6] : ""; + if (!out.empty()) { + nsparse::write_index(index.get(), const_cast(out.c_str())); std::ifstream file(out, std::ios::binary | std::ios::ate); std::cout << "index_bytes=" << file.tellg() << "\n"; } @@ -323,20 +358,20 @@ int run_baseline(int argc, char** argv) { int run_batched(int argc, char** argv) { if (argc < 9) { std::cerr << "batched " - " \n"; + " [index_type]\n"; return 2; } const std::string corpus_residency = argv[2]; const std::string csr = argv[3]; - const nsparse::SeismicClusterParameters params = { + const std::string index_type = argc >= 10 ? argv[9] : "seismic"; + nsparse::SeismicClusterParameters batched_params = { .lambda = std::atoi(argv[4]), .beta = std::atoi(argv[5]), .alpha = static_cast(std::atof(argv[6]))}; - const std::string out = std::string(argv[8]) + "/index.seismic.dat"; - nsparse::SeismicClusterParameters batched_params = params; + const std::string out = + std::string(argv[8]) + "/index." + index_type + ".dat"; batched_params.batch_clustering.batch_size = static_cast(std::atoi(argv[7])); - batched_params.batch_clustering.batch_file_output_path = out; // Which residency the corpus is held at is the point of the flag: // @@ -346,11 +381,12 @@ int run_batched(int argc, char** argv) { // mmap -- a native CSR, borrowed. Cheaper by the size of the corpus, but // not comparable to the baseline, because the saving is the // residency rather than the batching. - nsparse::SeismicIndex index(csr_dimension(csr), batched_params); + std::unique_ptr index = + make_index(index_type, csr_dimension(csr), batched_params, out); if (corpus_residency == "inmem") { - streaming_add(&index, csr); + streaming_add(index.get(), csr); } else if (corpus_residency == "mmap") { - index.read_csr(csr.c_str(), nsparse::Residency::kMmap); + index->read_csr(csr.c_str(), nsparse::Residency::kMmap); } else { std::cerr << "corpus residency must be inmem or mmap\n"; return 2; @@ -361,17 +397,21 @@ int run_batched(int argc, char** argv) { reset_vm_hwm(); PeakRssSampler sampler; const double started = now_seconds(); - // batch_file_output_path is set, so build() streams the index out rather - // than retaining it -- the same call an ordinary build makes. - index.build(); + // batch_file_output_path is set, so build() writes the index out as it goes + // rather than assembling it in memory -- the same call an ordinary build + // makes. + index->build(); const double build_s = now_seconds() - started; report("batched", - "corpus=" + corpus_residency + " batches=" + + "type=" + index_type + " corpus=" + corpus_residency + " batches=" + std::to_string(batched_params.batch_clustering.batch_size), build_s, load_hwm, start_anon, sampler); - std::ifstream file(out, std::ios::binary | std::ios::ate); - std::cout << "index_bytes=" << file.tellg() << "\n"; + // One window writes no file: it is an ordinary build, holding its own + // lists. + if (std::ifstream file(out, std::ios::binary | std::ios::ate); file) { + std::cout << "index_bytes=" << file.tellg() << "\n"; + } return 0; } diff --git a/nsparse/disk_seismic_index.cpp b/nsparse/disk_seismic_index.cpp index e375cc4..41bbe02 100644 --- a/nsparse/disk_seismic_index.cpp +++ b/nsparse/disk_seismic_index.cpp @@ -57,7 +57,10 @@ DiskSeismicIndex* DiskSeismicIndex::mmap_index(const IndexHeader& header, MmapCursor cursor(mmap_file.data(), mmap_file.size()); cursor.skip(pos); - // No extra header for the float index; the shared payload follows directly. + // No extra header for the float index, so the hook is a no-op and the + // shared payload follows directly; called anyway, so the two disk types + // read their payload through the same two steps. + index->read_mapped_payload_header(&cursor); index->load_mapped_payload(&cursor, std::move(mmap_file)); return index.release(); } diff --git a/nsparse/disk_seismic_index_base.cpp b/nsparse/disk_seismic_index_base.cpp index 4078a06..3590543 100644 --- a/nsparse/disk_seismic_index_base.cpp +++ b/nsparse/disk_seismic_index_base.cpp @@ -12,9 +12,11 @@ #include #include #include +#include #include #include #include +#include #include #include @@ -23,8 +25,11 @@ #include "nsparse/disk_seismic_search.h" #include "nsparse/id_selector.h" #include "nsparse/index.h" +#include "nsparse/io/file_io.h" +#include "nsparse/io/index_io.h" #include "nsparse/io/inline_forward_index_io.h" #include "nsparse/io/seismic_invlists_writer.h" +#include "nsparse/seismic_batched_build.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" #include "nsparse/utils/checks.h" @@ -32,6 +37,35 @@ #include "nsparse/utils/mmap_file.h" namespace nsparse { +namespace { + +// Where a batched build spills its clustered lists: alongside the index it is +// writing, so it lands on whatever disk the caller chose for the output. +constexpr const char* kSpillSuffix = ".lists"; + +// Deletes the spill on the way out, whether the build finished or threw: it is +// scratch the size of the clustered lists, and nothing outside build() knows it +// exists. On the success path the build has already displaced the mapping of +// it, which matters on Windows, where a mapped file cannot be unlinked. +class SpillFile { +public: + explicit SpillFile(std::string path) : path_(std::move(path)) {} + ~SpillFile() { + std::error_code ignored; + std::filesystem::remove(path_, ignored); + } + SpillFile(const SpillFile&) = delete; + SpillFile& operator=(const SpillFile&) = delete; + SpillFile(SpillFile&&) = delete; + SpillFile& operator=(SpillFile&&) = delete; + + [[nodiscard]] const std::string& path() const { return path_; } + +private: + std::string path_; +}; + +} // namespace DiskSeismicIndexBase::DiskSeismicIndexBase(int dim, SeismicClusterParameters parameter) @@ -60,11 +94,61 @@ void DiskSeismicIndexBase::add(idx_t n, const idx_t* indptr, } void DiskSeismicIndexBase::build() { + const SparseVectorsConfig config = { + .element_size = code_element_size(), + .dimension = static_cast(get_dimension())}; + const auto& batch = cluster_parameter_.batch_clustering; + if (batch.batch_size > 1 && !batch.batch_file_output_path.empty()) { + build_streamed(config, batch.batch_file_output_path); + return; + } + // A single window is an ordinary build: it holds its own lists, and writing + // them out only to map them back would be work for nothing. clustered_inverted_lists = detail::build_inverted_lists_clusters( - get_vectors(), - {.element_size = code_element_size(), - .dimension = static_cast(get_dimension())}, - cluster_parameter_); + get_vectors(), config, cluster_parameter_); +} + +void DiskSeismicIndexBase::build_streamed(const SparseVectorsConfig& config, + const std::string& out_path) { + // This payload cannot be streamed section by section the way the in-memory + // types' can. Theirs ends with its posting lists, so a window can be + // serialized and dropped; here the summaries are followed by an inline + // forward index whose blocks are laid out from the doc-id membership of + // every list, and that membership is not known until the last window is + // clustered. + // + // So the clustering still runs once, a window at a time, but into a spill; + // the lists come back borrowed from it, and the payload is written from + // that mapping. Neither phase holds more than one window of anonymous + // memory -- the forward index streams its blocks out as it lays them. + const SpillFile spill(out_path + kSpillSuffix); + clustered_inverted_lists = + detail::spill_clustered_lists(get_vectors(), config, cluster_parameter_, + spill.path(), &batch_mapped_file_); + + size_t payload_offset = 0; + { + FileIOWriter writer(const_cast(out_path.c_str())); + detail::write_header({.id = fourcc(id()), + .version = format_version(), + .dimension = get_dimension()}, + &writer); + // Taken from the writer rather than from a constant, so the offset the + // mapping below skips to is the one the header actually occupied. + payload_offset = writer.pos(); + write_index(&writer); + writer.close(); + } + + // Ends holding the index it wrote, rather than an object whose lists point + // into scratch that is about to be deleted. The output's mapping displaces + // the spill's, which is safe in that order: load_mapped_payload replaces + // the lists that borrowed from the spill before it commits the mapping. + MmapFile mapped(out_path); + MmapCursor cursor(mapped.data(), mapped.size()); + cursor.skip(payload_offset); + read_mapped_payload_header(&cursor); + load_mapped_payload(&cursor, std::move(mapped), &batch_mapped_file_); } auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr, @@ -162,9 +246,8 @@ void DiskSeismicIndexBase::write_index(IOWriter* 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 = code_element_size(), - .dimension = static_cast(dimension_)}); + SparseVectors empty_vectors({.element_size = code_element_size(), + .dimension = static_cast(dimension_)}); const SparseVectors& v = vectors_ != nullptr ? *vectors_ : empty_vectors; detail::InlineForwardIndex forward(clustered_inverted_lists, v); forward.serialize(io_writer); @@ -181,7 +264,8 @@ void DiskSeismicIndexBase::read_index(IOReader* /*io_reader*/, } void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, - MmapFile&& mapped) { + MmapFile&& mapped, + MmapFile* slot) { // Same order write_index wrote them (past any extra header the caller // already consumed): doc count, summaries, inline forward. num_vectors_ = cursor->read_scalar(); @@ -190,16 +274,18 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, detail::InlineForwardIndex forward; forward.mmap_deserialize(cursor); + // Before the mapping is committed below: whatever these replace may have + // been borrowing from what `slot` still holds -- a batched build's spill. clustered_inverted_lists = std::move(inv_list_writer.release()); fwd_ = std::move(forward); // 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. + // from `mapped`, which is alive here), let the concrete index reject a + // width mismatch before we commit. validate_mapped_payload(); - // mapped_file_ last: the summaries and the forward index borrow from it, and + // The mapping last: the summaries and the forward index borrow from it, and // moving it does not move the mapping. - mapped_file_ = std::move(mapped); + *slot = std::move(mapped); } } // namespace nsparse diff --git a/nsparse/disk_seismic_index_base.h b/nsparse/disk_seismic_index_base.h index be2f53f..6224a0e 100644 --- a/nsparse/disk_seismic_index_base.h +++ b/nsparse/disk_seismic_index_base.h @@ -12,6 +12,8 @@ #include #include +#include +#include #include #include "nsparse/cluster/inverted_list_clusters.h" @@ -28,11 +30,12 @@ namespace nsparse { // 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, -// and search scores the global top-k_prime blocks (a DiskSeismicSearchParameters -// sets k_prime). add / build / search / serialization / mmap loading are all -// here; the concrete indexes differ only in the stored value width and, for the -// quantized one, a leading quantization header and a score-decoding step, which -// they supply through the virtual hooks below. +// and search scores the global top-k_prime blocks (a +// DiskSeismicSearchParameters sets k_prime). add / build / search / +// serialization / mmap loading are all here; the concrete indexes differ only +// in the stored value width and, for the quantized one, a leading quantization +// header and a score-decoding step, which they supply through the virtual hooks +// below. // // mmap-only: load with read_index(file, kUseMmap); the copying read throws. class DiskSeismicIndexBase : public MmapIndex, public IndexIO { @@ -71,10 +74,12 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO { // code_element_size()-byte-per-value data. `scratch` backs the result when // encoding must allocate; the float index returns its input reinterpreted, // with no copy. - virtual const uint8_t* encode_values(const float* values, size_t nnz, - std::vector& scratch) const = 0; + virtual const uint8_t* encode_values( + const float* values, size_t nnz, + std::vector& scratch) const = 0; - // Encode a query batch the same way, honoring any per-search range override. + // Encode a query batch the same way, honoring any per-search range + // override. virtual const uint8_t* encode_query( const float* values, size_t nnz, const SearchParameters* search_parameters, @@ -89,22 +94,45 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO { // quantization header for the quantized index; nothing for the float one). virtual void write_payload_header(IOWriter* /*io_writer*/) const {} - // Reject a just-mapped payload whose stored width disagrees with this index. - // No-op for the float index, which stores a fixed width. Runs after the - // summaries and forward index are populated but before the mapping commits. + // The mirror of write_payload_header: consume that header off a mapping and + // adopt what it declares. Used by the mapped read and by a batched build + // reopening the file it just wrote, so the two cannot read the payload from + // different offsets. + virtual void read_mapped_payload_header(MmapCursor* /*cursor*/) {} + + // Reject a just-mapped payload whose stored width disagrees with this + // index. No-op for the float index, which stores a fixed width. Runs after + // the summaries and forward index are populated but before the mapping + // commits. virtual void validate_mapped_payload() const {} // Reads the shared payload (doc count, summaries, inline forward) from the - // cursor, validates it, and commits the mapping. Each concrete mmap_index - // reads its own extra header first, then calls this. - void load_mapped_payload(MmapCursor* cursor, MmapFile&& mapped); + // cursor, validates it, and commits the mapping into `slot`. Each concrete + // mmap_index reads its own extra header first, then calls this. + // + // `slot` is mapped_file_ for a read_index load, which owns nothing else, + // but batch_mapped_file_ for a batched build: there the corpus may still be + // borrowing from mapped_file_, and giving that up would leave the index + // unable to score anything. + void load_mapped_payload(MmapCursor* cursor, MmapFile&& mapped, + MmapFile* slot); + void load_mapped_payload(MmapCursor* cursor, MmapFile&& mapped) { + load_mapped_payload(cursor, std::move(mapped), &mapped_file_); + } - // Borrowed from by score_summaries_transposed / the inline forward index, so - // the concrete validate_mapped_payload can inspect their widths. + // Borrowed from by score_summaries_transposed / the inline forward index, + // so the concrete validate_mapped_payload can inspect their widths. std::vector clustered_inverted_lists; detail::InlineForwardIndex fwd_; private: + // build() with batch_file_output_path set: clusters one term window at a + // time into a spill file, writes the index out of that mapping, and ends + // borrowing its own output. See the definition for why this payload cannot + // be streamed section by section the way the in-memory ones are. + void build_streamed(const SparseVectorsConfig& config, + const std::string& out_path); + auto search(idx_t n, const idx_t* indptr, const term_t* indices, const float* values, int k, SearchParameters* search_parameters = nullptr) diff --git a/nsparse/disk_seismic_scalar_quantized_index.cpp b/nsparse/disk_seismic_scalar_quantized_index.cpp index 7eb4201..4aaf117 100644 --- a/nsparse/disk_seismic_scalar_quantized_index.cpp +++ b/nsparse/disk_seismic_scalar_quantized_index.cpp @@ -142,6 +142,14 @@ void DiskSeismicScalarQuantizedIndex::write_payload_header( io_writer->write(&vmax, sizeof(float), 1); } +void DiskSeismicScalarQuantizedIndex::read_mapped_payload_header( + MmapCursor* cursor) { + const auto sq_type = cursor->read_scalar(); + const auto vmin = cursor->read_scalar(); + const auto vmax = cursor->read_scalar(); + sq_ = make_scalar_quantizer(sq_type, vmin, vmax); +} + void DiskSeismicScalarQuantizedIndex::validate_mapped_payload() const { throw_if_forward_width_mismatch(fwd_, sq_); throw_if_summary_width_mismatch(clustered_inverted_lists, sq_); @@ -159,11 +167,7 @@ DiskSeismicScalarQuantizedIndex* DiskSeismicScalarQuantizedIndex::mmap_index( // The quantization header opens the payload; the shared loader reads the // rest and calls validate_mapped_payload() against this quantizer. - const auto sq_type = cursor.read_scalar(); - const auto vmin = cursor.read_scalar(); - const auto vmax = cursor.read_scalar(); - index->sq_ = make_scalar_quantizer(sq_type, vmin, vmax); - + index->read_mapped_payload_header(&cursor); index->load_mapped_payload(&cursor, std::move(mmap_file)); return index.release(); } diff --git a/nsparse/disk_seismic_scalar_quantized_index.h b/nsparse/disk_seismic_scalar_quantized_index.h index fd344da..f4c403d 100644 --- a/nsparse/disk_seismic_scalar_quantized_index.h +++ b/nsparse/disk_seismic_scalar_quantized_index.h @@ -90,6 +90,8 @@ class DiskSeismicScalarQuantizedIndex : public DiskSeismicIndexBase { // The quantization parameters that open this index's payload -- distinct // from the IndexHeader the file itself starts with. void write_payload_header(IOWriter* io_writer) const override; + // Adopts the quantizer a mapped payload declares, rejecting an unknown type. + void read_mapped_payload_header(MmapCursor* cursor) override; void validate_mapped_payload() const override; // The quantizer a query is encoded with: DiskSeismicSQSearchParameters diff --git a/nsparse/seismic_batched_build.cpp b/nsparse/seismic_batched_build.cpp index a025e5d..2fdc2e7 100644 --- a/nsparse/seismic_batched_build.cpp +++ b/nsparse/seismic_batched_build.cpp @@ -27,38 +27,41 @@ #include "nsparse/utils/mmap_file.h" namespace nsparse::detail { +namespace { -size_t write_seismic_index_batched( - const SparseVectors* vectors, const SparseVectorsConfig& config, - const SeismicClusterParameters& params, const IndexHeader& header, - const std::function& write_prefix, - const std::string& out_path) { - if (out_path.empty()) { - throw std::invalid_argument( - "write_seismic_index_batched: output path must not be empty"); +// Both entry points below refuse the same two inputs: nowhere to write to, and +// a corpus with no postings to write. +void throw_if_not_streamable(const SparseVectors* vectors, + const std::string& path, const char* who) { + if (path.empty()) { + throw std::invalid_argument(std::string(who) + + ": output path must not be empty"); } if (vectors == nullptr || vectors->num_vectors() == 0) { throw std::invalid_argument( - "write_seismic_index_batched: corpus is empty; there is nothing to " - "stream"); + std::string(who) + ": corpus is empty; there is nothing to stream"); } +} - // One writer for the whole file, windows serialized straight into it rather - // than spilled and concatenated: serialize() pads each array relative to - // the writer's current offset (see io/align.h), so bytes produced by a - // writer that started at 0 carry the wrong padding once appended at some - // other offset. Streaming through a single writer keeps pos() the true - // absolute offset. - FileIOWriter writer(const_cast(out_path.c_str())); - write_header(header, &writer); - write_prefix(&writer); - +// The posting-list section -- [count][list...], the layout +// SeismicInvertedListsWriter produces -- streamed into `writer` one window at a +// time, starting wherever the writer has reached. +// +// The writer is the one the whole file is being written through, rather than a +// per-window one whose output is concatenated: serialize() pads each array +// relative to the writer's current offset (see io/align.h), so bytes produced +// by a writer that started at 0 carry the wrong padding once appended at some +// other offset. Streaming through a single writer keeps pos() the true absolute +// offset. +void stream_clustered_lists(const SparseVectors* vectors, + const SparseVectorsConfig& config, + const SeismicClusterParameters& params, + IOWriter* writer) { // The list count, exactly where SeismicInvertedListsWriter::serialize puts // it. It is the whole dimension, known before any window is built, which is // what lets the lists be streamed after it rather than counted first. - const size_t lists_offset = writer.pos(); size_t n_lists = config.dimension; - writer.write(&n_lists, sizeof(size_t), 1); + writer->write(&n_lists, sizeof(size_t), 1); // Windows arrive in ascending term order, so appending each in turn // produces the same byte sequence as writing every list at once. @@ -70,24 +73,57 @@ size_t write_seismic_index_batched( // The layout carries no per-list offsets, so a gap or a repeat // would silently shift every list after it. throw std::runtime_error( - "write_seismic_index_batched: windows arrived out of " - "order"); + "stream_clustered_lists: windows arrived out of order"); } for (const auto& list : clusters) { - list.serialize(&writer); + list.serialize(writer); } next_term = term_begin + clusters.size(); // clusters freed on return, before the next window is built. }); if (next_term != config.dimension) { throw std::runtime_error( - "write_seismic_index_batched: wrote " + std::to_string(next_term) + + "stream_clustered_lists: wrote " + std::to_string(next_term) + " of " + std::to_string(config.dimension) + " posting lists"); } +} + +} // namespace + +size_t write_seismic_index_batched( + const SparseVectors* vectors, const SparseVectorsConfig& config, + const SeismicClusterParameters& params, const IndexHeader& header, + const std::function& write_prefix, + const std::string& out_path) { + throw_if_not_streamable(vectors, out_path, "write_seismic_index_batched"); + + FileIOWriter writer(const_cast(out_path.c_str())); + write_header(header, &writer); + write_prefix(&writer); + + const size_t lists_offset = writer.pos(); + stream_clustered_lists(vectors, config, params, &writer); writer.close(); return lists_offset; } +std::vector spill_clustered_lists( + const SparseVectors* vectors, const SparseVectorsConfig& config, + const SeismicClusterParameters& params, const std::string& path, + MmapFile* into) { + throw_if_not_streamable(vectors, path, "spill_clustered_lists"); + { + // Closed before the mapping is taken: the writer buffers, and what is + // not flushed is not in the file to map. + FileIOWriter writer(const_cast(path.c_str())); + stream_clustered_lists(vectors, config, params, &writer); + writer.close(); + } + // Offset 0: a spill is the section and nothing else, with no header for it + // to sit behind. + return map_streamed_lists(path, /*lists_offset=*/0, into); +} + std::vector map_streamed_lists(const std::string& path, size_t lists_offset, MmapFile* into) { diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h index b447c05..05f666c 100644 --- a/nsparse/seismic_batched_build.h +++ b/nsparse/seismic_batched_build.h @@ -36,12 +36,13 @@ namespace nsparse::detail { // holds, at whatever residency SparseVectors was given) plus one window. // // Reached through an index's build(), by setting -// SeismicClusterParameters::batch_clustering.batch_file_output_path. The index -// is then the file, not the object: nothing is retained to serve or to -// write_index afterwards. +// SeismicClusterParameters::batch_clustering.batch_file_output_path. The build +// then maps the file back (see map_streamed_lists) rather than dropping it, so +// it ends holding the index it wrote. // -// `header` and `write_prefix` are what make this work for every type in the -// family rather than just SEIS. `write_prefix` writes whatever the type puts +// `header` and `write_prefix` are what make this work for every type whose +// payload ends with its posting lists, rather than just SEIS. `write_prefix` +// writes whatever the type puts // between the header and its posting lists -- the forward vectors, and for a // quantizing index its quantization header first. The lists then follow in the // byte-for-byte layout SeismicInvertedListsWriter produces, so the file is an @@ -64,6 +65,31 @@ size_t write_seismic_index_batched( const std::function& write_prefix, const std::string& out_path); +// Streams every window's clustered posting lists to `path` and then maps them +// back, so the caller ends up holding all of them without two windows ever +// having been resident at once. +// +// For the index types whose payload ends with its posting lists, +// write_seismic_index_batched writes the index itself and there is nothing to +// spill. A DiskSeismic payload is not one of those: its summaries precede an +// inline forward index whose blocks are laid out from the doc-id membership of +// every list, so no window's lists can be dropped before the last window is +// clustered. What can be dropped is their *residency* -- which is what this is +// for. `path` gets the lists in the same [count][list...] layout +// SeismicInvertedListsWriter produces, doc ids included (an index's own section +// writes them empty; the forward index is what needs them here), and the +// returned lists borrow from the mapping handed to `into` rather than the heap. +// +// The spill is scratch, not an index: it carries no header, nothing else reads +// it, and deleting it is the caller's job. It must outlive the returned lists. +// +// Throws if the corpus is empty, for the same reason as +// write_seismic_index_batched: there would be no windows to stream. +std::vector spill_clustered_lists( + const SparseVectors* vectors, const SparseVectorsConfig& config, + const SeismicClusterParameters& params, const std::string& path, + MmapFile* into); + // Maps the file a streamed build just wrote and borrows its posting lists out // of it, so the build ends holding a usable index without ever having held all // of the lists at once. diff --git a/python_tests/test_seismic_batched_build.py b/python_tests/test_seismic_batched_build.py index 8120846..cd38e9f 100644 --- a/python_tests/test_seismic_batched_build.py +++ b/python_tests/test_seismic_batched_build.py @@ -62,18 +62,49 @@ def test_happy_case(batch_size, corpus, queries, oracle, tmp_path): assert recall_at_k(labels, want_labels) >= RECALL_FLOOR -@pytest.mark.parametrize("kind", ["seismic", "seismic_sq"]) +@pytest.mark.parametrize( + "kind", ["seismic", "seismic_sq", "disk_seismic", "disk_seismic_sq"] +) def test_matches_in_memory_build(kind, corpus, tmp_path): """At a fixed seed a streamed build is the in-memory build, byte for byte. - Parametrized over a float and a quantizing index, because the shared build - only needs the code width -- add() has already encoded the values. + Over all four types in the family: float and quantizing, since the shared + build only needs the code width (add() has already encoded the values), and + in-memory and disk-resident, which get there differently -- the disk types' + payload cannot be streamed section by section, so they spill their clustered + lists and write from that mapping instead. """ in_memory = tmp_path / "memory.idx" nsparse.write_index(make_index(f"{kind},{BASE}", corpus), str(in_memory)) batched = streamed(corpus, tmp_path / "batched.idx", 4, kind=kind) assert in_memory.read_bytes() == open(batched, "rb").read() + # The spill the disk types take is scratch, deleted with the build. + assert not (tmp_path / "batched.idx.lists").exists() + + +@pytest.mark.parametrize("kind", ["disk_seismic", "disk_seismic_sq"]) +def test_streamed_disk_index_is_searchable_after_build( + kind, corpus, queries, oracle, tmp_path +): + """A batched disk build serves from the file it wrote. + + Its summaries and its inline forward index are borrowed from that mapping, + so there is no reopening by path -- and nothing left pointing at the spill. + """ + index = nsparse.index_factory( + corpus.dim, + f"{kind},{BASE}|inverted_list_batch_size=8" + f"|batch_file_output_path={tmp_path / 'streamed.idx'}", + ) + add_corpus(index, corpus) + index.build() + + assert index.num_vectors() == corpus.n + params = nsparse.DiskSeismicSearchParameters(8, 200) + _, labels = search(index, queries, params=params) + want_labels, _ = oracle + assert recall_at_k(labels, want_labels) >= RECALL_FLOOR def test_streamed_index_is_searchable_after_build(corpus, queries, oracle, tmp_path): diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index c9492dd..bcc930b 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -26,6 +26,7 @@ #include "csr_interchange_test_util.h" #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/disk_seismic_index.h" +#include "nsparse/disk_seismic_scalar_quantized_index.h" #include "nsparse/index_factory.h" #include "nsparse/io/file_io.h" #include "nsparse/io/index_io.h" @@ -223,10 +224,7 @@ TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeAnInMemoryBuild) { } // The disk-resident types share the same build, so batch_size has to bound -// their intermediates too without changing what they produce. They have no -// streaming write yet -- their payload interleaves summaries with an inline -// forward index -// -- so this covers the half they do get. +// their intermediates too without changing what they produce. TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeADiskIndex) { Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/71); TempDir dir("disk"); @@ -246,6 +244,159 @@ TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeADiskIndex) { EXPECT_EQ(unbatched, build_disk(64, dir.file("b64.dat"))); } +// The disk types cannot stream their payload out window by window -- the inline +// forward index that follows their summaries is laid out from the doc-id +// membership of every list -- so they spill the lists instead and write the +// payload from that mapping. Same contract as the in-memory types all the same: +// at a fixed seed the file is what write_index would have produced. +TEST(SeismicBatchedBuild, StreamsADiskIndexIdenticallyToo) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/71); + TempDir dir("disk_streamed"); + const std::string mem_path = dir.file("mem.dat"); + const std::string streamed_path = dir.file("streamed.dat"); + + DiskSeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + write_index(&mem, const_cast(mem_path.c_str())); + + DiskSeismicIndex batched(corpus.dim, params_for(8, streamed_path, kSeed)); + batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + batched.build(); + + EXPECT_EQ(read_file(mem_path), read_file(streamed_path)); + // The spill is scratch: it must not outlive the build that took it. + EXPECT_FALSE(std::filesystem::exists(streamed_path + ".lists")); + // And the window count is still not a behaviour knob. + const std::string many_path = dir.file("many.dat"); + DiskSeismicIndex many(corpus.dim, params_for(64, many_path, kSeed)); + many.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + many.build(); + EXPECT_EQ(read_file(mem_path), read_file(many_path)); +} + +// The quantized disk index writes a quantization header before the shared +// payload, so a batched build has to lay that down and read it back to reopen +// its own file at the right offset. +TEST(SeismicBatchedBuild, StreamsAQuantizedDiskIndexIdenticallyToo) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/29); + TempDir dir("disk_sq"); + const std::string mem_path = dir.file("mem.dat"); + const std::string streamed_path = dir.file("streamed.dat"); + + DiskSeismicScalarQuantizedIndex mem(QuantizerType::QT_8bit, 0.0F, 3.0F, + params_for(1, "", kSeed), corpus.dim); + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + write_index(&mem, const_cast(mem_path.c_str())); + + DiskSeismicScalarQuantizedIndex batched(QuantizerType::QT_8bit, 0.0F, 3.0F, + params_for(8, streamed_path, kSeed), + corpus.dim); + batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + batched.build(); + + EXPECT_EQ(read_file(mem_path), read_file(streamed_path)); + // It reads back as the quantized disk type it claims to be, with the range + // it was built with -- the header the batched write had to reproduce. + std::unique_ptr reloaded(read_index( + const_cast(streamed_path.c_str()), IndexIoFlag::kUseMmap)); + EXPECT_EQ(reloaded->id(), DiskSeismicScalarQuantizedIndex::name); + EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n)); + const auto* sq_index = + dynamic_cast(reloaded.get()); + ASSERT_NE(sq_index, nullptr); + EXPECT_EQ(sq_index->get_scalar_quantizer().get_min(), 0.0F); + EXPECT_EQ(sq_index->get_scalar_quantizer().get_max(), 3.0F); +} + +// A batched disk build ends serving from the file it wrote: its summaries and +// its forward index are borrowed from that mapping, not from the spill it +// deleted. Against an unbatched build at the same seed, so identical results +// rather than merely close ones. +TEST(SeismicBatchedBuild, BatchedDiskBuildIsSearchableAfterBuild) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/17); + Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); + const int k = 10; + const auto n = static_cast(queries.n); + TempDir dir("disk_searchable"); + + const auto search_with = [&](Index& index) { + std::vector dist(n * k); + std::vector lab(n * k); + DiskSeismicSearchParameters params(/*cut=*/3, /*k_prime=*/50); + index.search(queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, dist.data(), lab.data(), + ¶ms); + return std::pair{dist, lab}; + }; + + // The unbatched reference has to be read back mapped: an unwritten disk + // index has no forward index to score from. + const std::string mem_path = dir.file("mem.dat"); + DiskSeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + write_index(&mem, const_cast(mem_path.c_str())); + std::unique_ptr reference( + read_index(const_cast(mem_path.c_str()), IndexIoFlag::kUseMmap)); + const auto [want_dist, want_lab] = search_with(*reference); + + DiskSeismicIndex batched(corpus.dim, + params_for(4, dir.file("b.dat"), kSeed)); + batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + batched.build(); + + // No reopening by path: build() mapped its own output back in. + EXPECT_EQ(batched.num_vectors(), static_cast(corpus.n)); + const auto [got_dist, got_lab] = search_with(batched); + EXPECT_EQ(got_lab, want_lab); + EXPECT_EQ(got_dist, want_dist); +} + +// The disk index's own reason for existing: a corpus that came from a mapping. +// Three mappings are then live at once -- the corpus, the spill, and the output +// -- and the build may give up only the spill. +TEST(SeismicBatchedBuild, BatchedDiskBuildKeepsTheCorpusMapping) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/13); + Corpus queries = make_corpus(/*n_docs=*/40, /*dim=*/200, /*seed=*/77); + const int k = 10; + const auto n = static_cast(queries.n); + TempDir dir("disk_mapped"); + + const std::string native = write_native_csr(corpus, dir.file("corpus.csr")); + const std::string out = dir.file("out.dat"); + DiskSeismicIndex index(corpus.dim, params_for(3, out, kSeed)); + index.read_csr(native.c_str(), Residency::kMmap); + index.build(); + + EXPECT_EQ(index.num_vectors(), static_cast(corpus.n)); + std::vector dist(n * k); + std::vector lab(n * k); + DiskSeismicSearchParameters params(/*cut=*/3, /*k_prime=*/50); + static_cast(index).search( + queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, dist.data(), lab.data(), ¶ms); + EXPECT_TRUE( + std::any_of(lab.begin(), lab.end(), [](idx_t id) { return id >= 0; })); + + // Same file a heap-resident corpus produces: residency is SparseVectors' + // business, not the build's. + DiskSeismicIndex owned(corpus.dim, + params_for(3, dir.file("owned.dat"), kSeed)); + owned.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + owned.build(); + EXPECT_EQ(read_file(out), read_file(dir.file("owned.dat"))); +} + // The generalization that matters: a quantizing index streams too, because the // codes in `vectors_` are already quantized by add() and the shared build only // needs their width. @@ -521,6 +672,11 @@ TEST(SeismicBatchedBuild, RejectsInvalidInput) { // cannot parse, so it is refused rather than written. SeismicIndex empty(corpus.dim, params_for(4, dir.file("empty.dat"), kSeed)); EXPECT_THROW(empty.build(), std::invalid_argument); + + // Same for a disk index, whose spill would map back to no lists at all. + DiskSeismicIndex empty_disk( + corpus.dim, params_for(4, dir.file("empty_disk.dat"), kSeed)); + EXPECT_THROW(empty_disk.build(), std::invalid_argument); } } // namespace nsparse From cb6a5ceaa22dc258501c0593e3d989feb1b8172e Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 3 Sep 2026 05:25:59 +0000 Subject: [PATCH 11/15] docs: keep the reasoning, drop the corpus it was measured on Comments and the guide carried figures from one benchmark run on one dataset -- corpus names, per-term posting counts, peak-memory tables, query-latency rows. They date the moment they are read: a reader cannot tell whether they still hold, and the numbers are not what any of it is explaining. What is worth keeping is why the code is shaped this way -- that term frequencies are skewed so windows are cut by cost, that clustered postings are far bulkier than scattered ones, that anonymous memory is the column batching moves. That survives here without a corpus attached. The guide keeps the shape to expect and how to measure it; measurements belong to a run, and a run belongs in its own report. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 129 +++++++++++-------------------- nsparse/cluster/kmeans_utils.cpp | 6 +- nsparse/seismic_common.cpp | 40 ++++------ 3 files changed, 62 insertions(+), 113 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 732ce24..4f99b7f 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -292,103 +292,60 @@ cannot change what is produced. ### Choosing `inverted_list_batch_size` Windows are cut to equal estimated *memory*, not equal width. Term frequencies are -heavily skewed — on msmarco base_full the heaviest term holds 5.7M postings -against a mean of 37K — and peak memory is set by the largest window, so an uneven -split wastes most of what batching could save. +heavily skewed, and peak memory is set by the largest window, so an uneven split +wastes most of what batching could save. A window has two memory peaks and the split has to weigh both. Filling it holds every posting of its terms; clustering it holds what survives pruning -(`min(count, lambda)` per term, 10% of that corpus's postings) as clusters and -summaries, roughly 16× bulkier per posting. Weighting either phase alone -unbalances the other, both measurably worse than weighting their sum — see -`make_windows` in `nsparse/seismic_common.cpp`, which records what each choice -measured. +(`min(count, lambda)` per term) as clusters and summaries, an order of magnitude +bulkier per posting. Weighting either phase alone unbalances the other, both worse +than weighting their sum — see `make_windows` in `nsparse/seismic_common.cpp`. `RssAnon` is the figure to watch, being what the process itself allocated; `RssFile` is pages it touched of a mapping, which the kernel can reclaim under -pressure. On base_full (8.8M docs, dim 30109, 1.12B non-zeros, λ=6000 β=400 -α=0.4) on a 36-core/68GB host, with the corpus **mapped in every row**, so each -figure is the build's own memory: - -| windows | peak RssAnon | peak RSS | peak RssFile | build time | -|---|---|---|---|---| -| 1 (unbatched) | 10281 MB | 16732 MB | 6454 MB | 105 s | -| 10 | 2816 MB | 12811 MB | 9820 MB | 111 s | -| 20 | 1471 MB | 11465 MB | 9185 MB | 130 s | -| 100 | 465 MB | 10396 MB | 9725 MB | 295 s | - -The unbatched row writes no file and maps nothing back, which is why its `RssFile` -is the corpus alone. Anonymous memory falls faster than 1/N — 3.7× at 10 windows -and 22× at 100 — -because the split comes from the real per-term costs rather than from term ids. At -100 windows the build allocates 465 MB while indexing 1.12 billion postings. Total -RSS barely moves, being dominated by page cache: the 6.45 GB mapped corpus plus -~3.4 GB of the index touched when `build()` maps it back in. That is the whole -reason to report the split — measured as RSS alone this looks like a 1.9× win -rather than a 22× one, and neither figure in that column is memory the process -would have to give up under pressure. +pressure. Read as total RSS the win looks far smaller than it is: an index that +maps its corpus, and then maps its own output back, keeps most of its residency in +page cache, which total RSS counts and pressure reclaims. Batching moves the +anonymous column, so report the split rather than the total. + +What to expect from the shape of it: anonymous memory falls faster than 1/N, +because the split comes from the real per-term costs rather than from term ids, +and build time is flat to around ten windows and then climbs, because every window +makes its own pass over the corpus. Ten to twenty windows is usually the useful +range — most of the memory saving for a few percent of build time. The floor is +what the build cannot batch: one window plus whatever the corpus itself costs. + +The disk-resident pair behaves the same way, with two differences. Build time +climbs sooner, since the spill is written and read back and the payload write is +another pass over the lists; and `lambda` is the knob to watch for output size, +because their inline forward index copies every pruned posting's whole doc +vector — at a large `lambda` that section alone can dwarf the rest of the index, +and the spill needs scratch disk beside it. `disk_seismic_sq` cannot map a float +CSR (it searches over codes), so its corpus stays on the heap and shows up in the +same column the build's own growth does; subtract the reported +`start_rss_anon_mb`. Mapping the finished lists back in is nearly free in the column that matters: -across those rows it costs about 4 s and leaves `RssAnon` unchanged (2816 MB at 10 -windows against 2820 MB without it). Borrowing a 14.9 GB index takes 0.19 s and -8 MB of anonymous memory, against 11.2 s and 13.9 GB to copy it — the cursor only -reads the size header before each array and skips the bulk, so it faults in about -a quarter of the file as reclaimable page cache. That mapping is separate from the -corpus's: an index that mapped its corpus with `read_csr` keeps doing so, since it -still scores from it. - -Build time is flat to around ten windows and then climbs, because every window -makes its own pass over the corpus. Ten to twenty is the useful range: 3.6–7.2× -less allocated memory for at most a few percent of build time. - -The disk-resident pair spills rather than streams, and it holds up the same way. -Same corpus and host, corpus mapped in every row, at λ=600 β=40 α=0.4 — a tenth -of the λ above, because the inline forward index copies every pruned posting's -whole doc vector, and at λ=6000 that section alone would be ~140 GB: - -| windows | peak RssAnon | build time | index | -|---|---|---|---| -| 1 (unbatched) | 8671 MB | 48 s | — | -| 10 | 1440 MB | 94 s | 16.9 GB | -| 20 | 774 MB | 210 s | 16.9 GB | -| 100 | 434 MB | 375 s | 16.9 GB | - -20× less allocated memory at 100 windows, 6× at 10. Build time climbs sooner than -it does for `seismic`: the spill is written and read back, and the payload write -is a second pass over the lists on top of the per-window corpus passes. The -unbatched row writes no index, so its build time is not comparable to the others' -either — it is the memory it is there for. - -`disk_seismic_sq` cannot map a float CSR (it searches over codes), so its corpus -sits on the heap and the figures include it: at 8-bit codes the corpus is 3242 MB -resident, and the build grows 6862 MB on top of that unbatched against 1087 MB at -10 windows. - -Those rows are comparable to each other but not to a build that loads the corpus -instead of mapping it, and the difference is not just the corpus. The same -1-window build measures 10274 MB mapped against 24084 MB with a streaming ingest — -13.8 GB more for a 6.86 GB corpus — because the ingest stages a second copy of it -(a 13.0 GB load peak) that the allocator retains rather than returning to the OS. -How much of that overlaps the build's own peak depends on how much the build then -asks for, so do not read a heap figure and a mapped figure as differing by a fixed -offset. - -All of the above is measured from the start of the build, with the corpus already -resident. Loading it costs more than holding it — a streaming ingest stages a -second copy — so a whole-process high-water mark would report the loader (13.0 GB -on this corpus, on the heap path) rather than the build, and hide everything below -it. Peak RSS comes from `VmHWM`, a kernel counter; the anon and file peaks have no +borrowing an index costs a fraction of a second and single-digit megabytes of +anonymous memory, against copying it, which costs its whole size — the cursor +reads the size header before each array and skips the bulk, so it faults in part +of the file as reclaimable page cache. That mapping is separate from the corpus's: +an index that mapped its corpus with `read_csr` keeps doing so, since it still +scores from it. + +Compare like with like. A build that loads the corpus is not comparable to one +that maps it, and the difference is more than the corpus: a streaming ingest +stages a second copy that the allocator retains rather than returning to the OS, +so do not read a heap figure and a mapped figure as differing by a fixed offset. +Measure from the start of the build with the corpus already resident, or a +whole-process high-water mark reports the loader and hides everything below it. +Peak RSS comes from `VmHWM`, a kernel counter; the anon and file peaks have no such counter and are sampled, so they are lower bounds and can disagree with `VmHWM` by a hair. -Query performance does not move, because the index is the same index. Two -independent unseeded builds, whole-corpus against 10 windows, over 6980 msmarco -dev queries at k=10, read in-memory: - -| index | QPS | p50 | p90 | p99 | recall@10 | -|---|---|---|---|---|---| -| whole corpus | 69332 | 0.290 ms | 0.601 ms | 1.000 ms | 0.8406 | -| 10 windows | 70198 | 0.296 ms | 0.616 ms | 1.055 ms | 0.8432 | +Query performance does not move, because the index is the same index: identical +byte for byte at a fixed seed, and indistinguishable in latency, QPS and recall +for the random-seeded default. Write the index to a real disk. On a tmpfs such as `/tmp` it is RAM, which defeats the point. diff --git a/nsparse/cluster/kmeans_utils.cpp b/nsparse/cluster/kmeans_utils.cpp index 1f323b7..5247182 100644 --- a/nsparse/cluster/kmeans_utils.cpp +++ b/nsparse/cluster/kmeans_utils.cpp @@ -101,9 +101,9 @@ CentroidIndex build_centroid_index( // centroid shares a term, instead of scoring the doc against a dense // dimension x n_clusters centroid matrix. Cost falls from // O(n_docs * doc_nnz * n_clusters) to O(shared postings), and scratch from -// dimension * n_clusters floats to the centroids' non-zeros — the dense matrix -// (255 MB at dimension=30522, beta=2087) was rebuilt for every posting list and -// dominated build time. +// dimension * n_clusters floats to the centroids' non-zeros — the dense matrix, +// hundreds of megabytes at a realistic dimension and cluster count, was rebuilt +// for every posting list and dominated build time. template void map_docs_to_clusters_typed(const SparseVectors* vectors, const std::vector& docs, diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp index 4e54384..a0333a8 100644 --- a/nsparse/seismic_common.cpp +++ b/nsparse/seismic_common.cpp @@ -36,40 +36,32 @@ struct TermWindow { // How much more a posting costs once clustered than while being scattered into // an inverted list, per unit. A window's memory has two peaks: filling it holds -// every posting of its terms, and clustering it holds the pruned survivors as -// clusters and summaries, which are far bulkier per posting -- on msmarco -// base_full the finished lists come to 14.9GB for 115M pruned postings, against -// 8 bytes each while filling. +// every posting of its terms at a few bytes each, and clustering it holds the +// pruned survivors as clusters and summaries, which are an order of magnitude +// bulkier per posting. // // Only the ratio matters, and only roughly: the cost curve is a shallow basin, -// so assuming 8x or 32x here instead of 16x costs about a fifth of the benefit -// and still beats weighting either phase alone. It is deliberately not derived -// from alpha/beta/dimension, which would be a model of summarize() that this -// does not need to be right about. +// so assuming 8x or 32x here instead of 16x costs a fraction of the benefit and +// still beats weighting either phase alone. It is deliberately not derived from +// alpha/beta/dimension, which would be a model of summarize() that this does not +// need to be right about. constexpr size_t kClusterCostRatio = 16; // Cuts [0, dimension) into at most `batches` windows of near-equal estimated // memory, from the exact per-term counts. // -// Equal width would not do, because term frequencies are heavily skewed: on -// msmarco base_full the heaviest term holds 5.7M postings against a mean of -// 37K, and the top 1% of terms hold 19% of them. Peak memory is set by the -// largest window, not the average one, so an uneven split wastes most of what -// batching could save. +// Equal width would not do, because term frequencies are heavily skewed: a +// natural-language corpus puts orders of magnitude more postings on its heaviest +// term than on its mean one. Peak memory is set by the largest window, not the +// average one, so an uneven split wastes most of what batching could save. // // What to even out is neither phase alone but their sum. Weighting raw counts // balances the fill and unbalances the clustering, which is the more expensive -// phase, and measured worse than equal width. Weighting min(count, lambda) -- -// what survives pruning, and so what clustering holds -- balances that phase -// perfectly but concentrates the heavy terms, leaving one window holding 3.7x -// the mean raw postings, which then becomes the peak. Weighting both together -// at their relative cost balances what is actually resident. Predicted peak of -// the largest window at 10 windows on base_full: -// -// equal width 2.76GB -// raw count 4.93GB -// min(count, lambda) 3.32GB -// both, as here 2.07GB +// phase. Weighting min(count, lambda) -- what survives pruning, and so what +// clustering holds -- balances that phase perfectly but concentrates the heavy +// terms, leaving one window holding several times the mean raw postings, which +// then becomes the peak. Weighting both together at their relative cost balances +// what is actually resident, and measures best of the four on a skewed corpus. // // Windows stay contiguous and ascending, which is what lets the clustered lists // be appended to a file as each window finishes: the layout carries no per-list From e69abbe41f053300d5769487d79fc732a87d8a75 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 3 Sep 2026 06:05:09 +0000 Subject: [PATCH 12/15] Spill windows to scratch instead of writing an index from build() batch_file_output_path is a directory the build may spill into, not an index to produce. Serializing an index is write_index's job, and a build that wrote one had two ways to make the same file -- the streamed writer here and each type's write_index -- which is a layout to keep in step for no benefit. So there is one mechanism now, the one the disk-resident pair already needed: cluster a window, spill it, drop it, and map the finished lists back out of the spill. Every type reaches it through MmapIndex::build_clustered_lists, which is also where the batching decision is made, so a type's build() is one call again. The spill is unlinked as soon as it is mapped, so nothing is left in the caller's directory and a crash mid-build cannot strand scratch; where a mapped file cannot be unlinked the index removes it when it goes. That deletes more than it adds: write_seismic_index_batched, the index-header exposure it needed (reverting 6affd30), the disk base's payload-header hook and mapping-slot parameter, and its second write of the payload. A window count with no directory to spill into now resolves to one window rather than being half-applied: it would bound the fill intermediate while the clustered lists accumulated for the whole corpus anyway, which is a corpus pass per window for a fraction of the peak. Unchanged: at a fixed seed the index write_index produces is byte-for-byte what a whole-corpus build would have produced, now asserted over all four types in one parametrized test, and build() still leaves an index that serves. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 107 +++-- benchmarks/batched_build_mem_bench.cpp | 49 +- nsparse/disk_seismic_index.cpp | 5 +- nsparse/disk_seismic_index_base.cpp | 111 +---- nsparse/disk_seismic_index_base.h | 60 +-- .../disk_seismic_scalar_quantized_index.cpp | 14 +- nsparse/disk_seismic_scalar_quantized_index.h | 2 - nsparse/io/index_io.cpp | 24 +- nsparse/io/index_io.h | 5 - nsparse/mmap_index.h | 61 ++- nsparse/seismic_batched_build.cpp | 109 ++--- nsparse/seismic_batched_build.h | 120 ++--- nsparse/seismic_common.cpp | 2 +- nsparse/seismic_common.h | 63 ++- nsparse/seismic_index.cpp | 27 +- nsparse/seismic_scalar_quantized_index.cpp | 31 +- python_tests/test_seismic_batched_build.py | 133 +++--- tests/seismic_batched_build_test.cpp | 451 ++++++++---------- 18 files changed, 596 insertions(+), 778 deletions(-) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 4f99b7f..4d6a93b 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -239,31 +239,34 @@ a separate entry point, so it is set in the factory description alongside | Option | Effect | |---|---| -| `inverted_list_batch_size=N` | Build in `N` term windows. Bounds the inverted-list intermediate to one window; the index is still built in memory as usual. | -| `batch_file_output_path=P` | With `N > 1`, write the index to `P` as it is built rather than assembling it in memory, so the clustered lists are never all resident either, then borrow them back from `P` by mapping it. Unused at `N <= 1`, which is an ordinary build and already holds its own lists. | - -`seismic` and `seismic_sq` end their payload with their posting lists, so each -window is serialized straight into `P` and dropped. The disk-resident pair cannot -be written that way: their summaries are followed by an inline forward index whose -blocks are laid out from the doc-id membership of *every* list, which is not known -until the last window is clustered. They spill the clustered lists to `P.lists` -instead, map them back, and write the payload from that mapping — same bound on -anonymous memory, at the cost of scratch disk the size of the lists. The spill is -deleted with the build; nothing else reads it. +| `inverted_list_batch_size=N` | Build in `N` term windows, bounding the inverted-list intermediate to one window. Ignored without `batch_file_output_path`: the clustered lists would accumulate for the whole corpus anyway, so the peak would barely move while the build paid a corpus pass per window. | +| `batch_file_output_path=P` | An existing directory the build may spill windows into. With `N > 1`, each window's clustered lists are written there and freed as they are produced, then borrowed back by mapping the spill, so the clustered lists are never all resident either. Unused at `N <= 1`, which is an ordinary build with nothing to spill. | + +So both together or neither: either knob alone leaves an ordinary whole-corpus +build. + +`P` is scratch, not output. `build()` writes no index and leaves nothing in that +directory — the spill is unlinked as soon as it is mapped, and the lists go on +being read from the mapping — so serializing an index is still `write_index`'s +job, and the index it writes is byte-for-byte what a whole-corpus build would +have produced. ```cpp auto* index = nsparse::index_factory( dimension, "seismic,lambda=6000|beta=400|alpha=0.4" - "|inverted_list_batch_size=10|batch_file_output_path=/data/index.dat"); + "|inverted_list_batch_size=10|batch_file_output_path=/scratch"); // Corpus residency is SparseVectors' business, not the build's: read_csr can // map a native-layout CSR instead of copying it, and the build is unchanged. index->read_csr("corpus.mcsr", nsparse::Residency::kMmap); -index->build(); // streams to /data/index.dat, then maps its lists back in +index->build(); // one window at a time, spilling to /scratch + +// Serialized the ordinary way, from the lists the build ended holding. +nsparse::write_index(index, "/data/index.dat"); -// Ready to serve, with no reopening by path: the posting lists are borrowed from -// the file just written, and the corpus is still borrowed from its own mapping. +// And servable as it stands: the posting lists are borrowed from the spill's +// mapping, and the corpus is still borrowed from its own. index->search(...); ``` @@ -275,20 +278,25 @@ nsparse.convert("corpus.csr", native) index = nsparse.index_factory( dim, "seismic,lambda=6000|beta=400|alpha=0.4" - "|inverted_list_batch_size=10|batch_file_output_path=/data/index.dat", + "|inverted_list_batch_size=10|batch_file_output_path=/scratch", ) index.read_csr(native, nsparse.Residency_kMmap) -index.build() # streams out, then maps its lists back in +index.build() # spills to /scratch +nsparse.write_index(index, "/data/index.dat") # the index file, as usual dists, labels = index.search(n, indptr, indices, values, k) ``` -The file is an ordinary index of its type — byte-for-byte what `write_index` -would have produced from the equivalent whole-corpus build. That is asserted -rather than assumed: at a fixed `seed` the two are compared as files, for all four +That the two are the same index is asserted rather than assumed: at a fixed +`seed` a batched build and a whole-corpus one are compared as files, for all four types. Each posting list's k-means seed comes from its own *global* term id and `lambda`/`beta` are resolved once from the whole corpus, so the window count cannot change what is produced. +What the spill does cost is disk while the index lives: an unlinked file still +occupies its blocks until the last mapping of it goes, so budget the clustered +lists' size on that filesystem for as long as the index object is around, on top +of whatever `write_index` then writes. + ### Choosing `inverted_list_batch_size` Windows are cut to equal estimated *memory*, not equal width. Term frequencies are @@ -301,12 +309,16 @@ every posting of its terms; clustering it holds what survives pruning bulkier per posting. Weighting either phase alone unbalances the other, both worse than weighting their sum — see `make_windows` in `nsparse/seismic_common.cpp`. +That ratio is also why the window count does nothing on its own: it bounds the +fill, and the clustered lists — the bulkier peak — are what a spill directory +bounds. + `RssAnon` is the figure to watch, being what the process itself allocated; `RssFile` is pages it touched of a mapping, which the kernel can reclaim under pressure. Read as total RSS the win looks far smaller than it is: an index that -maps its corpus, and then maps its own output back, keeps most of its residency in -page cache, which total RSS counts and pressure reclaims. Batching moves the -anonymous column, so report the split rather than the total. +maps its corpus, and then maps its spill back, keeps most of its residency in page +cache, which total RSS counts and pressure reclaims. Batching moves the anonymous +column, so report the split rather than the total. What to expect from the shape of it: anonymous memory falls faster than 1/N, because the split comes from the real per-term costs rather than from term ids, @@ -315,23 +327,19 @@ makes its own pass over the corpus. Ten to twenty windows is usually the useful range — most of the memory saving for a few percent of build time. The floor is what the build cannot batch: one window plus whatever the corpus itself costs. -The disk-resident pair behaves the same way, with two differences. Build time -climbs sooner, since the spill is written and read back and the payload write is -another pass over the lists; and `lambda` is the knob to watch for output size, -because their inline forward index copies every pruned posting's whole doc -vector — at a large `lambda` that section alone can dwarf the rest of the index, -and the spill needs scratch disk beside it. `disk_seismic_sq` cannot map a float -CSR (it searches over codes), so its corpus stays on the heap and shows up in the -same column the build's own growth does; subtract the reported -`start_rss_anon_mb`. - -Mapping the finished lists back in is nearly free in the column that matters: -borrowing an index costs a fraction of a second and single-digit megabytes of -anonymous memory, against copying it, which costs its whole size — the cursor -reads the size header before each array and skips the bulk, so it faults in part -of the file as reclaimable page cache. That mapping is separate from the corpus's: -an index that mapped its corpus with `read_csr` keeps doing so, since it still -scores from it. +All four types behave alike here, since they share the build. What differs is what +`write_index` then costs them: the disk-resident pair's inline forward index copies +every pruned posting's whole doc vector, so at a large `lambda` that section alone +can dwarf the rest of the index. And `disk_seismic_sq` cannot map a float CSR (it +searches over codes), so its corpus stays on the heap and shows up in the same +column the build's own growth does; subtract the reported `start_rss_anon_mb`. + +Borrowing the lists back from the spill is nearly free in the column that matters: +it costs a fraction of a second and single-digit megabytes of anonymous memory, +against copying them, which costs their whole size — the cursor reads the size +header before each array and skips the bulk, so it faults in part of the file as +reclaimable page cache. That mapping is separate from the corpus's: an index that +mapped its corpus with `read_csr` keeps doing so, since it still scores from it. Compare like with like. A build that loads the corpus is not comparable to one that maps it, and the difference is more than the corpus: a streaming ingest @@ -347,29 +355,30 @@ Query performance does not move, because the index is the same index: identical byte for byte at a fixed seed, and indistinguishable in latency, QPS and recall for the random-seeded default. -Write the index to a real disk. On a tmpfs such as `/tmp` it is RAM, which -defeats the point. +Point the spill directory at a real disk. On a tmpfs such as `/tmp` it is RAM, +which defeats the point. ### Measuring it `benchmarks/batched_build_mem_bench` reports peak RSS (`VmHWM`) and wall time for one configuration per process — `google-benchmark` measures throughput, and a -high-water mark is only clean in a process that has built nothing else. Use -`inmem` to compare against `baseline`: both then hold the corpus on the heap, so -the difference is the batching rather than the residency. +high-water mark is only clean in a process that has built nothing else. It times +`build()` alone; the index is written afterwards, outside the measurement, only to +report its size. Use `inmem` to compare against `baseline`: both then hold the +corpus on the heap, so the difference is the batching rather than the residency. ```bash cmake -S . -B build -DNSPARSE_ENABLE_BENCHMARKS=ON && cmake --build build -j B=./build/benchmarks/batched_build_mem_bench $B convert corpus.csr corpus.mcsr -$B baseline corpus.csr 6000 400 0.4 # whole-corpus build, for reference -$B batched inmem corpus.csr 6000 400 0.4 10 /data # 10 windows, same corpus residency -$B batched mmap corpus.mcsr 6000 400 0.4 10 /data # ... or with the corpus mapped +$B baseline corpus.csr 6000 400 0.4 # whole-corpus build, for reference +$B batched inmem corpus.csr 6000 400 0.4 10 /scratch # 10 windows, same corpus residency +$B batched mmap corpus.mcsr 6000 400 0.4 10 /scratch # ... or with the corpus mapped # Any type in the family, as its factory name. Both arms need it, and the # baseline's index path is positional -- pass "" to skip writing one. $B baseline corpus.csr 6000 400 0.4 "" disk_seismic -$B batched mmap corpus.mcsr 6000 400 0.4 10 /data disk_seismic +$B batched mmap corpus.mcsr 6000 400 0.4 10 /scratch disk_seismic ``` ## Python Bindings diff --git a/benchmarks/batched_build_mem_bench.cpp b/benchmarks/batched_build_mem_bench.cpp index d3f5503..efc7184 100644 --- a/benchmarks/batched_build_mem_bench.cpp +++ b/benchmarks/batched_build_mem_bench.cpp @@ -17,19 +17,18 @@ // batched_build_mem_bench baseline \ // [out_index] [index_type] // batched_build_mem_bench batched \ -// [index_type] +// [index_type] // // "convert" produces the native CSR the mapped read wants. "baseline" builds // the index whole (streaming add of the interchange CSR, like the other // benchmarks) -- the memory this feature exists to avoid. "batched" runs the -// term-batched build at the given batch count. +// term-batched build at the given batch count, spilling its windows into +// ; only the build is measured, and the index is written +// afterwards only to report its size. // -// `index_type` is any seismic-family factory name, default "seismic". The two -// disk-resident ones are worth measuring separately: they cannot stream their -// payload out window by window, so their batched build spills the clustered -// lists and writes from that mapping instead (see build_streamed), and whether -// that holds anonymous memory down is exactly what this reports. A quantized -// type gets the factory's default range, which is fine for a memory figure. +// `index_type` is any seismic-family factory name, default "seismic". A +// quantized type gets the factory's default range, which is fine for a memory +// figure. // // Compare "batched inmem" against "baseline": both hold the corpus on the heap // via the same streaming_add, so the difference between them is the batching @@ -38,8 +37,8 @@ // from the residency rather than from batching, so it is not the baseline's // counterpart. // -// Point at a real disk: on a tmpfs such as /tmp the index is RAM, and -// the numbers are meaningless. +// Point at a real disk: on a tmpfs such as /tmp the spill is RAM, +// and the numbers are meaningless. #include #include @@ -255,18 +254,19 @@ void streaming_add(nsparse::Index* index, const std::string& path) { // The index under test, named rather than constructed, so every type in the // family is reachable from the command line. The cluster knobs go through the // factory description, which is also how a caller sets them (see -// parse_cluster_params); `out_path` empty leaves the batched output path unset. +// parse_cluster_params); `scratch_dir` empty leaves the spill directory unset, +// which is what makes it an unbatched build. std::unique_ptr make_index( const std::string& index_type, int dimension, const nsparse::SeismicClusterParameters& params, - const std::string& out_path) { + const std::string& scratch_dir) { std::string desc = index_type + ",lambda=" + std::to_string(params.lambda) + "|beta=" + std::to_string(params.beta) + "|alpha=" + std::to_string(params.alpha) + "|inverted_list_batch_size=" + std::to_string(params.batch_clustering.batch_size); - if (!out_path.empty()) { - desc += "|batch_file_output_path=" + out_path; + if (!scratch_dir.empty()) { + desc += "|batch_file_output_path=" + scratch_dir; } return std::unique_ptr( nsparse::index_factory(dimension, desc.c_str())); @@ -358,7 +358,7 @@ int run_baseline(int argc, char** argv) { int run_batched(int argc, char** argv) { if (argc < 9) { std::cerr << "batched " - " [index_type]\n"; + " [index_type]\n"; return 2; } const std::string corpus_residency = argv[2]; @@ -368,8 +368,8 @@ int run_batched(int argc, char** argv) { .lambda = std::atoi(argv[4]), .beta = std::atoi(argv[5]), .alpha = static_cast(std::atof(argv[6]))}; - const std::string out = - std::string(argv[8]) + "/index." + index_type + ".dat"; + const std::string scratch_dir = argv[8]; + const std::string out = scratch_dir + "/index." + index_type + ".dat"; batched_params.batch_clustering.batch_size = static_cast(std::atoi(argv[7])); @@ -382,7 +382,7 @@ int run_batched(int argc, char** argv) { // not comparable to the baseline, because the saving is the // residency rather than the batching. std::unique_ptr index = - make_index(index_type, csr_dimension(csr), batched_params, out); + make_index(index_type, csr_dimension(csr), batched_params, scratch_dir); if (corpus_residency == "inmem") { streaming_add(index.get(), csr); } else if (corpus_residency == "mmap") { @@ -397,9 +397,9 @@ int run_batched(int argc, char** argv) { reset_vm_hwm(); PeakRssSampler sampler; const double started = now_seconds(); - // batch_file_output_path is set, so build() writes the index out as it goes - // rather than assembling it in memory -- the same call an ordinary build - // makes. + // batch_file_output_path is set, so build() clusters into windows and + // spills them there rather than holding them all -- the same call an + // ordinary build makes. index->build(); const double build_s = now_seconds() - started; @@ -407,8 +407,11 @@ int run_batched(int argc, char** argv) { "type=" + index_type + " corpus=" + corpus_residency + " batches=" + std::to_string(batched_params.batch_clustering.batch_size), build_s, load_hwm, start_anon, sampler); - // One window writes no file: it is an ordinary build, holding its own - // lists. + + // Written after the peaks are read, and only to report the size: build() + // produces no index file, so serializing one is the caller's step and not + // part of what is being measured. + nsparse::write_index(index.get(), const_cast(out.c_str())); if (std::ifstream file(out, std::ios::binary | std::ios::ate); file) { std::cout << "index_bytes=" << file.tellg() << "\n"; } diff --git a/nsparse/disk_seismic_index.cpp b/nsparse/disk_seismic_index.cpp index 41bbe02..e375cc4 100644 --- a/nsparse/disk_seismic_index.cpp +++ b/nsparse/disk_seismic_index.cpp @@ -57,10 +57,7 @@ DiskSeismicIndex* DiskSeismicIndex::mmap_index(const IndexHeader& header, MmapCursor cursor(mmap_file.data(), mmap_file.size()); cursor.skip(pos); - // No extra header for the float index, so the hook is a no-op and the - // shared payload follows directly; called anyway, so the two disk types - // read their payload through the same two steps. - index->read_mapped_payload_header(&cursor); + // No extra header for the float index; the shared payload follows directly. index->load_mapped_payload(&cursor, std::move(mmap_file)); return index.release(); } diff --git a/nsparse/disk_seismic_index_base.cpp b/nsparse/disk_seismic_index_base.cpp index 3590543..09182b1 100644 --- a/nsparse/disk_seismic_index_base.cpp +++ b/nsparse/disk_seismic_index_base.cpp @@ -12,11 +12,9 @@ #include #include #include -#include #include #include #include -#include #include #include @@ -25,11 +23,8 @@ #include "nsparse/disk_seismic_search.h" #include "nsparse/id_selector.h" #include "nsparse/index.h" -#include "nsparse/io/file_io.h" -#include "nsparse/io/index_io.h" #include "nsparse/io/inline_forward_index_io.h" #include "nsparse/io/seismic_invlists_writer.h" -#include "nsparse/seismic_batched_build.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" #include "nsparse/utils/checks.h" @@ -37,35 +32,6 @@ #include "nsparse/utils/mmap_file.h" namespace nsparse { -namespace { - -// Where a batched build spills its clustered lists: alongside the index it is -// writing, so it lands on whatever disk the caller chose for the output. -constexpr const char* kSpillSuffix = ".lists"; - -// Deletes the spill on the way out, whether the build finished or threw: it is -// scratch the size of the clustered lists, and nothing outside build() knows it -// exists. On the success path the build has already displaced the mapping of -// it, which matters on Windows, where a mapped file cannot be unlinked. -class SpillFile { -public: - explicit SpillFile(std::string path) : path_(std::move(path)) {} - ~SpillFile() { - std::error_code ignored; - std::filesystem::remove(path_, ignored); - } - SpillFile(const SpillFile&) = delete; - SpillFile& operator=(const SpillFile&) = delete; - SpillFile(SpillFile&&) = delete; - SpillFile& operator=(SpillFile&&) = delete; - - [[nodiscard]] const std::string& path() const { return path_; } - -private: - std::string path_; -}; - -} // namespace DiskSeismicIndexBase::DiskSeismicIndexBase(int dim, SeismicClusterParameters parameter) @@ -94,61 +60,10 @@ void DiskSeismicIndexBase::add(idx_t n, const idx_t* indptr, } void DiskSeismicIndexBase::build() { - const SparseVectorsConfig config = { - .element_size = code_element_size(), - .dimension = static_cast(get_dimension())}; - const auto& batch = cluster_parameter_.batch_clustering; - if (batch.batch_size > 1 && !batch.batch_file_output_path.empty()) { - build_streamed(config, batch.batch_file_output_path); - return; - } - // A single window is an ordinary build: it holds its own lists, and writing - // them out only to map them back would be work for nothing. - clustered_inverted_lists = detail::build_inverted_lists_clusters( - get_vectors(), config, cluster_parameter_); -} - -void DiskSeismicIndexBase::build_streamed(const SparseVectorsConfig& config, - const std::string& out_path) { - // This payload cannot be streamed section by section the way the in-memory - // types' can. Theirs ends with its posting lists, so a window can be - // serialized and dropped; here the summaries are followed by an inline - // forward index whose blocks are laid out from the doc-id membership of - // every list, and that membership is not known until the last window is - // clustered. - // - // So the clustering still runs once, a window at a time, but into a spill; - // the lists come back borrowed from it, and the payload is written from - // that mapping. Neither phase holds more than one window of anonymous - // memory -- the forward index streams its blocks out as it lays them. - const SpillFile spill(out_path + kSpillSuffix); - clustered_inverted_lists = - detail::spill_clustered_lists(get_vectors(), config, cluster_parameter_, - spill.path(), &batch_mapped_file_); - - size_t payload_offset = 0; - { - FileIOWriter writer(const_cast(out_path.c_str())); - detail::write_header({.id = fourcc(id()), - .version = format_version(), - .dimension = get_dimension()}, - &writer); - // Taken from the writer rather than from a constant, so the offset the - // mapping below skips to is the one the header actually occupied. - payload_offset = writer.pos(); - write_index(&writer); - writer.close(); - } - - // Ends holding the index it wrote, rather than an object whose lists point - // into scratch that is about to be deleted. The output's mapping displaces - // the spill's, which is safe in that order: load_mapped_payload replaces - // the lists that borrowed from the spill before it commits the mapping. - MmapFile mapped(out_path); - MmapCursor cursor(mapped.data(), mapped.size()); - cursor.skip(payload_offset); - read_mapped_payload_header(&cursor); - load_mapped_payload(&cursor, std::move(mapped), &batch_mapped_file_); + clustered_inverted_lists = build_clustered_lists( + {.element_size = code_element_size(), + .dimension = static_cast(get_dimension())}, + cluster_parameter_); } auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr, @@ -246,8 +161,9 @@ void DiskSeismicIndexBase::write_index(IOWriter* 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 = code_element_size(), - .dimension = static_cast(dimension_)}); + SparseVectors empty_vectors( + {.element_size = code_element_size(), + .dimension = static_cast(dimension_)}); const SparseVectors& v = vectors_ != nullptr ? *vectors_ : empty_vectors; detail::InlineForwardIndex forward(clustered_inverted_lists, v); forward.serialize(io_writer); @@ -264,8 +180,7 @@ void DiskSeismicIndexBase::read_index(IOReader* /*io_reader*/, } void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, - MmapFile&& mapped, - MmapFile* slot) { + MmapFile&& mapped) { // Same order write_index wrote them (past any extra header the caller // already consumed): doc count, summaries, inline forward. num_vectors_ = cursor->read_scalar(); @@ -274,18 +189,16 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, detail::InlineForwardIndex forward; forward.mmap_deserialize(cursor); - // Before the mapping is committed below: whatever these replace may have - // been borrowing from what `slot` still holds -- a batched build's spill. clustered_inverted_lists = std::move(inv_list_writer.release()); fwd_ = std::move(forward); // 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. + // from `mapped`, which is alive here), let the concrete index reject a width + // mismatch before we commit. validate_mapped_payload(); - // The mapping last: the summaries and the forward index borrow from it, and + // mapped_file_ last: the summaries and the forward index borrow from it, and // moving it does not move the mapping. - *slot = std::move(mapped); + mapped_file_ = std::move(mapped); } } // namespace nsparse diff --git a/nsparse/disk_seismic_index_base.h b/nsparse/disk_seismic_index_base.h index 6224a0e..be2f53f 100644 --- a/nsparse/disk_seismic_index_base.h +++ b/nsparse/disk_seismic_index_base.h @@ -12,8 +12,6 @@ #include #include -#include -#include #include #include "nsparse/cluster/inverted_list_clusters.h" @@ -30,12 +28,11 @@ namespace nsparse { // 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, -// and search scores the global top-k_prime blocks (a -// DiskSeismicSearchParameters sets k_prime). add / build / search / -// serialization / mmap loading are all here; the concrete indexes differ only -// in the stored value width and, for the quantized one, a leading quantization -// header and a score-decoding step, which they supply through the virtual hooks -// below. +// and search scores the global top-k_prime blocks (a DiskSeismicSearchParameters +// sets k_prime). add / build / search / serialization / mmap loading are all +// here; the concrete indexes differ only in the stored value width and, for the +// quantized one, a leading quantization header and a score-decoding step, which +// they supply through the virtual hooks below. // // mmap-only: load with read_index(file, kUseMmap); the copying read throws. class DiskSeismicIndexBase : public MmapIndex, public IndexIO { @@ -74,12 +71,10 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO { // code_element_size()-byte-per-value data. `scratch` backs the result when // encoding must allocate; the float index returns its input reinterpreted, // with no copy. - virtual const uint8_t* encode_values( - const float* values, size_t nnz, - std::vector& scratch) const = 0; + virtual const uint8_t* encode_values(const float* values, size_t nnz, + std::vector& scratch) const = 0; - // Encode a query batch the same way, honoring any per-search range - // override. + // Encode a query batch the same way, honoring any per-search range override. virtual const uint8_t* encode_query( const float* values, size_t nnz, const SearchParameters* search_parameters, @@ -94,45 +89,22 @@ class DiskSeismicIndexBase : public MmapIndex, public IndexIO { // quantization header for the quantized index; nothing for the float one). virtual void write_payload_header(IOWriter* /*io_writer*/) const {} - // The mirror of write_payload_header: consume that header off a mapping and - // adopt what it declares. Used by the mapped read and by a batched build - // reopening the file it just wrote, so the two cannot read the payload from - // different offsets. - virtual void read_mapped_payload_header(MmapCursor* /*cursor*/) {} - - // Reject a just-mapped payload whose stored width disagrees with this - // index. No-op for the float index, which stores a fixed width. Runs after - // the summaries and forward index are populated but before the mapping - // commits. + // Reject a just-mapped payload whose stored width disagrees with this index. + // No-op for the float index, which stores a fixed width. Runs after the + // summaries and forward index are populated but before the mapping commits. virtual void validate_mapped_payload() const {} // Reads the shared payload (doc count, summaries, inline forward) from the - // cursor, validates it, and commits the mapping into `slot`. Each concrete - // mmap_index reads its own extra header first, then calls this. - // - // `slot` is mapped_file_ for a read_index load, which owns nothing else, - // but batch_mapped_file_ for a batched build: there the corpus may still be - // borrowing from mapped_file_, and giving that up would leave the index - // unable to score anything. - void load_mapped_payload(MmapCursor* cursor, MmapFile&& mapped, - MmapFile* slot); - void load_mapped_payload(MmapCursor* cursor, MmapFile&& mapped) { - load_mapped_payload(cursor, std::move(mapped), &mapped_file_); - } + // cursor, validates it, and commits the mapping. Each concrete mmap_index + // reads its own extra header first, then calls this. + void load_mapped_payload(MmapCursor* cursor, MmapFile&& mapped); - // Borrowed from by score_summaries_transposed / the inline forward index, - // so the concrete validate_mapped_payload can inspect their widths. + // Borrowed from by score_summaries_transposed / the inline forward index, so + // the concrete validate_mapped_payload can inspect their widths. std::vector clustered_inverted_lists; detail::InlineForwardIndex fwd_; private: - // build() with batch_file_output_path set: clusters one term window at a - // time into a spill file, writes the index out of that mapping, and ends - // borrowing its own output. See the definition for why this payload cannot - // be streamed section by section the way the in-memory ones are. - void build_streamed(const SparseVectorsConfig& config, - const std::string& out_path); - auto search(idx_t n, const idx_t* indptr, const term_t* indices, const float* values, int k, SearchParameters* search_parameters = nullptr) diff --git a/nsparse/disk_seismic_scalar_quantized_index.cpp b/nsparse/disk_seismic_scalar_quantized_index.cpp index 4aaf117..7eb4201 100644 --- a/nsparse/disk_seismic_scalar_quantized_index.cpp +++ b/nsparse/disk_seismic_scalar_quantized_index.cpp @@ -142,14 +142,6 @@ void DiskSeismicScalarQuantizedIndex::write_payload_header( io_writer->write(&vmax, sizeof(float), 1); } -void DiskSeismicScalarQuantizedIndex::read_mapped_payload_header( - MmapCursor* cursor) { - const auto sq_type = cursor->read_scalar(); - const auto vmin = cursor->read_scalar(); - const auto vmax = cursor->read_scalar(); - sq_ = make_scalar_quantizer(sq_type, vmin, vmax); -} - void DiskSeismicScalarQuantizedIndex::validate_mapped_payload() const { throw_if_forward_width_mismatch(fwd_, sq_); throw_if_summary_width_mismatch(clustered_inverted_lists, sq_); @@ -167,7 +159,11 @@ DiskSeismicScalarQuantizedIndex* DiskSeismicScalarQuantizedIndex::mmap_index( // The quantization header opens the payload; the shared loader reads the // rest and calls validate_mapped_payload() against this quantizer. - index->read_mapped_payload_header(&cursor); + const auto sq_type = cursor.read_scalar(); + const auto vmin = cursor.read_scalar(); + const auto vmax = cursor.read_scalar(); + index->sq_ = make_scalar_quantizer(sq_type, vmin, vmax); + index->load_mapped_payload(&cursor, std::move(mmap_file)); return index.release(); } diff --git a/nsparse/disk_seismic_scalar_quantized_index.h b/nsparse/disk_seismic_scalar_quantized_index.h index f4c403d..fd344da 100644 --- a/nsparse/disk_seismic_scalar_quantized_index.h +++ b/nsparse/disk_seismic_scalar_quantized_index.h @@ -90,8 +90,6 @@ class DiskSeismicScalarQuantizedIndex : public DiskSeismicIndexBase { // The quantization parameters that open this index's payload -- distinct // from the IndexHeader the file itself starts with. void write_payload_header(IOWriter* io_writer) const override; - // Adopts the quantizer a mapped payload declares, rejecting an unknown type. - void read_mapped_payload_header(MmapCursor* cursor) override; void validate_mapped_payload() const override; // The quantizer a query is encoded with: DiskSeismicSQSearchParameters diff --git a/nsparse/io/index_io.cpp b/nsparse/io/index_io.cpp index aea89ec..d1ad869 100644 --- a/nsparse/io/index_io.cpp +++ b/nsparse/io/index_io.cpp @@ -111,6 +111,18 @@ std::string id_to_string(uint32_t id_val) { return chars; } +void write_header(const IndexHeader& header, IOWriter* io_writer) { + // write index type + uint32_t id_val = header.id; + io_writer->write(&id_val, sizeof(uint32_t), 1); + // write payload layout version + uint32_t version = header.version; + io_writer->write(&version, sizeof(uint32_t), 1); + // write dimension + int dimension = header.dimension; + io_writer->write(&dimension, sizeof(int), 1); +} + IndexHeader read_header(IOReader* io_reader) { IndexHeader header; io_reader->read(&header.id, sizeof(uint32_t), 1); @@ -159,18 +171,6 @@ void throw_if_version_unsupported(const IndexHeader& header, } // namespace namespace detail { -void write_header(const IndexHeader& header, IOWriter* io_writer) { - // write index type - uint32_t id_val = header.id; - io_writer->write(&id_val, sizeof(uint32_t), 1); - // write payload layout version - uint32_t version = header.version; - io_writer->write(&version, sizeof(uint32_t), 1); - // write dimension - int dimension = header.dimension; - io_writer->write(&dimension, sizeof(int), 1); -} - void write_index(Index* index, IOWriter* io_writer, bool keep_open) { auto* index_io = dynamic_cast(index); StreamCloser closer(io_writer, keep_open); diff --git a/nsparse/io/index_io.h b/nsparse/io/index_io.h index 519194a..f1a8159 100644 --- a/nsparse/io/index_io.h +++ b/nsparse/io/index_io.h @@ -19,11 +19,6 @@ enum IndexIoFlag { }; namespace detail { -// The fixed prefix every serialized index starts with. Exposed because a writer -// that streams a payload out itself, rather than through an Index, still has to -// lay the header out exactly the way read_header expects — see -// build_seismic_index_batched. -void write_header(const IndexHeader& header, IOWriter* io_writer); void write_index(Index* index, IOWriter* io_writer, bool keep_open); // `filename`, when given, lets an index that was written for mmap borrow from // the file instead of copying; without one the copying path is used, since a diff --git a/nsparse/mmap_index.h b/nsparse/mmap_index.h index 72d46a3..c800cd3 100644 --- a/nsparse/mmap_index.h +++ b/nsparse/mmap_index.h @@ -17,8 +17,13 @@ #include #include #include +#include +#include +#include +#include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/index.h" +#include "nsparse/seismic_batched_build.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/utils/checks.h" @@ -30,6 +35,28 @@ class MmapIndex : public Index { public: explicit MmapIndex(int dim = 0) : Index(dim) {} + ~MmapIndex() override { + // A batched build's spill outlives build(), because the posting lists + // borrow from it, so removing it falls to whoever holds them. Only + // reached where a mapped file cannot be unlinked (Windows); elsewhere + // spill_clustered_lists has already unlinked it and left this empty. + if (batch_scratch_path_.empty()) { + return; + } + // Unmapped here rather than left to the member's own destructor, which + // runs after this body: the file cannot go while it is still mapped. + // Whatever borrowed from it lives in a derived class, already + // destroyed. + batch_mapped_file_ = MmapFile{}; + std::error_code ignored; + std::filesystem::remove(batch_scratch_path_, ignored); + } + + MmapIndex(const MmapIndex&) = delete; + MmapIndex& operator=(const MmapIndex&) = delete; + MmapIndex(MmapIndex&&) = delete; + MmapIndex& operator=(MmapIndex&&) = delete; + void read_csr(const char* file_path, Residency residency = Residency::kInMemory) override { switch (residency) { @@ -49,6 +76,31 @@ class MmapIndex : public Index { } protected: + // The build every seismic-family type runs, so the batching decision is + // made once rather than per type. + // + // With both batch knobs set (see BatchClusteringOption) it clusters one + // term window at a time, spilling to scratch and borrowing the finished + // lists back from it, so the peak is one window rather than the whole + // corpus; otherwise it is the ordinary whole-corpus build. Either way the + // lists come back complete and in term order, and writing an index file + // stays write_index's job -- the spill is not one, and build() produces no + // file a caller keeps. + std::vector build_clustered_lists( + const SparseVectorsConfig& config, + const SeismicClusterParameters& params) { + if (params.batch_clustering.effective_batch_size() <= 1) { + return detail::build_inverted_lists_clusters(get_vectors(), config, + params); + } + detail::SpilledLists spilled = detail::spill_clustered_lists( + get_vectors(), config, params, + params.batch_clustering.batch_file_output_path, + &batch_mapped_file_); + batch_scratch_path_ = std::move(spilled.scratch_path); + return std::move(spilled.lists); + } + // The mapping borrowed buffers point into: a native CSR file via read_csr, // or a serialized index file. Those sources are mutually exclusive, so one // member serves both. @@ -62,8 +114,8 @@ class MmapIndex : public Index { // mapped_file_ when mapped. get_vectors() cannot tell the two apart. std::unique_ptr vectors_; - // A second mapping, for the file a batched build streams itself to and then - // borrows its posting lists back from. Separate from mapped_file_ rather + // A second mapping, for the spill a batched build wrote its clustered lists + // to and then borrows them back from. Separate from mapped_file_ rather // than replacing it, because the two coexist: the corpus may itself be a // mapping that vectors_ is still borrowing from, and giving that up would // leave the index unable to score anything. @@ -73,6 +125,11 @@ class MmapIndex : public Index { // first. MmapFile batch_mapped_file_; + // The spill behind batch_mapped_file_, when the platform would not let it + // be unlinked while mapped. Empty otherwise, which is the usual case. See + // the destructor. + std::string batch_scratch_path_; + private: // Values are borrowed at their stored width, so a quantizing index cannot // use this path. diff --git a/nsparse/seismic_batched_build.cpp b/nsparse/seismic_batched_build.cpp index 2fdc2e7..95a1df3 100644 --- a/nsparse/seismic_batched_build.cpp +++ b/nsparse/seismic_batched_build.cpp @@ -10,15 +10,16 @@ #include "nsparse/seismic_batched_build.h" #include -#include +#include +#include #include #include +#include #include #include #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/io/file_io.h" -#include "nsparse/io/index_io.h" #include "nsparse/io/io.h" #include "nsparse/io/seismic_invlists_writer.h" #include "nsparse/seismic_common.h" @@ -29,30 +30,26 @@ namespace nsparse::detail { namespace { -// Both entry points below refuse the same two inputs: nowhere to write to, and -// a corpus with no postings to write. -void throw_if_not_streamable(const SparseVectors* vectors, - const std::string& path, const char* who) { - if (path.empty()) { - throw std::invalid_argument(std::string(who) + - ": output path must not be empty"); - } - if (vectors == nullptr || vectors->num_vectors() == 0) { - throw std::invalid_argument( - std::string(who) + ": corpus is empty; there is nothing to stream"); - } +// A spill file of this build's own inside `dir`. Named uniquely rather than +// fixed, so concurrent builds sharing a scratch directory cannot overwrite each +// other's windows. +std::string scratch_file_path(const std::string& dir) { + const auto token = static_cast(std::random_device{}()) << 32U | + std::random_device{}(); + return (std::filesystem::path(dir) / + ("nsparse-clustered-lists-" + std::to_string(token) + ".tmp")) + .string(); } // The posting-list section -- [count][list...], the layout // SeismicInvertedListsWriter produces -- streamed into `writer` one window at a -// time, starting wherever the writer has reached. +// time. // -// The writer is the one the whole file is being written through, rather than a -// per-window one whose output is concatenated: serialize() pads each array +// One writer for the whole section, windows serialized straight into it rather +// than written separately and concatenated: serialize() pads each array // relative to the writer's current offset (see io/align.h), so bytes produced // by a writer that started at 0 carry the wrong padding once appended at some -// other offset. Streaming through a single writer keeps pos() the true absolute -// offset. +// other offset. void stream_clustered_lists(const SparseVectors* vectors, const SparseVectorsConfig& config, const SeismicClusterParameters& params, @@ -73,7 +70,7 @@ void stream_clustered_lists(const SparseVectors* vectors, // The layout carries no per-list offsets, so a gap or a repeat // would silently shift every list after it. throw std::runtime_error( - "stream_clustered_lists: windows arrived out of order"); + "spill_clustered_lists: windows arrived out of order"); } for (const auto& list : clusters) { list.serialize(writer); @@ -83,35 +80,31 @@ void stream_clustered_lists(const SparseVectors* vectors, }); if (next_term != config.dimension) { throw std::runtime_error( - "stream_clustered_lists: wrote " + std::to_string(next_term) + + "spill_clustered_lists: spilled " + std::to_string(next_term) + " of " + std::to_string(config.dimension) + " posting lists"); } } } // namespace -size_t write_seismic_index_batched( - const SparseVectors* vectors, const SparseVectorsConfig& config, - const SeismicClusterParameters& params, const IndexHeader& header, - const std::function& write_prefix, - const std::string& out_path) { - throw_if_not_streamable(vectors, out_path, "write_seismic_index_batched"); - - FileIOWriter writer(const_cast(out_path.c_str())); - write_header(header, &writer); - write_prefix(&writer); - - const size_t lists_offset = writer.pos(); - stream_clustered_lists(vectors, config, params, &writer); - writer.close(); - return lists_offset; -} +SpilledLists spill_clustered_lists(const SparseVectors* vectors, + const SparseVectorsConfig& config, + const SeismicClusterParameters& params, + const std::string& scratch_dir, + MmapFile* into) { + if (!std::filesystem::is_directory(scratch_dir)) { + throw std::invalid_argument( + "spill_clustered_lists: batch_file_output_path must be an existing " + "directory to spill into, got '" + + scratch_dir + "'"); + } + if (vectors == nullptr || vectors->num_vectors() == 0) { + throw std::invalid_argument( + "spill_clustered_lists: corpus is empty; there is nothing to " + "spill"); + } -std::vector spill_clustered_lists( - const SparseVectors* vectors, const SparseVectorsConfig& config, - const SeismicClusterParameters& params, const std::string& path, - MmapFile* into) { - throw_if_not_streamable(vectors, path, "spill_clustered_lists"); + const std::string path = scratch_file_path(scratch_dir); { // Closed before the mapping is taken: the writer buffers, and what is // not flushed is not in the file to map. @@ -119,28 +112,30 @@ std::vector spill_clustered_lists( stream_clustered_lists(vectors, config, params, &writer); writer.close(); } - // Offset 0: a spill is the section and nothing else, with no header for it - // to sit behind. - return map_streamed_lists(path, /*lists_offset=*/0, into); -} -std::vector map_streamed_lists(const std::string& path, - size_t lists_offset, - MmapFile* into) { MmapFile mapped(path); - // The cursor starts at 0 and skips, rather than mapping from lists_offset: - // absolute file offsets are what serialize() padded against, so a cursor - // that began part-way through would compute different padding and misread - // every array. Same reason mmap_index skips rather than offsets. + // The cursor starts at the section, which is the whole file: a spill has no + // header for it to sit behind. Absolute offsets are what serialize() padded + // against, and here they are the section's own. MmapCursor cursor(mapped.data(), mapped.size()); - cursor.skip(lists_offset); SeismicInvertedListsWriter lists; lists.mmap_deserialize(&cursor); - // Committed only once the walk succeeded, so a truncated file cannot leave - // the index holding lists that point into a mapping it never took. + // Unlinked now rather than when the index is done with it: the mapping + // keeps the bytes alive wherever unlinking an open file is allowed, so + // scratch cannot outlive the process even if it dies mid-build. Where it is + // not allowed, the path goes back to the caller to remove after the + // mapping. + std::error_code failed; + std::filesystem::remove(path, failed); + + SpilledLists spilled; + spilled.lists = std::move(lists.release()); + spilled.scratch_path = failed ? path : std::string(); + // Committed once the walk succeeded, so a truncated spill cannot leave the + // caller holding lists that point into a mapping it never took. *into = std::move(mapped); - return std::move(lists.release()); + return spilled; } } // namespace nsparse::detail diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h index 05f666c..32fc916 100644 --- a/nsparse/seismic_batched_build.h +++ b/nsparse/seismic_batched_build.h @@ -10,102 +10,62 @@ #ifndef SEISMIC_BATCHED_BUILD_H #define SEISMIC_BATCHED_BUILD_H -#include -#include #include #include #include "nsparse/cluster/inverted_list_clusters.h" -#include "nsparse/index.h" -#include "nsparse/io/io.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/utils/mmap_file.h" namespace nsparse::detail { -// Builds a seismic-family index and writes it straight to `out_path`, one term -// window at a time, without ever holding the whole index in memory. +// What a spilled build hands back. +struct SpilledLists { + // Every term's clustered posting list, in term order -- the same thing + // build_inverted_lists_clusters returns, borrowed from the spill's mapping + // rather than allocated. + std::vector lists; + // The spill file, when it is still on disk to be removed. Empty when it was + // unlinked as soon as it was mapped, which is what happens wherever a + // mapped file can be unlinked; where it cannot (Windows), whoever holds the + // mapping has to remove this after releasing it. + std::string scratch_path; +}; + +// Clusters the corpus one contiguous term window at a time, spilling each +// window's lists to a temporary file in `scratch_dir` and mapping them back, so +// the caller ends up holding every list without two windows ever having been +// resident at once. // // The usual build holds two whole-corpus intermediates -- the inverted lists -// and then the clustered posting lists -- so its peak memory scales with the -// corpus's non-zeros, and a corpus whose posting lists do not fit in RAM cannot -// be indexed at all. for_each_clustered_window bounds the first to one window; -// serializing each window and dropping it, which is what this does, bounds the -// second. What is left resident is the forward corpus (which the caller already -// holds, at whatever residency SparseVectors was given) plus one window. -// -// Reached through an index's build(), by setting -// SeismicClusterParameters::batch_clustering.batch_file_output_path. The build -// then maps the file back (see map_streamed_lists) rather than dropping it, so -// it ends holding the index it wrote. +// (every posting) and then the clustered posting lists -- so its peak memory +// scales with the corpus's non-zeros, and a corpus whose posting lists do not +// fit in RAM cannot be indexed at all. for_each_clustered_window bounds the +// first to one window; spilling each window's clusters and dropping them, which +// is what this does, bounds the second. What is left resident is the forward +// corpus (which the caller already holds, at whatever residency SparseVectors +// was given) plus one window. // -// `header` and `write_prefix` are what make this work for every type whose -// payload ends with its posting lists, rather than just SEIS. `write_prefix` -// writes whatever the type puts -// between the header and its posting lists -- the forward vectors, and for a -// quantizing index its quantization header first. The lists then follow in the -// byte-for-byte layout SeismicInvertedListsWriter produces, so the file is an -// ordinary index of that type: read it back with read_index, mapped or copying, -// exactly as if it had been built in memory and written with write_index. +// The spill is scratch, not an index: it carries the posting-list section and +// nothing else -- no index header, no forward vectors -- and nothing outside +// this build ever reads it. Writing an index file remains write_index's job, +// from the lists this returns; a build that borrows its lists from scratch +// serializes byte-for-byte what a whole-corpus build would have. // // Identical to the unbatched build for a fixed `params.seed`, and identical -// whatever batch_size is, because every list's k-means seed comes from its own -// global term id -- see for_each_clustered_window. -// -// Returns the absolute byte offset of the posting-list section, so the lists -// can be mapped back in without re-parsing everything before them -- see -// map_streamed_lists. -// -// Throws if the corpus is empty: there would be no windows to stream, and a -// header-only file is not a readable index. -size_t write_seismic_index_batched( - const SparseVectors* vectors, const SparseVectorsConfig& config, - const SeismicClusterParameters& params, const IndexHeader& header, - const std::function& write_prefix, - const std::string& out_path); - -// Streams every window's clustered posting lists to `path` and then maps them -// back, so the caller ends up holding all of them without two windows ever -// having been resident at once. -// -// For the index types whose payload ends with its posting lists, -// write_seismic_index_batched writes the index itself and there is nothing to -// spill. A DiskSeismic payload is not one of those: its summaries precede an -// inline forward index whose blocks are laid out from the doc-id membership of -// every list, so no window's lists can be dropped before the last window is -// clustered. What can be dropped is their *residency* -- which is what this is -// for. `path` gets the lists in the same [count][list...] layout -// SeismicInvertedListsWriter produces, doc ids included (an index's own section -// writes them empty; the forward index is what needs them here), and the -// returned lists borrow from the mapping handed to `into` rather than the heap. -// -// The spill is scratch, not an index: it carries no header, nothing else reads -// it, and deleting it is the caller's job. It must outlive the returned lists. -// -// Throws if the corpus is empty, for the same reason as -// write_seismic_index_batched: there would be no windows to stream. -std::vector spill_clustered_lists( - const SparseVectors* vectors, const SparseVectorsConfig& config, - const SeismicClusterParameters& params, const std::string& path, - MmapFile* into); - -// Maps the file a streamed build just wrote and borrows its posting lists out -// of it, so the build ends holding a usable index without ever having held all -// of the lists at once. -// -// Only the lists. The forward vectors in the file are a copy of ones the index -// already has, at whatever residency the caller chose for them, so re-reading -// them would be work for nothing -- and it is why the corpus mapping can be -// left alone rather than swapped out. `lists_offset` is what -// write_seismic_index_batched returned, which saves parsing past the vectors to -// find where the lists start. +// whatever the window count is, because every list's k-means seed comes from +// its own global term id -- see for_each_clustered_window. // -// The mapping is handed to `into`, which must outlive the returned lists: they -// point into it. -std::vector map_streamed_lists(const std::string& path, - size_t lists_offset, - MmapFile* into); +// `into` takes the mapping the returned lists borrow from, and so must outlive +// them. Throws if `scratch_dir` is not a directory, or if the corpus is empty: +// there would be no windows to spill, and an empty spill maps back to no lists +// at all. +SpilledLists spill_clustered_lists(const SparseVectors* vectors, + const SparseVectorsConfig& config, + const SeismicClusterParameters& params, + const std::string& scratch_dir, + MmapFile* into); } // namespace nsparse::detail diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp index a0333a8..9f447b4 100644 --- a/nsparse/seismic_common.cpp +++ b/nsparse/seismic_common.cpp @@ -288,7 +288,7 @@ void for_each_clustered_window(const SparseVectors* vectors, const size_t dim = config.dimension; const size_t batches = - std::max(1, std::min(params.batch_clustering.batch_size, dim)); + std::min(params.batch_clustering.effective_batch_size(), dim); const int lambda = calculate_lambda(params.lambda, vectors->num_vectors()); const ResolvedParameters resolved = { diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index 4d5401b..e0d2eed 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -10,12 +10,12 @@ #ifndef SEISMIC_COMMON_H #define SEISMIC_COMMON_H -#include #include +#include #include #include -#include #include +#include #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/cluster/random_kmeans.h" @@ -33,22 +33,41 @@ namespace nsparse { // corpus: the inverted lists (every posting) and then the clustered lists. Both // scale with the corpus's non-zeros, which is what puts a ceiling on the corpus // an index can be built from. Splitting the term space into `batch_size` -// contiguous windows and finishing one window before starting the next makes the -// first of those proportional to a window instead. +// contiguous windows and finishing one window before starting the next makes +// the first of those proportional to a window instead. // // `batch_file_output_path` bounds the second as well: with batch_size > 1, each -// window's clustered lists are serialized to that path and freed as they are -// produced, and the finished list section is then mapped back in, so the build -// never holds more than one window's worth and still ends with a usable index. -// See write_seismic_index_batched. +// window's clustered lists are spilled to a temporary file there and freed as +// they are produced, and the finished lists are then mapped back out of it, so +// the build never holds more than one window's worth. See +// spill_clustered_lists. +// +// The two knobs are therefore only useful together, and each is ignored without +// the other -- see effective_batch_size. struct BatchClusteringOption { // Contiguous term windows. <= 1 means one window, i.e. no batching. Clamped // to the dimension, since a window cannot be narrower than one term. size_t batch_size = 1; - // Where a batched build streams the index it produces. Used only when - // batch_size > 1: a single window is an ordinary build, which holds its own - // posting lists and so has nothing to stream or to map back. + // An existing directory the build may spill windows into. Scratch, not + // output: build() writes no index there and leaves nothing behind, and + // serializing an index remains write_index's job. Used only when batch_size + // > 1, a single window being an ordinary build with nothing to spill. std::string batch_file_output_path; + + // Windows the build should actually run, which is 1 unless there is + // somewhere to write them. + // + // Splitting the term space without a path to write to bounds only the first + // intermediate: the clustered lists still accumulate for the whole corpus, + // and they are the bulkier of the two by roughly kClusterCostRatio. So it + // buys a fraction of the peak while costing a corpus pass per window, and + // the peak stays where an unbatched build's is. Not worth doing quietly on + // a caller's behalf, so `batch_size` alone is honoured as no batching at + // all. + [[nodiscard]] size_t effective_batch_size() const { + return batch_file_output_path.empty() ? 1 + : std::max(1, batch_size); + } }; // Draw fresh entropy at build time, which makes the build unreproducible. Any @@ -74,8 +93,8 @@ constexpr float kDefaultBetaRatio = 0.1F; constexpr int kDefaultBeta = -1; constexpr float kDefaultAlpha = 0.4F; -// const rather than constexpr: BatchClusteringOption holds a std::string for the -// output path, which is not a literal type. +// const rather than constexpr: BatchClusteringOption holds a std::string for +// the output path, which is not a literal type. inline const SeismicClusterParameters kDefaultSeismicClusterParams = { .lambda = kDefaultLambda, .beta = kDefaultBeta, .alpha = kDefaultAlpha}; @@ -169,21 +188,22 @@ inline int calculate_beta(int beta, int lambda) { // // This is the one place the seismic family's build work lives: every index type // reaches it, either through build_inverted_lists_clusters below or through the -// streaming build in seismic_batched_build.h. The element width comes from +// spilling build in seismic_batched_build.h. The element width comes from // `config`, so a quantizing index gets the same treatment as a float one -- the // values in `vectors` are already encoded by the time they arrive here. // // `sink` receives the window's global first term and its lists, and must not // hold on to them: they are freed as soon as it returns, which is what bounds -// the memory. Windows come from params.batch_clustering.batch_size. +// the memory. Windows come from +// params.batch_clustering.effective_batch_size(), so a caller that set a batch +// size but no output path gets one window -- see BatchClusteringOption. // // Every window's lambda and beta are computed from the GLOBAL corpus, and every // list's k-means seed from its own GLOBAL term id, so the window count cannot // change what is produced -- see the batched-build tests, which assert file // equality against an unbatched build. -using ClusteredWindowSink = - std::function&& clusters)>; +using ClusteredWindowSink = std::function&& clusters)>; void for_each_clustered_window(const SparseVectors* vectors, const SparseVectorsConfig& config, @@ -191,9 +211,10 @@ void for_each_clustered_window(const SparseVectors* vectors, const ClusteredWindowSink& sink); // Every term's clustered posting list, in term order. The whole-corpus form of -// for_each_clustered_window: batch_size still bounds the inverted-list -// intermediate, but the result is retained in full, so this is bounded by the -// clustered lists rather than by a window. +// for_each_clustered_window, and so the unbatched build: the result is retained +// in full, which is what bounds this by the clustered lists rather than by a +// window. Callers reach it without an output path, which is why that is also +// the case where the window count is dropped. inline std::vector build_inverted_lists_clusters( const SparseVectors* vectors, const SparseVectorsConfig& config, const SeismicClusterParameters& seismic_cluster_params) { diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index 39d8109..e1e23bb 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -24,7 +24,6 @@ #include "nsparse/index.h" #include "nsparse/invlists/inverted_lists.h" #include "nsparse/io/seismic_invlists_writer.h" -#include "nsparse/seismic_batched_build.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" @@ -143,28 +142,10 @@ void SeismicIndex::add(idx_t n, const idx_t* indptr, const term_t* indices, } void SeismicIndex::build() { - const SparseVectorsConfig config = { - .element_size = kElementSize, - .dimension = static_cast(get_dimension())}; - const auto& batch = cluster_parameter_.batch_clustering; - if (batch.batch_size > 1 && !batch.batch_file_output_path.empty()) { - // Streamed to a file a window at a time, so the whole index is never - // resident, then borrowed back so this is a usable index. A single - // window takes the ordinary path below: it already holds its own lists, - // and writing them out only to map them back would be work for nothing. - const size_t lists_offset = detail::write_seismic_index_batched( - get_vectors(), config, cluster_parameter_, - {.id = fourcc(name), - .version = kFormatVersion, - .dimension = get_dimension()}, - [this](IOWriter* io_writer) { vectors_->serialize(io_writer); }, - batch.batch_file_output_path); - clustered_inverted_lists = detail::map_streamed_lists( - batch.batch_file_output_path, lists_offset, &batch_mapped_file_); - return; - } - clustered_inverted_lists = detail::build_inverted_lists_clusters( - get_vectors(), config, cluster_parameter_); + clustered_inverted_lists = build_clustered_lists( + {.element_size = kElementSize, + .dimension = static_cast(get_dimension())}, + cluster_parameter_); } auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index 213b656..1d918d2 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -27,7 +27,6 @@ #include "nsparse/invlists/inverted_lists.h" #include "nsparse/io/io.h" #include "nsparse/io/seismic_invlists_writer.h" -#include "nsparse/seismic_batched_build.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" @@ -192,30 +191,12 @@ ScalarQuantizer SeismicScalarQuantizedIndex::query_quantizer( } void SeismicScalarQuantizedIndex::build() { - const SparseVectorsConfig config = { - .element_size = sq_.bytes_per_value(), - .dimension = static_cast(get_dimension())}; - const auto& batch = cluster_parameter_.batch_clustering; - if (batch.batch_size > 1 && !batch.batch_file_output_path.empty()) { - // The quantization header comes first, exactly as write_index writes - // it; the codes in `vectors_` are already quantized, so the batched - // build needs no knowledge of the quantizer beyond its width. - const size_t lists_offset = detail::write_seismic_index_batched( - get_vectors(), config, cluster_parameter_, - {.id = fourcc(name), - .version = kFormatVersion, - .dimension = get_dimension()}, - [this](IOWriter* io_writer) { - write_quantization_header(io_writer); - vectors_->serialize(io_writer); - }, - batch.batch_file_output_path); - clustered_inverted_lists = detail::map_streamed_lists( - batch.batch_file_output_path, lists_offset, &batch_mapped_file_); - return; - } - clustered_inverted_lists = detail::build_inverted_lists_clusters( - get_vectors(), config, cluster_parameter_); + // The codes in `vectors_` are already quantized by add(), so the shared + // build needs no knowledge of the quantizer beyond its width. + clustered_inverted_lists = build_clustered_lists( + {.element_size = sq_.bytes_per_value(), + .dimension = static_cast(get_dimension())}, + cluster_parameter_); } auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, diff --git a/python_tests/test_seismic_batched_build.py b/python_tests/test_seismic_batched_build.py index cd38e9f..93daf31 100644 --- a/python_tests/test_seismic_batched_build.py +++ b/python_tests/test_seismic_batched_build.py @@ -9,10 +9,10 @@ Batching is a build option rather than a separate entry point, so there is nothing new to wrap: it is reached through the factory description, the same way -lambda and beta are. `inverted_list_batch_size` bounds the build's memory; -with more than one window, `batch_file_output_path` is where the index is streamed -as it is built, and its posting lists are then borrowed back from that file -- -which is the path a corpus too large for RAM needs. +lambda and beta are. `inverted_list_batch_size` splits the term space, and +`batch_file_output_path` is a directory the build may spill windows into -- +scratch, not output. build() writes no index and leaves nothing behind; the +index is serialized with write_index, exactly as an unbatched one is. """ import numpy as np @@ -32,23 +32,32 @@ RECALL_FLOOR = 0.80 -def streamed(corpus, out_path, batch_size, kind="seismic"): - """Build straight to `out_path`; returns nothing, the file is the index.""" +def scratch_dir(tmp_path, name="scratch"): + """An existing directory for the build to spill into.""" + path = tmp_path / name + path.mkdir(exist_ok=True) + return path + + +def batched_index(corpus, tmp_path, batch_size, kind="seismic", name="scratch"): + """A build whose windows are spilled to scratch, ready to serve or write.""" spec = ( f"{kind},{BASE}|inverted_list_batch_size={batch_size}" - f"|batch_file_output_path={out_path}" + f"|batch_file_output_path={scratch_dir(tmp_path, name)}" ) index = nsparse.index_factory(corpus.dim, spec) add_corpus(index, corpus) index.build() - return str(out_path) + return index -# Batching starts at 2: one window is an ordinary build and writes no file. +# Batching starts at 2: one window is an ordinary build with nothing to spill. @pytest.mark.parametrize("batch_size", [2, 4, 32]) def test_happy_case(batch_size, corpus, queries, oracle, tmp_path): - """build -> read back mapped -> query -> accuracy, at several splits.""" - path = streamed(corpus, tmp_path / "batched.idx", batch_size) + """build -> write_index -> read back mapped -> query -> accuracy.""" + path = str(tmp_path / "batched.idx") + nsparse.write_index(batched_index(corpus, tmp_path, batch_size), path) + index = nsparse.read_index(path, nsparse.kUseMmap) assert index.num_vectors() == corpus.n assert index.get_dimension() == corpus.dim @@ -66,72 +75,65 @@ def test_happy_case(batch_size, corpus, queries, oracle, tmp_path): "kind", ["seismic", "seismic_sq", "disk_seismic", "disk_seismic_sq"] ) def test_matches_in_memory_build(kind, corpus, tmp_path): - """At a fixed seed a streamed build is the in-memory build, byte for byte. + """At a fixed seed a batched build serializes to the in-memory build's file. Over all four types in the family: float and quantizing, since the shared build only needs the code width (add() has already encoded the values), and - in-memory and disk-resident, which get there differently -- the disk types' - payload cannot be streamed section by section, so they spill their clustered - lists and write from that mapping instead. + in-memory and disk-resident, which have different payloads to write from the + same lists. """ in_memory = tmp_path / "memory.idx" nsparse.write_index(make_index(f"{kind},{BASE}", corpus), str(in_memory)) - batched = streamed(corpus, tmp_path / "batched.idx", 4, kind=kind) - assert in_memory.read_bytes() == open(batched, "rb").read() - # The spill the disk types take is scratch, deleted with the build. - assert not (tmp_path / "batched.idx.lists").exists() + spilled = tmp_path / "batched.idx" + nsparse.write_index(batched_index(corpus, tmp_path, 4, kind=kind), str(spilled)) + assert in_memory.read_bytes() == spilled.read_bytes() -@pytest.mark.parametrize("kind", ["disk_seismic", "disk_seismic_sq"]) -def test_streamed_disk_index_is_searchable_after_build( - kind, corpus, queries, oracle, tmp_path -): - """A batched disk build serves from the file it wrote. - Its summaries and its inline forward index are borrowed from that mapping, - so there is no reopening by path -- and nothing left pointing at the spill. +def test_scratch_directory_is_left_empty(corpus, tmp_path): + """The spill is scratch: unlinked as soon as it is mapped. + + The lists stay readable from the mapping, so the index is still servable + while the directory the caller lent is already empty again. """ - index = nsparse.index_factory( - corpus.dim, - f"{kind},{BASE}|inverted_list_batch_size=8" - f"|batch_file_output_path={tmp_path / 'streamed.idx'}", - ) - add_corpus(index, corpus) - index.build() + scratch = scratch_dir(tmp_path) + index = batched_index(corpus, tmp_path, 8) + assert list(scratch.iterdir()) == [] assert index.num_vectors() == corpus.n - params = nsparse.DiskSeismicSearchParameters(8, 200) - _, labels = search(index, queries, params=params) - want_labels, _ = oracle - assert recall_at_k(labels, want_labels) >= RECALL_FLOOR + nsparse.write_index(index, str(tmp_path / "out.idx")) + assert (tmp_path / "out.idx").stat().st_size > 0 -def test_streamed_index_is_searchable_after_build(corpus, queries, oracle, tmp_path): - """build() leaves a usable index, not an empty object. +@pytest.mark.parametrize("kind", ["seismic", "disk_seismic"]) +def test_batched_index_is_searchable_after_build( + kind, corpus, queries, oracle, tmp_path +): + """build() leaves an index that serves, not an empty object. - The lists are borrowed back from the file it just wrote, so there is no + Its posting lists are borrowed from the spill's mapping, so there is no reopening by path and they are never copied onto the heap. """ - spec = ( - f"seismic,{BASE}|inverted_list_batch_size=8" - f"|batch_file_output_path={tmp_path / 'streamed.idx'}" - ) - index = nsparse.index_factory(corpus.dim, spec) - add_corpus(index, corpus) - index.build() - + index = batched_index(corpus, tmp_path, 8, kind=kind) assert index.num_vectors() == corpus.n - _, labels = search(index, queries) + + params = ( + nsparse.DiskSeismicSearchParameters(8, 200) + if kind == "disk_seismic" + else None + ) + _, labels = search(index, queries, params=params) want_labels, _ = oracle assert recall_at_k(labels, want_labels) >= RECALL_FLOOR def test_batch_size_alone_leaves_the_index_in_memory(corpus, queries, tmp_path): - """Without an output path, batching only bounds the build's intermediates. + """Without a scratch directory the window count is ignored, not half-applied. - The index is still usable in memory and still the same index -- this is the - path every index type gets from build(), including the disk-resident ones. + Splitting the term space with nowhere to spill the windows would leave the + clustered lists accumulating anyway, for a corpus pass per window, so the + build runs as one window and produces the same index it always did. """ unbatched = make_index(f"seismic,{BASE}", corpus) batched = make_index(f"seismic,{BASE}|inverted_list_batch_size=8", corpus) @@ -144,17 +146,14 @@ def test_batch_size_alone_leaves_the_index_in_memory(corpus, queries, tmp_path): def test_batch_count_is_not_observable(corpus, queries, tmp_path): - """The split is a memory knob: at a fixed seed it cannot change the results. - - Against an unbatched build, which writes its file the ordinary way since one - window streams nothing. - """ + """The split is a memory knob: at a fixed seed it cannot change the results.""" plain = tmp_path / "plain.idx" nsparse.write_index(make_index(f"seismic,{BASE}", corpus), str(plain)) - many = streamed(corpus, tmp_path / "many.idx", 16) + many = tmp_path / "many.idx" + nsparse.write_index(batched_index(corpus, tmp_path, 16), str(many)) want_d, want_l = search(nsparse.read_index(str(plain)), queries) - got_d, got_l = search(nsparse.read_index(many), queries) + got_d, got_l = search(nsparse.read_index(str(many)), queries) np.testing.assert_array_equal(got_l, want_l) np.testing.assert_allclose(got_d, want_d, rtol=1e-6, atol=1e-6) @@ -163,9 +162,23 @@ def test_rejects_dimension_smaller_than_corpus(corpus, tmp_path): """A term the declared dimension does not cover is an error, not a silent drop.""" spec = ( f"seismic,{BASE}|inverted_list_batch_size=4" - f"|batch_file_output_path={tmp_path / 'bad.idx'}" + f"|batch_file_output_path={scratch_dir(tmp_path)}" ) index = nsparse.index_factory(corpus.dim // 2, spec) add_corpus(index, corpus) with pytest.raises(ValueError): index.build() + + +def test_rejects_a_scratch_path_that_is_not_a_directory(corpus, tmp_path): + """Somewhere to spill is the caller's to provide.""" + not_a_dir = tmp_path / "regular-file" + not_a_dir.write_bytes(b"") + spec = ( + f"seismic,{BASE}|inverted_list_batch_size=4" + f"|batch_file_output_path={not_a_dir}" + ) + index = nsparse.index_factory(corpus.dim, spec) + add_corpus(index, corpus) + with pytest.raises(ValueError): + index.build() diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index bcc930b..cdb8b09 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -104,16 +104,25 @@ class TempDir { return path_ + "/" + name; } + // An existing, empty directory to spill into: what + // batch_file_output_path takes. + [[nodiscard]] std::string scratch( + const std::string& name = "scratch") const { + const std::string dir = file(name); + std::filesystem::create_directories(dir); + return dir; + } + private: std::string path_; }; SeismicClusterParameters params_for(size_t batch_size, - const std::string& out_path, int seed) { + const std::string& scratch_dir, int seed) { SeismicClusterParameters params = { .lambda = kLambda, .beta = kBeta, .alpha = kAlpha}; params.batch_clustering.batch_size = batch_size; - params.batch_clustering.batch_file_output_path = out_path; + params.batch_clustering.batch_file_output_path = scratch_dir; params.seed = seed; return params; } @@ -124,17 +133,21 @@ std::vector read_file(const std::string& path) { std::istreambuf_iterator()}; } -// A build streamed straight to `out`, through the index's own build(). -std::vector streamed(const Corpus& corpus, size_t batch_size, - const std::string& out, int seed = kSeed) { - SeismicIndex index(corpus.dim, params_for(batch_size, out, seed)); +// A batched build -- windows spilled to a scratch directory -- then written out +// the ordinary way. build() itself produces no file; write_index is still what +// serializes an index. +std::vector batched(const Corpus& corpus, size_t batch_size, + const TempDir& dir, const std::string& out, + int seed = kSeed) { + SeismicIndex index(corpus.dim, params_for(batch_size, dir.scratch(), seed)); index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); index.build(); + write_index(&index, const_cast(out.c_str())); return read_file(out); } -// The same corpus built the ordinary way and written with write_index. +// The same corpus built whole and written the same way. std::vector in_memory(const Corpus& corpus, const std::string& out, size_t batch_size = 1, int seed = kSeed) { SeismicIndex index(corpus.dim, params_for(batch_size, "", seed)); @@ -186,34 +199,71 @@ std::string write_native_csr(const Corpus& corpus, const std::string& path) { } // namespace -// The point of the seeding discipline: for a fixed seed a streamed build is not -// merely equivalent to build() + write_index, it is the same file. Each list's -// k-means seed comes from its own global term id, so neither the window a term -// landed in nor the order the threads reached it can leak into the output. -TEST(SeismicBatchedBuild, StreamedBuildIsByteIdenticalToInMemoryBuild) { +// The point of the seeding discipline: for a fixed seed a batched build is not +// merely equivalent to a whole-corpus one, it serializes to the same file. Each +// list's k-means seed comes from its own global term id, so neither the window +// a term landed in nor the order the threads reached it can leak into the +// output. +TEST(SeismicBatchedBuild, BatchedBuildIsByteIdenticalToInMemoryBuild) { Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/42); TempDir dir("identical"); EXPECT_EQ(in_memory(corpus, dir.file("mem.dat")), - streamed(corpus, /*batch_size=*/4, dir.file("streamed.dat"))); + batched(corpus, /*batch_size=*/4, dir, dir.file("batched.dat"))); } // The window count is a memory knob, not a behaviour knob: at a fixed seed -// every count has to produce the same file. -TEST(SeismicBatchedBuild, StreamedBuildIsIdenticalAcrossBatchCounts) { +// every count has to produce the same index. +TEST(SeismicBatchedBuild, BatchedBuildIsIdenticalAcrossBatchCounts) { Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/11); TempDir dir("counts"); - // batch_size <= 1 is an ordinary build, so its file comes from write_index. const auto one = in_memory(corpus, dir.file("b1.dat")); ASSERT_FALSE(one.empty()); - EXPECT_EQ(one, streamed(corpus, 2, dir.file("b2.dat"))); - EXPECT_EQ(one, streamed(corpus, 10, dir.file("b10.dat"))); + EXPECT_EQ(one, batched(corpus, 2, dir, dir.file("b2.dat"))); + EXPECT_EQ(one, batched(corpus, 10, dir, dir.file("b10.dat"))); // More windows than terms is clamped to one term each. - EXPECT_EQ(one, streamed(corpus, 1000, dir.file("b1000.dat"))); + EXPECT_EQ(one, batched(corpus, 1000, dir, dir.file("b1000.dat"))); } -// batch_size alone bounds the inverted-list intermediate and leaves the index -// in memory. That is the path every index type gets, including the two disk -// ones, so it must not change what build() produces either. +// The spill is scratch, and nothing outlives the build: it is unlinked as soon +// as it is mapped, so the directory the caller lent is empty again while the +// index is still serving from those very bytes. +TEST(SeismicBatchedBuild, LeavesNothingInTheScratchDirectory) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/3); + TempDir dir("scratch"); + const std::string scratch = dir.scratch(); + + SeismicIndex index(corpus.dim, params_for(8, scratch, kSeed)); + index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index.build(); + + EXPECT_TRUE(std::filesystem::is_empty(scratch)); + // And the lists are still readable, which is the point of unlinking rather + // than deleting: the mapping keeps the bytes alive. + write_index(&index, const_cast(dir.file("out.dat").c_str())); + EXPECT_EQ(read_file(dir.file("out.dat")), + in_memory(corpus, dir.file("mem.dat"))); +} + +// The two knobs are only useful together. A window count with nowhere to spill +// the windows bounds the fill intermediate while leaving the clustered lists to +// accumulate -- a corpus pass per window for a fraction of the peak -- so it is +// resolved to one window rather than honoured. +TEST(SeismicBatchedBuild, BatchSizeWithoutAScratchDirectoryIsOneWindow) { + BatchClusteringOption opt; + opt.batch_size = 64; + EXPECT_EQ(opt.effective_batch_size(), 1U); + + opt.batch_file_output_path = "/tmp/does-not-need-to-exist"; + EXPECT_EQ(opt.effective_batch_size(), 64U); + + // 0 is not a window count; a build always runs at least one. + opt.batch_size = 0; + EXPECT_EQ(opt.effective_batch_size(), 1U); +} + +// End to end, the resolution above is invisible: a batch size with no scratch +// directory builds exactly what an unbatched build does. TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeAnInMemoryBuild) { Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/5); TempDir dir("inmem_batched"); @@ -223,210 +273,45 @@ TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeAnInMemoryBuild) { EXPECT_EQ(unbatched, in_memory(corpus, dir.file("b64.dat"), 64)); } -// The disk-resident types share the same build, so batch_size has to bound -// their intermediates too without changing what they produce. -TEST(SeismicBatchedBuild, BatchSizeAloneDoesNotChangeADiskIndex) { - Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/71); - TempDir dir("disk"); - - auto build_disk = [&corpus](size_t batch_size, const std::string& out) { - DiskSeismicIndex index(corpus.dim, params_for(batch_size, "", kSeed)); - index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - index.build(); - write_index(&index, const_cast(out.c_str())); - return read_file(out); - }; +// Every type in the family reaches the same build, so each has to spill and +// come back with the index a whole-corpus build would have produced. +// Parametrized over all four: float and quantizing, in-memory and +// disk-resident. +class BatchedBuildEveryType : public testing::TestWithParam {}; - const auto unbatched = build_disk(1, dir.file("b1.dat")); - ASSERT_FALSE(unbatched.empty()); - EXPECT_EQ(unbatched, build_disk(8, dir.file("b8.dat"))); - EXPECT_EQ(unbatched, build_disk(64, dir.file("b64.dat"))); -} - -// The disk types cannot stream their payload out window by window -- the inline -// forward index that follows their summaries is laid out from the doc-id -// membership of every list -- so they spill the lists instead and write the -// payload from that mapping. Same contract as the in-memory types all the same: -// at a fixed seed the file is what write_index would have produced. -TEST(SeismicBatchedBuild, StreamsADiskIndexIdenticallyToo) { +TEST_P(BatchedBuildEveryType, SpilledBuildIsByteIdenticalToInMemoryBuild) { + const std::string kind = GetParam(); Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/71); - TempDir dir("disk_streamed"); - const std::string mem_path = dir.file("mem.dat"); - const std::string streamed_path = dir.file("streamed.dat"); - - DiskSeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); - mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - mem.build(); - write_index(&mem, const_cast(mem_path.c_str())); - - DiskSeismicIndex batched(corpus.dim, params_for(8, streamed_path, kSeed)); - batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - batched.build(); - - EXPECT_EQ(read_file(mem_path), read_file(streamed_path)); - // The spill is scratch: it must not outlive the build that took it. - EXPECT_FALSE(std::filesystem::exists(streamed_path + ".lists")); - // And the window count is still not a behaviour knob. - const std::string many_path = dir.file("many.dat"); - DiskSeismicIndex many(corpus.dim, params_for(64, many_path, kSeed)); - many.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - many.build(); - EXPECT_EQ(read_file(mem_path), read_file(many_path)); -} - -// The quantized disk index writes a quantization header before the shared -// payload, so a batched build has to lay that down and read it back to reopen -// its own file at the right offset. -TEST(SeismicBatchedBuild, StreamsAQuantizedDiskIndexIdenticallyToo) { - Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/29); - TempDir dir("disk_sq"); - const std::string mem_path = dir.file("mem.dat"); - const std::string streamed_path = dir.file("streamed.dat"); - - DiskSeismicScalarQuantizedIndex mem(QuantizerType::QT_8bit, 0.0F, 3.0F, - params_for(1, "", kSeed), corpus.dim); - mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - mem.build(); - write_index(&mem, const_cast(mem_path.c_str())); - - DiskSeismicScalarQuantizedIndex batched(QuantizerType::QT_8bit, 0.0F, 3.0F, - params_for(8, streamed_path, kSeed), - corpus.dim); - batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - batched.build(); - - EXPECT_EQ(read_file(mem_path), read_file(streamed_path)); - // It reads back as the quantized disk type it claims to be, with the range - // it was built with -- the header the batched write had to reproduce. - std::unique_ptr reloaded(read_index( - const_cast(streamed_path.c_str()), IndexIoFlag::kUseMmap)); - EXPECT_EQ(reloaded->id(), DiskSeismicScalarQuantizedIndex::name); - EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n)); - const auto* sq_index = - dynamic_cast(reloaded.get()); - ASSERT_NE(sq_index, nullptr); - EXPECT_EQ(sq_index->get_scalar_quantizer().get_min(), 0.0F); - EXPECT_EQ(sq_index->get_scalar_quantizer().get_max(), 3.0F); -} - -// A batched disk build ends serving from the file it wrote: its summaries and -// its forward index are borrowed from that mapping, not from the spill it -// deleted. Against an unbatched build at the same seed, so identical results -// rather than merely close ones. -TEST(SeismicBatchedBuild, BatchedDiskBuildIsSearchableAfterBuild) { - Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/17); - Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); - const int k = 10; - const auto n = static_cast(queries.n); - TempDir dir("disk_searchable"); - - const auto search_with = [&](Index& index) { - std::vector dist(n * k); - std::vector lab(n * k); - DiskSeismicSearchParameters params(/*cut=*/3, /*k_prime=*/50); - index.search(queries.n, queries.indptr.data(), queries.indices.data(), - queries.values.data(), k, dist.data(), lab.data(), - ¶ms); - return std::pair{dist, lab}; + TempDir dir("every_type"); + const std::string base = + "lambda=64|beta=6|alpha=0.4|seed=42|inverted_list_batch_size="; + + const auto build_and_write = [&](const std::string& spec, + const std::string& out) { + std::unique_ptr index(index_factory(corpus.dim, spec.c_str())); + index->add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index->build(); + write_index(index.get(), const_cast(out.c_str())); + return read_file(out); }; - // The unbatched reference has to be read back mapped: an unwritten disk - // index has no forward index to score from. - const std::string mem_path = dir.file("mem.dat"); - DiskSeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); - mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - mem.build(); - write_index(&mem, const_cast(mem_path.c_str())); - std::unique_ptr reference( - read_index(const_cast(mem_path.c_str()), IndexIoFlag::kUseMmap)); - const auto [want_dist, want_lab] = search_with(*reference); - - DiskSeismicIndex batched(corpus.dim, - params_for(4, dir.file("b.dat"), kSeed)); - batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - batched.build(); - - // No reopening by path: build() mapped its own output back in. - EXPECT_EQ(batched.num_vectors(), static_cast(corpus.n)); - const auto [got_dist, got_lab] = search_with(batched); - EXPECT_EQ(got_lab, want_lab); - EXPECT_EQ(got_dist, want_dist); -} - -// The disk index's own reason for existing: a corpus that came from a mapping. -// Three mappings are then live at once -- the corpus, the spill, and the output -// -- and the build may give up only the spill. -TEST(SeismicBatchedBuild, BatchedDiskBuildKeepsTheCorpusMapping) { - Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/13); - Corpus queries = make_corpus(/*n_docs=*/40, /*dim=*/200, /*seed=*/77); - const int k = 10; - const auto n = static_cast(queries.n); - TempDir dir("disk_mapped"); - - const std::string native = write_native_csr(corpus, dir.file("corpus.csr")); - const std::string out = dir.file("out.dat"); - DiskSeismicIndex index(corpus.dim, params_for(3, out, kSeed)); - index.read_csr(native.c_str(), Residency::kMmap); - index.build(); - - EXPECT_EQ(index.num_vectors(), static_cast(corpus.n)); - std::vector dist(n * k); - std::vector lab(n * k); - DiskSeismicSearchParameters params(/*cut=*/3, /*k_prime=*/50); - static_cast(index).search( - queries.n, queries.indptr.data(), queries.indices.data(), - queries.values.data(), k, dist.data(), lab.data(), ¶ms); - EXPECT_TRUE( - std::any_of(lab.begin(), lab.end(), [](idx_t id) { return id >= 0; })); - - // Same file a heap-resident corpus produces: residency is SparseVectors' - // business, not the build's. - DiskSeismicIndex owned(corpus.dim, - params_for(3, dir.file("owned.dat"), kSeed)); - owned.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - owned.build(); - EXPECT_EQ(read_file(out), read_file(dir.file("owned.dat"))); + const auto whole = + build_and_write(kind + "," + base + "1", dir.file("mem.dat")); + ASSERT_FALSE(whole.empty()); + const std::string scratch = dir.scratch(); + EXPECT_EQ(whole, build_and_write(kind + "," + base + + "8|batch_file_output_path=" + scratch, + dir.file("b8.dat"))); + EXPECT_EQ(whole, build_and_write(kind + "," + base + + "64|batch_file_output_path=" + scratch, + dir.file("b64.dat"))); + EXPECT_TRUE(std::filesystem::is_empty(scratch)); } -// The generalization that matters: a quantizing index streams too, because the -// codes in `vectors_` are already quantized by add() and the shared build only -// needs their width. -TEST(SeismicBatchedBuild, StreamsAQuantizedIndexIdenticallyToo) { - Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/23); - TempDir dir("sq"); - const std::string mem_path = dir.file("mem.dat"); - const std::string streamed_path = dir.file("streamed.dat"); - - SeismicScalarQuantizedIndex mem(QuantizerType::QT_8bit, 0.0F, 3.0F, - params_for(1, "", kSeed), corpus.dim); - mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - mem.build(); - write_index(&mem, const_cast(mem_path.c_str())); - - SeismicScalarQuantizedIndex batched(QuantizerType::QT_8bit, 0.0F, 3.0F, - params_for(4, streamed_path, kSeed), - corpus.dim); - batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - batched.build(); - - EXPECT_EQ(read_file(mem_path), read_file(streamed_path)); - // And it loads as the quantized type it claims to be. - std::unique_ptr reloaded( - read_index(const_cast(streamed_path.c_str()))); - EXPECT_EQ(reloaded->id(), SeismicScalarQuantizedIndex::name); - EXPECT_EQ(reloaded->num_vectors(), static_cast(corpus.n)); -} +INSTANTIATE_TEST_SUITE_P(AllSeismicTypes, BatchedBuildEveryType, + testing::Values("seismic", "seismic_sq", + "disk_seismic", "disk_seismic_sq")); // Windows are cut to equal posting counts, not equal width, so a term heavier // than a whole window's target has to be handled: it cannot be split, and it @@ -461,17 +346,17 @@ TEST(SeismicBatchedBuild, HandlesATermHeavierThanAWholeWindow) { TempDir dir("skewed"); const auto one = in_memory(skewed, dir.file("b1.dat")); ASSERT_FALSE(one.empty()); - // Every one of these has to cover all 64 terms exactly once, or the - // streamed file would be short and the write would refuse it. - EXPECT_EQ(one, streamed(skewed, 8, dir.file("b8.dat"))); - EXPECT_EQ(one, streamed(skewed, 32, dir.file("b32.dat"))); - EXPECT_EQ(one, streamed(skewed, 64, dir.file("b64.dat"))); - EXPECT_EQ(one, streamed(skewed, 200, dir.file("b200.dat"))); + // Every one of these has to cover all 64 terms exactly once, or the spill + // would be short and mapping it back would refuse it. + EXPECT_EQ(one, batched(skewed, 8, dir, dir.file("b8.dat"))); + EXPECT_EQ(one, batched(skewed, 32, dir, dir.file("b32.dat"))); + EXPECT_EQ(one, batched(skewed, 64, dir, dir.file("b64.dat"))); + EXPECT_EQ(one, batched(skewed, 200, dir, dir.file("b200.dat"))); } -// A batched build ends holding the index it wrote, so build() leaves something -// usable rather than an empty object. Against an unbatched build at the same -// seed: identical builds, so identical results. +// A batched build ends holding its lists, borrowed from the spill, so build() +// leaves an index that serves rather than an empty object. Against an unbatched +// build at the same seed: identical builds, so identical results. TEST(SeismicBatchedBuild, BatchedBuildIsSearchableAfterBuild) { Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/7); Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); @@ -495,22 +380,56 @@ TEST(SeismicBatchedBuild, BatchedBuildIsSearchableAfterBuild) { mem.build(); const auto [want_dist, want_lab] = search_with(mem); - SeismicIndex batched(corpus.dim, params_for(4, dir.file("b.dat"), kSeed)); - batched.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + SeismicIndex spilled(corpus.dim, params_for(4, dir.scratch(), kSeed)); + spilled.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + spilled.build(); + + EXPECT_EQ(spilled.num_vectors(), static_cast(corpus.n)); + const auto [got_dist, got_lab] = search_with(spilled); + EXPECT_EQ(got_lab, want_lab); + EXPECT_EQ(got_dist, want_dist); +} + +// A batched disk build serves too, from the corpus it already holds -- the +// forward index only exists once write_index has laid it out. +TEST(SeismicBatchedBuild, BatchedDiskBuildIsSearchableAfterBuild) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/17); + Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); + const int k = 10; + const auto n = static_cast(queries.n); + TempDir dir("disk_searchable"); + + const auto search_with = [&](Index& index) { + std::vector dist(n * k); + std::vector lab(n * k); + DiskSeismicSearchParameters params(/*cut=*/3, /*k_prime=*/50); + index.search(queries.n, queries.indptr.data(), queries.indices.data(), + queries.values.data(), k, dist.data(), lab.data(), + ¶ms); + return std::pair{dist, lab}; + }; + + DiskSeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); + mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + mem.build(); + const auto [want_dist, want_lab] = search_with(mem); + + DiskSeismicIndex spilled(corpus.dim, params_for(4, dir.scratch(), kSeed)); + spilled.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); - batched.build(); + spilled.build(); - // No reopening by path: build() mapped the lists back in. - EXPECT_EQ(batched.num_vectors(), static_cast(corpus.n)); - const auto [got_dist, got_lab] = search_with(batched); + EXPECT_EQ(spilled.num_vectors(), static_cast(corpus.n)); + const auto [got_dist, got_lab] = search_with(spilled); EXPECT_EQ(got_lab, want_lab); EXPECT_EQ(got_dist, want_dist); } // A corpus borrowed from a mapping is the case two mappings exist for: the -// corpus one, which vectors_ still scores from, and the one the build just -// wrote, which the posting lists borrow from. Neither may be given up for the -// other. +// corpus's, which vectors_ still scores from, and the spill's, which the +// posting lists borrow from. Neither may be given up for the other. TEST(SeismicBatchedBuild, KeepsTheCorpusMappingWhileBorrowingItsOwnLists) { Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/13); Corpus queries = make_corpus(/*n_docs=*/40, /*dim=*/200, /*seed=*/77); @@ -519,12 +438,12 @@ TEST(SeismicBatchedBuild, KeepsTheCorpusMappingWhileBorrowingItsOwnLists) { TempDir dir("mapped_reload"); const std::string native = write_native_csr(corpus, dir.file("corpus.csr")); - SeismicIndex index(corpus.dim, params_for(3, dir.file("out.dat"), kSeed)); + SeismicIndex index(corpus.dim, params_for(3, dir.scratch(), kSeed)); index.read_csr(native.c_str(), Residency::kMmap); index.build(); // Still serving: scoring reads the mapped corpus, the lists come from the - // second mapping. + // spill's mapping. EXPECT_EQ(index.num_vectors(), static_cast(corpus.n)); std::vector dist(n * k); std::vector lab(n * k); @@ -545,7 +464,7 @@ TEST(SeismicBatchedBuild, PerTermMembershipInvariantAcrossBatches) { const std::string one = dir.file("1.dat"); const std::string ten = dir.file("10.dat"); in_memory(corpus, one, /*batch_size=*/1, kRandomSeed); - streamed(corpus, 10, ten, kRandomSeed); + batched(corpus, 10, dir, ten, kRandomSeed); auto sets1 = per_term_doc_sets(one); auto sets10 = per_term_doc_sets(ten); @@ -559,14 +478,14 @@ TEST(SeismicBatchedBuild, PerTermMembershipInvariantAcrossBatches) { // Regression for the term_t (uint16) window-bound overflow: a dimension at the // 2^16 boundary must still build every term's list. Before the fix, size_t // window bounds cast to term_t wrapped mod 65536, so dim=65536 built nothing -// while n_lists claimed 65536 -> a corrupt, unloadable file. +// while n_lists claimed 65536 -> a corrupt, unloadable spill. TEST(SeismicBatchedBuild, HandlesDimensionAt65536) { const int dim = 65536; // term ids 0..65535 all fit term_t (uint16) Corpus corpus = make_corpus(/*n_docs=*/3000, dim, /*seed=*/5); TempDir dir("dim64k"); const std::string path = dir.file("index.dat"); - // Two windows, so the batched path is what produces the file. - streamed(corpus, 2, path); + // Two windows, so the spilling path is what produces the lists. + batched(corpus, 2, dir, path); // Must load without "unexpected end of index file". std::unique_ptr idx(read_index(const_cast(path.c_str()))); @@ -581,23 +500,23 @@ TEST(SeismicBatchedBuild, HandlesDimensionAt65536) { EXPECT_GT(total_docs, 0U); } -// A streamed index must serve correctly through the mapped read path -- the way -// a caller whose corpus did not fit in RAM is going to use it. +// A batched build's index must serve correctly through the mapped read path -- +// the way a caller whose corpus did not fit in RAM is going to use it. TEST(SeismicBatchedBuild, SearchThroughMappedReadMatchesInMemory) { Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/7); Corpus queries = make_corpus(/*n_docs=*/50, /*dim=*/200, /*seed=*/99); const int k = 10; TempDir dir("search"); - const std::string streamed_path = dir.file("streamed.dat"); + const std::string batched_path = dir.file("batched.dat"); SeismicIndex mem(corpus.dim, params_for(1, "", kSeed)); mem.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); mem.build(); - streamed(corpus, 4, streamed_path); + batched(corpus, 4, dir, batched_path); std::unique_ptr disk(read_index( - const_cast(streamed_path.c_str()), IndexIoFlag::kUseMmap)); + const_cast(batched_path.c_str()), IndexIoFlag::kUseMmap)); SeismicSearchParameters search_params(/*cut=*/3, /*heap_factor=*/1.0F); const auto n = static_cast(queries.n); @@ -625,12 +544,13 @@ TEST(SeismicBatchedBuild, MappedCorpusMatchesOwnedCorpus) { TempDir dir("mapped"); const std::string owned_path = dir.file("owned.dat"); const std::string mapped_path = dir.file("mapped.dat"); - streamed(corpus, 3, owned_path); + batched(corpus, 3, dir, owned_path); const std::string native = write_native_csr(corpus, dir.file("corpus.csr")); - SeismicIndex mapped(corpus.dim, params_for(3, mapped_path, kSeed)); + SeismicIndex mapped(corpus.dim, params_for(3, dir.scratch(), kSeed)); mapped.read_csr(native.c_str(), Residency::kMmap); mapped.build(); + write_index(&mapped, const_cast(mapped_path.c_str())); EXPECT_EQ(read_file(owned_path), read_file(mapped_path)); } @@ -644,15 +564,15 @@ TEST(SeismicBatchedBuild, FactoryDescriptionDrivesTheBatchedBuild) { const std::string spec = "seismic,lambda=64|beta=6|alpha=0.4|seed=42|" "inverted_list_batch_size=8|batch_file_output_path=" + - path; + dir.scratch("from_factory"); std::unique_ptr index(index_factory(corpus.dim, spec.c_str())); index->add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); index->build(); + write_index(index.get(), const_cast(path.c_str())); - ASSERT_TRUE(std::filesystem::exists(path)); - EXPECT_EQ(streamed(corpus, 8, dir.file("direct.dat")), read_file(path)); + EXPECT_EQ(batched(corpus, 8, dir, dir.file("direct.dat")), read_file(path)); } TEST(SeismicBatchedBuild, RejectsInvalidInput) { @@ -662,20 +582,27 @@ TEST(SeismicBatchedBuild, RejectsInvalidInput) { // A term the declared dimension does not cover. The mapped read path does // not range-check terms, so this is the build's own guard -- without it the // term would be silently dropped from the index. - SeismicIndex narrow(corpus.dim / 2, - params_for(1, dir.file("narrow.dat"), kSeed)); + SeismicIndex narrow(corpus.dim / 2, params_for(1, "", kSeed)); narrow.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), corpus.values.data()); EXPECT_THROW(narrow.build(), std::invalid_argument); - // Streaming an empty corpus would leave a header-only file that read_index - // cannot parse, so it is refused rather than written. - SeismicIndex empty(corpus.dim, params_for(4, dir.file("empty.dat"), kSeed)); + // Somewhere to spill is the caller's to provide: a path that is not a + // directory is refused rather than half-written. + SeismicIndex no_dir(corpus.dim, + params_for(4, dir.file("not-a-directory"), kSeed)); + no_dir.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + EXPECT_THROW(no_dir.build(), std::invalid_argument); + + // An empty corpus spills no windows at all, which would map back to no + // lists, so it is refused rather than silently producing an empty index. + SeismicIndex empty(corpus.dim, params_for(4, dir.scratch(), kSeed)); EXPECT_THROW(empty.build(), std::invalid_argument); - // Same for a disk index, whose spill would map back to no lists at all. - DiskSeismicIndex empty_disk( - corpus.dim, params_for(4, dir.file("empty_disk.dat"), kSeed)); + // Same for a disk index, which reaches the same build. + DiskSeismicIndex empty_disk(corpus.dim, + params_for(4, dir.scratch(), kSeed)); EXPECT_THROW(empty_disk.build(), std::invalid_argument); } From 51a0d1829dcfdf5ec231f1761740783dd5424a60 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 3 Sep 2026 06:24:57 +0000 Subject: [PATCH 13/15] Let the spill own its own file, and take the build out of MmapIndex Four things the previous shape got wrong. The spill's file and its mapping are one lifetime, so they are one type now -- ClusteredListsSpill, which unmaps then removes, in that order, and re-checks the name against the prefix and suffix it writes so it can only ever delete a file it created. An index holds one instead of a mapping plus a path, and MmapIndex needs no destructor. build_clustered_lists was a method on MmapIndex, which is a mapping owner rather than a builder. It is a free function beside the spill it uses, and each type's build() passes its own corpus and dimension. SparseVectorsConfig carried an element width the corpus already knows, so callers were passing the same fact twice and the build checked the two against each other. for_each_clustered_window now takes the corpus and the dimension -- which is the list count, and the one thing an empty corpus cannot supply -- and reads the width off SparseVectors. Plus a test that the caller's scratch directory is left alone: files that were there before the build survive it, including one named like a spill, which the build did not create and so must not remove. Comments trimmed to the reasons. Signed-off-by: Liyun Xiu --- nsparse/disk_seismic_index_base.cpp | 26 ++-- nsparse/mmap_index.h | 70 +---------- nsparse/seismic_batched_build.cpp | 133 ++++++++++++--------- nsparse/seismic_batched_build.h | 109 ++++++++++------- nsparse/seismic_common.cpp | 58 ++++----- nsparse/seismic_common.h | 85 +++++-------- nsparse/seismic_index.cpp | 7 +- nsparse/seismic_scalar_quantized_index.cpp | 11 +- tests/seismic_batched_build_test.cpp | 32 +++++ 9 files changed, 255 insertions(+), 276 deletions(-) diff --git a/nsparse/disk_seismic_index_base.cpp b/nsparse/disk_seismic_index_base.cpp index 09182b1..3b4cc7b 100644 --- a/nsparse/disk_seismic_index_base.cpp +++ b/nsparse/disk_seismic_index_base.cpp @@ -48,9 +48,9 @@ void DiskSeismicIndexBase::add(idx_t n, const idx_t* indptr, // 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 = element_size, - .dimension = static_cast(dimension_)}); + vectors_ = std::make_unique( + SparseVectorsConfig{.element_size = element_size, + .dimension = static_cast(dimension_)}); } std::vector scratch; const uint8_t* codes = encode_values(values, nnz, scratch); @@ -60,10 +60,9 @@ void DiskSeismicIndexBase::add(idx_t n, const idx_t* indptr, } void DiskSeismicIndexBase::build() { - clustered_inverted_lists = build_clustered_lists( - {.element_size = code_element_size(), - .dimension = static_cast(get_dimension())}, - cluster_parameter_); + clustered_inverted_lists = detail::build_clustered_lists( + get_vectors(), static_cast(get_dimension()), cluster_parameter_, + &batch_spill_); } auto DiskSeismicIndexBase::search(idx_t n, const idx_t* indptr, @@ -161,9 +160,8 @@ void DiskSeismicIndexBase::write_index(IOWriter* 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 = code_element_size(), - .dimension = static_cast(dimension_)}); + SparseVectors empty_vectors({.element_size = code_element_size(), + .dimension = static_cast(dimension_)}); const SparseVectors& v = vectors_ != nullptr ? *vectors_ : empty_vectors; detail::InlineForwardIndex forward(clustered_inverted_lists, v); forward.serialize(io_writer); @@ -192,12 +190,12 @@ void DiskSeismicIndexBase::load_mapped_payload(MmapCursor* cursor, clustered_inverted_lists = std::move(inv_list_writer.release()); fwd_ = std::move(forward); // 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. + // from `mapped`, which is alive here), let the concrete index reject a + // width mismatch before we commit. validate_mapped_payload(); - // mapped_file_ last: the summaries and the forward index borrow from it, and - // moving it does not move the mapping. + // mapped_file_ last: the summaries and the forward index borrow from it, + // and moving it does not move the mapping. mapped_file_ = std::move(mapped); } diff --git a/nsparse/mmap_index.h b/nsparse/mmap_index.h index c800cd3..2a2e117 100644 --- a/nsparse/mmap_index.h +++ b/nsparse/mmap_index.h @@ -17,11 +17,8 @@ #include #include #include -#include #include -#include -#include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/index.h" #include "nsparse/seismic_batched_build.h" #include "nsparse/seismic_common.h" @@ -35,28 +32,6 @@ class MmapIndex : public Index { public: explicit MmapIndex(int dim = 0) : Index(dim) {} - ~MmapIndex() override { - // A batched build's spill outlives build(), because the posting lists - // borrow from it, so removing it falls to whoever holds them. Only - // reached where a mapped file cannot be unlinked (Windows); elsewhere - // spill_clustered_lists has already unlinked it and left this empty. - if (batch_scratch_path_.empty()) { - return; - } - // Unmapped here rather than left to the member's own destructor, which - // runs after this body: the file cannot go while it is still mapped. - // Whatever borrowed from it lives in a derived class, already - // destroyed. - batch_mapped_file_ = MmapFile{}; - std::error_code ignored; - std::filesystem::remove(batch_scratch_path_, ignored); - } - - MmapIndex(const MmapIndex&) = delete; - MmapIndex& operator=(const MmapIndex&) = delete; - MmapIndex(MmapIndex&&) = delete; - MmapIndex& operator=(MmapIndex&&) = delete; - void read_csr(const char* file_path, Residency residency = Residency::kInMemory) override { switch (residency) { @@ -76,31 +51,6 @@ class MmapIndex : public Index { } protected: - // The build every seismic-family type runs, so the batching decision is - // made once rather than per type. - // - // With both batch knobs set (see BatchClusteringOption) it clusters one - // term window at a time, spilling to scratch and borrowing the finished - // lists back from it, so the peak is one window rather than the whole - // corpus; otherwise it is the ordinary whole-corpus build. Either way the - // lists come back complete and in term order, and writing an index file - // stays write_index's job -- the spill is not one, and build() produces no - // file a caller keeps. - std::vector build_clustered_lists( - const SparseVectorsConfig& config, - const SeismicClusterParameters& params) { - if (params.batch_clustering.effective_batch_size() <= 1) { - return detail::build_inverted_lists_clusters(get_vectors(), config, - params); - } - detail::SpilledLists spilled = detail::spill_clustered_lists( - get_vectors(), config, params, - params.batch_clustering.batch_file_output_path, - &batch_mapped_file_); - batch_scratch_path_ = std::move(spilled.scratch_path); - return std::move(spilled.lists); - } - // The mapping borrowed buffers point into: a native CSR file via read_csr, // or a serialized index file. Those sources are mutually exclusive, so one // member serves both. @@ -114,21 +64,13 @@ class MmapIndex : public Index { // mapped_file_ when mapped. get_vectors() cannot tell the two apart. std::unique_ptr vectors_; - // A second mapping, for the spill a batched build wrote its clustered lists - // to and then borrows them back from. Separate from mapped_file_ rather - // than replacing it, because the two coexist: the corpus may itself be a - // mapping that vectors_ is still borrowing from, and giving that up would - // leave the index unable to score anything. + // A batched build's spill, which its posting lists borrow from. Separate + // from mapped_file_ because the two coexist: the corpus may itself be a + // mapping vectors_ is still borrowing from. // - // Whatever borrows from this lives in the derived class, and derived - // members are destroyed before base ones, so the borrowers are always gone - // first. - MmapFile batch_mapped_file_; - - // The spill behind batch_mapped_file_, when the platform would not let it - // be unlinked while mapped. Empty otherwise, which is the usual case. See - // the destructor. - std::string batch_scratch_path_; + // Its borrowers live in the derived class, destroyed before base members, + // so they are always gone before the spill is released. + detail::ClusteredListsSpill batch_spill_; private: // Values are borrowed at their stored width, so a quantizing index cannot diff --git a/nsparse/seismic_batched_build.cpp b/nsparse/seismic_batched_build.cpp index 95a1df3..9066896 100644 --- a/nsparse/seismic_batched_build.cpp +++ b/nsparse/seismic_batched_build.cpp @@ -10,6 +10,7 @@ #include "nsparse/seismic_batched_build.h" #include +#include #include #include #include @@ -30,41 +31,42 @@ namespace nsparse::detail { namespace { -// A spill file of this build's own inside `dir`. Named uniquely rather than -// fixed, so concurrent builds sharing a scratch directory cannot overwrite each +// How a spill is named. Checked again before removal, so this code can only +// ever delete a file it wrote. +constexpr const char* kSpillPrefix = "nsparse-clustered-lists-"; +constexpr const char* kSpillSuffix = ".tmp"; + +// Unique per build, so builds sharing a scratch directory cannot overwrite each // other's windows. -std::string scratch_file_path(const std::string& dir) { - const auto token = static_cast(std::random_device{}()) << 32U | - std::random_device{}(); +std::string spill_path(const std::string& dir) { + std::random_device entropy; + const auto token = static_cast(entropy()) << 32U | entropy(); return (std::filesystem::path(dir) / - ("nsparse-clustered-lists-" + std::to_string(token) + ".tmp")) + (kSpillPrefix + std::to_string(token) + kSpillSuffix)) .string(); } +bool is_spill_path(const std::string& path) { + const std::string name = std::filesystem::path(path).filename().string(); + return name.starts_with(kSpillPrefix) && name.ends_with(kSpillSuffix); +} + // The posting-list section -- [count][list...], the layout -// SeismicInvertedListsWriter produces -- streamed into `writer` one window at a -// time. -// -// One writer for the whole section, windows serialized straight into it rather -// than written separately and concatenated: serialize() pads each array -// relative to the writer's current offset (see io/align.h), so bytes produced -// by a writer that started at 0 carry the wrong padding once appended at some -// other offset. -void stream_clustered_lists(const SparseVectors* vectors, - const SparseVectorsConfig& config, +// SeismicInvertedListsWriter produces -- streamed into `writer` a window at a +// time. One writer for the whole section: serialize() pads each array relative +// to the writer's offset (io/align.h), so separately written windows would +// carry the wrong padding once concatenated. +void stream_clustered_lists(const SparseVectors* vectors, size_t dimension, const SeismicClusterParameters& params, IOWriter* writer) { - // The list count, exactly where SeismicInvertedListsWriter::serialize puts - // it. It is the whole dimension, known before any window is built, which is - // what lets the lists be streamed after it rather than counted first. - size_t n_lists = config.dimension; + // The whole dimension, known before any window is built, which is what lets + // the lists follow it rather than be counted first. + size_t n_lists = dimension; writer->write(&n_lists, sizeof(size_t), 1); - // Windows arrive in ascending term order, so appending each in turn - // produces the same byte sequence as writing every list at once. size_t next_term = 0; for_each_clustered_window( - vectors, config, params, + vectors, dimension, params, [&](size_t term_begin, std::vector&& clusters) { if (term_begin != next_term) { // The layout carries no per-list offsets, so a gap or a repeat @@ -78,20 +80,54 @@ void stream_clustered_lists(const SparseVectors* vectors, next_term = term_begin + clusters.size(); // clusters freed on return, before the next window is built. }); - if (next_term != config.dimension) { - throw std::runtime_error( - "spill_clustered_lists: spilled " + std::to_string(next_term) + - " of " + std::to_string(config.dimension) + " posting lists"); + if (next_term != dimension) { + throw std::runtime_error("spill_clustered_lists: spilled " + + std::to_string(next_term) + " of " + + std::to_string(dimension) + " posting lists"); } } } // namespace -SpilledLists spill_clustered_lists(const SparseVectors* vectors, - const SparseVectorsConfig& config, - const SeismicClusterParameters& params, - const std::string& scratch_dir, - MmapFile* into) { +void ClusteredListsSpill::adopt(const std::string& path) { + MmapFile mapped(path); + std::error_code failed; + std::filesystem::remove(path, failed); + // Committed after the mapping succeeded, so a failed open leaves nothing + // half-owned. + release(); + mapping_ = std::move(mapped); + path_ = failed ? path : std::string(); +} + +void ClusteredListsSpill::release() { + if (path_.empty()) { + return; + } + const std::string path = std::move(path_); + path_.clear(); + mapping_ = MmapFile{}; + if (is_spill_path(path)) { + std::error_code ignored; + std::filesystem::remove(path, ignored); + } +} + +std::vector build_clustered_lists( + const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params, ClusteredListsSpill* spill) { + const auto& batch = params.batch_clustering; + if (batch.effective_batch_size() <= 1) { + return build_inverted_lists_clusters(vectors, dimension, params); + } + return spill_clustered_lists(vectors, dimension, params, + batch.batch_file_output_path, spill); +} + +std::vector spill_clustered_lists( + const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params, const std::string& scratch_dir, + ClusteredListsSpill* into) { if (!std::filesystem::is_directory(scratch_dir)) { throw std::invalid_argument( "spill_clustered_lists: batch_file_output_path must be an existing " @@ -104,38 +140,21 @@ SpilledLists spill_clustered_lists(const SparseVectors* vectors, "spill"); } - const std::string path = scratch_file_path(scratch_dir); + const std::string path = spill_path(scratch_dir); { - // Closed before the mapping is taken: the writer buffers, and what is - // not flushed is not in the file to map. + // Closed before the file is mapped: the writer buffers. FileIOWriter writer(const_cast(path.c_str())); - stream_clustered_lists(vectors, config, params, &writer); + stream_clustered_lists(vectors, dimension, params, &writer); writer.close(); } + into->adopt(path); - MmapFile mapped(path); - // The cursor starts at the section, which is the whole file: a spill has no - // header for it to sit behind. Absolute offsets are what serialize() padded - // against, and here they are the section's own. - MmapCursor cursor(mapped.data(), mapped.size()); + // The section is the whole file, and absolute offsets are the ones + // serialize() padded against, so the cursor starts where the writer did. + MmapCursor cursor(into->mapping().data(), into->mapping().size()); SeismicInvertedListsWriter lists; lists.mmap_deserialize(&cursor); - - // Unlinked now rather than when the index is done with it: the mapping - // keeps the bytes alive wherever unlinking an open file is allowed, so - // scratch cannot outlive the process even if it dies mid-build. Where it is - // not allowed, the path goes back to the caller to remove after the - // mapping. - std::error_code failed; - std::filesystem::remove(path, failed); - - SpilledLists spilled; - spilled.lists = std::move(lists.release()); - spilled.scratch_path = failed ? path : std::string(); - // Committed once the walk succeeded, so a truncated spill cannot leave the - // caller holding lists that point into a mapping it never took. - *into = std::move(mapped); - return spilled; + return std::move(lists.release()); } } // namespace nsparse::detail diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h index 32fc916..685fa1a 100644 --- a/nsparse/seismic_batched_build.h +++ b/nsparse/seismic_batched_build.h @@ -10,7 +10,9 @@ #ifndef SEISMIC_BATCHED_BUILD_H #define SEISMIC_BATCHED_BUILD_H +#include #include +#include #include #include "nsparse/cluster/inverted_list_clusters.h" @@ -20,52 +22,79 @@ namespace nsparse::detail { -// What a spilled build hands back. -struct SpilledLists { - // Every term's clustered posting list, in term order -- the same thing - // build_inverted_lists_clusters returns, borrowed from the spill's mapping - // rather than allocated. - std::vector lists; - // The spill file, when it is still on disk to be removed. Empty when it was - // unlinked as soon as it was mapped, which is what happens wherever a - // mapped file can be unlinked; where it cannot (Windows), whoever holds the - // mapping has to remove this after releasing it. - std::string scratch_path; +// The temporary file a batched build spilled its clustered lists to, and the +// mapping those lists borrow from. Owns both, so an index that holds one keeps +// its lists valid and cleans up after itself. +class ClusteredListsSpill { +public: + ClusteredListsSpill() = default; + ~ClusteredListsSpill() { release(); } + + ClusteredListsSpill(const ClusteredListsSpill&) = delete; + ClusteredListsSpill& operator=(const ClusteredListsSpill&) = delete; + ClusteredListsSpill(ClusteredListsSpill&& other) noexcept + : path_(std::move(other.path_)), mapping_(std::move(other.mapping_)) { + other.path_.clear(); + } + ClusteredListsSpill& operator=(ClusteredListsSpill&& other) noexcept { + if (this != &other) { + release(); + path_ = std::move(other.path_); + mapping_ = std::move(other.mapping_); + other.path_.clear(); + } + return *this; + } + + // Takes a spill this build just wrote and maps it. `path` is unlinked here + // when the platform allows it while mapped, so a crash cannot strand + // scratch; otherwise it is remembered and removed on release(). + void adopt(const std::string& path); + + [[nodiscard]] const MmapFile& mapping() const { return mapping_; } + +private: + // Unmaps, then removes the file -- in that order, since Windows cannot + // unlink a mapped file. Only ever a path adopt() created, and re-checked + // against the spill naming, because this deletes. + void release(); + + std::string path_; // empty unless a removal is still owed + MmapFile mapping_; }; -// Clusters the corpus one contiguous term window at a time, spilling each -// window's lists to a temporary file in `scratch_dir` and mapping them back, so -// the caller ends up holding every list without two windows ever having been -// resident at once. +// The build every seismic-family index runs, and the one place the batching +// decision is made. Clusters one contiguous term window at a time into `spill` +// when both batch knobs are set, whole-corpus otherwise; either way the lists +// come back complete and in term order. // -// The usual build holds two whole-corpus intermediates -- the inverted lists -// (every posting) and then the clustered posting lists -- so its peak memory -// scales with the corpus's non-zeros, and a corpus whose posting lists do not -// fit in RAM cannot be indexed at all. for_each_clustered_window bounds the -// first to one window; spilling each window's clusters and dropping them, which -// is what this does, bounds the second. What is left resident is the forward -// corpus (which the caller already holds, at whatever residency SparseVectors -// was given) plus one window. +// `dimension` is the index's, not the corpus's: it is the list count, and an +// empty corpus still has one. The element width comes from `vectors`. +std::vector build_clustered_lists( + const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params, ClusteredListsSpill* spill); + +// Clusters one term window at a time, spilling each window's lists into +// `scratch_dir` and freeing them, then maps the finished lists back out. // -// The spill is scratch, not an index: it carries the posting-list section and -// nothing else -- no index header, no forward vectors -- and nothing outside -// this build ever reads it. Writing an index file remains write_index's job, -// from the lists this returns; a build that borrows its lists from scratch -// serializes byte-for-byte what a whole-corpus build would have. +// A whole-corpus build holds two intermediates that scale with the corpus's +// non-zeros -- the inverted lists and then the clustered lists -- which is what +// puts a ceiling on the corpus an index can be built from. +// for_each_clustered_window bounds the first to a window; spilling bounds the +// second. What stays resident is the corpus plus one window. // -// Identical to the unbatched build for a fixed `params.seed`, and identical -// whatever the window count is, because every list's k-means seed comes from -// its own global term id -- see for_each_clustered_window. +// The spill is scratch, not an index: posting lists only, no header, read by +// nothing else. Serializing an index remains write_index's job, and at a fixed +// `params.seed` what it then writes is byte-for-byte what a whole-corpus build +// would have produced -- every list's k-means seed comes from its own global +// term id, so the window count cannot leak into the output. // -// `into` takes the mapping the returned lists borrow from, and so must outlive -// them. Throws if `scratch_dir` is not a directory, or if the corpus is empty: -// there would be no windows to spill, and an empty spill maps back to no lists -// at all. -SpilledLists spill_clustered_lists(const SparseVectors* vectors, - const SparseVectorsConfig& config, - const SeismicClusterParameters& params, - const std::string& scratch_dir, - MmapFile* into); +// Throws if `scratch_dir` is not a directory, or if the corpus is empty: there +// would be no windows to spill. +std::vector spill_clustered_lists( + const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params, const std::string& scratch_dir, + ClusteredListsSpill* into); } // namespace nsparse::detail diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp index 9f447b4..ca024fe 100644 --- a/nsparse/seismic_common.cpp +++ b/nsparse/seismic_common.cpp @@ -35,39 +35,29 @@ struct TermWindow { }; // How much more a posting costs once clustered than while being scattered into -// an inverted list, per unit. A window's memory has two peaks: filling it holds -// every posting of its terms at a few bytes each, and clustering it holds the -// pruned survivors as clusters and summaries, which are an order of magnitude -// bulkier per posting. -// -// Only the ratio matters, and only roughly: the cost curve is a shallow basin, -// so assuming 8x or 32x here instead of 16x costs a fraction of the benefit and -// still beats weighting either phase alone. It is deliberately not derived from -// alpha/beta/dimension, which would be a model of summarize() that this does not -// need to be right about. +// an inverted list. Only the ratio matters, and only roughly: the cost curve is +// a shallow basin, so 8x or 32x instead of 16x costs a fraction of the benefit +// and still beats weighting either phase alone. Deliberately not derived from +// alpha/beta/dimension, which would be a model of summarize() to keep correct. constexpr size_t kClusterCostRatio = 16; // Cuts [0, dimension) into at most `batches` windows of near-equal estimated // memory, from the exact per-term counts. // -// Equal width would not do, because term frequencies are heavily skewed: a -// natural-language corpus puts orders of magnitude more postings on its heaviest -// term than on its mean one. Peak memory is set by the largest window, not the -// average one, so an uneven split wastes most of what batching could save. +// Not equal width, because term frequencies are heavily skewed and the peak is +// set by the largest window, so an uneven split wastes most of what batching +// could save. // -// What to even out is neither phase alone but their sum. Weighting raw counts -// balances the fill and unbalances the clustering, which is the more expensive -// phase. Weighting min(count, lambda) -- what survives pruning, and so what -// clustering holds -- balances that phase perfectly but concentrates the heavy -// terms, leaving one window holding several times the mean raw postings, which -// then becomes the peak. Weighting both together at their relative cost balances -// what is actually resident, and measures best of the four on a skewed corpus. +// What to even out is both phases' sum, not either alone: weighting raw counts +// balances the fill and unbalances the costlier clustering, and weighting +// min(count, lambda) balances clustering but concentrates the heavy terms, +// whose fill then becomes the peak. // -// Windows stay contiguous and ascending, which is what lets the clustered lists -// be appended to a file as each window finishes: the layout carries no per-list -// offsets, so a list's position in the file is its term order. Grouping terms -// by a hash (term % batches) balances comparably, but scatters each window's -// terms across the file, which the streaming write cannot express. +// Windows stay contiguous and ascending, which is what lets each window's lists +// be appended to the spill as it finishes: the layout carries no per-list +// offsets, so a list's position in the file is its term order. Hashing terms +// into windows balances comparably but scatters them, which the append cannot +// express. // // Bounds are size_t, not term_t: dimension may be up to 65536 (term_t is // uint16), so a term_t window boundary would wrap and silently drop terms. @@ -272,21 +262,15 @@ std::vector cluster_window( } // namespace -void for_each_clustered_window(const SparseVectors* vectors, - const SparseVectorsConfig& config, +void for_each_clustered_window(const SparseVectors* vectors, size_t dimension, const SeismicClusterParameters& params, const ClusteredWindowSink& sink) { if (vectors == nullptr || vectors->num_vectors() == 0) { return; } - if (vectors->get_element_size() != config.element_size) { - throw std::invalid_argument( - "for_each_clustered_window: corpus element width does not match " - "the " - "index's"); - } - const size_t dim = config.dimension; + const size_t dim = dimension; + const size_t element_size = vectors->get_element_size(); const size_t batches = std::min(params.batch_clustering.effective_batch_size(), dim); @@ -312,8 +296,8 @@ void for_each_clustered_window(const SparseVectors* vectors, for (const TermWindow& window : make_windows( term_counts, static_cast(resolved.lambda), batches)) { - WindowLists lists(term_counts, window, config.element_size); - fill_from_corpus(*vectors, window, config.element_size, &lists); + WindowLists lists(term_counts, window, element_size); + fill_from_corpus(*vectors, window, element_size, &lists); sink(window.begin, cluster_window(*vectors, lists.seal(), window, resolved)); // The window's lists and clusters are freed here, before the next diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index e0d2eed..4cd07f4 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -29,41 +29,25 @@ namespace nsparse { // How a build bounds its own memory. // -// Clustering the whole term space at once holds two intermediates for the whole -// corpus: the inverted lists (every posting) and then the clustered lists. Both -// scale with the corpus's non-zeros, which is what puts a ceiling on the corpus -// an index can be built from. Splitting the term space into `batch_size` -// contiguous windows and finishing one window before starting the next makes -// the first of those proportional to a window instead. +// A whole-corpus build holds two intermediates that scale with the corpus's +// non-zeros -- the inverted lists, then the clustered lists -- which is what +// caps the corpus an index can be built from. Windowing bounds the first to a +// window; spilling each window's clusters bounds the second. // -// `batch_file_output_path` bounds the second as well: with batch_size > 1, each -// window's clustered lists are spilled to a temporary file there and freed as -// they are produced, and the finished lists are then mapped back out of it, so -// the build never holds more than one window's worth. See -// spill_clustered_lists. -// -// The two knobs are therefore only useful together, and each is ignored without -// the other -- see effective_batch_size. +// Both knobs are needed for either to help, so each is ignored without the +// other +// -- see effective_batch_size and build_clustered_lists. struct BatchClusteringOption { - // Contiguous term windows. <= 1 means one window, i.e. no batching. Clamped - // to the dimension, since a window cannot be narrower than one term. + // Contiguous term windows. <= 1 means no batching; clamped to the + // dimension, a window being at least one term. size_t batch_size = 1; - // An existing directory the build may spill windows into. Scratch, not - // output: build() writes no index there and leaves nothing behind, and - // serializing an index remains write_index's job. Used only when batch_size - // > 1, a single window being an ordinary build with nothing to spill. + // An existing directory to spill windows into. Scratch, not output: build() + // writes no index and leaves nothing behind. std::string batch_file_output_path; - // Windows the build should actually run, which is 1 unless there is - // somewhere to write them. - // - // Splitting the term space without a path to write to bounds only the first - // intermediate: the clustered lists still accumulate for the whole corpus, - // and they are the bulkier of the two by roughly kClusterCostRatio. So it - // buys a fraction of the peak while costing a corpus pass per window, and - // the peak stays where an unbatched build's is. Not worth doing quietly on - // a caller's behalf, so `batch_size` alone is honoured as no batching at - // all. + // Windows to actually run: 1 unless there is somewhere to spill them, since + // windowing alone leaves the bulkier intermediate whole-corpus and costs a + // corpus pass per window. [[nodiscard]] size_t effective_batch_size() const { return batch_file_output_path.empty() ? 1 : std::max(1, batch_size); @@ -184,43 +168,36 @@ inline int calculate_beta(int beta, int lambda) { } // Clusters and summarizes the posting lists of one term window at a time, -// handing each window to `sink` in ascending term order. +// handing each window to `sink` in ascending term order. The one place the +// seismic family's build work lives: every index type reaches it, through +// build_inverted_lists_clusters below or through build_clustered_lists. // -// This is the one place the seismic family's build work lives: every index type -// reaches it, either through build_inverted_lists_clusters below or through the -// spilling build in seismic_batched_build.h. The element width comes from -// `config`, so a quantizing index gets the same treatment as a float one -- the -// values in `vectors` are already encoded by the time they arrive here. +// `dimension` is the index's declared term space, which an empty corpus still +// has; the element width comes from `vectors`, already encoded by add(), so a +// quantizing index needs no special case. // -// `sink` receives the window's global first term and its lists, and must not -// hold on to them: they are freed as soon as it returns, which is what bounds -// the memory. Windows come from -// params.batch_clustering.effective_batch_size(), so a caller that set a batch -// size but no output path gets one window -- see BatchClusteringOption. +// `sink` must not hold on to what it is given: the window is freed as soon as +// it returns, which is what bounds the memory. Windows come from +// params.batch_clustering.effective_batch_size(). // // Every window's lambda and beta are computed from the GLOBAL corpus, and every // list's k-means seed from its own GLOBAL term id, so the window count cannot -// change what is produced -- see the batched-build tests, which assert file -// equality against an unbatched build. +// change what is produced. using ClusteredWindowSink = std::function&& clusters)>; -void for_each_clustered_window(const SparseVectors* vectors, - const SparseVectorsConfig& config, +void for_each_clustered_window(const SparseVectors* vectors, size_t dimension, const SeismicClusterParameters& params, const ClusteredWindowSink& sink); -// Every term's clustered posting list, in term order. The whole-corpus form of -// for_each_clustered_window, and so the unbatched build: the result is retained -// in full, which is what bounds this by the clustered lists rather than by a -// window. Callers reach it without an output path, which is why that is also -// the case where the window count is dropped. +// The whole-corpus form: every window retained, so this is bounded by the +// clustered lists rather than by one window. inline std::vector build_inverted_lists_clusters( - const SparseVectors* vectors, const SparseVectorsConfig& config, - const SeismicClusterParameters& seismic_cluster_params) { - std::vector clustered(config.dimension); + const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params) { + std::vector clustered(dimension); for_each_clustered_window( - vectors, config, seismic_cluster_params, + vectors, dimension, params, [&clustered](size_t term_begin, std::vector&& window) { std::move(window.begin(), window.end(), diff --git a/nsparse/seismic_index.cpp b/nsparse/seismic_index.cpp index e1e23bb..d0fd025 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -142,10 +142,9 @@ void SeismicIndex::add(idx_t n, const idx_t* indptr, const term_t* indices, } void SeismicIndex::build() { - clustered_inverted_lists = build_clustered_lists( - {.element_size = kElementSize, - .dimension = static_cast(get_dimension())}, - cluster_parameter_); + clustered_inverted_lists = detail::build_clustered_lists( + get_vectors(), static_cast(get_dimension()), cluster_parameter_, + &batch_spill_); } auto SeismicIndex::search(idx_t n, const idx_t* indptr, const term_t* indices, diff --git a/nsparse/seismic_scalar_quantized_index.cpp b/nsparse/seismic_scalar_quantized_index.cpp index 1d918d2..100f702 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -191,12 +191,11 @@ ScalarQuantizer SeismicScalarQuantizedIndex::query_quantizer( } void SeismicScalarQuantizedIndex::build() { - // The codes in `vectors_` are already quantized by add(), so the shared - // build needs no knowledge of the quantizer beyond its width. - clustered_inverted_lists = build_clustered_lists( - {.element_size = sq_.bytes_per_value(), - .dimension = static_cast(get_dimension())}, - cluster_parameter_); + // add() has already quantized the values, so the shared build reads their + // width off the corpus and needs to know nothing of the quantizer. + clustered_inverted_lists = detail::build_clustered_lists( + get_vectors(), static_cast(get_dimension()), cluster_parameter_, + &batch_spill_); } auto SeismicScalarQuantizedIndex::search(idx_t n, const idx_t* indptr, diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index cdb8b09..6ed05cb 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -245,6 +245,38 @@ TEST(SeismicBatchedBuild, LeavesNothingInTheScratchDirectory) { in_memory(corpus, dir.file("mem.dat"))); } +// The build deletes its own spill and nothing else. The directory is the +// caller's, so whatever else lives in it -- including a file named like a +// spill, which the build did not create -- has to still be there afterwards. +TEST(SeismicBatchedBuild, LeavesOtherFilesInTheScratchDirectoryAlone) { + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/3); + TempDir dir("scratch_others"); + const std::string scratch = dir.scratch(); + const std::array bystanders = { + scratch + "/keep-me.txt", scratch + "/index.dat", + scratch + "/nsparse-clustered-lists-999.tmp"}; + for (const std::string& path : bystanders) { + std::ofstream(path, std::ios::binary) << "not the build's"; + } + + { + SeismicIndex index(corpus.dim, params_for(8, scratch, kSeed)); + index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index.build(); + // Destroyed here, which is when a spill that could not be unlinked + // while mapped is removed. + } + + for (const std::string& path : bystanders) { + EXPECT_TRUE(std::filesystem::exists(path)) << path; + EXPECT_EQ(std::filesystem::file_size(path), 15U) << path; + } + EXPECT_EQ(std::distance(std::filesystem::directory_iterator(scratch), + std::filesystem::directory_iterator{}), + static_cast(bystanders.size())); +} + // The two knobs are only useful together. A window count with nowhere to spill // the windows bounds the fill intermediate while leaving the clustered lists to // accumulate -- a corpus pass per window for a fraction of the peak -- so it is From 922c10db1ad73d10bee65c81fd165830d6b2a29b Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 3 Sep 2026 06:37:52 +0000 Subject: [PATCH 14/15] tests: assert when the spill is gone, not how it went The scratch-directory test read the spill's removal as immediate, which it is only where a mapped file can be unlinked. On Windows it cannot, so the spill sits there until the index releases it -- exactly what the code says it does, and what the CI failure reported. So the assertion is now the contract: the directory is empty once the index is gone, and until then the spill is at most its sole occupant. The lists stay readable throughout on both, which is the part that matters. Signed-off-by: Liyun Xiu --- python_tests/test_seismic_batched_build.py | 12 ++++++--- tests/seismic_batched_build_test.cpp | 29 ++++++++++++++-------- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/python_tests/test_seismic_batched_build.py b/python_tests/test_seismic_batched_build.py index 93daf31..856a31b 100644 --- a/python_tests/test_seismic_batched_build.py +++ b/python_tests/test_seismic_batched_build.py @@ -92,19 +92,23 @@ def test_matches_in_memory_build(kind, corpus, tmp_path): def test_scratch_directory_is_left_empty(corpus, tmp_path): - """The spill is scratch: unlinked as soon as it is mapped. + """The spill is scratch: whatever is left of it goes with the index. - The lists stay readable from the mapping, so the index is still servable - while the directory the caller lent is already empty again. + Its lists stay readable while the index lives, so it is servable and + writable throughout; when the spill is removed (immediately where a mapped + file can be unlinked, on release where it cannot) is the platform's business. """ scratch = scratch_dir(tmp_path) index = batched_index(corpus, tmp_path, 8) - assert list(scratch.iterdir()) == [] assert index.num_vectors() == corpus.n + assert len(list(scratch.iterdir())) <= 1 nsparse.write_index(index, str(tmp_path / "out.idx")) assert (tmp_path / "out.idx").stat().st_size > 0 + del index + assert list(scratch.iterdir()) == [] + @pytest.mark.parametrize("kind", ["seismic", "disk_seismic"]) def test_batched_index_is_searchable_after_build( diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index 6ed05cb..2fcc386 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -224,23 +224,32 @@ TEST(SeismicBatchedBuild, BatchedBuildIsIdenticalAcrossBatchCounts) { EXPECT_EQ(one, batched(corpus, 1000, dir, dir.file("b1000.dat"))); } -// The spill is scratch, and nothing outlives the build: it is unlinked as soon -// as it is mapped, so the directory the caller lent is empty again while the -// index is still serving from those very bytes. +// The spill is scratch: whatever is left of it goes with the index, and while +// the index lives its lists stay readable from the mapping either way. +// +// When it goes is the platform's business, not the contract's -- unlinked the +// moment it is mapped where that is allowed, removed on release where it is not +// (Windows) -- so this asserts the directory is empty once the index is gone, +// and only that the spill is the sole occupant before then. TEST(SeismicBatchedBuild, LeavesNothingInTheScratchDirectory) { Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/3); TempDir dir("scratch"); const std::string scratch = dir.scratch(); - SeismicIndex index(corpus.dim, params_for(8, scratch, kSeed)); - index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), - corpus.values.data()); - index.build(); + { + SeismicIndex index(corpus.dim, params_for(8, scratch, kSeed)); + index.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + index.build(); + + EXPECT_LE(std::distance(std::filesystem::directory_iterator(scratch), + std::filesystem::directory_iterator{}), + 1); + // Serialized from lists that are still borrowed from the spill. + write_index(&index, const_cast(dir.file("out.dat").c_str())); + } EXPECT_TRUE(std::filesystem::is_empty(scratch)); - // And the lists are still readable, which is the point of unlinking rather - // than deleting: the mapping keeps the bytes alive. - write_index(&index, const_cast(dir.file("out.dat").c_str())); EXPECT_EQ(read_file(dir.file("out.dat")), in_memory(corpus, dir.file("mem.dat"))); } From 35861411e5d8ce8b9da45e8f508aaa619852d40c Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 3 Sep 2026 07:16:33 +0000 Subject: [PATCH 15/15] Own the spill from creation, so a failed build removes it The spill was written first and adopted afterwards, which left everything in between unguarded: the counting pass rejecting a term outside the dimension, a window arriving short, a full disk, a failed flush, an allocation failing in a build that exists for memory-tight corpora. Any of those left a half-written file in the caller's scratch directory. The first of them is a case the tests already cover, so this was reachable, not theoretical. ClusteredListsSpill::write_and_map now creates the file, hands it to the caller's writer, and maps it, owning it throughout: the path is recorded before the file exists, and any throw releases it on the way out. Unlinking after the map still covers a crash after the write; a crash during it is what the distinctive name is for. Signed-off-by: Liyun Xiu --- nsparse/seismic_batched_build.cpp | 45 ++++++++++++++++++---------- nsparse/seismic_batched_build.h | 20 +++++++++---- tests/seismic_batched_build_test.cpp | 19 ++++++++++++ 3 files changed, 62 insertions(+), 22 deletions(-) diff --git a/nsparse/seismic_batched_build.cpp b/nsparse/seismic_batched_build.cpp index 9066896..b939c55 100644 --- a/nsparse/seismic_batched_build.cpp +++ b/nsparse/seismic_batched_build.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -89,15 +90,32 @@ void stream_clustered_lists(const SparseVectors* vectors, size_t dimension, } // namespace -void ClusteredListsSpill::adopt(const std::string& path) { - MmapFile mapped(path); - std::error_code failed; - std::filesystem::remove(path, failed); - // Committed after the mapping succeeded, so a failed open leaves nothing - // half-owned. +void ClusteredListsSpill::write_and_map( + const std::string& dir, + const std::function& write_section) { release(); - mapping_ = std::move(mapped); - path_ = failed ? path : std::string(); + // Owned before it exists, so every path out of here removes it. + path_ = spill_path(dir); + try { + { + // Closed before the file is mapped: the writer buffers, and close() + // is what reports a failed flush. + FileIOWriter writer(const_cast(path_.c_str())); + write_section(&writer); + writer.close(); + } + mapping_ = MmapFile(path_); + } catch (...) { + release(); + throw; + } + // Unlinked now rather than on release: the mapping keeps the bytes alive + // wherever that is allowed, so a crash cannot strand scratch either. + std::error_code failed; + std::filesystem::remove(path_, failed); + if (!failed) { + path_.clear(); + } } void ClusteredListsSpill::release() { @@ -140,14 +158,9 @@ std::vector spill_clustered_lists( "spill"); } - const std::string path = spill_path(scratch_dir); - { - // Closed before the file is mapped: the writer buffers. - FileIOWriter writer(const_cast(path.c_str())); - stream_clustered_lists(vectors, dimension, params, &writer); - writer.close(); - } - into->adopt(path); + into->write_and_map(scratch_dir, [&](IOWriter* writer) { + stream_clustered_lists(vectors, dimension, params, writer); + }); // The section is the whole file, and absolute offsets are the ones // serialize() padded against, so the cursor starts where the writer did. diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h index 685fa1a..4237305 100644 --- a/nsparse/seismic_batched_build.h +++ b/nsparse/seismic_batched_build.h @@ -11,11 +11,13 @@ #define SEISMIC_BATCHED_BUILD_H #include +#include #include #include #include #include "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/io/io.h" #include "nsparse/seismic_common.h" #include "nsparse/sparse_vectors.h" #include "nsparse/utils/mmap_file.h" @@ -46,17 +48,23 @@ class ClusteredListsSpill { return *this; } - // Takes a spill this build just wrote and maps it. `path` is unlinked here - // when the platform allows it while mapped, so a crash cannot strand - // scratch; otherwise it is remembered and removed on release(). - void adopt(const std::string& path); + // Creates a spill in `dir`, hands it to `write_section`, and maps what was + // written. + // + // The file is owned from the moment it is created, so a throw anywhere in + // `write_section` -- a corpus the dimension does not cover, a full disk, a + // failed flush -- takes the partial spill with it. Once mapped it is + // unlinked where the platform allows that while mapped, which also covers a + // crash; where it does not, release() removes it. + void write_and_map(const std::string& dir, + const std::function& write_section); [[nodiscard]] const MmapFile& mapping() const { return mapping_; } private: // Unmaps, then removes the file -- in that order, since Windows cannot - // unlink a mapped file. Only ever a path adopt() created, and re-checked - // against the spill naming, because this deletes. + // unlink a mapped file. Only ever a path write_and_map created, and + // re-checked against the spill naming, because this deletes. void release(); std::string path_; // empty unless a removal is still owed diff --git a/tests/seismic_batched_build_test.cpp b/tests/seismic_batched_build_test.cpp index 2fcc386..d83f191 100644 --- a/tests/seismic_batched_build_test.cpp +++ b/tests/seismic_batched_build_test.cpp @@ -254,6 +254,25 @@ TEST(SeismicBatchedBuild, LeavesNothingInTheScratchDirectory) { in_memory(corpus, dir.file("mem.dat"))); } +// A build that throws part-way through spilling must not leave the half-written +// spill behind. The corpus here has a term the declared dimension does not +// cover, which the counting pass rejects after the spill file has been created. +TEST(SeismicBatchedBuild, RemovesAPartialSpillWhenTheBuildThrows) { + Corpus corpus = make_corpus(/*n_docs=*/50, /*dim=*/32, /*seed=*/1); + TempDir dir("partial"); + const std::string scratch = dir.scratch(); + + { + SeismicIndex narrow(corpus.dim / 2, params_for(4, scratch, kSeed)); + narrow.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + EXPECT_THROW(narrow.build(), std::invalid_argument); + EXPECT_TRUE(std::filesystem::is_empty(scratch)) + << "a failed build left scratch behind"; + } + EXPECT_TRUE(std::filesystem::is_empty(scratch)); +} + // The build deletes its own spill and nothing else. The directory is the // caller's, so whatever else lives in it -- including a file named like a // spill, which the build did not create -- has to still be there afterwards.