From 1051810eb81a66f7d473039a5264e9b896cbbea5 Mon Sep 17 00:00:00 2001 From: Sheharyar Ahmad Date: Tue, 5 Nov 2024 19:05:39 +0500 Subject: [PATCH 1/3] fixed pgvectorivvflat cli reranking key error bug --- vectordb_bench/backend/clients/pgvector/cli.py | 3 +++ vectordb_bench/backend/clients/pgvector/config.py | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/vectordb_bench/backend/clients/pgvector/cli.py b/vectordb_bench/backend/clients/pgvector/cli.py index cde125cc..bf94a727 100644 --- a/vectordb_bench/backend/clients/pgvector/cli.py +++ b/vectordb_bench/backend/clients/pgvector/cli.py @@ -136,6 +136,9 @@ def PgVectorIVFFlat( lists=parameters["lists"], probes=parameters["probes"], quantization_type=parameters["quantization_type"], + reranking=parameters["reranking"], + reranking_metric=parameters["reranking_metric"], + quantized_fetch_limit=parameters["quantized_fetch_limit"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/pgvector/config.py b/vectordb_bench/backend/clients/pgvector/config.py index 31d832f1..16d54744 100644 --- a/vectordb_bench/backend/clients/pgvector/config.py +++ b/vectordb_bench/backend/clients/pgvector/config.py @@ -168,6 +168,9 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig): maintenance_work_mem: Optional[str] = None max_parallel_workers: Optional[int] = None quantization_type: Optional[str] = None + reranking: Optional[bool] = None + quantized_fetch_limit: Optional[int] = None + reranking_metric: Optional[str] = None def index_param(self) -> PgVectorIndexParam: index_parameters = {"lists": self.lists} @@ -187,6 +190,9 @@ def index_param(self) -> PgVectorIndexParam: def search_param(self) -> PgVectorSearchParam: return { "metric_fun_op": self.parse_metric_fun_op(), + "reranking": self.reranking, + "reranking_metric_fun_op": self.parse_reranking_metric_fun_op(), + "quantized_fetch_limit": self.quantized_fetch_limit, } def session_param(self) -> PgVectorSessionCommands: From bb8b84c8b49dcd6d9a521231e0d76de25274e421 Mon Sep 17 00:00:00 2001 From: Sheharyar Ahmad Date: Tue, 5 Nov 2024 23:05:36 +0500 Subject: [PATCH 2/3] set default value of quantized_fetch_limit to 100 in case of ivfflat, because ef_search is not a param of ivfflat. --- vectordb_bench/backend/clients/pgvector/cli.py | 5 +++-- vectordb_bench/cli/cli.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vectordb_bench/backend/clients/pgvector/cli.py b/vectordb_bench/backend/clients/pgvector/cli.py index bf94a727..ce98806a 100644 --- a/vectordb_bench/backend/clients/pgvector/cli.py +++ b/vectordb_bench/backend/clients/pgvector/cli.py @@ -18,11 +18,12 @@ from vectordb_bench.backend.clients import DB - def set_default_quantized_fetch_limit(ctx, param, value): if ctx.params.get("reranking") and value is None: # ef_search is the default value for quantized_fetch_limit as it's bound by ef_search. - return ctx.params["ef_search"] + # 100 (arbitrary) is default value for quantized_fetch_limit for IVFFlat. + default_value = ctx.params["ef_search"] if ctx.command.name == "pgvectorhnsw" else 100 + return default_value return value class PgVectorTypedDict(CommonTypedDict): diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index edce758e..c28b64f4 100644 --- a/vectordb_bench/cli/cli.py +++ b/vectordb_bench/cli/cli.py @@ -414,7 +414,7 @@ class HNSWBaseRequiredTypedDict(TypedDict): class HNSWFlavor1(HNSWBaseTypedDict): ef_search: Annotated[ - Optional[int], click.option("--ef-search", type=int, help="hnsw ef-search") + Optional[int], click.option("--ef-search", type=int, help="hnsw ef-search", is_eager=True) ] From 9c422d9247c0a88e919f33a7b0ae853773d925bb Mon Sep 17 00:00:00 2001 From: Sheharyar Ahmad Date: Tue, 5 Nov 2024 23:13:47 +0500 Subject: [PATCH 3/3] update comment --- vectordb_bench/backend/clients/pgvector/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vectordb_bench/backend/clients/pgvector/cli.py b/vectordb_bench/backend/clients/pgvector/cli.py index ce98806a..43385ee6 100644 --- a/vectordb_bench/backend/clients/pgvector/cli.py +++ b/vectordb_bench/backend/clients/pgvector/cli.py @@ -21,7 +21,7 @@ def set_default_quantized_fetch_limit(ctx, param, value): if ctx.params.get("reranking") and value is None: # ef_search is the default value for quantized_fetch_limit as it's bound by ef_search. - # 100 (arbitrary) is default value for quantized_fetch_limit for IVFFlat. + # 100 is default value for quantized_fetch_limit for IVFFlat. default_value = ctx.params["ef_search"] if ctx.command.name == "pgvectorhnsw" else 100 return default_value return value