Compare commits

..

39 Commits

Author SHA1 Message Date
Andrew Kane
dd127da62a Added comments [skip ci] 2024-09-22 10:45:09 -07:00
Andrew Kane
9c59d09a97 Updated comment [skip ci] 2024-09-22 10:35:51 -07:00
Andrew Kane
c5b6769bed Updated variable name [skip ci] 2024-09-22 10:33:44 -07:00
Andrew Kane
d7152448f9 Improved code [skip ci] 2024-09-22 10:24:36 -07:00
Andrew Kane
08420b6404 Updated comment [skip ci] 2024-09-22 10:10:55 -07:00
Andrew Kane
3fba403f8f Updated comment [skip ci] 2024-09-22 10:07:30 -07:00
Andrew Kane
f53e65fb6a Updated comment [skip ci] 2024-09-22 09:31:16 -07:00
Andrew Kane
12015dbe53 Updated comment [skip ci] 2024-09-22 09:18:25 -07:00
Andrew Kane
794ce3a54c Updated comment [skip ci] 2024-09-22 09:10:05 -07:00
Andrew Kane
be51204a45 Only use some bits [skip ci] 2024-09-22 08:53:41 -07:00
Andrew Kane
e78ce02dc9 Fixed CI [skip ci] 2024-09-22 01:52:14 -07:00
Andrew Kane
7412ee6cee Use smaller batch size for better performance 2024-09-22 00:00:02 -07:00
Andrew Kane
ba0196ba10 Improved logging [skip ci] 2024-09-21 23:20:23 -07:00
Andrew Kane
64b0e6359b Updated message [skip ci] 2024-09-21 23:06:30 -07:00
Andrew Kane
0e83213e88 Removed logging [skip ci] 2024-09-21 22:56:35 -07:00
Andrew Kane
f1c7494c35 Added sorting [skip ci] 2024-09-21 22:56:03 -07:00
Andrew Kane
e07d17fe38 Limit work_mem [skip ci] 2024-09-21 22:37:03 -07:00
Andrew Kane
96ff1b992f Added comment [skip ci] 2024-09-21 16:21:02 -07:00
Andrew Kane
569bf93355 Updated comment [skip ci] 2024-09-21 16:18:56 -07:00
Andrew Kane
72aa53ef54 Only free if streaming [skip ci] 2024-09-21 16:15:00 -07:00
Andrew Kane
4773931b11 Improved code [skip ci] 2024-09-21 12:28:12 -07:00
Andrew Kane
b7564dfc19 Improved logging [skip ci] 2024-09-21 12:23:05 -07:00
Andrew Kane
7dc1b36adc Merge branch 'master' into hnsw-streaming 2024-09-21 12:16:40 -07:00
Andrew Kane
d5e8fc96a5 Changed HnswPairingHeapNode to HnswSearchCandidate to reduce allocations and improve code 2024-09-21 12:07:44 -07:00
Andrew Kane
5266c208a6 Reduced calls to AddToVisited [skip ci] 2024-09-21 00:32:31 -07:00
Andrew Kane
a0fddf2d1d Added todo [skip ci] 2024-09-20 16:24:20 -07:00
Andrew Kane
63d5d121a3 Added HNSW_BENCH define [skip ci] 2024-09-20 16:20:27 -07:00
Andrew Kane
39d2ef624a Merge branch 'master' into hnsw-streaming 2024-09-20 15:24:27 -07:00
Andrew Kane
6d2af6d3f9 Improved code [skip ci] 2024-09-20 15:21:57 -07:00
Andrew Kane
aa0b7ddf70 Removed code for pg12 [skip ci] 2024-09-20 15:19:50 -07:00
Andrew Kane
d499ead3c6 Removed todo [skip ci] 2024-09-20 15:19:23 -07:00
Andrew Kane
f20f5e28b8 Added versioning to tuples 2024-09-20 15:12:39 -07:00
Andrew Kane
3dde18a883 Fixed CI 2024-09-20 14:13:33 -07:00
Andrew Kane
0af1c7dd55 Merge branch 'master' into hnsw-streaming 2024-09-20 13:57:23 -07:00
Andrew Kane
f6ebc5d708 Fixed warning [skip ci] 2024-09-20 13:56:27 -07:00
Andrew Kane
c6493415b2 Added test for streaming recall 2024-09-20 13:04:08 -07:00
Andrew Kane
aaff3de409 Free memory [skip ci] 2024-09-19 15:32:22 -07:00
Andrew Kane
4d1c6ff955 Merge branch 'master' into hnsw-streaming 2024-09-19 14:52:34 -07:00
Andrew Kane
af1727775d Added streaming option for HNSW [skip ci] 2024-09-18 14:55:58 -07:00
7 changed files with 425 additions and 77 deletions

View File

@@ -18,6 +18,7 @@
#endif
int hnsw_ef_search;
bool hnsw_streaming;
int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind;
@@ -68,6 +69,13 @@ HnswInit(void)
"Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Figure out name */
DefineCustomBoolVariable("hnsw.streaming", "Use streaming mode",
NULL, &hnsw_streaming,
HNSW_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Add option for limiting iterative search */
MarkGUCPrefixReserved("hnsw");
}
@@ -126,6 +134,8 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
/* Account for number of tuples (or entry level), m, and ef_search */
costs.numIndexTuples = (entryLevel + 2) * m;
/* TODO Adjust for selectivity for iterative scans */
genericcostestimate(root, path, loop_count, &costs);
/* Use total cost since most work happens before first tuple is returned */

View File

@@ -12,6 +12,10 @@
#include "utils/sampling.h"
#include "vector.h"
#ifdef HNSW_BENCH
#include "portability/instr_time.h"
#endif
#define HNSW_MAX_DIM 2000
#define HNSW_MAX_NNZ 1000
@@ -42,6 +46,7 @@
#define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1
#define HNSW_MAX_EF_SEARCH 1000
#define HNSW_DEFAULT_STREAMING false
/* Tuple types */
#define HNSW_ELEMENT_TUPLE_TYPE 1
@@ -68,6 +73,21 @@
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
#define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page))
#ifdef HNSW_BENCH
#define HnswBench(name, code) \
do { \
instr_time start; \
instr_time duration; \
INSTR_TIME_SET_CURRENT(start); \
(code); \
INSTR_TIME_SET_CURRENT(duration); \
INSTR_TIME_SUBTRACT(duration, start); \
elog(INFO, "%s: %.3f ms", name, INSTR_TIME_GET_MILLISEC(duration)); \
} while (0)
#else
#define HnswBench(name, code) (code)
#endif
#if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
@@ -106,6 +126,7 @@
/* Variables */
extern int hnsw_ef_search;
extern bool hnsw_streaming;
extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData;
@@ -156,12 +177,16 @@ struct HnswNeighborArray
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
};
typedef struct HnswPairingHeapNode
typedef struct HnswSearchCandidate
{
HnswCandidate *inner;
pairingheap_node c_node;
pairingheap_node w_node;
} HnswPairingHeapNode;
HnswElementPtr element;
float distance;
} HnswSearchCandidate;
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
/* HNSW index options */
typedef struct HnswOptions
@@ -325,11 +350,23 @@ typedef struct HnswNeighborTupleData
typedef HnswNeighborTupleData * HnswNeighborTuple;
typedef union
{
struct pointerhash_hash *pointers;
struct offsethash_hash *offsets;
struct tidhash_hash *tids;
} visited_hash;
typedef struct HnswScanOpaqueData
{
const HnswTypeInfo *typeInfo;
bool first;
List *w;
visited_hash v;
pairingheap *discarded;
Datum q;
int m;
int64 tuples;
MemoryContext tmpCtx;
/* Support functions */
@@ -375,14 +412,14 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInit(void);
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited);
HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);

View File

@@ -26,6 +26,9 @@ GetScanItems(IndexScanDesc scan, Datum q)
/* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint);
so->q = q;
so->m = m;
if (entryPoint == NULL)
return NIL;
@@ -33,11 +36,44 @@ GetScanItems(IndexScanDesc scan, Datum q)
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true);
ep = w;
}
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, &so->discarded, true);
}
/*
* Resume scan at ground level with discarded candidates
*/
static List *
ResumeScanItems(IndexScanDesc scan)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
List *ep = NIL;
char *base = NULL;
int batch_size = hnsw_ef_search;
if (pairingheap_is_empty(so->discarded))
return NIL;
/* Get next batch of candidates */
for (int i = 0; i < batch_size; i++)
{
HnswSearchCandidate *hc;
if (pairingheap_is_empty(so->discarded))
break;
hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
ep = lappend(ep, hc);
}
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false);
}
/*
@@ -103,7 +139,13 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
if (!so->first)
{
pairingheap_reset(so->discarded);
tidhash_reset(so->v.tids);
}
so->first = true;
so->tuples = 0;
MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0)
@@ -153,7 +195,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
so->w = GetScanItems(scan, value);
HnswBench("scan iteration", so->w = GetScanItems(scan, value));
/* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
@@ -161,24 +203,83 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->first = false;
#if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(so->tmpCtx, false) / (1024 * 1024));
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
#endif
}
while (list_length(so->w) > 0)
for (;;)
{
char *base = NULL;
HnswCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element);
HnswSearchCandidate *hc;
HnswElement element;
ItemPointer heaptid;
if (list_length(so->w) == 0)
{
if (!hnsw_streaming)
break;
/* Prevent scans from consuming too much memory */
if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
{
if (pairingheap_is_empty(so->discarded))
{
ereport(NOTICE,
(errmsg("hnsw iterative search exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase work_mem to scan more tuples.")));
break;
}
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
else
{
/*
* Locking ensures when neighbors are read, the elements they
* reference will not be deleted (and replaced) during the
* iteration.
*
* Elements loaded into memory on previous iterations may have
* been deleted (and replaced), so when reading neighbors, the
* element version must be checked.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
HnswBench("scan iteration", so->w = ResumeScanItems(scan));
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
#if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
#endif
}
if (list_length(so->w) == 0)
break;
}
hc = llast(so->w);
element = HnswPtrAccess(base, hc->element);
/* Move to next element if no valid heap TIDs */
if (element->heaptidsLength == 0)
{
so->w = list_delete_last(so->w);
/* Mark memory as free for next iteration */
if (hnsw_streaming)
{
pfree(element);
pfree(hc);
}
continue;
}
so->tuples++;
heaptid = &element->heaptids[--element->heaptidsLength];
MemoryContextSwitchTo(oldCtx);

View File

@@ -100,13 +100,6 @@ hash_offset(Size offset)
#define SH_DEFINE
#include "lib/simplehash.h"
typedef union
{
pointerhash_hash *pointers;
offsethash_hash *offsets;
tidhash_hash *tids;
} visited_hash;
typedef union
{
HnswElement element;
@@ -613,10 +606,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
/*
* Create a candidate for the entry point
*/
HnswCandidate *
HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, hc->element, entryPoint);
if (index == NULL)
@@ -626,19 +619,31 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
return hc;
}
#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner)
#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner)
/*
* Compare candidate distances
*/
static int
CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance)
if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance)
return 1;
if (HnswGetPairingHeapCandidateConst(c_node, a)->distance > HnswGetPairingHeapCandidateConst(c_node, b)->distance)
if (HnswGetSearchCandidateConst(c_node, a)->distance > HnswGetSearchCandidateConst(c_node, b)->distance)
return -1;
return 0;
}
/*
* Compare discarded candidate distances
*/
static int
CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return 1;
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
return -1;
return 0;
@@ -650,27 +655,15 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
static int
CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance)
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return -1;
if (HnswGetPairingHeapCandidateConst(w_node, a)->distance > HnswGetPairingHeapCandidateConst(w_node, b)->distance)
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
return 1;
return 0;
}
/*
* Create a pairing heap node for a candidate
*/
static HnswPairingHeapNode *
CreatePairingHeapNode(HnswCandidate * c)
{
HnswPairingHeapNode *node = palloc(sizeof(HnswPairingHeapNode));
node->inner = c;
return node;
}
/*
* Init visited
*/
@@ -814,13 +807,13 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper
*/
List *
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited)
{
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0;
visited_hash v;
visited_hash vh;
ListCell *lc2;
HnswNeighborArray *localNeighborhood = NULL;
Size neighborhoodSize = 0;
@@ -828,7 +821,19 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength;
InitVisited(base, &v, index, ef, m);
if (v == NULL)
{
v = &vh;
initVisited = true;
}
if (initVisited)
{
InitVisited(base, v, index, ef, m);
if (discarded != NULL)
*discarded = pairingheap_allocate(CompareNearestDiscardedCandidates, NULL);
}
/* Create local memory for neighborhood if needed */
if (index == NULL)
@@ -840,15 +845,14 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Add entry points to v, C, and W */
foreach(lc2, ep)
{
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2);
bool found;
HnswPairingHeapNode *node;
AddToVisited(base, &v, hc->element, index, &found);
if (initVisited)
AddToVisited(base, v, hc->element, index, &found);
node = CreatePairingHeapNode(hc);
pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
pairingheap_add(C, &hc->c_node);
pairingheap_add(W, &hc->w_node);
/*
* Do not count elements being deleted towards ef when vacuuming. It
@@ -861,8 +865,8 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
while (!pairingheap_is_empty(C))
{
HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C));
HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C));
HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
HnswElement cElement;
if (c->distance > f->distance)
@@ -871,27 +875,23 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
cElement = HnswPtrAccess(base, c->element);
if (index == NULL)
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize);
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize);
else
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, &v, index, m, lm, lc);
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc);
for (int i = 0; i < unvisitedLength; i++)
{
HnswElement eElement;
HnswCandidate *e;
HnswPairingHeapNode *node;
HnswSearchCandidate *e;
float eDistance;
bool alwaysAdd = wlen < ef;
f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
if (index == NULL)
{
eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
if (!(eDistance < f->distance || alwaysAdd))
continue;
}
else
{
@@ -901,26 +901,33 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Avoid any allocations if not adding */
eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
if (eElement == NULL)
continue;
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
}
Assert(!eElement->deleted);
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
{
if (discarded != NULL)
{
/* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(*discarded, &e->w_node);
}
continue;
}
/* Make robust to issues */
if (eElement->level < lc)
continue;
/* Create a new candidate */
e = palloc(sizeof(HnswCandidate));
e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
node = CreatePairingHeapNode(e);
pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node);
/*
* Do not count elements being deleted towards ef when vacuuming.
@@ -933,7 +940,12 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* No need to decrement wlen */
if (wlen > ef)
pairingheap_remove_first(W);
{
HnswSearchCandidate *d = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
if (discarded != NULL)
pairingheap_add(*discarded, &d->w_node);
}
}
}
}
@@ -941,7 +953,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Add each element of W to w */
while (!pairingheap_is_empty(W))
{
HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W));
HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
w = lappend(w, hc);
}
@@ -1306,7 +1318,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
/* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--)
{
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
ep = w;
}
@@ -1322,16 +1334,27 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
{
int lm = HnswGetLayerM(m, lc);
List *neighbors;
List *lw;
List *lw = NIL;
ListCell *lc2;
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement);
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
/* Convert search candidates to candidates */
foreach(lc2, w)
{
HnswSearchCandidate *sc = lfirst(lc2);
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
hc->element = sc->element;
hc->distance = sc->distance;
lw = lappend(lw, hc);
}
/* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */
if (index != NULL)
lw = RemoveElements(base, w, skipElement);
else
lw = w;
lw = RemoveElements(base, lw, skipElement);
/*
* Candidates are sorted, but not deterministically. Could set

View File

@@ -529,7 +529,10 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Increment version */
/* This is used to avoid incorrect reads for iterative scans */
/* Reserve some bits for future use */
etup->version++;
if (etup->version > 15)
etup->version = 1;
ntup->version = etup->version;
/*

View File

@@ -0,0 +1,43 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_l2_ops)
));
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
is($count, 10);
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET work_mem = '2MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
like($stderr, qr/iterative search exceeded work_mem after \d+ tuples/);
done_testing();

View File

@@ -0,0 +1,131 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my @cs = (100, 1000);
sub test_recall
{
my ($c, $ef_search, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.streaming = on;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.streaming = on;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my @expected_ids = split("\n", $expected[$i]);
my %expected_set = map { $_ => 1 } @expected_ids;
foreach (@actual_ids)
{
if (exists($expected_set{$_}))
{
$correct++;
}
}
$total += $limit;
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
}
# Check each index type
my @operators = ("<->", "<=>");
my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
foreach (@cs)
{
my $c = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
WITH top AS (
SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
)
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
));
push(@expected, $res);
}
if ($c == 100)
{
test_recall($c, 40, 0.99, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 40, 0.99, $operator);
}
else
{
test_recall($c, 40, 0.99, $operator);
}
}
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();