From 25834ea0a00a0b309692157aab08ba5666fe084a Mon Sep 17 00:00:00 2001 From: sophisid Date: Wed, 21 Oct 2026 12:12:05 +0200 Subject: [PATCH] Make Coconut streaming disk-backed with bounded LSM updates --- CONTINUITY.md | 17 + README.md | 47 +- benchmark/CMakeLists.txt | 2 + benchmark/bm_Coconut_Streaming.cpp | 32 + demos/demo_Coconut_Streaming.cpp | 8 + demos/demo_Coconut_Streaming.py | 11 +- docs/coconut-disk-streaming.md | 60 ++ docs/coconut-paper-implementation.md | 121 +++ docs/daisy-architecture.puml | 11 + docs/daisy-components.puml | 3 + docs/how-to-contribute.md | 8 + docs/test-framework-guide.md | 15 + lib/algos/CMakeLists.txt | 3 + lib/algos/Coconut.cpp | 1047 ++++++++++++++++---------- lib/algos/Coconut.hpp | 129 ++-- lib/algos/CoconutBounds.hpp | 46 ++ lib/daisy.hpp | 2 + pybinds/setup.cpp | 70 +- tests/CMakeLists.txt | 5 + tests/test_Coconut_Range.cpp | 11 + tests/test_Coconut_Streaming.cpp | 97 +++ tests/test_coconut_disk.py | 241 ++++++ tests/test_coconut_topk.py | 2 +- 23 files changed, 1511 insertions(+), 477 deletions(-) create mode 100644 benchmark/bm_Coconut_Streaming.cpp create mode 100644 docs/coconut-disk-streaming.md create mode 100644 docs/coconut-paper-implementation.md create mode 100644 lib/algos/CoconutBounds.hpp create mode 100644 tests/test_coconut_disk.py diff --git a/CONTINUITY.md b/CONTINUITY.md index 4e9c627..26bc4f9 100644 --- a/CONTINUITY.md +++ b/CONTINUITY.md @@ -211,3 +211,20 @@ Update `CONTINUITY.md` when: - Python binding structure changes materially - Tests, demos, or benchmarks gain a new required pattern - A major known rough edge is fixed or a new one appears + + +## Coconut materialized disk streaming + +Coconut uses a bounded raw memtable (default 1M records), external-sort bulk build, +immutable sorted disk runs and background binary compaction. Resident SAX remains +O(N); raw vectors persist to disk. Include waitForMerges() when timing completed +insertion. Temporary indexes have no reopen/WAL/recovery or concurrent caller +write/search support. Gaussian SAX only; inputs are never normalized internally. + +Top-k uses per-run sortable-key seeds, safe SAX pruning and independent exact +heaps, sharing bounded candidate reads across configurable groups of up to 64 +queries (default 32). Recent-window top-k filters stable insertion IDs. Python +bindings expose layout/storage/merge/query counters; BUILD_COCONUT remains optional. +Tests cover strict physical neighbors through flushes/merges, batch/scalar parity, +lookup-bound parity, ranges, corruption and API validation. See +`docs/coconut-disk-streaming.md` and `docs/coconut-paper-implementation.md`. diff --git a/README.md b/README.md index 52cc243..60dc7eb 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ The following table summarizes the key features of each algorithm: | **[Hercules](https://helios2.mi.parisdescartes.fr/~themisp/publications/pvldb22-hercules.pdf)** | In-memory hierarchical similarity search using EAPCA and SAX-based pruning | | **[DumpyOS](https://helios2.mi.parisdescartes.fr/~themisp/publications/vldbj24-dumpyos.pdf)** | In-memory scalable data series similarity search using an adaptive multi-ary iSAX index | | **[FreSH](http://publications.ics.forth.gr/tech-reports/2023/2023.TR489_FreSh_A_LockFree_Data_Series_Index.pdf)** | In-memory lock-free parallel similarity search using an iSAX index (SRDS 2023) | -| **[COCONUT](http://www.vldb.org/pvldb/vol11/p677-kondylakis.pdf)** | Sortable-SAX index built bottom-up; supports both static datasets and **streaming** (incremental) inserts (PVLDB 2018) | +| **[COCONUT](https://arxiv.org/abs/2006.11474)** | Paper-based materialized Coconut-LSM-Full: bottom-up disk index, streaming inserts, resident-summary SIMS and recent-window top-k | | **[FAISS](https://github.com/facebookresearch/faiss) flat** | Exact flat baseline backed by `faiss::IndexFlatL2`; static build plus streaming inserts with no retraining. Optional: build with `-DBUILD_FAISS=ON` | ### Incremental streaming inserts @@ -91,18 +91,45 @@ data and the index stays exact after every insert. It supports `L2_SQUARED` only Streaming updates are not concurrent with queries. Inserts can reallocate the owned database, so callers should not retain a pointer returned by `getDatabase()` across them. -`Coconut` uses the upstream ordered approximate search before exact top-k refinement: -it locates the query's sortable-SAX key in a resident sorted directory and seeds a -heap with nearby distinct records. Exact search then checks the remaining SAX bounds. -Top-k and range queries read candidate leaves as complete blocks into private worker -buffers, avoiding per-vector file seeks. The ordered lookup uses the resident directory -instead of importing Berkeley DB's disk B-tree engine; raw vectors remain in leaf files. -Search supports squared L2 only. Inserts append to an unmerged in-memory buffer; -Coconut-LSM flush/merge handling and an out-of-core key directory remain unimplemented. +`Coconut` now uses bounded disk streaming: a RAM memtable of at most +`getBufferCapacity()` records (default 1,000,000) flushes to sorted disk runs. +Binary LSM compaction merges adjacent temporal runs in the background, including +the raw vectors. Include `waitForMerges()` after `insertBatch()` when measuring +completed insertion time (sorting, writes, `fsync` and compaction). Initial build +uses external sorting; SIMS keeps packed SAX summaries resident and parallelizes +bounds over resident summaries before ordered raw refinement. Up to 32 independent +queries share bounded candidate reads; each retains its own exact heap and bound. +Sparse reads coalesce short gaps, while dense reads share blocks of at most 8 MiB. +Ordered approximate seeding precedes exact SAX-bound refinement across every run +and any partially full memtable. Candidate raw vectors are read in leaf blocks. +Top-k and inclusive range search support squared L2 with Gaussian SAX breakpoints; +vectors are never normalized internally. Configure layout before building: + +```python +index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) +index.setBufferCapacity(1_000_000) # records, not the external insertion batch size +index.setIndexDirectory("/path/to/storage") # unique temporary child directory +index.buildIndex(initial_vectors) +index.insertBatch(new_vectors) # full tables persist and schedule compaction +index.waitForMerges() # include this barrier in completed-insert timing +index.setQueryBatchSize(32) # share candidate reads across independent queries +ids, distances = index.searchIndexRecent(queries, 1, 50_000_000) # latest records +print(index.getDiskRecords(), index.getBufferedRecords(), index.getMergeCount()) +``` + +This implements the paper's materialized Coconut-LSM-Full algorithms, with explicit +[implementation choices and tests](docs/coconut-paper-implementation.md); it is +not the authors' binary or an upstream performance reproduction. POSIX disk storage is required. The OS may +cache disk pages, so measure cgroup memory as well as the bounded explicit buffers. +Indexes are temporary and removed on destruction; crash recovery, reopening, +deletion and concurrent caller writes/reads are unsupported. Recent-window queries +use insertion counts, not arbitrary user timestamps. Resident SAX uses O(N) RAM; bound/read scratch is bounded independently of N; the raw vectors and bounded insertion buffer have separate budgets. +See [the implementation notes](docs/coconut-disk-streaming.md). Earlier results +from the unmerged RAM insertion buffer do not evaluate disk-based streaming. ### Range (distance-r) queries -Every algorithm except Coconut's streaming-only paths answers range queries through +Algorithms that implement range search answer range queries through `SearchConfig`. Instead of a fixed `k`, each query returns however many series fall within the radius, so the results come back as one vector per query: diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 9455cb2..47893d4 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -147,6 +147,8 @@ endif() # ////// COCONUT ////// if(BUILD_COCONUT) + add_executable(bm_Coconut_Streaming bm_Coconut_Streaming.cpp) + target_link_libraries(bm_Coconut_Streaming PRIVATE benchmark::benchmark dino_lib commons_lib) add_executable(bm_Coconut_L2Square bm_Coconut_L2Square.cpp bm_utils.cpp diff --git a/benchmark/bm_Coconut_Streaming.cpp b/benchmark/bm_Coconut_Streaming.cpp new file mode 100644 index 0000000..3bc7aff --- /dev/null +++ b/benchmark/bm_Coconut_Streaming.cpp @@ -0,0 +1,32 @@ +#include +#include "../commons/dataloaders.hpp" +#include "../lib/algos/Coconut.hpp" +#include +#include + +// Match the existing streaming benchmark pattern. Unlike an enqueue-only +// measurement, wall time includes every flush and background compaction. +static void BM_Coconut_Streaming(benchmark::State &state) +{ + constexpr daisy::idx_t dim = 96, initial = 4096, queries = 4; + const auto batch = static_cast(state.range(0)); + std::unique_ptr data(loadRandomData(initial + batch, dim, 100, true)); + std::unique_ptr query(loadRandomData(queries, dim, 50, true)); + daisy::Coconut index(daisy::DistanceType::L2_SQUARED); + index.setBufferCapacity(512); index.setNumThreads(2); + index.setQueryBatchSize(static_cast(state.range(1))); + index.buildIndex(data.get(), initial, dim); + std::vector ids(queries); std::vector distances(queries); + for (auto _ : state) + { + index.insertBatch(data.get() + initial * dim, batch); + index.waitForMerges(); + index.searchIndex(query.get(), queries, 1, ids.data(), distances.data()); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(static_cast(state.iterations()) * batch); + state.counters["n_database_final"] = static_cast(index.getNDatabase()); + state.counters["resident_summary_bytes"] = static_cast(index.getResidentSummaryBytes()); +} +BENCHMARK(BM_Coconut_Streaming)->Args({1024, 1})->Args({1024, 32})->Args({8192, 32})->Iterations(5)->UseRealTime()->Unit(benchmark::kMillisecond); +BENCHMARK_MAIN(); diff --git a/demos/demo_Coconut_Streaming.cpp b/demos/demo_Coconut_Streaming.cpp index 9ae059b..aa9e358 100644 --- a/demos/demo_Coconut_Streaming.cpp +++ b/demos/demo_Coconut_Streaming.cpp @@ -19,6 +19,9 @@ int main() float *query = loadRandomData(n_query, dim, 50, true); daisy::Coconut coconut(daisy::DistanceType::L2_SQUARED); + coconut.setBufferCapacity(10000); // auto flush and background binary disk merges + coconut.setNumThreads(2); + coconut.setQueryBatchSize(32); // shared reads; exact per-query SAX bounds/heaps // 1. Build on the initial batch. coconut.buildIndex(stream, initial, dim); @@ -42,9 +45,14 @@ int main() for (int s = 0; s < steps; s++) { coconut.insertBatch(stream + (size_t)seen * dim, batch); + coconut.waitForMerges(); // include pending compaction in completed-insert timing seen += batch; printf("Streamed a batch; index now holds %llu series.\n", seen); + printf("Disk records: %llu; buffered: %llu; merges: %llu\n", + coconut.getDiskRecords(), coconut.getBufferedRecords(), coconut.getMergeCount()); run_query("after insert"); + coconut.searchIndexRecent(query, n_query, 1, batch, I, D); + printf("Most recent batch, query 0 top-1: %llu (%.1f)\n", I[0], D[0]); } delete[] stream; diff --git a/demos/demo_Coconut_Streaming.py b/demos/demo_Coconut_Streaming.py index b684267..314d501 100644 --- a/demos/demo_Coconut_Streaming.py +++ b/demos/demo_Coconut_Streaming.py @@ -23,7 +23,9 @@ def main(): query = np.random.randn(n_query, dim).astype(np.float32) index = Coconut(DistanceType.L2_SQUARED) - index.setNumThreads(1) + index.setNumThreads(2) + index.setQueryBatchSize(32) # shared bounded candidate reads; independent exact heaps + index.setBufferCapacity(10_000) # bounded RAM table; full buffers persist and schedule merges index.buildIndex(stream[:initial]) # build on the initial batch def show(when): @@ -35,9 +37,14 @@ def show(when): seen = initial for s in range(steps): index.insertBatch(stream[seen:seen + batch]) # stream the next batch + index.waitForMerges() # include this barrier when timing completed insertions seen += batch - print(f"streamed a batch; index now holds {seen} series") + print(f"streamed {seen} series; disk={index.getDiskRecords()}, " + f"RAM buffer={index.getBufferedRecords()}, merges={index.getMergeCount()}") show("after insert") + recent_ids, _ = index.searchIndexRecent(query, 1, batch) + print("top-1 in the most recent batch:", recent_ids[:, 0]) + print("SIMS phases and reads:", index.getLastSearchStats()) if __name__ == "__main__": main() diff --git a/docs/coconut-disk-streaming.md b/docs/coconut-disk-streaming.md new file mode 100644 index 0000000..34c6022 --- /dev/null +++ b/docs/coconut-disk-streaming.md @@ -0,0 +1,60 @@ +# Coconut disk streaming + +`Coconut` follows the materialized sorted-run design of +[Coconut-LSM](https://arxiv.org/abs/2006.11474), adapted to DaiSy's exact top-k and +inclusive range APIs. It is a new implementation of that storage/update policy, +not an import of the 2018 Berkeley DB executable or a claim of performance parity. + +* A configurable memtable holds at most `setBufferCapacity(records)` raw vectors + plus SAX summaries and sortable keys. Sorting adds one bounded permutation. +* Full tables are sorted and persisted as immutable runs with stable insertion + IDs. The initial build also processes bounded chunks, including file sources. +* Each run has a fixed-width `[sortable key][SAX word][ID]` directory and a raw + float32 data file in the corresponding order. Directories are reclaimable + read-only file mappings; no full vector database or heap key directory is kept. +* Equal-level runs are merged with sequential bounded input/output buffers, + using binary size tiers in a background worker. A new run is written, flushed and synced before its + input runs are retired. Bulk build consolidates its sorted chunks into a run. +* Queries seed a shared bound from each run's sortable-key neighborhood, then + refine every surviving SAX summary and inspect the remaining bounded memtable. + Top-k uses resident SAX, exact query interval lookups and bounded shared + candidate reads across independent queries. Returned IDs refer to the original stream. +* `insertBatch` includes automatic flushes and schedules asynchronous merges. + Include `waitForMerges()` in completed-insert timing; `flush()` can persist an + incomplete table. `getDiskRecords`, `getBufferedRecords`, `getDiskBytes`, + `getRunCount`, `getFlushCount`, `getMergeCount`, and `getPeakBufferedRecords` + provide evidence that inserted data are actually reaching disk. + +There is no WAL, recoverable manifest, restart/reopen support, deletion or arbitrary-timestamp +window API. Writes and searches must not run concurrently. The temporary directory +is removed when its instance is destroyed. A failed batch can have accepted a +prefix before an I/O failure; inspect the count and treat the error as a failed +operation, not an atomic transaction. Invalid shapes/non-finite vectors are rejected. +Changing SAX/leaf/buffer/storage layout after building is rejected. Gaussian SAX +is supported; selecting adaptive breakpoints is explicitly rejected in C++. +No additional dependency is introduced; `BUILD_COCONUT=OFF` still excludes it. + +## Memory and measurements + +Raw vectors persist in run files; the memtable is bounded by record count. +Resident SAX summaries cost one byte per segment per persisted record, so they +still require O(N) RAM. Mapped directories, OS file cache and summaries retained +by in-flight merges also count toward memory usage. Record cgroup/RSS alongside +explicit buffer counters; this is not a constant-memory index. + +Freeze PAA segments, leaf size, query batch size and memtable capacity across +streaming stages. Include `waitForMerges()` after `insertBatch()` inside the timed +operation, then audit live disk/buffer counters. Do not charge outstanding merge +work to the subsequent query. Build and batch loading are separate measurements. +Query phase/read counters describe logical requests, not physical device I/O. + +See [the paper mapping and optimization details](coconut-paper-implementation.md). + +## Validation + +Validation includes strict float64 physical-neighbor checks after repeated flushes +and merges, stable IDs, inclusive ranges, partial buffers, file-backed builds, +strided Python inputs, corruption detection and optional-build exclusion. The +legacy Astronomy range fixture has non-finite rows 10321 and 40380; those nine +parameterized cases explicitly skip the invalid fixture rather than filtering it +or claiming exact range accuracy on NaNs. Finite-input rejection has its own test. diff --git a/docs/coconut-paper-implementation.md b/docs/coconut-paper-implementation.md new file mode 100644 index 0000000..f8d8698 --- /dev/null +++ b/docs/coconut-paper-implementation.md @@ -0,0 +1,121 @@ +# Coconut-LSM-Full implementation against the paper + +Source: Kondylakis, Dayan, Zoumpatianos and Palpanas, +[Coconut: Sortable Summarizations for Scalable Indexes over Static and Streaming Data Series](https://arxiv.org/abs/2006.11474), +[author PDF](https://nivdayan.github.io/Coconut_VLDBJ_CR.pdf). +The PDF used on 2026-10-19 is saved with the isolated build; its SHA256 is +`52da8f081c167ee8c75c303d7764194f226c9caac4e366b5d916e30c16c870be`. + +This is a paper-based implementation of the **materialized Coconut-LSM-Full** +variant, using the existing public `Coconut` name. It is not the authors' binary, +a bit-for-bit reproduction of their storage engine, or a performance-parity claim. +The nonmaterialized variant, Coconut-Trie and the static Berkeley DB executable +are different implementations and are not silently substituted here. + +## Traceability + +| Paper operation | DaiSy implementation | Validation | +|---|---|---| +| Sortable SAX | Existing bit-interleaved sortable SAX; Gaussian regions, no implicit normalization | Prior upstream-key comparison; exact result tests | +| Bottom-up build, Algorithms 3/6 | Sort initial bounded chunks once, then bounded-fan-in external merge; no repeated insertion-style compactions during bulk loading | More than 64 initial chunks, file input, partial chunks | +| Balanced indexed run | Evenly packed adjacent leaves (at least half full), bottom-up 128-way fence-key B+tree; ordered keys/IDs and sorted raw files | More than 128² leaves, duplicate keys and both query-key endpoints | +| Approximate seed, Algorithm 4 | Tree traversal to insertion point, leaf-sized neighboring window in every eligible run, one global best-so-far heap | k=1/10/100/1000, ties, empty build, distinct stable IDs | +| Resident summaries, Algorithms 5/7 | Packed SAX snapshot in each immutable run, exactly PAA-segments bytes per record. Original SAX is retained rather than undoing the invertible key permutation for every bound | Resident bytes equal persisted-record count times segments | +| Parallel bounds, Algorithm 7 | OpenMP bounds across resident summaries and independent queries; bounded block arrays (original scalar-query schedule available with batch size 1) | One/four threads return exactly the same integer-data IDs/distances | +| Skip sequential refinement | Iterate run positions in order; skip prunable records, read candidate leaves into one reusable buffer, update shared query bound | Physical float64 ranks and returned-vector distances; read/phase counters | +| LSM updates, Section 4.4/Figure 7 | Bounded memtable, durable sorted flushes, asynchronous adjacent-run binary merges (size ratio 2), bounded pending runs/backpressure | Queries while worker replaces runs, barrier and file accounting, sync/async modes | +| Bounded temporal partitioning, Section 5.3 | Runs carry contiguous insertion-ID intervals. Skip wholly older runs; filter the boundary run by original IDs. `searchIndexRecent` accepts the count of latest records | Windows 1 through entire stream; old runs skipped; partial memtable; ties | + +## Storage, memory and timing decisions + +* The materialized variant rewrites raw vectors during compaction. Every published + replacement is fully written and synced before its inputs are retired. +* Readers retain immutable shared run snapshots, so background replacement cannot + remove their files. Caller-issued inserts and queries must still be serialized. +* Background merging defaults on; `setAsyncMerges(false)` selects blocking updates. + `waitForMerges()` waits for all pending work and surfaces background I/O errors. + Time `insertBatch` **and this barrier** for completed-batch comparisons. Enqueue + latency alone is not a comparable insertion time. A partial table stays queryable + in RAM; call `flush()` to persist it, then the barrier to finish compaction. +* Initial chunk merging has fan-in at most 64 to bound file descriptors and I/O + buffers. Inputs beyond that limit require additional external-sort passes. This + is an engineering limit rather than a promise of two passes at every scale. +* The RAM memtable remains bounded, but SIMS deliberately has O(N) resident SAX: + 16 GB of SAX at 1B records with 16 segments, plus bounded block-bound arrays and small tree directories. In-flight merges/query snapshots can + temporarily retain both old and new summaries. `getResidentSummaryBytes()` + reports the currently published runs; RSS/cgroup usage captures transients too. +* Keys and IDs are mmap-backed; raw leaf reads use `pread`. OS file cache counts + toward the memory budget. Explicit resident summaries are never described as + the whole raw database in RAM. `getDiskBytes()` counts live logical run files; + concurrent compaction creates temporary additional files. Audit physical bytes + only after `waitForMerges()` and after outstanding queries return. +* `getLastSearchStats()` gives last top-k call totals: resident summaries examined, + logical leaf reads/bytes requested (not physical device counters), runs searched/ + skipped, and approximate/bounds/refinement phase wall times. Cold summary + construction is part of build/merge, not hidden inside query timing. +* `setLeafSize()` controls the approximate radius and maximum records per leaf. + Fanout 128, split files and POSIX I/O are documented implementation choices; + the paper does not fix the source-level layout for reproductions. + +## Bounded batch SIMS optimization (2026-10-20) + +The current default groups at most 32 queries (`setQueryBatchSize(1..64)`). Each +query still traverses the Coconut run trees for its own seed and keeps its own +exact heap, SAX lower bounds and stable-ID tie order. Summaries are evaluated in +ordered blocks, using the current per-query bound. Qualifying positions across +the query group share reads; no raw block is read unless it contains a candidate. +Sparse positions are coalesced across gaps of at most 4 KiB. Dense candidates +share a contiguous read of at most 8 MiB. Refinement rechecks every query's bound +before a distance computation. This changes scheduling and I/O granularity, not +the sortable-SAX, LSM, pruning or exact-search logic. It is an engineering +optimization of paper-based SIMS, not the paper's original sequential-query +schedule. No unconditional exhaustive-search fallback is introduced. + +Lower-bound arrays are now bounded by query count times block capacity, rather +than total run size. A block has at most 32,768 records and at most 8 MiB of raw +vectors (or one vector for dimensions exceeding that size). Scratch counters +cover raw/bounds/mask buffers; heaps and resident summaries are separate. Query +results from different groups never share heap state. Caller writes/reads remain +serialized; immutable run snapshots still protect background compactions. + +The per-query interval-distance lookup (`CoconutBounds.hpp`) retains the reference +helper's double squared contributions and float accumulation order. A partial +nonnegative sum may stop only once its scaled bound exceeds the query bound; +otherwise it returns the same full bound. A dedicated C++ test checks exact float +equality against the existing SAX helper for 1-8 bits and 1/3/16/32/48/96 segments, +including equal and adjacent-float thresholds. This is a faster evaluation of the +same Coconut lower bound, not a learned/approximate bound. Dense candidate blocks +(at least half their positions qualify across the group) coalesce into one read; +sparse blocks retain the 4 KiB gap limit. The lookup tables are included in the +reported scratch-array memory. + +New counters distinguish candidate/query pairs (`candidate_records`), unique +qualifying positions per group (`unique_candidate_records`), actual distance +computations, seed computations/records read, raw records read and peak scratch +array bytes. To evaluate block-read amplification, exclude seed reads and divide +raw records read by the unique qualifying positions. Logical read counters are +not physical device I/O measurements. PAA segment count is an ordinary Coconut +parameter; tuning it changes summaries and memory needs, never the input vectors +or the exact squared-L2 answers. Freeze that parameter in every measured config. + +## API extensions and limits + +Top-k generalizes the paper's 1-NN with a unique-ID heap and deterministic equal- +distance ID ordering. Inclusive range search remains an existing DaiSy extension. +Recency uses insertion order, not arbitrary user-provided wall-clock timestamps; +`searchIndexRecent(queries, k, window_size)` returns original stream IDs, with +existing max-float/ID-zero padding if fewer than k records qualify. It is a +Coconut-specific API; it does not change the common SearchConfig contract. +Windowed range/approximate-only APIs are not exposed. Raw input is never normalized; +the paper's z-normalized data and our unchanged DEEP embeddings are distinct +workloads, so performance claims must state the actual data preprocessing. + +Indexes remain temporary POSIX files with no reopen, recovery/WAL, deletion, or +concurrent caller reads/writes. Failed inserts may accept a prefix; a failed +background merge is surfaced and the index must be rebuilt. Optional build guards +remain unchanged. Neither this implementation nor its exactness tests guarantees +that Coconut will outperform FAISS/MESSI on the chosen workload. + +Performance evaluations must freeze binary/source and configuration, include the +compaction barrier in insertion timing, validate each prefix against independent +ground truth, and retain phase/read/memory diagnostics. diff --git a/docs/daisy-architecture.puml b/docs/daisy-architecture.puml index 861008b..3d4e3b2 100644 --- a/docs/daisy-architecture.puml +++ b/docs/daisy-architecture.puml @@ -27,6 +27,16 @@ package lib { class ParIS as PARIS class Sing as SING class Odyssey as ODYSSEY + class Coconut as COCONUT + COCONUT : + setBufferCapacity(records) + COCONUT : + insertBatch(data, n) + COCONUT : + getDiskRecords() + COCONUT : + waitForMerges() + COCONUT : + searchIndexRecent(window) + COCONUT : resident SAX / bounded batch SIMS + COCONUT : + setQueryBatchSize(queries) + class "Sorted disk runs" as CoconutRuns + COCONUT *-- CoconutRuns : bounded flush / background binary merge class DataSource as DS <> DS : + nextRecord() bool @@ -42,6 +52,7 @@ package lib { SSA <|-- PARIS : extends SSA <|-- SING : extends SSA <|-- ODYSSEY : extends + SSA <|-- COCONUT : extends DS <|-- IMDS : extends DS <|-- FDS : extends SSA --> DS : buildIndex() diff --git a/docs/daisy-components.puml b/docs/daisy-components.puml index 9529984..2e5a33f 100644 --- a/docs/daisy-components.puml +++ b/docs/daisy-components.puml @@ -22,6 +22,9 @@ lib -right-> pybinds : BUILD_PYTHON\n→ modulo Python note bottom of lib Algoritmi: Bruteforce, LbBruteforce, MESSI, ParIS, SING (CUDA), Odyssey (MPI) + Coconut: bounded memtable -> sorted disk runs + Background LSM merges, resident SAX / shared-read batch SIMS + Bottom-up run trees, BTP recent-window top-k (POSIX) end note @enduml diff --git a/docs/how-to-contribute.md b/docs/how-to-contribute.md index 6794b9e..20fc0d6 100644 --- a/docs/how-to-contribute.md +++ b/docs/how-to-contribute.md @@ -18,6 +18,14 @@ Here is a (non-exhaustive) mockup of our future and ongoing goals: ## Prerequisites +For disk-streaming algorithms, correctness alone is insufficient: exercise buffer +flushes and compaction in tests, verify actual disk-record counts and bounded RAM +buffers, and include background-compaction barriers in completed-insertion +measurements. Resident summaries are a separate O(N) memory cost. Validate query +read-sharing against independent scalar-query answers; record candidate counts +and logical/physical reads separately. See +[Coconut disk streaming](coconut-disk-streaming.md) for the current adapter contract. + Before starting, ensure you have: - A working development environment with CMake, C++17 or later, and Python 3.7+ - Familiarize yourself with the existing algorithms in `lib/algos/` (e.g., `Bruteforce.hpp/cpp`, `Messi.hpp/cpp`) diff --git a/docs/test-framework-guide.md b/docs/test-framework-guide.md index dbe488a..5e39419 100644 --- a/docs/test-framework-guide.md +++ b/docs/test-framework-guide.md @@ -68,3 +68,18 @@ ID/distance pairs, rejects duplicates and validates the full-dataset reference against the published SIFT ground truth. Distances of raw integer descriptors are compared exactly; equally distant alternative IDs are allowed. The dataset and FAISS are not required for the regular C++ tests. + + +## Coconut disk streaming + +With `BUILD_COCONUT=ON`, C++ tests cover static and streaming top-k, inclusive +ranges, exact physical neighbor IDs after disk compaction, recent-window queries +and SAX-bound lookup parity at adjacent-float thresholds. Non-finite Astronomy +range fixtures are explicitly skipped; finite-input rejection is tested separately. + +With `BUILD_PYTHON=ON`, `test_coconut_topk_py` (NumPy and FAISS required) and +`test_coconut_disk_py` (NumPy required) additionally exercise file-backed builds, +bounded memtables, sync/async compaction, stable IDs, scalar versus batched read +sharing, odd dimensions, strided arrays, corrupted files and invalid inputs. +Run them with `ctest --test-dir build -R '[Cc]oconut' --output-on-failure`. +`BUILD_COCONUT=OFF` excludes the implementation, bindings, tests and benchmark. diff --git a/lib/algos/CMakeLists.txt b/lib/algos/CMakeLists.txt index a930d67..7437571 100644 --- a/lib/algos/CMakeLists.txt +++ b/lib/algos/CMakeLists.txt @@ -29,12 +29,15 @@ endif() # ////// COCONUT ////// if(BUILD_COCONUT) + find_package(Threads REQUIRED) + target_link_libraries(dino_lib PRIVATE Threads::Threads) if(DEBUG_MSG) message(STATUS "Adding Coconut.cpp to dino_lib sources.") endif() target_sources(dino_lib PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/Coconut.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/CoconutBounds.hpp ) target_compile_definitions(dino_lib PUBLIC COCONUT_ENABLED) else() diff --git a/lib/algos/Coconut.cpp b/lib/algos/Coconut.cpp index 5ff6967..d4fc661 100644 --- a/lib/algos/Coconut.cpp +++ b/lib/algos/Coconut.cpp @@ -1,460 +1,731 @@ #include "Coconut.hpp" - +#include "CoconutBounds.hpp" #include "../isax/SAX.hpp" -#include "../isax/iSAXTypes.hpp" - +#include +#include #include -#include +#include #include -#include #include -#include +#include +#include +#include +#include #include -#include +#include +#include +#include +#include +#include +#include +#include +#include #include #include namespace daisy { - - Coconut::Coconut(DistanceType distance_type) - : SimilaritySearchAlgorithm(distance_type) +namespace +{ + size_t checkedSize(idx_t n, size_t width) { + if (width && n > std::numeric_limits::max() / width) + throw std::runtime_error("Coconut index size overflow"); + return (size_t)n * width; } - - Coconut::~Coconut() + void checkFinite(const float *data, idx_t n, idx_t dim) { - cleanupLeafFiles(); + if (n && !data) throw std::runtime_error("Coconut vector data is null"); + for (size_t i = 0, count = checkedSize(n, (size_t)dim); i < count; ++i) + if (!std::isfinite(data[i])) throw std::runtime_error("Coconut requires finite vectors"); } - - std::string Coconut::leafPath(int leaf_no) const + float squaredL2(const float *a, const float *b, idx_t dim, float bound) { - return index_dir_ + "/leaf." + std::to_string(leaf_no); + // The shared SIMD kernel consumes full groups of eight. Handle a tail + // locally rather than reading into the next vector for other dimensions. + const int full = (int)dim / 8 * 8; + float distance = full ? ts_euclidean_distance_SIMD(const_cast(a), const_cast(b), full, bound) : 0; + for (idx_t d = full; d < dim && distance <= bound; ++d) + { float delta = a[d] - b[d]; distance += delta * delta; } + return distance; } - - // Compare two sortable-SAX keys lexicographically over their bytes (COCONUT's order). - static bool key_less(const std::vector &a, const std::vector &b) + struct File { - for (size_t i = 0; i < a.size(); i++) + FILE *p = nullptr; + std::vector buffer; + File(const std::string &path, const char *mode) : buffer(1024 * 1024) { - if (a[i] != b[i]) - return a[i] < b[i]; + p = fopen(path.c_str(), mode); + if (!p) throw std::runtime_error("Coconut failed to open index leaf/run: " + path); + setvbuf(p, buffer.data(), _IOFBF, buffer.size()); } - return false; - } - - void Coconut::buildIndex(DataSource *data_source) + ~File() { if (p) fclose(p); } + File(const File &) = delete; + void write(const void *data, size_t bytes) + { + if (fwrite(data, 1, bytes, p) != bytes) throw std::runtime_error("Coconut failed to write index leaf/run"); + } + void read(void *data, size_t bytes) + { + if (fread(data, 1, bytes, p) != bytes) throw std::runtime_error("Coconut failed to read index leaf/run"); + } + void finish() + { + if (fflush(p) || fsync(fileno(p))) throw std::runtime_error("Coconut failed to persist index leaf/run"); + FILE *closed = p; p = nullptr; + if (fclose(closed)) throw std::runtime_error("Coconut failed to close index leaf/run"); + } + }; + // Fixed-size directory records: [sortable key][SAX word][stable ID]. + // The directory is mmap-backed and reclaimable, not a vector of O(N) + // heap allocations. Raw vectors are in a separate corresponding run file. + using Clock = std::chrono::steady_clock; + double elapsed(Clock::time_point start) { return std::chrono::duration(Clock::now() - start).count(); } + struct Run { - this->dim = data_source->getDim(); - this->n_database = data_source->getTotalRecords(); - if (this->n_database == 0) + idx_t n = 0, first_id = 0, end_id = 0, leaf_count = 0, leaf_base = 0, large_leaves = 0; + int segments = 0; + static constexpr idx_t fanout = 128; + std::vector summaries; + // Bottom-up packed B+tree: each level stores one maximum key per child. + std::vector> tree; + size_t stride = 0, raw_stride = 0; + unsigned level = 0; + std::string meta_path, raw_path; + int meta_fd = -1, raw_fd = -1; + const unsigned char *map = nullptr; + ~Run() { - data_source->reset(); - idx_t count = 0; - float *dummy = new float[this->dim]; - while (data_source->nextRecord(dummy)) - count++; - delete[] dummy; - this->n_database = count; - data_source->reset(); + if (map) munmap(const_cast(map), (size_t)n * stride); + if (meta_fd >= 0) close(meta_fd); + if (raw_fd >= 0) close(raw_fd); + if (!meta_path.empty()) unlink(meta_path.c_str()); + if (!raw_path.empty()) unlink(raw_path.c_str()); } - - const int seg = this->paa_segments; - const int bits = this->sax_cardinality; - if (seg <= 0 || this->dim == 0 || this->dim % (idx_t)seg != 0 || - bits <= 0 || bits > 8 || this->leaf_capacity <= 0) - throw std::runtime_error("Coconut requires positive leaf/PAA sizes, 1-8 SAX bits, and a dimension divisible by the PAA size"); - this->ts_values_per_segment_ = (int)this->dim / seg; - this->sax_alphabet_cardinality_ = (int)std::pow(2, bits); - this->mindist_sqrt_ = (float)this->dim / (float)seg; - this->max_cardinalities_.assign(seg, (sax_type)bits); - - // Borrow a contiguous input only during the build. Copy stream-only sources; - // raw series are owned by the leaf files once this method returns. - data_source->reset(); - const float *raw = data_source->rawPointer(); - std::vector buffer; - if (!raw) + const unsigned char *record(idx_t pos) const { return map + (size_t)pos * stride; } + idx_t id(idx_t pos, int seg) const + { + idx_t result; memcpy(&result, record(pos) + 2 * seg, sizeof(result)); return result; + } + void validate() const + { + struct stat metadata, raw; + if (fstat(meta_fd, &metadata) || fstat(raw_fd, &raw) || + (uint64_t)metadata.st_size != checkedSize(n, stride) || + (uint64_t)raw.st_size != checkedSize(n, raw_stride)) + throw std::runtime_error("Coconut truncated or unreadable index leaf/run"); + } + idx_t leafBegin(idx_t leaf) const { return leaf * leaf_base + std::min(leaf, large_leaves); } + idx_t leafOf(idx_t pos) const + { + idx_t wide = (leaf_base + 1) * large_leaves; + return pos < wide ? pos / (leaf_base + 1) : large_leaves + (pos - wide) / leaf_base; + } + idx_t locate(const sax_type *key) const { - buffer.resize((size_t)this->n_database * this->dim); - std::vector rec(this->dim); - idx_t idx = 0; - while (idx < this->n_database && data_source->nextRecord(rec.data())) + idx_t selected = 0; + for (size_t l = tree.size(); l; --l) { - std::copy(rec.begin(), rec.end(), buffer.data() + (size_t)idx * this->dim); - idx++; + const auto &keys = tree[l - 1]; + idx_t begin = selected * fanout, end = std::min(begin + fanout, keys.size() / segments); + idx_t lo = begin, hi = end; + while (lo < hi) + { + idx_t mid = lo + (hi - lo) / 2; + if (memcmp(keys.data() + (size_t)mid * segments, key, segments) < 0) lo = mid + 1; + else hi = mid; + } + selected = std::min(lo, end - 1); } - if (idx != this->n_database) - throw std::runtime_error("Coconut datasource ended before all records were read"); - raw = buffer.data(); + idx_t lo = leafBegin(selected), hi = leafBegin(selected + 1); + while (lo < hi) + { + idx_t mid = lo + (hi - lo) / 2; + if (memcmp(record(mid), key, segments) < 0) lo = mid + 1; else hi = mid; + } + return lo; } - - // 1. Compute SAX + sortable SAX for every series. - records_.resize(this->n_database); - for (idx_t i = 0; i < this->n_database; i++) + void open(int seg, int leaf_size) { - Record &r = records_[i]; - r.series_id = i; - r.sax.resize(seg); - r.key.resize(seg); - sax_from_ts(const_cast(raw + (size_t)i * this->dim), r.sax.data(), - this->ts_values_per_segment_, seg, - this->sax_alphabet_cardinality_, bits); - sortable_sax_from_sax(r.sax.data(), r.key.data(), seg, bits); + meta_fd = ::open(meta_path.c_str(), O_RDONLY); raw_fd = ::open(raw_path.c_str(), O_RDONLY); + if (meta_fd < 0 || raw_fd < 0) throw std::runtime_error("Coconut failed to open index leaf/run"); + validate(); + void *p = mmap(nullptr, checkedSize(n, stride), PROT_READ, MAP_SHARED, meta_fd, 0); + if (p == MAP_FAILED) throw std::runtime_error("Coconut cannot map its on-disk key directory"); + map = static_cast(p); + segments = seg; + // Algorithm 7 scans a packed, resident SAX snapshot, not mapped keys/IDs. + summaries.resize(checkedSize(n, (size_t)seg)); + for (idx_t i = 0; i < n; ++i) + memcpy(summaries.data() + (size_t)i * seg, record(i) + seg, seg); + leaf_count = 1 + (n - 1) / (idx_t)leaf_size; + leaf_base = n / leaf_count; large_leaves = n % leaf_count; + tree.emplace_back(checkedSize(leaf_count, (size_t)seg)); + for (idx_t i = 0; i < leaf_count; ++i) + memcpy(tree[0].data() + (size_t)i * seg, record(leafBegin(i + 1) - 1), seg); + while (tree.back().size() / seg > fanout) + { + const auto &lower = tree.back(); idx_t children = lower.size() / seg; + std::vector upper(checkedSize(1 + (children - 1) / fanout, (size_t)seg)); + for (idx_t i = 0; i < upper.size() / seg; ++i) + memcpy(upper.data() + (size_t)i * seg, lower.data() + (size_t)(std::min(children, (i + 1) * fanout) - 1) * seg, seg); + tree.push_back(std::move(upper)); + } } - - // 2. Sort records by sortable-SAX key (bottom-up ordering). - std::sort(records_.begin(), records_.end(), - [](const Record &a, const Record &b) { return key_less(a.key, b.key); }); - - // 3. Write on-disk leaf files: consecutive sorted records, each = [inv-SAX][raw TS]. - static std::atomic uid{0}; - this->index_dir_ = "/tmp/daisy_coconut_" + std::to_string((unsigned long)getpid()) + - "_" + std::to_string(uid.fetch_add(1)); - mkdir(this->index_dir_.c_str(), 0777); - - FILE *lf = nullptr; - int cur_leaf = -1; - for (size_t p = 0; p < records_.size(); p++) + void leaf(idx_t begin, idx_t count, std::vector &page) const + { + size_t bytes = checkedSize(count, raw_stride), got = 0; + page.resize(bytes / sizeof(float)); + while (got < bytes) + { + ssize_t result = pread(raw_fd, reinterpret_cast(page.data()) + got, + bytes - got, (off_t)checkedSize(begin, raw_stride) + got); + if (result < 0 && errno == EINTR) continue; + if (result <= 0) throw std::runtime_error("Coconut failed to read index leaf/run"); + got += (size_t)result; + } + } + }; +} + struct Coconut::DiskState + { + idx_t dim, capacity, count = 0, buffer_first = 0; + int seg, bits, values, alphabet, leaf_size; + size_t stride; + std::string directory; + std::atomic serial{0}; + idx_t flushes = 0, peak_buffered = 0; + std::atomic merges{0}; + bool bulk_loading = true, async_merges = true, stopping = false, merging = false; + mutable std::mutex mutex; + std::condition_variable changed; + std::thread merger; + std::exception_ptr merge_error; + std::vector data; + std::vector sax, keys, cardinalities; + std::vector> runs; + DiskState(idx_t d, int s, int b, int leaf, idx_t cap, const std::string &parent, bool async) + : dim(d), capacity(cap), seg(s), bits(b), values((int)d / s), alphabet(1 << b), + leaf_size(leaf), stride(2 * (size_t)s + sizeof(idx_t)), cardinalities(s, (sax_type)b) + { + std::string root = parent.empty() ? std::filesystem::temp_directory_path().string() : parent; + std::filesystem::create_directories(root); + std::string pattern = root + "/daisy_coconut_" + std::to_string((unsigned long)getpid()) + "_XXXXXX"; + std::vector name(pattern.begin(), pattern.end()); name.push_back(0); + if (!mkdtemp(name.data())) throw std::runtime_error("Coconut cannot create its disk directory"); + directory = name.data(); async_merges = async; + } + ~DiskState() + { + { std::lock_guard lock(mutex); stopping = true; } + changed.notify_all(); if (merger.joinable()) merger.join(); + runs.clear(); + std::error_code ignored; std::filesystem::remove_all(directory, ignored); + } + idx_t buffered() const { return data.size() / (size_t)dim; } + std::shared_ptr newRun(idx_t n, unsigned level) + { + auto run = std::make_shared(); run->n = n; run->level = level; + run->stride = stride; run->raw_stride = checkedSize(dim, sizeof(float)); + std::string prefix = directory + "/run." + std::to_string(serial++); + run->meta_path = prefix + ".meta"; run->raw_path = prefix + ".raw"; return run; + } + std::shared_ptr merge(const std::vector> &inputs, unsigned level) { - int leaf_no = (int)(p / (size_t)this->leaf_capacity); - if (leaf_no != cur_leaf) + idx_t total = 0; + for (const auto &run : inputs) { run->validate(); total += run->n; } + auto result = newRun(total, level); + result->first_id = inputs.front()->first_id; result->end_id = inputs.front()->end_id; + for (const auto &run : inputs) + { result->first_id = std::min(result->first_id, run->first_id); result->end_id = std::max(result->end_id, run->end_id); } + File meta(result->meta_path, "wb"), raw(result->raw_path, "wb"); + std::vector> readers; + for (const auto &run : inputs) readers.emplace_back(new File(run->raw_path, "rb")); + std::vector positions(inputs.size(), 0); + auto later = [&](size_t a, size_t b) { + int cmp = memcmp(inputs[a]->record(positions[a]), inputs[b]->record(positions[b]), seg); + return cmp ? cmp > 0 : inputs[a]->id(positions[a], seg) > inputs[b]->id(positions[b], seg); + }; + std::priority_queue, decltype(later)> heap(later); + for (size_t i = 0; i < inputs.size(); ++i) heap.push(i); + std::vector ts(dim); + while (!heap.empty()) { - if (lf) fclose(lf); - lf = fopen(leafPath(leaf_no).c_str(), "wb"); - cur_leaf = leaf_no; + size_t i = heap.top(); heap.pop(); + meta.write(inputs[i]->record(positions[i]), stride); + readers[i]->read(ts.data(), checkedSize(dim, sizeof(float))); + raw.write(ts.data(), checkedSize(dim, sizeof(float))); + if (++positions[i] < inputs[i]->n) heap.push(i); } - if (lf) + meta.finish(); raw.finish(); result->open(seg, leaf_size); ++merges; return result; + } + std::pair mergePair() const + { + // Only adjacent temporal intervals merge; this preserves BTP recency. + for (size_t i = 0; i + 1 < runs.size(); ++i) + if (runs[i]->level == runs[i + 1]->level) + return {i, i + 1}; + return {runs.size(), runs.size()}; + } + void compactLoop() + { + std::unique_lock lock(mutex); + while (true) { - fwrite(records_[p].key.data(), sizeof(sax_type), seg, lf); - fwrite(raw + (size_t)records_[p].series_id * this->dim, - sizeof(float), this->dim, lf); + changed.wait(lock, [&] { return stopping || mergePair().first < runs.size(); }); + if (stopping) return; + auto pair = mergePair(); + auto a = runs[pair.first], b = runs[pair.second]; merging = true; + lock.unlock(); + std::shared_ptr replacement; + try { replacement = merge({a, b}, a->level + 1); } + catch (...) + { + lock.lock(); merge_error = std::current_exception(); merging = false; + changed.notify_all(); return; + } + lock.lock(); + // Atomically publish a complete replacement, retaining old files + // while any query snapshot still refers to the two input runs. + auto it = std::find(runs.begin(), runs.end(), a); + auto pos = it - runs.begin(); + runs.erase(it, it + 2); runs.insert(runs.begin() + pos, std::move(replacement)); + a.reset(); b.reset(); merging = false; changed.notify_all(); } } - if (lf) fclose(lf); - // buffer freed here: raw series now live on disk in the leaf files. - - this->next_id_ = this->n_database; // streaming inserts get ids after the built set + void waitForMerges() + { + if (bulk_loading) return; + std::unique_lock lock(mutex); + changed.wait(lock, [&] { return merge_error || (!merging && mergePair().first == runs.size()); }); + if (merge_error) std::rethrow_exception(merge_error); + } + std::vector> snapshot() const + { + std::lock_guard lock(mutex); + if (merge_error) std::rethrow_exception(merge_error); + return runs; + } + void publish(std::shared_ptr run) + { + { std::lock_guard lock(mutex); runs.push_back(std::move(run)); } + if (bulk_loading) return; + changed.notify_all(); + if (!async_merges) { waitForMerges(); return; } + // Bound pending runs even when the writer outruns disk compaction. + unsigned levels = 1; idx_t chunks = count / capacity; + while (chunks > 1) { chunks /= 2; ++levels; } + std::unique_lock lock(mutex); + changed.wait(lock, [&] { return merge_error || runs.size() <= 2 * (size_t)(levels + 1); }); + if (merge_error) std::rethrow_exception(merge_error); + } + void flush() + { + idx_t n = buffered(); if (!n) return; + std::vector order(n); std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](idx_t a, idx_t b) { + int cmp = memcmp(keys.data() + (size_t)a * seg, keys.data() + (size_t)b * seg, seg); + return cmp ? cmp < 0 : a < b; + }); + auto run = newRun(n, 0); run->first_id = buffer_first; run->end_id = buffer_first + n; + File meta(run->meta_path, "wb"), raw(run->raw_path, "wb"); + for (idx_t slot : order) + { + idx_t id = buffer_first + slot; + meta.write(keys.data() + (size_t)slot * seg, seg); + meta.write(sax.data() + (size_t)slot * seg, seg); meta.write(&id, sizeof(id)); + raw.write(data.data() + checkedSize(slot, (size_t)dim), checkedSize(dim, sizeof(float))); + } + meta.finish(); raw.finish(); run->open(seg, leaf_size); + publish(std::move(run)); + data.clear(); sax.clear(); keys.clear(); buffer_first = count; ++flushes; + } + void append(const float *input, idx_t n, int threads) + { + if (n > std::numeric_limits::max() - count) throw std::runtime_error("Coconut record count overflow"); + snapshot(); // Surface any failed background merge before accepting more records. + while (n) + { + if (buffered() == capacity) flush(); + idx_t first = buffered(), take = std::min(n, capacity - first); + if (!first) buffer_first = count; + size_t raw_size = checkedSize(first + take, (size_t)dim); + size_t key_size = checkedSize(first + take, (size_t)seg); + // Explicit reserve avoids vector's geometric growth exceeding + // the declared memtable capacity after many small inserts. + if (data.capacity() < raw_size) data.reserve(raw_size); + if (sax.capacity() < key_size) sax.reserve(key_size); + if (keys.capacity() < key_size) keys.reserve(key_size); + data.resize(raw_size); sax.resize(key_size); keys.resize(key_size); + memcpy(data.data() + checkedSize(first, (size_t)dim), input, checkedSize(take, checkedSize(dim, sizeof(float)))); +#pragma omp parallel num_threads(std::max(1, threads)) + { + std::vector paa(seg); +#pragma omp for schedule(static) + for (idx_t i = first; i < first + take; ++i) + { + paa_from_ts(data.data() + checkedSize(i, (size_t)dim), paa.data(), seg, values); + sax_from_paa(paa.data(), sax.data() + (size_t)i * seg, seg, alphabet, bits); + sortable_sax_from_sax(sax.data() + (size_t)i * seg, keys.data() + (size_t)i * seg, seg, bits); + } + } + count += take; n -= take; input += checkedSize(take, (size_t)dim); + peak_buffered = std::max(peak_buffered, buffered()); + if (buffered() == capacity) flush(); + } + } + void consolidateBuild() + { + // Algorithm 6: sort initial chunks once, then merge them; do not run + // incremental binary compactions repeatedly during initial loading. + unsigned level = 0; idx_t size = capacity; + while (size < count && size <= std::numeric_limits::max() / 2) { size *= 2; ++level; } + constexpr size_t max_fanin = 64; + while (runs.size() > 1) + { + std::vector> next; + for (size_t i = 0; i < runs.size(); i += max_fanin) + { + size_t end = std::min(runs.size(), i + max_fanin); + std::vector> group(runs.begin() + i, runs.begin() + end); + next.push_back(group.size() > 1 ? merge(group, level) : group[0]); + } + runs = std::move(next); + } + if (!runs.empty()) runs[0]->level = level; + bulk_loading = false; + merger = std::thread([this] { compactLoop(); }); + } + float bound(const float *paa, const sax_type *word) const + { + return minidist_paa_to_isax(const_cast(paa), const_cast(word), + const_cast(cardinalities.data()), (sax_type)bits, alphabet, + seg, MINVAL, MAXVAL, (float)values); + } + void validateRuns() const { for (const auto &run : snapshot()) run->validate(); } + }; + + Coconut::Coconut(DistanceType type) : SimilaritySearchAlgorithm(type) {} + Coconut::~Coconut() = default; + void Coconut::requireMutableConfig() const + { if (state_) throw std::runtime_error("Coconut layout settings cannot change after buildIndex"); } + void Coconut::requireBuilt() const + { if (!state_) throw std::runtime_error("Coconut requires an initial buildIndex first"); } + void Coconut::setPaaSegments(int n) { requireMutableConfig(); paa_segments = n; } + void Coconut::setSaxCardinality(int n) { requireMutableConfig(); sax_cardinality = n; } + void Coconut::setLeafSize(int n) { requireMutableConfig(); leaf_capacity_ = n; } + void Coconut::setBufferCapacity(idx_t n) { requireMutableConfig(); buffer_capacity_ = n; } + void Coconut::setQueryBatchSize(idx_t queries) + { + if (!queries || queries > 64) throw std::runtime_error("Coconut query batch size must be between 1 and 64"); + query_batch_size_ = queries; } - - void Coconut::insert(const float *series) + void Coconut::setAsyncMerges(bool enabled) { requireMutableConfig(); async_merges_ = enabled; } + void Coconut::waitForMerges() { requireBuilt(); state_->waitForMerges(); } + idx_t Coconut::getResidentSummaryBytes() const { - if (this->paa_segments <= 0 || this->ts_values_per_segment_ <= 0) - throw std::runtime_error("Coconut::insert requires an initial buildIndex first"); - - const int seg = this->paa_segments; - Record r; - r.series_id = this->next_id_++; - r.sax.resize(seg); - r.key.resize(seg); - sax_from_ts(const_cast(series), r.sax.data(), - this->ts_values_per_segment_, seg, - this->sax_alphabet_cardinality_, this->sax_cardinality); - sortable_sax_from_sax(r.sax.data(), r.key.data(), seg, this->sax_cardinality); - - buffer_records_.push_back(std::move(r)); - buffer_data_.insert(buffer_data_.end(), series, series + this->dim); - this->n_database++; + idx_t bytes = 0; if (state_) for (const auto &run : state_->snapshot()) bytes += run->summaries.size(); + return bytes; } - - void Coconut::insertBatch(const float *data, idx_t n) + void Coconut::setIndexDirectory(const std::string &p) { requireMutableConfig(); storage_parent_ = p; } + void Coconut::setNormalized(bool normalized) { - buffer_data_.reserve(buffer_data_.size() + (size_t)n * this->dim); - for (idx_t i = 0; i < n; i++) - insert(data + (size_t)i * this->dim); + if (!normalized) throw std::runtime_error("Coconut supports Gaussian breakpoints only; input is never normalized"); + SimilaritySearchAlgorithm::setNormalized(true); } - - bool Coconut::readLeaf(size_t leaf_no, std::vector &page) const + std::string Coconut::getIndexDirectory() const { return state_ ? state_->directory : std::string(); } + idx_t Coconut::getBufferedRecords() const { return state_ ? state_->buffered() : 0; } + idx_t Coconut::getDiskRecords() const { return n_database - getBufferedRecords(); } + idx_t Coconut::getDiskBytes() const + { return state_ ? getDiskRecords() * (state_->stride + state_->dim * sizeof(float)) : 0; } + idx_t Coconut::getRunCount() const { return state_ ? state_->snapshot().size() : 0; } + idx_t Coconut::getFlushCount() const { return state_ ? state_->flushes : 0; } + idx_t Coconut::getMergeCount() const { return state_ ? state_->merges.load() : 0; } + idx_t Coconut::getPeakBufferedRecords() const { return state_ ? state_->peak_buffered : 0; } + + void Coconut::buildIndex(DataSource *source) { - const size_t begin = leaf_no * (size_t)this->leaf_capacity; - if (begin >= records_.size()) - return false; - const size_t count = std::min((size_t)this->leaf_capacity, records_.size() - begin); - const size_t record_bytes = (size_t)this->paa_segments * sizeof(sax_type) + - (size_t)this->dim * sizeof(float); - page.resize(count * record_bytes); - std::unique_ptr file(fopen(leafPath((int)leaf_no).c_str(), "rb"), &fclose); - return file && fread(page.data(), 1, page.size(), file.get()) == page.size(); + if (!source) throw std::runtime_error("Coconut datasource is null"); + idx_t d = source->getDim(); + if (!d || d > (idx_t)std::numeric_limits::max() || paa_segments <= 0 || d % paa_segments || + sax_cardinality < 1 || sax_cardinality > 8 || leaf_capacity_ <= 0 || !buffer_capacity_) + throw std::runtime_error("Coconut requires positive leaf/PAA sizes, buffer capacity, 1-8 SAX bits, and a dimension divisible by PAA size"); + checkedSize(buffer_capacity_, checkedSize(d, sizeof(float))); + auto fresh = std::make_unique(d, paa_segments, sax_cardinality, leaf_capacity_, buffer_capacity_, storage_parent_, async_merges_); + set_active_breakpoints(nullptr, nullptr); source->reset(); + idx_t expected = source->getTotalRecords(); const float *raw = source->rawPointer(); + if (raw) + for (idx_t start = 0; start < expected; start += buffer_capacity_) + { + idx_t take = std::min(buffer_capacity_, expected - start); + checkFinite(raw + checkedSize(start, (size_t)d), take, d); + fresh->append(raw + checkedSize(start, (size_t)d), take, num_threads); + } + else + { + std::vector chunk(checkedSize(buffer_capacity_, (size_t)d)); + while (!expected || fresh->count < expected) + { + idx_t got = 0, limit = expected ? std::min(buffer_capacity_, expected - fresh->count) : buffer_capacity_; + while (got < limit && source->nextRecord(chunk.data() + checkedSize(got, (size_t)d))) ++got; + if (!got) break; + checkFinite(chunk.data(), got, d); fresh->append(chunk.data(), got, num_threads); + if (got < limit) break; + } + if (expected && fresh->count != expected) + throw std::runtime_error("Coconut datasource ended before all records were read"); + } + fresh->flush(); fresh->consolidateBuild(); + dim = d; n_database = fresh->count; state_ = std::move(fresh); } - - void Coconut::copySeriesFromLeaf(const std::vector &page, size_t slot, float *out) const + void Coconut::insert(const float *data) { insertBatch(data, 1); } + void Coconut::insertBatch(const float *data, idx_t n) { - const size_t key_bytes = (size_t)this->paa_segments * sizeof(sax_type); - const size_t series_bytes = (size_t)this->dim * sizeof(float); - // Copy instead of casting: a key length need not leave the raw floats aligned. - memcpy(out, page.data() + slot * (key_bytes + series_bytes) + key_bytes, series_bytes); + requireBuilt(); checkFinite(data, n, dim); set_active_breakpoints(nullptr, nullptr); + try { state_->append(data, n, num_threads); } + catch (...) { n_database = state_->count; throw; } + n_database = state_->count; } + void Coconut::flush() { requireBuilt(); state_->flush(); } - void Coconut::searchIndex(const float *query, const idx_t n_query, const idx_t k, idx_t *I, float *D) + void Coconut::searchIndex(const float *query, idx_t n_query, idx_t k, idx_t *I, float *D) + { searchTopK(query, n_query, k, 0, I, D); } + + void Coconut::searchIndexRecent(const float *query, idx_t n_query, idx_t k, idx_t window, idx_t *I, float *D) { - const int seg = this->paa_segments; - if (this->distance_type != DistanceType::L2_SQUARED) - throw std::runtime_error("Coconut top-k search only supports L2_SQUARED"); - if (this->ts_values_per_segment_ <= 0) - throw std::runtime_error("Coconut top-k search requires an initial buildIndex first"); - if (k == 0) - throw std::runtime_error("Coconut top-k search requires k > 0"); - const size_t cap = (size_t)this->leaf_capacity; - const size_t n_leaves = (records_.size() + cap - 1) / cap; - std::atomic io_failed{false}; - const int workers = (int)std::max(1, std::min(std::max(1, num_threads), n_query)); + requireBuilt(); + if (!window) throw std::runtime_error("Coconut recent window must be positive"); + searchTopK(query, n_query, k, window < n_database ? n_database - window : 0, I, D); + } -#pragma omp parallel for num_threads(workers) schedule(dynamic, 1) - for (idx_t q = 0; q < n_query; q++) + void Coconut::searchTopK(const float *query, idx_t n_query, idx_t k, idx_t first_id, idx_t *I, float *D) + { + requireBuilt(); + if (distance_type != DistanceType::L2_SQUARED) throw std::runtime_error("Coconut top-k search only supports L2_SQUARED"); + if (!k) throw std::runtime_error("Coconut top-k search requires k > 0"); + checkFinite(query, n_query, dim); set_active_breakpoints(nullptr, nullptr); + const auto &s = *state_; auto runs = s.snapshot(); + for (const auto &run : runs) run->validate(); + last_stats_ = {}; + // Batch SIMS keeps Coconut's seed -> SAX bounds -> exact refinement. + // Queries share bounded ordered reads, never their heaps or pruning bound. + struct alignas(64) Work { - std::vector paa(seg); - std::vector ts(this->dim); - std::vector page; - const float *qts = query + (size_t)q * this->dim; - paa_from_ts(qts, paa.data(), seg, this->ts_values_per_segment_); - - // k-BSF as a max-heap on distance (top = current worst of the best k). - std::priority_queue> bsf; - - // Upstream Coconut first seeks the query key and visits nearby records. - // Our sorted, resident key directory replaces DB_SET_RANGE/DB_PREV/DB_NEXT. - // For top-k the seed must contain k distinct records whenever they exist. - std::vector query_sax(seg), query_key(seg); - sax_from_paa(paa.data(), query_sax.data(), seg, - this->sax_alphabet_cardinality_, this->sax_cardinality); - sortable_sax_from_sax(query_sax.data(), query_key.data(), seg, this->sax_cardinality); - auto match = std::lower_bound(records_.begin(), records_.end(), query_key, - [](const Record &record, const std::vector &key) { - return key_less(record.key, key); - }); - const size_t center = (size_t)(match - records_.begin()); - const size_t window = std::min(records_.size(), std::max(cap, (size_t)k)); - const size_t seed_begin = std::min(center > window / 2 ? center - window / 2 : 0, - records_.size() - window); - const size_t seed_end = seed_begin + window; - size_t cached_leaf = n_leaves; - - auto refine = [&](size_t p) { - const size_t leaf_no = p / cap; - if (cached_leaf != leaf_no) - { - if (!readLeaf(leaf_no, page)) - { - io_failed.store(true); - return false; - } - cached_leaf = leaf_no; - } - copySeriesFromLeaf(page, p % cap, ts.data()); - float kth = bsf.size() >= (size_t)k ? bsf.top().first : FLT_MAX; - float dist = ts_euclidean_distance_SIMD(const_cast(qts), ts.data(), - (int)this->dim, kth); - if (bsf.size() < (size_t)k) - bsf.push({dist, records_[p].series_id}); - else if (dist < bsf.top().first) - { - bsf.pop(); - bsf.push({dist, records_[p].series_id}); - } - return true; + const float *query = nullptr; + std::vector paa; + coconut_detail::QueryBounds lookup; + std::vector key; + std::vector> windows; + std::priority_queue> best; + idx_t candidates = 0, evaluations = 0, seeds = 0; + }; + const idx_t max_raw_bytes = 8 * 1024 * 1024; + const size_t raw_stride = checkedSize(dim, sizeof(float)); + const idx_t block_capacity = std::min(32768, std::max(1, max_raw_bytes / raw_stride)); + for (idx_t base = 0; base < n_query; base += query_batch_size_) + { + idx_t count = std::min(query_batch_size_, n_query - base); + const int workers = std::max(1, num_threads); + std::vector work(count); + std::vector page; + page.reserve(checkedSize(block_capacity, (size_t)dim)); + auto refine = [&](Work &w, const float *ts, idx_t id, bool seed) { + if (id < first_id) return; + ++w.evaluations; if (seed) ++w.seeds; + float kth = w.best.size() >= k ? w.best.top().first : FLT_MAX; + float distance = squaredL2(w.query, ts, dim, std::nextafter(kth, FLT_MAX)); + std::pair candidate(distance, id); + if (w.best.size() < k) w.best.push(candidate); + else if (candidate < w.best.top()) { w.best.pop(); w.best.push(candidate); } }; - for (size_t p = seed_begin; p < seed_end; p++) - if (!refine(p)) - break; - if (io_failed.load()) - continue; - - // Exact refinement still considers every unseeded SAX summary. A surviving - // candidate loads its whole leaf once, avoiding per-record seeks and reads. - for (size_t leaf_no = 0; leaf_no < n_leaves; leaf_no++) + auto read = [&](const std::shared_ptr &run, idx_t begin, idx_t records, bool seed = false) { + run->leaf(begin, records, page); + ++last_stats_.leaf_reads; last_stats_.raw_bytes_read += records * raw_stride; + last_stats_.raw_records_read += records; + if (seed) last_stats_.seed_raw_records_read += records; + }; + auto started = Clock::now(); + for (idx_t q = 0; q < count; ++q) { - const size_t begin = leaf_no * cap; - const size_t end = std::min(begin + cap, records_.size()); - for (size_t p = begin; p < end; p++) + auto &w = work[q]; w.query = query + checkedSize(base + q, (size_t)dim); + w.paa.resize(s.seg); w.key.resize(s.seg); w.windows.resize(runs.size()); + std::vector word(s.seg); + paa_from_ts(const_cast(w.query), w.paa.data(), s.seg, s.values); + w.lookup.reset(w.paa.data(), s.seg, s.bits, (float)s.values); + sax_from_paa(w.paa.data(), word.data(), s.seg, s.alphabet, s.bits); + sortable_sax_from_sax(word.data(), w.key.data(), s.seg, s.bits); + for (size_t r = 0; r < runs.size(); ++r) { - if (p >= seed_begin && p < seed_end) - continue; // each seeded ID has already been considered exactly once - float kth = (bsf.size() >= (size_t)k) ? bsf.top().first : FLT_MAX; - - float mindist = minidist_paa_to_isax( - paa.data(), - const_cast(records_[p].sax.data()), - const_cast(this->max_cardinalities_.data()), - (sax_type)this->sax_cardinality, - this->sax_alphabet_cardinality_, - seg, MINVAL, MAXVAL, this->mindist_sqrt_); - - if (mindist >= kth) - continue; // lower bound already worse than the k-th best -> prune - - if (!refine(p)) - break; + const auto &run = runs[r]; + if (run->end_id <= first_id) { ++last_stats_.runs_skipped; continue; } + ++last_stats_.runs_searched; + idx_t target = run->locate(w.key.data()); + idx_t window = std::min(run->n, std::max(s.leaf_size, k)); + idx_t begin = std::min(target > window / 2 ? target - window / 2 : 0, run->n - window); + w.windows[r] = {begin, begin + window}; + for (idx_t first = begin; first < begin + window;) + { + idx_t end = std::min(begin + window, std::min(first + block_capacity, run->leafBegin(run->leafOf(first) + 1))); + read(run, first, end - first, true); + for (idx_t pos = first; pos < end; ++pos) + refine(w, page.data() + checkedSize(pos - first, (size_t)dim), run->id(pos, s.seg), true); + first = end; + } } - if (io_failed.load()) - break; + for (idx_t b = 0; b < s.buffered(); ++b) + refine(w, s.data.data() + checkedSize(b, (size_t)dim), s.buffer_first + b, true); } - if (io_failed.load()) - continue; - - // Scan the in-memory streaming buffer (raw TS already in memory). - for (size_t b = 0; b < buffer_records_.size(); b++) + last_stats_.approximate_seconds += elapsed(started); + // Scratch is independent of N: one raw block plus query-count x block + // bounds/masks. SAX summaries remain the only resident O(N) search data. + std::vector bounds(checkedSize(count, (size_t)block_capacity)); + std::vector candidate(bounds.size()), active(block_capacity); + page.reserve(checkedSize(block_capacity, (size_t)dim)); + size_t lookup_bytes = 0; for (const auto &w : work) lookup_bytes += w.lookup.bytes(); + last_stats_.peak_scratch_bytes = std::max(last_stats_.peak_scratch_bytes, + bounds.size() * sizeof(float) + candidate.size() + active.size() + page.capacity() * sizeof(float) + lookup_bytes); + for (size_t r = 0; r < runs.size(); ++r) { - float kth = (bsf.size() >= (size_t)k) ? bsf.top().first : FLT_MAX; - float mindist = minidist_paa_to_isax( - paa.data(), - const_cast(buffer_records_[b].sax.data()), - const_cast(this->max_cardinalities_.data()), - (sax_type)this->sax_cardinality, - this->sax_alphabet_cardinality_, - seg, MINVAL, MAXVAL, this->mindist_sqrt_); - if (mindist >= kth) - continue; - float *bts = buffer_data_.data() + b * (size_t)this->dim; - float dist = ts_euclidean_distance_SIMD(const_cast(qts), bts, - (int)this->dim, kth); - if (bsf.size() < (size_t)k) - bsf.push({dist, buffer_records_[b].series_id}); - else if (dist < bsf.top().first) + const auto &run = runs[r]; if (run->end_id <= first_id) continue; + last_stats_.summaries_examined += count * run->n; + const bool boundary_run = run->first_id < first_id; + for (idx_t begin = 0; begin < run->n; begin += block_capacity) { - bsf.pop(); - bsf.push({dist, buffer_records_[b].series_id}); + idx_t records = std::min(block_capacity, run->n - begin); + started = Clock::now(); + std::vector threshold(count); + for (idx_t q = 0; q < count; ++q) + threshold[q] = work[q].best.size() >= k ? work[q].best.top().first : FLT_MAX; +#pragma omp parallel for collapse(2) num_threads(workers) schedule(static) + for (idx_t q = 0; q < count; ++q) + for (idx_t i = 0; i < records; ++i) + { + const auto &w = work[q]; size_t row = (size_t)q * block_capacity; + idx_t pos = begin + i; + float bound = w.lookup(run->summaries.data() + (size_t)pos * s.seg, threshold[q]); + bounds[row + i] = bound; + bool survives = bound <= threshold[q] && !(pos >= w.windows[r].first && pos < w.windows[r].second); + if (survives && boundary_run) survives = run->id(pos, s.seg) >= first_id; + candidate[row + i] = survives; + } + for (idx_t q = 0; q < count; ++q) + for (idx_t i = 0; i < records; ++i) + work[q].candidates += candidate[(size_t)q * block_capacity + i]; + last_stats_.bounds_seconds += elapsed(started); + started = Clock::now(); + std::fill_n(active.data(), records, 0); + for (idx_t q = 0; q < count; ++q) + { + const auto *mask = candidate.data() + (size_t)q * block_capacity; + for (idx_t i = 0; i < records; ++i) active[i] |= mask[i]; + } + idx_t active_count = 0; + for (idx_t i = 0; i < records; ++i) active_count += active[i] != 0; + last_stats_.unique_candidate_records += active_count; + // Dense blocks trade at most 2x unique-candidate bytes for + // a single sequential read; sparse blocks retain short gaps. + size_t max_gap = active_count * 2 >= records ? records * raw_stride : 4096; + // Sparse positions are coalesced only across <=4 KiB gaps. + // Dense candidate blocks become one bounded read reused by + // every query in this tile. No unconditional raw full scan. + idx_t cursor = 0; + while (cursor < records) + { + while (cursor < records && !active[cursor]) ++cursor; + if (cursor == records) break; + idx_t first = cursor, last = cursor++; + while (cursor < records) + { + if (active[cursor]) last = cursor; + else if ((cursor - last) * raw_stride > max_gap) break; + ++cursor; + } + idx_t end = last + 1; + read(run, begin + first, end - first); + std::exception_ptr error; +#pragma omp parallel for num_threads(workers) if(count > 1) schedule(static) + for (idx_t q = 0; q < count; ++q) + { + try + { + auto &w = work[q]; size_t row = (size_t)q * block_capacity; + for (idx_t i = first; i < end; ++i) + { + if (!candidate[row + i]) continue; + float kth = w.best.size() >= k ? w.best.top().first : FLT_MAX; + if (bounds[row + i] > kth) continue; + refine(w, page.data() + checkedSize(i - first, (size_t)dim), run->id(begin + i, s.seg), false); + } + } + catch (...) + { +#pragma omp critical(coconut_batch_query_error) + { if (!error) error = std::current_exception(); } + } + } + if (error) std::rethrow_exception(error); + } + last_stats_.refinement_seconds += elapsed(started); } } - - // Emit in ascending distance order. - idx_t got = (idx_t)bsf.size(); - for (idx_t j = got; j > 0; j--) - { - D[q * k + (j - 1)] = bsf.top().first; - I[q * k + (j - 1)] = bsf.top().second; - bsf.pop(); - } - for (idx_t j = got; j < k; j++) // pad if fewer than k available + for (idx_t q = 0; q < count; ++q) { - D[q * k + j] = FLT_MAX; - I[q * k + j] = 0; + auto &w = work[q]; last_stats_.candidate_records += w.candidates; + last_stats_.distance_evaluations += w.evaluations; last_stats_.seed_evaluations += w.seeds; + idx_t got = w.best.size(), offset = (base + q) * k; + for (idx_t j = got; j; --j) + { D[offset + j - 1] = w.best.top().first; I[offset + j - 1] = w.best.top().second; w.best.pop(); } + for (idx_t j = got; j < k; ++j) { I[offset + j] = 0; D[offset + j] = FLT_MAX; } } } - if (io_failed.load()) - throw std::runtime_error("Coconut failed to open or read an index leaf"); } - void Coconut::searchIndex(const float *query, idx_t n_query, const SearchConfig &config, - std::vector> &I, - std::vector> &D) + std::vector> &I, std::vector> &D) { if (config.type == QueryType::TOP_K) + { SimilaritySearchAlgorithm::searchIndex(query, n_query, config, I, D); return; } + requireBuilt(); + if (distance_type != DistanceType::L2_SQUARED) throw std::runtime_error("Coconut range search only supports L2_SQUARED"); + if (std::isnan(config.r) || config.r < 0) throw std::runtime_error("Coconut radius must be nonnegative"); + checkFinite(query, n_query, dim); state_->validateRuns(); set_active_breakpoints(nullptr, nullptr); + I.assign(n_query, {}); D.assign(n_query, {}); + const auto &s = *state_; auto runs = s.snapshot(); std::exception_ptr error; + int workers = (int)std::max(1, std::min(std::max(1, num_threads), n_query)); +#pragma omp parallel for num_threads(workers) schedule(dynamic, 1) + for (idx_t q = 0; q < n_query; ++q) { - SimilaritySearchAlgorithm::searchIndex(query, n_query, config, I, D); - return; - } - if (this->distance_type != DistanceType::L2_SQUARED) - throw std::runtime_error("Coconut range search only supports L2_SQUARED."); - - const int seg = this->paa_segments; - const float r = config.r; - // Early-abandon just past r: a partial sum can then never tie with r and be mistaken - // for a hit, while every series actually within the radius is still summed in full. - const float abandon_bound = std::nextafter(r, FLT_MAX); - - I.assign(n_query, {}); - D.assign(n_query, {}); - std::atomic io_failed{false}; - -#pragma omp parallel for num_threads(num_threads) - for (idx_t q = 0; q < n_query; q++) - { - const float *qts = query + (size_t)q * this->dim; - std::vector paa(seg); - std::vector ts(this->dim); - std::vector> hits; - std::vector page; - - paa_from_ts(qts, paa.data(), seg, this->ts_values_per_segment_); - - // Read a candidate leaf as one block, then refine candidates in memory. - const size_t cap = (size_t)this->leaf_capacity; - const size_t n_leaves = (records_.size() + cap - 1) / cap; - std::vector candidates; - - for (size_t leaf_no = 0; leaf_no < n_leaves; leaf_no++) + try { - const size_t begin = leaf_no * cap; - const size_t end = std::min(begin + cap, records_.size()); - - candidates.clear(); - for (size_t p = begin; p < end; p++) - { - float mindist = minidist_paa_to_isax( - paa.data(), - const_cast(records_[p].sax.data()), - const_cast(this->max_cardinalities_.data()), - (sax_type)this->sax_cardinality, - this->sax_alphabet_cardinality_, - seg, MINVAL, MAXVAL, this->mindist_sqrt_); - - if (mindist <= r) // lower bound outside the radius -> prune - candidates.push_back(p); - } - - if (candidates.empty()) - continue; - - if (!readLeaf(leaf_no, page)) - { - io_failed.store(true); - break; - } - for (size_t p : candidates) - { - copySeriesFromLeaf(page, p - begin, ts.data()); - float dist = ts_euclidean_distance_SIMD(const_cast(qts), ts.data(), - (int)this->dim, abandon_bound); - if (dist <= r) - hits.emplace_back(dist, records_[p].series_id); - } - if (io_failed.load()) - break; - } - - // Scan the in-memory streaming buffer (raw TS already in memory). - for (size_t b = 0; b < buffer_records_.size(); b++) - { - float mindist = minidist_paa_to_isax( - paa.data(), - const_cast(buffer_records_[b].sax.data()), - const_cast(this->max_cardinalities_.data()), - (sax_type)this->sax_cardinality, - this->sax_alphabet_cardinality_, - seg, MINVAL, MAXVAL, this->mindist_sqrt_); - - if (mindist > r) - continue; - - float *bts = buffer_data_.data() + b * (size_t)this->dim; - float dist = ts_euclidean_distance_SIMD(const_cast(qts), bts, - (int)this->dim, abandon_bound); - if (dist <= r) - hits.emplace_back(dist, buffer_records_[b].series_id); + const float *qts = query + checkedSize(q, (size_t)dim); + std::vector paa(s.seg), page; paa_from_ts(const_cast(qts), paa.data(), s.seg, s.values); + std::vector> hits; + auto refine = [&](const float *ts, idx_t id) { + float distance = squaredL2(qts, ts, dim, std::nextafter(config.r, FLT_MAX)); + if (distance <= config.r) hits.emplace_back(distance, id); + }; + for (const auto &run : runs) + for (idx_t leaf = 0; leaf < run->leaf_count; ++leaf) + { + idx_t begin = run->leafBegin(leaf), end = run->leafBegin(leaf + 1); std::vector candidates; + for (idx_t p = begin; p < end; ++p) + if (s.bound(paa.data(), run->summaries.data() + (size_t)p * s.seg) <= config.r) candidates.push_back(p); + if (candidates.empty()) continue; + run->leaf(begin, end - begin, page); + for (idx_t p : candidates) + refine(page.data() + checkedSize(p - begin, (size_t)dim), run->id(p, s.seg)); + } + for (idx_t b = 0; b < s.buffered(); ++b) + if (s.bound(paa.data(), s.sax.data() + (size_t)b * s.seg) <= config.r) + refine(s.data.data() + checkedSize(b, (size_t)dim), s.buffer_first + b); + std::sort(hits.begin(), hits.end()); + for (const auto &hit : hits) { D[q].push_back(hit.first); I[q].push_back(hit.second); } } - - std::sort(hits.begin(), hits.end()); - I[q].resize(hits.size()); - D[q].resize(hits.size()); - for (size_t j = 0; j < hits.size(); j++) + catch (...) { - D[q][j] = hits[j].first; - I[q][j] = hits[j].second; +#pragma omp critical(coconut_query_error) + { if (!error) error = std::current_exception(); } } } - if (io_failed.load()) - throw std::runtime_error("Coconut failed to open or read an index leaf"); + if (error) std::rethrow_exception(error); } - - void Coconut::cleanupLeafFiles() - { - if (index_dir_.empty()) - return; - int nleaves = (int)((records_.size() + leaf_capacity - 1) / (size_t)std::max(1, leaf_capacity)); - for (int l = 0; l < nleaves; l++) - remove(leafPath(l).c_str()); - rmdir(index_dir_.c_str()); - index_dir_.clear(); - } - } diff --git a/lib/algos/Coconut.hpp b/lib/algos/Coconut.hpp index ecb77ef..f141ade 100644 --- a/lib/algos/Coconut.hpp +++ b/lib/algos/Coconut.hpp @@ -1,91 +1,82 @@ #ifndef COCONUT_HPP #define COCONUT_HPP - #include "SimilaritySearchAlgorithm.hpp" - -#include "../isax/iSAXTypes.hpp" - -#include +#include #include -#include -#include namespace daisy { + struct CoconutSearchStats + { + idx_t summaries_examined = 0, leaf_reads = 0, raw_bytes_read = 0; + idx_t runs_searched = 0, runs_skipped = 0; + idx_t candidate_records = 0, unique_candidate_records = 0; + idx_t distance_evaluations = 0, seed_evaluations = 0, raw_records_read = 0, seed_raw_records_read = 0; + idx_t peak_scratch_bytes = 0; + double approximate_seconds = 0, bounds_seconds = 0, refinement_seconds = 0; + }; - // COCONUT (Kondylakis et al., VLDB'18): a "sortable SAX" key (bit-interleaved SAX) orders - // series so the index can be built bottom-up by sorting and extended by streaming inserts. - // buildIndex sorts records onto on-disk leaves ([inv-SAX][raw TS] each); searchIndex does - // exact kNN or range search via SAX MINDIST + L2 refinement, parallel across queries. - // Top-k starts with an ordered approximate search around the query's sortable-SAX key, - // followed by exact SAX-bound refinement. Each query caches a complete leaf privately. - // The sorted in-memory key directory provides ordered lookup without Berkeley DB. - // Follow-ups: out-of-core key directory/sort, LSM merges, equi-depth breakpoints. + // Paper-based Coconut-LSM-Full: bounded RAM memtable, resident summaries, + // sorted disk runs, binary compaction and file-backed key directories. + // Writes and queries are not concurrent; the index is temporary, not a + // recoverable database. Gaussian SAX never normalizes the input vectors. class Coconut : public SimilaritySearchAlgorithm { public: - Coconut(DistanceType distance_type); - - void setNumThreads(int num_threads) override { this->num_threads = num_threads; } + explicit Coconut(DistanceType type); + ~Coconut() override; + Coconut(const Coconut &) = delete; + Coconut &operator=(const Coconut &) = delete; + void setNumThreads(int n) override { num_threads = n; } int getNumThreads() const { return num_threads; } - int getPaaSegments() const { return paa_segments; } - void setPaaSegments(int v) { paa_segments = v; } + void setPaaSegments(int n); int getSaxCardinality() const { return sax_cardinality; } - void setSaxCardinality(int v) { sax_cardinality = v; } - int getLeafSize() const { return leaf_capacity; } - void setLeafSize(int v) { leaf_capacity = v; } - + void setSaxCardinality(int n); + int getLeafSize() const { return leaf_capacity_; } + void setLeafSize(int n); + idx_t getBufferCapacity() const { return buffer_capacity_; } + void setBufferCapacity(idx_t records); + void setIndexDirectory(const std::string &parent); + std::string getIndexDirectory() const; + idx_t getBufferedRecords() const; + idx_t getDiskRecords() const; + idx_t getDiskBytes() const; + idx_t getRunCount() const; + idx_t getFlushCount() const; + idx_t getMergeCount() const; + idx_t getPeakBufferedRecords() const; + void setQueryBatchSize(idx_t queries); + idx_t getQueryBatchSize() const { return query_batch_size_; } + void setAsyncMerges(bool enabled); + bool getAsyncMerges() const { return async_merges_; } + void waitForMerges(); + idx_t getResidentSummaryBytes() const; + CoconutSearchStats getLastSearchStats() const { return last_stats_; } + // BTP recency is expressed as a count of most recently inserted vectors. + void searchIndexRecent(const float *query, idx_t n_query, idx_t k, idx_t window, + idx_t *I, float *D); + void setNormalized(bool normalized) override; using SimilaritySearchAlgorithm::buildIndex; - void buildIndex(DataSource *data_source) override; - void buildIndex(const std::string &filename, idx_t dim, idx_t n_database = 0) override - { - throw std::runtime_error("Coconut requires in-memory data. Use buildIndex(database, n_database, dim)."); - } - - void searchIndex(const float *query, const idx_t n_query, const idx_t k, idx_t *I, float *D) override; - - // Range search: every series within radius r of the query, in ascending distance order. + void buildIndex(DataSource *source) override; + void searchIndex(const float *query, idx_t n_query, idx_t k, idx_t *I, float *D) override; void searchIndex(const float *query, idx_t n_query, const SearchConfig &config, - std::vector> &I, - std::vector> &D) override; - - // Streaming: add series to a live index (needs buildIndex first). New series go into - // an in-memory buffer that queries scan alongside the on-disk index. + std::vector> &I, std::vector> &D) override; void insert(const float *series) override; void insertBatch(const float *data, idx_t n) override; - - ~Coconut() override; - + void flush(); private: - // After sorting, record p lives in leaf p/leaf_capacity at slot p%leaf_capacity. - struct Record - { - std::vector sax; // SAX word (for MINDIST) - std::vector key; // sortable SAX (defines the order) - idx_t series_id; // original index in the input - }; - - std::vector records_; // main index, sorted by sortable-SAX key (TS on disk) - std::string index_dir_; // unique per-instance dir for leaf files - int leaf_capacity = 1000; // records per leaf file - - // Streaming buffer: inserted series kept in memory (summaries + raw TS), scanned by queries. - std::vector buffer_records_; - std::vector buffer_data_; // buffer_records_.size() * dim - idx_t next_id_ = 0; - - int ts_values_per_segment_ = 0; // dim / paa_segments - int sax_alphabet_cardinality_ = 0; // 2^sax_cardinality - float mindist_sqrt_ = 0.0f; // (dim / paa_segments), the MINDIST scale factor - std::vector max_cardinalities_; // all = sax_cardinality - - std::string leafPath(int leaf_no) const; - bool readLeaf(size_t leaf_no, std::vector &page) const; - void copySeriesFromLeaf(const std::vector &page, size_t slot, float *out) const; - void cleanupLeafFiles(); + struct DiskState; + std::unique_ptr state_; + int leaf_capacity_ = 1000; + idx_t buffer_capacity_ = 1000000; + std::string storage_parent_; + bool async_merges_ = true; + idx_t query_batch_size_ = 32; + CoconutSearchStats last_stats_; + void searchTopK(const float *query, idx_t n_query, idx_t k, idx_t first_id, idx_t *I, float *D); + void requireMutableConfig() const; + void requireBuilt() const; }; - } - #endif diff --git a/lib/algos/CoconutBounds.hpp b/lib/algos/CoconutBounds.hpp new file mode 100644 index 0000000..df1cee1 --- /dev/null +++ b/lib/algos/CoconutBounds.hpp @@ -0,0 +1,46 @@ +#ifndef DAISY_COCONUT_BOUNDS_HPP +#define DAISY_COCONUT_BOUNDS_HPP +#include "../isax/SAX.hpp" +#include +#include + +namespace daisy { namespace coconut_detail { +// Full-cardinality SAX interval distances for a single query. Double table +// entries preserve the reference helper's pow(float_delta, 2) contributions; +// the float accumulator keeps its original per-segment rounding/order. +class QueryBounds +{ + int segments_ = 0, alphabet_ = 0; + float ratio_ = 1; + std::vector contribution_; +public: + void reset(const float *paa, int segments, int bits, float ratio) + { + segments_ = segments; alphabet_ = 1 << bits; ratio_ = ratio; + const int offset = (alphabet_ - 1) * (alphabet_ - 2) / 2; + contribution_.resize((size_t)segments * alphabet_); + for (int d = 0; d < segments; ++d) + for (int symbol = 0; symbol < alphabet_; ++symbol) + { + float low = symbol ? sax_breakpoints[offset + symbol - 1] : MINVAL; + float high = symbol == alphabet_ - 1 ? MAXVAL : sax_breakpoints[offset + symbol]; + float delta = low > paa[d] ? low - paa[d] : high < paa[d] ? high - paa[d] : 0; + contribution_[(size_t)d * alphabet_ + symbol] = std::pow(delta, 2); + } + } + float operator()(const sax_type *word, float limit) const + { + float distance = 0; + for (int d = 0; d < segments_; ++d) + { + distance += contribution_[(size_t)d * alphabet_ + word[d]]; + // All contributions are nonnegative. A partial lower bound above + // the current bound safely prunes without evaluating later segments. + if ((d & 3) == 3 && ratio_ * distance > limit) return ratio_ * distance; + } + return ratio_ * distance; + } + size_t bytes() const { return contribution_.capacity() * sizeof(double); } +}; +}} +#endif diff --git a/lib/daisy.hpp b/lib/daisy.hpp index 8ad5da0..ae8ea48 100644 --- a/lib/daisy.hpp +++ b/lib/daisy.hpp @@ -3,7 +3,9 @@ #include "algos/Bruteforce.hpp" #include "algos/LbBruteforce.hpp" +#ifdef COCONUT_ENABLED #include "algos/Coconut.hpp" +#endif #include "algos/Messi.hpp" #include "algos/ParIS.hpp" #include "algos/Sing.hpp" diff --git a/pybinds/setup.cpp b/pybinds/setup.cpp index 2c1636c..71a2be8 100644 --- a/pybinds/setup.cpp +++ b/pybinds/setup.cpp @@ -380,7 +380,7 @@ PYBIND11_MODULE(_core, m) ////// COCONUT ////// #ifdef COCONUT_ENABLED - pybind11::class_(m, "Coconut", "COCONUT sortable-SAX index (static bottom-up build + streaming insert)") + pybind11::class_(m, "Coconut", "Paper-based Coconut-LSM-Full with resident summaries and within-query parallel SIMS") .def(pybind11::init(), "Create a new Coconut with the given distance metric") .def("getPaaSegments", &daisy::Coconut::getPaaSegments, "Get the number of PAA segments") @@ -390,9 +390,55 @@ PYBIND11_MODULE(_core, m) .def("setPaaSegments", &daisy::Coconut::setPaaSegments, "Set the number of PAA segments (must divide the series length)") .def("setSaxCardinality", &daisy::Coconut::setSaxCardinality, "Set the SAX bit cardinality") .def("setLeafSize", &daisy::Coconut::setLeafSize, "Set the records-per-leaf-file capacity") + .def("setBufferCapacity", &daisy::Coconut::setBufferCapacity, "Set the maximum RAM memtable records before building; full buffers persist and schedule compaction") + .def("getBufferCapacity", &daisy::Coconut::getBufferCapacity) + .def("setIndexDirectory", &daisy::Coconut::setIndexDirectory, "Place the temporary disk index in a unique child of this parent directory") + .def("getIndexDirectory", &daisy::Coconut::getIndexDirectory) + .def("getBufferedRecords", &daisy::Coconut::getBufferedRecords) + .def("getDiskRecords", &daisy::Coconut::getDiskRecords) + .def("getDiskBytes", &daisy::Coconut::getDiskBytes) + .def("getRunCount", &daisy::Coconut::getRunCount) + .def("getFlushCount", &daisy::Coconut::getFlushCount) + .def("getMergeCount", &daisy::Coconut::getMergeCount) + .def("getPeakBufferedRecords", &daisy::Coconut::getPeakBufferedRecords) + .def("setAsyncMerges", &daisy::Coconut::setAsyncMerges, "Enable background compaction before building (default true)") + .def("getAsyncMerges", &daisy::Coconut::getAsyncMerges) + .def("setQueryBatchSize", &daisy::Coconut::setQueryBatchSize, "Share bounded ordered candidate reads across 1-64 independent queries (default 32)") + .def("getQueryBatchSize", &daisy::Coconut::getQueryBatchSize) + .def("waitForMerges", &daisy::Coconut::waitForMerges, pybind11::call_guard(), + "Wait for pending compactions; include this barrier in durable insertion timing") + .def("getResidentSummaryBytes", &daisy::Coconut::getResidentSummaryBytes) + .def("getLastSearchStats", [](const daisy::Coconut &self) { + auto s = self.getLastSearchStats(); pybind11::dict d; + d["summaries_examined"] = s.summaries_examined; d["leaf_reads"] = s.leaf_reads; + d["raw_bytes_read"] = s.raw_bytes_read; d["runs_searched"] = s.runs_searched; + d["runs_skipped"] = s.runs_skipped; d["approximate_seconds"] = s.approximate_seconds; + d["candidate_records"] = s.candidate_records; d["unique_candidate_records"] = s.unique_candidate_records; + d["distance_evaluations"] = s.distance_evaluations; d["seed_evaluations"] = s.seed_evaluations; + d["seed_raw_records_read"] = s.seed_raw_records_read; + d["raw_records_read"] = s.raw_records_read; d["peak_scratch_bytes"] = s.peak_scratch_bytes; + d["bounds_seconds"] = s.bounds_seconds; d["refinement_seconds"] = s.refinement_seconds; + return d; + }, "Phase times and logical read counters for the last top-k call, including all queries") + .def("searchIndexRecent", [](daisy::Coconut &self, + pybind11::array_t query, + daisy::idx_t k, daisy::idx_t window) { + auto buf = query.request(); + if (buf.ndim != 2 || buf.shape[1] != static_cast(self.getDim())) + throw std::runtime_error("Coconut query dimension mismatch"); + if (!k) throw std::runtime_error("Coconut top-k search requires k > 0"); + pybind11::array_t ids({static_cast(buf.shape[0]), static_cast(k)}); + pybind11::array_t distances({static_cast(buf.shape[0]), static_cast(k)}); + self.searchIndexRecent(static_cast(buf.ptr), buf.shape[0], k, window, ids.mutable_data(), distances.mutable_data()); + return pybind11::make_tuple(ids, distances); + }, pybind11::arg("queries"), pybind11::arg("k"), pybind11::arg("window_size"), + "Exact top-k over the most recent window_size insertions; returns original stream IDs") + .def("getNDatabase", &daisy::Coconut::getNDatabase) + .def("getDim", &daisy::Coconut::getDim) + .def("flush", &daisy::Coconut::flush, "Persist the partially full memtable; inserts already flush automatically at capacity") // Static bottom-up build from a 2D float32 numpy array. - .def("buildIndex", [](daisy::Coconut &self, pybind11::array_t db) + .def("buildIndex", [](daisy::Coconut &self, pybind11::array_t db) { pybind11::buffer_info buf = db.request(); if (buf.ndim != 2) @@ -401,28 +447,37 @@ PYBIND11_MODULE(_core, m) daisy::idx_t d = buf.shape[1]; daisy::InMemoryDataSource data_source(static_cast(buf.ptr), n, d); self.buildIndex(&data_source); }, "Build the index from a 2D float32 numpy array") + .def("buildIndex", [](daisy::Coconut &self, const std::string &filename, daisy::idx_t dim, daisy::idx_t n) + { self.buildIndex(filename, dim, n); }, pybind11::arg("filename"), pybind11::arg("dim"), + pybind11::arg("n_database") = 0, "Build from a raw float32 vector file using bounded memory") // Streaming: insert one series (1D) or a batch (2D) into a live index. - .def("insert", [](daisy::Coconut &self, pybind11::array_t series) + .def("insert", [](daisy::Coconut &self, pybind11::array_t series) { pybind11::buffer_info buf = series.request(); if (buf.ndim != 1) throw std::runtime_error("insert expects a 1D float32 array"); + if (buf.shape[0] != static_cast(self.getDim())) + throw std::runtime_error("Coconut insert dimension mismatch"); self.insert(static_cast(buf.ptr)); }, "Insert a single series (1D float32 array) into the live index") - .def("insertBatch", [](daisy::Coconut &self, pybind11::array_t batch) + .def("insertBatch", [](daisy::Coconut &self, pybind11::array_t batch) { pybind11::buffer_info buf = batch.request(); if (buf.ndim != 2) throw std::runtime_error("insertBatch expects a 2D float32 array"); + if (buf.shape[1] != static_cast(self.getDim())) + throw std::runtime_error("Coconut insertBatch dimension mismatch"); self.insertBatch(static_cast(buf.ptr), buf.shape[0]); }, "Insert a batch of series (2D float32 array) into the live index") // kNN search returning (indices, distances). - .def("searchIndex", [](daisy::Coconut &self, pybind11::array_t query, daisy::idx_t k) + .def("searchIndex", [](daisy::Coconut &self, pybind11::array_t query, daisy::idx_t k) { pybind11::buffer_info buf = query.request(); if (buf.ndim != 2) throw std::runtime_error("Query array must be 2D"); + if (buf.shape[1] != static_cast(self.getDim())) + throw std::runtime_error("Coconut query dimension mismatch"); std::vector indices(buf.shape[0] * k); std::vector distances(buf.shape[0] * k); @@ -439,11 +494,13 @@ PYBIND11_MODULE(_core, m) pybind11::format_descriptor::format(), 2, std::vector{static_cast(buf.shape[0]), static_cast(k)}, std::vector{static_cast(sizeof(float) * k), static_cast(sizeof(float))}))); }, "kNN search: returns (indices, distances)") - .def("searchIndex", [](daisy::Coconut &self, pybind11::array_t query, daisy::SearchConfig config) + .def("searchIndex", [](daisy::Coconut &self, pybind11::array_t query, daisy::SearchConfig config) { pybind11::buffer_info query_buf = query.request(); if (query_buf.ndim != 2) throw std::runtime_error("Query array must be 2D"); + if (query_buf.shape[1] != static_cast(self.getDim())) + throw std::runtime_error("Coconut query dimension mismatch"); daisy::idx_t n_query = query_buf.shape[0]; std::vector> I; std::vector> D; @@ -451,7 +508,6 @@ PYBIND11_MODULE(_core, m) return pybind11::make_tuple(I, D); }, "Search using SearchConfig (top-k or range) and return (indices, distances)"); #endif - ////// MESSI ////// pybind11::class_(m, "Messi", "MESSI (Multi-Queue Efficient SAX Similarity Index) algorithm for time series similarity search") // Constructor diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0ccf2e3..11e332a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -841,6 +841,11 @@ if(BUILD_PYTHON) set_tests_properties(test_coconut_topk_py PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_SOURCE_DIR}" ) + add_test(NAME test_coconut_disk_py + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_coconut_disk.py + WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}) + set_tests_properties(test_coconut_disk_py PROPERTIES + ENVIRONMENT "PYTHONPATH=${CMAKE_SOURCE_DIR}") endif() endif() endif() diff --git a/tests/test_Coconut_Range.cpp b/tests/test_Coconut_Range.cpp index 5d1a508..0988604 100644 --- a/tests/test_Coconut_Range.cpp +++ b/tests/test_Coconut_Range.cpp @@ -2,10 +2,21 @@ #include "../lib/algos/Coconut.hpp" #include "../commons/test_bm_utils.hpp" #include "../commons/paramSetup.hpp" +#include +#include TEST_P(CoconutRangeParameterizedTest, AllConfigurations) { const RangeTestConfig &config = GetParam(); + // Historical Astronomy fixture contains two non-finite series. It cannot + // establish exact L2 range correctness under the finite-input contract. + // Detect invalid fixture data explicitly; do not silently filter/reindex it. + daisy::FileDataSource fixture(config.dataset_path.c_str(), config.dim, config.n_database); + std::vector record(config.dim); + while (fixture.nextRecord(record.data())) + if (std::any_of(record.begin(), record.end(), [](float x) { return !std::isfinite(x); })) + GTEST_SKIP() << "Non-finite range fixture: " << config.dataset_path + << "; finite-input rejection is tested separately"; for (int i = 0; i < 3; ++i) { daisy::Coconut search(daisy::DistanceType::L2_SQUARED); runSSTRange(&search, config); diff --git a/tests/test_Coconut_Streaming.cpp b/tests/test_Coconut_Streaming.cpp index 376f1f8..b444f8f 100644 --- a/tests/test_Coconut_Streaming.cpp +++ b/tests/test_Coconut_Streaming.cpp @@ -3,6 +3,7 @@ #include "test_utils.hpp" #include "../lib/algos/Coconut.hpp" +#include "../lib/algos/CoconutBounds.hpp" #include #include @@ -100,6 +101,102 @@ TEST(CoconutTest, StreamingKnnMatchesBruteForce) } } +TEST(CoconutTest, DiskMergesPreserveExactPhysicalNeighbors) +{ + constexpr int dim = 12, total = 123, k = 10; + std::vector data(total * dim), query(3 * dim); + for (int i = 0; i < total; ++i) + for (int d = 0; d < dim; ++d) + data[i * dim + d] = (float)((i * 17 + d * 13) % 101 - 50); + for (int q = 0; q < 3; ++q) + std::copy_n(data.data() + (q == 0 ? 0 : q == 1 ? 41 : 122) * dim, dim, query.data() + q * dim); + daisy::Coconut index(daisy::DistanceType::L2_SQUARED); + index.setPaaSegments(3); index.setBufferCapacity(7); index.setLeafSize(5); + index.setNumThreads(3); index.buildIndex(data.data(), 19, dim); + int previous = 19; + for (int seen : {19, 37, 83, 123}) + { + if (seen > previous) index.insertBatch(data.data() + previous * dim, seen - previous); + EXPECT_LE(index.getPeakBufferedRecords(), 7u); + EXPECT_LT(index.getBufferedRecords(), 7u); + EXPECT_EQ(index.getDiskRecords() + index.getBufferedRecords(), (idx_t)seen); + std::vector ids(3 * k); std::vector distances(3 * k); + index.searchIndex(query.data(), 3, k, ids.data(), distances.data()); + for (int q = 0; q < 3; ++q) + { + std::vector> expected; + for (int i = 0; i < seen; ++i) + { + float squared = 0; + for (int d = 0; d < dim; ++d) + { + float delta = data[i * dim + d] - query[q * dim + d]; squared += delta * delta; + } + expected.emplace_back(squared, i); + } + std::sort(expected.begin(), expected.end()); + for (int rank = 0; rank < k; ++rank) + { + EXPECT_EQ(distances[q * k + rank], expected[rank].first); + EXPECT_EQ(ids[q * k + rank], expected[rank].second); + } + } + previous = seen; + } + index.flush(); + EXPECT_EQ(index.getBufferedRecords(), 0u); + EXPECT_EQ(index.getDiskRecords(), (idx_t)total); + index.waitForMerges(); + EXPECT_GT(index.getMergeCount(), 0u); + EXPECT_EQ(index.getResidentSummaryBytes(), (idx_t)total * 3); + std::vector recent_ids(3 * k); std::vector recent_distances(3 * k); + index.searchIndexRecent(query.data(), 3, k, 17, recent_ids.data(), recent_distances.data()); + for (int q = 0; q < 3; ++q) + { + std::vector> expected; + for (int i = total - 17; i < total; ++i) + { + float distance = 0; + for (int d = 0; d < dim; ++d) + { float delta = data[i * dim + d] - query[q * dim + d]; distance += delta * delta; } + expected.emplace_back(distance, i); + } + std::sort(expected.begin(), expected.end()); + for (int j = 0; j < k; ++j) + { EXPECT_EQ(recent_ids[q * k + j], expected[j].second); EXPECT_EQ(recent_distances[q * k + j], expected[j].first); } + } +} + +TEST(CoconutTest, QueryBoundLookupMatchesReferenceAndPrunesSafely) +{ + std::mt19937 rng(112677); + std::uniform_real_distribution values(-4, 4); + for (int segments : {1, 3, 16, 32, 48, 96}) + for (int bits = 1; bits <= 8; ++bits) + { + int alphabet = 1 << bits; + std::vector paa(segments); + std::vector word(segments), cardinalities(segments, bits); + for (auto &v : paa) v = values(rng); + daisy::coconut_detail::QueryBounds lookup; + lookup.reset(paa.data(), segments, bits, 6); + for (int trial = 0; trial < 512; ++trial) + { + for (auto &v : word) v = rng() % alphabet; + float expected = daisy::minidist_paa_to_isax(paa.data(), word.data(), cardinalities.data(), + bits, alphabet, segments, MINVAL, MAXVAL, 6, sax_breakpoints); + EXPECT_EQ(lookup(word.data(), FLT_MAX), expected); + for (float limit : {0.0f, expected * .5f, std::nextafter(expected, 0.0f), expected, + std::nextafter(expected, FLT_MAX)}) + { + float actual = lookup(word.data(), limit); + if (expected <= limit) EXPECT_EQ(actual, expected); + else { EXPECT_GT(actual, limit); EXPECT_LE(actual, expected); } + } + } + } +} + int main(int argc, char **argv) { ::testing::InitGoogleTest(&argc, argv); diff --git a/tests/test_coconut_disk.py b/tests/test_coconut_disk.py new file mode 100644 index 0000000..3c0f482 --- /dev/null +++ b/tests/test_coconut_disk.py @@ -0,0 +1,241 @@ +"""Exact physical results after bounded disk flushes and LSM compactions.""" +import gc +import math +from pathlib import Path +import tempfile +import unittest +import numpy as np +import daisy + + +@unittest.skipUnless(hasattr(daisy, 'Coconut'), 'Coconut is disabled') +class CoconutDiskTest(unittest.TestCase): + def check_topk(self, index, data, query): + exact = np.sum((data.astype(np.float64)[None] - query.astype(np.float64)[:, None]) ** 2, axis=2) + for k in (1, 10, min(100, len(data)), min(1000, len(data))): + ids, distances = index.searchIndex(query, k) + self.assertTrue(np.all(ids < len(data))) + self.assertTrue(all(len(set(row)) == k for row in ids)) + physical = np.take_along_axis(exact, ids.astype(np.int64), axis=1) + expected = np.sort(exact, axis=1)[:, :k] + allowance = 2 * np.abs(np.spacing(expected.astype(np.float32)).astype(np.float64)) + self.assertTrue(np.all(np.abs(np.sort(physical, axis=1) - expected) <= allowance)) + np.testing.assert_allclose(distances, physical, rtol=2e-6, atol=2e-5) + + def test_many_flushes_merges_partial_buffers_and_stable_ids(self): + rng = np.random.default_rng(677) + data = rng.normal(size=(777, 96)).astype(np.float32) + queries = np.concatenate((data[[0, 91, 512, 776]], rng.normal(size=(5, 96)).astype(np.float32))) + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setBufferCapacity(23) + index.setLeafSize(7) + index.setNumThreads(3) + index.buildIndex(data[:91]) + previous = 91 + for n in (91, 199, 402, 777): + if n > previous: + added = data[previous:n].copy() + index.insertBatch(added) + added.fill(999) + index.waitForMerges() + self.assertEqual(index.getNDatabase(), n) + self.assertLess(index.getBufferedRecords(), 23) + self.assertLessEqual(index.getPeakBufferedRecords(), 23) + self.assertEqual(index.getDiskRecords() + index.getBufferedRecords(), n) + self.assertEqual(index.getDiskBytes(), index.getDiskRecords() * (96 * 4 + 16 * 2 + 8)) + self.assertLessEqual(index.getRunCount(), 1 + math.ceil(math.log2(index.getFlushCount() + 1))) + actual_bytes = sum(p.stat().st_size for p in Path(index.getIndexDirectory()).iterdir()) + self.assertEqual(actual_bytes, index.getDiskBytes()) + self.check_topk(index, data[:n], queries) + previous = n + self.assertGreater(index.getMergeCount(), 0) + index.flush() + index.waitForMerges() + self.assertEqual(index.getBufferedRecords(), 0) + self.assertEqual(index.getDiskRecords(), len(data)) + self.check_topk(index, data, queries) + + def test_file_build_odd_key_size_range_and_cleanup(self): + rng = np.random.default_rng(99) + # Integer pairs give exact float32 sums and inclusive-radius boundaries. + data = rng.integers(-3, 4, size=(197, 12)).astype(np.float32) + queries = data[[0, 49, 110, 196]].copy() + with tempfile.TemporaryDirectory() as parent: + source = Path(parent) / 'source.bin' + data[:49].tofile(source) + original = source.read_bytes() + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setIndexDirectory(parent) + index.setBufferCapacity(13) + index.setPaaSegments(3) + index.setLeafSize(5) + index.buildIndex(str(source), 12) + index.insert(data[49]) + index.insertBatch(data[50:]) + self.assertEqual(source.read_bytes(), original) + self.check_topk(index, data, queries) + for radius in (0., 16., 10000.): + config = daisy.SearchConfig(); config.type = daisy.QueryType.RANGE; config.r = radius + ids, distances = index.searchIndex(queries, config) + for q, row_i, row_d in zip(queries, ids, distances): + exact = np.sum((data.astype(np.float64) - q) ** 2, axis=1) + self.assertEqual(set(row_i), set(np.flatnonzero(exact <= radius))) + np.testing.assert_array_equal(row_d, np.sort(exact[exact <= radius])) + directory = Path(index.getIndexDirectory()) + del index; gc.collect() + self.assertFalse(directory.exists()) + self.assertTrue(source.exists()) + + def test_strided_arrays_shapes_and_nonfinite_insert(self): + data = np.arange(600, dtype=np.float32).reshape(50, 12) / 100 + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setPaaSegments(3); index.setBufferCapacity(7) + index.buildIndex(data[:20:2]) + index.insertBatch(data[21::2]) + expected = np.concatenate((data[:20:2], data[21::2])) + self.check_topk(index, expected, data[::7]) + n = index.getNDatabase() + with self.assertRaisesRegex(RuntimeError, 'dimension mismatch'): index.insertBatch(np.zeros((2, 11), np.float32)) + with self.assertRaisesRegex(RuntimeError, 'dimension mismatch'): index.insert(np.zeros(11, np.float32)) + with self.assertRaisesRegex(RuntimeError, 'dimension mismatch'): index.searchIndex(np.zeros((1, 11), np.float32), 1) + with self.assertRaisesRegex(RuntimeError, 'finite'): index.insert(np.full(12, np.nan, np.float32)) + self.assertEqual(index.getNDatabase(), n) + with self.assertRaisesRegex(RuntimeError, 'cannot change'): index.setBufferCapacity(3) + + def test_corrupt_directory_rejected_before_mapping_access(self): + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + data = np.zeros((20, 64), np.float32) + index.setBufferCapacity(7); index.buildIndex(data) + next(Path(index.getIndexDirectory()).glob('*.meta')).write_bytes(b'') + with self.assertRaisesRegex(RuntimeError, 'index leaf/run'): index.searchIndex(data[:1], 1) + + def test_paper_sims_resident_summaries_and_multilevel_tree(self): + # >128**2 leaves exercises three packed B+tree levels, duplicate keys, + # both query-key endpoints, and external-sort fan-in beyond 64 runs. + rng = np.random.default_rng(718) + data = rng.integers(-4, 5, size=(33001, 12)).astype(np.float32) + query = np.concatenate((data[[0, 1024, 33000]], np.full((1, 12), -30, np.float32), + np.full((1, 12), 30, np.float32))) + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setPaaSegments(3); index.setLeafSize(2); index.setBufferCapacity(101) + index.buildIndex(data) + self.assertEqual(index.getRunCount(), 1) + self.assertEqual(index.getResidentSummaryBytes(), len(data) * 3) + exact = np.sum((query.astype(np.float64)[:, None] - data.astype(np.float64)) ** 2, axis=2) + expected = np.argsort(exact, axis=1, kind='stable')[:, :10] + for threads in (1, 4): + index.setNumThreads(threads) + ids, distances = index.searchIndex(query, 10) + np.testing.assert_array_equal(ids, expected) + np.testing.assert_array_equal(distances, np.take_along_axis(exact, expected, axis=1)) + stats = index.getLastSearchStats() + self.assertEqual(stats['summaries_examined'], len(data) * len(query)) + self.assertEqual(stats['runs_searched'], len(query)) + self.assertGreater(stats['raw_bytes_read'], 0) + self.assertGreater(stats['bounds_seconds'], 0) + + def test_paper_async_compaction_live_snapshot_and_barrier(self): + rng = np.random.default_rng(270) + data = rng.integers(-3, 4, size=(2049, 12)).astype(np.float32) + query = data[[1, 127, 777, 2048]].copy() + for async_merges in (False, True): + with self.subTest(async_merges=async_merges): + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setPaaSegments(3); index.setBufferCapacity(17); index.setLeafSize(9) + index.setAsyncMerges(async_merges); index.setNumThreads(3) + index.buildIndex(data[:131]) + previous = 131 + for n in (198, 513, 1028, 2049): + index.insertBatch(data[previous:n]) + # No barrier: queries must be exact even if old runs are + # being merged/replaced concurrently by the worker. + exact = np.sum((query[:, None].astype(np.float64) - data[None, :n]) ** 2, axis=2) + expected = np.argsort(exact, axis=1, kind='stable')[:, :10] + ids, distances = index.searchIndex(query, 10) + np.testing.assert_array_equal(ids, expected) + np.testing.assert_array_equal(distances, np.take_along_axis(exact, expected, axis=1)) + previous = n + index.flush(); index.waitForMerges() + self.assertEqual(index.getResidentSummaryBytes(), len(data) * 3) + self.assertEqual(index.getDiskBytes(), sum(p.stat().st_size for p in Path(index.getIndexDirectory()).iterdir())) + self.assertLessEqual(index.getRunCount(), 1 + math.ceil(math.log2(index.getFlushCount() + 1))) + with self.assertRaisesRegex(RuntimeError, 'cannot change'): index.setAsyncMerges(False) + directory = Path(index.getIndexDirectory()) + del index; gc.collect() + self.assertFalse(directory.exists()) + + def test_paper_btp_recent_windows_skip_old_runs_and_preserve_ids(self): + rng = np.random.default_rng(102) + data = rng.integers(-2, 3, size=(773, 12)).astype(np.float32) + data[-1] = data[0] # newest exact match must win inside a small window + query = data[[0, 20, 772]].copy() + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setPaaSegments(3); index.setLeafSize(7); index.setBufferCapacity(16) + index.setNumThreads(4); index.buildIndex(data[:128]); index.insertBatch(data[128:]); index.waitForMerges() + for window in (1, 5, 17, 200, 773, 1000): + first = max(0, len(data) - window) + exact = np.sum((query[:, None].astype(np.float64) - data[None, first:]) ** 2, axis=2) + expected = np.argsort(exact, axis=1, kind='stable')[:, :10] + first + ids, distances = index.searchIndexRecent(query, 10, window) + count = min(10, len(data) - first) + np.testing.assert_array_equal(ids[:, :count], expected) + np.testing.assert_array_equal(distances[:, :count], np.take_along_axis(exact, expected - first, axis=1)) + self.assertTrue(np.all(distances[:, count:] == np.finfo(np.float32).max)) + if window <= 17: + self.assertGreater(index.getLastSearchStats()['runs_skipped'], 0) + with self.assertRaisesRegex(RuntimeError, 'window must be positive'): + index.searchIndexRecent(query, 1, 0) + with self.assertRaisesRegex(RuntimeError, 'dimension mismatch'): + index.searchIndexRecent(np.zeros((1, 11), np.float32), 1, 5) + + def test_batch_sims_matches_serial_and_exact_with_tail_and_windows(self): + rng = np.random.default_rng(1120) + data = rng.integers(-3, 4, size=(2049, 12)).astype(np.float32) + query = data[rng.integers(0, len(data), size=65)].copy() + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setPaaSegments(3); index.setBufferCapacity(97); index.setLeafSize(17) + index.setNumThreads(3); index.buildIndex(data[:513]); index.insertBatch(data[513:]) + for window in (len(data), 29): + first = max(0, len(data) - window) + exact = np.sum((query[:, None].astype(np.float64) - data[None, first:]) ** 2, axis=2) + for batch_size in (1, 4, 32, 64): + index.setQueryBatchSize(batch_size) + for k in (1, 10, 31): + ids, distances = index.searchIndexRecent(query, k, window) + n = min(k, len(data) - first) + expected = np.argsort(exact, axis=1, kind='stable')[:, :n] + np.testing.assert_array_equal(ids[:, :n], expected + first) + np.testing.assert_array_equal(distances[:, :n], np.take_along_axis(exact, expected, axis=1)) + self.assertTrue(np.all(distances[:, n:] == np.finfo(np.float32).max)) + stats = index.getLastSearchStats() + self.assertLessEqual(stats['distance_evaluations'] - stats['seed_evaluations'], stats['candidate_records']) + self.assertLessEqual(stats['unique_candidate_records'], stats['raw_records_read']) + self.assertLess(stats['peak_scratch_bytes'], 20 * 1024 * 1024) + for bad in (0, 65): + with self.assertRaisesRegex(RuntimeError, 'between 1 and 64'): index.setQueryBatchSize(bad) + + def test_batch_sims_shares_reads_without_replacing_sax_pruning(self): + rng = np.random.default_rng(728) + data = rng.normal(size=(4096, 12)).astype(np.float32) + query = np.repeat(np.zeros((1, 12), np.float32), 17, axis=0) + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setPaaSegments(3); index.setLeafSize(17); index.setBufferCapacity(512) + index.setNumThreads(4); index.buildIndex(data) + index.setQueryBatchSize(1); before = index.searchIndex(query, 10); serial = index.getLastSearchStats() + index.setQueryBatchSize(32); after = index.searchIndex(query, 10); batched = index.getLastSearchStats() + np.testing.assert_array_equal(before[0], after[0]); np.testing.assert_array_equal(before[1], after[1]) + self.assertLess(batched['raw_bytes_read'], serial['raw_bytes_read'] / 4) + self.assertEqual(batched['candidate_records'], serial['candidate_records']) + # An exact seed prunes the distant cluster: no unconditional raw scan. + data.fill(1); data[0].fill(0); index.buildIndex(data) + ids, distances = index.searchIndex(query, 1) + np.testing.assert_array_equal(ids, np.zeros((17, 1), np.uint64)) + np.testing.assert_array_equal(distances, np.zeros((17, 1), np.float32)) + stats = index.getLastSearchStats() + self.assertEqual(stats['candidate_records'], 0) + self.assertEqual(stats['distance_evaluations'], stats['seed_evaluations']) + self.assertLess(stats['raw_bytes_read'], len(data) * 12 * 4) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_coconut_topk.py b/tests/test_coconut_topk.py index 726f71e..1f94cca 100644 --- a/tests/test_coconut_topk.py +++ b/tests/test_coconut_topk.py @@ -122,7 +122,7 @@ def test_unreadable_leaf_raises(self): index.setNumThreads(4) directories = set(Path('/tmp').glob(f'daisy_coconut_{os.getpid()}_*')) - before self.assertEqual(len(directories), 1) - leaf = next(iter(directories)) / 'leaf.0' + leaf = next(next(iter(directories)).glob('run.*.raw')) # A truncated file must never produce distances from an uninitialized buffer. leaf.write_bytes(b'') with self.assertRaisesRegex(RuntimeError, 'index leaf'):