Improve Coconut top-k seeding and leaf reads

.gitignore

@@ -128,3 +128,4 @@ venv.bak/
 !.vscode/extensions.json
 
 daisy_exact_search.egg-info/
+/scripts/experiments/

CONTINUITY.md

@@ -186,6 +186,22 @@ The repository still relies heavily on pattern matching against existing impleme
 - If you add a build flag, document it immediately.
 - If you notice a current inconsistency, record it in docs or fix it in the same change if safe.
 
+## Coconut Query And Storage Behavior
+
+Coconut top-k uses ordered query-key seeding followed by exact SAX-bound
+refinement, with independent queries distributed across the requested workers.
+Seeded records are excluded from the later scan to keep returned IDs distinct.
+Top-k and range load candidate leaves into private byte buffers, handle partial
+last leaves, and report missing/truncated reads after the parallel region.
+Contiguous inputs are borrowed only during the initial build; completed leaves
+own the raw vectors. Top-k and range support squared L2 only.
+`tests/test_coconut_topk.py` covers query workers, streaming, ties, seed-directory
+endpoints, oversized k, input lifetime, and failed reads; CMake registers it
+only with `BUILD_PYTHON` and `BUILD_COCONUT` enabled.
+The resident sorted directory replaces upstream Berkeley DB lookup, and the
+insert buffer is still unmerged: this is not a complete Coconut-LSM port.
+See `docs/coconut-upstream-audit.md` for the pinned upstream comparison.
+
 ## When This File Should Be Updated
 
 Update `CONTINUITY.md` when:

README.md

@@ -85,6 +85,16 @@ 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.
+See [the upstream comparison](docs/coconut-upstream-audit.md) for details.
+
 ### Range (distance-r) queries
 
 Every algorithm except Coconut's streaming-only paths answers range queries through

lib/algos/Coconut.cpp

@@ -8,6 +8,7 @@
 #include <cmath>
 #include <cstdio>
 #include <cstring>
+#include <memory>
 #include <queue>
 #include <string>
 #include <sys/stat.h>
@@ -60,22 +61,32 @@ namespace daisy
 
         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);
 
-        // Read the whole database into a temporary buffer (freed once leaves are on disk).
-        std::vector<float> buffer((size_t)this->n_database * this->dim);
+        // 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<float> buffer;
+        if (!raw)
         {
-            float *rec = new float[this->dim];
+            buffer.resize((size_t)this->n_database * this->dim);
+            std::vector<float> rec(this->dim);
             idx_t idx = 0;
-            while (data_source->nextRecord(rec) && idx < this->n_database)
+            while (idx < this->n_database && data_source->nextRecord(rec.data()))
             {
-                std::copy(rec, rec + this->dim, buffer.data() + (size_t)idx * this->dim);
+                std::copy(rec.begin(), rec.end(), buffer.data() + (size_t)idx * this->dim);
                 idx++;
             }
-            delete[] rec;
+            if (idx != this->n_database)
+                throw std::runtime_error("Coconut datasource ended before all records were read");
+            raw = buffer.data();
         }
 
         // 1. Compute SAX + sortable SAX for every series.
@@ -86,7 +97,7 @@ namespace daisy
             r.series_id = i;
             r.sax.resize(seg);
             r.key.resize(seg);
-            sax_from_ts(buffer.data() + (size_t)i * this->dim, r.sax.data(),
+            sax_from_ts(const_cast<float *>(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);
@@ -116,7 +127,7 @@ namespace daisy
             if (lf)
             {
                 fwrite(records_[p].key.data(), sizeof(sax_type), seg, lf);
-                fwrite(buffer.data() + (size_t)records_[p].series_id * this->dim,
+                fwrite(raw + (size_t)records_[p].series_id * this->dim,
                        sizeof(float), this->dim, lf);
             }
         }
@@ -153,68 +164,132 @@ namespace daisy
             insert(data + (size_t)i * this->dim);
     }
 
-    void Coconut::readSeriesAt(FILE *lf, int slot, float *out) const
+    bool Coconut::readLeaf(size_t leaf_no, std::vector<unsigned char> &page) const
     {
-        const long record_bytes = (long)this->paa_segments * (long)sizeof(sax_type) +
-                                   (long)this->dim * (long)sizeof(float);
-        // seek past this slot's inv-SAX to its raw TS
-        fseek(lf, (long)slot * record_bytes + (long)this->paa_segments * (long)sizeof(sax_type), SEEK_SET);
-        size_t got = fread(out, sizeof(float), this->dim, lf);
-        (void)got;
+        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, decltype(&fclose)> file(fopen(leafPath((int)leaf_no).c_str(), "rb"), &fclose);
+        return file && fread(page.data(), 1, page.size(), file.get()) == page.size();
     }
 
-    void Coconut::readSeriesFromLeaf(int leaf_no, int slot, float *out) const
+    void Coconut::copySeriesFromLeaf(const std::vector<unsigned char> &page, size_t slot, float *out) const
     {
-        FILE *lf = fopen(leafPath(leaf_no).c_str(), "rb");
-        if (!lf)
-            return;
-        readSeriesAt(lf, slot, out);
-        fclose(lf);
+        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);
     }
 
     void Coconut::searchIndex(const float *query, const idx_t n_query, const idx_t k, idx_t *I, float *D)
     {
         const int seg = this->paa_segments;
-        std::vector<float> paa(seg);
-        std::vector<float> ts(this->dim);
-
+        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<bool> io_failed{false};
+        const int workers = (int)std::max<idx_t>(1, std::min<idx_t>(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++)
         {
+            std::vector<float> paa(seg);
+            std::vector<float> ts(this->dim);
+            std::vector<unsigned char> 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<std::pair<float, idx_t>> bsf;
 
-            for (size_t p = 0; p < records_.size(); p++)
-            {
-                float kth = (bsf.size() >= (size_t)k) ? bsf.top().first : FLT_MAX;
-
-                float mindist = minidist_paa_to_isax(
-                    paa.data(),
-                    const_cast<sax_type *>(records_[p].sax.data()),
-                    const_cast<sax_type *>(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
-
-                int leaf_no = (int)(p / (size_t)this->leaf_capacity);
-                int slot = (int)(p % (size_t)this->leaf_capacity);
-                readSeriesFromLeaf(leaf_no, slot, ts.data());
-
+            // 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<sax_type> 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<sax_type> &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<float *>(qts), ts.data(),
-                                                        (int)this->dim, kth);
+                                                       (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;
+            };
+            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++)
+            {
+                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++)
+                {
+                    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<sax_type *>(records_[p].sax.data()),
+                        const_cast<sax_type *>(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;
+                }
+                if (io_failed.load())
+                    break;
             }
+            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++)
@@ -255,6 +330,8 @@ namespace daisy
                 I[q * k + j] = 0;
             }
         }
+        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,
@@ -277,6 +354,7 @@ namespace daisy
 
         I.assign(n_query, {});
         D.assign(n_query, {});
+        std::atomic<bool> io_failed{false};
 
 #pragma omp parallel for num_threads(num_threads)
         for (idx_t q = 0; q < n_query; q++)
@@ -285,12 +363,11 @@ namespace daisy
             std::vector<float> paa(seg);
             std::vector<float> ts(this->dim);
             std::vector<std::pair<float, idx_t>> hits;
+            std::vector<unsigned char> page;
 
             paa_from_ts(qts, paa.data(), seg, this->ts_values_per_segment_);
 
-            // Leaf at a time: a wide radius leaves many candidates per leaf, and opening the
-            // leaf file once for all of them keeps the scan from degenerating into one
-            // fopen/fclose per record the way a k-NN query's tight BSF never exposes.
+            // 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<size_t> candidates;
@@ -318,18 +395,21 @@ namespace daisy
                 if (candidates.empty())
                     continue;
 
-                FILE *lf = fopen(leafPath((int)leaf_no).c_str(), "rb");
-                if (!lf)
-                    continue;
+                if (!readLeaf(leaf_no, page))
+                {
+                    io_failed.store(true);
+                    break;
+                }
                 for (size_t p : candidates)
                 {
-                    readSeriesAt(lf, (int)(p - begin), ts.data());
+                    copySeriesFromLeaf(page, p - begin, ts.data());
                     float dist = ts_euclidean_distance_SIMD(const_cast<float *>(qts), ts.data(),
                                                             (int)this->dim, abandon_bound);
                     if (dist <= r)
                         hits.emplace_back(dist, records_[p].series_id);
                 }
-                fclose(lf);
+                if (io_failed.load())
+                    break;
             }
 
             // Scan the in-memory streaming buffer (raw TS already in memory).
@@ -362,6 +442,8 @@ namespace daisy
                 I[q][j] = hits[j].second;
             }
         }
+        if (io_failed.load())
+            throw std::runtime_error("Coconut failed to open or read an index leaf");
     }
 
     void Coconut::cleanupLeafFiles()

lib/algos/Coconut.hpp

@@ -16,8 +16,11 @@ namespace daisy
     // 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. Follow-ups: paged B-tree,
-    // out-of-core sort, LSM merge of the insert buffer, equi-depth breakpoints.
+    // 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.
     class Coconut : public SimilaritySearchAlgorithm
     {
     public:
@@ -78,8 +81,8 @@ namespace daisy
         std::vector<sax_type> max_cardinalities_;  // all = sax_cardinality
 
         std::string leafPath(int leaf_no) const;
-        void readSeriesFromLeaf(int leaf_no, int slot, float *out) const;
-        void readSeriesAt(FILE *lf, int slot, float *out) const;  // slot of an already-open leaf
+        bool readLeaf(size_t leaf_no, std::vector<unsigned char> &page) const;
+        void copySeriesFromLeaf(const std::vector<unsigned char> &page, size_t slot, float *out) const;
         void cleanupLeafFiles();
     };
 

tests/CMakeLists.txt

@@ -832,6 +832,16 @@ if(BUILD_PYTHON)
         set_tests_properties(test_streaming_py PROPERTIES
             ENVIRONMENT "PYTHONPATH=${CMAKE_SOURCE_DIR}"
         )
+        if(BUILD_COCONUT)
+            add_test(
+                NAME test_coconut_topk_py
+                COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_coconut_topk.py
+                WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
+            )
+            set_tests_properties(test_coconut_topk_py PROPERTIES
+                ENVIRONMENT "PYTHONPATH=${CMAKE_SOURCE_DIR}"
+            )
+        endif()
     endif()
 endif()