Compare commits

..

31 Commits

Author SHA1 Message Date
Andrew Kane
70785ff0fe Merge branch 'master' into hqann2 2024-10-10 19:06:26 -07:00
Andrew Kane
35b252a3e3 Switched to macos-13 on CI [skip ci] 2024-10-10 18:33:35 -07:00
Andrew Kane
2832e746f0 Use get_float8_infinity instead of INFINITY
Co-authored-by: "Jonathan S. Katz" <jkatz@amazon.com>
2024-10-10 18:16:39 -07:00
Andrew Kane
961cb17d80 Added iterative search for HNSW [skip ci] 2024-10-10 18:14:39 -07:00
Andrew Kane
c91ed7b2c3 Added iterative search for IVFFlat [skip ci] 2024-10-10 18:12:27 -07:00
Andrew Kane
48fe70c219 Fixed sortdesc for multiple attributes [skip ci] 2024-10-10 15:22:23 -07:00
Andrew Kane
29908405ab List assignment does not change with scaling [skip ci] 2024-10-10 15:16:08 -07:00
Andrew Kane
08d0340655 Improved IVFFlat vacuum test [skip ci] 2024-10-10 14:24:26 -07:00
Andrew Kane
7d2eb49c2a Fixed list assignment for cosine distance and non-normalized vectors for IVFFlat 2024-10-10 12:54:13 -07:00
Andrew Kane
772ab69de6 Updated IVFFlat to support multiple attributes (not enabled yet) 2024-10-10 12:34:03 -07:00
Andrew Kane
e13e9a9614 Improved scoping [skip ci] 2024-10-10 12:19:08 -07:00
Andrew Kane
e2fab306ac Added tupdesc to buildstate [skip ci] 2024-10-10 12:18:41 -07:00
Andrew Kane
edc2126a4a Renamed tupdesc to sortdesc [skip ci] 2024-10-10 12:16:57 -07:00
Andrew Kane
2cbd08b6c0 Moved unions and macros [skip ci] 2024-10-10 09:41:26 -07:00
Andrew Kane
a73fe9e9e1 Use same TupleDesc [skip ci] 2024-10-10 02:01:23 -07:00
Andrew Kane
2ffbb2500c Merge branch 'master' into hqann2 2024-10-09 23:59:11 -07:00
Andrew Kane
fa6782985a Added HnswQuery struct for query data 2024-10-09 23:45:47 -07:00
Andrew Kane
46209740ce Merge branch 'master' into hqann2 2024-10-09 23:26:57 -07:00
Andrew Kane
32ab27d72a Added HnswSupport struct for support functions 2024-10-09 23:10:26 -07:00
Andrew Kane
d650d23da8 Merge branch 'master' into hqann2 2024-10-09 22:06:29 -07:00
Andrew Kane
064db12de7 Moved procinfo initialization for inserts [skip ci] 2024-10-09 21:59:21 -07:00
Andrew Kane
ad755d19ca Merge branch 'master' into hqann2 2024-10-09 21:52:45 -07:00
Andrew Kane
45a6eef9e0 Improved variable name [skip ci] 2024-10-09 21:52:10 -07:00
Andrew Kane
f39b515ce7 Merge branch 'master' into hqann2 2024-10-09 21:50:53 -07:00
Andrew Kane
17266ed409 Use inMemory for conditionals 2024-10-09 21:49:32 -07:00
Andrew Kane
ab88612c83 Always use tuple for duplicates [skip ci] 2024-10-09 21:30:10 -07:00
Andrew Kane
356f175359 Merge branch 'master' into hqann2 2024-10-09 21:25:30 -07:00
Andrew Kane
a98534e5ab DRY HNSW procinfo 2024-10-09 21:03:18 -07:00
Andrew Kane
57c05c59a2 DRY code for forming index value 2024-10-09 20:50:17 -07:00
Andrew Kane
a49a2b9480 Improved test [skip ci] 2024-10-09 19:54:06 -07:00
Andrew Kane
3b61384367 Added HnswGetElementTupleSize method [skip ci] 2024-10-09 19:37:09 -07:00
21 changed files with 977 additions and 321 deletions

View File

@@ -49,7 +49,7 @@ jobs:
- postgres: 16 - postgres: 16
os: macos-14 os: macos-14
- postgres: 14 - postgres: 14
os: macos-12 os: macos-13
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1

View File

@@ -1,6 +1,7 @@
## 0.8.0 (unreleased) ## 0.8.0 (unreleased)
- Added support for inline filtering with HNSW - Added support for inline filtering with HNSW
- Added support for iterative index scans
- Added casts for arrays to `sparsevec` - Added casts for arrays to `sparsevec`
- Improved cost estimation - Improved cost estimation
- Improved performance of HNSW inserts and on-disk index builds - Improved performance of HNSW inserts and on-disk index builds

View File

@@ -18,7 +18,16 @@
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x) #define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif #endif
static const struct config_enum_entry hnsw_iterative_search_options[] = {
{"off", HNSW_ITERATIVE_SEARCH_OFF, false},
{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
{NULL, 0, false}
};
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_iterative_search_max_tuples;
int hnsw_iterative_search;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -69,6 +78,15 @@ 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);
DefineCustomEnumVariable("hnsw.iterative_search", "Sets iterative search",
NULL, &hnsw_iterative_search,
HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Ensure ivfflat.max_probes uses same value for "all" */
DefineCustomIntVariable("hnsw.iterative_search_max_tuples", "Sets the max number of candidates to visit for iterative search",
"-1 means all", &hnsw_iterative_search_max_tuples,
-1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw"); MarkGUCPrefixReserved("hnsw");
} }

View File

@@ -89,6 +89,9 @@
/* Ensure fits on page and in uint8 */ /* Ensure fits on page and in uint8 */
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255) #define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255)
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
#define HnswGetValue(base, element) PointerGetDatum(HnswPtrAccess(base, (element)->value)) #define HnswGetValue(base, element) PointerGetDatum(HnswPtrAccess(base, (element)->value))
#if PG_VERSION_NUM < 140005 #if PG_VERSION_NUM < 140005
@@ -109,8 +112,17 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_iterative_search;
extern int hnsw_iterative_search_max_tuples;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeSearchType
{
HNSW_ITERATIVE_SEARCH_OFF,
HNSW_ITERATIVE_SEARCH_RELAXED,
HNSW_ITERATIVE_SEARCH_STRICT
} HnswIterativeSearchType;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;
@@ -133,6 +145,7 @@ 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;
@@ -243,6 +256,20 @@ typedef struct HnswTypeInfo
void (*checkValue) (Pointer v); void (*checkValue) (Pointer v);
} HnswTypeInfo; } HnswTypeInfo;
typedef struct HnswSupport
{
FmgrInfo *procinfo[2];
FmgrInfo *normprocinfo;
Oid *collation;
} HnswSupport;
typedef struct HnswQuery
{
Datum value;
IndexTuple itup;
ScanKeyData *keyData;
} HnswQuery;
typedef struct HnswBuildState typedef struct HnswBuildState
{ {
/* Info */ /* Info */
@@ -262,9 +289,7 @@ typedef struct HnswBuildState
double reltuples; double reltuples;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo[2]; HnswSupport support;
FmgrInfo *normprocinfo;
Oid *collation;
/* Variables */ /* Variables */
HnswGraph graphData; HnswGraph graphData;
@@ -314,10 +339,10 @@ typedef struct HnswElementTupleData
uint8 type; uint8 type;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint8 unused; uint8 version;
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused2; uint16 unused;
Vector data; Vector data;
} HnswElementTupleData; } HnswElementTupleData;
@@ -326,24 +351,41 @@ typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData typedef struct HnswNeighborTupleData
{ {
uint8 type; uint8 type;
uint8 unused; uint8 version;
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 union
{
HnswElement element;
ItemPointerData indextid;
} HnswUnvisited;
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;
HnswQuery q;
int m;
int64 tuples;
double previousDistance;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo[2]; HnswSupport support;
FmgrInfo *normprocinfo;
Oid *collation;
} HnswScanOpaqueData; } HnswScanOpaqueData;
typedef HnswScanOpaqueData * HnswScanOpaque; typedef HnswScanOpaqueData * HnswScanOpaque;
@@ -361,8 +403,7 @@ typedef struct HnswVacuumState
int efConstruction; int efConstruction;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo[2]; HnswSupport support;
Oid *collation;
/* Variables */ /* Variables */
struct tidhash_hash *deleted; struct tidhash_hash *deleted;
@@ -378,35 +419,36 @@ typedef struct HnswVacuumState
int HnswGetM(Relation index); int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index); int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
void HnswInitSupport(HnswSupport * support, Relation index);
Datum HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value); Datum HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value);
bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value); bool HnswCheckNorm(HnswSupport * support, 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, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, int m, bool inserting, HnswElement skipElement, bool inMemory); List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
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, bool inMemory); void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation rel, FmgrInfo **procinfo, Oid *collation, bool loadVec, bool inMemory); HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec, bool inMemory);
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);
HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator); HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc); void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building); bool HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc);
void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m, bool checkExisting, bool building); void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index);
void HnswLoadElement(HnswElement element, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance); void HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple); void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple);
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo **procinfo, Oid *collation); void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);
bool HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc);
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc); bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
void HnswInitLockTranche(void); void HnswInitLockTranche(void);
const HnswTypeInfo *HnswGetTypeInfo(Relation index); const HnswTypeInfo *HnswGetTypeInfo(Relation index);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc); PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
void HnswInitProcinfo(FmgrInfo **procinfo, Oid **collation, Relation index);
Size HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple); Size HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple);
bool HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc); bool HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc);

View File

@@ -330,7 +330,6 @@ static bool
FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc) FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc)
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup); IndexTuple itup = HnswPtrAccess(base, element->itup);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
@@ -338,18 +337,9 @@ FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, Tuple
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element); HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
if (useIndexTuple)
{
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc)) if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false; return false;
}
else
{
/* Exit early since ordered by distance */
if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
return false;
}
/* Check for space */ /* Check for space */
if (AddDuplicateInMemory(element, neighborElement)) if (AddDuplicateInMemory(element, neighborElement))
@@ -375,7 +365,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
* Update neighbors * Update neighbors
*/ */
static void static void
UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m) UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswElement e, int m)
{ {
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
@@ -397,7 +387,7 @@ UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *co
Assert(neighborElement); Assert(neighborElement);
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE); LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, procinfo, collation); HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, support);
LWLockRelease(&neighborElement->lock); LWLockRelease(&neighborElement->lock);
} }
} }
@@ -407,7 +397,7 @@ UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *co
* Update graph in memory * Update graph in memory
*/ */
static void static void
UpdateGraphInMemory(FmgrInfo **procinfo, Oid *collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate) UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
{ {
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
@@ -420,7 +410,7 @@ UpdateGraphInMemory(FmgrInfo **procinfo, Oid *collation, HnswElement element, in
AddElementInMemory(base, graph, element); AddElementInMemory(base, graph, element);
/* Update neighbors */ /* Update neighbors */
UpdateNeighborsInMemory(base, buildstate->index, procinfo, collation, element, m); UpdateNeighborsInMemory(base, buildstate->index, support, element, m);
/* Update entry point if needed (already have lock) */ /* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -434,9 +424,8 @@ static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element) InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
FmgrInfo **procinfo = buildstate->procinfo;
Oid *collation = buildstate->collation;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswSupport *support = &buildstate->support;
HnswElement entryPoint; HnswElement entryPoint;
LWLock *entryLock = &graph->entryLock; LWLock *entryLock = &graph->entryLock;
LWLock *entryWaitLock = &graph->entryWaitLock; LWLock *entryWaitLock = &graph->entryWaitLock;
@@ -468,10 +457,10 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false, true); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, true);
/* Update graph in memory */ /* Update graph in memory */
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate); UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate);
/* Release entry lock */ /* Release entry lock */
LWLockRelease(entryLock); LWLockRelease(entryLock);
@@ -483,38 +472,24 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
static bool static bool
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate) InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
{ {
const HnswTypeInfo *typeInfo = buildstate->typeInfo;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswElement element; HnswElement element;
HnswAllocator *allocator = &buildstate->allocator; HnswAllocator *allocator = &buildstate->allocator;
Size valueSize; HnswSupport *support = &buildstate->support;
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock; LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
TupleDesc tupdesc = buildstate->tupdesc; TupleDesc tupdesc = buildstate->tupdesc;
IndexTuple itup; IndexTuple itup;
Size itupSize; Size itupSize;
IndexTuple itupPtr; IndexTuple itupShared;
bool unused;
/* Detoast once for all calls */ /* Form index tuple */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); if (!HnswFormIndexTuple(&itup, values, isnull, buildstate->typeInfo, support, tupdesc))
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation[0], value))
return false; return false;
value = HnswNormValue(typeInfo, buildstate->collation[0], value); /* Get tuple size */
} itupSize = IndexTupleSize(itup);
/* Get datum size */
valueSize = VARSIZE_ANY(DatumGetPointer(value));
/* Ensure graph not flushed when inserting */ /* Ensure graph not flushed when inserting */
LWLockAcquire(flushLock, LW_SHARED); LWLockAcquire(flushLock, LW_SHARED);
@@ -524,7 +499,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
{ {
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true); return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
} }
/* /*
@@ -556,22 +531,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true); return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
} }
/* Ok, we can proceed to allocate the element */ /* Ok, we can proceed to allocate the element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator); element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
itupShared = HnswAlloc(allocator, itupSize);
if (useIndexTuple)
{
/* TODO fix */
values[0] = value;
itup = index_form_tuple(tupdesc, values, isnull);
itupSize = IndexTupleSize(itup);
itupPtr = HnswAlloc(allocator, itupSize);
}
else
valuePtr = HnswAlloc(allocator, valueSize);
/* /*
* We have now allocated the space needed for the element, so we don't * We have now allocated the space needed for the element, so we don't
@@ -580,20 +545,10 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
*/ */
LWLockRelease(&graph->allocatorLock); LWLockRelease(&graph->allocatorLock);
/* Copy the datum */ /* Copy the tuple */
if (useIndexTuple) memcpy(itupShared, itup, itupSize);
{ HnswPtrStore(base, element->itup, itupShared);
bool unused; HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupShared, 1, tupdesc, &unused)));
memcpy(itupPtr, itup, itupSize);
HnswPtrStore(base, element->itup, itupPtr);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupPtr, 1, tupdesc, &unused)));
}
else
{
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
}
/* Create a lock for the element */ /* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id); LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
@@ -753,8 +708,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->indtuples = 0; buildstate->indtuples = 0;
/* Get support functions */ /* Get support functions */
HnswInitProcinfo(buildstate->procinfo, &buildstate->collation, index); HnswInitSupport(&buildstate->support, index);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
InitGraph(&buildstate->graphData, NULL, (Size) maintenance_work_mem * 1024L); InitGraph(&buildstate->graphData, NULL, (Size) maintenance_work_mem * 1024L);
buildstate->graph = &buildstate->graphData; buildstate->graph = &buildstate->graphData;

View File

@@ -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) 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)
{ {
OffsetNumber offno; OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page); OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -98,6 +98,7 @@ 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)
@@ -153,6 +154,7 @@ 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;
bool useIndexTuple = HnswUseIndexTuple(index); bool useIndexTuple = HnswUseIndexTuple(index);
@@ -203,7 +205,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)) if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
{ {
if (nbuf != buf) if (nbuf != buf)
{ {
@@ -213,6 +215,10 @@ 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;
} }
@@ -369,7 +375,7 @@ HnswLoadNeighbors(HnswElement element, Relation index, int m, int lm, int lc)
* Load elements for insert * Load elements for insert
*/ */
static void static void
LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, int *idx, Relation index, FmgrInfo **procinfo, Oid *collation) LoadElementsForInsert(HnswNeighborArray * neighbors, HnswQuery * q, int *idx, Relation index, HnswSupport * support)
{ {
char *base = NULL; char *base = NULL;
@@ -380,7 +386,7 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, i
double distance; double distance;
bool matches; bool matches;
HnswLoadElement(element, &distance, &matches, &q, qtup, NULL, index, procinfo, collation, true, NULL); HnswLoadElement(element, &distance, &matches, q, index, support, true, NULL);
hc->distance = distance; hc->distance = distance;
/* Prune element if being deleted */ /* Prune element if being deleted */
@@ -396,7 +402,7 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, i
* Get update index * Get update index
*/ */
static int static int
GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, MemoryContext updateCtx) GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, HnswSupport * support, MemoryContext updateCtx)
{ {
char *base = NULL; char *base = NULL;
int idx = -1; int idx = -1;
@@ -421,13 +427,16 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
idx = -2; idx = -2;
else else
{ {
Datum q = HnswGetValue(base, element); HnswQuery q;
IndexTuple qtup = HnswPtrAccess(base, element->itup);;
LoadElementsForInsert(neighbors, q, qtup, &idx, index, procinfo, collation); q.value = HnswGetValue(base, element);
q.itup = HnswPtrAccess(base, element->itup);
q.keyData = NULL;
LoadElementsForInsert(neighbors, &q, &idx, index, support);
if (idx == -1) if (idx == -1)
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, procinfo, collation); HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, support);
} }
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -532,7 +541,7 @@ UpdateNeighborOnDisk(HnswElement element, HnswElement newElement, int idx, int m
* Update neighbors * Update neighbors
*/ */
void void
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m, bool checkExisting, bool building) HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building)
{ {
char *base = NULL; char *base = NULL;
@@ -555,7 +564,7 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, H
HnswElement neighborElement = HnswPtrAccess(base, hc->element); HnswElement neighborElement = HnswPtrAccess(base, hc->element);
int idx; int idx;
idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation, updateCtx); idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, support, updateCtx);
/* New element was not selected as a neighbor */ /* New element was not selected as a neighbor */
if (idx == -1) if (idx == -1)
@@ -628,13 +637,12 @@ AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool bu
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
FindDuplicateOnDisk(Relation index, HnswElement element, bool building) FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDesc tupdesc)
{ {
char *base = NULL; char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup); IndexTuple itup = HnswPtrAccess(base, element->itup);
TupleDesc tupdesc = RelationGetDescr(index);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
@@ -665,12 +673,12 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
* Update graph on disk * Update graph on disk
*/ */
static void static void
UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building) UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building, TupleDesc tupdesc)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Look for duplicate */ /* Look for duplicate */
if (FindDuplicateOnDisk(index, element, building)) if (FindDuplicateOnDisk(index, element, building, tupdesc))
return; return;
/* Add element */ /* Add element */
@@ -681,7 +689,7 @@ UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswEleme
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building); HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building); HnswUpdateNeighborsOnDisk(index, support, element, m, false, building);
/* Update entry point if needed */ /* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -692,18 +700,15 @@ UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswEleme
* Insert a tuple into the index * Insert a tuple into the index
*/ */
bool bool
HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building) HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc)
{ {
HnswElement entryPoint; HnswElement entryPoint;
HnswElement element; HnswElement element;
int m; int m;
int efConstruction = HnswGetEfConstruction(index); int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo[2];
Oid *collation;
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
char *base = NULL; char *base = NULL;
bool unused;
HnswInitProcinfo(procinfo, &collation, index);
/* /*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts * Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -716,25 +721,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
/* Create an element */ /* Create an element */
element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL); element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
if (HnswUseIndexTuple(index))
{
/* TODO no toast */
TupleDesc tupdesc = RelationGetDescr(index);
IndexTuple itup;
bool unused;
/* TODO fix */
values[0] = value;
itup = index_form_tuple(tupdesc, values, isnull);
HnswPtrStore(base, element->itup, itup); HnswPtrStore(base, element->itup, itup);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused))); HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
}
else
HnswPtrStore(base, element->value, DatumGetPointer(value));
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
{ {
@@ -750,10 +740,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false, false); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, false);
/* Update graph on disk */ /* Update graph on disk */
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building); UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building, tupdesc);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -765,31 +755,20 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
* Insert a tuple into the index * Insert a tuple into the index
*/ */
static void static void
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid) HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
{ {
Datum value; IndexTuple itup;
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index); const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
FmgrInfo *normprocinfo; TupleDesc tupdesc = RelationGetDescr(index);
Oid *collation = index->rd_indcollation; HnswSupport support;
/* Detoast once for all calls */ HnswInitSupport(&support, index);
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Check value */ /* Form index tuple */
if (typeInfo->checkValue != NULL) if (!HnswFormIndexTuple(&itup, values, isnull, typeInfo, &support, tupdesc))
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswCheckNorm(normprocinfo, collation[0], value))
return; return;
value = HnswNormValue(typeInfo, collation[0], value); HnswInsertTupleOnDisk(index, &support, itup, heaptid, false, tupdesc);
}
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false);
} }
/* /*

View File

@@ -5,41 +5,78 @@
#include "pgstat.h" #include "pgstat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "storage/lmgr.h" #include "storage/lmgr.h"
#include "utils/float.h"
#include "utils/memutils.h" #include "utils/memutils.h"
/* /*
* Algorithm 5 from paper * Algorithm 5 from paper
*/ */
static List * static List *
GetScanItems(IndexScanDesc scan, Datum q) GetScanItems(IndexScanDesc scan, Datum value)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation; Relation index = scan->indexRelation;
FmgrInfo **procinfo = so->procinfo; HnswSupport *support = &so->support;
Oid *collation = so->collation;
List *ep; List *ep;
List *w; List *w;
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL; char *base = NULL;
bool inMemory = false; bool inMemory = false;
ScanKeyData *keyData = scan->keyData; HnswQuery *q = &so->q;
q->value = value;
q->itup = NULL;
q->keyData = scan->keyData;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
so->m = m;
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, NULL, keyData, index, procinfo, collation, false, inMemory)); ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false, inMemory));
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
{ {
w = HnswSearchLayer(base, q, NULL, keyData, ep, 1, lc, index, procinfo, collation, m, false, NULL, inMemory); w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, inMemory, NULL, NULL, true, NULL);
ep = w; ep = w;
} }
return HnswSearchLayer(base, q, NULL, keyData, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, inMemory); return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, inMemory, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
}
/*
* Resume scan at ground level with discarded candidates
*/
static List *
ResumeScanItems(IndexScanDesc scan)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
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 *sc;
if (pairingheap_is_empty(so->discarded))
break;
sc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
ep = lappend(ep, sc);
}
return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, false, &so->v, &so->discarded, false, &so->tuples);
} }
/* /*
@@ -62,8 +99,8 @@ GetScanValue(IndexScanDesc scan)
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value))); Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
/* Normalize if needed */ /* Normalize if needed */
if (so->normprocinfo != NULL) if (so->support.normprocinfo != NULL)
value = HnswNormValue(so->typeInfo, so->collation[0], value); value = HnswNormValue(so->typeInfo, so->support.collation[0], value);
} }
return value; return value;
@@ -83,13 +120,14 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData)); so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->typeInfo = HnswGetTypeInfo(index); so->typeInfo = HnswGetTypeInfo(index);
so->first = true; so->first = true;
so->v.tids = NULL;
so->discarded = NULL;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context", "Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
/* Set support functions */ /* Set support functions */
HnswInitProcinfo(so->procinfo, &so->collation, index); HnswInitSupport(&so->support, index);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
scan->opaque = so; scan->opaque = so;
@@ -104,7 +142,15 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
if (so->v.tids != NULL)
tidhash_reset(so->v.tids);
if (so->discarded != NULL)
pairingheap_reset(so->discarded);
so->first = true; so->first = true;
so->tuples = 0;
so->previousDistance = -get_float8_infinity();
MemoryContextReset(so->tmpCtx); MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
@@ -166,22 +212,100 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
#endif #endif
} }
while (list_length(so->w) > 0) for (;;)
{ {
char *base = NULL; char *base = NULL;
HnswSearchCandidate *hc = llast(so->w); HnswSearchCandidate *sc;
HnswElement element = HnswPtrAccess(base, hc->element); HnswElement element;
ItemPointer heaptid; ItemPointer heaptid;
if (list_length(so->w) == 0)
{
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_OFF)
break;
/* Empty index */
if (so->discarded == NULL)
break;
/* Reached max number of additional tuples */
if (hnsw_iterative_search_max_tuples != -1 && so->tuples >= hnsw_iterative_search_max_tuples)
{
if (pairingheap_is_empty(so->discarded))
break;
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
/* Prevent scans from consuming too much memory */
else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
{
if (pairingheap_is_empty(so->discarded))
{
ereport(DEBUG1,
(errmsg("hnsw index scan 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);
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;
}
sc = llast(so->w);
element = HnswPtrAccess(base, sc->element);
/* Move to next element if no valid heap TIDs */ /* Move to next element if no valid heap TIDs */
if (!hc->matches || element->heaptidsLength == 0) if (!sc->matches || 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_iterative_search != HNSW_ITERATIVE_SEARCH_OFF)
{
pfree(element);
pfree(sc);
}
continue; continue;
} }
heaptid = &element->heaptids[--element->heaptidsLength]; heaptid = &element->heaptids[--element->heaptidsLength];
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_STRICT)
{
if (sc->distance < so->previousDistance)
continue;
so->previousDistance = sc->distance;
}
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;

View File

@@ -100,19 +100,6 @@ 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
{
HnswElement element;
ItemPointerData indextid;
} HnswUnvisited;
/* /*
* Get the max number of connections in an upper layer for each element in the index * Get the max number of connections in an upper layer for each element in the index
*/ */
@@ -154,17 +141,18 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
} }
/* /*
* Init procinfo * Init support functions
*/ */
void void
HnswInitProcinfo(FmgrInfo **procinfo, Oid **collation, Relation index) HnswInitSupport(HnswSupport * support, Relation index)
{ {
procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); support->procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
if (IndexRelationGetNumberOfKeyAttributes(index) > 1) if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC); support->procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC);
*collation = index->rd_indcollation; support->collation = index->rd_indcollation;
support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
} }
/* /*
@@ -204,9 +192,9 @@ HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
* Check if non-zero norm * Check if non-zero norm
*/ */
bool bool
HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value) HnswCheckNorm(HnswSupport * support, Datum value)
{ {
return DatumGetFloat8(FunctionCall1Coll(procinfo, collation, value)) > 0; return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation[0], value)) > 0;
} }
/* /*
@@ -322,6 +310,8 @@ 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);
@@ -465,6 +455,39 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
/*
* Form index tuple
*/
bool
HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc)
{
Datum newValues[2];
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
if (support->normprocinfo != NULL)
{
if (!HnswCheckNorm(support, value))
return false;
value = HnswNormValue(typeInfo, support->collation[0], value);
}
newValues[0] = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
*out = index_form_tuple(tupdesc, newValues, isnull);
return true;
}
/* /*
* Set element tuple, except for neighbor info * Set element tuple, except for neighbor info
*/ */
@@ -474,6 +497,7 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool
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)
@@ -528,6 +552,7 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
} }
ntup->count = idx; ntup->count = idx;
ntup->version = e->version;
} }
/* /*
@@ -538,6 +563,7 @@ 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;
@@ -593,14 +619,14 @@ AttributeDistance(double e)
* Calculate the distance between values * Calculate the distance between values
*/ */
static double static double
HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool *matches) HnswGetDistance(IndexTuple itup, Datum vec, HnswQuery * q, Relation index, HnswSupport * support, bool *matches)
{ {
double g; double g;
if (DatumGetPointer(q) == NULL) if (DatumGetPointer(q->value) == NULL)
g = 0; g = 0;
else else
g = DatumGetFloat8(FunctionCall2Coll(procinfo[0], collation[0], q, vec)); g = DatumGetFloat8(FunctionCall2Coll(support->procinfo[0], support->collation[0], q->value, vec));
Assert(PointerIsValid(matches)); Assert(PointerIsValid(matches));
*matches = true; *matches = true;
@@ -611,14 +637,14 @@ HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyDat
double e = 0.0; double e = 0.0;
TupleDesc tupdesc = RelationGetDescr(index); TupleDesc tupdesc = RelationGetDescr(index);
if (keyData) if (q->keyData)
{ {
/* TODO need to pass length of key data */ /* TODO need to pass length of key data */
int keyCount = 1; int keyCount = 1;
for (int i = 0; i < keyCount; i++) for (int i = 0; i < keyCount; i++)
{ {
ScanKey key = &keyData[i]; ScanKey key = &q->keyData[i];
bool isnull; bool isnull;
Datum value = index_getattr(itup, key->sk_attno, tupdesc, &isnull); Datum value = index_getattr(itup, key->sk_attno, tupdesc, &isnull);
bool attnull = key->sk_flags & SK_ISNULL; bool attnull = key->sk_flags & SK_ISNULL;
@@ -633,7 +659,7 @@ HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyDat
} }
else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument))) else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument)))
{ {
double ei = fabs(DatumGetFloat8(FunctionCall2Coll(procinfo[key->sk_attno - 1], collation[key->sk_attno - 1], value, key->sk_argument))); double ei = fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[key->sk_attno - 1], support->collation[key->sk_attno - 1], value, key->sk_argument)));
if (ei > 0) if (ei > 0)
e += ei; e += ei;
@@ -647,7 +673,7 @@ HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyDat
return w * g + AttributeDistance(e); return w * g + AttributeDistance(e);
} }
else if (qtup) else if (q->itup)
{ {
int keyCount = IndexRelationGetNumberOfKeyAttributes(index) - 1; int keyCount = IndexRelationGetNumberOfKeyAttributes(index) - 1;
@@ -656,7 +682,7 @@ HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyDat
bool isnull; bool isnull;
bool attnull; bool attnull;
Datum value = index_getattr(itup, i + 2, tupdesc, &isnull); Datum value = index_getattr(itup, i + 2, tupdesc, &isnull);
Datum value2 = index_getattr(qtup, i + 2, tupdesc, &attnull); Datum value2 = index_getattr(q->itup, i + 2, tupdesc, &attnull);
if (isnull || attnull) if (isnull || attnull)
{ {
@@ -664,7 +690,7 @@ HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyDat
e += 1000; e += 1000;
} }
else else
e += fabs(DatumGetFloat8(FunctionCall2Coll(procinfo[i + 1], collation[i + 1], value, value2))); e += fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[i + 1], support->collation[i + 1], value, value2)));
} }
return w * g + AttributeDistance(e); return w * g + AttributeDistance(e);
@@ -678,7 +704,7 @@ HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyDat
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
static void static void
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance, HnswElement * element) HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -712,7 +738,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
value = PointerGetDatum(&etup->data); value = PointerGetDatum(&etup->data);
} }
*distance = HnswGetDistance(itup, value, *q, qtup, keyData, index, procinfo, collation, matches); *distance = HnswGetDistance(itup, value, q, index, support, matches);
} }
/* Load element */ /* Load element */
@@ -731,42 +757,39 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
void void
HnswLoadElement(HnswElement element, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance) HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
{ {
HnswLoadElementImpl(element->blkno, element->offno, distance, matches, q, qtup, keyData, index, procinfo, collation, loadVec, maxDistance, &element); HnswLoadElementImpl(element->blkno, element->offno, distance, matches, q, index, support, loadVec, maxDistance, &element);
} }
/* /*
* Get the distance for an element * Get the distance for an element
*/ */
static double static double
GetElementDistance(char *base, HnswElement element, bool *matches, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation) GetElementDistance(char *base, HnswElement element, bool *matches, HnswQuery * q, Relation index, HnswSupport * support)
{ {
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup); IndexTuple itup = HnswPtrAccess(base, element->itup);
return HnswGetDistance(itup, value, q, qtup, keyData, index, procinfo, collation, matches); return HnswGetDistance(itup, value, q, index, support, matches);
} }
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswSearchCandidate * HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, bool inMemory) HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, bool inMemory)
{ {
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate)); HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, sc->element, entryPoint); HnswPtrStore(base, sc->element, entryPoint);
if (inMemory) if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, qtup, keyData, index, procinfo, collation); sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, index, support);
else else
HnswLoadElement(entryPoint, &sc->distance, &sc->matches, &q, qtup, keyData, index, procinfo, collation, loadVec, NULL); HnswLoadElement(entryPoint, &sc->distance, &sc->matches, q, index, support, loadVec, NULL);
return sc; return sc;
} }
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
/* /*
* Compare candidate distances * Compare candidate distances
*/ */
@@ -782,6 +805,21 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
return 0; return 0;
} }
/*
* Compare discarded candidate distances
*/
static int
CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return 1;
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
return -1;
return 0;
}
/* /*
* Compare candidate distances * Compare candidate distances
*/ */
@@ -902,8 +940,11 @@ HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation i
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno)); ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
/* Ensure expected neighbors */ /*
if (ntup->count != (element->level + 2) * m) * Ensure the neighbor tuple has not been deleted or replaced between
* index scan iterations
*/
if (ntup->version != element->version || ntup->count != (element->level + 2) * m)
{ {
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
return false; return false;
@@ -949,13 +990,13 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * List *
HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, int m, bool inserting, HnswElement skipElement, bool inMemory) HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
{ {
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 v; visited_hash vh;
ListCell *lc2; ListCell *lc2;
HnswNeighborArray *localNeighborhood = NULL; HnswNeighborArray *localNeighborhood = NULL;
Size neighborhoodSize = 0; Size neighborhoodSize = 0;
@@ -963,9 +1004,21 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited)); HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength; int unvisitedLength;
uint64 additional = 0; uint64 additional = 0;
uint64 maxAdditional = keyData && lc == 0 ? 10000 : 0; uint64 maxAdditional = q->keyData && lc == 0 ? 10000 : 0;
InitVisited(base, &v, inMemory, ef, m); if (v == NULL)
{
v = &vh;
initVisited = true;
}
if (initVisited)
{
InitVisited(base, v, inMemory, 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 (inMemory) if (inMemory)
@@ -980,7 +1033,13 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2); HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2);
bool found; bool found;
AddToVisited(base, &v, sc->element, inMemory, &found); if (initVisited)
{
AddToVisited(base, v, sc->element, inMemory, &found);
if (tuples != NULL)
(*tuples)++;
}
pairingheap_add(C, &sc->c_node); pairingheap_add(C, &sc->c_node);
pairingheap_add(W, &sc->w_node); pairingheap_add(W, &sc->w_node);
@@ -1010,9 +1069,12 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
cElement = HnswPtrAccess(base, c->element); cElement = HnswPtrAccess(base, c->element);
if (inMemory) if (inMemory)
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);
if (tuples != NULL)
(*tuples) += unvisitedLength;
for (int i = 0; i < unvisitedLength; i++) for (int i = 0; i < unvisitedLength; i++)
{ {
@@ -1027,7 +1089,7 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
if (inMemory) if (inMemory)
{ {
eElement = unvisited[i].element; eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, &eMatches, q, qtup, keyData, index, procinfo, collation); eDistance = GetElementDistance(base, eElement, &eMatches, q, index, support);
} }
else else
{ {
@@ -1037,16 +1099,25 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
/* Avoid any allocations if not adding */ /* Avoid any allocations if not adding */
eElement = NULL; eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, &q, qtup, keyData, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement); HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
if (eElement == NULL) if (eElement == NULL)
continue; continue;
} }
if (!(eDistance < f->distance || alwaysAdd)) if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
continue; {
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);
}
Assert(!eElement->deleted); continue;
}
/* Make robust to issues */ /* Make robust to issues */
if (eElement->level < lc) if (eElement->level < lc)
@@ -1075,7 +1146,12 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
/* 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);
}
} }
} }
} }
@@ -1143,13 +1219,16 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
* Check if an element is closer to q than any element from R * Check if an element is closer to q than any element from R
*/ */
static bool static bool
CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, FmgrInfo **procinfo, Oid *collation) CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswSupport * support)
{ {
HnswElement eElement = HnswPtrAccess(base, e->element); HnswElement eElement = HnswPtrAccess(base, e->element);
Datum eValue = HnswGetValue(base, eElement); HnswQuery q;
IndexTuple etup = HnswPtrAccess(base, eElement->itup);
ListCell *lc2; ListCell *lc2;
q.value = HnswGetValue(base, eElement);
q.itup = HnswPtrAccess(base, eElement->itup);
q.keyData = NULL;
foreach(lc2, r) foreach(lc2, r)
{ {
HnswCandidate *ri = lfirst(lc2); HnswCandidate *ri = lfirst(lc2);
@@ -1157,7 +1236,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, FmgrI
Datum riValue = HnswGetValue(base, riElement); Datum riValue = HnswGetValue(base, riElement);
IndexTuple ritup = HnswPtrAccess(base, riElement->itup); IndexTuple ritup = HnswPtrAccess(base, riElement->itup);
bool matches; bool matches;
float distance = HnswGetDistance(etup, eValue, riValue, ritup, NULL, index, procinfo, collation, &matches); float distance = HnswGetDistance(ritup, riValue, &q, index, support, &matches);
if (distance <= e->distance) if (distance <= e->distance)
return false; return false;
@@ -1170,7 +1249,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, FmgrI
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo, Oid *collation, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates) SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
@@ -1204,7 +1283,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
/* Use previous state of r and wd to skip work when possible */ /* Use previous state of r and wd to skip work when possible */
if (mustCalculate) if (mustCalculate)
e->closer = CheckElementCloser(base, e, r, index, procinfo, collation); e->closer = CheckElementCloser(base, e, r, index, support);
else if (list_length(added) > 0) else if (list_length(added) > 0)
{ {
/* Keep Valgrind happy for in-memory, parallel builds */ /* Keep Valgrind happy for in-memory, parallel builds */
@@ -1217,8 +1296,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
*/ */
if (e->closer) if (e->closer)
{ {
e->closer = CheckElementCloser(base, e, added, index, procinfo, collation); e->closer = CheckElementCloser(base, e, added, index, support);
if (!e->closer) if (!e->closer)
removedAny = true; removedAny = true;
} }
@@ -1230,7 +1308,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
*/ */
if (removedAny) if (removedAny)
{ {
e->closer = CheckElementCloser(base, e, r, index, procinfo, collation); e->closer = CheckElementCloser(base, e, r, index, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1238,7 +1316,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
} }
else if (e == newCandidate) else if (e == newCandidate)
{ {
e->closer = CheckElementCloser(base, e, r, index, procinfo, collation); e->closer = CheckElementCloser(base, e, r, index, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1289,7 +1367,7 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
* Update connections * Update connections
*/ */
void void
HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo **procinfo, Oid *collation) HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support)
{ {
HnswCandidate newHc; HnswCandidate newHc;
@@ -1315,7 +1393,7 @@ HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newE
c = lappend(c, &neighbors->items[i]); c = lappend(c, &neighbors->items[i]);
c = lappend(c, &newHc); c = lappend(c, &newHc);
SelectNeighbors(base, c, lm, index, procinfo, collation, &neighbors->closerSet, &newHc, &pruned, true); SelectNeighbors(base, c, lm, index, support, &neighbors->closerSet, &newHc, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1386,17 +1464,20 @@ PrecomputeHash(char *base, HnswElement element)
* Algorithm 1 from paper * Algorithm 1 from paper
*/ */
void void
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo **procinfo, Oid *collation, int m, int efConstruction, bool existing, bool inMemory) HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory)
{ {
List *ep; List *ep;
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = HnswGetValue(base, element); HnswQuery q;
IndexTuple qtup = HnswPtrAccess(base, element->itup);
ScanKeyData *keyData = NULL;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
q.value = HnswGetValue(base, element);
q.itup = HnswPtrAccess(base, element->itup);
q.keyData = NULL;
/* Precompute hash */ /* Precompute hash */
if (inMemory) if (inMemory)
PrecomputeHash(base, element); PrecomputeHash(base, element);
@@ -1406,13 +1487,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
return; return;
/* Get entry point and level */ /* Get entry point and level */
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, qtup, keyData, index, procinfo, collation, true, inMemory)); ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true, inMemory));
entryLevel = entryPoint->level; entryLevel = entryPoint->level;
/* 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, qtup, keyData, ep, 1, lc, index, procinfo, collation, m, true, skipElement, inMemory); w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
ep = w; ep = w;
} }
@@ -1431,7 +1512,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *lw = NIL; List *lw = NIL;
ListCell *lc2; ListCell *lc2;
w = HnswSearchLayer(base, q, qtup, keyData, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, inMemory); w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
/* Convert search candidates to candidates */ /* Convert search candidates to candidates */
foreach(lc2, w) foreach(lc2, w)
@@ -1455,7 +1536,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
* sortCandidates to true for in-memory builds to enable closer * sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance. * caching, but there does not seem to be a difference in performance.
*/ */
neighbors = SelectNeighbors(base, lw, lm, index, procinfo, collation, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false); neighbors = SelectNeighbors(base, lw, lm, index, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);
AddConnections(base, element, neighbors, lc); AddConnections(base, element, neighbors, lc);

View File

@@ -184,13 +184,12 @@ static void
RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswElement entryPoint) RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswElement entryPoint)
{ {
Relation index = vacuumstate->index; Relation index = vacuumstate->index;
HnswSupport *support = &vacuumstate->support;
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
int m = vacuumstate->m; int m = vacuumstate->m;
int efConstruction = vacuumstate->efConstruction; int efConstruction = vacuumstate->efConstruction;
FmgrInfo **procinfo = vacuumstate->procinfo;
Oid *collation = vacuumstate->collation;
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup; HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m); Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
@@ -205,7 +204,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptidsLength = 0; element->heaptidsLength = 0;
/* Find neighbors for element, skipping itself */ /* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true, false); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true, false);
/* Zero memory for each element */ /* Zero memory for each element */
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE); MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -229,7 +228,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false); HnswUpdateNeighborsOnDisk(index, support, element, m, true, false);
} }
/* /*
@@ -239,6 +238,7 @@ static void
RepairGraphEntryPoint(HnswVacuumState * vacuumstate) RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
{ {
Relation index = vacuumstate->index; Relation index = vacuumstate->index;
HnswSupport *support = &vacuumstate->support;
HnswElement highestPoint = &vacuumstate->highestPoint; HnswElement highestPoint = &vacuumstate->highestPoint;
HnswElement entryPoint; HnswElement entryPoint;
MemoryContext oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx); MemoryContext oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx);
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock); LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */ /* Load element */
HnswLoadElement(highestPoint, NULL, NULL, NULL, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL); HnswLoadElement(highestPoint, NULL, NULL, NULL, index, support, true, NULL);
/* Repair if needed */ /* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint)) if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* is outdated, this can remove connections at higher levels in * is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine. * the graph until they are repaired, but this should be fine.
*/ */
HnswLoadElement(entryPoint, NULL, NULL, NULL, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL); HnswLoadElement(entryPoint, NULL, NULL, NULL, index, support, true, NULL);
if (NeedsUpdated(vacuumstate, entryPoint)) if (NeedsUpdated(vacuumstate, entryPoint))
{ {
@@ -535,6 +535,14 @@ 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
@@ -581,12 +589,13 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
vacuumstate->callback_state = callback_state; vacuumstate->callback_state = callback_state;
vacuumstate->efConstruction = HnswGetEfConstruction(index); vacuumstate->efConstruction = HnswGetEfConstruction(index);
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD); vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
HnswInitProcinfo(vacuumstate->procinfo, &vacuumstate->collation, index);
vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE); vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context", "Hnsw vacuum temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
HnswInitSupport(&vacuumstate->support, index);
/* Get m from metapage */ /* Get m from metapage */
HnswGetMetaPageInfo(index, &vacuumstate->m, NULL); HnswGetMetaPageInfo(index, &vacuumstate->m, NULL);

View File

@@ -138,7 +138,7 @@ SampleRows(IvfflatBuildState * buildstate)
* Add tuple to sort * Add tuple to sort
*/ */
static void static void
AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState * buildstate) AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, IvfflatBuildState * buildstate)
{ {
double distance; double distance;
double minDistance = DBL_MAX; double minDistance = DBL_MAX;
@@ -184,6 +184,11 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
slot->tts_isnull[1] = false; slot->tts_isnull[1] = false;
slot->tts_values[2] = value; slot->tts_values[2] = value;
slot->tts_isnull[2] = false; slot->tts_isnull[2] = false;
for (int i = 1; i < buildstate->tupdesc->natts; i++)
{
slot->tts_values[2 + i] = values[i];
slot->tts_isnull[2 + i] = isnull[i];
}
ExecStoreVirtualTuple(slot); ExecStoreVirtualTuple(slot);
/* /*
@@ -215,7 +220,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add tuple to sort */ /* Add tuple to sort */
AddTupleToSort(index, tid, values, buildstate); AddTupleToSort(index, tid, values, isnull, buildstate);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -226,19 +231,20 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
* Get index tuple from sort state * Get index tuple from sort state
*/ */
static inline void static inline void
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list) GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, Datum *values, bool *isnull, IndexTuple *itup, int *list)
{ {
Datum value;
bool isnull;
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL)) if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{ {
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull)); bool unused;
value = slot_getattr(slot, 3, &isnull);
*list = DatumGetInt32(slot_getattr(slot, 1, &unused));
for (int i = 0; i < tupdesc->natts; i++)
values[i] = slot_getattr(slot, 3 + i, &isnull[i]);
/* Form the index tuple */ /* Form the index tuple */
*itup = index_form_tuple(tupdesc, &value, &isnull); *itup = index_form_tuple(tupdesc, values, isnull);
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull))); (*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &unused)));
} }
else else
*list = -1; *list = -1;
@@ -254,14 +260,16 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
IndexTuple itup = NULL; /* silence compiler warning */ IndexTuple itup = NULL; /* silence compiler warning */
int64 inserted = 0; int64 inserted = 0;
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple); TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = RelationGetDescr(index); TupleDesc tupdesc = buildstate->tupdesc;
Datum *values = palloc(tupdesc->natts * sizeof(Datum));
bool *isnull = palloc(tupdesc->natts * sizeof(bool));
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD); pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples); pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples);
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list); GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
for (int i = 0; i < buildstate->centers->length; i++) for (int i = 0; i < buildstate->centers->length; i++)
{ {
@@ -297,7 +305,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted); pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted);
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list); GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
} }
insertPage = BufferGetBlockNumber(buf); insertPage = BufferGetBlockNumber(buf);
@@ -307,6 +315,9 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
/* Set the start and insert pages */ /* Set the start and insert pages */
IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum); IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
} }
pfree(values);
pfree(isnull);
} }
/* /*
@@ -319,6 +330,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->index = index; buildstate->index = index;
buildstate->indexInfo = indexInfo; buildstate->indexInfo = indexInfo;
buildstate->typeInfo = IvfflatGetTypeInfo(index); buildstate->typeInfo = IvfflatGetTypeInfo(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->lists = IvfflatGetLists(index); buildstate->lists = IvfflatGetLists(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod; buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
@@ -356,12 +368,13 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
errmsg("dimensions must be greater than one for this opclass"))); errmsg("dimensions must be greater than one for this opclass")));
/* Create tuple description for sorting */ /* Create tuple description for sorting */
buildstate->tupdesc = CreateTemplateTupleDesc(3); buildstate->sortdesc = CreateTemplateTupleDesc(2 + buildstate->tupdesc->natts);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0); for (int i = 0; i < buildstate->tupdesc->natts; i++)
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) (3 + i), NULL, buildstate->tupdesc->attrs[i].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions)); buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions));
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists); buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
@@ -633,7 +646,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo); InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen); memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
buildstate.centers->length = buildstate.centers->maxlen; buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = InitBuildSortState(buildstate.tupdesc, sortmem, coordinate); ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, sortmem, coordinate);
buildstate.sortstate = ivfspool->sortstate; buildstate.sortstate = ivfspool->sortstate;
scan = table_beginscan_parallel(ivfspool->heap, scan = table_beginscan_parallel(ivfspool->heap,
ParallelTableScanFromIvfflatShared(ivfshared)); ParallelTableScanFromIvfflatShared(ivfshared));
@@ -950,7 +963,7 @@ AssignTuples(IvfflatBuildState * buildstate)
} }
/* Begin serial/leader tuplesort */ /* Begin serial/leader tuplesort */
buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate); buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate);
/* Add tuples to sort */ /* Add tuples to sort */
if (buildstate->heap != NULL) if (buildstate->heap != NULL)

View File

@@ -17,8 +17,16 @@
#endif #endif
int ivfflat_probes; int ivfflat_probes;
int ivfflat_iterative_search;
int ivfflat_iterative_search_max_probes;
static relopt_kind ivfflat_relopt_kind; static relopt_kind ivfflat_relopt_kind;
static const struct config_enum_entry ivfflat_iterative_search_options[] = {
{"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false},
{"on", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false},
{NULL, 0, false}
};
/* /*
* Initialize index options and variables * Initialize index options and variables
*/ */
@@ -33,6 +41,14 @@ IvfflatInit(void)
"Valid range is 1..lists.", &ivfflat_probes, "Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
DefineCustomEnumVariable("ivfflat.iterative_search", "Sets whether to use iterative search",
NULL, &ivfflat_iterative_search,
IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
DefineCustomIntVariable("ivfflat.iterative_search_max_probes", "Sets the max number of probes for iterative search",
"Zero sets to the number of lists", &ivfflat_iterative_search_max_probes,
0, 0, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat"); MarkGUCPrefixReserved("ivfflat");
} }

View File

@@ -80,6 +80,14 @@
/* Variables */ /* Variables */
extern int ivfflat_probes; extern int ivfflat_probes;
extern int ivfflat_iterative_search;
extern int ivfflat_iterative_search_max_probes;
typedef enum IvfflatIterativeSearchType
{
IVFFLAT_ITERATIVE_SEARCH_OFF,
IVFFLAT_ITERATIVE_SEARCH_RELAXED
} IvfflatIterativeSearchType;
typedef struct VectorArrayData typedef struct VectorArrayData
{ {
@@ -165,6 +173,7 @@ typedef struct IvfflatBuildState
Relation index; Relation index;
IndexInfo *indexInfo; IndexInfo *indexInfo;
const IvfflatTypeInfo *typeInfo; const IvfflatTypeInfo *typeInfo;
TupleDesc tupdesc;
/* Settings */ /* Settings */
int dimensions; int dimensions;
@@ -198,7 +207,7 @@ typedef struct IvfflatBuildState
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc tupdesc; TupleDesc sortdesc;
TupleTableSlot *slot; TupleTableSlot *slot;
/* Memory */ /* Memory */
@@ -247,8 +256,10 @@ typedef struct IvfflatScanOpaqueData
{ {
const IvfflatTypeInfo *typeInfo; const IvfflatTypeInfo *typeInfo;
int probes; int probes;
int maxProbes;
int dimensions; int dimensions;
bool first; bool first;
Datum value;
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
@@ -265,6 +276,8 @@ typedef struct IvfflatScanOpaqueData
/* Lists */ /* Lists */
pairingheap *listQueue; pairingheap *listQueue;
BlockNumber *listPages;
int listIndex;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */ IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
} IvfflatScanOpaqueData; } IvfflatScanOpaqueData;

View File

@@ -78,6 +78,8 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
ListInfo listInfo; ListInfo listInfo;
BlockNumber originalInsertPage; BlockNumber originalInsertPage;
TupleDesc tupdesc = RelationGetDescr(index);
Datum *newValues = palloc(tupdesc->natts * sizeof(Datum));
/* Detoast once for all calls */ /* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -98,12 +100,16 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
IvfflatGetMetaPageInfo(index, NULL, NULL); IvfflatGetMetaPageInfo(index, NULL, NULL);
/* Find the insert page - sets the page and list info */ /* Find the insert page - sets the page and list info */
FindInsertPage(index, values, &insertPage, &listInfo); FindInsertPage(index, &value, &insertPage, &listInfo);
Assert(BlockNumberIsValid(insertPage)); Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage; originalInsertPage = insertPage;
newValues[0] = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
/* Form tuple */ /* Form tuple */
itup = index_form_tuple(RelationGetDescr(index), &value, isnull); itup = index_form_tuple(tupdesc, newValues, isnull);
itup->t_tid = *heap_tid; itup->t_tid = *heap_tid;
/* Get tuple size */ /* Get tuple size */

View File

@@ -65,7 +65,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
/* Use procinfo from the index instead of scan key for performance */ /* Use procinfo from the index instead of scan key for performance */
distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value)); distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
if (listCount < so->probes) if (listCount < so->maxProbes)
{ {
IvfflatScanList *scanlist; IvfflatScanList *scanlist;
@@ -78,7 +78,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
pairingheap_add(so->listQueue, &scanlist->ph_node); pairingheap_add(so->listQueue, &scanlist->ph_node);
/* Calculate max distance */ /* Calculate max distance */
if (listCount == so->probes) if (listCount == so->maxProbes)
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance; maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
} }
else if (distance < maxDistance) else if (distance < maxDistance)
@@ -102,6 +102,11 @@ GetScanLists(IndexScanDesc scan, Datum value)
UnlockReleaseBuffer(cbuf); UnlockReleaseBuffer(cbuf);
} }
for (int i = listCount - 1; i >= 0; i--)
so->listPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
Assert(pairingheap_is_empty(so->listQueue));
} }
/* /*
@@ -114,11 +119,14 @@ GetScanItems(IndexScanDesc scan, Datum value)
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0; double tuples = 0;
TupleTableSlot *slot = so->vslot; TupleTableSlot *slot = so->vslot;
int batchProbes = 0;
tuplesort_reset(so->sortstate);
/* Search closest probes lists */ /* Search closest probes lists */
while (!pairingheap_is_empty(so->listQueue)) while (so->listIndex < so->maxProbes && (++batchProbes) <= so->probes)
{ {
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage; BlockNumber searchPage = so->listPages[so->listIndex++];
/* Search all entry pages for list */ /* Search all entry pages for list */
while (BlockNumberIsValid(searchPage)) while (BlockNumberIsValid(searchPage))
@@ -166,13 +174,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
} }
} }
if (tuples < 100) if (tuples < 100 && ivfflat_iterative_search == IVFFLAT_ITERATIVE_SEARCH_OFF)
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("index scan found few tuples"), (errmsg("index scan found few tuples"),
errdetail("Index may have been created with little data."), errdetail("Index may have been created with little data."),
errhint("Recreate the index and possibly decrease lists."))); errhint("Recreate the index and possibly decrease lists.")));
tuplesort_performsort(so->sortstate); tuplesort_performsort(so->sortstate);
#if defined(IVFFLAT_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
} }
/* /*
@@ -240,6 +252,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
int lists; int lists;
int dimensions; int dimensions;
int probes = ivfflat_probes; int probes = ivfflat_probes;
int maxProbes;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
@@ -249,10 +262,21 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
if (probes > lists) if (probes > lists)
probes = lists; probes = lists;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList)); if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF)
{
if (ivfflat_iterative_search_max_probes == 0)
maxProbes = lists;
else
maxProbes = Min(ivfflat_iterative_search_max_probes, lists);
}
else
maxProbes = probes;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList));
so->typeInfo = IvfflatGetTypeInfo(index); so->typeInfo = IvfflatGetTypeInfo(index);
so->first = true; so->first = true;
so->probes = probes; so->probes = probes;
so->maxProbes = maxProbes;
so->dimensions = dimensions; so->dimensions = dimensions;
/* Set support functions */ /* Set support functions */
@@ -280,6 +304,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->bas = GetAccessStrategy(BAS_BULKREAD); so->bas = GetAccessStrategy(BAS_BULKREAD);
so->listQueue = pairingheap_allocate(CompareLists, scan); so->listQueue = pairingheap_allocate(CompareLists, scan);
so->listPages = palloc(maxProbes * sizeof(BlockNumber));
so->listIndex = 0;
scan->opaque = so; scan->opaque = so;
@@ -294,11 +320,9 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
if (!so->first)
tuplesort_reset(so->sortstate);
so->first = true; so->first = true;
pairingheap_reset(so->listQueue); pairingheap_reset(so->listQueue);
so->listIndex = 0;
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData)); memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -314,6 +338,8 @@ bool
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir) ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
ItemPointer heaptid;
bool isnull;
/* /*
* Index can be used to scan backward, but Postgres doesn't support * Index can be used to scan backward, but Postgres doesn't support
@@ -341,20 +367,20 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IvfflatBench("GetScanLists", GetScanLists(scan, value)); IvfflatBench("GetScanLists", GetScanLists(scan, value));
IvfflatBench("GetScanItems", GetScanItems(scan, value)); IvfflatBench("GetScanItems", GetScanItems(scan, value));
so->first = false; so->first = false;
so->value = value;
#if defined(IVFFLAT_MEMORY) /* TODO clean up if we allocated a new value */
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
/* Clean up if we allocated a new value */
if (value != scan->orderByData->sk_argument)
pfree(DatumGetPointer(value));
} }
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL)) while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{ {
bool isnull; if (so->listIndex == so->maxProbes)
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull)); return false;
IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
}
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;
@@ -362,9 +388,6 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
return true; return true;
} }
return false;
}
/* /*
* End a scan and release resources * End a scan and release resources
*/ */
@@ -374,6 +397,7 @@ ivfflatendscan(IndexScanDesc scan)
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
pairingheap_free(so->listQueue); pairingheap_free(so->listQueue);
pfree(so->listPages);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas); FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc); FreeTupleDesc(so->tupdesc);

View File

@@ -26,7 +26,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
Page cpage; Page cpage;
OffsetNumber coffno; OffsetNumber coffno;
OffsetNumber cmaxoffno; OffsetNumber cmaxoffno;
BlockNumber startPages[MaxOffsetNumber]; BlockNumber listPages[MaxOffsetNumber];
ListInfo listInfo; ListInfo listInfo;
cbuf = ReadBuffer(index, blkno); cbuf = ReadBuffer(index, blkno);
@@ -40,7 +40,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
{ {
IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno)); IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));
startPages[coffno - FirstOffsetNumber] = list->startPage; listPages[coffno - FirstOffsetNumber] = list->startPage;
} }
listInfo.blkno = blkno; listInfo.blkno = blkno;
@@ -50,7 +50,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno)) for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
{ {
BlockNumber searchPage = startPages[coffno - FirstOffsetNumber]; BlockNumber searchPage = listPages[coffno - FirstOffsetNumber];
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
/* Iterate over entry pages */ /* Iterate over entry pages */

View File

@@ -6,13 +6,7 @@ use Test::More;
my $dim = 3; my $dim = 3;
my @r = (); my $array_sql = join(",", ('random()') x $dim);
for (1 .. $dim)
{
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
my $array_sql = join(", ", @r);
# Initialize node # Initialize node
my $node = PostgreSQL::Test::Cluster->new('node'); my $node = PostgreSQL::Test::Cluster->new('node');
@@ -23,19 +17,20 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;"); $node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));"); $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres", $node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;" "INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
); );
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);"); $node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
# Get size # Get size
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');"); my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
# Store values
$node->safe_psql("postgres", "CREATE TABLE tmp AS SELECT * FROM tst;");
# Delete all, vacuum, and insert same data # Delete all, vacuum, and insert same data
$node->safe_psql("postgres", "DELETE FROM tst;"); $node->safe_psql("postgres", "DELETE FROM tst;");
$node->safe_psql("postgres", "VACUUM tst;"); $node->safe_psql("postgres", "VACUUM tst;");
$node->safe_psql("postgres", $node->safe_psql("postgres", "INSERT INTO tst SELECT * FROM tmp;");
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
# Check size # Check size
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');"); my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");

View File

@@ -0,0 +1,54 @@
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 PRIMARY KEY, 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", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 10;
SET ivfflat.iterative_search = on;
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);
foreach ((30, 50, 70))
{
my $max_probes = $_;
my $expected = $max_probes / 10;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 10;
SET ivfflat.iterative_search = on;
SET ivfflat.iterative_search_max_probes = $max_probes;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
done_testing();

View File

@@ -0,0 +1,125 @@
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 @cs = (100, 1000);
sub test_recall
{
my ($c, $probes, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SET ivfflat.iterative_search = 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 ivfflat.probes = $probes;
SET ivfflat.iterative_search = 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(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
);
# Generate queries
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
push(@queries, "[$r1,$r2,$r3]");
}
# 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", "CREATE INDEX idx ON tst USING ivfflat (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, 1, 0.58, $operator);
test_recall($c, 10, 0.98, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 1, 0.80, $operator);
}
else
{
test_recall($c, 1, 0.88, $operator);
}
}
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();

View File

@@ -0,0 +1,67 @@
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 PRIMARY KEY, 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.iterative_search = 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);
foreach ((30000, 50000, 70000))
{
my $max_tuples = $_;
my $expected = $max_tuples / 10000;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_search = on;
SET hnsw.iterative_search_max_tuples = $max_tuples;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_search = on;
SET client_min_messages = debug1;
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/hnsw index scan exceeded work_mem after \d+ tuples/);
done_testing();

View File

@@ -0,0 +1,131 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my @cs = (100, 1000);
sub test_recall
{
my ($c, $ef_search, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_search = 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.iterative_search = 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();

View File

@@ -11,7 +11,7 @@ my @expected;
my $limit = 20; my $limit = 20;
my $dim = 3; my $dim = 3;
my $array_sql = join(",", ('random()') x $dim); my $array_sql = join(",", ('random()') x $dim);
my $nc = 50; my $nc = 1000;
sub test_recall sub test_recall
{ {
@@ -79,15 +79,19 @@ for (1 .. 20)
@expected = (); @expected = ();
for my $i (0 .. $#queries) for my $i (0 .. $#queries)
{ {
my $res = $node->safe_psql("postgres", "SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <-> '$queries[$i]' LIMIT $limit;"); my $res = $node->safe_psql("postgres", "SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <=> '$queries[$i]' LIMIT $limit;");
push(@expected, $res); push(@expected, $res);
} }
# Add index # Add index
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c);"); $node->safe_psql("postgres", qq(
SET maintenance_work_mem = '256MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, c);
));
# Test recall # Test recall
test_recall(0.99, '<->'); test_recall(0.99, '<=>');
# Test vacuum # Test vacuum
$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;"); $node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
@@ -97,13 +101,13 @@ $node->safe_psql("postgres", "VACUUM tst;");
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);"); my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);");
like($stderr, qr/first column must be a vector/); like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_l2_ops);"); ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_cosine_ops);");
like($stderr, qr/first column must be a vector/); like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c, c);"); ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, c, c);");
like($stderr, qr/index cannot have more than two columns/); like($stderr, qr/index cannot have more than two columns/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, v vector_l2_ops);"); ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, v vector_cosine_ops);");
like($stderr, qr/column 2 cannot be a vector/); like($stderr, qr/column 2 cannot be a vector/);
done_testing(); done_testing();