From 60ff6f33e2bce29c06fa9f17be1725a83375474b Mon Sep 17 00:00:00 2001 From: "min.tian" Date: Thu, 18 Sep 2025 10:54:48 +0800 Subject: [PATCH] fix: remove psycopg dependency from rate_runner Signed-off-by: min.tian --- vectordb_bench/backend/clients/api.py | 1 + vectordb_bench/backend/runner/rate_runner.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index a727cabc1..fdc445618 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -135,6 +135,7 @@ class VectorDB(ABC): "The filtering types supported by the VectorDB Client, default only non-filter" supported_filter_types: list[FilterOp] = [FilterOp.NonFilter] + name: str = "" @classmethod def filter_supported(cls, filters: Filter) -> bool: diff --git a/vectordb_bench/backend/runner/rate_runner.py b/vectordb_bench/backend/runner/rate_runner.py index ca66c58a8..163d50689 100644 --- a/vectordb_bench/backend/runner/rate_runner.py +++ b/vectordb_bench/backend/runner/rate_runner.py @@ -7,7 +7,6 @@ from vectordb_bench import config from vectordb_bench.backend.clients import api -from vectordb_bench.backend.clients.pgvector.pgvector import PgVector from vectordb_bench.backend.dataset import DataSetIterator from vectordb_bench.backend.utils import time_it @@ -48,7 +47,7 @@ def _insert_embeddings(db: api.VectorDB, emb: list[list[float]], metadata: list[ msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times" raise RuntimeError(msg) from None - if isinstance(db, PgVector): + if db.name == "PgVector": # pgvector is not thread-safe for concurrent insert, # so we need to copy the db object, make sure each thread has its own connection db_copy = deepcopy(db)