faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/NNDescent.cpp

@@ -10,10 +10,10 @@
 #include <faiss/impl/NNDescent.h>
 
 #include <mutex>
-#include <string>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/VisitedTable.h>
 
 namespace faiss {
 
@@ -23,7 +23,7 @@ namespace nndescent {
 
 void gen_random(std::mt19937& rng, int* addr, const int size, const int N);
 
-Nhood::Nhood(int l, int s, std::mt19937& rng, int N) {
+Nhood::Nhood(int /* l */, int s, std::mt19937& rng, int N) {
     M = s;
     nn_new.resize(s * 2);
     gen_random(rng, nn_new.data(), (int)nn_new.size(), N);
@@ -31,25 +31,46 @@ Nhood::Nhood(int l, int s, std::mt19937& rng, int N) {
 
 /// Copy operator
 Nhood& Nhood::operator=(const Nhood& other) {
-    M = other.M;
-    std::copy(
-            other.nn_new.begin(),
-            other.nn_new.end(),
-            std::back_inserter(nn_new));
-    nn_new.reserve(other.nn_new.capacity());
-    pool.reserve(other.pool.capacity());
+    if (this != &other) {
+        M = other.M;
+        nn_new = other.nn_new;
+        nn_old = other.nn_old;
+        rnn_new = other.rnn_new;
+        rnn_old = other.rnn_old;
+        pool = other.pool;
+    }
     return *this;
 }
 
 /// Copy constructor
-Nhood::Nhood(const Nhood& other) {
-    M = other.M;
-    std::copy(
-            other.nn_new.begin(),
-            other.nn_new.end(),
-            std::back_inserter(nn_new));
-    nn_new.reserve(other.nn_new.capacity());
-    pool.reserve(other.pool.capacity());
+Nhood::Nhood(const Nhood& other)
+        : pool(other.pool),
+          M(other.M),
+          nn_old(other.nn_old),
+          nn_new(other.nn_new),
+          rnn_old(other.rnn_old),
+          rnn_new(other.rnn_new) {}
+
+/// Move constructor
+Nhood::Nhood(Nhood&& other) noexcept
+        : pool(std::move(other.pool)),
+          M(other.M),
+          nn_old(std::move(other.nn_old)),
+          nn_new(std::move(other.nn_new)),
+          rnn_old(std::move(other.rnn_old)),
+          rnn_new(std::move(other.rnn_new)) {}
+
+/// Move assignment operator
+Nhood& Nhood::operator=(Nhood&& other) noexcept {
+    if (this != &other) {
+        M = other.M;
+        nn_new = std::move(other.nn_new);
+        nn_old = std::move(other.nn_old);
+        rnn_new = std::move(other.rnn_new);
+        rnn_old = std::move(other.rnn_old);
+        pool = std::move(other.pool);
+    }
+    return *this;
 }
 
 /// Insert a point into the candidate pool
@@ -58,7 +79,7 @@ void Nhood::insert(int id, float dist) {
     if (dist > pool.front().distance) {
         return;
     }
-    for (int i = 0; i < pool.size(); i++) {
+    for (size_t i = 0; i < pool.size(); i++) {
         if (id == pool[i].id) {
             return;
         }
@@ -90,6 +111,22 @@ void Nhood::join(C callback) const {
 }
 
 void gen_random(std::mt19937& rng, int* addr, const int size, const int N) {
+    FAISS_THROW_IF_NOT_FMT(
+            size > 0 && size <= N,
+            "gen_random: size (%d) must be > 0 and <= N (%d)",
+            size,
+            N);
+    if (size == N) {
+        // Special case: return all indices in random order
+        for (int i = 0; i < size; ++i) {
+            addr[i] = i;
+        }
+        for (int i = size - 1; i > 0; --i) {
+            int j = rng() % (i + 1);
+            std::swap(addr[i], addr[j]);
+        }
+        return;
+    }
     for (int i = 0; i < size; ++i) {
         addr[i] = rng() % (N - size);
     }
@@ -153,7 +190,7 @@ using namespace nndescent;
 
 constexpr int NUM_EVAL_POINTS = 100;
 
-NNDescent::NNDescent(const int d, const int K) : K(K), d(d) {
+NNDescent::NNDescent(const int d_in, const int K_in) : K(K_in), d(d_in) {
     L = K + 50;
 }
 
@@ -197,7 +234,7 @@ void NNDescent::update() {
         auto& nn = graph[n];
         std::sort(nn.pool.begin(), nn.pool.end());
 
-        if (nn.pool.size() > L) {
+        if (nn.pool.size() > static_cast<size_t>(L)) {
             nn.pool.resize(L);
         }
         nn.pool.reserve(L); // keep the pool size be L
@@ -238,7 +275,7 @@ void NNDescent::update() {
                     // the candidate pool of the other side
                     if (nn.distance > other.pool.back().distance) {
                         LockGuard guard(other.lock);
-                        if (other.rnn_new.size() < R) {
+                        if (other.rnn_new.size() < static_cast<size_t>(R)) {
                             other.rnn_new.push_back(n);
                         } else {
                             int pos = rng() % R;
@@ -254,7 +291,7 @@ void NNDescent::update() {
                     // the candidate pool of the other side
                     if (nn.distance > other.pool.back().distance) {
                         LockGuard guard(other.lock);
-                        if (other.rnn_old.size() < R) {
+                        if (other.rnn_old.size() < static_cast<size_t>(R)) {
                             other.rnn_old.push_back(n);
                         } else {
                             int pos = rng() % R;
@@ -280,7 +317,7 @@ void NNDescent::update() {
 
         nn_new.insert(nn_new.end(), rnn_new.begin(), rnn_new.end());
         nn_old.insert(nn_old.end(), rnn_old.begin(), rnn_old.end());
-        if (nn_old.size() > R * 2) {
+        if (nn_old.size() > static_cast<size_t>(R * 2)) {
             nn_old.resize(R * 2);
             nn_old.reserve(R * 2);
         }
@@ -294,7 +331,7 @@ void NNDescent::nndescent(DistanceComputer& qdis, bool verbose) {
     int num_eval_points = std::min(NUM_EVAL_POINTS, ntotal);
     std::vector<int> eval_points(num_eval_points);
     std::vector<std::vector<int>> acc_eval_set(num_eval_points);
-    std::mt19937 rng(random_seed * 6577 + omp_get_thread_num());
+    std::mt19937 rng(random_seed * 6577);
     gen_random(rng, eval_points.data(), eval_points.size(), ntotal);
     generate_eval_set(qdis, eval_points, acc_eval_set, ntotal);
     for (int it = 0; it < iter; it++) {
@@ -315,7 +352,7 @@ void NNDescent::generate_eval_set(
         std::vector<std::vector<int>>& v,
         int N) {
 #pragma omp parallel for
-    for (int i = 0; i < c.size(); i++) {
+    for (int i = 0; i < static_cast<int>(c.size()); i++) {
         std::vector<Neighbor> tmp;
         for (int j = 0; j < N; j++) {
             if (c[i] == j) {
@@ -488,7 +525,7 @@ void NNDescent::search(
             ++k;
         }
     }
-    for (size_t i = 0; i < topk; i++) {
+    for (int i = 0; i < topk; i++) {
         indices[i] = retset[i].id;
         dists[i] = retset[i].distance;
     }

data/vendor/faiss/faiss/impl/NNDescent.h

@@ -54,8 +54,8 @@ struct Neighbor {
     bool flag;
 
     Neighbor() = default;
-    Neighbor(int id, float distance, bool f)
-            : id(id), distance(distance), flag(f) {}
+    Neighbor(int id_in, float distance_in, bool f)
+            : id(id_in), distance(distance_in), flag(f) {}
 
     inline bool operator<(const Neighbor& other) const {
         return distance < other.distance;
@@ -80,6 +80,10 @@ struct Nhood {
 
     Nhood(const Nhood& other);
 
+    Nhood(Nhood&& other) noexcept;
+
+    Nhood& operator=(Nhood&& other) noexcept;
+
     void insert(int id, float dist);
 
     template <typename C>

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -8,11 +8,13 @@
 #include <faiss/impl/NSG.h>
 
 #include <algorithm>
+#include <limits>
 #include <memory>
 #include <mutex>
 #include <stack>
 
 #include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/VisitedTable.h>
 
 namespace faiss {
 
@@ -41,8 +43,8 @@ struct Neighbor {
     bool flag;
 
     Neighbor() = default;
-    Neighbor(int id, float distance, bool f)
-            : id(id), distance(distance), flag(f) {}
+    Neighbor(int id_in, float distance_in, bool f)
+            : id(id_in), distance(distance_in), flag(f) {}
 
     inline bool operator<(const Neighbor& other) const {
         return distance < other.distance;
@@ -54,7 +56,7 @@ struct Node {
     float distance;
 
     Node() = default;
-    Node(int id, float distance) : id(id), distance(distance) {}
+    Node(int id_in, float distance_in) : id(id_in), distance(distance_in) {}
 
     inline bool operator<(const Node& other) const {
         return distance < other.distance;
@@ -109,10 +111,9 @@ inline int insert_into_pool(Neighbor* addr, int K, Neighbor nn) {
 
 using namespace nsg;
 
-NSG::NSG(int R) : R(R), rng(0x0903) {
+NSG::NSG(int R_in) : R(R_in), rng(0x0903) {
     L = R + 32;
     C = R + 100;
-    srand(0x1998);
 }
 
 void NSG::search(
@@ -130,7 +131,7 @@ void NSG::search(
     search_on_graph<false>(
             *final_graph, dis, vt, enterpoint, pool_size, retset, tmp);
 
-    for (size_t i = 0; i < k; i++) {
+    for (int i = 0; i < k; i++) {
         I[i] = retset[i].id;
         D[i] = retset[i].distance;
     }
@@ -178,7 +179,7 @@ void NSG::build(
     is_built = true;
 
     if (verbose) {
-        int max = 0, min = 1e6;
+        int max = 0, min = std::numeric_limits<int>::max();
         double avg = 0;
 
         for (int i = 0; i < n; i++) {
@@ -233,7 +234,7 @@ void NSG::init_graph(Index* storage, const nsg::Graph<idx_t>& knn_graph) {
     std::unique_ptr<DistanceComputer> dis(storage_distance_computer(storage));
 
     dis->set_query(center.get());
-    VisitedTable vt(ntotal);
+    VisitedTable vt(ntotal, use_visited_hashset);
 
     // Do not collect the visited nodes
     search_on_graph<false>(knn_graph, *dis, vt, ep, L, retset, tmpset);
@@ -255,16 +256,18 @@ void NSG::search_on_graph(
     retset.resize(pool_size + 1);
     std::vector<int> init_ids(pool_size);
 
+    vt.reserve(pool_size);
+
     int num_ids = 0;
     std::vector<index_t> neighbors(graph.K);
     size_t nneigh = graph.get_neighbors(ep, neighbors.data());
-    for (int i = 0; i < init_ids.size() && i < nneigh; i++) {
+    for (size_t i = 0; i < init_ids.size() && i < nneigh; i++) {
         int id = (int)neighbors[i];
         if (id >= ntotal) {
             continue;
         }
 
-        init_ids[i] = id;
+        init_ids[num_ids] = id;
         vt.set(id);
         num_ids += 1;
     }
@@ -280,7 +283,7 @@ void NSG::search_on_graph(
         vt.set(id);
     }
 
-    for (int i = 0; i < init_ids.size(); i++) {
+    for (size_t i = 0; i < init_ids.size(); i++) {
         int id = init_ids[i];
 
         float dist = dis(id);
@@ -302,9 +305,9 @@ void NSG::search_on_graph(
             int n = retset[k].id;
 
             size_t nneigh_for_n = graph.get_neighbors(n, neighbors.data());
-            for (int m = 0; m < nneigh_for_n; m++) {
+            for (size_t m = 0; m < nneigh_for_n; m++) {
                 int id = neighbors[m];
-                if (id > ntotal || vt.get(id)) {
+                if (id >= ntotal || vt.get(id)) {
                     continue;
                 }
                 vt.set(id);
@@ -341,7 +344,7 @@ void NSG::link(
         std::vector<Node> pool;
         std::vector<Neighbor> tmp;
 
-        VisitedTable vt(ntotal);
+        VisitedTable vt(ntotal, use_visited_hashset);
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
 
@@ -396,16 +399,30 @@ void NSG::sync_prune(
 
     std::vector<Node> result;
 
+    if (pool.empty()) {
+        for (int i = 0; i < R; i++) {
+            graph.at(q, i).id = EMPTY_ID;
+        }
+        return;
+    }
+
     int start = 0;
     if (pool[start].id == q) {
         start++;
     }
+    if (start >= static_cast<int>(pool.size())) {
+        for (int i = 0; i < R; i++) {
+            graph.at(q, i).id = EMPTY_ID;
+        }
+        return;
+    }
     result.push_back(pool[start]);
 
-    while (result.size() < R && (++start) < pool.size() && start < C) {
+    while (result.size() < static_cast<size_t>(R) &&
+           (++start) < static_cast<int>(pool.size()) && start < C) {
         auto& p = pool[start];
         bool occlude = false;
-        for (int t = 0; t < result.size(); t++) {
+        for (size_t t = 0; t < result.size(); t++) {
             if (p.id == result[t].id) {
                 occlude = true;
                 break;
@@ -421,8 +438,8 @@ void NSG::sync_prune(
         }
     }
 
-    for (size_t i = 0; i < R; i++) {
-        if (i < result.size()) {
+    for (int i = 0; i < R; i++) {
+        if (static_cast<size_t>(i) < result.size()) {
             graph.at(q, i).id = result[i].id;
             graph.at(q, i).distance = result[i].distance;
         } else {
@@ -436,7 +453,7 @@ void NSG::add_reverse_links(
         std::vector<std::mutex>& locks,
         DistanceComputer& dis,
         nsg::Graph<Node>& graph) {
-    for (size_t i = 0; i < R; i++) {
+    for (int i = 0; i < R; i++) {
         if (graph.at(q, i).id == EMPTY_ID) {
             break;
         }
@@ -465,17 +482,18 @@ void NSG::add_reverse_links(
         }
 
         tmp_pool.push_back(sn);
-        if (tmp_pool.size() > R) {
+        if (tmp_pool.size() > static_cast<size_t>(R)) {
             std::vector<Node> result;
             int start = 0;
             std::sort(tmp_pool.begin(), tmp_pool.end());
             result.push_back(tmp_pool[start]);
 
-            while (result.size() < R && (++start) < tmp_pool.size()) {
+            while (result.size() < static_cast<size_t>(R) &&
+                   (++start) < static_cast<int>(tmp_pool.size())) {
                 auto& p = tmp_pool[start];
                 bool occlude = false;
 
-                for (int t = 0; t < result.size(); t++) {
+                for (size_t t = 0; t < result.size(); t++) {
                     if (p.id == result[t].id) {
                         occlude = true;
                         break;
@@ -494,7 +512,7 @@ void NSG::add_reverse_links(
 
             {
                 LockGuard guard(locks[des]);
-                for (int t = 0; t < result.size(); t++) {
+                for (size_t t = 0; t < result.size(); t++) {
                     graph.at(des, t) = result[t];
                 }
             }
@@ -513,8 +531,8 @@ void NSG::add_reverse_links(
 
 int NSG::tree_grow(Index* storage, std::vector<int>& degrees) {
     int root = enterpoint;
-    VisitedTable vt(ntotal);
-    VisitedTable vt2(ntotal);
+    VisitedTable vt(ntotal, use_visited_hashset);
+    VisitedTable vt2(ntotal, use_visited_hashset);
 
     int num_attached = 0;
     int cnt = 0;
@@ -613,7 +631,7 @@ int NSG::attach_unlinked(
 
     int node;
     bool found = false;
-    for (int i = 0; i < pool.size(); i++) {
+    for (size_t i = 0; i < pool.size(); i++) {
         node = pool[i].id;
         if (degrees[node] < R && node != id) {
             found = true;

data/vendor/faiss/faiss/impl/NSG.h

@@ -9,12 +9,12 @@
 
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <vector>
 
 #include <omp.h>
 
 #include <faiss/Index.h>
-#include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
@@ -38,6 +38,7 @@ namespace faiss {
  */
 
 struct DistanceComputer; // from AuxIndexStructures
+struct VisitedTable;
 
 namespace nsg {
 
@@ -59,18 +60,20 @@ struct Graph {
     bool own_fields; ///< the underlying data owned by itself or not
 
     // construct from a known graph
-    Graph(node_t* data, int N, int K)
-            : data(data), K(K), N(N), own_fields(false) {}
+    Graph(node_t* data_in, int N_in, int K_in)
+            : data(data_in), K(K_in), N(N_in), own_fields(false) {}
 
     // construct an empty graph
     // NOTE: the newly allocated data needs to be destroyed at destruction time
-    Graph(int N, int K) : K(K), N(N), own_fields(true) {
-        data = new node_t[N * K];
+    Graph(int N_in, int K_in) : K(K_in), N(N_in), own_fields(true) {
+        size_t total = faiss::mul_no_overflow(
+                (size_t)N_in, (size_t)K_in, "Graph allocation");
+        data = new node_t[total];
     }
 
     // copy constructor
     Graph(const Graph& g) : Graph(g.N, g.K) {
-        memcpy(data, g.data, N * K * sizeof(node_t));
+        memcpy(data, g.data, (size_t)N * (size_t)K * sizeof(node_t));
     }
 
     // release the allocated memory if needed
@@ -82,21 +85,25 @@ struct Graph {
 
     // access the j-th neighbor of node i
     inline node_t at(int i, int j) const {
-        return data[i * K + j];
+        FAISS_CHECK_RANGE_DEBUG(i, 0, N);
+        FAISS_CHECK_RANGE_DEBUG(j, 0, K);
+        return data[(size_t)i * K + j];
     }
 
     // access the j-th neighbor of node i by reference
     inline node_t& at(int i, int j) {
-        return data[i * K + j];
+        FAISS_CHECK_RANGE_DEBUG(i, 0, N);
+        FAISS_CHECK_RANGE_DEBUG(j, 0, K);
+        return data[(size_t)i * K + j];
     }
 
     // get all neighbors of node i (used during search only)
     virtual size_t get_neighbors(int i, node_t* neighbors) const {
         for (int j = 0; j < K; j++) {
-            if (data[i * K + j] < 0) {
+            if (data[(size_t)i * K + j] < 0) {
                 return j;
             }
-            neighbors[j] = data[i * K + j];
+            neighbors[j] = data[(size_t)i * K + j];
         }
         return K;
     }
@@ -122,6 +129,9 @@ struct NSG {
     // search-time parameters
     int search_L = 16; ///< length of the search path
 
+    // See impl/VisitedTable.h.
+    std::optional<bool> use_visited_hashset;
+
     int enterpoint; ///< enterpoint
 
     std::shared_ptr<nsg::Graph<int32_t>> final_graph; ///< NSG graph structure

data/vendor/faiss/faiss/impl/Panorama.cpp

@@ -10,23 +10,77 @@
 #include <algorithm>
 #include <cmath>
 #include <cstring>
-#include <vector>
+
+#include <faiss/impl/FaissAssert.h>
 
 namespace faiss {
 
+namespace {
+
+/// Helper function to compute cumulative sums by iterating backwards through
+/// levels. This is the core logic shared by compute_cumulative_sums and
+/// compute_query_cum_sums.
+template <typename OffsetFunc>
+inline void compute_cum_sums_impl(
+        const float* vector,
+        float* output,
+        size_t d,
+        size_t n_levels,
+        size_t level_width_floats,
+        OffsetFunc&& get_offset) {
+    // Iterate backwards through levels, accumulating sum as we go.
+    // This avoids computing the suffix sum for each vector, which takes
+    // extra memory.
+    float sum = 0.0f;
+
+    for (int level = n_levels - 1; level >= 0; level--) {
+        size_t start_idx = level * level_width_floats;
+        size_t end_idx = std::min(
+                (level + 1) * level_width_floats, static_cast<size_t>(d));
+
+        for (size_t j = start_idx; j < end_idx; j++) {
+            sum += vector[j] * vector[j];
+        }
+
+        output[get_offset(level)] = std::sqrt(sum);
+    }
+
+    output[get_offset(n_levels)] = 0.0f;
+}
+
+} // namespace
+
 /**************************************************************
  * Panorama structure implementation
  **************************************************************/
 
-Panorama::Panorama(size_t code_size, size_t n_levels, size_t batch_size)
-        : code_size(code_size), n_levels(n_levels), batch_size(batch_size) {
+Panorama::Panorama(
+        size_t code_size_in,
+        size_t n_levels_in,
+        size_t batch_size_in)
+        : code_size(code_size_in),
+          n_levels(n_levels_in),
+          batch_size(batch_size_in) {
     set_derived_values();
 }
 
 void Panorama::set_derived_values() {
+    FAISS_THROW_IF_NOT_MSG(n_levels > 0, "Panorama: n_levels must be > 0");
     this->d = code_size / sizeof(float);
+    FAISS_THROW_IF_NOT_MSG(n_levels <= d, "Panorama: n_levels must be <= d");
     this->level_width_floats = ((d + n_levels - 1) / n_levels);
     this->level_width = this->level_width_floats * sizeof(float);
+    size_t n_real_levels = d / level_width_floats;
+    if (d > n_real_levels * level_width_floats) {
+        n_real_levels++;
+    }
+    if (this->n_levels != n_real_levels) {
+        fprintf(stderr,
+                "WARNING truncating nlevels from %zu to %zu\n",
+                this->n_levels,
+                n_real_levels);
+        this->n_levels = n_real_levels;
+    }
 }
 
 /**
@@ -69,64 +123,34 @@ void Panorama::compute_cumulative_sums(
         float* cumsum_base,
         size_t offset,
         size_t n_entry,
-        const float* vectors) {
-    std::vector<float> suffix_sums(d + 1);
-
+        const float* vectors) const {
     for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) {
         size_t current_pos = offset + entry_idx;
         size_t batch_no = current_pos / batch_size;
         size_t pos_in_batch = current_pos % batch_size;
 
         const float* vector = vectors + entry_idx * d;
-
-        // Compute suffix sums of squared values.
-        suffix_sums[d] = 0.0f;
-        for (int j = d - 1; j >= 0; j--) {
-            float squared_val = vector[j] * vector[j];
-            suffix_sums[j] = suffix_sums[j + 1] + squared_val;
-        }
-
-        // Store cumulative sums in batch-oriented layout.
         size_t cumsum_batch_offset = batch_no * batch_size * (n_levels + 1);
 
-        for (size_t level = 0; level < n_levels; level++) {
-            size_t start_idx = level * level_width_floats;
-            size_t cumsum_offset =
-                    cumsum_batch_offset + level * batch_size + pos_in_batch;
-            if (start_idx < d) {
-                cumsum_base[cumsum_offset] = std::sqrt(suffix_sums[start_idx]);
-            } else {
-                cumsum_base[cumsum_offset] = 0.0f;
-            }
-        }
-
-        // Last level sum is always 0.
-        size_t cumsum_offset =
-                cumsum_batch_offset + n_levels * batch_size + pos_in_batch;
-        cumsum_base[cumsum_offset] = 0.0f;
+        auto get_offset = [&](size_t level) {
+            return cumsum_batch_offset + level * batch_size + pos_in_batch;
+        };
+
+        compute_cum_sums_impl(
+                vector,
+                cumsum_base,
+                d,
+                n_levels,
+                level_width_floats,
+                get_offset);
     }
 }
 
 void Panorama::compute_query_cum_sums(const float* query, float* query_cum_sums)
         const {
-    std::vector<float> suffix_sums(d + 1);
-    suffix_sums[d] = 0.0f;
-
-    for (int j = d - 1; j >= 0; j--) {
-        float squared_val = query[j] * query[j];
-        suffix_sums[j] = suffix_sums[j + 1] + squared_val;
-    }
-
-    for (size_t level = 0; level < n_levels; level++) {
-        size_t start_idx = level * level_width_floats;
-        if (start_idx < d) {
-            query_cum_sums[level] = std::sqrt(suffix_sums[start_idx]);
-        } else {
-            query_cum_sums[level] = 0.0f;
-        }
-    }
-
-    query_cum_sums[n_levels] = 0.0f;
+    auto get_offset = [](size_t level) { return level; };
+    compute_cum_sums_impl(
+            query, query_cum_sums, d, n_levels, level_width_floats, get_offset);
 }
 
 void Panorama::reconstruct(idx_t key, float* recons, const uint8_t* codes_base)
@@ -139,12 +163,12 @@ void Panorama::reconstruct(idx_t key, float* recons, const uint8_t* codes_base)
 
     for (size_t level = 0; level < n_levels; level++) {
         size_t level_offset = level * level_width * batch_size;
+        size_t actual_level_width =
+                std::min(level_width, code_size - level * level_width);
         const uint8_t* src = codes_base + batch_offset + level_offset +
-                pos_in_batch * level_width;
+                pos_in_batch * actual_level_width;
         uint8_t* dest = recons_buffer + level * level_width;
-        size_t copy_size =
-                std::min(level_width, code_size - level * level_width);
-        memcpy(dest, src, copy_size);
+        memcpy(dest, src, actual_level_width);
     }
 }