Timetravel implementation

Author: Nikos Papailiou

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

@@ -29,7 +29,7 @@ def __init__(
         self,
         uri: str,
         config: Optional[Mapping[str, Any]] = None,
-        timestamp: int = None,
+        timestamp=None,
     ):
         # If the user passes a tiledb python Config object convert to a dictionary
         if isinstance(config, tiledb.Config):
@@ -40,13 +40,14 @@ def __init__(
         self.ctx = Ctx(config)
         self.group = tiledb.Group(self.uri, "r", ctx=tiledb.Ctx(config))
         self.storage_version = self.group.meta.get("storage_version", "0.1")
-        self.update_arrays_uri = None
+        updates_array_name = storage_formats[self.storage_version][
+            "UPDATES_ARRAY_NAME"
+        ]
+        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", "[]")))
-        print(f"ingestion_timestamps: {self.ingestion_timestamps}")
-        self.base_array_timestamp = self.ingestion_timestamps[len(self.ingestion_timestamps)-1]
-        print(f"base_array_timestamp: {self.base_array_timestamp}")
+        self.latest_ingestion_timestamp = self.ingestion_timestamps[len(self.ingestion_timestamps)-1]
+        self.base_array_timestamp = self.latest_ingestion_timestamp
         self.query_base_array = True
         self.update_array_timestamp = (self.base_array_timestamp+1, None)
         if timestamp is not None:
@@ -70,13 +71,14 @@ def __init__(
                 self.update_array_timestamp = (self.base_array_timestamp+1, timestamp)
             else:
                 raise TypeError("Unexpected argument type for 'timestamp' keyword argument")
-        print(f"base_array_timestamp: {self.base_array_timestamp}")
-        print(f"update_array_timestamp: {self.update_array_timestamp}")
         self.thread_executor = futures.ThreadPoolExecutor()
 
     def query(self, queries: np.ndarray, k, **kwargs):
-        if self.update_arrays_uri is None:
-            return self.query_internal(queries, k, **kwargs)
+        if not tiledb.array_exists(self.updates_array_uri):
+            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((queries.shape[0], k), MAX_UINT64)
 
         # Query with updates
         # Perform the queries in parallel
@@ -87,13 +89,17 @@ def query(self, queries: np.ndarray, k, **kwargs):
             queries,
             k,
             self.dtype,
-            self.update_arrays_uri,
+            self.updates_array_uri,
             int(os.cpu_count() / 2),
             self.update_array_timestamp,
         )
-        internal_results_d, internal_results_i = self.query_internal(
-            queries, retrieval_k, **kwargs
-        )
+        if self.query_base_array:
+            internal_results_d, internal_results_i = self.query_internal(
+                queries, retrieval_k, **kwargs
+            )
+        else:
+            internal_results_d = np.full((queries.shape[0], k), MAX_FLOAT_32)
+            internal_results_i = np.full((queries.shape[0], k), MAX_UINT64)
         addition_results_d, addition_results_i, updated_ids = future.result()
 
         # Filter updated vectors
@@ -142,11 +148,11 @@ def query(self, queries: np.ndarray, k, **kwargs):
 
     @staticmethod
     def query_additions(
-        queries: np.ndarray, k, dtype, update_arrays_uri, nthreads=8, timestamp=None
+        queries: np.ndarray, k, dtype, updates_array_uri, nthreads=8, timestamp=None
     ):
         assert queries.dtype == np.float32
         additions_vectors, additions_external_ids, updated_ids = Index.read_additions(
-            update_arrays_uri, timestamp
+            updates_array_uri, timestamp
         )
         if additions_vectors is None:
             return None, None, updated_ids
@@ -162,10 +168,10 @@ def query_additions(
         return np.transpose(np.array(d)), np.transpose(np.array(i)), updated_ids
 
     @staticmethod
-    def read_additions(update_arrays_uri, timestamp=None) -> (np.ndarray, np.array):
-        if update_arrays_uri is None:
+    def read_additions(updates_array_uri, timestamp=None) -> (np.ndarray, np.array):
+        if updates_array_uri is None:
             return None, None, np.array([], np.uint64)
-        updates_array = tiledb.open(update_arrays_uri, mode="r", timestamp=timestamp)
+        updates_array = tiledb.open(updates_array_uri, mode="r", timestamp=timestamp)
         q = updates_array.query(attrs=("vector",), coords=True)
         data = q[:]
         updates_array.close()
@@ -215,22 +221,22 @@ def delete_batch(self, external_ids: np.array, timestamp: int = None):
         self.consolidate_update_fragments()
 
     def consolidate_update_fragments(self):
-        fragments_info = tiledb.array_fragments(self.update_arrays_uri)
+        fragments_info = tiledb.array_fragments(self.updates_array_uri)
         if len(fragments_info) > 10:
-            tiledb.consolidate(self.update_arrays_uri)
-            tiledb.vacuum(self.update_arrays_uri)
+            tiledb.consolidate(self.updates_array_uri)
+            tiledb.vacuum(self.updates_array_uri)
 
     def get_updates_uri(self):
-        return self.update_arrays_uri
+        return self.updates_array_uri
 
     def open_updates_array(self, timestamp: int = None):
-        if self.update_arrays_uri is None:
+        if timestamp is not None and timestamp <= self.latest_ingestion_timestamp:
+            raise ValueError(f"Updates at a timestamp before the latest_ingestion_timestamp are not supported. "
+                             f"timestamp: {timestamp}, latest_ingestion_timestamp: {self.latest_ingestion_timestamp}")
+        if not tiledb.array_exists(self.updates_array_uri):
             updates_array_name = storage_formats[self.storage_version][
                 "UPDATES_ARRAY_NAME"
             ]
-            updates_array_uri = f"{self.group.uri}/{updates_array_name}"
-            if tiledb.array_exists(updates_array_uri):
-                raise RuntimeError(f"Array {updates_array_uri} already exists.")
             external_id_dim = tiledb.Dim(
                 name="external_id",
                 domain=(0, MAX_UINT64 - 1),
@@ -244,27 +250,32 @@ def open_updates_array(self, timestamp: int = None):
                 attrs=[vector_attr],
                 allows_duplicates=False,
             )
-            tiledb.Array.create(updates_array_uri, updates_schema)
+            tiledb.Array.create(self.updates_array_uri, updates_schema)
             self.group.close()
             self.group = tiledb.Group(self.uri, "w", ctx=tiledb.Ctx(self.config))
-            self.group.add(updates_array_uri, name=updates_array_name)
+            self.group.add(self.updates_array_uri, name=updates_array_name)
             self.group.close()
             self.group = tiledb.Group(self.uri, "r", ctx=tiledb.Ctx(self.config))
-            self.update_arrays_uri = updates_array_uri
         if timestamp is None:
             timestamp = int(time.time() * 1000)
-        return tiledb.open(self.update_arrays_uri, mode="w", timestamp=timestamp)
+        return tiledb.open(self.updates_array_uri, mode="w", timestamp=timestamp)
 
     def consolidate_updates(self):
         from tiledb.vector_search.ingestion import ingest
 
+        fragments_info = tiledb.array_fragments(self.updates_array_uri, ctx=tiledb.Ctx(self.config))
+        max_timestamp = self.base_array_timestamp
+        for fragment_info in fragments_info:
+            if fragment_info.timestamp_range[1] > max_timestamp:
+                max_timestamp = fragment_info.timestamp_range[1]
         new_index = ingest(
             index_type=self.index_type,
             index_uri=self.uri,
             size=self.size,
             source_uri=self.db_uri,
             external_ids_uri=self.ids_uri,
-            updates_uri=self.update_arrays_uri,
+            updates_uri=self.updates_array_uri,
+            index_timestamp=max_timestamp,
+            config=self.config,
         )
-        new_index.update_arrays_uri = self.update_arrays_uri
         return new_index

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

@@ -1785,8 +1785,14 @@ def consolidate_and_vacuum(
         ingestion_timestamps = list(json.loads(group.meta.get("ingestion_timestamps", "[]")))
         if partitions == -1:
             partitions = int(group.meta.get("partitions", "-1"))
+
+        if len(ingestion_timestamps) > 0:
+            previous_ingestion_timestamp = ingestion_timestamps[len(ingestion_timestamps)-1]
+            if index_timestamp <= previous_ingestion_timestamp:
+                raise ValueError(f"New ingestion timestamp: {index_timestamp} can't be smaller that the latest ingestion "
+                                 f"timestamp: {previous_ingestion_timestamp}")
+
         ingestion_timestamps.append(index_timestamp)
-        print(f"ingestion_timestamps: {ingestion_timestamps}")
         group.close()
         group = tiledb.Group(index_group_uri, "w")
 

apis/python/test/common.py

@@ -162,9 +162,13 @@ def create_array(path: str, data):
         A[:] = data
 
 
-def accuracy(result, gt, external_ids_offset=0, updated_ids=None):
+def accuracy(result, gt, external_ids_offset=0, updated_ids=None, only_updated_ids=False):
     found = 0
     total = 0
+    if updated_ids is not None:
+        updated_ids_rev = {}
+        for updated_id in updated_ids:
+            updated_ids_rev[updated_ids[updated_id]] = updated_id
     for i in range(len(result)):
         if external_ids_offset != 0:
             temp_result = []
@@ -173,7 +177,12 @@ def accuracy(result, gt, external_ids_offset=0, updated_ids=None):
         elif updated_ids is not None:
             temp_result = []
             for j in range(len(result[i])):
-                uid = updated_ids.get(result[i][j])
+                if result[i][j] in updated_ids:
+                    raise ValueError(f"Found updated id {result[i][j]} in query results.")
+                if only_updated_ids:
+                    if result[i][j] not in updated_ids_rev:
+                        raise ValueError(f"Found not_updated_id {result[i][j]} in query results while expecting only_updated_ids.")
+                uid = updated_ids_rev.get(result[i][j])
                 if uid is not None:
                     temp_result.append(int(uid))
                 else:

apis/python/test/test_ingestion.py

@@ -288,11 +288,11 @@ def test_ivf_flat_ingestion_with_updates(tmp_path):
     dataset_dir = os.path.join(tmp_path, "dataset")
     index_uri = os.path.join(tmp_path, "array")
     k = 10
-    size = 100000
-    partitions = 100
+    size = 1000
+    partitions = 10
     dimensions = 128
     nqueries = 100
-    nprobe = 20
+    nprobe = 10
     data = create_random_dataset_u8(nb=size, d=dimensions, nq=nqueries, k=k, path=dataset_dir)
     dtype = np.uint8
 
@@ -303,24 +303,23 @@ def test_ivf_flat_ingestion_with_updates(tmp_path):
         index_uri=index_uri,
         source_uri=os.path.join(dataset_dir, "data.u8bin"),
         partitions=partitions,
-        input_vectors_per_work_item=int(size / 10),
     )
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
-    assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+    assert accuracy(result, gt_i) == 1.0
 
     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 + update_ids_offset)
-        updated_ids[i + update_ids_offset] = i
+        updated_ids[i] = i + update_ids_offset
 
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
-    assert accuracy(result, gt_i, updated_ids=updated_ids) > MINIMUM_ACCURACY
+    assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
 
     index = index.consolidate_updates()
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
-    assert accuracy(result, gt_i, updated_ids=updated_ids) > MINIMUM_ACCURACY
+    assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
 
 def test_ivf_flat_ingestion_with_batch_updates(tmp_path):
     dataset_dir = os.path.join(tmp_path, "dataset")
@@ -330,7 +329,7 @@ def test_ivf_flat_ingestion_with_batch_updates(tmp_path):
     partitions = 100
     dimensions = 128
     nqueries = 100
-    nprobe = 20
+    nprobe = 100
     data = create_random_dataset_u8(nb=size, d=dimensions, nq=nqueries, k=k, path=dataset_dir)
     dtype = np.uint8
 
@@ -344,18 +343,18 @@ def test_ivf_flat_ingestion_with_batch_updates(tmp_path):
         input_vectors_per_work_item=int(size / 10),
     )
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
-    assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+    assert accuracy(result, gt_i) > 0.99
 
     update_ids = {}
     updated_ids = {}
     update_ids_offset = MAX_UINT64 - size
     for i in range(0, 100000, 2):
-        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')
+        updated_ids[i] = i + update_ids_offset
+        update_ids[i + update_ids_offset] = i
+    external_ids = np.zeros((len(updated_ids) * 2), dtype=np.uint64)
+    updates = np.empty((len(updated_ids) * 2), dtype='O')
     id = 0
-    for prev_id, new_id in update_ids.items():
+    for prev_id, new_id in updated_ids.items():
         external_ids[id] = prev_id
         updates[id] = np.array([], dtype=dtype)
         id += 1
@@ -365,9 +364,102 @@ def test_ivf_flat_ingestion_with_batch_updates(tmp_path):
 
     index.update_batch(vectors=updates, external_ids=external_ids)
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
-    assert accuracy(result, gt_i, updated_ids=updated_ids) > MINIMUM_ACCURACY
+    assert accuracy(result, gt_i, updated_ids=updated_ids) > 0.99
 
     index = index.consolidate_updates()
     _, result = index.query(query_vectors, k=k, nprobe=nprobe)
-    assert accuracy(result, gt_i, updated_ids=updated_ids) > MINIMUM_ACCURACY
+    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")
+    k = 10
+    size = 1000
+    partitions = 10
+    dimensions = 128
+    nqueries = 100
+    nprobe = 10
+    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, nprobe=nprobe)
+    assert accuracy(result, gt_i) == 1.0
 
+    update_ids_offset = MAX_UINT64-size
+    updated_ids = {}
+    for i in range(2, 102):
+        index.delete(external_id=i, timestamp=i)
+        index.update(vector=data[i].astype(dtype), external_id=i + update_ids_offset, timestamp=i)
+        updated_ids[i] = i + update_ids_offset
+
+    index = IVFFlatIndex(uri=index_uri, timestamp=101)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(0, 101))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(2, 101))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert 0.05 <= accuracy(result, gt_i, updated_ids=updated_ids, only_updated_ids=True) <= 0.15
+
+    # Timetravel with partial read from updates table
+    updated_ids_part = {}
+    for i in range(2, 52):
+        updated_ids_part[i] = i + update_ids_offset
+    index = IVFFlatIndex(uri=index_uri, timestamp=51)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids_part) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(0, 51))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids_part) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(2, 51))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert 0.02 <= accuracy(result, gt_i, updated_ids=updated_ids, only_updated_ids=True) <= 0.07
+
+    # Timetravel at previous ingestion timestamp
+    index = IVFFlatIndex(uri=index_uri, timestamp=1)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i) == 1.0
+
+    # Consolidate updates
+    index = index.consolidate_updates()
+    index = IVFFlatIndex(uri=index_uri, timestamp=101)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(0, 101))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(2, 101))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert 0.05 <= accuracy(result, gt_i, updated_ids=updated_ids, only_updated_ids=True) <= 0.15
+
+    # Timetravel with partial read from updates table
+    updated_ids_part = {}
+    for i in range(2, 52):
+        updated_ids_part[i] = i + update_ids_offset
+    index = IVFFlatIndex(uri=index_uri, timestamp=51)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids_part) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(0, 51))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i, updated_ids=updated_ids_part) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(2, 51))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert 0.02 <= accuracy(result, gt_i, updated_ids=updated_ids, only_updated_ids=True) <= 0.07
+
+    # Timetravel at previous ingestion timestamp
+    index = IVFFlatIndex(uri=index_uri, timestamp=1)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i) == 1.0
+    index = IVFFlatIndex(uri=index_uri, timestamp=(0, 1))
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
+    assert accuracy(result, gt_i) == 1.0