Fix bug with adjacency_row_index_uri() array type being incorrect, remove adjacency_row_index_type from API (#405)

Author: jparismorgan

apis/python/src/tiledb/vector_search/type_erased_module.cc

@@ -364,9 +364,6 @@ void init_type_erased_module(py::module_& m) {
           py::arg("storage_version") = "")
       .def("feature_type_string", &IndexVamana::feature_type_string)
       .def("id_type_string", &IndexVamana::id_type_string)
-      .def(
-          "adjacency_row_index_type_string",
-          &IndexVamana::adjacency_row_index_type_string)
       .def("dimensions", &IndexVamana::dimensions)
       .def_static(
           "clear_history",

apis/python/src/tiledb/vector_search/vamana_index.py

@@ -163,7 +163,6 @@ def create(
     index = vspy.IndexVamana(
         feature_type=np.dtype(vector_type).name,
         id_type=np.dtype(np.uint64).name,
-        adjacency_row_index_type=np.dtype(np.uint64).name,
         dimensions=dimensions,
     )
     # TODO(paris): Run all of this with a single C++ call.

apis/python/test/test_type_erased_module.py

@@ -266,29 +266,21 @@ def test_construct_IndexVamana():
     a = vspy.IndexVamana()
     assert a.feature_type_string() == "any"
     assert a.id_type_string() == "uint32"
-    assert a.adjacency_row_index_type_string() == "uint32"
     assert a.dimensions() == 0
 
     a = vspy.IndexVamana(feature_type="float32")
     assert a.feature_type_string() == "float32"
     assert a.id_type_string() == "uint32"
-    assert a.adjacency_row_index_type_string() == "uint32"
     assert a.dimensions() == 0
 
-    a = vspy.IndexVamana(
-        feature_type="uint8", id_type="uint64", adjacency_row_index_type="int64"
-    )
+    a = vspy.IndexVamana(feature_type="uint8", id_type="uint64")
     assert a.feature_type_string() == "uint8"
     assert a.id_type_string() == "uint64"
-    assert a.adjacency_row_index_type_string() == "int64"
     assert a.dimensions() == 0
 
-    a = vspy.IndexVamana(
-        feature_type="float32", id_type="int64", adjacency_row_index_type="uint64"
-    )
+    a = vspy.IndexVamana(feature_type="float32", id_type="int64")
     assert a.feature_type_string() == "float32"
     assert a.id_type_string() == "int64"
-    assert a.adjacency_row_index_type_string() == "uint64"
     assert a.dimensions() == 0
 
 
@@ -299,13 +291,11 @@ def test_construct_IndexVamana_with_empty_vector(tmp_path):
     dimensions = 128
     feature_type = "float32"
     id_type = "uint64"
-    adjacency_row_index_type = "uint64"
 
     # First create an empty index.
     a = vspy.IndexVamana(
         feature_type=feature_type,
         id_type=id_type,
-        adjacency_row_index_type=adjacency_row_index_type,
         dimensions=dimensions,
     )
     empty_vector = vspy.FeatureVectorArray(dimensions, 0, feature_type, id_type)
@@ -335,9 +325,7 @@ def test_inplace_build_query_IndexVamana():
     opt_l = 100
     k_nn = 10
 
-    a = vspy.IndexVamana(
-        id_type="uint32", adjacency_row_index_type="uint32", feature_type="float32"
-    )
+    a = vspy.IndexVamana(id_type="uint32", feature_type="float32")
 
     training_set = vspy.FeatureVectorArray(ctx, siftsmall_inputs_uri)
     assert training_set.feature_type_string() == "float32"

src/include/api/vamana_index.h

@@ -90,7 +90,6 @@ class IndexVamana {
       const std::optional<IndexOptions>& config = std::nullopt) {
     feature_datatype_ = TILEDB_ANY;
     id_datatype_ = TILEDB_UINT32;
-    adjacency_row_index_datatype_ = TILEDB_UINT32;
 
     if (config) {
       for (auto&& c : *config) {
@@ -108,8 +107,6 @@ class IndexVamana {
           feature_datatype_ = string_to_datatype(value);
         } else if (key == "id_type") {
           id_datatype_ = string_to_datatype(value);
-        } else if (key == "adjacency_row_index_type") {
-          adjacency_row_index_datatype_ = string_to_datatype(value);
         } else {
           throw std::runtime_error("Invalid index config key: " + key);
         }
@@ -135,12 +132,7 @@ class IndexVamana {
       const tiledb::Context& ctx,
       const URI& group_uri,
       std::optional<TemporalPolicy> temporal_policy = std::nullopt) {
-    read_types(
-        ctx,
-        group_uri,
-        &feature_datatype_,
-        &id_datatype_,
-        &adjacency_row_index_datatype_);
+    read_types(ctx, group_uri, &feature_datatype_, &id_datatype_);
 
     auto type = std::tuple{
         feature_datatype_, id_datatype_, adjacency_row_index_datatype_};
@@ -250,16 +242,10 @@ class IndexVamana {
       uint64_t timestamp) {
     tiledb_datatype_t feature_datatype{TILEDB_ANY};
     tiledb_datatype_t id_datatype{TILEDB_ANY};
-    tiledb_datatype_t adjacency_row_index_datatype{TILEDB_ANY};
-    read_types(
-        ctx,
-        group_uri,
-        &feature_datatype,
-        &id_datatype,
-        &adjacency_row_index_datatype);
-
-    auto type =
-        std::tuple{feature_datatype, id_datatype, adjacency_row_index_datatype};
+    read_types(ctx, group_uri, &feature_datatype, &id_datatype);
+
+    auto type = std::tuple{
+        feature_datatype, id_datatype, adjacency_row_index_datatype_};
     if (clear_history_dispatch_table.find(type) ==
         clear_history_dispatch_table.end()) {
       throw std::runtime_error("Unsupported datatype combination");
@@ -307,29 +293,17 @@ class IndexVamana {
     return datatype_to_string(id_datatype_);
   }
 
-  constexpr auto adjacency_row_index_type() const {
-    return adjacency_row_index_datatype_;
-  }
-
-  inline auto adjacency_row_index_type_string() const {
-    return datatype_to_string(adjacency_row_index_datatype_);
-  }
-
  private:
   static void read_types(
       const tiledb::Context& ctx,
       const std::string& group_uri,
       tiledb_datatype_t* feature_datatype,
-      tiledb_datatype_t* id_datatype,
-      tiledb_datatype_t* adjacency_row_index_datatype) {
+      tiledb_datatype_t* id_datatype) {
     using metadata_element =
         std::tuple<std::string, tiledb_datatype_t*, tiledb_datatype_t>;
     std::vector<metadata_element> metadata{
         {"feature_datatype", feature_datatype, TILEDB_UINT32},
-        {"id_datatype", id_datatype, TILEDB_UINT32},
-        {"adjacency_row_index_datatype",
-         adjacency_row_index_datatype,
-         TILEDB_UINT32}};
+        {"id_datatype", id_datatype, TILEDB_UINT32}};
 
     tiledb::Group read_group(ctx, group_uri, TILEDB_READ, ctx.config());
 
@@ -536,7 +510,8 @@ class IndexVamana {
   size_t b_backtrack_ = 0;
   tiledb_datatype_t feature_datatype_{TILEDB_ANY};
   tiledb_datatype_t id_datatype_{TILEDB_ANY};
-  tiledb_datatype_t adjacency_row_index_datatype_{TILEDB_ANY};
+  static constexpr tiledb_datatype_t adjacency_row_index_datatype_{
+      TILEDB_UINT64};
   std::unique_ptr<index_base> index_;
 };
 

src/include/index/index_metadata.h

@@ -33,7 +33,7 @@
  *  "base_sizes",            // (json) list
  *  "dataset_type",          // "vector_search"
  *  "dtype",                 // "float32", etc (Python dtype names)
- *  "index_type",            // "FLAT", "IVF_FLAT", "VAMANA"
+ *  "index_type",            // "FLAT", "IVF_FLAT", "VAMANA", "IVF_PQ"
  *  "ingestion_timestamps",  // (json) list
  *  "storage_version",       // "0.3"
  *  "temp_size",             // TILEDB_INT64 or TILEDB_FLOAT64

src/include/index/vamana_group.h

@@ -237,12 +237,12 @@ class vamana_index_group : public base_index_group<index_type> {
      * @todo Make this table-driven
      *************************************************************************/
     metadata_.adjacency_scores_datatype_ =
-        type_to_tiledb_v<typename index_type::score_type>;
+        type_to_tiledb_v<typename index_type::adjacency_scores_type>;
     metadata_.adjacency_row_index_datatype_ =
         type_to_tiledb_v<typename index_type::adjacency_row_index_type>;
 
     metadata_.adjacency_scores_type_str_ =
-        type_to_string_v<typename index_type::score_type>;
+        type_to_string_v<typename index_type::adjacency_scores_type>;
     metadata_.adjacency_row_index_type_str_ =
         type_to_string_v<typename index_type::adjacency_row_index_type>;
 
@@ -281,7 +281,7 @@ class vamana_index_group : public base_index_group<index_type> {
     tiledb_helpers::add_to_group(
         write_group, this->ids_uri(), this->ids_array_name());
 
-    create_empty_for_vector<typename index_type::score_type>(
+    create_empty_for_vector<typename index_type::adjacency_scores_type>(
         cached_ctx_,
         adjacency_scores_uri(),
         default_domain,
@@ -299,7 +299,7 @@ class vamana_index_group : public base_index_group<index_type> {
     tiledb_helpers::add_to_group(
         write_group, adjacency_ids_uri(), adjacency_ids_array_name());
 
-    create_empty_for_vector<typename index_type::id_type>(
+    create_empty_for_vector<typename index_type::adjacency_row_index_type>(
         cached_ctx_,
         adjacency_row_index_uri(),
         default_domain,

src/include/index/vamana_index.h

@@ -98,17 +98,22 @@ auto medoid(auto&& P, Distance distance = Distance{}) {
 }
 
 /**
- * @brief Index class for vamana search
- * @tparam feature_type Type of the elements in the feature vectors
- * @tparam id_type Type of the ids of the feature vectors
+ * @brief The Vamana index.
+ *
+ * @tparam FeatureType Type of the elements in the feature vectors.
+ * @tparam IdType Type of the ids of the feature vectors.
+ * @tparam AdjacencyRowIndexType Types of the indexes used in the graph.
  */
-template <class FeatureType, class IdType, class IndexType = uint64_t>
+template <
+    class FeatureType,
+    class IdType,
+    class AdjacencyRowIndexType = uint64_t>
 class vamana_index {
  public:
   using feature_type = FeatureType;
   using id_type = IdType;
-  using adjacency_row_index_type = IndexType;
-  using score_type = float;
+  using adjacency_row_index_type = AdjacencyRowIndexType;
+  using adjacency_scores_type = float;
 
   using group_type = vamana_index_group<vamana_index>;
   using metadata_type = vamana_index_metadata;
@@ -136,7 +141,7 @@ class vamana_index {
   uint64_t num_edges_{0};
 
   /** The graph representing the index over `feature_vectors_` */
-  ::detail::graph::adj_list<score_type, id_type> graph_;
+  ::detail::graph::adj_list<adjacency_scores_type, id_type> graph_;
 
   /*
    * The medoid of the feature vectors -- the vector in the set that is closest
@@ -250,7 +255,7 @@ class vamana_index {
      ****************************************************************************/
     graph_ = ::detail::graph::adj_list<feature_type, id_type>(num_vectors_);
 
-    auto adj_scores = read_vector<score_type>(
+    auto adj_scores = read_vector<adjacency_scores_type>(
         group_->cached_ctx(),
         group_->adjacency_scores_uri(),
         0,
@@ -467,7 +472,7 @@ class vamana_index {
 
     auto top_k = ColMajorMatrix<id_type>(k_nn, ::num_vectors(queries));
     auto top_k_scores =
-        ColMajorMatrix<score_type>(k_nn, ::num_vectors(queries));
+        ColMajorMatrix<adjacency_scores_type>(k_nn, ::num_vectors(queries));
 
     for (size_t i = 0; i < num_vectors(queries); ++i) {
       auto&& [tk_scores, tk, V] = ::best_first_O2(
@@ -494,7 +499,7 @@ class vamana_index {
 
     auto top_k = ColMajorMatrix<id_type>(k_nn, ::num_vectors(queries));
     auto top_k_scores =
-        ColMajorMatrix<score_type>(k_nn, ::num_vectors(queries));
+        ColMajorMatrix<adjacency_scores_type>(k_nn, ::num_vectors(queries));
 
     for (size_t i = 0; i < num_vectors(queries); ++i) {
       auto&& [tk_scores, tk, V] = ::best_first_O3(
@@ -521,7 +526,7 @@ class vamana_index {
 
     auto top_k = ColMajorMatrix<id_type>(k_nn, ::num_vectors(queries));
     auto top_k_scores =
-        ColMajorMatrix<score_type>(k_nn, ::num_vectors(queries));
+        ColMajorMatrix<adjacency_scores_type>(k_nn, ::num_vectors(queries));
 
     for (size_t i = 0; i < num_vectors(queries); ++i) {
       auto&& [tk_scores, tk, V] = ::best_first_O4(
@@ -548,7 +553,7 @@ class vamana_index {
 
     auto top_k = ColMajorMatrix<id_type>(k_nn, ::num_vectors(queries));
     auto top_k_scores =
-        ColMajorMatrix<score_type>(k_nn, ::num_vectors(queries));
+        ColMajorMatrix<adjacency_scores_type>(k_nn, ::num_vectors(queries));
 
     for (size_t i = 0; i < num_vectors(queries); ++i) {
       auto&& [tk_scores, tk, V] = ::best_first_O5(
@@ -588,7 +593,8 @@ class vamana_index {
     // L = std::min<size_t>(L, l_build_);
 
     auto top_k = ColMajorMatrix<id_type>(k, ::num_vectors(query_set));
-    auto top_k_scores = ColMajorMatrix<score_type>(k, ::num_vectors(query_set));
+    auto top_k_scores =
+        ColMajorMatrix<adjacency_scores_type>(k, ::num_vectors(query_set));
 
 #if 0
     // Parallelized implementation -- we stay single-threaded for now
@@ -768,7 +774,7 @@ class vamana_index {
         false,
         temporal_policy_);
 
-    auto adj_scores = Vector<score_type>(graph_.num_edges());
+    auto adj_scores = Vector<adjacency_scores_type>(graph_.num_edges());
     auto adj_ids = Vector<id_type>(graph_.num_edges());
     auto adj_index =
         Vector<adjacency_row_index_type>(graph_.num_vertices() + 1);

src/include/index/vamana_metadata.h

@@ -117,10 +117,6 @@ class vamana_index_metadata
       {"alpha_min", &alpha_min_, TILEDB_FLOAT32, false},
       {"alpha_max", &alpha_max_, TILEDB_FLOAT32, false},
       {"medoid", &medoid_, TILEDB_UINT64, false},
-      {"adjacency_row_index_datatype",
-       &adjacency_row_index_datatype_,
-       TILEDB_UINT32,
-       false},
   };
 
   void clear_history_impl(uint64_t timestamp) {