mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-23 04:20:56 +08:00
Compare commits
2 Commits
hnsw-strea
...
time-budge
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dad5dffe8e | ||
|
|
06b641407f |
13
src/hnsw.c
13
src/hnsw.c
@@ -18,7 +18,7 @@
|
|||||||
#endif
|
#endif
|
||||||
|
|
||||||
int hnsw_ef_search;
|
int hnsw_ef_search;
|
||||||
bool hnsw_streaming;
|
int hnsw_time_budget;
|
||||||
int hnsw_lock_tranche_id;
|
int hnsw_lock_tranche_id;
|
||||||
static relopt_kind hnsw_relopt_kind;
|
static relopt_kind hnsw_relopt_kind;
|
||||||
|
|
||||||
@@ -69,12 +69,9 @@ HnswInit(void)
|
|||||||
"Valid range is 1..1000.", &hnsw_ef_search,
|
"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);
|
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
|
||||||
|
|
||||||
/* TODO Figure out name */
|
DefineCustomIntVariable("hnsw.time_budget", "Sets the time budget for search",
|
||||||
DefineCustomBoolVariable("hnsw.streaming", "Use streaming mode",
|
NULL, &hnsw_time_budget,
|
||||||
NULL, &hnsw_streaming,
|
-1, -1, INT_MAX, PGC_USERSET, GUC_UNIT_MS, NULL, NULL, NULL);
|
||||||
HNSW_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL);
|
|
||||||
|
|
||||||
/* TODO Add option for limiting iterative search */
|
|
||||||
|
|
||||||
MarkGUCPrefixReserved("hnsw");
|
MarkGUCPrefixReserved("hnsw");
|
||||||
}
|
}
|
||||||
@@ -134,8 +131,6 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
|||||||
/* Account for number of tuples (or entry level), m, and ef_search */
|
/* Account for number of tuples (or entry level), m, and ef_search */
|
||||||
costs.numIndexTuples = (entryLevel + 2) * m;
|
costs.numIndexTuples = (entryLevel + 2) * m;
|
||||||
|
|
||||||
/* TODO Adjust for selectivity for iterative scans */
|
|
||||||
|
|
||||||
genericcostestimate(root, path, loop_count, &costs);
|
genericcostestimate(root, path, loop_count, &costs);
|
||||||
|
|
||||||
/* Use total cost since most work happens before first tuple is returned */
|
/* Use total cost since most work happens before first tuple is returned */
|
||||||
|
|||||||
56
src/hnsw.h
56
src/hnsw.h
@@ -12,10 +12,6 @@
|
|||||||
#include "utils/sampling.h"
|
#include "utils/sampling.h"
|
||||||
#include "vector.h"
|
#include "vector.h"
|
||||||
|
|
||||||
#ifdef HNSW_BENCH
|
|
||||||
#include "portability/instr_time.h"
|
|
||||||
#endif
|
|
||||||
|
|
||||||
#define HNSW_MAX_DIM 2000
|
#define HNSW_MAX_DIM 2000
|
||||||
#define HNSW_MAX_NNZ 1000
|
#define HNSW_MAX_NNZ 1000
|
||||||
|
|
||||||
@@ -46,7 +42,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
|
|
||||||
|
|
||||||
/* Tuple types */
|
/* Tuple types */
|
||||||
#define HNSW_ELEMENT_TUPLE_TYPE 1
|
#define HNSW_ELEMENT_TUPLE_TYPE 1
|
||||||
@@ -73,21 +68,6 @@
|
|||||||
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
|
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
|
||||||
#define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(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
|
#if PG_VERSION_NUM >= 150000
|
||||||
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
|
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
|
||||||
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
|
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
|
||||||
@@ -126,7 +106,7 @@
|
|||||||
|
|
||||||
/* Variables */
|
/* Variables */
|
||||||
extern int hnsw_ef_search;
|
extern int hnsw_ef_search;
|
||||||
extern bool hnsw_streaming;
|
extern int hnsw_time_budget;
|
||||||
extern int hnsw_lock_tranche_id;
|
extern int hnsw_lock_tranche_id;
|
||||||
|
|
||||||
typedef struct HnswElementData HnswElementData;
|
typedef struct HnswElementData HnswElementData;
|
||||||
@@ -150,7 +130,6 @@ struct HnswElementData
|
|||||||
uint8 heaptidsLength;
|
uint8 heaptidsLength;
|
||||||
uint8 level;
|
uint8 level;
|
||||||
uint8 deleted;
|
uint8 deleted;
|
||||||
uint8 version;
|
|
||||||
uint32 hash;
|
uint32 hash;
|
||||||
HnswNeighborsPtr neighbors;
|
HnswNeighborsPtr neighbors;
|
||||||
BlockNumber blkno;
|
BlockNumber blkno;
|
||||||
@@ -177,16 +156,12 @@ struct HnswNeighborArray
|
|||||||
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
|
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
|
||||||
};
|
};
|
||||||
|
|
||||||
typedef struct HnswSearchCandidate
|
typedef struct HnswPairingHeapNode
|
||||||
{
|
{
|
||||||
|
HnswCandidate *inner;
|
||||||
pairingheap_node c_node;
|
pairingheap_node c_node;
|
||||||
pairingheap_node w_node;
|
pairingheap_node w_node;
|
||||||
HnswElementPtr element;
|
} HnswPairingHeapNode;
|
||||||
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 */
|
/* HNSW index options */
|
||||||
typedef struct HnswOptions
|
typedef struct HnswOptions
|
||||||
@@ -331,10 +306,10 @@ typedef struct HnswElementTupleData
|
|||||||
uint8 type;
|
uint8 type;
|
||||||
uint8 level;
|
uint8 level;
|
||||||
uint8 deleted;
|
uint8 deleted;
|
||||||
uint8 version;
|
uint8 unused;
|
||||||
ItemPointerData heaptids[HNSW_HEAPTIDS];
|
ItemPointerData heaptids[HNSW_HEAPTIDS];
|
||||||
ItemPointerData neighbortid;
|
ItemPointerData neighbortid;
|
||||||
uint16 unused;
|
uint16 unused2;
|
||||||
Vector data;
|
Vector data;
|
||||||
} HnswElementTupleData;
|
} HnswElementTupleData;
|
||||||
|
|
||||||
@@ -343,31 +318,20 @@ typedef HnswElementTupleData * HnswElementTuple;
|
|||||||
typedef struct HnswNeighborTupleData
|
typedef struct HnswNeighborTupleData
|
||||||
{
|
{
|
||||||
uint8 type;
|
uint8 type;
|
||||||
uint8 version;
|
uint8 unused;
|
||||||
uint16 count;
|
uint16 count;
|
||||||
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
|
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
|
||||||
} HnswNeighborTupleData;
|
} HnswNeighborTupleData;
|
||||||
|
|
||||||
typedef HnswNeighborTupleData * HnswNeighborTuple;
|
typedef HnswNeighborTupleData * HnswNeighborTuple;
|
||||||
|
|
||||||
typedef union
|
|
||||||
{
|
|
||||||
struct pointerhash_hash *pointers;
|
|
||||||
struct offsethash_hash *offsets;
|
|
||||||
struct tidhash_hash *tids;
|
|
||||||
} visited_hash;
|
|
||||||
|
|
||||||
typedef struct HnswScanOpaqueData
|
typedef struct HnswScanOpaqueData
|
||||||
{
|
{
|
||||||
const HnswTypeInfo *typeInfo;
|
const HnswTypeInfo *typeInfo;
|
||||||
bool first;
|
bool first;
|
||||||
List *w;
|
List *w;
|
||||||
visited_hash v;
|
|
||||||
pairingheap *discarded;
|
|
||||||
Datum q;
|
|
||||||
int m;
|
|
||||||
int64 tuples;
|
|
||||||
MemoryContext tmpCtx;
|
MemoryContext tmpCtx;
|
||||||
|
instr_time start;
|
||||||
|
|
||||||
/* Support functions */
|
/* Support functions */
|
||||||
FmgrInfo *procinfo;
|
FmgrInfo *procinfo;
|
||||||
@@ -412,14 +376,14 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
|
|||||||
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
|
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
|
||||||
void HnswInitPage(Buffer buf, Page page);
|
void HnswInitPage(Buffer buf, Page page);
|
||||||
void HnswInit(void);
|
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, visited_hash * v, pairingheap **discarded, bool initVisited);
|
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, instr_time *start);
|
||||||
HnswElement HnswGetEntryPoint(Relation index);
|
HnswElement HnswGetEntryPoint(Relation index);
|
||||||
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
|
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
|
||||||
void *HnswAlloc(HnswAllocator * allocator, Size size);
|
void *HnswAlloc(HnswAllocator * allocator, Size size);
|
||||||
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
|
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
|
||||||
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
|
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);
|
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
|
||||||
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
|
HnswCandidate *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 HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
|
||||||
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
|
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
|
||||||
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
|
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
|
||||||
|
|||||||
@@ -36,7 +36,7 @@ GetInsertPage(Relation index)
|
|||||||
* Check for a free offset
|
* Check for a free offset
|
||||||
*/
|
*/
|
||||||
static bool
|
static bool
|
||||||
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion)
|
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
|
||||||
{
|
{
|
||||||
OffsetNumber offno;
|
OffsetNumber offno;
|
||||||
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
|
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
|
||||||
@@ -98,7 +98,6 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
|
|||||||
{
|
{
|
||||||
*freeOffno = offno;
|
*freeOffno = offno;
|
||||||
*freeNeighborOffno = neighborOffno;
|
*freeNeighborOffno = neighborOffno;
|
||||||
*tupleVersion = etup->version;
|
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
else if (*nbuf != buf)
|
else if (*nbuf != buf)
|
||||||
@@ -154,7 +153,6 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
|
|||||||
OffsetNumber freeOffno = InvalidOffsetNumber;
|
OffsetNumber freeOffno = InvalidOffsetNumber;
|
||||||
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
|
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
|
||||||
BlockNumber newInsertPage = InvalidBlockNumber;
|
BlockNumber newInsertPage = InvalidBlockNumber;
|
||||||
uint8 tupleVersion;
|
|
||||||
char *base = NULL;
|
char *base = NULL;
|
||||||
|
|
||||||
/* Calculate sizes */
|
/* Calculate sizes */
|
||||||
@@ -204,7 +202,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
|
|||||||
}
|
}
|
||||||
|
|
||||||
/* Next, try space from a deleted element */
|
/* Next, try space from a deleted element */
|
||||||
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
|
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
|
||||||
{
|
{
|
||||||
if (nbuf != buf)
|
if (nbuf != buf)
|
||||||
{
|
{
|
||||||
@@ -214,10 +212,6 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
|
|||||||
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
|
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
/* Set tuple version */
|
|
||||||
etup->version = tupleVersion;
|
|
||||||
ntup->version = tupleVersion;
|
|
||||||
|
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
120
src/hnswscan.c
120
src/hnswscan.c
@@ -3,6 +3,7 @@
|
|||||||
#include "access/relscan.h"
|
#include "access/relscan.h"
|
||||||
#include "hnsw.h"
|
#include "hnsw.h"
|
||||||
#include "pgstat.h"
|
#include "pgstat.h"
|
||||||
|
#include "portability/instr_time.h"
|
||||||
#include "storage/bufmgr.h"
|
#include "storage/bufmgr.h"
|
||||||
#include "storage/lmgr.h"
|
#include "storage/lmgr.h"
|
||||||
#include "utils/memutils.h"
|
#include "utils/memutils.h"
|
||||||
@@ -22,13 +23,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
|
|||||||
int m;
|
int m;
|
||||||
HnswElement entryPoint;
|
HnswElement entryPoint;
|
||||||
char *base = NULL;
|
char *base = NULL;
|
||||||
|
instr_time *start = hnsw_time_budget == -1 ? NULL : &so->start;
|
||||||
|
|
||||||
/* Get m and entry point */
|
/* Get m and entry point */
|
||||||
HnswGetMetaPageInfo(index, &m, &entryPoint);
|
HnswGetMetaPageInfo(index, &m, &entryPoint);
|
||||||
|
|
||||||
so->q = q;
|
|
||||||
so->m = m;
|
|
||||||
|
|
||||||
if (entryPoint == NULL)
|
if (entryPoint == NULL)
|
||||||
return NIL;
|
return NIL;
|
||||||
|
|
||||||
@@ -36,44 +35,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
|
|||||||
|
|
||||||
for (int lc = entryPoint->level; lc >= 1; lc--)
|
for (int lc = entryPoint->level; lc >= 1; lc--)
|
||||||
{
|
{
|
||||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true);
|
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, start);
|
||||||
ep = w;
|
ep = w;
|
||||||
}
|
}
|
||||||
|
|
||||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, &so->discarded, true);
|
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, start);
|
||||||
}
|
|
||||||
|
|
||||||
/*
|
|
||||||
* 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);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/*
|
/*
|
||||||
@@ -126,6 +92,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
|
|||||||
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||||
so->collation = index->rd_indcollation[0];
|
so->collation = index->rd_indcollation[0];
|
||||||
|
|
||||||
|
if (hnsw_time_budget != -1)
|
||||||
|
INSTR_TIME_SET_CURRENT(so->start);
|
||||||
|
|
||||||
scan->opaque = so;
|
scan->opaque = so;
|
||||||
|
|
||||||
return scan;
|
return scan;
|
||||||
@@ -139,13 +108,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
|
|||||||
{
|
{
|
||||||
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
|
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
|
||||||
|
|
||||||
if (!so->first)
|
|
||||||
{
|
|
||||||
pairingheap_reset(so->discarded);
|
|
||||||
tidhash_reset(so->v.tids);
|
|
||||||
}
|
|
||||||
so->first = true;
|
so->first = true;
|
||||||
so->tuples = 0;
|
|
||||||
MemoryContextReset(so->tmpCtx);
|
MemoryContextReset(so->tmpCtx);
|
||||||
|
|
||||||
if (keys && scan->numberOfKeys > 0)
|
if (keys && scan->numberOfKeys > 0)
|
||||||
@@ -195,7 +158,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
|||||||
*/
|
*/
|
||||||
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
||||||
|
|
||||||
HnswBench("scan iteration", so->w = GetScanItems(scan, value));
|
so->w = GetScanItems(scan, value);
|
||||||
|
|
||||||
/* Release shared lock */
|
/* Release shared lock */
|
||||||
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
||||||
@@ -203,83 +166,24 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
|||||||
so->first = false;
|
so->first = false;
|
||||||
|
|
||||||
#if defined(HNSW_MEMORY)
|
#if defined(HNSW_MEMORY)
|
||||||
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
|
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(so->tmpCtx, false) / (1024 * 1024));
|
||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
for (;;)
|
while (list_length(so->w) > 0)
|
||||||
{
|
{
|
||||||
char *base = NULL;
|
char *base = NULL;
|
||||||
HnswSearchCandidate *hc;
|
HnswCandidate *hc = llast(so->w);
|
||||||
HnswElement element;
|
HnswElement element = HnswPtrAccess(base, hc->element);
|
||||||
ItemPointer heaptid;
|
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 */
|
/* Move to next element if no valid heap TIDs */
|
||||||
if (element->heaptidsLength == 0)
|
if (element->heaptidsLength == 0)
|
||||||
{
|
{
|
||||||
so->w = list_delete_last(so->w);
|
so->w = list_delete_last(so->w);
|
||||||
|
|
||||||
/* Mark memory as free for next iteration */
|
|
||||||
if (hnsw_streaming)
|
|
||||||
{
|
|
||||||
pfree(element);
|
|
||||||
pfree(hc);
|
|
||||||
}
|
|
||||||
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
so->tuples++;
|
|
||||||
|
|
||||||
heaptid = &element->heaptids[--element->heaptidsLength];
|
heaptid = &element->heaptids[--element->heaptidsLength];
|
||||||
|
|
||||||
MemoryContextSwitchTo(oldCtx);
|
MemoryContextSwitchTo(oldCtx);
|
||||||
|
|||||||
184
src/hnswutils.c
184
src/hnswutils.c
@@ -9,6 +9,7 @@
|
|||||||
#include "fmgr.h"
|
#include "fmgr.h"
|
||||||
#include "hnsw.h"
|
#include "hnsw.h"
|
||||||
#include "lib/pairingheap.h"
|
#include "lib/pairingheap.h"
|
||||||
|
#include "portability/instr_time.h"
|
||||||
#include "sparsevec.h"
|
#include "sparsevec.h"
|
||||||
#include "storage/bufmgr.h"
|
#include "storage/bufmgr.h"
|
||||||
#include "utils/datum.h"
|
#include "utils/datum.h"
|
||||||
@@ -100,6 +101,13 @@ hash_offset(Size offset)
|
|||||||
#define SH_DEFINE
|
#define SH_DEFINE
|
||||||
#include "lib/simplehash.h"
|
#include "lib/simplehash.h"
|
||||||
|
|
||||||
|
typedef union
|
||||||
|
{
|
||||||
|
pointerhash_hash *pointers;
|
||||||
|
offsethash_hash *offsets;
|
||||||
|
tidhash_hash *tids;
|
||||||
|
} visited_hash;
|
||||||
|
|
||||||
typedef union
|
typedef union
|
||||||
{
|
{
|
||||||
HnswElement element;
|
HnswElement element;
|
||||||
@@ -246,8 +254,6 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
|
|||||||
|
|
||||||
element->level = level;
|
element->level = level;
|
||||||
element->deleted = 0;
|
element->deleted = 0;
|
||||||
/* Start at one to make it easier to find issues */
|
|
||||||
element->version = 1;
|
|
||||||
|
|
||||||
HnswInitNeighbors(base, element, m, allocator);
|
HnswInitNeighbors(base, element, m, allocator);
|
||||||
|
|
||||||
@@ -400,7 +406,6 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
|
|||||||
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
|
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
|
||||||
etup->level = element->level;
|
etup->level = element->level;
|
||||||
etup->deleted = 0;
|
etup->deleted = 0;
|
||||||
etup->version = element->version;
|
|
||||||
for (int i = 0; i < HNSW_HEAPTIDS; i++)
|
for (int i = 0; i < HNSW_HEAPTIDS; i++)
|
||||||
{
|
{
|
||||||
if (i < element->heaptidsLength)
|
if (i < element->heaptidsLength)
|
||||||
@@ -443,7 +448,6 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
|
|||||||
}
|
}
|
||||||
|
|
||||||
ntup->count = idx;
|
ntup->count = idx;
|
||||||
ntup->version = e->version;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/*
|
/*
|
||||||
@@ -517,7 +521,6 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
|
|||||||
{
|
{
|
||||||
element->level = etup->level;
|
element->level = etup->level;
|
||||||
element->deleted = etup->deleted;
|
element->deleted = etup->deleted;
|
||||||
element->version = etup->version;
|
|
||||||
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
|
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
|
||||||
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
|
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
|
||||||
element->heaptidsLength = 0;
|
element->heaptidsLength = 0;
|
||||||
@@ -606,10 +609,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
|
|||||||
/*
|
/*
|
||||||
* Create a candidate for the entry point
|
* Create a candidate for the entry point
|
||||||
*/
|
*/
|
||||||
HnswSearchCandidate *
|
HnswCandidate *
|
||||||
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
|
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
|
||||||
{
|
{
|
||||||
HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate));
|
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
|
||||||
|
|
||||||
HnswPtrStore(base, hc->element, entryPoint);
|
HnswPtrStore(base, hc->element, entryPoint);
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
@@ -619,31 +622,19 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
|
|||||||
return hc;
|
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
|
* Compare candidate distances
|
||||||
*/
|
*/
|
||||||
static int
|
static int
|
||||||
CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
||||||
{
|
{
|
||||||
if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance)
|
if (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance)
|
||||||
return 1;
|
return 1;
|
||||||
|
|
||||||
if (HnswGetSearchCandidateConst(c_node, a)->distance > HnswGetSearchCandidateConst(c_node, b)->distance)
|
if (HnswGetPairingHeapCandidateConst(c_node, a)->distance > HnswGetPairingHeapCandidateConst(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 -1;
|
||||||
|
|
||||||
return 0;
|
return 0;
|
||||||
@@ -655,15 +646,27 @@ CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_n
|
|||||||
static int
|
static int
|
||||||
CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
||||||
{
|
{
|
||||||
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
|
if (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance)
|
||||||
return -1;
|
return -1;
|
||||||
|
|
||||||
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
|
if (HnswGetPairingHeapCandidateConst(w_node, a)->distance > HnswGetPairingHeapCandidateConst(w_node, b)->distance)
|
||||||
return 1;
|
return 1;
|
||||||
|
|
||||||
return 0;
|
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
|
* Init visited
|
||||||
*/
|
*/
|
||||||
@@ -764,30 +767,20 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
|
|||||||
int start;
|
int start;
|
||||||
ItemPointerData indextids[HNSW_MAX_M * 2];
|
ItemPointerData indextids[HNSW_MAX_M * 2];
|
||||||
|
|
||||||
*unvisitedLength = 0;
|
|
||||||
|
|
||||||
buf = ReadBuffer(index, element->neighborPage);
|
buf = ReadBuffer(index, element->neighborPage);
|
||||||
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
||||||
page = BufferGetPage(buf);
|
page = BufferGetPage(buf);
|
||||||
|
|
||||||
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
|
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
|
||||||
|
start = (element->level - lc) * m;
|
||||||
/*
|
|
||||||
* Ensure the neighbor tuple has not been deleted or replaced between
|
|
||||||
* index scan iterations
|
|
||||||
*/
|
|
||||||
if (ntup->version != element->version)
|
|
||||||
{
|
|
||||||
UnlockReleaseBuffer(buf);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
/* Copy to minimize lock time */
|
/* Copy to minimize lock time */
|
||||||
start = (element->level - lc) * m;
|
|
||||||
memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData));
|
memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData));
|
||||||
|
|
||||||
UnlockReleaseBuffer(buf);
|
UnlockReleaseBuffer(buf);
|
||||||
|
|
||||||
|
*unvisitedLength = 0;
|
||||||
|
|
||||||
for (int i = 0; i < lm; i++)
|
for (int i = 0; i < lm; i++)
|
||||||
{
|
{
|
||||||
ItemPointer indextid = &indextids[i];
|
ItemPointer indextid = &indextids[i];
|
||||||
@@ -807,13 +800,13 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
|
|||||||
* Algorithm 2 from paper
|
* Algorithm 2 from paper
|
||||||
*/
|
*/
|
||||||
List *
|
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)
|
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, instr_time *start)
|
||||||
{
|
{
|
||||||
List *w = NIL;
|
List *w = NIL;
|
||||||
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
||||||
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
|
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
|
||||||
int wlen = 0;
|
int wlen = 0;
|
||||||
visited_hash vh;
|
visited_hash v;
|
||||||
ListCell *lc2;
|
ListCell *lc2;
|
||||||
HnswNeighborArray *localNeighborhood = NULL;
|
HnswNeighborArray *localNeighborhood = NULL;
|
||||||
Size neighborhoodSize = 0;
|
Size neighborhoodSize = 0;
|
||||||
@@ -821,19 +814,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
|
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
|
||||||
int unvisitedLength;
|
int unvisitedLength;
|
||||||
|
|
||||||
if (v == NULL)
|
InitVisited(base, &v, index, ef, m);
|
||||||
{
|
|
||||||
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 */
|
/* Create local memory for neighborhood if needed */
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
@@ -845,14 +826,15 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
/* Add entry points to v, C, and W */
|
/* Add entry points to v, C, and W */
|
||||||
foreach(lc2, ep)
|
foreach(lc2, ep)
|
||||||
{
|
{
|
||||||
HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2);
|
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
|
||||||
bool found;
|
bool found;
|
||||||
|
HnswPairingHeapNode *node;
|
||||||
|
|
||||||
if (initVisited)
|
AddToVisited(base, &v, hc->element, index, &found);
|
||||||
AddToVisited(base, v, hc->element, index, &found);
|
|
||||||
|
|
||||||
pairingheap_add(C, &hc->c_node);
|
node = CreatePairingHeapNode(hc);
|
||||||
pairingheap_add(W, &hc->w_node);
|
pairingheap_add(C, &node->c_node);
|
||||||
|
pairingheap_add(W, &node->w_node);
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Do not count elements being deleted towards ef when vacuuming. It
|
* Do not count elements being deleted towards ef when vacuuming. It
|
||||||
@@ -865,28 +847,40 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
|
|
||||||
while (!pairingheap_is_empty(C))
|
while (!pairingheap_is_empty(C))
|
||||||
{
|
{
|
||||||
HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C));
|
HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C));
|
||||||
HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
|
HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
|
||||||
HnswElement cElement;
|
HnswElement cElement;
|
||||||
|
|
||||||
|
/* Check time budget */
|
||||||
|
if (start != NULL)
|
||||||
|
{
|
||||||
|
instr_time duration;
|
||||||
|
|
||||||
|
INSTR_TIME_SET_CURRENT(duration);
|
||||||
|
INSTR_TIME_SUBTRACT(duration, *start);
|
||||||
|
if (INSTR_TIME_GET_MILLISEC(duration) >= hnsw_time_budget)
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
if (c->distance > f->distance)
|
if (c->distance > f->distance)
|
||||||
break;
|
break;
|
||||||
|
|
||||||
cElement = HnswPtrAccess(base, c->element);
|
cElement = HnswPtrAccess(base, c->element);
|
||||||
|
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize);
|
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize);
|
||||||
else
|
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++)
|
for (int i = 0; i < unvisitedLength; i++)
|
||||||
{
|
{
|
||||||
HnswElement eElement;
|
HnswElement eElement;
|
||||||
HnswSearchCandidate *e;
|
HnswCandidate *e;
|
||||||
|
HnswPairingHeapNode *node;
|
||||||
float eDistance;
|
float eDistance;
|
||||||
bool alwaysAdd = wlen < ef;
|
bool alwaysAdd = wlen < ef;
|
||||||
|
|
||||||
f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
|
f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
|
||||||
|
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
{
|
{
|
||||||
@@ -901,33 +895,29 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
|
|
||||||
/* Avoid any allocations if not adding */
|
/* Avoid any allocations if not adding */
|
||||||
eElement = NULL;
|
eElement = NULL;
|
||||||
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
|
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
|
||||||
|
|
||||||
|
if (eElement == NULL)
|
||||||
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
|
if (!(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;
|
continue;
|
||||||
}
|
|
||||||
|
Assert(!eElement->deleted);
|
||||||
|
|
||||||
/* Make robust to issues */
|
/* Make robust to issues */
|
||||||
if (eElement->level < lc)
|
if (eElement->level < lc)
|
||||||
continue;
|
continue;
|
||||||
|
|
||||||
/* Create a new candidate */
|
/* Create a new candidate */
|
||||||
e = palloc(sizeof(HnswSearchCandidate));
|
e = palloc(sizeof(HnswCandidate));
|
||||||
HnswPtrStore(base, e->element, eElement);
|
HnswPtrStore(base, e->element, eElement);
|
||||||
e->distance = eDistance;
|
e->distance = eDistance;
|
||||||
pairingheap_add(C, &e->c_node);
|
|
||||||
pairingheap_add(W, &e->w_node);
|
node = CreatePairingHeapNode(e);
|
||||||
|
pairingheap_add(C, &node->c_node);
|
||||||
|
pairingheap_add(W, &node->w_node);
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Do not count elements being deleted towards ef when vacuuming.
|
* Do not count elements being deleted towards ef when vacuuming.
|
||||||
@@ -940,12 +930,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
|
|
||||||
/* No need to decrement wlen */
|
/* No need to decrement wlen */
|
||||||
if (wlen > ef)
|
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);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -953,7 +938,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
/* Add each element of W to w */
|
/* Add each element of W to w */
|
||||||
while (!pairingheap_is_empty(W))
|
while (!pairingheap_is_empty(W))
|
||||||
{
|
{
|
||||||
HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
|
HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W));
|
||||||
|
|
||||||
w = lappend(w, hc);
|
w = lappend(w, hc);
|
||||||
}
|
}
|
||||||
@@ -1318,7 +1303,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
|||||||
/* 1st phase: greedy search to insert level */
|
/* 1st phase: greedy search to insert level */
|
||||||
for (int lc = entryLevel; lc >= level + 1; lc--)
|
for (int lc = entryLevel; lc >= level + 1; lc--)
|
||||||
{
|
{
|
||||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
|
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL);
|
||||||
ep = w;
|
ep = w;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1334,27 +1319,16 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
|||||||
{
|
{
|
||||||
int lm = HnswGetLayerM(m, lc);
|
int lm = HnswGetLayerM(m, lc);
|
||||||
List *neighbors;
|
List *neighbors;
|
||||||
List *lw = NIL;
|
List *lw;
|
||||||
ListCell *lc2;
|
|
||||||
|
|
||||||
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
|
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL);
|
||||||
|
|
||||||
/* 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 */
|
/* Elements being deleted or skipped can help with search */
|
||||||
/* but should be removed before selecting neighbors */
|
/* but should be removed before selecting neighbors */
|
||||||
if (index != NULL)
|
if (index != NULL)
|
||||||
lw = RemoveElements(base, lw, skipElement);
|
lw = RemoveElements(base, w, skipElement);
|
||||||
|
else
|
||||||
|
lw = w;
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Candidates are sorted, but not deterministically. Could set
|
* Candidates are sorted, but not deterministically. Could set
|
||||||
|
|||||||
@@ -527,14 +527,6 @@ MarkDeleted(HnswVacuumState * vacuumstate)
|
|||||||
for (int i = 0; i < ntup->count; i++)
|
for (int i = 0; i < ntup->count; i++)
|
||||||
ItemPointerSetInvalid(&ntup->indextids[i]);
|
ItemPointerSetInvalid(&ntup->indextids[i]);
|
||||||
|
|
||||||
/* 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;
|
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* We modified the tuples in place, no need to call
|
* We modified the tuples in place, no need to call
|
||||||
* PageIndexTupleOverwrite
|
* PageIndexTupleOverwrite
|
||||||
|
|||||||
@@ -1,43 +0,0 @@
|
|||||||
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();
|
|
||||||
@@ -1,131 +0,0 @@
|
|||||||
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();
|
|
||||||
Reference in New Issue
Block a user