liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-23 04:20:56 +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;
}
static void
ivfflatunguardbatch(IndexScanDesc scan, IndexScanBatch batch)
{
}
/*
* Define index handler
*
@@ -229,9 +234,9 @@ ivfflathandler(PG_FUNCTION_ARGS)
.amadjustmembers = NULL,
.ambeginscan = ivfflatbeginscan,
.amrescan = ivfflatrescan,
.amgettuple = ivfflatgettuple,
.amgetbatch = NULL,
.amunguardbatch = NULL,
.amgettuple = NULL,
.amgetbatch = ivfflatgetbatch,
.amunguardbatch = ivfflatunguardbatch,
.amkillitemsbatch = NULL,
.amgettransform = NULL,
.amgetbitmap = NULL,

2
src/ivfflat.h
View File

@@ -344,7 +344,7 @@ IndexBulkDeleteResult *ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteR
IndexBulkDeleteResult *ivfflatvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats);
IndexScanDesc ivfflatbeginscan(Relation index, int nkeys, 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);
#endif

45
src/ivfscan.c
View File

@@ -3,6 +3,7 @@
#include <float.h>
#include "access/genam.h"
#include "access/indexbatch.h"
#include "access/itup.h"
#include "access/relscan.h"
#include "access/tupdesc.h"
@@ -261,6 +262,11 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
MemoryContext oldCtx;
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 */
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
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IndexScanBatch
ivfflatgetbatch(IndexScanDesc scan, IndexScanBatch priorbatch, ScanDirection dir)
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
IndexScanBatch batch = indexam_util_alloc_batch(scan);
ItemPointer heaptid;
bool isnull;
int nitems = 0;
/*
* 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))
{
if (so->listIndex == so->maxProbes)
return false;
{
indexam_util_release_batch(scan, batch);
return NULL;
}
IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
}
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
for (;;)
{
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
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;
}
batch->firstItem = 0;
batch->lastItem = nitems - 1;
batch->dir = ForwardScanDirection;
return batch;
}
/*