Add more Python and C++ unit tests around metadata (#281)

Author: jparismorgan

apis/python/test/test_index.py

@@ -2,11 +2,14 @@
 import pytest
 from array_paths import *
 from common import *
+import json
 
 import tiledb.vector_search.index as ind
 from tiledb.vector_search import Index
 from tiledb.vector_search import flat_index
 from tiledb.vector_search import ivf_flat_index
+from tiledb.vector_search.index import create_metadata
+from tiledb.vector_search.index import DATASET_TYPE
 from tiledb.vector_search.flat_index import FlatIndex
 from tiledb.vector_search.ingestion import ingest
 from tiledb.vector_search.ivf_flat_index import IVFFlatIndex
@@ -18,18 +21,49 @@ def query_and_check(index, queries, k, expected, **kwargs):
         result_d, result_i = index.query(queries, k=k, **kwargs)
         assert expected.issubset(set(result_i[0]))
 
+def check_default_metadata(uri, expected_vector_type, expected_storage_version, expected_index_type):
+    group = tiledb.Group(uri, "r", ctx=tiledb.Ctx(None))
+    assert "dataset_type" in group.meta
+    assert group.meta["dataset_type"] == DATASET_TYPE
+    assert type(group.meta["dataset_type"]) == str
+
+    assert "dtype" in group.meta
+    assert group.meta["dtype"] == np.dtype(expected_vector_type).name
+    assert type(group.meta["dtype"]) == str
+
+    assert "storage_version" in group.meta
+    assert group.meta["storage_version"] == expected_storage_version
+    assert type(group.meta["storage_version"]) == str
+
+    assert "index_type" in group.meta
+    assert group.meta["index_type"] == expected_index_type
+    assert type(group.meta["index_type"]) == str
+
+    assert "base_sizes" in group.meta
+    assert group.meta["base_sizes"] == json.dumps([0])
+    assert type(group.meta["base_sizes"]) == str
+
+    assert "ingestion_timestamps" in group.meta
+    assert group.meta["ingestion_timestamps"] == json.dumps([0])
+    assert type(group.meta["ingestion_timestamps"]) == str
+
+    assert "has_updates" in group.meta
+    assert group.meta["has_updates"] == False
+    assert type(group.meta["has_updates"]) == np.int64
 
 def test_flat_index(tmp_path):
     uri = os.path.join(tmp_path, "array")
-    index = flat_index.create(uri=uri, dimensions=3, vector_type=np.dtype(np.uint8))
+    vector_type = np.dtype(np.uint8)
+    index = flat_index.create(uri=uri, dimensions=3, vector_type=vector_type)
     query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {ind.MAX_UINT64})
+    check_default_metadata(uri, vector_type, STORAGE_VERSION, "FLAT")
 
     update_vectors = np.empty([5], dtype=object)
-    update_vectors[0] = np.array([0, 0, 0], dtype=np.dtype(np.uint8))
-    update_vectors[1] = np.array([1, 1, 1], dtype=np.dtype(np.uint8))
-    update_vectors[2] = np.array([2, 2, 2], dtype=np.dtype(np.uint8))
-    update_vectors[3] = np.array([3, 3, 3], dtype=np.dtype(np.uint8))
-    update_vectors[4] = np.array([4, 4, 4], dtype=np.dtype(np.uint8))
+    update_vectors[0] = np.array([0, 0, 0], dtype=vector_type)
+    update_vectors[1] = np.array([1, 1, 1], dtype=vector_type)
+    update_vectors[2] = np.array([2, 2, 2], dtype=vector_type)
+    update_vectors[3] = np.array([3, 3, 3], dtype=vector_type)
+    update_vectors[4] = np.array([4, 4, 4], dtype=vector_type)
     index.update_batch(vectors=update_vectors, external_ids=np.array([0, 1, 2, 3, 4]))
     query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {1, 2, 3})
 
@@ -43,8 +77,8 @@ def test_flat_index(tmp_path):
     query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {0, 2, 4})
 
     update_vectors = np.empty([2], dtype=object)
-    update_vectors[0] = np.array([1, 1, 1], dtype=np.dtype(np.uint8))
-    update_vectors[1] = np.array([3, 3, 3], dtype=np.dtype(np.uint8))
+    update_vectors[0] = np.array([1, 1, 1], dtype=vector_type)
+    update_vectors[1] = np.array([3, 3, 3], dtype=vector_type)
     index.update_batch(vectors=update_vectors, external_ids=np.array([1, 3]))
     query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {1, 2, 3})
 
@@ -61,9 +95,9 @@ def test_flat_index(tmp_path):
 def test_ivf_flat_index(tmp_path):
     partitions = 10
     uri = os.path.join(tmp_path, "array")
-
+    vector_type = np.dtype(np.uint8)
     index = ivf_flat_index.create(
-        uri=uri, dimensions=3, vector_type=np.dtype(np.uint8), partitions=partitions
+        uri=uri, dimensions=3, vector_type=vector_type, partitions=partitions
     )
     query_and_check(
         index,
@@ -72,13 +106,14 @@ def test_ivf_flat_index(tmp_path):
         {ind.MAX_UINT64},
         nprobe=partitions,
     )
+    check_default_metadata(uri, vector_type, STORAGE_VERSION, "IVF_FLAT")
 
     update_vectors = np.empty([5], dtype=object)
-    update_vectors[0] = np.array([0, 0, 0], dtype=np.dtype(np.uint8))
-    update_vectors[1] = np.array([1, 1, 1], dtype=np.dtype(np.uint8))
-    update_vectors[2] = np.array([2, 2, 2], dtype=np.dtype(np.uint8))
-    update_vectors[3] = np.array([3, 3, 3], dtype=np.dtype(np.uint8))
-    update_vectors[4] = np.array([4, 4, 4], dtype=np.dtype(np.uint8))
+    update_vectors[0] = np.array([0, 0, 0], dtype=vector_type)
+    update_vectors[1] = np.array([1, 1, 1], dtype=vector_type)
+    update_vectors[2] = np.array([2, 2, 2], dtype=vector_type)
+    update_vectors[3] = np.array([3, 3, 3], dtype=vector_type)
+    update_vectors[4] = np.array([4, 4, 4], dtype=vector_type)
     index.update_batch(vectors=update_vectors, external_ids=np.array([0, 1, 2, 3, 4]))
 
     query_and_check(
@@ -102,8 +137,8 @@ def test_ivf_flat_index(tmp_path):
     )
 
     update_vectors = np.empty([2], dtype=object)
-    update_vectors[0] = np.array([1, 1, 1], dtype=np.dtype(np.uint8))
-    update_vectors[1] = np.array([3, 3, 3], dtype=np.dtype(np.uint8))
+    update_vectors[0] = np.array([1, 1, 1], dtype=vector_type)
+    update_vectors[1] = np.array([3, 3, 3], dtype=vector_type)
     index.update_batch(vectors=update_vectors, external_ids=np.array([1, 3]))
     query_and_check(
         index, np.array([[2, 2, 2]], dtype=np.float32), 3, {1, 2, 3}, nprobe=partitions
@@ -251,3 +286,17 @@ def test_index_with_incorrect_num_of_query_columns_in_single_vector_query(tmp_pa
         # TODO:  This also throws a TypeError for incorrect dimension
         with pytest.raises(TypeError):
             index.query(np.array([1, 1, 1], dtype=np.float32), k=3)
+
+def test_create_metadata(tmp_path):
+    uri = os.path.join(tmp_path, "array")
+
+    # Create the metadata at the specified URI.
+    dimensions = 3
+    vector_type: np.dtype = np.dtype(np.uint8)
+    index_type: str = "IVF_FLAT"
+    storage_version: str = STORAGE_VERSION
+    group_exists: bool = False
+    create_metadata(uri, dimensions, vector_type, index_type, storage_version, group_exists)
+    
+    # Check it contains the default metadata.
+    check_default_metadata(uri, vector_type, storage_version, index_type)

src/include/index/vamana_index.h

@@ -39,15 +39,15 @@
 #include <queue>
 #include <unordered_set>
 
+#include "detail/graph/adj_list.h"
+#include "detail/graph/graph_utils.h"
+#include "detail/linalg/vector.h"
+#include "index/vamana_group.h"
 #include "scoring.h"
 #include "stats.h"
 #include "utils/fixed_min_heap.h"
 #include "utils/print_types.h"
 
-#include "detail/graph/adj_list.h"
-#include "detail/graph/graph_utils.h"
-#include "index/vamana_group.h"
-
 #include <tiledb/tiledb>
 
 #include <tiledb/group_experimental.h>

src/include/test/unit_api_ivf_flat_index.cc

@@ -201,7 +201,7 @@ TEST_CASE(
     auto nt = num_vectors(t);
     auto recall = ((double)intersections) / ((double)nt * k_nn);
     if (nprobe == 32) {
-      CHECK(recall == 1.0);
+      CHECK(std::abs(recall - 1.0) <= 1e-3);
     } else if (nprobe == 8) {
       CHECK(recall > 0.925);
     }
@@ -305,7 +305,7 @@ TEST_CASE(
     CHECK(nt == nv);
     auto recall = ((double)intersections_a) / ((double)nt * k_nn);
     if (nprobe == 32) {
-      CHECK(recall == 1.0);
+      CHECK(std::abs(recall - 1.0) <= 1e-3);
     } else if (nprobe == 8) {
       CHECK(recall > 0.925);
     }

src/include/test/unit_backwards_compatibility.cc

@@ -53,6 +53,8 @@ TEST_CASE(
     "backwards_compatibility: test_query_old_indices",
     "[backwards_compatibility]") {
   tiledb::Context ctx;
+  tiledb::Config cfg;
+
   std::string datasets_path = backwards_compatibility_root / "data";
   auto base =
       read_bin_local<siftsmall_feature_type>(ctx, siftmicro_inputs_file);
@@ -93,8 +95,11 @@ TEST_CASE(
           index_uri.find("0.0.17") != std::string::npos) {
         continue;
       }
+
+      auto read_group = tiledb::Group(ctx, index_uri, TILEDB_READ, cfg);
       std::vector<float> expected_distances(query_indices.size(), 0.0);
       if (index_uri.find("ivf_flat") != std::string::npos) {
+        // First check that we can query the index.
         auto index = IndexIVFFlat(ctx, index_uri);
         auto&& [scores, ids] =
             index.query_infinite_ram(queries_feature_vector_array, 1, 10);
@@ -111,6 +116,10 @@ TEST_CASE(
           CHECK(ids_span[0][i] == query_indices[i]);
           CHECK(scores_span[0][i] == 0);
         }
+
+        // Next check that we can load the metadata.
+        auto metadata = ivf_flat_index_metadata();
+        metadata.load_metadata(read_group);
       } else if (index_uri.find("flat") != std::string::npos) {
         // TODO(paris): Fix flat_l2_index and re-enable. Right now it just tries
         // to load the URI as a tdbMatrix.

src/include/test/unit_ivf_flat_metadata.cc

@@ -35,6 +35,8 @@
 #include <string>
 
 #include "array_defs.h"
+#include "detail/linalg/tdb_matrix.h"
+#include "index/ivf_flat_index.h"
 #include "index/ivf_flat_metadata.h"
 
 TEST_CASE("ivf_flat_metadata: test test", "[ivf_flat_metadata]") {
@@ -55,6 +57,31 @@ TEST_CASE(
   y.dump();
 }
 
+TEST_CASE(
+    "ivf_flat_metadata: load metadata from index", "[ivf_flat_metadata]") {
+  tiledb::Context ctx;
+  tiledb::Config cfg;
+
+  std::string uri =
+      (std::filesystem::temp_directory_path() / "tmp_ivf_index").string();
+  auto training_vectors =
+      tdbColMajorPreLoadMatrix<float>(ctx, siftsmall_inputs_uri, 0);
+  auto idx = ivf_flat_index<float, uint32_t, uint32_t>(100, 1);
+  idx.train(training_vectors, kmeans_init::kmeanspp);
+  idx.add(training_vectors);
+  idx.write_index(ctx, uri, true);
+
+  auto read_group = tiledb::Group(ctx, uri, TILEDB_READ, cfg);
+
+  auto x = ivf_flat_index_metadata();
+  x.load_metadata(read_group);
+
+  // Compare two constructed objects.
+  ivf_flat_index_metadata y;
+  y.load_metadata(read_group);
+  CHECK(x.compare_metadata(y));
+}
+
 TEST_CASE("ivf_flat_metadata: open group", "[ivf_flat_metadata]") {
   tiledb::Context ctx;
   tiledb::Config cfg;

src/include/test/unit_vamana_metadata.cc

@@ -33,6 +33,8 @@
 #include <tiledb/tiledb>
 #include <vector>
 #include "array_defs.h"
+#include "detail/linalg/tdb_matrix.h"
+#include "index/vamana_index.h"
 #include "index/vamana_metadata.h"
 
 std::vector<std::tuple<std::string, std::string>> expected_str{
@@ -88,6 +90,30 @@ TEST_CASE("vamana_metadata: default constructor dump", "[vamana_metadata]") {
   }
 }
 
+TEST_CASE("vamana_metadata: load metadata from index", "[vamana_metadata]") {
+  tiledb::Context ctx;
+  tiledb::Config cfg;
+
+  std::string uri =
+      (std::filesystem::temp_directory_path() / "tmp_vamana_index").string();
+  auto training_vectors =
+      tdbColMajorPreLoadMatrix<float>(ctx, siftsmall_inputs_uri);
+  auto idx =
+      vamana_index<float, uint64_t>(num_vectors(training_vectors), 20, 40, 30);
+  idx.train(training_vectors);
+  idx.add(training_vectors);
+  idx.write_index(ctx, uri, true);
+
+  auto read_group = tiledb::Group(ctx, uri, TILEDB_READ, cfg);
+  auto x = vamana_index_metadata();
+  x.load_metadata(read_group);
+
+  // Compare two constructed objects.
+  vamana_index_metadata y;
+  y.load_metadata(read_group);
+  CHECK(x.compare_metadata(y));
+}
+
 // @todo More vamana groups (from "real" data) to test with
 TEST_CASE("vamana_metadata: open group", "[vamana_metadata]") {
   bool debug = false;