Guard IVF-Flat cosine coarse search against zero-norm centers

cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh

@@ -171,7 +171,11 @@ void search_impl(raft::resources const& handle,
       [=] __device__(const uint32_t idx, const float dist) {
         const auto query   = idx / n_lists;
         const auto cluster = idx % n_lists;
-        return dist / (q_norm_ptr[query] * index_center_norm_ptr[cluster]);
+        // A list whose k-means cluster ended up empty has a zero center. Define the cosine
+        // distance to a zero-norm vector as 1 (similarity 0), as IVF-SQ does: a NaN here is never
+        // selected, and with fewer finite lists than n_probes the probes get padded with list 0.
+        const auto denom = q_norm_ptr[query] * index_center_norm_ptr[cluster];
+        return denom > 0.0f ? dist / denom : 0.0f;
       },
       raft::make_const_mdspan(distance_buffer_dev_view));
   }

cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu

@@ -10,10 +10,13 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/init.cuh>
 #include <raft/matrix/init.cuh>
+#include <raft/random/rng.cuh>
 
+#include <algorithm>
 #include <numeric>
 #include <optional>
 #include <sstream>
+#include <vector>
 
 namespace cuvs::neighbors::ivf_flat {
 
@@ -48,6 +51,83 @@ TEST(IVFFlatSerialization, StreamRoundTrip)
   EXPECT_EQ(loaded.n_lists(), index.n_lists());
 }
 
+// k-means gives an empty cluster a zero center. With the cosine metric, the coarse search must
+// still rank such lists instead of producing NaN distances: when fewer than n_probes lists have a
+// finite distance, the NaN columns are never selected and the remaining probe slots are padded
+// with list 0, which then is scanned several times and returns duplicated neighbors.
+TEST(AnnIVFFlatTest, CosineZeroNormCenters)
+{
+  raft::resources handle;
+  auto stream = raft::resource::get_cuda_stream(handle);
+
+  constexpr int64_t n_rows        = 10000;
+  constexpr int64_t n_queries     = 500;
+  constexpr int64_t dim           = 8;
+  constexpr uint32_t n_lists      = 1024;
+  constexpr uint32_t n_zero_lists = 1000;
+  constexpr uint32_t n_probes     = 40;
+  constexpr uint32_t k            = 16;
+
+  auto data    = raft::make_device_matrix<float, int64_t>(handle, n_rows, dim);
+  auto queries = raft::make_device_matrix<float, int64_t>(handle, n_queries, dim);
+  raft::random::RngState rng(1234ULL);
+  raft::random::uniform(handle, rng, data.data_handle(), data.size(), 0.1f, 2.0f);
+  raft::random::uniform(handle, rng, queries.data_handle(), queries.size(), 0.1f, 2.0f);
+
+  index_params params;
+  params.n_lists           = n_lists;
+  params.metric            = cuvs::distance::DistanceType::CosineExpanded;
+  params.add_data_on_build = false;
+  auto idx                 = build(handle, params, raft::make_const_mdspan(data.view()));
+
+  // Zero the centers of lists 1..n_zero_lists, as if k-means had left these clusters empty. Fewer
+  // lists than n_probes keep a non-zero center, and list 0 is one of them.
+  std::vector<float> centers_h(idx.centers().size());
+  raft::update_host(centers_h.data(), idx.centers().data_handle(), centers_h.size(), stream);
+  raft::resource::sync_stream(handle);
+  std::fill(centers_h.begin() + dim, centers_h.begin() + (n_zero_lists + 1) * dim, 0.0f);
+  raft::update_device(idx.centers().data_handle(), centers_h.data(), centers_h.size(), stream);
+  extend(handle, raft::make_const_mdspan(data.view()), std::nullopt, &idx);
+
+  // All data rows are in the lists with a non-zero center, and all of those lists are probed.
+  search_params sparams;
+  sparams.n_probes = n_probes;
+  auto neighbors   = raft::make_device_matrix<int64_t, int64_t>(handle, n_queries, k);
+  auto distances   = raft::make_device_matrix<float, int64_t>(handle, n_queries, k);
+  search(handle,
+         sparams,
+         idx,
+         raft::make_const_mdspan(queries.view()),
+         neighbors.view(),
+         distances.view());
+
+  auto naive_neighbors = raft::make_device_matrix<int64_t, int64_t>(handle, n_queries, k);
+  auto naive_distances = raft::make_device_matrix<float, int64_t>(handle, n_queries, k);
+  cuvs::neighbors::naive_knn<float, float, int64_t>(handle,
+                                                    naive_distances.data_handle(),
+                                                    naive_neighbors.data_handle(),
+                                                    queries.data_handle(),
+                                                    data.data_handle(),
+                                                    n_queries,
+                                                    n_rows,
+                                                    dim,
+                                                    k,
+                                                    params.metric);
+
+  std::vector<int64_t> neighbors_h(n_queries * k), naive_neighbors_h(n_queries * k);
+  std::vector<float> distances_h(n_queries * k), naive_distances_h(n_queries * k);
+  raft::update_host(neighbors_h.data(), neighbors.data_handle(), neighbors_h.size(), stream);
+  raft::update_host(distances_h.data(), distances.data_handle(), distances_h.size(), stream);
+  raft::update_host(
+    naive_neighbors_h.data(), naive_neighbors.data_handle(), naive_neighbors_h.size(), stream);
+  raft::update_host(
+    naive_distances_h.data(), naive_distances.data_handle(), naive_distances_h.size(), stream);
+  raft::resource::sync_stream(handle);
+
+  EXPECT_TRUE(eval_neighbours(
+    naive_neighbors_h, neighbors_h, naive_distances_h, distances_h, n_queries, k, 0.001, 0.99));
+}
+
 TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity)
 {
   raft::resources handle;