liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-26 22:02:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added support for amgetbatch for IVFFlat [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2026-06-10 13:47:36 -07:00
parent 91717c81b3
commit 8a8bc90d6d
3 changed files with 45 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/ivfflat.c
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
* *
@@ -229,9 +234,9 @@ ivfflathandler(PG_FUNCTION_ARGS)
.amadjustmembers = NULL, .amadjustmembers = NULL,
.ambeginscan = ivfflatbeginscan, .ambeginscan = ivfflatbeginscan,
.amrescan = ivfflatrescan, .amrescan = ivfflatrescan,
.amgettuple = ivfflatgettuple, .amgettuple = NULL,
.amgetbatch = NULL, .amgetbatch = ivfflatgetbatch,
.amunguardbatch = NULL, .amunguardbatch = ivfflatunguardbatch,
.amkillitemsbatch = NULL, .amkillitemsbatch = NULL,
.amgettransform = NULL, .amgettransform = NULL,
.amgetbitmap = NULL, .amgetbitmap = NULL,

2
src/ivfflat.h
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

45
src/ivfscan.c
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,36 @@ 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 (;;)
{
scan->xs_heaptid = *heaptid; heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
scan->xs_recheck = false;
scan->xs_recheckorderby = false; batch->items[nitems].tableTid = *heaptid;
return true; 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;
}
batch->firstItem = 0;
batch->lastItem = nitems - 1;
batch->dir = ForwardScanDirection;
return batch;
} }
/* /*