Compare commits

...

10 Commits

Author SHA1 Message Date
Andrew Kane
89dc41ce2b Fixed crash with debugging disabled [skip ci] 2026-06-10 14:31:54 -07:00
Andrew Kane
06ab41c094 Removed extra whitespace [skip ci] 2026-06-10 13:49:39 -07:00
Andrew Kane
8a8bc90d6d Added support for amgetbatch for IVFFlat [skip ci] 2026-06-10 13:47:36 -07:00
Andrew Kane
91717c81b3 Improved batch size flexibility [skip ci] 2026-06-10 12:10:20 -07:00
Andrew Kane
98a7f94608 Added comment [skip ci] 2026-06-10 11:59:38 -07:00
Andrew Kane
7db3d67b2d Re-added check for iterative scans [skip ci] 2026-06-10 11:58:19 -07:00
Andrew Kane
1dd39f66de Minimized diff [skip ci] 2026-06-10 11:56:52 -07:00
Andrew Kane
cd431c51ca Removed unneeded check [skip ci] 2026-06-10 11:55:31 -07:00
Andrew Kane
73890f2f70 Fixed comment [skip ci] 2026-06-10 11:46:15 -07:00
Andrew Kane
18980387ec Added support for amgetbatch [skip ci] 2026-06-10 11:31:57 -07:00
6 changed files with 116 additions and 42 deletions

View File

@@ -258,6 +258,11 @@ hnswvalidate(Oid opclassoid)
return true; return true;
} }
static void
hnswunguardbatch(IndexScanDesc scan, IndexScanBatch batch)
{
}
/* /*
* Define index handler * Define index handler
* *
@@ -279,6 +284,7 @@ hnswhandler(PG_FUNCTION_ARGS)
.amconsistentequality = false, .amconsistentequality = false,
.amconsistentordering = false, .amconsistentordering = false,
.amcanbackward = false, .amcanbackward = false,
.amcanmarkpos = false,
.amcanunique = false, .amcanunique = false,
.amcanmulticol = false, .amcanmulticol = false,
.amoptionalkey = true, .amoptionalkey = true,
@@ -311,11 +317,14 @@ hnswhandler(PG_FUNCTION_ARGS)
.amadjustmembers = NULL, .amadjustmembers = NULL,
.ambeginscan = hnswbeginscan, .ambeginscan = hnswbeginscan,
.amrescan = hnswrescan, .amrescan = hnswrescan,
.amgettuple = hnswgettuple, .amgettuple = NULL,
.amgetbatch = hnswgetbatch,
.amunguardbatch = hnswunguardbatch,
.amkillitemsbatch = NULL,
.amgettransform = NULL,
.amgetbitmap = NULL, .amgetbitmap = NULL,
.amendscan = hnswendscan, .amendscan = hnswendscan,
.ammarkpos = NULL, .amposreset = NULL,
.amrestrpos = NULL,
.amestimateparallelscan = NULL, .amestimateparallelscan = NULL,
.aminitparallelscan = NULL, .aminitparallelscan = NULL,
.amparallelrescan = NULL, .amparallelrescan = NULL,

View File

@@ -465,7 +465,7 @@ IndexBulkDeleteResult *hnswbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResu
IndexBulkDeleteResult *hnswvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats); IndexBulkDeleteResult *hnswvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats);
IndexScanDesc hnswbeginscan(Relation index, int nkeys, int norderbys); IndexScanDesc hnswbeginscan(Relation index, int nkeys, int norderbys);
void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, 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); void hnswendscan(IndexScanDesc scan);
static inline HnswNeighborArray * static inline HnswNeighborArray *

View File

@@ -3,6 +3,7 @@
#include <limits.h> #include <limits.h>
#include "access/genam.h" #include "access/genam.h"
#include "access/indexbatch.h"
#include "access/relscan.h" #include "access/relscan.h"
#include "hnsw.h" #include "hnsw.h"
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
@@ -135,6 +136,11 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
double maxMemory; double maxMemory;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
scan->maxitemsbatch = hnsw_ef_search * HNSW_HEAPTIDS;
/* unused but must 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 = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->typeInfo = HnswGetTypeInfo(index); 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 IndexScanBatch
hnswgettuple(IndexScanDesc scan, ScanDirection dir) hnswgetbatch(IndexScanDesc scan, IndexScanBatch priorbatch, ScanDirection dir)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
IndexScanBatch batch = indexam_util_alloc_batch(scan);
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx); MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
/* /*
@@ -245,6 +252,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
HnswSearchCandidate *sc; HnswSearchCandidate *sc;
HnswElement element; HnswElement element;
ItemPointer heaptid; ItemPointer heaptid;
int nitems = 0;
if (list_length(so->w) == 0) if (list_length(so->w) == 0)
{ {
@@ -290,44 +298,64 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
break; break;
} }
sc = llast(so->w); while (list_length(so->w) != 0)
element = HnswPtrAccess(base, sc->element);
/* Move to next element if no valid heap TIDs */
if (element->heaptidsLength == 0)
{ {
so->w = list_delete_last(so->w); sc = llast(so->w);
element = HnswPtrAccess(base, sc->element);
/* Mark memory as free for next iteration */ /* Move to next element if no valid heap TIDs */
if (hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF) if (element->heaptidsLength == 0)
{ {
pfree(element); so->w = list_delete_last(so->w);
pfree(sc);
/* Mark memory as free for next iteration */
if (hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF)
{
pfree(element);
pfree(sc);
}
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++;
/* Keep batch size flexible */
if (nitems == scan->maxitemsbatch)
break;
}
/* Needed for strict iterative scans */
if (nitems == 0)
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;
}
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false; scan->xs_recheck = false;
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true;
batch->firstItem = 0;
batch->lastItem = nitems - 1;
batch->dir = ForwardScanDirection;
return batch;
} }
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
return false; indexam_util_release_batch(scan, batch);
return NULL;
} }
/* /*

View File

@@ -175,6 +175,11 @@ ivfflatvalidate(Oid opclassoid)
return true; return true;
} }
static void
ivfflatunguardbatch(IndexScanDesc scan, IndexScanBatch batch)
{
}
/* /*
* Define index handler * Define index handler
* *
@@ -196,6 +201,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
.amconsistentequality = false, .amconsistentequality = false,
.amconsistentordering = false, .amconsistentordering = false,
.amcanbackward = false, .amcanbackward = false,
.amcanmarkpos = false,
.amcanunique = false, .amcanunique = false,
.amcanmulticol = false, .amcanmulticol = false,
.amoptionalkey = true, .amoptionalkey = true,
@@ -228,11 +234,14 @@ ivfflathandler(PG_FUNCTION_ARGS)
.amadjustmembers = NULL, .amadjustmembers = NULL,
.ambeginscan = ivfflatbeginscan, .ambeginscan = ivfflatbeginscan,
.amrescan = ivfflatrescan, .amrescan = ivfflatrescan,
.amgettuple = ivfflatgettuple, .amgettuple = NULL,
.amgetbatch = ivfflatgetbatch,
.amunguardbatch = ivfflatunguardbatch,
.amkillitemsbatch = NULL,
.amgettransform = NULL,
.amgetbitmap = NULL, .amgetbitmap = NULL,
.amendscan = ivfflatendscan, .amendscan = ivfflatendscan,
.ammarkpos = NULL, .amposreset = NULL,
.amrestrpos = NULL,
.amestimateparallelscan = NULL, .amestimateparallelscan = NULL,
.aminitparallelscan = NULL, .aminitparallelscan = NULL,
.amparallelrescan = NULL, .amparallelrescan = NULL,

View File

@@ -344,7 +344,7 @@ IndexBulkDeleteResult *ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteR
IndexBulkDeleteResult *ivfflatvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats); IndexBulkDeleteResult *ivfflatvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats);
IndexScanDesc ivfflatbeginscan(Relation index, int nkeys, int norderbys); IndexScanDesc ivfflatbeginscan(Relation index, int nkeys, int norderbys);
void ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys); void ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys);
bool ivfflatgettuple(IndexScanDesc scan, ScanDirection dir); IndexScanBatch ivfflatgetbatch(IndexScanDesc scan, IndexScanBatch priorbatch, ScanDirection dir);
void ivfflatendscan(IndexScanDesc scan); void ivfflatendscan(IndexScanDesc scan);
#endif #endif

View File

@@ -3,6 +3,7 @@
#include <float.h> #include <float.h>
#include "access/genam.h" #include "access/genam.h"
#include "access/indexbatch.h"
#include "access/itup.h" #include "access/itup.h"
#include "access/relscan.h" #include "access/relscan.h"
#include "access/tupdesc.h" #include "access/tupdesc.h"
@@ -261,6 +262,11 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
MemoryContext oldCtx; MemoryContext oldCtx;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
scan->maxitemsbatch = 1000;
/* unused but must be > 0 */
scan->batch_index_opaque_static = MAXALIGN(1);
scan->batch_index_opaque_dyn = 0;
scan->batch_tuples_workspace = 0;
/* Get lists and dimensions from metapage */ /* Get lists and dimensions from metapage */
IvfflatGetMetaPageInfo(index, &lists, &dimensions); IvfflatGetMetaPageInfo(index, &lists, &dimensions);
@@ -348,14 +354,16 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
} }
/* /*
* Fetch the next tuple in the given scan * Fetch the next batch in the given scan
*/ */
bool IndexScanBatch
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir) ivfflatgetbatch(IndexScanDesc scan, IndexScanBatch priorbatch, ScanDirection dir)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
IndexScanBatch batch = indexam_util_alloc_batch(scan);
ItemPointer heaptid; ItemPointer heaptid;
bool isnull; bool isnull;
int nitems = 0;
/* /*
* Index can be used to scan backward, but Postgres doesn't support * Index can be used to scan backward, but Postgres doesn't support
@@ -393,17 +401,37 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL)) while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{ {
if (so->listIndex == so->maxProbes) if (so->listIndex == so->maxProbes)
return false; {
indexam_util_release_batch(scan, batch);
return NULL;
}
IvfflatBench("GetScanItems", GetScanItems(scan, so->value)); IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
} }
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull)); for (;;)
{
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
batch->items[nitems].tableTid = *heaptid;
batch->items[nitems].indexOffset = -1;
batch->items[nitems].tupleOffset = 0;
nitems++;
if (nitems == scan->maxitemsbatch)
break;
if (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
break;
}
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false; scan->xs_recheck = false;
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true;
batch->firstItem = 0;
batch->lastItem = nitems - 1;
batch->dir = ForwardScanDirection;
return batch;
} }
/* /*