range queries for coconut

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<std::vector<idx_t>> &I,
+                              std::vector<std::vector<float>> &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<float> paa(seg);
+            std::vector<float> ts(this->dim);
+            std::vector<std::pair<float, idx_t>> 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<size_t> 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<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 <= 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<float *>(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<sax_type *>(buffer_records_[b].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 > r)
+                    continue;
+
+                float *bts = buffer_data_.data() + b * (size_t)this->dim;
+                float dist = ts_euclidean_distance_SIMD(const_cast<float *>(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())

lib/algos/Coconut.hpp

@@ -5,6 +5,7 @@
 
 #include "../isax/iSAXTypes.hpp"
 
+#include <cstdio>
 #include <string>
 #include <vector>
 #include <cfloat>
@@ -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<std::vector<idx_t>> &I,
+                         std::vector<std::vector<float>> &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();
     };
 

pybinds/setup.cpp

@@ -317,7 +317,17 @@ PYBIND11_MODULE(_core, m)
                     distances.data(), sizeof(float),
                     pybind11::format_descriptor<float>::format(), 2,
                     std::vector<pybind11::ssize_t>{static_cast<pybind11::ssize_t>(buf.shape[0]), static_cast<pybind11::ssize_t>(k)},
-                    std::vector<pybind11::ssize_t>{static_cast<pybind11::ssize_t>(sizeof(float) * k), static_cast<pybind11::ssize_t>(sizeof(float))}))); }, "kNN search: returns (indices, distances)");
+                    std::vector<pybind11::ssize_t>{static_cast<pybind11::ssize_t>(sizeof(float) * k), static_cast<pybind11::ssize_t>(sizeof(float))}))); }, "kNN search: returns (indices, distances)")
+        .def("searchIndex", [](daisy::Coconut &self, pybind11::array_t<float> 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<std::vector<daisy::idx_t>> I;
+            std::vector<std::vector<float>> D;
+            self.searchIndex(static_cast<float *>(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_<daisy::Messi>(m, "Messi", "MESSI (Multi-Queue Efficient SAX Similarity Index) algorithm for time series similarity search")

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()

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<RangeTestConfig> &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();
+}

tests/test_utils.hpp

@@ -280,4 +280,14 @@ class SofaRangeParameterizedTest : public SimilaritySearchTest,
     static void TearDownTestSuite() {}
 };
 
+class CoconutRangeParameterizedTest : public SimilaritySearchTest,
+                                      public ::testing::WithParamInterface<RangeTestConfig>
+{
+protected:
+    using SimilaritySearchTest::runSSTRange;
+
+    static void SetUpTestSuite() {}
+    static void TearDownTestSuite() {}
+};
+
 #endif