Merge pull request #148 from TileDB-Inc/npapa/timetravel

Timetravel implementation

Author: NikolaosPapailiou

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

@@ -22,29 +22,35 @@ def __init__(
         self,
         uri: str,
         config: Optional[Mapping[str, Any]] = None,
+        timestamp=None,
     ):
-        super().__init__(uri=uri, config=config)
+        super().__init__(uri=uri, config=config, timestamp=timestamp)
         self.index_type = "FLAT"
         self._index = None
         self.db_uri = self.group[storage_formats[self.storage_version]["PARTS_ARRAY_NAME"] + self.index_version].uri
         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
         # that the external_ids were the position of the vector in the array.
         if storage_formats[self.storage_version]["IDS_ARRAY_NAME"] + self.index_version in self.group:
             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._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

@@ -33,9 +33,10 @@ def __init__(
         self,
         uri: str,
         config: Optional[Mapping[str, Any]] = None,
+        timestamp=None,
         memory_budget: int = -1,
     ):
-        super().__init__(uri=uri, config=config)
+        super().__init__(uri=uri, config=config, timestamp=timestamp)
         self.index_type = "IVF_FLAT"
         self.db_uri = self.group[
             storage_formats[self.storage_version]["PARTS_ARRAY_NAME"] + self.index_version
@@ -52,9 +53,9 @@ def __init__(
         self.memory_budget = memory_budget
 
         self._centroids = load_as_matrix(
-            self.centroids_uri, ctx=self.ctx, config=config
+            self.centroids_uri, ctx=self.ctx, config=config, timestamp=self.base_array_timestamp
         )
-        self._index = read_vector_u64(self.ctx, self.index_array_uri, 0, 0)
+        self._index = read_vector_u64(self.ctx, self.index_array_uri, 0, 0, self.base_array_timestamp)
 
 
         dtype = self.group.meta.get("dtype", None)
@@ -73,13 +74,15 @@ 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:
-            self._db = load_as_matrix(self.db_uri, ctx=self.ctx, config=config, size=self.size)
-            self._ids = read_vector_u64(self.ctx, self.ids_uri, 0, self.size)
+            self._db = load_as_matrix(self.db_uri, ctx=self.ctx, config=config, size=self.size, timestamp=self.base_array_timestamp)
+            self._ids = read_vector_u64(self.ctx, self.ids_uri, 0, self.size, self.base_array_timestamp)
 
     def query_internal(
         self,
@@ -157,6 +160,7 @@ def query_internal(
                     nthreads=nthreads,
                     ctx=self.ctx,
                     use_nuv_implementation=use_nuv_implementation,
+                    timestamp=self.base_array_timestamp,
                 )
 
             return np.transpose(np.array(d)), np.transpose(np.array(i))
@@ -229,19 +233,34 @@ def dist_qv_udf(
             indices: np.array,
             k_nn: int,
             config: Optional[Mapping[str, Any]] = None,
+            timestamp: int = 0,
         ):
             queries_m = array_to_matrix(np.transpose(query_vectors))
-            r = dist_qv(
-                dtype=dtype,
-                parts_uri=parts_uri,
-                ids_uri=ids_uri,
-                query_vectors=queries_m,
-                active_partitions=active_partitions,
-                active_queries=active_queries,
-                indices=indices,
-                k_nn=k_nn,
-                ctx=Ctx(config),
-            )
+            if timestamp == 0:
+                r = dist_qv(
+                    dtype=dtype,
+                    parts_uri=parts_uri,
+                    ids_uri=ids_uri,
+                    query_vectors=queries_m,
+                    active_partitions=active_partitions,
+                    active_queries=active_queries,
+                    indices=indices,
+                    k_nn=k_nn,
+                    ctx=Ctx(config),
+                )
+            else:
+                r = dist_qv(
+                    dtype=dtype,
+                    parts_uri=parts_uri,
+                    ids_uri=ids_uri,
+                    query_vectors=queries_m,
+                    active_partitions=active_partitions,
+                    active_queries=active_queries,
+                    indices=indices,
+                    k_nn=k_nn,
+                    ctx=Ctx(config),
+                    timestamp=timestamp,
+                )
             results = []
             for q in range(len(r)):
                 tmp_results = []
@@ -308,6 +327,7 @@ def dist_qv_udf(
                     indices=np.array(self._index),
                     k_nn=k,
                     config=config,
+                    timestamp=self.base_array_timestamp,
                     resource_class="large",
                     image_name="3.9-vectorsearch",
                 )

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

@@ -150,7 +150,8 @@ static void declare_qv_query_heap_finite_ram(py::module& m, const std::string& s
          size_t nprobe,
          size_t k_nn,
          size_t upper_bound,
-         size_t nthreads) -> py::tuple { //std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
+         size_t nthreads,
+         uint64_t timestamp) -> py::tuple { //std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
 
         auto r = detail::ivf::qv_query_heap_finite_ram<T, Id_Type>(
             ctx,
@@ -162,7 +163,8 @@ static void declare_qv_query_heap_finite_ram(py::module& m, const std::string& s
             nprobe,
             k_nn,
             upper_bound,
-            nthreads);
+            nthreads,
+            timestamp);
         return make_python_pair(std::move(r));
         }, py::keep_alive<1,2>());
 }
@@ -204,7 +206,8 @@ static void declare_nuv_query_heap_finite_ram(py::module& m, const std::string&
          size_t nprobe,
          size_t k_nn,
          size_t upper_bound,
-         size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> { // TODO change return type
+         size_t nthreads,
+         uint64_t timestamp) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> { // TODO change return type
 
         auto r = detail::ivf::nuv_query_heap_finite_ram_reg_blocked<T, Id_Type>(
             ctx,
@@ -216,7 +219,8 @@ static void declare_nuv_query_heap_finite_ram(py::module& m, const std::string&
             nprobe,
             k_nn,
             upper_bound,
-            nthreads);
+            nthreads,
+            timestamp);
         return r;
         }, py::keep_alive<1,2>());
 }
@@ -234,7 +238,8 @@ static void declare_ivf_index(py::module& m, const std::string& suffix) {
         const std::string& id_uri,
         size_t start_pos,
         size_t end_pos,
-        size_t nthreads) -> int {
+        size_t nthreads,
+        uint64_t timestamp) -> int {
             return detail::ivf::ivf_index<T, uint64_t, float>(
                 ctx,
                 db,
@@ -246,7 +251,8 @@ static void declare_ivf_index(py::module& m, const std::string& suffix) {
                 id_uri,
                 start_pos,
                 end_pos,
-                nthreads);
+                nthreads,
+                timestamp);
         }, py::keep_alive<1,2>());
 }
 
@@ -263,7 +269,8 @@ static void declare_ivf_index_tdb(py::module& m, const std::string& suffix) {
         const std::string& id_uri,
         size_t start_pos,
         size_t end_pos,
-        size_t nthreads) -> int {
+        size_t nthreads,
+        uint64_t timestamp) -> int {
             return detail::ivf::ivf_index<T, uint64_t, float>(
                 ctx,
                 db_uri,
@@ -275,7 +282,8 @@ static void declare_ivf_index_tdb(py::module& m, const std::string& suffix) {
                 id_uri,
                 start_pos,
                 end_pos,
-                nthreads);
+                nthreads,
+                timestamp);
         }, py::keep_alive<1,2>());
 }
 
@@ -302,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>(),  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);
@@ -378,7 +386,8 @@ static void declare_dist_qv(py::module& m, const std::string& suffix) {
         std::vector<std::vector<int>>& active_queries,
         std::vector<shuffled_ids_type>& indices,
         const std::string& id_uri,
-        size_t k_nn
+        size_t k_nn,
+        uint64_t timestamp
         /* size_t nthreads TODO: optional arg w/ fallback to C++ default arg */
         ) { /* TODO return type */
             return detail::ivf::dist_qv_finite_ram_part<T, shuffled_ids_type>(
@@ -389,7 +398,8 @@ static void declare_dist_qv(py::module& m, const std::string& suffix) {
                 active_queries,
                 indices,
                 id_uri,
-                k_nn);
+                k_nn,
+                timestamp);
         }, py::keep_alive<1,2>());
 }
 
@@ -448,8 +458,26 @@ PYBIND11_MODULE(_tiledbvspy, m) {
     declareStdVector<size_t>(m, "szt");
   }
 
-  m.def("read_vector_u32", &read_vector<uint32_t>, "Read a vector from TileDB");
-  m.def("read_vector_u64", &read_vector<uint64_t>, "Read a vector from TileDB");
+  m.def("read_vector_u32",
+    [](const tiledb::Context& ctx,
+        const std::string& uri,
+        size_t start_pos,
+        size_t end_pos,
+        uint64_t timestamp) -> std::vector<uint32_t> {
+      auto r = read_vector<uint32_t>(ctx, uri, start_pos, end_pos, timestamp);
+      return r;
+    });
+  m.def("read_vector_u64",
+    [](const tiledb::Context& ctx,
+        const std::string& uri,
+        size_t start_pos,
+        size_t end_pos,
+        uint64_t timestamp) -> std::vector<uint64_t> {
+      auto r = read_vector<uint64_t>(ctx, uri, start_pos, end_pos, timestamp);
+      return r;
+    });
+//  m.def("read_vector_u32", &read_vector<uint32_t>, "Read a vector from TileDB");
+//  m.def("read_vector_u64", &read_vector<uint64_t>, "Read a vector from TileDB");
 
   m.def("_create_vector_u64", []() {
     auto v = std::vector<uint64_t>(10);

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

@@ -12,7 +12,8 @@ def load_as_matrix(
     path: str,
     ctx: "Ctx" = None,
     config: Optional[Mapping[str, Any]] = None,
-    size: int = 0
+    size: int = 0,
+    timestamp: int = 0,
 ):
     """
     Load array as Matrix class
@@ -36,17 +37,17 @@ 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)
+        m = tdbColMajorMatrix_f32(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.float64:
-        m = tdbColMajorMatrix_f64(ctx, path, size)
+        m = tdbColMajorMatrix_f64(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.int32:
-        m = tdbColMajorMatrix_i32(ctx, path, size)
+        m = tdbColMajorMatrix_i32(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.int32:
-        m = tdbColMajorMatrix_i64(ctx, path, size)
+        m = tdbColMajorMatrix_i64(ctx, path, 0, 0, 0, size, timestamp)
     elif dtype == np.uint8:
-        m = tdbColMajorMatrix_u8(ctx, path, size)
+        m = tdbColMajorMatrix_u8(ctx, path, 0, 0, 0, size, timestamp)
     # elif dtype == np.uint64:
-    #     return tdbColMajorMatrix_u64(ctx, path, size)
+    #     return tdbColMajorMatrix_u64(ctx, path, size, timestamp)
     else:
         raise ValueError("Unsupported Matrix dtype: {}".format(a.attr(0).dtype))
     m.load()
@@ -148,6 +149,7 @@ def ivf_index_tdb(
     start: int = 0,
     end: int = 0,
     nthreads: int = 0,
+    timestamp: int = 0,
     config: Dict = None,
 ):
     if config is None:
@@ -156,7 +158,7 @@ def ivf_index_tdb(
         ctx = Ctx(config)
 
     args = tuple(
-        [ctx, db_uri, external_ids_uri, deleted_ids, centroids_uri, parts_uri, index_array_uri, id_uri, start, end, nthreads]
+        [ctx, db_uri, external_ids_uri, deleted_ids, centroids_uri, parts_uri, index_array_uri, id_uri, start, end, nthreads, timestamp]
     )
 
     if dtype == np.float32:
@@ -179,6 +181,7 @@ def ivf_index(
     start: int = 0,
     end: int = 0,
     nthreads: int = 0,
+    timestamp: int = 0,
     config: Dict = None,
 ):
     if config is None:
@@ -187,7 +190,7 @@ def ivf_index(
         ctx = Ctx(config)
 
     args = tuple(
-        [ctx, db, external_ids, deleted_ids, centroids_uri, parts_uri, index_array_uri, id_uri, start, end, nthreads]
+        [ctx, db, external_ids, deleted_ids, centroids_uri, parts_uri, index_array_uri, id_uri, start, end, nthreads, timestamp]
     )
 
     if dtype == np.float32:
@@ -280,6 +283,7 @@ def ivf_query(
     nthreads: int,
     ctx: "Ctx" = None,
     use_nuv_implementation: bool = False,
+    timestamp: int = 0,
 ):
     """
     Run IVF vector query using a memory budget
@@ -308,6 +312,8 @@ def ivf_query(
         Number of theads
     ctx: Ctx
         Tiledb Context
+    timestamp: int
+        Read timestamp
     """
     if ctx is None:
         ctx = Ctx({})
@@ -324,6 +330,7 @@ def ivf_query(
             k_nn,
             memory_budget,
             nthreads,
+            timestamp,
         ]
     )
 
@@ -360,6 +367,7 @@ def dist_qv(
     indices: np.array,
     k_nn: int,
     ctx: "Ctx" = None,
+    timestamp: int = 0,
 ):
     if ctx is None:
         ctx = Ctx({})
@@ -373,6 +381,7 @@ def dist_qv(
             StdVector_u64(indices),
             ids_uri,
             k_nn,
+            timestamp,
         ]
     )
     if dtype == np.float32: