Merge pull request #143 from TileDB-Inc/npapa/fix-uint64-ids

Fix query errors in tables with updates

Author: NikolaosPapailiou

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

@@ -70,6 +70,12 @@ def query(self, queries: np.ndarray, k, **kwargs):
                 if res in updated_ids:
                     internal_results_d[query_id, res_id] = MAX_FLOAT_32
                     internal_results_i[query_id, res_id] = MAX_UINT64
+                if (
+                    internal_results_d[query_id, res_id] == 0
+                    and internal_results_i[query_id, res_id] == 0
+                ):
+                    internal_results_d[query_id, res_id] = MAX_FLOAT_32
+                    internal_results_i[query_id, res_id] = MAX_UINT64
                 res_id += 1
             query_id += 1
         sort_index = np.argsort(internal_results_d, axis=1)

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

@@ -695,11 +695,11 @@ def read_additions(
         logger.debug(
             "Reading additions vectors"
         )
-        updates_array = tiledb.open(updates_uri, mode="r")
-        q = updates_array.query(attrs=('vector',), coords=True)
-        data = q[:]
-        additions_filter = [len(item) > 0 for item in data["vector"]]
-        return np.vstack(data["vector"][additions_filter]), data["external_id"][additions_filter]
+        with tiledb.open(updates_uri, mode="r") as updates_array:
+            q = updates_array.query(attrs=('vector',), coords=True)
+            data = q[:]
+            additions_filter = [len(item) > 0 for item in data["vector"]]
+            return np.vstack(data["vector"][additions_filter]), data["external_id"][additions_filter]
 
     def read_updated_ids(
         updates_uri: str,
@@ -713,10 +713,10 @@ def read_updated_ids(
         logger.debug(
             "Reading updated vector ids"
         )
-        updates_array = tiledb.open(updates_uri, mode="r")
-        q = updates_array.query(attrs=('vector',), coords=True)
-        data = q[:]
-        return data["external_id"]
+        with tiledb.open(updates_uri, mode="r") as updates_array:
+            q = updates_array.query(attrs=('vector',), coords=True)
+            data = q[:]
+            return data["external_id"]
 
     def read_input_vectors(
         source_uri: str,
@@ -1729,10 +1729,15 @@ def consolidate_and_vacuum(
         config: Optional[Mapping[str, Any]] = None,
     ):
         group = tiledb.Group(index_group_uri)
-        if INPUT_VECTORS_ARRAY_NAME in group:
-            tiledb.Array.delete_array(group[INPUT_VECTORS_ARRAY_NAME].uri)
-        if EXTERNAL_IDS_ARRAY_NAME in group:
-            tiledb.Array.delete_array(group[EXTERNAL_IDS_ARRAY_NAME].uri)
+        try:
+            if INPUT_VECTORS_ARRAY_NAME in group:
+                tiledb.Array.delete_array(group[INPUT_VECTORS_ARRAY_NAME].uri)
+            if EXTERNAL_IDS_ARRAY_NAME in group:
+                tiledb.Array.delete_array(group[EXTERNAL_IDS_ARRAY_NAME].uri)
+        except tiledb.TileDBError as err:
+            message = str(err)
+            if "does not exist" not in message:
+                raise err
         modes = ["fragment_meta", "commits", "array_meta"]
         for mode in modes:
             conf = tiledb.Config(config)

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

@@ -402,7 +402,7 @@ static void declare_vq_query_heap(py::module& m, const std::string& suffix) {
            const std::vector<uint64_t> &ids,
            int k,
            size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
-          auto r = detail::flat::vq_query_heap(data, query_vectors, ids, k, nthreads);
+          auto r = detail::flat::vq_query_heap<tdbColMajorMatrix<T>, ColMajorMatrix<float>, uint64_t>(data, query_vectors, ids, k, nthreads);
           return r;
         });
 }
@@ -415,7 +415,7 @@ static void declare_vq_query_heap_pyarray(py::module& m, const std::string& suff
            const std::vector<uint64_t> &ids,
            int k,
            size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
-          auto r = detail::flat::vq_query_heap(data, query_vectors, ids, k, nthreads);
+          auto r = detail::flat::vq_query_heap<ColMajorMatrix<T>, ColMajorMatrix<float>, uint64_t>(data, query_vectors, ids, k, nthreads);
           return r;
         });
 }

apis/python/test/test_ingestion.py

@@ -9,6 +9,7 @@
 from tiledb.cloud.dag import Mode
 
 MINIMUM_ACCURACY = 0.85
+MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
 
 
 def test_flat_ingestion_u8(tmp_path):
@@ -307,11 +308,12 @@ def test_ivf_flat_ingestion_with_updates(tmp_path):
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
+    update_ids_offset = MAX_UINT64-size
     updated_ids = {}
     for i in range(100):
         index.delete(external_id=i)
-        index.update(vector=data[i].astype(dtype), external_id=i + 1000000)
-        updated_ids[i + 1000000] = i
+        index.update(vector=data[i].astype(dtype), external_id=i + update_ids_offset)
+        updated_ids[i + update_ids_offset] = i
 
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i, updated_ids=updated_ids) > MINIMUM_ACCURACY
@@ -346,9 +348,10 @@ def test_ivf_flat_ingestion_with_batch_updates(tmp_path):
 
     update_ids = {}
     updated_ids = {}
+    update_ids_offset = MAX_UINT64 - size
     for i in range(0, 100000, 2):
-        update_ids[i] = i + 1000000
-        updated_ids[i + 1000000] = i
+        update_ids[i] = i + update_ids_offset
+        updated_ids[i + update_ids_offset] = i
     external_ids = np.zeros((len(update_ids) * 2), dtype=np.uint64)
     updates = np.empty((len(update_ids) * 2), dtype='O')
     id = 0

src/include/detail/flat/vq.h

@@ -83,10 +83,10 @@ auto vq_query_heap(
     unsigned nthreads) {
   // @todo Need to get the total number of queries, not just the first block
   // @todo Use Matrix here rather than vector of vectors
-  std::vector<std::vector<fixed_min_pair_heap<float, unsigned>>> scores(
+  std::vector<std::vector<fixed_min_pair_heap<float, Index>>> scores(
       nthreads,
-      std::vector<fixed_min_pair_heap<float, unsigned>>(
-          size(q), fixed_min_pair_heap<float, unsigned>(k_nn)));
+      std::vector<fixed_min_pair_heap<float, Index>>(
+          size(q), fixed_min_pair_heap<float, Index>(k_nn)));
 
   unsigned size_q = size(q);
   auto par = stdx::execution::indexed_parallel_policy{nthreads};
@@ -184,10 +184,10 @@ auto vq_query_heap_tiled(
     unsigned nthreads) {
   // @todo Need to get the total number of queries, not just the first block
   // @todo Use Matrix here rather than vector of vectors
-  std::vector<std::vector<fixed_min_pair_heap<float, unsigned>>> scores(
+  std::vector<std::vector<fixed_min_pair_heap<float, Index>>> scores(
       nthreads,
-      std::vector<fixed_min_pair_heap<float, unsigned>>(
-          size(q), fixed_min_pair_heap<float, unsigned>(k_nn)));
+      std::vector<fixed_min_pair_heap<float, Index>>(
+          size(q), fixed_min_pair_heap<float, Index>(k_nn)));
 
   unsigned size_q = size(q);
   auto par = stdx::execution::indexed_parallel_policy{nthreads};
@@ -261,10 +261,10 @@ auto vq_query_heap_2(
     unsigned nthreads) {
   // @todo Need to get the total number of queries, not just the first block
   // @todo Use Matrix here rather than vector of vectors
-  std::vector<std::vector<fixed_min_pair_heap<float, size_t>>> scores(
+  std::vector<std::vector<fixed_min_pair_heap<float, Index>>> scores(
       nthreads,
-      std::vector<fixed_min_pair_heap<float, size_t>>(
-          size(q), fixed_min_pair_heap<float, size_t>(k_nn)));
+      std::vector<fixed_min_pair_heap<float, Index>>(
+          size(q), fixed_min_pair_heap<float, Index>(k_nn)));
 
   unsigned size_q = size(q);
   auto par = stdx::execution::indexed_parallel_policy{nthreads};