diff --git a/lib/algos/Coconut.cpp b/lib/algos/Coconut.cpp index 44a57d0..811090f 100644 --- a/lib/algos/Coconut.cpp +++ b/lib/algos/Coconut.cpp @@ -153,17 +153,22 @@ namespace daisy insert(data + (size_t)i * this->dim); } - void Coconut::readSeriesFromLeaf(int leaf_no, int slot, float *out) const + void Coconut::readSeriesAt(FILE *lf, int slot, float *out) const { - FILE *lf = fopen(leafPath(leaf_no).c_str(), "rb"); - if (!lf) - return; 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; + } + + void Coconut::readSeriesFromLeaf(int leaf_no, int slot, float *out) const + { + FILE *lf = fopen(leafPath(leaf_no).c_str(), "rb"); + if (!lf) + return; + readSeriesAt(lf, slot, out); fclose(lf); } @@ -252,6 +257,113 @@ namespace daisy } } + void Coconut::searchIndex(const float *query, idx_t n_query, const SearchConfig &config, + std::vector> &I, + std::vector> &D) + { + if (config.type == QueryType::TOP_K) + { + 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, {}); + +#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; + + 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. + 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++) + { + 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; + + FILE *lf = fopen(leafPath((int)leaf_no).c_str(), "rb"); + if (!lf) + continue; + for (size_t p : candidates) + { + readSeriesAt(lf, (int)(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); + } + + // 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); + } + + 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++) + { + D[q][j] = hits[j].first; + I[q][j] = hits[j].second; + } + } + } + void Coconut::cleanupLeafFiles() { if (index_dir_.empty()) diff --git a/lib/algos/Coconut.hpp b/lib/algos/Coconut.hpp index 7e23304..d8bf8a9 100644 --- a/lib/algos/Coconut.hpp +++ b/lib/algos/Coconut.hpp @@ -5,6 +5,7 @@ #include "../isax/iSAXTypes.hpp" +#include #include #include #include @@ -15,8 +16,8 @@ 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 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. Follow-ups: paged B-tree, + // out-of-core sort, LSM merge of the insert buffer, equi-depth breakpoints. class Coconut : public SimilaritySearchAlgorithm { public: @@ -41,6 +42,11 @@ namespace daisy 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 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. void insert(const float *series) override; @@ -73,6 +79,7 @@ namespace daisy 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 void cleanupLeafFiles(); }; diff --git a/pybinds/setup.cpp b/pybinds/setup.cpp index c7ad9b6..480795f 100644 --- a/pybinds/setup.cpp +++ b/pybinds/setup.cpp @@ -317,7 +317,17 @@ PYBIND11_MODULE(_core, m) distances.data(), sizeof(float), 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)"); + 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) + { + pybind11::buffer_info query_buf = query.request(); + if (query_buf.ndim != 2) + throw std::runtime_error("Query array must be 2D"); + daisy::idx_t n_query = query_buf.shape[0]; + std::vector> I; + std::vector> D; + self.searchIndex(static_cast(query_buf.ptr), n_query, config, I, D); + return pybind11::make_tuple(I, D); }, "Search using SearchConfig (top-k or range) and return (indices, distances)"); ////// MESSI ////// pybind11::class_(m, "Messi", "MESSI (Multi-Queue Efficient SAX Similarity Index) algorithm for time series similarity search") diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 4587d98..9136ae7 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1019,6 +1019,13 @@ if(BUILD_COCONUT) target_include_directories(test_Coconut_Streaming PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../lib ${CMAKE_CURRENT_SOURCE_DIR}/../commons) gtest_discover_tests(test_Coconut_Streaming WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}) + + # ////// COCONUT Range ////// + add_executable(test_Coconut_Range test_Coconut_Range.cpp test_utils.cpp) + target_link_libraries(test_Coconut_Range PRIVATE GTest::gtest_main dino_lib commons_lib stdc++fs) + target_include_directories(test_Coconut_Range PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../lib ${CMAKE_CURRENT_SOURCE_DIR}/../commons) + gtest_discover_tests(test_Coconut_Range WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}) elseif(DEBUG_MSG) message(STATUS "BUILD_COCONUT is OFF. Skipping Coconut tests.") endif() diff --git a/tests/test_Coconut_Range.cpp b/tests/test_Coconut_Range.cpp new file mode 100644 index 0000000..5d1a508 --- /dev/null +++ b/tests/test_Coconut_Range.cpp @@ -0,0 +1,29 @@ +#include "test_utils.hpp" +#include "../lib/algos/Coconut.hpp" +#include "../commons/test_bm_utils.hpp" +#include "../commons/paramSetup.hpp" + +TEST_P(CoconutRangeParameterizedTest, AllConfigurations) +{ + const RangeTestConfig &config = GetParam(); + for (int i = 0; i < 3; ++i) { + daisy::Coconut search(daisy::DistanceType::L2_SQUARED); + runSSTRange(&search, config); + } +} + +INSTANTIATE_TEST_SUITE_P( + CoconutRangeTests, + CoconutRangeParameterizedTest, + ::testing::ValuesIn(range_test_configs), + [](const ::testing::TestParamInfo &info) { + return info.param.name + "_r" + std::to_string((int)info.param.r_value) + + "_thread" + std::to_string(info.param.thread_count) + + "_idx" + std::to_string(info.index); + }); + +int main(int argc, char **argv) +{ + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tests/test_utils.hpp b/tests/test_utils.hpp index 6c393b8..0e1e327 100644 --- a/tests/test_utils.hpp +++ b/tests/test_utils.hpp @@ -280,4 +280,14 @@ class SofaRangeParameterizedTest : public SimilaritySearchTest, static void TearDownTestSuite() {} }; +class CoconutRangeParameterizedTest : public SimilaritySearchTest, + public ::testing::WithParamInterface +{ +protected: + using SimilaritySearchTest::runSSTRange; + + static void SetUpTestSuite() {} + static void TearDownTestSuite() {} +}; + #endif