Skip to content

Commit 961cb17

Browse files
ankaneankane
committed
Added iterative search for HNSW [skip ci]
1 parent c91ed7b commit 961cb17

File tree

9 files changed

+457
-30
lines changed

9 files changed

+457
-30
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
## 0.8.0 (unreleased)
22

3+
- Added support for iterative index scans
34
- Added casts for arrays to `sparsevec`
45
- Improved cost estimation
56
- Improved performance of HNSW inserts and on-disk index builds

‎src/hnsw.c

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,16 @@
1818
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
1919
#endif
2020

21+
static const struct config_enum_entry hnsw_iterative_search_options[] = {
22+
{"off", HNSW_ITERATIVE_SEARCH_OFF, false},
23+
{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
24+
{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
25+
{NULL, 0, false}
26+
};
27+
2128
int hnsw_ef_search;
29+
int hnsw_iterative_search_max_tuples;
30+
int hnsw_iterative_search;
2231
int hnsw_lock_tranche_id;
2332
static relopt_kind hnsw_relopt_kind;
2433

@@ -69,6 +78,15 @@ HnswInit(void)
6978
"Valid range is 1..1000.", &hnsw_ef_search,
7079
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
7180

81+
DefineCustomEnumVariable("hnsw.iterative_search", "Sets iterative search",
82+
NULL, &hnsw_iterative_search,
83+
HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
84+
85+
/* TODO Ensure ivfflat.max_probes uses same value for "all" */
86+
DefineCustomIntVariable("hnsw.iterative_search_max_tuples", "Sets the max number of candidates to visit for iterative search",
87+
"-1 means all", &hnsw_iterative_search_max_tuples,
88+
-1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
89+
7290
MarkGUCPrefixReserved("hnsw");
7391
}
7492

‎src/hnsw.h

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,17 @@
109109

110110
/* Variables */
111111
extern int hnsw_ef_search;
112+
extern int hnsw_iterative_search;
113+
extern int hnsw_iterative_search_max_tuples;
112114
extern int hnsw_lock_tranche_id;
113115

116+
typedef enum HnswIterativeSearchType
117+
{
118+
HNSW_ITERATIVE_SEARCH_OFF,
119+
HNSW_ITERATIVE_SEARCH_RELAXED,
120+
HNSW_ITERATIVE_SEARCH_STRICT
121+
} HnswIterativeSearchType;
122+
114123
typedef struct HnswElementData HnswElementData;
115124
typedef struct HnswNeighborArray HnswNeighborArray;
116125

@@ -132,6 +141,7 @@ struct HnswElementData
132141
uint8 heaptidsLength;
133142
uint8 level;
134143
uint8 deleted;
144+
uint8 version;
135145
uint32 hash;
136146
HnswNeighborsPtr neighbors;
137147
BlockNumber blkno;
@@ -319,10 +329,10 @@ typedef struct HnswElementTupleData
319329
uint8 type;
320330
uint8 level;
321331
uint8 deleted;
322-
uint8 unused;
332+
uint8 version;
323333
ItemPointerData heaptids[HNSW_HEAPTIDS];
324334
ItemPointerData neighbortid;
325-
uint16 unused2;
335+
uint16 unused;
326336
Vector data;
327337
} HnswElementTupleData;
328338

@@ -331,7 +341,7 @@ typedef HnswElementTupleData * HnswElementTuple;
331341
typedef struct HnswNeighborTupleData
332342
{
333343
uint8 type;
334-
uint8 unused;
344+
uint8 version;
335345
uint16 count;
336346
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
337347
} HnswNeighborTupleData;
@@ -356,6 +366,12 @@ typedef struct HnswScanOpaqueData
356366
const HnswTypeInfo *typeInfo;
357367
bool first;
358368
List *w;
369+
visited_hash v;
370+
pairingheap *discarded;
371+
HnswQuery q;
372+
int m;
373+
int64 tuples;
374+
double previousDistance;
359375
MemoryContext tmpCtx;
360376

361377
/* Support functions */
@@ -399,7 +415,7 @@ bool HnswCheckNorm(HnswSupport * support, Datum value);
399415
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
400416
void HnswInitPage(Buffer buf, Page page);
401417
void HnswInit(void);
402-
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement);
418+
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
403419
HnswElement HnswGetEntryPoint(Relation index);
404420
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
405421
void *HnswAlloc(HnswAllocator * allocator, Size size);

‎src/hnswinsert.c

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ GetInsertPage(Relation index)
3636
* Check for a free offset
3737
*/
3838
static bool
39-
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
39+
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)
4040
{
4141
OffsetNumber offno;
4242
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -98,6 +98,7 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
9898
{
9999
*freeOffno = offno;
100100
*freeNeighborOffno = neighborOffno;
101+
*tupleVersion = etup->version;
101102
return true;
102103
}
103104
else if (*nbuf != buf)
@@ -153,6 +154,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
153154
OffsetNumber freeOffno = InvalidOffsetNumber;
154155
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
155156
BlockNumber newInsertPage = InvalidBlockNumber;
157+
uint8 tupleVersion;
156158
char *base = NULL;
157159

158160
/* Calculate sizes */
@@ -202,7 +204,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
202204
}
203205

204206
/* Next, try space from a deleted element */
205-
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
207+
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
206208
{
207209
if (nbuf != buf)
208210
{
@@ -212,6 +214,10 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
212214
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
213215
}
214216

217+
/* Set tuple version */
218+
etup->version = tupleVersion;
219+
ntup->version = tupleVersion;
220+
215221
break;
216222
}
217223

‎src/hnswscan.c

Lines changed: 131 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include "postgres.h"
22

3+
#include <math.h>
4+
35
#include "access/relscan.h"
46
#include "hnsw.h"
57
#include "pgstat.h"
@@ -21,25 +23,57 @@ GetScanItems(IndexScanDesc scan, Datum value)
2123
int m;
2224
HnswElement entryPoint;
2325
char *base = NULL;
24-
HnswQuery q;
25-
26-
q.value = value;
26+
HnswQuery *q = &so->q;
2727

2828
/* Get m and entry point */
2929
HnswGetMetaPageInfo(index, &m, &entryPoint);
3030

31+
q->value = value;
32+
so->m = m;
33+
3134
if (entryPoint == NULL)
3235
return NIL;
3336

34-
ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, false));
37+
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false));
3538

3639
for (int lc = entryPoint->level; lc >= 1; lc--)
3740
{
38-
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, false, NULL);
41+
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL);
3942
ep = w;
4043
}
4144

42-
return HnswSearchLayer(base, &q, ep, hnsw_ef_search, 0, index, support, m, false, NULL);
45+
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
46+
}
47+
48+
/*
49+
* Resume scan at ground level with discarded candidates
50+
*/
51+
static List *
52+
ResumeScanItems(IndexScanDesc scan)
53+
{
54+
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
55+
Relation index = scan->indexRelation;
56+
List *ep = NIL;
57+
char *base = NULL;
58+
int batch_size = hnsw_ef_search;
59+
60+
if (pairingheap_is_empty(so->discarded))
61+
return NIL;
62+
63+
/* Get next batch of candidates */
64+
for (int i = 0; i < batch_size; i++)
65+
{
66+
HnswSearchCandidate *sc;
67+
68+
if (pairingheap_is_empty(so->discarded))
69+
break;
70+
71+
sc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
72+
73+
ep = lappend(ep, sc);
74+
}
75+
76+
return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
4377
}
4478

4579
/*
@@ -83,6 +117,8 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
83117
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
84118
so->typeInfo = HnswGetTypeInfo(index);
85119
so->first = true;
120+
so->v.tids = NULL;
121+
so->discarded = NULL;
86122
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
87123
"Hnsw scan temporary context",
88124
ALLOCSET_DEFAULT_SIZES);
@@ -103,7 +139,15 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
103139
{
104140
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
105141

142+
if (so->v.tids != NULL)
143+
tidhash_reset(so->v.tids);
144+
145+
if (so->discarded != NULL)
146+
pairingheap_reset(so->discarded);
147+
106148
so->first = true;
149+
so->tuples = 0;
150+
so->previousDistance = -INFINITY;
107151
MemoryContextReset(so->tmpCtx);
108152

109153
if (keys && scan->numberOfKeys > 0)
@@ -165,22 +209,100 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
165209
#endif
166210
}
167211

168-
while (list_length(so->w) > 0)
212+
for (;;)
169213
{
170214
char *base = NULL;
171-
HnswSearchCandidate *sc = llast(so->w);
172-
HnswElement element = HnswPtrAccess(base, sc->element);
215+
HnswSearchCandidate *sc;
216+
HnswElement element;
173217
ItemPointer heaptid;
174218

219+
if (list_length(so->w) == 0)
220+
{
221+
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_OFF)
222+
break;
223+
224+
/* Empty index */
225+
if (so->discarded == NULL)
226+
break;
227+
228+
/* Reached max number of additional tuples */
229+
if (hnsw_iterative_search_max_tuples != -1 && so->tuples >= hnsw_iterative_search_max_tuples)
230+
{
231+
if (pairingheap_is_empty(so->discarded))
232+
break;
233+
234+
/* Return remaining tuples */
235+
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
236+
}
237+
/* Prevent scans from consuming too much memory */
238+
else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
239+
{
240+
if (pairingheap_is_empty(so->discarded))
241+
{
242+
ereport(DEBUG1,
243+
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
244+
errhint("Increase work_mem to scan more tuples.")));
245+
246+
break;
247+
}
248+
249+
/* Return remaining tuples */
250+
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
251+
}
252+
else
253+
{
254+
/*
255+
* Locking ensures when neighbors are read, the elements they
256+
* reference will not be deleted (and replaced) during the
257+
* iteration.
258+
*
259+
* Elements loaded into memory on previous iterations may have
260+
* been deleted (and replaced), so when reading neighbors, the
261+
* element version must be checked.
262+
*/
263+
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
264+
265+
so->w = ResumeScanItems(scan);
266+
267+
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
268+
269+
#if defined(HNSW_MEMORY)
270+
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
271+
#endif
272+
}
273+
274+
if (list_length(so->w) == 0)
275+
break;
276+
}
277+
278+
sc = llast(so->w);
279+
element = HnswPtrAccess(base, sc->element);
280+
175281
/* Move to next element if no valid heap TIDs */
176282
if (element->heaptidsLength == 0)
177283
{
178284
so->w = list_delete_last(so->w);
285+
286+
/* Mark memory as free for next iteration */
287+
if (hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF)
288+
{
289+
pfree(element);
290+
pfree(sc);
291+
}
292+
179293
continue;
180294
}
181295

182296
heaptid = &element->heaptids[--element->heaptidsLength];
183297

298+
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_STRICT)
299+
{
300+
if (sc->distance < so->previousDistance)
301+
continue;
302+
303+
so->previousDistance = sc->distance;
304+
}
305+
184306
MemoryContextSwitchTo(oldCtx);
185307

186308
scan->xs_heaptid = *heaptid;

0 commit comments

Comments
 (0)