Add metadata tests for type-erased indexes in Python (#334)

Author: jparismorgan

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

@@ -9,6 +9,7 @@
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
 from tiledb.vector_search.storage_formats import validate_storage_version
 from tiledb.vector_search.utils import add_to_group
+from tiledb.vector_search.utils import is_type_erased_index
 
 
 class TrainingSamplingPolicy(enum.Enum):
@@ -290,9 +291,6 @@ def ingest(
     DEFAULT_IMG_NAME = "3.9-vectorsearch"
     MAX_INT32 = 2**31 - 1
 
-    def is_type_erased_index():
-        return index_type == "VAMANA"
-
     class SourceType(enum.Enum):
         """SourceType of input vectors"""
 
@@ -755,7 +753,7 @@ def create_arrays(
 
         # Note that we don't create type-erased indexes (i.e. Vamana) here. Instead we create them
         # at very start of ingest() in C++.
-        elif not is_type_erased_index():
+        elif not is_type_erased_index(index_type):
             raise ValueError(f"Not supported index_type {index_type}")
 
     def read_external_ids(
@@ -2825,7 +2823,7 @@ def consolidate_and_vacuum(
         temp_size = int(group.meta.get("temp_size", "0"))
         group.close()
 
-        if not is_type_erased_index():
+        if not is_type_erased_index(index_type):
             # For type-erased indexes (i.e. Vamana), we update this metadata in the write_index()
             # call during create_ingestion_dag(), so don't do it here.
             group = tiledb.Group(index_group_uri, "w")

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

@@ -5,6 +5,10 @@
 import tiledb
 
 
+def is_type_erased_index(index_type: str) -> bool:
+    return index_type == "VAMANA"
+
+
 def add_to_group(group, uri, name):
     """
     Adds an object to a group. Automatically infers whether to use a relative path or absolute path.

apis/python/test/test_index.py

@@ -1,4 +1,5 @@
 import json
+import time
 
 import numpy as np
 import pytest
@@ -15,6 +16,7 @@
 from tiledb.vector_search.index import create_metadata
 from tiledb.vector_search.ingestion import ingest
 from tiledb.vector_search.ivf_flat_index import IVFFlatIndex
+from tiledb.vector_search.utils import is_type_erased_index
 from tiledb.vector_search.utils import load_fvecs
 from tiledb.vector_search.vamana_index import VamanaIndex
 
@@ -55,16 +57,31 @@ def check_default_metadata(
     assert type(group.meta["index_type"]) == str
 
     assert "base_sizes" in group.meta
-    assert group.meta["base_sizes"] == json.dumps([0])
+    if is_type_erased_index(expected_index_type):
+        # NOTE(paris): Type-erased indexes have two values upon creation.
+        assert group.meta["base_sizes"] == "[0,0]"
+    else:
+        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])
+    if is_type_erased_index(expected_index_type):
+        # NOTE(paris): Type-erased indexes have two values upon creation.
+        ingestion_timestamps = json.loads(group.meta["ingestion_timestamps"])
+        assert len(ingestion_timestamps) == 2
+        assert ingestion_timestamps[0] == 0
+        current_time_ms = int(time.time() * 1000)
+        assert ingestion_timestamps[1] < current_time_ms
+        assert ingestion_timestamps[1] > current_time_ms - 1000 * 5
+    else:
+        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"] == 0
-    assert type(group.meta["has_updates"]) == np.int64
+    if not is_type_erased_index(expected_index_type):
+        # NOTE(paris): Type-erased indexes do not write has_updates.
+        assert "has_updates" in group.meta
+        assert group.meta["has_updates"] == 0
+        assert type(group.meta["has_updates"]) == np.int64
 
 
 def test_flat_index(tmp_path):
@@ -213,6 +230,7 @@ def test_vamana_index(tmp_path):
     query_and_check_distances(
         index, queries, 1, [[ind.MAX_FLOAT_32]], [[ind.MAX_UINT64]]
     )
+    check_default_metadata(uri, vector_type, STORAGE_VERSION, "VAMANA")
 
     update_vectors = np.empty([5], dtype=object)
     update_vectors[0] = np.array([0, 0, 0], dtype=np.dtype(np.float32))

src/include/index/index_metadata.h

@@ -112,7 +112,7 @@ class base_index_metadata {
   using metadata_string_check_type =
       std::tuple<std::string, std::string&, bool>;
   std::vector<metadata_string_check_type> metadata_string_checks{
-      // name, member_variable, default, expected, required
+      // name, member_variable, required
       {"dataset_type", dataset_type_, true},
       {"storage_version", storage_version_, true},
       {"dtype", dtype_, false},
@@ -127,11 +127,8 @@ class base_index_metadata {
   using metadata_arithmetic_check_type =
       std::tuple<std::string, void*, tiledb_datatype_t, bool>;
   std::vector<metadata_arithmetic_check_type> metadata_arithmetic_checks{
+      // name, member_variable, type, required
       {"temp_size", &temp_size_, TILEDB_INT64, true},
-      //{"index_kind",
-      // nstatic_cast<IndexMetadata*>(this)->index_kind_,
-      // TILEDB_UINT64,
-      // false},
       {"dimension", &dimension_, TILEDB_UINT32, false},
       {"feature_datatype", &feature_datatype_, TILEDB_UINT32, false},
       {"id_datatype", &id_datatype_, TILEDB_UINT32, false},
@@ -339,54 +336,6 @@ class base_index_metadata {
     }
   }
 
-#if 0
-  /**************************************************************************
-   * Getters and setters
-   **************************************************************************/
-  std::string base_sizes_str() const {
-    return base_sizes_str_;
-  }
-  auto& base_sizes() {
-    return base_sizes_;
-  }
-  auto& base_sizes() const {
-    return base_sizes_;
-  }
-  void set_base_sizes(const std::vector<base_sizes_type>& base_sizes) {
-    base_sizes_ = base_sizes;
-  }
-  auto storage_version() const {
-    return storage_version_;
-  }
-  auto& storage_version()  {
-    return storage_version_;
-  }
-  auto dtype() const {
-    return dtype_;
-  }
-  auto& dtype()  {
-    return dtype_;
-  }
-  auto feature_datatype() const {
-    return feature_datatype_;
-  }
-  auto& feature_datatype()  {
-    return feature_datatype_;
-  }
-  auto id_datatype() const {
-    return id_datatype_;
-  }
-  auto& id_datatype()  {
-    return id_datatype_;
-  }
-  auto px_datatype() const {
-    return px_datatype_;
-  }
-  auto& px_datatype()  {
-    return px_datatype_;
-  }
-#endif
-
   /**************************************************************************
    * Helpful functions for debugging, testing, etc
    **************************************************************************/