Configure memory budget for distributed OOC queries (#462)

Configure memory budget for distributed OOC queries.

This is a no-op until the library change gets released in TileDB cloud.

Author: NikolaosPapailiou

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

@@ -371,33 +371,23 @@ def dist_qv_udf(
             k_nn: int,
             config: Optional[Mapping[str, Any]] = None,
             timestamp: int = 0,
+            memory_budget: int = -1,
         ):
             queries_m = array_to_matrix(np.transpose(query_vectors))
-            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,
-                )
+            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,
+                # TODO(nikos) add this after the library change gets released in TileDB cloud
+                # upper_bound=0 if memory_budget == -1 else memory_budget,
+            )
             results = []
             for q in range(len(r)):
                 tmp_results = []
@@ -464,6 +454,7 @@ def dist_qv_udf(
                     k_nn=k,
                     config=config,
                     timestamp=self.base_array_timestamp,
+                    memory_budget=self.memory_budget,
                     resource_class="large"
                     if (not resources and not resource_class)
                     else resource_class,

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

@@ -496,10 +496,10 @@ static void declare_dist_qv(py::module& m, const std::string& suffix) {
          std::vector<indices_type>& indices,             // 5
          const std::string& id_uri,
          size_t k_nn,
-         uint64_t timestamp
+         uint64_t timestamp,
+         size_t upper_bound
          /* size_t nthreads TODO: optional arg w/ fallback to C++ default arg */
       ) { /* TODO return type */
-          size_t upper_bound{0};
           auto nthreads = std::thread::hardware_concurrency();
 
           return detail::ivf::dist_qv_finite_ram_part<T, shuffled_ids_type>(
@@ -511,7 +511,8 @@ static void declare_dist_qv(py::module& m, const std::string& suffix) {
               indices,
               id_uri,
               k_nn,
-              timestamp);
+              timestamp,
+              upper_bound);
       },
       py::keep_alive<1, 2>());
   m.def(
@@ -525,11 +526,11 @@ static void declare_dist_qv(py::module& m, const std::string& suffix) {
          std::vector<shuffled_ids_type>& indices,
          const std::string& id_uri,
          size_t k_nn,
-         uint64_t timestamp
+         uint64_t timestamp,
+         size_t upper_bound
          /* size_t nthreads @todo: optional arg w/ fallback to C++ default arg
           */
       ) { /* @todo: return type */
-          size_t upper_bound{0};
           auto nthreads = std::thread::hardware_concurrency();
           auto temporal_policy{
               (timestamp == 0) ? TemporalPolicy() :
@@ -564,7 +565,8 @@ static void declare_dist_qv(py::module& m, const std::string& suffix) {
               indices,
               id_uri,
               k_nn,
-              timestamp);
+              timestamp,
+              upper_bound);
       },
       py::keep_alive<1, 2>());
 }

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

@@ -448,6 +448,7 @@ def dist_qv(
     k_nn: int,
     ctx: "Ctx" = None,
     timestamp: int = 0,
+    upper_bound: int = 0,
 ):
     if ctx is None:
         ctx = vspy.Ctx({})
@@ -462,6 +463,7 @@ def dist_qv(
             ids_uri,
             k_nn,
             timestamp,
+            upper_bound,
         ]
     )