Share max int / float values and index lists in Python (#396)

Author: jparismorgan

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

@@ -13,10 +13,11 @@
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
 from tiledb.vector_search.storage_formats import storage_formats
 from tiledb.vector_search.storage_formats import validate_storage_version
+from tiledb.vector_search.utils import MAX_FLOAT32
+from tiledb.vector_search.utils import MAX_INT32
+from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import add_to_group
 
-MAX_INT32 = np.iinfo(np.dtype("int32")).max
-MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 TILE_SIZE_BYTES = 128000000  # 128MB
 INDEX_TYPE = "FLAT"
 
@@ -116,8 +117,8 @@ def query_internal(
         # - typecheck queries
         # - add all the options and query strategies
         if self.size == 0:
-            return np.full((queries.shape[0], k), index.MAX_FLOAT_32), np.full(
-                (queries.shape[0], k), index.MAX_UINT64
+            return np.full((queries.shape[0], k), MAX_FLOAT32), np.full(
+                (queries.shape[0], k), MAX_UINT64
             )
 
         assert queries.dtype == np.float32

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

@@ -7,11 +7,10 @@
 from tiledb.vector_search import _tiledbvspy as vspy
 from tiledb.vector_search.module import *
 from tiledb.vector_search.storage_formats import storage_formats
+from tiledb.vector_search.utils import MAX_FLOAT32
+from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import add_to_group
 
-MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
-MAX_INT32 = np.iinfo(np.dtype("int32")).max
-MAX_FLOAT_32 = np.finfo(np.dtype("float32")).max
 DATASET_TYPE = "vector_search"
 
 
@@ -190,7 +189,7 @@ def query(self, queries: np.ndarray, k: int, **kwargs):
                 if self.query_base_array:
                     return self.query_internal(queries, k, **kwargs)
                 else:
-                    return np.full((queries.shape[0], k), MAX_FLOAT_32), np.full(
+                    return np.full((queries.shape[0], k), MAX_FLOAT32), np.full(
                         (queries.shape[0], k), MAX_UINT64
                     )
 
@@ -213,7 +212,7 @@ def query(self, queries: np.ndarray, k: int, **kwargs):
                 queries, retrieval_k, **kwargs
             )
         else:
-            internal_results_d = np.full((queries.shape[0], k), MAX_FLOAT_32)
+            internal_results_d = np.full((queries.shape[0], k), MAX_FLOAT32)
             internal_results_i = np.full((queries.shape[0], k), MAX_UINT64)
         addition_results_d, addition_results_i, updated_ids = future.result()
 
@@ -223,7 +222,7 @@ def query(self, queries: np.ndarray, k: int, **kwargs):
             res_id = 0
             for res in query:
                 if res in updated_ids:
-                    internal_results_d[query_id, res_id] = MAX_FLOAT_32
+                    internal_results_d[query_id, res_id] = MAX_FLOAT32
                     internal_results_i[query_id, res_id] = MAX_UINT64
                 res_id += 1
             query_id += 1
@@ -243,7 +242,7 @@ def query(self, queries: np.ndarray, k: int, **kwargs):
                     addition_results_d[query_id, res_id] == 0
                     and addition_results_i[query_id, res_id] == 0
                 ):
-                    addition_results_d[query_id, res_id] = MAX_FLOAT_32
+                    addition_results_d[query_id, res_id] = MAX_FLOAT32
                     addition_results_i[query_id, res_id] = MAX_UINT64
                 res_id += 1
             query_id += 1

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

@@ -35,10 +35,11 @@
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
 from tiledb.vector_search.storage_formats import storage_formats
 from tiledb.vector_search.storage_formats import validate_storage_version
+from tiledb.vector_search.utils import MAX_FLOAT32
+from tiledb.vector_search.utils import MAX_INT32
+from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import add_to_group
 
-MAX_INT32 = np.iinfo(np.dtype("int32")).max
-MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 TILE_SIZE_BYTES = 64000000  # 64MB
 INDEX_TYPE = "IVF_FLAT"
 
@@ -215,8 +216,8 @@ def query_internal(
             If provided, this is the number of workers to use for the query execution.
         """
         if self.size == 0:
-            return np.full((queries.shape[0], k), index.MAX_FLOAT_32), np.full(
-                (queries.shape[0], k), index.MAX_UINT64
+            return np.full((queries.shape[0], k), MAX_FLOAT32), np.full(
+                (queries.shape[0], k), MAX_UINT64
             )
 
         if mode != Mode.BATCH and resources:

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

@@ -7,6 +7,10 @@
 import tiledb
 from tiledb.vector_search import _tiledbvspy as vspy
 
+MAX_INT32 = np.iinfo(np.dtype("int32")).max
+MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
+MAX_FLOAT32 = np.finfo(np.dtype("float32")).max
+
 
 def is_type_erased_index(index_type: str) -> bool:
     return index_type == "VAMANA"

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

@@ -21,9 +21,10 @@
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
 from tiledb.vector_search.storage_formats import storage_formats
 from tiledb.vector_search.storage_formats import validate_storage_version
+from tiledb.vector_search.utils import MAX_FLOAT32
+from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import to_temporal_policy
 
-MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 INDEX_TYPE = "VAMANA"
 
 
@@ -97,8 +98,8 @@ def query_internal(
         """
         warnings.warn("The Vamana index is not yet supported, please use with caution.")
         if self.size == 0:
-            return np.full((queries.shape[0], k), index.MAX_FLOAT_32), np.full(
-                (queries.shape[0], k), index.MAX_UINT64
+            return np.full((queries.shape[0], k), MAX_FLOAT32), np.full(
+                (queries.shape[0], k), MAX_UINT64
             )
 
         assert queries.dtype == np.float32

apis/python/test/common.py

@@ -8,8 +8,19 @@
 
 import tiledb
 from tiledb.cloud import groups
+from tiledb.vector_search.flat_index import FlatIndex
+from tiledb.vector_search.ivf_flat_index import IVFFlatIndex
 from tiledb.vector_search.storage_formats import STORAGE_VERSION
 from tiledb.vector_search.storage_formats import storage_formats
+from tiledb.vector_search.vamana_index import VamanaIndex
+
+INDEXES = ["FLAT", "IVF_FLAT", "VAMANA"]
+INDEX_CLASSES = [FlatIndex, IVFFlatIndex, VamanaIndex]
+INDEX_FILES = [
+    tiledb.vector_search.flat_index,
+    tiledb.vector_search.ivf_flat_index,
+    tiledb.vector_search.vamana_index,
+]
 
 
 def xbin_mmap(fname, dtype):

apis/python/test/test_index.py

@@ -7,7 +7,6 @@
 from common import *
 from common import load_metadata
 
-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
@@ -17,6 +16,8 @@
 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 MAX_FLOAT32
+from tiledb.vector_search.utils import MAX_UINT64
 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
@@ -78,7 +79,7 @@ def test_flat_index(tmp_path):
     uri = os.path.join(tmp_path, "array")
     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})
+    query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {MAX_UINT64})
     check_default_metadata(uri, vector_type, STORAGE_VERSION, "FLAT")
 
     update_vectors = np.empty([5], dtype=object)
@@ -136,7 +137,7 @@ def test_ivf_flat_index(tmp_path):
         index,
         np.array([[2, 2, 2]], dtype=np.float32),
         3,
-        {ind.MAX_UINT64},
+        {MAX_UINT64},
         nprobe=partitions,
     )
     check_default_metadata(uri, vector_type, STORAGE_VERSION, "IVF_FLAT")
@@ -221,12 +222,12 @@ def test_vamana_index_simple(tmp_path):
     # Create the index.
     index = vamana_index.create(uri=uri, dimensions=dimensions, vector_type=vector_type)
     assert index.get_dimensions() == dimensions
-    query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {ind.MAX_UINT64})
+    query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {MAX_UINT64})
 
     # Open the index.
     index = VamanaIndex(uri=uri)
     assert index.get_dimensions() == dimensions
-    query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {ind.MAX_UINT64})
+    query_and_check(index, np.array([[2, 2, 2]], dtype=np.float32), 3, {MAX_UINT64})
 
     vfs = tiledb.VFS()
     assert vfs.dir_size(uri) > 0
@@ -254,11 +255,9 @@ def test_vamana_index(tmp_path):
     distances, ids = index.query(queries, k=1)
     assert distances.shape == (1, 1)
     assert ids.shape == (1, 1)
-    assert distances[0][0] == ind.MAX_FLOAT_32
-    assert ids[0][0] == ind.MAX_UINT64
-    query_and_check_distances(
-        index, queries, 1, [[ind.MAX_FLOAT_32]], [[ind.MAX_UINT64]]
-    )
+    assert distances[0][0] == MAX_FLOAT32
+    assert ids[0][0] == MAX_UINT64
+    query_and_check_distances(index, queries, 1, [[MAX_FLOAT32]], [[MAX_UINT64]])
     check_default_metadata(uri, vector_type, STORAGE_VERSION, "VAMANA")
 
     update_vectors = np.empty([5], dtype=object)

apis/python/test/test_ingestion.py

@@ -15,21 +15,13 @@
 from tiledb.vector_search.module import array_to_matrix
 from tiledb.vector_search.module import kmeans_fit
 from tiledb.vector_search.module import kmeans_predict
+from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import is_type_erased_index
 from tiledb.vector_search.utils import load_fvecs
 from tiledb.vector_search.utils import metadata_to_list
 from tiledb.vector_search.vamana_index import VamanaIndex
 
 MINIMUM_ACCURACY = 0.85
-MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
-
-INDEXES = ["FLAT", "IVF_FLAT", "VAMANA"]
-INDEX_CLASSES = [FlatIndex, IVFFlatIndex, VamanaIndex]
-INDEX_FILES = [
-    tiledb.vector_search.flat_index,
-    tiledb.vector_search.ivf_flat_index,
-    tiledb.vector_search.vamana_index,
-]
 
 
 def query_and_check_equals(index, queries, expected_result_d, expected_result_i):

apis/python/test/test_object_index.py

@@ -11,7 +11,6 @@
 from tiledb.vector_search.object_readers import ObjectReader
 
 EMBED_DIM = 4
-INDEXES = ["FLAT", "IVF_FLAT", "VAMANA"]
 
 
 # TestEmbedding with vectors of EMBED_DIM size with all values being the id of the object
@@ -236,6 +235,8 @@ def df_filter(row):
 
 
 def test_object_index(tmp_path):
+    from common import INDEXES
+
     for index_type in INDEXES:
         index_uri = os.path.join(tmp_path, f"object_index_{index_type}")
         reader = TestReader(

src/include/test/unit_ivf_pq_index.cc

@@ -288,8 +288,12 @@ TEST_CASE("ivf_index write and read", "[ivf_index]") {
   size_t nthreads = 1;
 
   tiledb::Context ctx;
+  tiledb::VFS vfs(ctx);
   std::string ivf_index_uri =
       (std::filesystem::temp_directory_path() / "tmp_ivf_index").string();
+  if (vfs.is_dir(ivf_index_uri)) {
+    vfs.remove_dir(ivf_index_uri);
+  }
   auto training_set = tdbColMajorMatrix<float>(ctx, siftsmall_inputs_uri, 0);
   load(training_set);
 
@@ -302,6 +306,9 @@ TEST_CASE("ivf_index write and read", "[ivf_index]") {
   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);
   auto idx2 = ivf_pq_index<float, uint32_t, uint32_t>(ctx, ivf_index_uri);
   idx2.read_index_infinite();