Initial work to get IVF PQ working like Vamana does (#390)

Author: jparismorgan

src/include/api/ivf_pq_index.h

@@ -185,7 +185,7 @@ class IndexVamana {
     // Create a new index. Note that we may have already loaded an existing
     // index by URI. In that case, we have updated our local state (i.e.
     // l_build_, r_max_degree_, b_backtrack_), but we should also use the
-    // timestamp from that already loaded index..
+    // timestamp from that already loaded index.
     index_ = dispatch_table.at(type)(
         training_set.num_vectors(),
         l_build_,

src/include/api/vamana_index.h

@@ -497,6 +497,8 @@ auto qv_partition(
   std::vector<id_type> top_k(num_vectors(q));
 
   auto par = stdx::execution::indexed_parallel_policy{(size_t)nthreads};
+  // For each query vector, find the closest vector in the database (i.e. in
+  // centroids).
   stdx::range_for_each(
       std::move(par), q, [&, size_db](auto&& qvec, auto&& n = 0, auto&& j = 0) {
         score_type min_score = std::numeric_limits<score_type>::max();

src/include/detail/flat/qv.h

@@ -89,7 +89,12 @@ auto partition_ivf_flat_index(
     Distance distance = Distance{}) {
   scoped_timer _{tdb_func__};
 
-  assert(::num_vectors(centroids) >= nprobe);
+  if (::num_vectors(centroids) < nprobe) {
+    throw std::invalid_argument(
+        "nprobe (" + std::to_string(nprobe) +
+        ") must be less than the number of centroids (" +
+        std::to_string(::num_vectors(centroids)) + ")");
+  }
 
   size_t num_queries = num_vectors(query);
 

src/include/detail/ivf/partition.h

@@ -155,17 +155,11 @@ class PartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       , part_index_(num_parts + 1)
       , num_vectors_{::num_vectors(training_set)}
       , num_parts_{num_parts} {
-    if (num_vectors_ == 0) {
-      throw std::invalid_argument("training_set cannot be empty.");
-    }
     if (size(part_labels) != ::num_vectors(training_set)) {
       throw std::invalid_argument(
           "The number of part_labels must equal the number of vectors in the "
           "training_set.");
     }
-    if (num_parts <= 0) {
-      throw std::invalid_argument("num_parts should be greater than 0.");
-    }
 
     auto degrees = std::vector<size_t>(num_parts);
 
@@ -180,7 +174,11 @@ class PartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       size_t bin = part_labels[i];
       size_t ibin = part_index_[bin];
 
-      ids_[ibin] = i;
+      if constexpr (feature_vector_array_with_ids<F>) {
+        ids_[ibin] = training_set.id(i);
+      } else {
+        ids_[ibin] = i;
+      }
 
       assert(ibin < this->num_cols());
       for (size_t j = 0; j < dimensions(training_set); ++j) {

src/include/detail/linalg/partitioned_matrix.h

@@ -237,6 +237,11 @@ class flat_pq_index {
 
   template <class SubDistance = cached_sub_sum_of_squares_distance>
   auto add(const ColMajorMatrix<feature_type>& feature_vectors) {
+    // These will be encoded. We will still have the same number of vectors, but
+    // now each will have num_subspaces_ dimensions instead of the original
+    // dimensions_. This is because we will chunk up the original vector into
+    // num_subspaces_ chunks and then for each chunk the vector will get an ID
+    // which maps to a set of numbers which are stored in centroids_.
     pq_vectors_ =
         ColMajorMatrix<code_type>(num_subspaces_, num_vectors(feature_vectors));
 

src/include/index/flat_pq_index.h

@@ -55,6 +55,8 @@ enum class IndexKind { FlatL2, IVFFlat, Vamana, IVFPQ };
   return index_kind_strings[static_cast<int>(kind)];
 }
 
+enum class QueryType { FiniteRAM, InfiniteRAM };
+
 /******************************************************************************
  * Static info for arrays associated with an index group
  ******************************************************************************/

src/include/index/index_defs.h

@@ -129,6 +129,25 @@ class ivf_pq_group : public base_index_group<index_type> {
     }
   }
 
+  void clear_history_impl(uint64_t timestamp) {
+    tiledb::Array::delete_fragments(
+        cached_ctx_, cluster_centroids_uri(), 0, timestamp);
+    tiledb::Array::delete_fragments(
+        cached_ctx_, flat_ivf_centroids_uri(), 0, timestamp);
+    tiledb::Array::delete_fragments(
+        cached_ctx_, pq_ivf_centroids_uri(), 0, timestamp);
+    tiledb::Array::delete_fragments(cached_ctx_, ivf_index_uri(), 0, timestamp);
+    tiledb::Array::delete_fragments(cached_ctx_, ivf_ids_uri(), 0, timestamp);
+    tiledb::Array::delete_fragments(
+        cached_ctx_, pq_ivf_vectors_uri(), 0, timestamp);
+    for (size_t i = 0; i < this->get_num_subspaces(); ++i) {
+      std::string this_table_uri =
+          distance_tables_uri() + "_" + std::to_string(i);
+      tiledb::Array::delete_fragments(
+          cached_ctx_, this_table_uri, 0, timestamp);
+    }
+  }
+
   /*****************************************************************************
    * Partitioning / repartitioning history information
    ****************************************************************************/
@@ -269,9 +288,9 @@ class ivf_pq_group : public base_index_group<index_type> {
         type_to_string_v<typename index_type::indices_type>;
 
     // Initialize IVF related metadata
-    metadata_.ingestion_timestamps_ = {0};
-    metadata_.base_sizes_ = {0};
-    metadata_.partition_history_ = {0};
+    metadata_.ingestion_timestamps_ = {};
+    metadata_.base_sizes_ = {};
+    metadata_.partition_history_ = {};
     metadata_.temp_size_ = 0;
     metadata_.dimensions_ = this->get_dimensions();