8000 Added HnswQuery struct for query data · postgrespro/pgvector@fa67829 · GitHub
[go: up one dir, main page]
Skip to content

Commit fa67829

Browse files
ankaneankane
committed
Added HnswQuery struct for query data
1 parent 32ab27d commit fa67829

File tree

4 files changed

+37
-25
lines changed

4 files changed

+37
-25
lines changed

‎src/hnsw.h

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,11 @@ typedef struct HnswSupport
244244
Oid collation;
245245
} HnswSupport;
246246

247+
typedef struct HnswQuery
248+
{
249+
Datum value;
250+
} HnswQuery;
251+
247252
typedef struct HnswBuildState
248253
{
249254
/* Info */
@@ -378,14 +383,14 @@ bool HnswCheckNorm(HnswSupport * support, Datum value);
378383
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
379384
void HnswInitPage(Buffer buf, Page page);
380385
void HnswInit(void);
381-
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement);
386+
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement);
382387
HnswElement HnswGetEntryPoint(Relation index);
383388
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
384389
void *HnswAlloc(HnswAllocator * allocator, Size size);
385390
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
386391
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
387392
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing);
388-
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, HnswSupport * support, bool loadVec);
393+
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec);
389394
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
390395
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
391396
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
@@ -394,7 +399,7 @@ void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator *
394399
bool HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building);
395400
void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building);
396401
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
397-
void HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
402+
void HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
398403
bool HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support);
399404
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
400405
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);

‎src/hnswinsert.c

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,7 @@ HnswLoadNeighbors(HnswElement element, Relation index, int m, int lm, int lc)
368368
* Load elements for insert
369369
*/
370370
static void
371-
LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, int *idx, Relation index, HnswSupport * support)
371+
LoadElementsForInsert(HnswNeighborArray * neighbors, HnswQuery * q, int *idx, Relation index, HnswSupport * support)
372372
{
373373
char *base = NULL;
374374

@@ -378,7 +378,7 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, int *idx, Relation
378378
HnswElement element = HnswPtrAccess(base, hc->element);
379379
double distance;
380380

381-
HnswLoadElement(element, &distance, &q, index, support, true, NULL);
381+
HnswLoadElement(element, &distance, q, index, support, true, NULL);
382382
hc->distance = distance;
383383

384384
/* Prune element if being deleted */
@@ -419,9 +419,11 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
419419
idx = -2;
420420
else
421421
{
422-
Datum q = HnswGetValue(base, element);
422+
HnswQuery q;
423423

424-
LoadElementsForInsert(neighbors, q, &idx, index, support);
424+
q.value = HnswGetValue(base, element);
425+
426+
LoadElementsForInsert(neighbors, &q, &idx, index, support);
425427

426428
if (idx == -1)
427429
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, support);

‎src/hnswscan.c

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
* Algorithm 5 from paper
1212
*/
1313
static List *
14-
GetScanItems(IndexScanDesc scan, Datum q)
14+
GetScanItems(IndexScanDesc scan, Datum value)
1515
{
1616
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
1717
Relation index = scan->indexRelation;
@@ -21,22 +21,25 @@ GetScanItems(IndexScanDesc scan, Datum q)
2121
int m;
2222
HnswElement entryPoint;
2323
char *base = NULL;
24+
HnswQuery q;
25+
26+
q.value = value;
2427

2528
/* Get m and entry point */
2629
HnswGetMetaPageInfo(index, &m, &entryPoint);
2730

2831
if (entryPoint == NULL)
2932
return NIL;
3033

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

3336
for (int lc = entryPoint->level; lc >= 1; lc--)
3437
{
35-
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL);
38+
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, false, NULL);
3639
ep = w;
3740
}
3841

39-
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL);
42+
return HnswSearchLayer(base, &q, ep, hnsw_ef_search, 0, index, support, m, false, NULL);
4043
}
4144

4245
/*

‎src/hnswutils.c

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -533,7 +533,7 @@ HnswGetDistance(Datum a, Datum b, HnswSupport * support)
533533
* Load an element and optionally get its distance from q
534534
*/
535535
static void
536-
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
536+
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
537537
{
538538
Buffer buf;
539539
Page page;
@@ -551,10 +551,10 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Dat
551551
/* Calculate distance */
552552
if (distance != NULL)
553553
{
554-
if (DatumGetPointer(*q) == NULL)
554+
if (DatumGetPointer(q->value) == NULL)
555555
*distance = 0;
556556
else
557-
*distance = HnswGetDistance(*q, PointerGetDatum(&etup->data), support);
557+
*distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support);
558558
}
559559

560560
/* Load element */
@@ -573,7 +573,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Dat
573573
* Load an element and optionally get its distance from q
574574
*/
575575
void
576-
HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
576+
HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
577577
{
578578
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element);
579579
}
@@ -582,18 +582,18 @@ HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index,
582582
* Get the distance for an element
583583
*/
584584
static double
585-
GetElementDistance(char *base, HnswElement element, Datum q, HnswSupport * support)
585+
GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport * support)
586586
{
587587
Datum value = HnswGetValue(base, element);
588588

589-
return HnswGetDistance(q, value, support);
589+
return HnswGetDistance(q->value, value, support);
590590
}
591591

592592
/*
593593
* Create a candidate for the entry point
594594
*/
595595
HnswSearchCandidate *
596-
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, HnswSupport * support, bool loadVec)
596+
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
597597
{
598598
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
599599
bool inMemory = index == NULL;
@@ -602,7 +602,7 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
602602
if (inMemory)
603603
sc->distance = GetElementDistance(base, entryPoint, q, support);
604604
else
605-
HnswLoadElement(entryPoint, &sc->distance, &q, index, support, loadVec, NULL);
605+
HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL);
606606
return sc;
607607
}
608608

@@ -791,7 +791,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
791791
* Algorithm 2 from paper
792792
*/
793793
List *
794-
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement)
794+
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool i EED3 nserting, HnswElement skipElement)
795795
{
796796
List *w = NIL;
797797
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -873,7 +873,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, H
873873

874874
/* Avoid any allocations if not adding */
875875
eElement = NULL;
876-
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, support, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
876+
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
877877

878878
if (eElement == NULL)
879879
continue;
@@ -1220,10 +1220,12 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
12201220
List *w;
12211221
int level = element->level;
12221222
int entryLevel;
1223-
Datum q = HnswGetValue(base, element);
1223+
HnswQuery q;
12241224
HnswElement skipElement = existing ? element : NULL;
12251225
bool inMemory = index == NULL;
12261226

1227+
q.value = HnswGetValue(base, element);
1228+
12271229
/* Precompute hash */
12281230
if (inMemory)
12291231
PrecomputeHash(base, element);
@@ -1233,13 +1235,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
12331235
return;
12341236

12351237
/* Get entry point and level */
1236-
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, true));
1238+
ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true));
12371239
entryLevel = entryPoint->level;
12381240

12391241
/* 1st phase: greedy search to insert level */
12401242
for (int lc = entryLevel; lc >= level + 1; lc--)
12411243
{
1242-
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, true, skipElement);
1244+
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement);
12431245
ep = w;
12441246
}
12451247

@@ -1258,7 +1260,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
12581260
List *lw = NIL;
12591261
ListCell *lc2;
12601262

1261-
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, support, m, true, skipElement);
1263+
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement);
12621264

12631265
/* Convert search candidates to candidates */
12641266
foreach(lc2, w)

0 commit comments

Comments
 (0)
0