liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-11 15:16:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use DSA

Browse Source
This commit is contained in:
Andrew Kane
2024-01-16 21:59:11 -08:00
parent d801a843f4
commit a7b6be73e1
2 changed files with 30 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/hnsw.h
View File

@@ -213,6 +213,7 @@ typedef struct HnswShared
Oid indexrelid; Oid indexrelid;
bool isconcurrent; bool isconcurrent;
int scantuplesortstates; int scantuplesortstates;
dsa_pointer dsaptr;
/* Worker progress */ /* Worker progress */
ConditionVariable workersdonecv; ConditionVariable workersdonecv;

42
src/hnswbuild.c
View File

@@ -768,6 +768,8 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
char *sharedquery; char *sharedquery;
HnswSpool *hnswspool; HnswSpool *hnswspool;
HnswShared *hnswshared; HnswShared *hnswshared;
dsa_area *dsa;
void *dsaspace;
char *hnswarea; char *hnswarea;
Relation heapRel; Relation heapRel;
Relation indexRel; Relation indexRel;
@@ -809,7 +811,9 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
hnswspool->heap = heapRel; hnswspool->heap = heapRel;
hnswspool->index = indexRel; hnswspool->index = indexRel;
hnswarea = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_AREA, false); dsaspace = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_AREA, false);
dsa = dsa_attach_in_place(dsaspace, seg);
hnswarea = dsa_get_address(dsa, hnswshared->dsaptr);
/* Perform inserts */ /* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, hnswarea, false); HnswParallelScanAndInsert(hnswspool, hnswshared, hnswarea, false);
@@ -889,8 +893,10 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
Snapshot snapshot; Snapshot snapshot;
Size esthnswshared; Size esthnswshared;
Size esthnswarea; Size esthnswarea;
Size estother; Size estdsaspace;
HnswShared *hnswshared; HnswShared *hnswshared;
void *dsaspace;
dsa_area *dsa;
char *hnswarea; char *hnswarea;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader)); HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true; bool leaderparticipates = true;
@@ -921,15 +927,9 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot); esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared); shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
/* Leave space for other objects in shared memory */ /* Start with a minimal DSA so InitializeParallelDSM does not fail */
/* Docker has a default limit of 64 MB for shm_size */ estdsaspace = dsa_minimum_size();
/* which happens to be the default value of maintenance_work_mem */ shm_toc_estimate_chunk(&pcxt->estimator, estdsaspace);
esthnswarea = maintenance_work_mem * 1024L;
estother = 2 * 1024 * 1024;
if (esthnswarea > estother)
esthnswarea -= estother;
shm_toc_estimate_chunk(&pcxt->estimator, esthnswarea);
shm_toc_estimate_keys(&pcxt->estimator, 2); shm_toc_estimate_keys(&pcxt->estimator, 2);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */ /* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
@@ -975,12 +975,28 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot); heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
#endif #endif
hnswarea = (char *) shm_toc_allocate(pcxt->toc, esthnswarea); dsaspace = shm_toc_allocate(pcxt->toc, estdsaspace);
dsa = dsa_create_in_place(dsaspace, estdsaspace, hnsw_lock_tranche_id, pcxt->seg);
esthnswarea = maintenance_work_mem * 1024L;
hnswshared->dsaptr = dsa_allocate_extended(dsa, esthnswarea, DSA_ALLOC_HUGE | DSA_ALLOC_NO_OOM);
/* If not enough shared memory, back out (do serial build) */
if (!DsaPointerIsValid(hnswshared->dsaptr))
{
if (IsMVCCSnapshot(snapshot))
UnregisterSnapshot(snapshot);
DestroyParallelContext(pcxt);
ExitParallelMode();
return;
}
hnswarea = dsa_get_address(dsa, hnswshared->dsaptr);
/* Report less than allocated so never fails */ /* Report less than allocated so never fails */
InitGraph(&hnswshared->graphData, hnswarea, esthnswarea - 1024 * 1024); InitGraph(&hnswshared->graphData, hnswarea, esthnswarea - 1024 * 1024);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared); shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_AREA, hnswarea); shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_AREA, dsaspace);
/* Store query string for workers */ /* Store query string for workers */
if (debug_query_string) if (debug_query_string)