Cleanup IVF PQ test and add better logging on failure (#406)

Author: jparismorgan

src/include/index/ivf_pq_index.h

@@ -1543,7 +1543,7 @@ class ivf_pq_index {
     if (::num_vectors(lhs) != ::num_vectors(rhs) ||
         ::dimensions(lhs) != ::dimensions(rhs)) {
       std::cout << "num_vectors(lhs) != num_vectors(rhs) || dimensions(lhs) != "
-                   "dimensions(rhs)n"
+                   "dimensions(rhs)"
                 << std::endl;
       std::cout << "num_vectors(lhs): " << ::num_vectors(lhs)
                 << " num_vectors(rhs): " << ::num_vectors(rhs) << std::endl;
@@ -1578,13 +1578,32 @@ class ivf_pq_index {
    * @return
    */
   template <feature_vector L, feature_vector R>
-  auto compare_feature_vectors(const L& lhs, const R& rhs) const {
+  auto compare_feature_vectors(
+      const L& lhs, const R& rhs, const std::string& msg = "") const {
     if (::dimensions(lhs) != ::dimensions(rhs)) {
-      std::cout << "dimensions(lhs) != dimensions(rhs) (" << ::dimensions(lhs)
+      std::cout << "[ivf_pq_index@compare_feature_vectors] " << msg
+                << " dimensions(lhs) != dimensions(rhs) (" << ::dimensions(lhs)
                 << " != " << ::dimensions(rhs) << ")" << std::endl;
       return false;
     }
-    return std::equal(begin(lhs), end(lhs), begin(rhs));
+    auto equal = std::equal(begin(lhs), end(lhs), begin(rhs));
+    if (!equal) {
+      std::cout << "[ivf_pq_index@compare_feature_vectors] " << msg
+                << " failed the equality check." << std::endl;
+      auto printed = 0;
+      for (size_t i = 0; i < ::dimensions(lhs); ++i) {
+        if (lhs[i] != rhs[i]) {
+          std::cout << "  lhs[" << i << "]: " << lhs[i] << " - rhs[" << i
+                    << "]: " << rhs[i] << std::endl;
+          printed++;
+        }
+        if (printed > 50) {
+          std::cout << "  ..." << std::endl;
+          break;
+        }
+      }
+    }
+    return equal;
   }
 
   auto compare_cluster_centroids(const ivf_pq_index& rhs) const {
@@ -1611,7 +1630,8 @@ class ivf_pq_index {
     }
     return compare_feature_vectors(
         partitioned_pq_vectors_->indices(),
-        rhs.partitioned_pq_vectors_->indices());
+        rhs.partitioned_pq_vectors_->indices(),
+        "partitioned_pq_vectors_->indices()");
   }
 
   auto compare_ivf_ids(const ivf_pq_index& rhs) const {
@@ -1622,7 +1642,9 @@ class ivf_pq_index {
       return false;
     }
     return compare_feature_vectors(
-        partitioned_pq_vectors_->ids(), rhs.partitioned_pq_vectors_->ids());
+        partitioned_pq_vectors_->ids(),
+        rhs.partitioned_pq_vectors_->ids(),
+        "partitioned_pq_vectors_->ids()");
   }
 
   auto compare_pq_ivf_vectors(const ivf_pq_index& rhs) const {

src/include/test/unit_ivf_pq_index.cc

@@ -310,25 +310,22 @@ TEST_CASE("ivf_index write and read", "[ivf_pq_index]") {
     vfs.remove_dir(ivf_index_uri);
   }
 
+  // Create and write an index.
   auto training_set = tdbColMajorMatrix<float>(ctx, siftsmall_inputs_uri, 0);
   load(training_set);
   std::vector<siftsmall_ids_type> ids(num_vectors(training_set));
   std::iota(begin(ids), end(ids), 0);
   auto idx = ivf_pq_index<float, uint32_t, uint32_t>(
-      /*dimension,*/ nlist, num_subspaces, max_iters, nthreads);
+      nlist, num_subspaces, max_iters, nthreads);
   idx.train_ivf(training_set, kmeans_init::kmeanspp);
   idx.add(training_set, ids);
-  ivf_index_uri =
-      (std::filesystem::temp_directory_path() / "second_tmp_ivf_index")
-          .string();
-  if (vfs.is_dir(ivf_index_uri)) {
-    vfs.remove_dir(ivf_index_uri);
-  }
-
   idx.write_index(ctx, ivf_index_uri);
+
+  // Load it from URI.
   auto idx2 = ivf_pq_index<float, uint32_t, uint32_t>(ctx, ivf_index_uri);
   idx2.read_index_infinite();
 
+  // Check that the two indexes are the same.
   CHECK(idx.compare_cached_metadata(idx2));
   CHECK(idx.compare_cached_metadata(idx2));
   CHECK(idx.compare_cluster_centroids(idx2));