Keep track of array sizes

Author: Nikos Papailiou

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

@@ -31,11 +31,15 @@ def __init__(
         schema = tiledb.ArraySchema.load(
             self.db_uri, ctx=tiledb.Ctx(self.config)
         )
-        self.size = schema.domain.dim(1).domain[1]+1
+        if self.base_size == -1:
+            self.size = schema.domain.dim(1).domain[1] + 1
+        else:
+            self.size = self.base_size
         self._db = load_as_matrix(
             self.db_uri,
             ctx=self.ctx,
             config=config,
+            size=self.size,
             timestamp=self.base_array_timestamp,
         )
         # Check for existence of ids array. Previous versions were not using external_ids in the ingestion assuming
@@ -44,8 +48,9 @@ def __init__(
             self.ids_uri = self.group[
                 storage_formats[self.storage_version]["IDS_ARRAY_NAME"] + self.index_version
             ].uri
-            self._ids = read_vector_u64(self.ctx, self.ids_uri, 0, 0, self.base_array_timestamp)
+            self._ids = read_vector_u64(self.ctx, self.ids_uri, 0, self.size, self.base_array_timestamp)
         else:
+            self.ids_uri = ""
             self._ids = StdVector_u64(np.arange(self.size).astype(np.uint64))
 
         dtype = self.group.meta.get("dtype", None)

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

@@ -46,7 +46,15 @@ def __init__(
         self.updates_array_uri = f"{self.group.uri}/{updates_array_name}"
         self.index_version = self.group.meta.get("index_version", "")
         self.ingestion_timestamps = list(json.loads(self.group.meta.get("ingestion_timestamps", "[]")))
-        self.latest_ingestion_timestamp = self.ingestion_timestamps[len(self.ingestion_timestamps)-1]
+        if len(self.ingestion_timestamps) > 0:
+            self.latest_ingestion_timestamp = self.ingestion_timestamps[len(self.ingestion_timestamps)-1]
+        else:
+            self.latest_ingestion_timestamp = MAX_UINT64
+        self.base_sizes = list(json.loads(self.group.meta.get("base_sizes", "[]")))
+        if len(self.base_sizes) > 0:
+            self.base_size = self.base_sizes[len(self.ingestion_timestamps)-1]
+        else:
+            self.base_size = -1
         self.base_array_timestamp = self.latest_ingestion_timestamp
         self.query_base_array = True
         self.update_array_timestamp = (self.base_array_timestamp+1, None)
@@ -59,15 +67,20 @@ def __init__(
                         self.query_base_array = False
                         self.update_array_timestamp = timestamp
                     else:
+                        self.base_size = self.base_sizes[0]
                         self.base_array_timestamp = self.ingestion_timestamps[0]
                         self.update_array_timestamp = (self.base_array_timestamp+1, timestamp[1])
                 else:
+                    self.base_size = self.base_sizes[0]
                     self.base_array_timestamp = self.ingestion_timestamps[0]
                     self.update_array_timestamp = (self.base_array_timestamp+1, timestamp[1])
             elif isinstance(timestamp, int):
+                i = 0
                 for ingestion_timestamp in self.ingestion_timestamps:
                     if ingestion_timestamp <= timestamp:
                         self.base_array_timestamp = ingestion_timestamp
+                        self.base_size = self.base_sizes[i]
+                    i += 1
                 self.update_array_timestamp = (self.base_array_timestamp+1, timestamp)
             else:
                 raise TypeError("Unexpected argument type for 'timestamp' keyword argument")
@@ -274,6 +287,7 @@ def consolidate_updates(self):
             size=self.size,
             source_uri=self.db_uri,
             external_ids_uri=self.ids_uri,
+            external_ids_type="TILEDB_ARRAY",
             updates_uri=self.updates_array_uri,
             index_timestamp=max_timestamp,
             config=self.config,

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

@@ -341,7 +341,7 @@ def create_arrays(
                 logger.debug("Creating ids array")
                 ids_array_rows_dim = tiledb.Dim(
                     name="rows",
-                    domain=(0, size - 1),
+                    domain=(0, MAX_INT32),
                     tile=int(size / partitions),
                     dtype=np.dtype(np.int32),
                 )
@@ -373,7 +373,7 @@ def create_arrays(
                 )
                 parts_array_cols_dim = tiledb.Dim(
                     name="cols",
-                    domain=(0, size - 1),
+                    domain=(0, MAX_INT32),
                     tile=int(size / partitions),
                     dtype=np.dtype(np.int32),
                 )
@@ -1784,6 +1784,7 @@ def consolidate_and_vacuum(
                 raise err
         group = tiledb.Group(index_group_uri, "r")
         ingestion_timestamps = list(json.loads(group.meta.get("ingestion_timestamps", "[]")))
+        base_sizes = list(json.loads(group.meta.get("base_sizes", "[]")))
         if partitions == -1:
             partitions = int(group.meta.get("partitions", "-1"))
 
@@ -1820,6 +1821,7 @@ def consolidate_and_vacuum(
             size = in_size
         if size > in_size:
             size = in_size
+        base_sizes.append(size)
         logger.debug("Input dataset size %d", size)
         logger.debug("Input dataset dimensions %d", dimensions)
         logger.debug("Vector dimension type %s", vector_type)
@@ -1840,6 +1842,7 @@ def consolidate_and_vacuum(
         group.meta["partitions"] = partitions
         group.meta["storage_version"] = STORAGE_VERSION
         group.meta["ingestion_timestamps"] = json.dumps(ingestion_timestamps)
+        group.meta["base_sizes"] = json.dumps(base_sizes)
 
         if external_ids is not None:
             external_ids_uri = write_external_ids(

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

@@ -74,8 +74,10 @@ def __init__(
             )
             self.partitions = schema.domain.dim("cols").domain[1] + 1
 
-        self.size = self._index[self.partitions]
-
+        if self.base_size == -1:
+            self.size = self._index[self.partitions]
+        else:
+            self.size = self.base_size
 
         # TODO pass in a context
         if self.memory_budget == -1:

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

@@ -310,7 +310,7 @@ static void declareColMajorMatrixSubclass(py::module& mod,
   // TODO auto-namify
   PyTMatrix cls(mod, (name + suffix).c_str(), py::buffer_protocol());
 
-  cls.def(py::init<const Ctx&, std::string, size_t, uint64_t>(),  py::keep_alive<1,2>());
+  cls.def(py::init<const Ctx&, std::string, size_t, size_t, size_t, size_t, uint64_t>(),  py::keep_alive<1,2>());
 
   if constexpr (std::is_same<P, tdbColMajorMatrix<T>>::value) {
     cls.def("load", &TMatrix::load);

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

@@ -37,15 +37,15 @@ def load_as_matrix(
     a = tiledb.ArraySchema.load(path, ctx=tiledb.Ctx(config))
     dtype = a.attr(0).dtype
     if dtype == np.float32:
-        m = tdbColMajorMatrix_f32(ctx, path, size, timestamp)
+        m = tdbColMajorMatrix_f32(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.float64:
-        m = tdbColMajorMatrix_f64(ctx, path, size, timestamp)
+        m = tdbColMajorMatrix_f64(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.int32:
-        m = tdbColMajorMatrix_i32(ctx, path, size, timestamp)
+        m = tdbColMajorMatrix_i32(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.int32:
-        m = tdbColMajorMatrix_i64(ctx, path, size, timestamp)
+        m = tdbColMajorMatrix_i64(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.uint8:
-        m = tdbColMajorMatrix_u8(ctx, path, size, timestamp)
+        m = tdbColMajorMatrix_u8(ctx, path, 0, 0, 0, size, timestamp)
     # elif dtype == np.uint64:
     #     return tdbColMajorMatrix_u64(ctx, path, size, timestamp)
     else:

apis/python/test/test_ingestion.py

@@ -370,6 +370,7 @@ def test_ivf_flat_ingestion_with_batch_updates(tmp_path):
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i, updated_ids=updated_ids) > 0.99
 
+
 def test_ivf_flat_ingestion_with_updates_and_timetravel(tmp_path):
     dataset_dir = os.path.join(tmp_path, "dataset")
     index_uri = os.path.join(tmp_path, "array")
@@ -463,3 +464,42 @@ def test_ivf_flat_ingestion_with_updates_and_timetravel(tmp_path):
     index = IVFFlatIndex(uri=index_uri, timestamp=(0, 1))
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) == 1.0
+
+
+def test_ivf_flat_ingestion_with_additions_and_timetravel(tmp_path):
+    dataset_dir = os.path.join(tmp_path, "dataset")
+    index_uri = os.path.join(tmp_path, "array")
+    k = 100
+    size = 100
+    partitions = 1
+    dimensions = 128
+    nqueries = 1
+    nprobe = 1
+    data = create_random_dataset_u8(nb=size, d=dimensions, nq=nqueries, k=k, path=dataset_dir)
+    dtype = np.uint8
+
+    query_vectors = get_queries(dataset_dir, dtype=dtype)
+    gt_i, gt_d = get_groundtruth(dataset_dir, k)
+    index = ingest(
+        index_type="IVF_FLAT",
+        index_uri=index_uri,
+        source_uri=os.path.join(dataset_dir, "data.u8bin"),
+        partitions=partitions,
+        index_timestamp=1,
+    )
+    _, result = index.query(query_vectors, k=k)
+    assert accuracy(result, gt_i) == 1.0
+
+    update_ids_offset = MAX_UINT64-size
+    updated_ids = {}
+    for i in range(100):
+        index.update(vector=data[i].astype(dtype), external_id=i + update_ids_offset, timestamp=i+2)
+        updated_ids[i] = i + update_ids_offset
+
+    index = IVFFlatIndex(uri=index_uri)
+    _, result = index.query(query_vectors, k=k)
+    assert 0.45 < accuracy(result, gt_i) < 0.55
+
+    index = index.consolidate_updates()
+    _, result = index.query(query_vectors, k=k)
+    assert 0.45 < accuracy(result, gt_i) < 0.55

src/include/detail/ivf/dist_qv.h

@@ -80,7 +80,7 @@ auto dist_qv_finite_ram_part(
   if (nthreads == 0) {
     nthreads = std::thread::hardware_concurrency();
   }
-  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
 
   using score_type = float;
   using parts_type =

src/include/detail/ivf/index.h

@@ -69,7 +69,7 @@ int ivf_index(
   if (nthreads == 0) {
     nthreads = std::thread::hardware_concurrency();
   }
-  auto read_temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto read_temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
   auto write_temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
   auto centroids = tdbColMajorMatrix<centroids_type>(ctx, centroids_uri, 0, read_temporal_policy);
   centroids.load();
@@ -191,14 +191,14 @@ int ivf_index(
     size_t end_pos = 0,
     size_t nthreads = 0,
     uint64_t timestamp = 0) {
-  auto db = tdbColMajorMatrix<T>(ctx, db_uri, 0, 0, start_pos, end_pos);
+  auto db = tdbColMajorMatrix<T>(ctx, db_uri, 0, 0, start_pos, end_pos, timestamp);
   db.load();
   std::vector<ids_type> external_ids;
   if (external_ids_uri.empty()) {
     external_ids = std::vector<ids_type>(db.num_cols());
     std::iota(begin(external_ids), end(external_ids), start_pos);
   } else {
-    auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+    auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
     external_ids =
         read_vector<ids_type>(ctx, external_ids_uri, start_pos, end_pos, temporal_policy);
   }
@@ -231,7 +231,7 @@ int ivf_index(
     size_t end_pos = 0,
     size_t nthreads = 0,
     uint64_t timestamp = 0) {
-  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
   auto db = tdbColMajorMatrix<T>(ctx, db_uri, 0, 0, start_pos, end_pos, temporal_policy);
   db.load();
   return ivf_index<T, ids_type, centroids_type>(
@@ -268,7 +268,7 @@ int ivf_index(
     external_ids = std::vector<ids_type>(db.num_cols());
     std::iota(begin(external_ids), end(external_ids), start_pos);
   } else {
-    auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+    auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
     external_ids = read_vector<ids_type>(ctx, external_ids_uri, start_pos, end_pos, temporal_policy);
   }
   return ivf_index<T, ids_type, centroids_type>(

src/include/detail/ivf/qv.h

@@ -422,7 +422,7 @@ auto nuv_query_heap_infinite_ram_reg_blocked(
     size_t nthreads,
     uint64_t timestamp = 0) {
   scoped_timer _{tdb_func__};
-  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
 
   // Read the shuffled database and ids
   // @todo To this more systematically
@@ -647,7 +647,7 @@ auto qv_query_heap_finite_ram(
     size_t nthreads,
     uint64_t timestamp = 0) {
   tiledb::Context ctx;
-  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
 
   auto centroids = tdbColMajorMatrix<centroids_type>(ctx, centroids_uri, 0, temporal_policy);
   centroids.load();
@@ -722,7 +722,7 @@ auto qv_query_heap_finite_ram(
     size_t nthreads,
     uint64_t timestamp) {
   scoped_timer _{tdb_func__};
-  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
 
   using score_type = float;
   using indices_type =
@@ -1194,7 +1194,7 @@ auto nuv_query_heap_finite_ram_reg_blocked(
     size_t nthreads,
     uint64_t timestamp = 0) {
   scoped_timer _{tdb_func__ + " " + part_uri};
-  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp+1);
+  auto temporal_policy = (timestamp == 0) ? tiledb::TemporalPolicy() : tiledb::TemporalPolicy(tiledb::TimeTravel, timestamp);
 
   // Check that the size of the indices vector is correct
   assert(size(indices) == centroids.num_cols() + 1);