change when prewarm run

timescale · May 1, 2024 · 2401aef · 2401aef
1 parent d994751
commit 2401aef
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ann_benchmarks/algorithms/tsvector/module.py b/ann_benchmarks/algorithms/tsvector/module.py
@@ -356,9 +356,6 @@ def fit(self, X: numpy.array) -> None:
             chunks = self.list_chunks(conn)
         if len(chunks) > 0:
             self.index_chunks(chunks)
-        with self._pool.connection() as conn:
-            self.prewarm_heap(conn)
-            self.prewarm_index(conn)
 
     def set_query_arguments(self, query_search_list_size, query_rescore):
         self._query_search_list_size = query_search_list_size
@@ -367,6 +364,9 @@ def set_query_arguments(self, query_search_list_size, query_rescore):
         self._pool.close()
         self._pool = None
         self.start_pool()
+        with self._pool.connection() as conn:
+            self.prewarm_heap(conn)
+            self.prewarm_index(conn)
 
     def get_memory_usage(self) -> Optional[float]:
         return psutil.Process().memory_info().rss / 1024