liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-14 00:26:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Changed option to enum [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-09-29 13:39:53 -07:00
parent 0047630baf
commit 1b337ad97d
3 changed files with 35 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/hnsw.c
View File

@@ -18,9 +18,18 @@
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x) #define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif #endif
static const struct config_enum_entry hnsw_streaming_options[] = {
{"off", HNSW_STREAMING_OFF, false},
{"strict", HNSW_STREAMING_STRICT, false},
{"relaxed", HNSW_STREAMING_RELAXED, false},
/* TODO Change to strict before merging */
{"on", HNSW_STREAMING_RELAXED, false},
{NULL, 0, false}
};
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_ef_stream; int hnsw_ef_stream;
bool hnsw_streaming; int hnsw_streaming;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -72,9 +81,9 @@ HnswInit(void)
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Figure out name */ /* TODO Figure out name */
DefineCustomBoolVariable("hnsw.streaming", "Use streaming mode", DefineCustomEnumVariable("hnsw.streaming", "Use streaming mode",
NULL, &hnsw_streaming, NULL, &hnsw_streaming,
HNSW_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_STREAMING_OFF, hnsw_streaming_options, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Figure out name */ /* TODO Figure out name */
/* TODO Use same value as ivfflat.max_probes for "all" */ /* TODO Use same value as ivfflat.max_probes for "all" */

11
src/hnsw.h
View File

@@ -46,7 +46,6 @@
#define HNSW_DEFAULT_EF_SEARCH 40 #define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1 #define HNSW_MIN_EF_SEARCH 1
#define HNSW_MAX_EF_SEARCH 1000 #define HNSW_MAX_EF_SEARCH 1000
#define HNSW_DEFAULT_STREAMING false
#define HNSW_DEFAULT_EF_STREAM -1 #define HNSW_DEFAULT_EF_STREAM -1
#define HNSW_MIN_EF_STREAM -1 #define HNSW_MIN_EF_STREAM -1
#define HNSW_MAX_EF_STREAM INT_MAX #define HNSW_MAX_EF_STREAM INT_MAX
@@ -130,9 +129,16 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_ef_stream; extern int hnsw_ef_stream;
extern bool hnsw_streaming; extern int hnsw_streaming;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswStreamingMode
{
HNSW_STREAMING_OFF,
HNSW_STREAMING_STRICT,
HNSW_STREAMING_RELAXED
} HnswStreamingMode;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;
@@ -371,6 +377,7 @@ typedef struct HnswScanOpaqueData
Datum q; Datum q;
int m; int m;
int64 tuples; int64 tuples;
double previousDistance;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */

17
src/hnswscan.c
View File

@@ -1,5 +1,7 @@
#include "postgres.h" #include "postgres.h"
#include <float.h>
#include "access/relscan.h" #include "access/relscan.h"
#include "hnsw.h" #include "hnsw.h"
#include "pgstat.h" #include "pgstat.h"
@@ -40,7 +42,7 @@ GetScanItems(IndexScanDesc scan, Datum q)
ep = w; ep = w;
} }
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming ? &so->discarded : NULL, true, &so->tuples); return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming != HNSW_STREAMING_OFF ? &so->discarded : NULL, true, &so->tuples);
} }
/* /*
@@ -149,6 +151,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
so->first = true; so->first = true;
so->tuples = 0; so->tuples = 0;
so->previousDistance = -INFINITY;
MemoryContextReset(so->tmpCtx); MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
@@ -219,7 +222,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (list_length(so->w) == 0) if (list_length(so->w) == 0)
{ {
if (!hnsw_streaming) if (hnsw_streaming == HNSW_STREAMING_OFF)
break; break;
/* Empty index */ /* Empty index */
@@ -285,7 +288,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);
/* Mark memory as free for next iteration */ /* Mark memory as free for next iteration */
if (hnsw_streaming) if (hnsw_streaming != HNSW_STREAMING_OFF)
{ {
pfree(element); pfree(element);
pfree(hc); pfree(hc);
@@ -296,6 +299,14 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
heaptid = &element->heaptids[--element->heaptidsLength]; heaptid = &element->heaptids[--element->heaptidsLength];
if (hnsw_streaming == HNSW_STREAMING_STRICT)
{
if (hc->distance < so->previousDistance)
continue;
so->previousDistance = hc->distance;
}
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;