diff --git a/src/ivfflat.c b/src/ivfflat.c index 9e29724..d8b4332 100644 --- a/src/ivfflat.c +++ b/src/ivfflat.c @@ -17,6 +17,7 @@ #endif int ivfflat_probes; +int ivfflat_max_probes; bool ivfflat_streaming; static relopt_kind ivfflat_relopt_kind; @@ -34,6 +35,10 @@ IvfflatInit(void) "Valid range is 1..lists.", &ivfflat_probes, IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); + DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative scans", + NULL, &ivfflat_max_probes, + -1, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); + DefineCustomBoolVariable("ivfflat.streaming", "Use streaming mode", NULL, &ivfflat_streaming, IVFFLAT_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL); diff --git a/src/ivfflat.h b/src/ivfflat.h index 5826360..bf4e4a0 100644 --- a/src/ivfflat.h +++ b/src/ivfflat.h @@ -81,6 +81,7 @@ /* Variables */ extern int ivfflat_probes; +extern int ivfflat_max_probes; extern bool ivfflat_streaming; typedef struct VectorArrayData diff --git a/src/ivfscan.c b/src/ivfscan.c index d2f871e..aeaffee 100644 --- a/src/ivfscan.c +++ b/src/ivfscan.c @@ -262,7 +262,15 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys) if (probes > lists) probes = lists; - maxProbes = ivfflat_streaming ? lists : probes; + if (ivfflat_streaming) + { + if (ivfflat_max_probes == -1) + maxProbes = lists; + else + maxProbes = ivfflat_max_probes; + } + else + maxProbes = probes; so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList)); so->typeInfo = IvfflatGetTypeInfo(index);