diff --git a/.gitignore b/.gitignore index 5c02a6e..517129b 100644 --- a/.gitignore +++ b/.gitignore @@ -128,3 +128,4 @@ venv.bak/ !.vscode/extensions.json daisy_exact_search.egg-info/ +/scripts/experiments/ diff --git a/CONTINUITY.md b/CONTINUITY.md index 71fc81b..f422215 100644 --- a/CONTINUITY.md +++ b/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: diff --git a/README.md b/README.md index 29e1351..62670b5 100644 --- a/README.md +++ b/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 diff --git a/lib/algos/Coconut.cpp b/lib/algos/Coconut.cpp index 811090f..5ff6967 100644 --- a/lib/algos/Coconut.cpp +++ b/lib/algos/Coconut.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -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 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 buffer; + if (!raw) { - float *rec = new float[this->dim]; + buffer.resize((size_t)this->n_database * this->dim); + std::vector 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(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,60 +164,86 @@ 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 &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(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 &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 paa(seg); - std::vector 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 io_failed{false}; + const 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++) { + 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; - 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(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 - - 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 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); + (int)this->dim, kth); if (bsf.size() < (size_t)k) bsf.push({dist, records_[p].series_id}); else if (dist < bsf.top().first) @@ -214,7 +251,45 @@ namespace daisy 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(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; + } + 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 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 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_); - // 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 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(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() diff --git a/lib/algos/Coconut.hpp b/lib/algos/Coconut.hpp index d8bf8a9..ecb77ef 100644 --- a/lib/algos/Coconut.hpp +++ b/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 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 &page) const; + void copySeriesFromLeaf(const std::vector &page, size_t slot, float *out) const; void cleanupLeafFiles(); }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6167f04..3d80fe5 100644 --- a/tests/CMakeLists.txt +++ b/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() diff --git a/tests/test_coconut_topk.py b/tests/test_coconut_topk.py new file mode 100644 index 0000000..726f71e --- /dev/null +++ b/tests/test_coconut_topk.py @@ -0,0 +1,138 @@ +"""Coconut exact top-k across query workers, partial leaves and streaming inserts.""" +import os +from pathlib import Path +import unittest + +import numpy as np +import faiss +import daisy + + +@unittest.skipUnless(hasattr(daisy, 'Coconut'), 'Coconut is disabled') +class CoconutTopKTest(unittest.TestCase): + def test_seed_boundaries_ties_and_partial_leaves(self): + # Keys at either end of the directory; k exceeds the leaf size; ties and + # best matches arriving in the streaming buffer must retain distinct IDs. + data = np.repeat(np.arange(-20, 21, dtype=np.float32)[:, None], 64, axis=1) + query = np.repeat(np.array([-100, -20, 0, 20, 100], dtype=np.float32)[:, None], 64, axis=1) + inserted = np.concatenate((query, query)) + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setLeafSize(7) + index.setNumThreads(3) + index.buildIndex(data) + for add in (False, True): + if add: + index.insertBatch(inserted) + data = np.concatenate((data, inserted)) + reference = faiss.IndexFlatL2(64) + reference.add(data) + for k in (1, 10, len(data), len(data) + 3): + expected, _ = reference.search(query, k) + ids, distances = index.searchIndex(query, k) + count = min(k, len(data)) + np.testing.assert_allclose(distances, expected, rtol=2e-4, atol=2e-5) + self.assertTrue(all(len(set(row[:count])) == count for row in ids)) + actual = np.sum((data[ids[:, :count].astype(np.int64)] - query[:, None]) ** 2, axis=2) + np.testing.assert_allclose(distances[:, :count], actual, rtol=2e-4, atol=2e-5) + config = daisy.SearchConfig() + config.type = daisy.QueryType.RANGE + config.r = 64.0 + ids, distances = index.searchIndex(query, config) + for q, result_ids, result_dist in zip(query, ids, distances): + exact = np.sum((data - q) ** 2, axis=1) + self.assertEqual(set(result_ids), set(np.flatnonzero(exact <= config.r))) + np.testing.assert_allclose(result_dist, np.sort(exact[exact <= config.r])) + + def test_empty_initial_index_and_invalid_configuration(self): + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.buildIndex(np.empty((0, 64), dtype=np.float32)) + data = np.ones((5, 64), dtype=np.float32) + index.insertBatch(data) + ids, distances = index.searchIndex(data[:1], 5) + self.assertEqual(len(set(ids[0])), 5) + np.testing.assert_array_equal(distances, np.zeros((1, 5), dtype=np.float32)) + with self.assertRaisesRegex(RuntimeError, 'k > 0'): + index.searchIndex(data[:1], 0) + bad = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + bad.setLeafSize(0) + with self.assertRaisesRegex(RuntimeError, 'positive leaf/PAA sizes'): + bad.buildIndex(data) + dtw = daisy.Coconut(daisy.DistanceType.DTW) + dtw.buildIndex(data) + with self.assertRaisesRegex(RuntimeError, 'only supports L2_SQUARED'): + dtw.searchIndex(data[:1], 1) + + def test_build_does_not_retain_borrowed_input(self): + rng = np.random.default_rng(891) + data = rng.normal(size=(1500, 96)).astype(np.float32) + query = data[[7, 503, 1499]].copy() + reference = faiss.IndexFlatL2(96) + reference.add(data) + expected, _ = reference.search(query, 10) + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setNumThreads(3) + index.buildIndex(data) + # A successful build owns the raw series in leaves, even when it borrows + # the input while writing them. Later changes to the array cannot affect it. + data.fill(123.0) + del data + _, distances = index.searchIndex(query, 10) + np.testing.assert_allclose(distances, expected, rtol=2e-4, atol=2e-5) + + def test_parallel_streaming_matches_faiss(self): + rng = np.random.default_rng(194) + data = rng.normal(size=(1600, 64)).astype(np.float32) + # Queries include exact matches in both the built leaves and insert buffer. + queries = np.concatenate((data[[0, 701, 900, 1499]], rng.normal(size=(12, 64)).astype(np.float32))) + faiss.omp_set_num_threads(1) + for leaf_size in (127, 1000): + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + index.setLeafSize(leaf_size) + index.buildIndex(data[:900]) + for stage in (900, 1200, 1600): + if stage == 1200: + index.insertBatch(data[900:1200]) + elif stage == 1600: + index.insert(data[1200]) + index.insertBatch(data[1201:1600]) + reference = faiss.IndexFlatL2(64) + reference.add(data[:stage]) + for k in (1, 10, min(1000, stage)): + expected, _ = reference.search(queries, k) + for workers in (1, 4): + with self.subTest(leaf_size=leaf_size, stage=stage, k=k, workers=workers): + index.setNumThreads(workers) + ids, distances = index.searchIndex(queries, k) + np.testing.assert_allclose(distances, expected, rtol=2e-4, atol=2e-5) + self.assertTrue(all(len(set(row)) == k for row in ids)) + actual = np.sum((data[ids.astype(np.int64)] - queries[:, None, :]) ** 2, axis=2) + np.testing.assert_allclose(distances, actual, rtol=2e-4, atol=2e-5) + config = daisy.SearchConfig() + config.type = daisy.QueryType.TOP_K + config.k = 10 + _, distances = index.searchIndex(queries, config) + expected, _ = reference.search(queries, 10) + np.testing.assert_allclose(distances, expected, rtol=2e-4, atol=2e-5) + + def test_unreadable_leaf_raises(self): + before = set(Path('/tmp').glob(f'daisy_coconut_{os.getpid()}_*')) + index = daisy.Coconut(daisy.DistanceType.L2_SQUARED) + data = np.zeros((20, 64), dtype=np.float32) + index.buildIndex(data) + 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' + # A truncated file must never produce distances from an uninitialized buffer. + leaf.write_bytes(b'') + with self.assertRaisesRegex(RuntimeError, 'index leaf'): + index.searchIndex(data[:4], 1) + config = daisy.SearchConfig() + config.type = daisy.QueryType.RANGE + config.r = 1.0 + with self.assertRaisesRegex(RuntimeError, 'index leaf'): + index.searchIndex(data[:4], config) + + +if __name__ == '__main__': + unittest.main()