Check finite and infinite IVF_PQ queries return the same ids and distances, fix count_intersections() to not modify inputs, update read_index_finite() to return data (#509)

Author: jparismorgan

src/include/api/feature_vector_array.h

@@ -543,4 +543,203 @@ auto count_intersections(
   }
 }
 
+bool are_equal(
+    const FeatureVectorArray& a,
+    const FeatureVectorArray& b,
+    double epsilon = 0.0) {
+  if (a.feature_type() != b.feature_type()) {
+    std::cout << "[feature_vector_array@are_equal] Feature types do not match: "
+              << a.feature_type_string() << " vs " << b.feature_type_string()
+              << std::endl;
+    return false;
+  }
+  if (a.feature_size() != b.feature_size()) {
+    std::cout << "[feature_vector_array@are_equal] Feature sizes do not match: "
+              << a.feature_size() << " vs " << b.feature_size() << std::endl;
+    return false;
+  }
+  if (a.num_vectors() != b.num_vectors()) {
+    std::cout
+        << "[feature_vector_array@are_equal] Number of vectors do not match: "
+        << a.num_vectors() << " vs " << b.num_vectors() << std::endl;
+    return false;
+  }
+  if (a.dimensions() != b.dimensions()) {
+    std::cout << "[feature_vector_array@are_equal] Number of dimensions do not "
+                 "match: "
+              << a.dimensions() << " vs " << b.dimensions() << std::endl;
+    return false;
+  }
+
+  if (a.ids_type() != b.ids_type()) {
+    std::cout << "[feature_vector_array@are_equal] IDs types do not match: "
+              << a.ids_type_string() << " vs " << b.ids_type_string()
+              << std::endl;
+    return false;
+  }
+  if (a.ids_size() != b.ids_size()) {
+    std::cout << "[feature_vector_array@are_equal] IDs sizes do not match: "
+              << a.ids_size() << " vs " << b.ids_size() << std::endl;
+    return false;
+  }
+  if (a.num_ids() != b.num_ids()) {
+    std::cout << "[feature_vector_array@are_equal] Number of IDs do not match: "
+              << a.num_ids() << " vs " << b.num_ids() << std::endl;
+    return false;
+  }
+
+  if (a.extents() != b.extents()) {
+    std::cout << "[feature_vector_array@are_equal] Extents do not match: "
+              << "A: {" << a.extents()[0] << ", " << a.dimensions() << "}, "
+              << "B: {" << b.extents()[0] << ", " << b.dimensions() << "}"
+              << std::endl;
+    return false;
+  }
+
+  // Compare the data of the feature vectors with optional epsilon
+  auto compare_data = [&epsilon](
+                          const auto* data_a, const auto* data_b, size_t size) {
+    if (epsilon > 0.0) {
+      for (size_t i = 0; i < size; ++i) {
+        if (std::abs(
+                static_cast<double>(data_a[i]) -
+                static_cast<double>(data_b[i])) > epsilon) {
+          std::cout
+              << "[feature_vector_array@are_equal] Data mismatch at index " << i
+              << ": " << data_a[i] << " vs " << data_b[i]
+              << " (epsilon: " << epsilon << ")" << std::endl;
+          return false;
+        }
+      }
+    } else {
+      for (size_t i = 0; i < size; ++i) {
+        if (data_a[i] != data_b[i]) {
+          std::cout
+              << "[feature_vector_array@are_equal] Data mismatch at index " << i
+              << ": " << data_a[i] << " vs " << data_b[i] << std::endl;
+          return false;
+        }
+      }
+    }
+    return true;
+  };
+
+  switch (a.feature_type()) {
+    case TILEDB_FLOAT32: {
+      const auto* data_a = static_cast<const float*>(a.data());
+      const auto* data_b = static_cast<const float*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type FLOAT32"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    case TILEDB_INT8: {
+      const auto* data_a = static_cast<const int8_t*>(a.data());
+      const auto* data_b = static_cast<const int8_t*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type INT8"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    case TILEDB_UINT8: {
+      const auto* data_a = static_cast<const uint8_t*>(a.data());
+      const auto* data_b = static_cast<const uint8_t*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type UINT8"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    case TILEDB_INT32: {
+      const auto* data_a = static_cast<const int32_t*>(a.data());
+      const auto* data_b = static_cast<const int32_t*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type INT32"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    case TILEDB_UINT32: {
+      const auto* data_a = static_cast<const uint32_t*>(a.data());
+      const auto* data_b = static_cast<const uint32_t*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type UINT32"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    case TILEDB_INT64: {
+      const auto* data_a = static_cast<const int64_t*>(a.data());
+      const auto* data_b = static_cast<const int64_t*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type INT64"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    case TILEDB_UINT64: {
+      const auto* data_a = static_cast<const uint64_t*>(a.data());
+      const auto* data_b = static_cast<const uint64_t*>(b.data());
+      if (!compare_data(data_a, data_b, a.num_vectors() * a.dimensions())) {
+        std::cout << "[feature_vector_array@are_equal] Feature vector data "
+                     "comparison failed for type UINT64"
+                  << std::endl;
+        return false;
+      }
+      break;
+    }
+    default:
+      std::cout << "[feature_vector_array@are_equal] Unsupported feature "
+                   "vector attribute type"
+                << std::endl;
+      throw std::runtime_error("Unsupported attribute type");
+  }
+
+  // If the arrays have IDs, compare the IDs as well
+  if (a.ids_type() != TILEDB_ANY && b.ids_type() != TILEDB_ANY) {
+    switch (a.ids_type()) {
+      case TILEDB_UINT32: {
+        const auto* ids_a = static_cast<const uint32_t*>(a.ids());
+        const auto* ids_b = static_cast<const uint32_t*>(b.ids());
+        if (!compare_data(ids_a, ids_b, a.num_ids())) {
+          std::cout << "[feature_vector_array@are_equal] ID comparison failed "
+                       "for type UINT32"
+                    << std::endl;
+          return false;
+        }
+        break;
+      }
+      case TILEDB_UINT64: {
+        const auto* ids_a = static_cast<const uint64_t*>(a.ids());
+        const auto* ids_b = static_cast<const uint64_t*>(b.ids());
+        if (!compare_data(ids_a, ids_b, a.num_ids())) {
+          std::cout << "[feature_vector_array@are_equal] ID comparison failed "
+                       "for type UINT64"
+                    << std::endl;
+          return false;
+        }
+        break;
+      }
+      default:
+        throw std::runtime_error("Unsupported ID type");
+    }
+  }
+
+  return true;
+}
+
 #endif  // TILEDB_API_FEATURE_VECTOR_ARRAY_H

src/include/detail/linalg/matrix.h

@@ -442,8 +442,10 @@ constexpr auto SubMatrix(
 template <class Matrix>
 void debug_matrix(
     const Matrix& matrix, const std::string& msg = "", size_t max_size = 10) {
-  auto rowsEnd = std::min(dimensions(matrix), static_cast<size_t>(max_size));
-  auto colsEnd = std::min(num_vectors(matrix), static_cast<size_t>(max_size));
+  auto rowsEnd = std::min(
+      dimensions(matrix), static_cast<typename Matrix::size_type>(max_size));
+  auto colsEnd = std::min(
+      num_vectors(matrix), static_cast<typename Matrix::size_type>(max_size));
 
   std::cout << "# " << msg << " (" << dimensions(matrix) << " rows x "
             << num_vectors(matrix) << " cols) ("

src/include/detail/linalg/matrix_with_ids.h

@@ -218,8 +218,12 @@ void debug_matrix_with_ids(
     const MatrixWithIds& matrix,
     const std::string& msg = "",
     size_t max_size = 10) {
-  auto rowsEnd = std::min(dimensions(matrix), static_cast<size_t>(max_size));
-  auto colsEnd = std::min(num_vectors(matrix), static_cast<size_t>(max_size));
+  auto rowsEnd = std::min(
+      dimensions(matrix),
+      static_cast<typename MatrixWithIds::size_type>(max_size));
+  auto colsEnd = std::min(
+      num_vectors(matrix),
+      static_cast<typename MatrixWithIds::size_type>(max_size));
 
   debug_matrix(matrix, msg, max_size);
 

src/include/detail/linalg/partitioned_matrix.h

@@ -242,7 +242,7 @@ void debug_partitioned_matrix(
   debug_matrix(matrix, msg, max_size);
 
   std::cout << "# ids: [";
-  auto end = std::min(matrix.ids().size(), static_cast<size_t>(max_size));
+  auto end = std::min(matrix.num_vectors(), static_cast<size_t>(max_size));
   for (size_t i = 0; i < end; ++i) {
     std::cout << matrix.ids()[i];
     if (i != matrix.ids().size() - 1) {

src/include/index/ivf_pq_index.h

@@ -452,8 +452,6 @@ class ivf_pq_index {
       auto sub_begin = subspace * dimensions_ / num_subspaces_;
       auto sub_end = (subspace + 1) * dimensions_ / num_subspaces_;
 
-      auto local_sub_distance = SubDistance{sub_begin, sub_end};
-
       // @todo Make choice of kmeans init configurable
       sub_kmeans_random_init(
           training_set, cluster_centroids_, sub_begin, sub_end);
@@ -909,8 +907,7 @@ class ivf_pq_index {
     auto&& [active_partitions, active_queries] =
         detail::ivf::partition_ivf_flat_index<indices_type>(
             flat_ivf_centroids_, query_vectors, nprobe, num_threads_);
-
-    partitioned_pq_vectors_ = std::make_unique<tdb_pq_storage_type>(
+    auto partitioned_pq_vectors = std::make_unique<tdb_pq_storage_type>(
         group_->cached_ctx(),
         group_->pq_ivf_vectors_uri(),
         group_->pq_ivf_indices_uri(),
@@ -923,7 +920,9 @@ class ivf_pq_index {
     // NB: We don't load the partitioned_pq_vectors here. We will load them
     // when we do the query.
     return std::make_tuple(
-        std::move(active_partitions), std::move(active_queries));
+        std::move(active_partitions),
+        std::move(active_queries),
+        std::move(partitioned_pq_vectors));
   }
 
   /**
@@ -1236,21 +1235,17 @@ class ivf_pq_index {
           "run if you're loading the index by URI. Please open it by URI and "
           "try again. If you just wrote the index, open it up again by URI.");
     }
-    if (partitioned_pq_vectors_) {
-      // We did an infinite query before this. Reset so we can load again.
-      partitioned_pq_vectors_.reset();
-    }
     if (::num_vectors(flat_ivf_centroids_) < nprobe) {
       nprobe = ::num_vectors(flat_ivf_centroids_);
     }
-    auto&& [active_partitions, active_queries] =
+    auto&& [active_partitions, active_queries, partitioned_pq_vectors] =
         read_index_finite(query_vectors, nprobe, upper_bound);
     auto query_to_pq_centroid_distance_tables =
         std::move(*generate_query_to_pq_centroid_distance_tables<
                   Q,
                   ColMajorMatrix<float>>(query_vectors));
     return detail::ivf::query_finite_ram(
-        *partitioned_pq_vectors_,
+        *partitioned_pq_vectors,
         query_to_pq_centroid_distance_tables,
         active_queries,
         k_nn,

src/include/scoring.h

@@ -680,36 +680,39 @@ bool validate_top_k(TK& top_k, const G& g) {
 
 template <feature_vector_array U, feature_vector_array V>
 auto count_intersections(const U& I, const V& groundtruth, size_t k_nn) {
-  // print_types(I, groundtruth);
-
   size_t total_intersected = 0;
 
   if constexpr (feature_vector_array<std::remove_cvref_t<decltype(I)>>) {
     for (size_t i = 0; i < I.num_cols(); ++i) {
-      std::sort(begin(I[i]), end(I[i]));
-      std::sort(begin(groundtruth[i]), begin(groundtruth[i]) + k_nn);
+      std::vector<std::decay_t<decltype(I[i][0])>> sorted_I(
+          begin(I[i]), end(I[i]));
+      std::vector<std::decay_t<decltype(groundtruth[i][0])>> sorted_groundtruth(
+          begin(groundtruth[i]), begin(groundtruth[i]) + k_nn);
 
-      // @todo remove -- for debugging only
-      std::vector<size_t> x(begin(I[i]), end(I[i]));
-      std::vector<size_t> y(begin(groundtruth[i]), end(groundtruth[i]));
+      std::sort(begin(sorted_I), end(sorted_I));
+      std::sort(begin(sorted_groundtruth), end(sorted_groundtruth));
 
       total_intersected += std::set_intersection(
-          begin(I[i]),
-          end(I[i]),
-          begin(groundtruth[i]),
-          /*end(groundtruth[i]*/ begin(groundtruth[i]) + k_nn,
+          begin(sorted_I),
+          end(sorted_I),
+          begin(sorted_groundtruth),
+          end(sorted_groundtruth),
           assignment_counter{});
     }
   } else {
     if constexpr (feature_vector<std::remove_cvref_t<decltype(I)>>) {
-      std::sort(begin(I), end(I));
-      std::sort(begin(groundtruth), begin(groundtruth) + k_nn);
+      std::vector<std::decay_t<decltype(I[0])>> sorted_I(begin(I), end(I));
+      std::vector<std::decay_t<decltype(groundtruth[0])>> sorted_groundtruth(
+          begin(groundtruth), begin(groundtruth) + k_nn);
+
+      std::sort(begin(sorted_I), end(sorted_I));
+      std::sort(begin(sorted_groundtruth), end(sorted_groundtruth));
 
       total_intersected += std::set_intersection(
-          begin(I),
-          end(I),
-          begin(groundtruth),
-          /*end(groundtruth)*/ begin(groundtruth) + k_nn,
+          begin(sorted_I),
+          end(sorted_I),
+          begin(sorted_groundtruth),
+          end(sorted_groundtruth),
           assignment_counter{});
     } else {
       static_assert(
@@ -718,7 +721,7 @@ auto count_intersections(const U& I, const V& groundtruth, size_t k_nn) {
     }
   }
   return total_intersected;
-};
+}
 
 #if defined(TILEDB_VS_ENABLE_BLAS) && 0