mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 20:15:46 +08:00
Compare commits
5 Commits
time-budge
...
ivfflat-se
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
61b2ddeb86 | ||
|
|
97cf990e0f | ||
|
|
55dc735e1a | ||
|
|
be4e9a9df2 | ||
|
|
d5e8fc96a5 |
@@ -18,7 +18,6 @@
|
|||||||
#endif
|
#endif
|
||||||
|
|
||||||
int hnsw_ef_search;
|
int hnsw_ef_search;
|
||||||
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,10 +68,6 @@ 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);
|
||||||
|
|
||||||
DefineCustomIntVariable("hnsw.time_budget", "Sets the time budget for search",
|
|
||||||
NULL, &hnsw_time_budget,
|
|
||||||
-1, -1, INT_MAX, PGC_USERSET, GUC_UNIT_MS, NULL, NULL, NULL);
|
|
||||||
|
|
||||||
MarkGUCPrefixReserved("hnsw");
|
MarkGUCPrefixReserved("hnsw");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
13
src/hnsw.h
13
src/hnsw.h
@@ -106,7 +106,6 @@
|
|||||||
|
|
||||||
/* Variables */
|
/* Variables */
|
||||||
extern int hnsw_ef_search;
|
extern int hnsw_ef_search;
|
||||||
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;
|
||||||
@@ -156,12 +155,13 @@ struct HnswNeighborArray
|
|||||||
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
|
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
|
||||||
};
|
};
|
||||||
|
|
||||||
typedef struct HnswPairingHeapNode
|
typedef struct HnswSearchCandidate
|
||||||
{
|
{
|
||||||
HnswCandidate *inner;
|
|
||||||
pairingheap_node c_node;
|
pairingheap_node c_node;
|
||||||
pairingheap_node w_node;
|
pairingheap_node w_node;
|
||||||
} HnswPairingHeapNode;
|
HnswElementPtr element;
|
||||||
|
float distance;
|
||||||
|
} HnswSearchCandidate;
|
||||||
|
|
||||||
/* HNSW index options */
|
/* HNSW index options */
|
||||||
typedef struct HnswOptions
|
typedef struct HnswOptions
|
||||||
@@ -331,7 +331,6 @@ typedef struct HnswScanOpaqueData
|
|||||||
bool first;
|
bool first;
|
||||||
List *w;
|
List *w;
|
||||||
MemoryContext tmpCtx;
|
MemoryContext tmpCtx;
|
||||||
instr_time start;
|
|
||||||
|
|
||||||
/* Support functions */
|
/* Support functions */
|
||||||
FmgrInfo *procinfo;
|
FmgrInfo *procinfo;
|
||||||
@@ -376,14 +375,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, instr_time *start);
|
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
|
||||||
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);
|
||||||
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 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);
|
||||||
|
|||||||
@@ -3,7 +3,6 @@
|
|||||||
#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"
|
||||||
@@ -23,7 +22,6 @@ 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);
|
||||||
@@ -35,11 +33,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, start);
|
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
|
||||||
ep = w;
|
ep = w;
|
||||||
}
|
}
|
||||||
|
|
||||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, start);
|
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
|
||||||
}
|
}
|
||||||
|
|
||||||
/*
|
/*
|
||||||
@@ -92,9 +90,6 @@ 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;
|
||||||
@@ -166,14 +161,14 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
|||||||
so->first = false;
|
so->first = false;
|
||||||
|
|
||||||
#if defined(HNSW_MEMORY)
|
#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
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
while (list_length(so->w) > 0)
|
while (list_length(so->w) > 0)
|
||||||
{
|
{
|
||||||
char *base = NULL;
|
char *base = NULL;
|
||||||
HnswCandidate *hc = llast(so->w);
|
HnswSearchCandidate *hc = llast(so->w);
|
||||||
HnswElement element = HnswPtrAccess(base, hc->element);
|
HnswElement element = HnswPtrAccess(base, hc->element);
|
||||||
ItemPointer heaptid;
|
ItemPointer heaptid;
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,6 @@
|
|||||||
#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"
|
||||||
@@ -609,10 +608,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
|
|||||||
/*
|
/*
|
||||||
* Create a candidate for the entry point
|
* Create a candidate for the entry point
|
||||||
*/
|
*/
|
||||||
HnswCandidate *
|
HnswSearchCandidate *
|
||||||
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)
|
||||||
{
|
{
|
||||||
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
|
HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate));
|
||||||
|
|
||||||
HnswPtrStore(base, hc->element, entryPoint);
|
HnswPtrStore(base, hc->element, entryPoint);
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
@@ -622,8 +621,8 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
|
|||||||
return hc;
|
return hc;
|
||||||
}
|
}
|
||||||
|
|
||||||
#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner)
|
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
|
||||||
#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner)
|
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Compare candidate distances
|
* Compare candidate distances
|
||||||
@@ -631,10 +630,10 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
|
|||||||
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 (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance)
|
if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance)
|
||||||
return 1;
|
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 -1;
|
||||||
|
|
||||||
return 0;
|
return 0;
|
||||||
@@ -646,27 +645,15 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
|
|||||||
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 (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance)
|
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
|
||||||
return -1;
|
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 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
|
||||||
*/
|
*/
|
||||||
@@ -800,7 +787,7 @@ 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, instr_time *start)
|
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
|
||||||
{
|
{
|
||||||
List *w = NIL;
|
List *w = NIL;
|
||||||
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
||||||
@@ -826,15 +813,13 @@ 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)
|
||||||
{
|
{
|
||||||
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
|
HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2);
|
||||||
bool found;
|
bool found;
|
||||||
HnswPairingHeapNode *node;
|
|
||||||
|
|
||||||
AddToVisited(base, &v, hc->element, index, &found);
|
AddToVisited(base, &v, hc->element, index, &found);
|
||||||
|
|
||||||
node = CreatePairingHeapNode(hc);
|
pairingheap_add(C, &hc->c_node);
|
||||||
pairingheap_add(C, &node->c_node);
|
pairingheap_add(W, &hc->w_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
|
||||||
@@ -847,21 +832,10 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
|
|
||||||
while (!pairingheap_is_empty(C))
|
while (!pairingheap_is_empty(C))
|
||||||
{
|
{
|
||||||
HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C));
|
HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C));
|
||||||
HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
|
HnswSearchCandidate *f = HnswGetSearchCandidate(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;
|
||||||
|
|
||||||
@@ -875,12 +849,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
for (int i = 0; i < unvisitedLength; i++)
|
for (int i = 0; i < unvisitedLength; i++)
|
||||||
{
|
{
|
||||||
HnswElement eElement;
|
HnswElement eElement;
|
||||||
HnswCandidate *e;
|
HnswSearchCandidate *e;
|
||||||
HnswPairingHeapNode *node;
|
|
||||||
float eDistance;
|
float eDistance;
|
||||||
bool alwaysAdd = wlen < ef;
|
bool alwaysAdd = wlen < ef;
|
||||||
|
|
||||||
f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
|
f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
|
||||||
|
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
{
|
{
|
||||||
@@ -911,13 +884,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
|||||||
continue;
|
continue;
|
||||||
|
|
||||||
/* Create a new candidate */
|
/* Create a new candidate */
|
||||||
e = palloc(sizeof(HnswCandidate));
|
e = palloc(sizeof(HnswSearchCandidate));
|
||||||
HnswPtrStore(base, e->element, eElement);
|
HnswPtrStore(base, e->element, eElement);
|
||||||
e->distance = eDistance;
|
e->distance = eDistance;
|
||||||
|
pairingheap_add(C, &e->c_node);
|
||||||
node = CreatePairingHeapNode(e);
|
pairingheap_add(W, &e->w_node);
|
||||||
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.
|
||||||
@@ -938,7 +909,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))
|
||||||
{
|
{
|
||||||
HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W));
|
HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
|
||||||
|
|
||||||
w = lappend(w, hc);
|
w = lappend(w, hc);
|
||||||
}
|
}
|
||||||
@@ -1303,7 +1274,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);
|
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
|
||||||
ep = w;
|
ep = w;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1319,16 +1290,27 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
|||||||
{
|
{
|
||||||
int lm = HnswGetLayerM(m, lc);
|
int lm = HnswGetLayerM(m, lc);
|
||||||
List *neighbors;
|
List *neighbors;
|
||||||
List *lw;
|
List *lw = NIL;
|
||||||
|
ListCell *lc2;
|
||||||
|
|
||||||
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL);
|
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement);
|
||||||
|
|
||||||
|
/* 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, w, skipElement);
|
lw = RemoveElements(base, lw, skipElement);
|
||||||
else
|
|
||||||
lw = w;
|
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Candidates are sorted, but not deterministically. Could set
|
* Candidates are sorted, but not deterministically. Could set
|
||||||
|
|||||||
@@ -57,6 +57,39 @@ ivfflatbuildphasename(int64 phasenum)
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/*
|
||||||
|
* Estimate the number of probes needed to get requested tuples
|
||||||
|
*/
|
||||||
|
static int
|
||||||
|
EstimateProbes(PlannerInfo *root, IndexPath *path, int lists)
|
||||||
|
{
|
||||||
|
double selectivity = 1;
|
||||||
|
ListCell *lc;
|
||||||
|
double tuplesPerList;
|
||||||
|
|
||||||
|
/* Cannot estimate without limit */
|
||||||
|
/* limit_tuples includes offset */
|
||||||
|
if (root->limit_tuples < 0)
|
||||||
|
return 0;
|
||||||
|
|
||||||
|
/* Get the selectivity of non-index conditions */
|
||||||
|
foreach(lc, path->indexinfo->indrestrictinfo)
|
||||||
|
{
|
||||||
|
RestrictInfo *rinfo = lfirst(lc);
|
||||||
|
|
||||||
|
/* Skip DEFAULT_INEQ_SEL since it may be a distance filter */
|
||||||
|
if (rinfo->norm_selec >= 0 && rinfo->norm_selec <= 1 && rinfo->norm_selec != (Selectivity) DEFAULT_INEQ_SEL)
|
||||||
|
selectivity *= rinfo->norm_selec;
|
||||||
|
}
|
||||||
|
|
||||||
|
tuplesPerList = path->indexinfo->tuples * selectivity / (double) lists;
|
||||||
|
if (tuplesPerList == 0)
|
||||||
|
return lists;
|
||||||
|
|
||||||
|
/* No need to cap at this point */
|
||||||
|
return root->limit_tuples / tuplesPerList;
|
||||||
|
}
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Estimate the cost of an index scan
|
* Estimate the cost of an index scan
|
||||||
*/
|
*/
|
||||||
@@ -68,6 +101,7 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
|||||||
{
|
{
|
||||||
GenericCosts costs;
|
GenericCosts costs;
|
||||||
int lists;
|
int lists;
|
||||||
|
int probes;
|
||||||
double ratio;
|
double ratio;
|
||||||
double spc_seq_page_cost;
|
double spc_seq_page_cost;
|
||||||
Relation index;
|
Relation index;
|
||||||
@@ -89,8 +123,10 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
|||||||
IvfflatGetMetaPageInfo(index, &lists, NULL);
|
IvfflatGetMetaPageInfo(index, &lists, NULL);
|
||||||
index_close(index, NoLock);
|
index_close(index, NoLock);
|
||||||
|
|
||||||
|
probes = Max(ivfflat_probes, EstimateProbes(root, path, lists));
|
||||||
|
|
||||||
/* Get the ratio of lists that we need to visit */
|
/* Get the ratio of lists that we need to visit */
|
||||||
ratio = ((double) ivfflat_probes) / lists;
|
ratio = ((double) probes) / lists;
|
||||||
if (ratio > 1.0)
|
if (ratio > 1.0)
|
||||||
ratio = 1.0;
|
ratio = 1.0;
|
||||||
|
|
||||||
|
|||||||
@@ -253,8 +253,9 @@ typedef struct IvfflatScanOpaqueData
|
|||||||
/* Sorting */
|
/* Sorting */
|
||||||
Tuplesortstate *sortstate;
|
Tuplesortstate *sortstate;
|
||||||
TupleDesc tupdesc;
|
TupleDesc tupdesc;
|
||||||
TupleTableSlot *slot;
|
TupleTableSlot *vslot;
|
||||||
bool isnull;
|
TupleTableSlot *mslot;
|
||||||
|
BufferAccessStrategy bas;
|
||||||
|
|
||||||
/* Support functions */
|
/* Support functions */
|
||||||
FmgrInfo *procinfo;
|
FmgrInfo *procinfo;
|
||||||
|
|||||||
@@ -15,16 +15,19 @@
|
|||||||
#include "utils/memutils.h"
|
#include "utils/memutils.h"
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
|
||||||
|
#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* Compare list distances
|
* Compare list distances
|
||||||
*/
|
*/
|
||||||
static int
|
static int
|
||||||
CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
||||||
{
|
{
|
||||||
if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance)
|
if (GetScanListConst(a)->distance > GetScanListConst(b)->distance)
|
||||||
return 1;
|
return 1;
|
||||||
|
|
||||||
if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
|
if (GetScanListConst(a)->distance < GetScanListConst(b)->distance)
|
||||||
return -1;
|
return -1;
|
||||||
|
|
||||||
return 0;
|
return 0;
|
||||||
@@ -76,14 +79,14 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
|||||||
|
|
||||||
/* Calculate max distance */
|
/* Calculate max distance */
|
||||||
if (listCount == so->probes)
|
if (listCount == so->probes)
|
||||||
maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
|
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
|
||||||
}
|
}
|
||||||
else if (distance < maxDistance)
|
else if (distance < maxDistance)
|
||||||
{
|
{
|
||||||
IvfflatScanList *scanlist;
|
IvfflatScanList *scanlist;
|
||||||
|
|
||||||
/* Remove */
|
/* Remove */
|
||||||
scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
|
scanlist = GetScanList(pairingheap_remove_first(so->listQueue));
|
||||||
|
|
||||||
/* Reuse */
|
/* Reuse */
|
||||||
scanlist->startPage = list->startPage;
|
scanlist->startPage = list->startPage;
|
||||||
@@ -91,7 +94,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
|||||||
pairingheap_add(so->listQueue, &scanlist->ph_node);
|
pairingheap_add(so->listQueue, &scanlist->ph_node);
|
||||||
|
|
||||||
/* Update max distance */
|
/* Update max distance */
|
||||||
maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
|
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -110,19 +113,12 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
|||||||
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
||||||
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
|
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
|
||||||
double tuples = 0;
|
double tuples = 0;
|
||||||
TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
|
TupleTableSlot *slot = so->vslot;
|
||||||
|
|
||||||
/*
|
|
||||||
* Reuse same set of shared buffers for scan
|
|
||||||
*
|
|
||||||
* See postgres/src/backend/storage/buffer/README for description
|
|
||||||
*/
|
|
||||||
BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
|
|
||||||
|
|
||||||
/* Search closest probes lists */
|
/* Search closest probes lists */
|
||||||
while (!pairingheap_is_empty(so->listQueue))
|
while (!pairingheap_is_empty(so->listQueue))
|
||||||
{
|
{
|
||||||
BlockNumber searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
|
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
|
||||||
|
|
||||||
/* Search all entry pages for list */
|
/* Search all entry pages for list */
|
||||||
while (BlockNumberIsValid(searchPage))
|
while (BlockNumberIsValid(searchPage))
|
||||||
@@ -131,7 +127,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
|||||||
Page page;
|
Page page;
|
||||||
OffsetNumber maxoffno;
|
OffsetNumber maxoffno;
|
||||||
|
|
||||||
buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
|
buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, so->bas);
|
||||||
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
||||||
page = BufferGetPage(buf);
|
page = BufferGetPage(buf);
|
||||||
maxoffno = PageGetMaxOffsetNumber(page);
|
maxoffno = PageGetMaxOffsetNumber(page);
|
||||||
@@ -170,8 +166,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
FreeAccessStrategy(bas);
|
|
||||||
|
|
||||||
if (tuples < 100)
|
if (tuples < 100)
|
||||||
ereport(DEBUG1,
|
ereport(DEBUG1,
|
||||||
(errmsg("index scan found few tuples"),
|
(errmsg("index scan found few tuples"),
|
||||||
@@ -274,7 +268,16 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
|
|||||||
/* Prep sort */
|
/* Prep sort */
|
||||||
so->sortstate = InitScanSortState(so->tupdesc);
|
so->sortstate = InitScanSortState(so->tupdesc);
|
||||||
|
|
||||||
so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
|
/* Need separate slots for puttuple and gettuple */
|
||||||
|
so->vslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
|
||||||
|
so->mslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
|
||||||
|
|
||||||
|
/*
|
||||||
|
* Reuse same set of shared buffers for scan
|
||||||
|
*
|
||||||
|
* See postgres/src/backend/storage/buffer/README for description
|
||||||
|
*/
|
||||||
|
so->bas = GetAccessStrategy(BAS_BULKREAD);
|
||||||
|
|
||||||
so->listQueue = pairingheap_allocate(CompareLists, scan);
|
so->listQueue = pairingheap_allocate(CompareLists, scan);
|
||||||
|
|
||||||
@@ -348,9 +351,10 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
|
|||||||
pfree(DatumGetPointer(value));
|
pfree(DatumGetPointer(value));
|
||||||
}
|
}
|
||||||
|
|
||||||
if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
|
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
|
||||||
{
|
{
|
||||||
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
|
bool isnull;
|
||||||
|
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
|
||||||
|
|
||||||
scan->xs_heaptid = *heaptid;
|
scan->xs_heaptid = *heaptid;
|
||||||
scan->xs_recheck = false;
|
scan->xs_recheck = false;
|
||||||
@@ -371,6 +375,10 @@ ivfflatendscan(IndexScanDesc scan)
|
|||||||
|
|
||||||
pairingheap_free(so->listQueue);
|
pairingheap_free(so->listQueue);
|
||||||
tuplesort_end(so->sortstate);
|
tuplesort_end(so->sortstate);
|
||||||
|
FreeAccessStrategy(so->bas);
|
||||||
|
FreeTupleDesc(so->tupdesc);
|
||||||
|
|
||||||
|
/* TODO Free vslot and mslot without freeing TupleDesc */
|
||||||
|
|
||||||
pfree(so);
|
pfree(so);
|
||||||
scan->opaque = NULL;
|
scan->opaque = NULL;
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ use PostgreSQL::Test::Utils;
|
|||||||
use Test::More;
|
use Test::More;
|
||||||
|
|
||||||
my $dim = 3;
|
my $dim = 3;
|
||||||
my $nc = 50;
|
my $nc = 100;
|
||||||
my $limit = 20;
|
my $limit = 20;
|
||||||
|
|
||||||
my $array_sql = join(",", ('random()') x $dim);
|
my $array_sql = join(",", ('random()') x $dim);
|
||||||
@@ -102,15 +102,13 @@ $node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
|
|||||||
$explain = $node->safe_psql("postgres", qq(
|
$explain = $node->safe_psql("postgres", qq(
|
||||||
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
|
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
|
||||||
));
|
));
|
||||||
# TODO Use attribute index
|
like($explain, qr/Index Scan on attribute_idx/);
|
||||||
like($explain, qr/Index Scan using idx/);
|
|
||||||
|
|
||||||
# Test partial index
|
# Test partial index
|
||||||
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5) WHERE (c = $c);");
|
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5) WHERE (c = $c);");
|
||||||
$explain = $node->safe_psql("postgres", qq(
|
$explain = $node->safe_psql("postgres", qq(
|
||||||
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
|
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
|
||||||
));
|
));
|
||||||
# TODO Use partial index
|
like($explain, qr/Index Scan using partial_idx/);
|
||||||
like($explain, qr/Index Scan using idx/);
|
|
||||||
|
|
||||||
done_testing();
|
done_testing();
|
||||||
|
|||||||
Reference in New Issue
Block a user