Skip to content

Commit 631831b

Browse files
alwayslove2013XuanYang-cn
authored andcommitted
feat(pgvector): add new label/int-filter support and refactor client
Signed-off-by: min.tian <[email protected]>
1 parent 522b6e0 commit 631831b

File tree

2 files changed

+123
-88
lines changed

2 files changed

+123
-88
lines changed

vectordb_bench/backend/clients/pgvector/config.py

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,25 @@ class PgVectorConfigDict(TypedDict):
2121

2222

2323
class PgVectorConfig(DBConfig):
24-
user_name: SecretStr = SecretStr("postgres")
24+
user_name: SecretStr = "postgres"
2525
password: SecretStr
2626
host: str = "localhost"
2727
port: int = 5432
28-
db_name: str
28+
db_name: str = "vectordb"
29+
table_name: str = "vdbbench_table_test"
2930

3031
def to_dict(self) -> PgVectorConfigDict:
31-
user_str = self.user_name.get_secret_value()
32+
user_str = self.user_name.get_secret_value() if isinstance(self.user_name, SecretStr) else self.user_name
3233
pwd_str = self.password.get_secret_value()
3334
return {
34-
"host": self.host,
35-
"port": self.port,
36-
"dbname": self.db_name,
37-
"user": user_str,
38-
"password": pwd_str,
35+
"connect_config": {
36+
"host": self.host,
37+
"port": self.port,
38+
"dbname": self.db_name,
39+
"user": user_str,
40+
"password": pwd_str,
41+
},
42+
"table_name": self.table_name,
3943
}
4044

4145

@@ -59,6 +63,10 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
5963
metric_type: MetricType | None = None
6064
create_index_before_load: bool = False
6165
create_index_after_load: bool = True
66+
# Scan more of the index to get enough results for filter-cases.
67+
# Options: "strict_order" (order by distance), "relaxed_order" (slightly out of order but better recall)
68+
# See: https://github.com/pgvector/pgvector?tab=readme-ov-file#iterative-index-scans
69+
iterative_scan: str = "relaxed_order"
6270

6371
def parse_metric(self) -> str:
6472
d = {
@@ -205,7 +213,7 @@ def search_param(self) -> PgVectorSearchParam:
205213
}
206214

207215
def session_param(self) -> PgVectorSessionCommands:
208-
session_parameters = {"ivfflat.probes": self.probes}
216+
session_parameters = {"ivfflat.probes": self.probes, "ivfflat.iterative_scan": self.iterative_scan}
209217
return {"session_options": self._optionally_build_set_options(session_parameters)}
210218

211219

@@ -255,7 +263,7 @@ def search_param(self) -> PgVectorSearchParam:
255263
}
256264

257265
def session_param(self) -> PgVectorSessionCommands:
258-
session_parameters = {"hnsw.ef_search": self.ef_search}
266+
session_parameters = {"hnsw.ef_search": self.ef_search, "hnsw.iterative_scan": self.iterative_scan}
259267
return {"session_options": self._optionally_build_set_options(session_parameters)}
260268

261269

0 commit comments

Comments
 (0)