Fix segfault in finite ram queries

Author: Nikos Papailiou

apis/python/test/test_ingestion.py

@@ -72,7 +72,7 @@ def test_ivf_flat_ingestion_u8(tmp_path):
     )
     assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
 
-    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype)
+    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
     result = np.transpose(
         index_ram.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
@@ -114,7 +114,7 @@ def test_ivf_flat_ingestion_f32(tmp_path):
     )
     assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
 
-    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype)
+    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
     result = np.transpose(
         index_ram.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
@@ -159,7 +159,7 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
     )
     assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
 
-    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype)
+    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
     result = np.transpose(
         index_ram.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )

src/include/detail/ivf/qv.h

@@ -609,16 +609,16 @@ auto nuv_query_heap_finite_ram(
     _i.start();
 
     size_t parts_per_thread =
-        (size(active_partitions) + nthreads - 1) / nthreads;
+        (shuffled_db.num_col_parts() + nthreads - 1) / nthreads;
 
     std::vector<std::future<void>> futs;
     futs.reserve(nthreads);
 
     for (size_t n = 0; n < nthreads; ++n) {
       auto first_part =
-          std::min<size_t>(n * parts_per_thread, size(active_partitions));
+          std::min<size_t>(n * parts_per_thread, shuffled_db.num_col_parts());
       auto last_part =
-          std::min<size_t>((n + 1) * parts_per_thread, size(active_partitions));
+          std::min<size_t>((n + 1) * parts_per_thread, shuffled_db.num_col_parts());
 
       if (first_part != last_part) {
         futs.emplace_back(std::async(
@@ -794,16 +794,16 @@ auto qv_query_heap_finite_ram(
 
     // size_t block_size = (size(active_partitions) + nthreads - 1) / nthreads;
     size_t parts_per_thread =
-        (size(active_partitions) + nthreads - 1) / nthreads;
+        (shuffled_db.num_col_parts() + nthreads - 1) / nthreads;
 
     std::vector<std::future<void>> futs;
     futs.reserve(nthreads);
 
     for (size_t n = 0; n < nthreads; ++n) {
       auto first_part =
-          std::min<size_t>(n * parts_per_thread, size(active_partitions));
+          std::min<size_t>(n * parts_per_thread, shuffled_db.num_col_parts());
       auto last_part =
-          std::min<size_t>((n + 1) * parts_per_thread, size(active_partitions));
+          std::min<size_t>((n + 1) * parts_per_thread, shuffled_db.num_col_parts());
 
       if (first_part != last_part) {
         futs.emplace_back(