From 18980387ec49d69012181ce5a6ed19616ba92602 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 10 Jun 2026 11:31:57 -0700 Subject: [PATCH] Added support for amgetbatch [skip ci] --- src/hnsw.c | 15 ++++++++-- src/hnsw.h | 2 +- src/hnswscan.c | 80 +++++++++++++++++++++++++++++++------------------- src/ivfflat.c | 8 +++-- 4 files changed, 68 insertions(+), 37 deletions(-) diff --git a/src/hnsw.c b/src/hnsw.c index a3f5822..0c07cfd 100644 --- a/src/hnsw.c +++ b/src/hnsw.c @@ -258,6 +258,11 @@ hnswvalidate(Oid opclassoid) return true; } +static void +hnswunguardbatch(IndexScanDesc scan, IndexScanBatch batch) +{ +} + /* * Define index handler * @@ -279,6 +284,7 @@ hnswhandler(PG_FUNCTION_ARGS) .amconsistentequality = false, .amconsistentordering = false, .amcanbackward = false, + .amcanmarkpos = false, .amcanunique = false, .amcanmulticol = false, .amoptionalkey = true, @@ -311,11 +317,14 @@ hnswhandler(PG_FUNCTION_ARGS) .amadjustmembers = NULL, .ambeginscan = hnswbeginscan, .amrescan = hnswrescan, - .amgettuple = hnswgettuple, + .amgettuple = NULL, + .amgetbatch = hnswgetbatch, + .amunguardbatch = hnswunguardbatch, + .amkillitemsbatch = NULL, + .amgettransform = NULL, .amgetbitmap = NULL, .amendscan = hnswendscan, - .ammarkpos = NULL, - .amrestrpos = NULL, + .amposreset = NULL, .amestimateparallelscan = NULL, .aminitparallelscan = NULL, .amparallelrescan = NULL, diff --git a/src/hnsw.h b/src/hnsw.h index 65a74e8..a2c8ade 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -465,7 +465,7 @@ IndexBulkDeleteResult *hnswbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResu IndexBulkDeleteResult *hnswvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats); IndexScanDesc hnswbeginscan(Relation index, int nkeys, int norderbys); void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys); -bool hnswgettuple(IndexScanDesc scan, ScanDirection dir); +IndexScanBatch hnswgetbatch(IndexScanDesc scan, IndexScanBatch priorbatch, ScanDirection dir); void hnswendscan(IndexScanDesc scan); static inline HnswNeighborArray * diff --git a/src/hnswscan.c b/src/hnswscan.c index 32cdb7d..a574a01 100644 --- a/src/hnswscan.c +++ b/src/hnswscan.c @@ -3,6 +3,7 @@ #include #include "access/genam.h" +#include "access/indexbatch.h" #include "access/relscan.h" #include "hnsw.h" #include "lib/pairingheap.h" @@ -135,6 +136,11 @@ hnswbeginscan(Relation index, int nkeys, int norderbys) double maxMemory; scan = RelationGetIndexScan(index, nkeys, norderbys); + scan->maxitemsbatch = hnsw_ef_search * HNSW_HEAPTIDS; + /* unused but be > 0 */ + scan->batch_index_opaque_static = MAXALIGN(1); + scan->batch_index_opaque_dyn = 0; + scan->batch_tuples_workspace = 0; so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData)); so->typeInfo = HnswGetTypeInfo(index); @@ -184,12 +190,13 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no } /* - * Fetch the next tuple in the given scan + * Fetch the next batch in the given scan */ -bool -hnswgettuple(IndexScanDesc scan, ScanDirection dir) +IndexScanBatch +hnswgetbatch(IndexScanDesc scan, IndexScanBatch priorbatch, ScanDirection dir) { HnswScanOpaque so = (HnswScanOpaque) scan->opaque; + IndexScanBatch batch = indexam_util_alloc_batch(scan); MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx); /* @@ -242,9 +249,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir) for (;;) { char *base = NULL; - HnswSearchCandidate *sc; - HnswElement element; - ItemPointer heaptid; + int nitems = 0; if (list_length(so->w) == 0) { @@ -290,44 +295,57 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir) break; } - sc = llast(so->w); - element = HnswPtrAccess(base, sc->element); - - /* Move to next element if no valid heap TIDs */ - if (element->heaptidsLength == 0) + while (list_length(so->w) != 0) { - so->w = list_delete_last(so->w); + HnswSearchCandidate *sc = llast(so->w); + HnswElement element = HnswPtrAccess(base, sc->element); + ItemPointer heaptid; - /* Mark memory as free for next iteration */ - if (hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF) + /* Move to next element if no valid heap TIDs */ + if (element->heaptidsLength == 0) { - pfree(element); - pfree(sc); + so->w = list_delete_last(so->w); + + /* Mark memory as free for next iteration */ + if (hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF) + { + pfree(element); + pfree(sc); + } + + continue; } - continue; + heaptid = &element->heaptids[--element->heaptidsLength]; + + if (hnsw_iterative_scan == HNSW_ITERATIVE_SCAN_STRICT) + { + if (sc->distance < so->previousDistance) + continue; + + so->previousDistance = sc->distance; + } + + batch->items[nitems].tableTid = *heaptid; + batch->items[nitems].indexOffset = -1; + batch->items[nitems].tupleOffset = 0; + nitems++; } - heaptid = &element->heaptids[--element->heaptidsLength]; - - if (hnsw_iterative_scan == HNSW_ITERATIVE_SCAN_STRICT) - { - if (sc->distance < so->previousDistance) - continue; - - so->previousDistance = sc->distance; - } + if (nitems == 0) + break; MemoryContextSwitchTo(oldCtx); - scan->xs_heaptid = *heaptid; - scan->xs_recheck = false; - scan->xs_recheckorderby = false; - return true; + batch->firstItem = 0; + batch->lastItem = nitems - 1; + batch->dir = ForwardScanDirection; + return batch; } MemoryContextSwitchTo(oldCtx); - return false; + indexam_util_release_batch(scan, batch); + return NULL; } /* diff --git a/src/ivfflat.c b/src/ivfflat.c index 8489d93..c4a03ec 100644 --- a/src/ivfflat.c +++ b/src/ivfflat.c @@ -196,6 +196,7 @@ ivfflathandler(PG_FUNCTION_ARGS) .amconsistentequality = false, .amconsistentordering = false, .amcanbackward = false, + .amcanmarkpos = false, .amcanunique = false, .amcanmulticol = false, .amoptionalkey = true, @@ -229,10 +230,13 @@ ivfflathandler(PG_FUNCTION_ARGS) .ambeginscan = ivfflatbeginscan, .amrescan = ivfflatrescan, .amgettuple = ivfflatgettuple, + .amgetbatch = NULL, + .amunguardbatch = NULL, + .amkillitemsbatch = NULL, + .amgettransform = NULL, .amgetbitmap = NULL, .amendscan = ivfflatendscan, - .ammarkpos = NULL, - .amrestrpos = NULL, + .amposreset = NULL, .amestimateparallelscan = NULL, .aminitparallelscan = NULL, .amparallelrescan = NULL,