diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 269bd45..4d6a93b 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -224,6 +224,163 @@ 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, 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=/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(); // 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"); + +// 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(...); +``` + +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=/scratch", +) +index.read_csr(native, nsparse.Residency_kMmap) +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) +``` + +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 +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) 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`. + +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 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, +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. + +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 +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: identical +byte for byte at a fixed seed, and indistinguishable in latency, QPS and recall +for the random-seeded default. + +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. 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 /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 /scratch disk_seismic +``` + ## Python Bindings ### Build Python Bindings 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..efc7184 --- /dev/null +++ b/benchmarks/batched_build_mem_bench.cpp @@ -0,0 +1,441 @@ +/** + * 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] [index_type] +// batched_build_mem_bench batched \ +// [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, 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". 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 +// 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 spill is RAM, +// and the numbers are meaningless. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#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" +#include "nsparse/utils/csr_layout.h" + +namespace { + +// 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; + } + 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. +// +// 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; + } +} + +// 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); `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& 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 (!scratch_dir.empty()) { + desc += "|batch_file_output_path=" + scratch_dir; + } + 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); + 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), 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, 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()) + // 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"; +} + +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] [index_type]\n"; + return 2; + } + const std::string csr = argv[2]; + nsparse::SeismicClusterParameters params = { + .lambda = std::atoi(argv[3]), + .beta = std::atoi(argv[4]), + .alpha = static_cast(std::atof(argv[5]))}; + 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:"); + 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(); + const double build_s = now_seconds() - started; + 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"; + } + return 0; +} + +int run_batched(int argc, char** argv) { + if (argc < 9) { + std::cerr << "batched " + " [index_type]\n"; + return 2; + } + const std::string corpus_residency = argv[2]; + const std::string csr = argv[3]; + 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 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])); + + // 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. + std::unique_ptr index = + 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") { + 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(); + const long start_anon = read_status_kib("RssAnon:"); + reset_vm_hwm(); + PeakRssSampler sampler; + const double started = now_seconds(); + // 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; + + report("batched", + "type=" + index_type + " corpus=" + corpus_residency + " batches=" + + std::to_string(batched_params.batch_clustering.batch_size), + build_s, load_hwm, start_anon, sampler); + + // 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"; + } + 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; +} 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/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/disk_seismic_index_base.cpp b/nsparse/disk_seismic_index_base.cpp index 4078a06..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,11 +60,9 @@ void DiskSeismicIndexBase::add(idx_t n, const idx_t* indptr, } void DiskSeismicIndexBase::build() { - clustered_inverted_lists = detail::build_inverted_lists_clusters( - get_vectors(), - {.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, @@ -162,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); @@ -193,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/index_factory.cpp b/nsparse/index_factory.cpp index 2beeef6..946ff7c 100644 --- a/nsparse/index_factory.cpp +++ b/nsparse/index_factory.cpp @@ -55,10 +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")), - .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..2a2e117 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 @@ -24,13 +17,23 @@ #include #include #include +#include + +#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" +#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 +64,14 @@ class MmapIndex : public Index { // mapped_file_ when mapped. get_vectors() cannot tell the two apart. std::unique_ptr vectors_; + // 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. + // + // 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 // use this path. @@ -105,7 +116,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 +145,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 new file mode 100644 index 0000000..b939c55 --- /dev/null +++ b/nsparse/seismic_batched_build.cpp @@ -0,0 +1,173 @@ +/** + * 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 "nsparse/cluster/inverted_list_clusters.h" +#include "nsparse/io/file_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 { +namespace { + +// 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 spill_path(const std::string& dir) { + std::random_device entropy; + const auto token = static_cast(entropy()) << 32U | entropy(); + return (std::filesystem::path(dir) / + (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` 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 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); + + size_t next_term = 0; + for_each_clustered_window( + 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 + // would silently shift every list after it. + throw std::runtime_error( + "spill_clustered_lists: 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 != dimension) { + throw std::runtime_error("spill_clustered_lists: spilled " + + std::to_string(next_term) + " of " + + std::to_string(dimension) + " posting lists"); + } +} + +} // namespace + +void ClusteredListsSpill::write_and_map( + const std::string& dir, + const std::function& write_section) { + release(); + // 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() { + 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 " + "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"); + } + + 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. + MmapCursor cursor(into->mapping().data(), into->mapping().size()); + SeismicInvertedListsWriter lists; + lists.mmap_deserialize(&cursor); + return std::move(lists.release()); +} + +} // namespace nsparse::detail diff --git a/nsparse/seismic_batched_build.h b/nsparse/seismic_batched_build.h new file mode 100644 index 0000000..4237305 --- /dev/null +++ b/nsparse/seismic_batched_build.h @@ -0,0 +1,109 @@ +/** + * 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 +#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" + +namespace nsparse::detail { + +// 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; + } + + // 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 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 + MmapFile mapping_; +}; + +// 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. +// +// `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. +// +// 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. +// +// 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. +// +// 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 + +#endif // SEISMIC_BATCHED_BUILD_H diff --git a/nsparse/seismic_common.cpp b/nsparse/seismic_common.cpp new file mode 100644 index 0000000..ca024fe --- /dev/null +++ b/nsparse/seismic_common.cpp @@ -0,0 +1,308 @@ +/** + * 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; } +}; + +// How much more a posting costs once clustered than while being scattered into +// 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. +// +// 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 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 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. +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 count + kClusterCostRatio * std::min(count, lambda); + }; + size_t remaining_load = 0; + for (size_t count : term_counts) { + remaining_load += load_of(count); + } + + std::vector windows; + 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; +} + +// 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, size_t dimension, + const SeismicClusterParameters& params, + const ClusteredWindowSink& sink) { + if (vectors == nullptr || vectors->num_vectors() == 0) { + return; + } + + 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); + + 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( + term_counts, static_cast(resolved.lambda), batches)) { + 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 + // 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..4cd07f4 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -10,11 +10,12 @@ #ifndef SEISMIC_COMMON_H #define SEISMIC_COMMON_H +#include #include #include #include -#include #include +#include #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/cluster/random_kmeans.h" @@ -26,9 +27,31 @@ namespace nsparse { +// How a build bounds its own memory. +// +// 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. +// +// 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 no batching; clamped to the + // dimension, a window being at least one term. size_t batch_size = 1; + // 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 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); + } }; // Draw fresh entropy at build time, which makes the build unreproducible. Any @@ -39,8 +62,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 +77,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 +167,43 @@ 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. 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. +// +// `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` 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. +using ClusteredWindowSink = std::function&& clusters)>; + +void for_each_clustered_window(const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params, + const ClusteredWindowSink& sink); + +// 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) { - // 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; + const SparseVectors* vectors, size_t dimension, + const SeismicClusterParameters& params) { + std::vector clustered(dimension); + for_each_clustered_window( + vectors, dimension, 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..d0fd025 100644 --- a/nsparse/seismic_index.cpp +++ b/nsparse/seismic_index.cpp @@ -142,11 +142,9 @@ 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_)); + 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_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 b15f9a5..100f702 100644 --- a/nsparse/seismic_scalar_quantized_index.cpp +++ b/nsparse/seismic_scalar_quantized_index.cpp @@ -191,11 +191,11 @@ 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_)); + // 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, @@ -262,14 +262,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); @@ -285,9 +285,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); @@ -315,11 +314,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(); @@ -431,7 +431,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 new file mode 100644 index 0000000..856a31b --- /dev/null +++ b/python_tests/test_seismic_batched_build.py @@ -0,0 +1,188 @@ +# 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` 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 +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 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={scratch_dir(tmp_path, name)}" + ) + index = nsparse.index_factory(corpus.dim, spec) + add_corpus(index, corpus) + index.build() + return index + + +# 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 -> 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 + + 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", "disk_seismic", "disk_seismic_sq"] +) +def test_matches_in_memory_build(kind, corpus, tmp_path): + """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 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)) + + 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() + + +def test_scratch_directory_is_left_empty(corpus, tmp_path): + """The spill is scratch: whatever is left of it goes with the index. + + 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 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( + kind, corpus, queries, oracle, tmp_path +): + """build() leaves an index that serves, not an empty object. + + 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. + """ + index = batched_index(corpus, tmp_path, 8, kind=kind) + assert index.num_vectors() == corpus.n + + 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 a scratch directory the window count is ignored, not half-applied. + + 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) + 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.""" + plain = tmp_path / "plain.idx" + nsparse.write_index(make_index(f"seismic,{BASE}", corpus), str(plain)) + 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(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) + + +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={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/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..d83f191 --- /dev/null +++ b/tests/seismic_batched_build_test.cpp @@ -0,0 +1,669 @@ +/** + * 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 + +#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" +#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; + +// `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; +}; + +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.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 + // 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; + } + + // 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& 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 = scratch_dir; + 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 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 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)); + 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 as a native-layout CSR: what a mapped read consumes. +std::string write_native_csr(const Corpus& corpus, const std::string& path) { + 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; +} + +} // namespace + +// 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")), + 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 index. +TEST(SeismicBatchedBuild, BatchedBuildIsIdenticalAcrossBatchCounts) { + Corpus corpus = make_corpus(/*n_docs=*/2000, /*dim=*/200, /*seed=*/11); + TempDir dir("counts"); + const auto one = in_memory(corpus, dir.file("b1.dat")); + ASSERT_FALSE(one.empty()); + 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, batched(corpus, 1000, dir, dir.file("b1000.dat"))); +} + +// 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(); + + 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)); + EXPECT_EQ(read_file(dir.file("out.dat")), + 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. +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 +// 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"); + 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)); +} + +// 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 {}; + +TEST_P(BatchedBuildEveryType, SpilledBuildIsByteIdenticalToInMemoryBuild) { + const std::string kind = GetParam(); + Corpus corpus = make_corpus(/*n_docs=*/1500, /*dim=*/200, /*seed=*/71); + 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); + }; + + 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)); +} + +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 +// 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.n = corpus.n; + 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 = 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 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 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); + 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 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()); + 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 corpus borrowed from a mapping is the case two mappings exist for: the +// 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); + 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.scratch(), kSeed)); + index.read_csr(native.c_str(), Residency::kMmap); + index.build(); + + // Still serving: scoring reads the mapped corpus, the lists come from the + // spill's 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. +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"); + in_memory(corpus, one, /*batch_size=*/1, kRandomSeed); + batched(corpus, 10, dir, 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 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 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()))); + 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 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 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(); + batched(corpus, 4, dir, batched_path); + + std::unique_ptr disk(read_index( + const_cast(batched_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"); + 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, 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)); +} + +// 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=" + + 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())); + + EXPECT_EQ(batched(corpus, 8, dir, 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, "", kSeed)); + narrow.add(corpus.n, corpus.indptr.data(), corpus.indices.data(), + corpus.values.data()); + EXPECT_THROW(narrow.build(), std::invalid_argument); + + // 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, which reaches the same build. + DiskSeismicIndex empty_disk(corpus.dim, + params_for(4, dir.scratch(), kSeed)); + EXPECT_THROW(empty_disk.build(), std::invalid_argument); +} + +} // namespace nsparse