vectordb-bench 0.0.30__py3-none-any.whl → 1.0.0__py3-none-any.whl

vectordb_bench/backend/runner/rate_runner.py

@@ -30,78 +30,94 @@ class RatedMultiThreadingInsertRunner:
         self.insert_rate = rate
         self.batch_rate = rate // config.NUM_PER_BATCH
 
-    def send_insert_task(self, db: api.VectorDB, emb: list[list[float]], metadata: list[str]):
-        db.insert_embeddings(emb, metadata)
+        self.executing_futures = []
+        self.sig_idx = 0
+
+    def send_insert_task(self, db: api.VectorDB, emb: list[list[float]], metadata: list[str], retry_idx: int = 0):
+        _, error = db.insert_embeddings(emb, metadata)
+        if error is not None:
+            log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {error}")
+            retry_idx += 1
+            if retry_idx <= config.MAX_INSERT_RETRY:
+                time.sleep(retry_idx)
+                self.send_insert_task(db, emb=emb, metadata=metadata, retry_idx=retry_idx)
+            else:
+                msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
+                raise RuntimeError(msg) from None
 
     @time_it
     def run_with_rate(self, q: mp.Queue):
         with ThreadPoolExecutor(max_workers=mp.cpu_count()) as executor:
-            executing_futures = []
 
             @time_it
             def submit_by_rate() -> bool:
                 rate = self.batch_rate
                 for data in self.dataset:
                     emb, metadata = get_data(data, self.normalize)
-                    executing_futures.append(
-                        executor.submit(self.send_insert_task, self.db, emb, metadata),
-                    )
+                    self.executing_futures.append(executor.submit(self.send_insert_task, self.db, emb, metadata))
                     rate -= 1
 
                     if rate == 0:
                         return False
                 return rate == self.batch_rate
 
+            def check_and_send_signal(wait_interval: float, finished: bool = False):
+                try:
+                    done, not_done = concurrent.futures.wait(
+                        self.executing_futures,
+                        timeout=wait_interval,
+                        return_when=concurrent.futures.FIRST_EXCEPTION,
+                    )
+                    _ = [fut.result() for fut in done]
+                    if len(not_done) > 0:
+                        self.executing_futures = list(not_done)
+                    else:
+                        self.executing_futures = []
+
+                    self.sig_idx += len(done)
+                    while self.sig_idx >= self.batch_rate:
+                        self.sig_idx -= self.batch_rate
+                        if self.sig_idx < self.batch_rate and len(not_done) == 0 and finished:
+                            q.put(True, block=True)
+                        else:
+                            q.put(False, block=False)
+
+                except Exception as e:
+                    log.warning(f"task error, terminating, err={e}")
+                    q.put(None, block=True)
+                    executor.shutdown(wait=True, cancel_futures=True)
+                    raise e from None
+
+            time_per_batch = config.TIME_PER_BATCH
             with self.db.init():
+                start_time = time.perf_counter()
+                round_idx = 0
+
                 while True:
-                    start_time = time.perf_counter()
-                    finished, elapsed_time = submit_by_rate()
-                    if finished is True:
-                        q.put(True, block=True)
-                        log.info(f"End of dataset, left unfinished={len(executing_futures)}")
-                        break
-
-                    q.put(False, block=False)
-                    wait_interval = 1 - elapsed_time if elapsed_time < 1 else 0.001
-
-                    try:
-                        done, not_done = concurrent.futures.wait(
-                            executing_futures,
-                            timeout=wait_interval,
-                            return_when=concurrent.futures.FIRST_EXCEPTION,
-                        )
-
-                        if len(not_done) > 0:
-                            log.warning(
-                                f"Failed to finish all tasks in 1s, [{len(not_done)}/{len(executing_futures)}] "
-                                f"tasks are not done, waited={wait_interval:.2f}, trying to wait in the next round"
+                    if len(self.executing_futures) > 200:
+                        log.warning("Skip data insertion this round. There are 200+ unfinished insertion tasks.")
+                    else:
+                        finished, elapsed_time = submit_by_rate()
+                        if finished is True:
+                            log.info(
+                                f"End of dataset, left unfinished={len(self.executing_futures)}, num_round={round_idx}"
                             )
-                            executing_futures = list(not_done)
-                        else:
-                            log.debug(
-                                f"Finished {len(executing_futures)} insert-{config.NUM_PER_BATCH} "
-                                f"task in 1s, wait_interval={wait_interval:.2f}"
+                            break
+                        if elapsed_time >= 1.5:
+                            log.warning(
+                                f"Submit insert tasks took {elapsed_time}s, expected 1s, "
+                                f"indicating potential resource limitations on the client machine.",
                             )
-                            executing_futures = []
-                    except Exception as e:
-                        log.warning(f"task error, terminating, err={e}")
-                        q.put(None, block=True)
-                        executor.shutdown(wait=True, cancel_futures=True)
-                        raise e from e
 
-                    dur = time.perf_counter() - start_time
-                    if dur < 1:
-                        time.sleep(1 - dur)
+                    check_and_send_signal(wait_interval=0.001, finished=False)
+                    dur = time.perf_counter() - start_time - round_idx * time_per_batch
+                    if dur < time_per_batch:
+                        time.sleep(time_per_batch - dur)
+                    round_idx += 1
 
                 # wait for all tasks in executing_futures to complete
-                if len(executing_futures) > 0:
-                    try:
-                        done, _ = concurrent.futures.wait(
-                            executing_futures,
-                            return_when=concurrent.futures.FIRST_EXCEPTION,
-                        )
-                    except Exception as e:
-                        log.warning(f"task error, terminating, err={e}")
-                        q.put(None, block=True)
-                        executor.shutdown(wait=True, cancel_futures=True)
-                        raise e from e
+                while len(self.executing_futures) > 0:
+                    check_and_send_signal(wait_interval=1, finished=True)
+                    round_idx += 1
+
+                log.info(f"Finish all streaming insertion, num_round={round_idx}")

vectordb_bench/backend/runner/read_write_runner.py

@@ -1,13 +1,18 @@
 import concurrent
+import concurrent.futures
 import logging
 import math
 import multiprocessing as mp
+import time
 from collections.abc import Iterable
 
 import numpy as np
 
 from vectordb_bench.backend.clients import api
 from vectordb_bench.backend.dataset import DatasetManager
+from vectordb_bench.backend.filter import Filter, non_filter
+from vectordb_bench.backend.utils import time_it
+from vectordb_bench.metric import Metric
 
 from .mp_runner import MultiProcessingSearchRunner
 from .rate_runner import RatedMultiThreadingInsertRunner
@@ -24,35 +29,39 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
         insert_rate: int = 1000,
         normalize: bool = False,
         k: int = 100,
-        filters: dict | None = None,
+        filters: Filter = non_filter,
         concurrencies: Iterable[int] = (1, 15, 50),
-        search_stage: Iterable[float] = (
+        search_stages: Iterable[float] = (
             0.5,
             0.6,
             0.7,
             0.8,
             0.9,
         ),  # search from insert portion, 0.0 means search from the start
+        optimize_after_write: bool = True,
         read_dur_after_write: int = 300,  # seconds, search duration when insertion is done
         timeout: float | None = None,
     ):
         self.insert_rate = insert_rate
         self.data_volume = dataset.data.size
 
-        for stage in search_stage:
+        for stage in search_stages:
             assert 0.0 <= stage < 1.0, "each search stage should be in [0.0, 1.0)"
-        self.search_stage = sorted(search_stage)
+        self.search_stages = sorted(search_stages)
+        self.optimize_after_write = optimize_after_write
         self.read_dur_after_write = read_dur_after_write
 
         log.info(
-            f"Init runner, concurencys={concurrencies}, search_stage={search_stage}, "
-            f"stage_search_dur={read_dur_after_write}"
+            f"Init runner, concurencys={concurrencies}, search_stages={self.search_stages}, "
+            f"stage_search_dur={read_dur_after_write}",
         )
 
-        test_emb = np.stack(dataset.test_data["emb"])
         if normalize:
+            test_emb = np.array(dataset.test_data)
             test_emb = test_emb / np.linalg.norm(test_emb, axis=1)[:, np.newaxis]
-        test_emb = test_emb.tolist()
+            test_emb = test_emb.tolist()
+        else:
+            test_emb = dataset.test_data
 
         MultiProcessingSearchRunner.__init__(
             self,
@@ -74,8 +83,10 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
             test_data=test_emb,
             ground_truth=dataset.gt_data,
             k=k,
+            filters=filters,
         )
 
+    @time_it
     def run_optimize(self):
         """Optimize needs to run in differenct process for pymilvus schema recursion problem"""
         with self.db.init():
@@ -83,49 +94,102 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
             self.db.optimize(data_size=self.data_volume)
             log.info("Search after write - Optimize finished")
 
-    def run_search(self):
+    def run_search(self, perc: int):
         log.info("Search after write - Serial search start")
+        test_time = round(time.perf_counter(), 4)
         res, ssearch_dur = self.serial_search_runner.run()
         recall, ndcg, p99_latency = res
         log.info(
-            f"Search after write - Serial search - recall={recall}, ndcg={ndcg}, p99={p99_latency}, "
-            f"dur={ssearch_dur:.4f}",
+            f"Search after write - Serial search - recall={recall}, ndcg={ndcg}, "
+            f"p99={p99_latency}, dur={ssearch_dur:.4f}",
+        )
+        log.info(
+            f"Search after wirte - Conc search start, dur for each conc={self.read_dur_after_write}",
         )
-        log.info(f"Search after wirte - Conc search start, dur for each conc={self.read_dur_after_write}")
-        max_qps = self.run_by_dur(self.read_dur_after_write)
+        max_qps, conc_failed_rate = self.run_by_dur(self.read_dur_after_write)
         log.info(f"Search after wirte - Conc search finished, max_qps={max_qps}")
 
-        return (max_qps, recall, ndcg, p99_latency)
+        return [(perc, test_time, max_qps, recall, ndcg, p99_latency, conc_failed_rate)]
 
-    def run_read_write(self):
-        with mp.Manager() as m:
-            q = m.Queue()
-            with concurrent.futures.ProcessPoolExecutor(
-                mp_context=mp.get_context("spawn"),
-                max_workers=2,
-            ) as executor:
-                read_write_futures = []
-                read_write_futures.append(executor.submit(self.run_with_rate, q))
-                read_write_futures.append(executor.submit(self.run_search_by_sig, q))
+    def run_read_write(self) -> Metric:
+        """
+        Test search performance with a fixed insert rate.
+        - Insert requests are sent to VectorDB at a fixed rate within a dedicated insert process pool.
+          - if the database cannot promptly process these requests, the process pool will accumulate insert tasks.
+        - Search Tests are categorized into three types:
+          - streaming_search: Initiates a new search test upon receiving a signal that the inserted data has
+          reached the search_stage.
+          - streaming_end_search: initiates a new search test after all data has been inserted.
+          - optimized_search (optional): After the streaming_end_search, optimizes and initiates a search test.
+        """
+        m = Metric()
+        with mp.Manager() as mp_manager:
+            q = mp_manager.Queue()
+            with concurrent.futures.ProcessPoolExecutor(mp_context=mp.get_context("spawn"), max_workers=2) as executor:
+                insert_future = executor.submit(self.run_with_rate, q)
+                streaming_search_future = executor.submit(self.run_search_by_sig, q)
 
                 try:
-                    for f in concurrent.futures.as_completed(read_write_futures):
-                        res = f.result()
-                        log.info(f"Result = {res}")
+                    start_time = time.perf_counter()
+                    _, m.insert_duration = insert_future.result()
+                    streaming_search_res = streaming_search_future.result()
+                    if streaming_search_res is None:
+                        streaming_search_res = []
+
+                    streaming_end_search_future = executor.submit(self.run_search, 100)
+                    streaming_end_search_res = streaming_end_search_future.result()
 
                     # Wait for read_write_futures finishing and do optimize and search
-                    op_future = executor.submit(self.run_optimize)
-                    op_future.result()
+                    if self.optimize_after_write:
+                        op_future = executor.submit(self.run_optimize)
+                        _, m.optimize_duration = op_future.result()
+                        log.info(f"Optimize cost {m.optimize_duration}s")
+                        optimized_search_future = executor.submit(self.run_search, 110)
+                        optimized_search_res = optimized_search_future.result()
+                    else:
+                        log.info("Skip optimization and search")
+                        optimized_search_res = []
 
-                    search_future = executor.submit(self.run_search)
-                    last_res = search_future.result()
+                    r = [*streaming_search_res, *streaming_end_search_res, *optimized_search_res]
+                    m.st_search_stage_list = [d[0] for d in r]
+                    m.st_search_time_list = [round(d[1] - start_time, 4) for d in r]
+                    m.st_max_qps_list_list = [d[2] for d in r]
+                    m.st_recall_list = [d[3] for d in r]
+                    m.st_ndcg_list = [d[4] for d in r]
+                    m.st_serial_latency_p99_list = [d[5] for d in r]
+                    m.st_conc_failed_rate_list = [d[6] for d in r]
 
-                    log.info(f"Max QPS after optimze and search: {last_res}")
                 except Exception as e:
                     log.warning(f"Read and write error: {e}")
                     executor.shutdown(wait=True, cancel_futures=True)
-                    raise e from e
-        log.info("Concurrent read write all done")
+                    # raise e
+        m.st_ideal_insert_duration = math.ceil(self.data_volume / self.insert_rate)
+        log.info(f"Concurrent read write all done, results: {m}")
+        return m
+
+    def get_each_conc_search_dur(self, ssearch_dur: float, cur_stage: float, next_stage: float) -> float:
+        # Search duration for non-last search stage is carefully calculated.
+        # If duration for each concurrency is less than 30s, runner will raise error.
+        total_dur_between_stages = self.data_volume * (next_stage - cur_stage) // self.insert_rate
+        csearch_dur = total_dur_between_stages - ssearch_dur
+
+        # Try to leave room for init process executors
+        if csearch_dur > 60:
+            csearch_dur -= 30
+        elif csearch_dur > 30:
+            csearch_dur -= 15
+        else:
+            csearch_dur /= 2
+
+        each_conc_search_dur = round(csearch_dur / len(self.concurrencies), 4)
+        if each_conc_search_dur < 30:
+            warning_msg = (
+                f"Results might be inaccurate, duration[{csearch_dur:.4f}] left for conc-search is too short, "
+                f"total available dur={total_dur_between_stages}, serial_search_cost={ssearch_dur}, "
+                f"each_conc_search_dur={each_conc_search_dur}."
+            )
+            log.warning(warning_msg)
+        return each_conc_search_dur
 
     def run_search_by_sig(self, q: mp.Queue):
         """
@@ -149,7 +213,7 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
                 start += 1
             return True
 
-        for idx, stage in enumerate(self.search_stage):
+        for idx, stage in enumerate(self.search_stages):
             target_batch = int(total_batch * stage)
             perc = int(stage * 100)
 
@@ -159,41 +223,34 @@ class ReadWriteRunner(MultiProcessingSearchRunner, RatedMultiThreadingInsertRunn
                 return None
 
             log.info(f"Insert {perc}% done, total batch={total_batch}")
-            log.info(f"[{target_batch}/{total_batch}] Serial search - {perc}% start")
-            res, ssearch_dur = self.serial_search_runner.run()
-            recall, ndcg, p99_latency = res
-            log.info(
-                f"[{target_batch}/{total_batch}] Serial search - {perc}% done, recall={recall}, "
-                f"ndcg={ndcg}, p99={p99_latency}, dur={ssearch_dur:.4f}"
-            )
-
-            # Search duration for non-last search stage is carefully calculated.
-            # If duration for each concurrency is less than 30s, runner will raise error.
-            if idx < len(self.search_stage) - 1:
-                total_dur_between_stages = self.data_volume * (self.search_stage[idx + 1] - stage) // self.insert_rate
-                csearch_dur = total_dur_between_stages - ssearch_dur
-
-                # Try to leave room for init process executors
-                csearch_dur = csearch_dur - 30 if csearch_dur > 60 else csearch_dur
+            test_time = round(time.perf_counter(), 4)
+            max_qps, recall, ndcg, p99_latency, conc_failed_rate = 0, 0, 0, 0, 0
+            try:
+                log.info(f"[{target_batch}/{total_batch}] Serial search - {perc}% start")
+                res, ssearch_dur = self.serial_search_runner.run()
+                ssearch_dur = round(ssearch_dur, 4)
+                recall, ndcg, p99_latency = res
+                log.info(
+                    f"[{target_batch}/{total_batch}] Serial search - {perc}% done, "
+                    f"recall={recall}, ndcg={ndcg}, p99={p99_latency}, dur={ssearch_dur}"
+                )
 
-                each_conc_search_dur = csearch_dur / len(self.concurrencies)
-                if each_conc_search_dur < 30:
-                    warning_msg = (
-                        f"Results might be inaccurate, duration[{csearch_dur:.4f}] left for conc-search is too short, "
-                        f"total available dur={total_dur_between_stages}, serial_search_cost={ssearch_dur}."
+                each_conc_search_dur = self.get_each_conc_search_dur(
+                    ssearch_dur,
+                    cur_stage=stage,
+                    next_stage=self.search_stages[idx + 1] if idx < len(self.search_stages) - 1 else 1.0,
+                )
+                if each_conc_search_dur > 10:
+                    log.info(
+                        f"[{target_batch}/{total_batch}] Concurrent search - {perc}% start, "
+                        f"dur={each_conc_search_dur:.4f}"
                     )
-                    log.warning(warning_msg)
-
-            # The last stage
-            else:
-                each_conc_search_dur = 60
-
-            log.info(
-                f"[{target_batch}/{total_batch}] Concurrent search - {perc}% start, dur={each_conc_search_dur:.4f}"
-            )
-            max_qps = self.run_by_dur(each_conc_search_dur)
-            result.append((perc, max_qps, recall, ndcg, p99_latency))
-
+                    max_qps, conc_failed_rate = self.run_by_dur(each_conc_search_dur)
+                else:
+                    log.warning(f"Skip concurrent tests, each_conc_search_dur={each_conc_search_dur} less than 10s.")
+            except Exception as e:
+                log.warning(f"Streaming Search Failed at stage={stage}. Exception: {e}")
+            result.append((perc, test_time, max_qps, recall, ndcg, p99_latency, conc_failed_rate))
             start_batch = target_batch
 
         # Drain the queue

vectordb_bench/backend/runner/serial_runner.py

@@ -6,10 +6,10 @@ import time
 import traceback
 
 import numpy as np
-import pandas as pd
 import psutil
 
 from vectordb_bench.backend.dataset import DatasetManager
+from vectordb_bench.backend.filter import Filter, FilterOp, non_filter
 
 from ... import config
 from ...metric import calc_ndcg, calc_recall, get_ideal_dcg
@@ -18,8 +18,7 @@ from .. import utils
 from ..clients import api
 
 NUM_PER_BATCH = config.NUM_PER_BATCH
-LOAD_MAX_TRY_COUNT = 10
-WAITTING_TIME = 60
+LOAD_MAX_TRY_COUNT = config.LOAD_MAX_TRY_COUNT
 
 log = logging.getLogger(__name__)
 
@@ -30,12 +29,26 @@ class SerialInsertRunner:
         db: api.VectorDB,
         dataset: DatasetManager,
         normalize: bool,
+        filters: Filter = non_filter,
         timeout: float | None = None,
     ):
         self.timeout = timeout if isinstance(timeout, int | float) else None
         self.dataset = dataset
         self.db = db
         self.normalize = normalize
+        self.filters = filters
+
+    def retry_insert(self, db: api.VectorDB, retry_idx: int = 0, **kwargs):
+        _, error = db.insert_embeddings(**kwargs)
+        if error is not None:
+            log.warning(f"Insert Failed, try_idx={retry_idx}, Exception: {error}")
+            retry_idx += 1
+            if retry_idx <= config.MAX_INSERT_RETRY:
+                time.sleep(retry_idx)
+                self.retry_insert(db, retry_idx=retry_idx, **kwargs)
+            else:
+                msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
+                raise RuntimeError(msg) from None
 
     def task(self) -> int:
         count = 0
@@ -43,9 +56,9 @@ class SerialInsertRunner:
             log.info(f"({mp.current_process().name:16}) Start inserting embeddings in batch {config.NUM_PER_BATCH}")
             start = time.perf_counter()
             for data_df in self.dataset:
-                all_metadata = data_df["id"].tolist()
+                all_metadata = data_df[self.dataset.data.train_id_field].tolist()
 
-                emb_np = np.stack(data_df["emb"])
+                emb_np = np.stack(data_df[self.dataset.data.train_vector_field])
                 if self.normalize:
                     log.debug("normalize the 100k train data")
                     all_embeddings = (emb_np / np.linalg.norm(emb_np, axis=1)[:, np.newaxis]).tolist()
@@ -54,12 +67,25 @@ class SerialInsertRunner:
                 del emb_np
                 log.debug(f"batch dataset size: {len(all_embeddings)}, {len(all_metadata)}")
 
+                labels_data = None
+                if self.filters.type == FilterOp.StrEqual:
+                    if self.dataset.data.scalar_labels_file_separated:
+                        labels_data = self.dataset.scalar_labels[self.filters.label_field][all_metadata].to_list()
+                    else:
+                        labels_data = data_df[self.filters.label_field].tolist()
+
                 insert_count, error = self.db.insert_embeddings(
                     embeddings=all_embeddings,
                     metadata=all_metadata,
+                    labels_data=labels_data,
                 )
                 if error is not None:
-                    raise error
+                    self.retry_insert(
+                        self.db,
+                        embeddings=all_embeddings,
+                        metadata=all_metadata,
+                        labels_data=labels_data,
+                    )
 
                 assert insert_count == len(all_metadata)
                 count += insert_count
@@ -101,7 +127,7 @@ class SerialInsertRunner:
                     already_insert_count += insert_count
                     if error is not None:
                         retry_count += 1
-                        time.sleep(WAITTING_TIME)
+                        time.sleep(10)
 
                         log.info(f"Failed to insert data, try {retry_count} time")
                         if retry_count >= LOAD_MAX_TRY_COUNT:
@@ -149,8 +175,8 @@ class SerialInsertRunner:
         # only 1 file
         data_df = next(iter(self.dataset))
         all_embeddings, all_metadata = (
-            np.stack(data_df["emb"]).tolist(),
-            data_df["id"].tolist(),
+            np.stack(data_df[self.dataset.data.train_vector_field]).tolist(),
+            data_df[self.dataset.data.train_id_field].tolist(),
         )
 
         start_time = time.perf_counter()
@@ -188,9 +214,9 @@ class SerialSearchRunner:
         self,
         db: api.VectorDB,
         test_data: list[list[float]],
-        ground_truth: pd.DataFrame,
+        ground_truth: list[list[int]],
         k: int = 100,
-        filters: dict | None = None,
+        filters: Filter = non_filter,
     ):
         self.db = db
         self.k = k
@@ -202,35 +228,42 @@ class SerialSearchRunner:
             self.test_data = test_data
         self.ground_truth = ground_truth
 
-    def search(self, args: tuple[list, pd.DataFrame]) -> tuple[float, float, float]:
+    def _get_db_search_res(self, emb: list[float], retry_idx: int = 0) -> list[int]:
+        try:
+            results = self.db.search_embedding(emb, self.k)
+        except Exception as e:
+            log.warning(f"Serial search failed, retry_idx={retry_idx}, Exception: {e}")
+            if retry_idx < config.MAX_SEARCH_RETRY:
+                return self._get_db_search_res(emb=emb, retry_idx=retry_idx + 1)
+
+            msg = f"Serial search failed and retried more than {config.MAX_SEARCH_RETRY} times"
+            raise RuntimeError(msg) from e
+
+        return results
+
+    def search(self, args: tuple[list, list[list[int]]]) -> tuple[float, float, float]:
         log.info(f"{mp.current_process().name:14} start search the entire test_data to get recall and latency")
         with self.db.init():
+            self.db.prepare_filter(self.filters)
             test_data, ground_truth = args
             ideal_dcg = get_ideal_dcg(self.k)
 
             log.debug(f"test dataset size: {len(test_data)}")
-            if ground_truth is not None:
-                log.debug(f"ground truth size: {ground_truth.columns}, shape: {ground_truth.shape}")
+            log.debug(f"ground truth size: {len(ground_truth)}")
 
             latencies, recalls, ndcgs = [], [], []
             for idx, emb in enumerate(test_data):
                 s = time.perf_counter()
                 try:
-                    results = self.db.search_embedding(
-                        emb,
-                        self.k,
-                        self.filters,
-                    )
-
+                    results = self._get_db_search_res(emb)
                 except Exception as e:
                     log.warning(f"VectorDB search_embedding error: {e}")
-                    traceback.print_exc(chain=True)
                     raise e from None
 
                 latencies.append(time.perf_counter() - s)
 
                 if ground_truth is not None:
-                    gt = ground_truth["neighbors_id"][idx]
+                    gt = ground_truth[idx]
                     recalls.append(calc_recall(self.k, gt[: self.k], results))
                     ndcgs.append(calc_ndcg(gt[: self.k], results, ideal_dcg))
                 else:
@@ -253,7 +286,7 @@ class SerialSearchRunner:
             f"cost={cost}s, "
             f"queries={len(latencies)}, "
             f"avg_recall={avg_recall}, "
-            f"avg_ndcg={avg_ndcg},"
+            f"avg_ndcg={avg_ndcg}, "
             f"avg_latency={avg_latency}, "
             f"p99={p99}"
         )