Added bound option

CHANGELOG.md

@@ -1,8 +1,3 @@
-## 0.2.6 (unreleased)
-
-- Significantly improved index query performance
-- Improved performance of index creation for Postgres < 12
-
 ## 0.2.5 (2022-02-11)
 
 - Reduced memory usage during index creation

src/ivfbuild.c

@@ -112,13 +112,13 @@ SampleRows(IvfflatBuildState * buildstate)
 
 #if PG_VERSION_NUM >= 120000
 		table_index_build_range_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-									 false, true, false, targblock, 1, SampleCallback, (void *) buildstate, NULL);
+									 false, true, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
 #elif PG_VERSION_NUM >= 110000
 		IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-								false, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
+								true, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
 #else
 		IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-								false, true, targblock, 1, SampleCallback, (void *) buildstate);
+								true, true, targblock, 1, SampleCallback, (void *) buildstate);
 #endif
 	}
 }
@@ -167,18 +167,18 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 
 #ifdef IVFFLAT_KMEANS_DEBUG
 	buildstate->inertia += minDistance;
-	buildstate->listSums[closestCenter] += minDistance;
-	buildstate->listCounts[closestCenter]++;
 #endif
 
 	/* Create a virtual tuple */
 	ExecClearTuple(slot);
 	slot->tts_values[0] = Int32GetDatum(closestCenter);
 	slot->tts_isnull[0] = false;
-	slot->tts_values[1] = PointerGetDatum(tid);
+	slot->tts_values[1] = Int32GetDatum(ItemPointerGetBlockNumberNoCheck(tid));
 	slot->tts_isnull[1] = false;
-	slot->tts_values[2] = value;
+	slot->tts_values[2] = Int32GetDatum(ItemPointerGetOffsetNumberNoCheck(tid));
 	slot->tts_isnull[2] = false;
+	slot->tts_values[3] = value;
+	slot->tts_isnull[3] = false;
 	ExecStoreVirtualTuple(slot);
 
 	/*
@@ -200,6 +200,8 @@ GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot,
 {
 	Datum		value;
 	bool		isnull;
+	int			tupblk;
+	int			tupoff;
 
 #if PG_VERSION_NUM >= 100000
 	if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
@@ -208,11 +210,13 @@ GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot,
 #endif
 	{
 		*list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
-		value = slot_getattr(slot, 3, &isnull);
+		tupblk = DatumGetInt32(slot_getattr(slot, 2, &isnull));
+		tupoff = DatumGetInt32(slot_getattr(slot, 3, &isnull));
+		value = slot_getattr(slot, 4, &isnull);
 
 		/* Form the index tuple */
 		*itup = index_form_tuple(tupdesc, &value, &isnull);
-		(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull)));
+		ItemPointerSet(&(*itup)->t_tid, tupblk, tupoff);
 	}
 	else
 		*list = -1;
@@ -321,16 +325,17 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 
 	/* Create tuple description for sorting */
 #if PG_VERSION_NUM >= 120000
-	buildstate->tupdesc = CreateTemplateTupleDesc(3);
+	buildstate->tupdesc = CreateTemplateTupleDesc(4);
 #else
-	buildstate->tupdesc = CreateTemplateTupleDesc(3, false);
+	buildstate->tupdesc = CreateTemplateTupleDesc(4, false);
 #endif
 	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "blkno", INT4OID, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "offset", INT4OID, -1, 0);
 #if PG_VERSION_NUM >= 110000
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 4, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
 #else
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0]->atttypid, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 4, "vector", RelationGetDescr(index)->attrs[0]->atttypid, -1, 0);
 #endif
 
 #if PG_VERSION_NUM >= 120000
@@ -347,8 +352,6 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 
 #ifdef IVFFLAT_KMEANS_DEBUG
 	buildstate->inertia = 0;
-	buildstate->listSums = palloc0(sizeof(double) * buildstate->lists);
-	buildstate->listCounts = palloc0(sizeof(int) * buildstate->lists);
 #endif
 }
 
@@ -361,11 +364,6 @@ FreeBuildState(IvfflatBuildState * buildstate)
 	pfree(buildstate->centers);
 	pfree(buildstate->listInfo);
 	pfree(buildstate->normvec);
-
-#ifdef IVFFLAT_KMEANS_DEBUG
-	pfree(buildstate->listSums);
-	pfree(buildstate->listCounts);
-#endif
 }
 
 /*
@@ -473,51 +471,6 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
 	pfree(list);
 }
 
-/*
- * Print k-means metrics
- */
-#ifdef IVFFLAT_KMEANS_DEBUG
-static void
-PrintKmeansMetrics(IvfflatBuildState * buildstate)
-{
-	elog(INFO, "inertia: %.3e", buildstate->inertia);
-
-	/* Calculate Davies-Bouldin index */
-	if (buildstate->lists > 1)
-	{
-		double		db = 0.0;
-
-		/* Calculate average distance */
-		for (int i = 0; i < buildstate->lists; i++)
-		{
-			if (buildstate->listCounts[i] > 0)
-				buildstate->listSums[i] /= buildstate->listCounts[i];
-		}
-
-		for (int i = 0; i < buildstate->lists; i++)
-		{
-			double		max = 0.0;
-			double		distance;
-
-			for (int j = 0; j < buildstate->lists; j++)
-			{
-				if (j == i)
-					continue;
-
-				distance = DatumGetFloat8(FunctionCall2Coll(buildstate->procinfo, buildstate->collation, PointerGetDatum(VectorArrayGet(buildstate->centers, i)), PointerGetDatum(VectorArrayGet(buildstate->centers, j))));
-				distance = (buildstate->listSums[i] + buildstate->listSums[j]) / distance;
-
-				if (distance > max)
-					max = distance;
-			}
-			db += max;
-		}
-		db /= buildstate->lists;
-		elog(INFO, "davies-bouldin: %.3f", db);
-	}
-}
-#endif
-
 /*
  * Create entry pages
  */
@@ -556,7 +509,7 @@ CreateEntryPages(IvfflatBuildState * buildstate, ForkNumber forkNum)
 	tuplesort_performsort(buildstate->sortstate);
 
 #ifdef IVFFLAT_KMEANS_DEBUG
-	PrintKmeansMetrics(buildstate);
+	elog(INFO, "inertia: %.3e", buildstate->inertia);
 #endif
 
 	/* Insert */

src/ivfflat.c

@@ -13,6 +13,7 @@
 #endif
 
 int			ivfflat_probes;
+int			ivfflat_bound;
 static relopt_kind ivfflat_relopt_kind;
 
 /*
@@ -32,6 +33,10 @@ _PG_init(void)
 	DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
 							"Valid range is 1..lists.", &ivfflat_probes,
 							1, 1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
+
+	DefineCustomIntVariable("ivfflat.bound", "Sets the max results from index (experimental)",
+							NULL, &ivfflat_bound,
+							0, 0, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
 }
 
 /*

src/ivfflat.h

@@ -62,8 +62,14 @@
 #define IvfflatBench(name, code) (code)
 #endif
 
+#if PG_VERSION_NUM < 100000
+#define ItemPointerGetBlockNumberNoCheck ItemPointerGetBlockNumber
+#define ItemPointerGetOffsetNumberNoCheck ItemPointerGetOffsetNumber
+#endif
+
 /* Variables */
 extern int	ivfflat_probes;
+extern int	ivfflat_bound;
 
 typedef struct VectorArrayData
 {
@@ -117,8 +123,6 @@ typedef struct IvfflatBuildState
 
 #ifdef IVFFLAT_KMEANS_DEBUG
 	double		inertia;
-	double	   *listSums;
-	int		   *listCounts;
 #endif
 
 	/* Sampling */
@@ -162,33 +166,16 @@ typedef IvfflatListData * IvfflatList;
 
 typedef struct IvfflatScanList
 {
-	pairingheap_node ph_node;
 	BlockNumber startPage;
 	double		distance;
 }			IvfflatScanList;
 
-typedef struct IvfflatScanItem
-{
-	pairingheap_node ph_node;
-	BlockNumber searchPage;
-	double		distance;
-	ItemPointerData tid;
-}			IvfflatScanItem;
-
 typedef struct IvfflatScanOpaqueData
 {
 	int			probes;
-	int			stage;
+	bool		first;
 	Buffer		buf;
 
-	/* Items */
-	int			maxItems;
-	int			itemCount;
-	pairingheap *itemQueue;
-	IvfflatScanItem *items;
-	IvfflatScanItem **sortedItems;
-	bool		heapFull;
-
 	/* Sorting */
 	Tuplesortstate *sortstate;
 	TupleDesc	tupdesc;
@@ -200,9 +187,6 @@ typedef struct IvfflatScanOpaqueData
 	FmgrInfo   *normprocinfo;
 	Oid			collation;
 
-	/* Lists */
-	pairingheap *listQueue;
-	IvfflatScanList **sortedLists;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;
 

src/ivfscan.c

@@ -1,7 +1,5 @@
 #include "postgres.h"
 
-#include <float.h>
-
 #include "access/relscan.h"
 #include "ivfflat.h"
 #include "miscadmin.h"
@@ -19,32 +17,19 @@
  * Compare list distances
  */
 static int
-CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
+CompareLists(const void *a, const void *b)
 {
-	if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance)
+	double		diff = (((IvfflatScanList *) a)->distance - ((IvfflatScanList *) b)->distance);
+
+	if (diff > 0)
 		return 1;
 
-	if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
+	if (diff < 0)
 		return -1;
 
 	return 0;
 }
 
-/*
- * Compare item distances
- */
-static int
-CompareItems(const pairingheap_node *a, const pairingheap_node *b, void *arg)
-{
-	if (((const IvfflatScanItem *) a)->distance > ((const IvfflatScanItem *) b)->distance)
-		return 1;
-
-	if (((const IvfflatScanItem *) a)->distance < ((const IvfflatScanItem *) b)->distance)
-		return -1;
-
-	return ItemPointerCompare(&((IvfflatScanItem *) a)->tid, &((IvfflatScanItem *) b)->tid);
-}
-
 /*
  * Get lists and sort by distance
  */
@@ -59,10 +44,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
 	BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO;
 	int			listCount = 0;
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
-	int			i;
 	double		distance;
-	IvfflatScanList *scanlist;
-	double		maxDistance = DBL_MAX;
 
 	/* Search all list pages */
 	while (BlockNumberIsValid(nextblkno))
@@ -80,33 +62,9 @@ GetScanLists(IndexScanDesc scan, Datum value)
 			/* Use procinfo from the index instead of scan key for performance */
 			distance = DatumGetFloat8(FunctionCall2Coll(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
 
-			if (listCount < so->probes)
-			{
-				scanlist = &so->lists[listCount];
-				scanlist->startPage = list->startPage;
-				scanlist->distance = distance;
-				listCount++;
-
-				/* Add to heap */
-				pairingheap_add(so->listQueue, &scanlist->ph_node);
-
-				/* Calculate max distance */
-				if (listCount == so->probes)
-					maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
-			}
-			else if (distance < maxDistance)
-			{
-				/* Remove */
-				scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
-
-				/* Reuse */
-				scanlist->startPage = list->startPage;
-				scanlist->distance = distance;
-				pairingheap_add(so->listQueue, &scanlist->ph_node);
-
-				/* Update max distance */
-				maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
-			}
+			so->lists[listCount].startPage = list->startPage;
+			so->lists[listCount].distance = distance;
+			listCount++;
 		}
 
 		nextblkno = IvfflatPageGetOpaque(cpage)->nextblkno;
@@ -114,139 +72,12 @@ GetScanLists(IndexScanDesc scan, Datum value)
 		UnlockReleaseBuffer(cbuf);
 	}
 
-	for (i = 0; i < so->probes; i++)
-		so->sortedLists[i] = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
+	/* Sort by distance */
+	/* TODO Use heap for performance */
+	qsort(so->lists, listCount, sizeof(IvfflatScanList), CompareLists);
 
-	Assert(pairingheap_is_empty(so->listQueue));
-}
-
-/*
- * Get items
- */
-static void
-GetScanItemsQuick(IndexScanDesc scan, Datum value)
-{
-	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
-	Buffer		buf;
-	Page		page;
-	IndexTuple	itup;
-	BlockNumber searchPage;
-	OffsetNumber offno;
-	OffsetNumber maxoffno;
-	Datum		datum;
-	bool		isnull;
-	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
-	int			i;
-	double		distance;
-	IvfflatScanItem *scanitem;
-	double		maxDistance = DBL_MAX;
-
-	/*
-	 * Reuse same set of shared buffers for scan
-	 *
-	 * See postgres/src/backend/storage/buffer/README for description
-	 */
-	BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
-
-	/* Search closest probes lists */
-	for (i = 0; i < so->probes; i++)
-	{
-		/* Read closest lists first for performance */
-		searchPage = so->sortedLists[i]->startPage;
-
-		/* Search all entry pages for list */
-		while (BlockNumberIsValid(searchPage))
-		{
-			buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
-			LockBuffer(buf, BUFFER_LOCK_SHARE);
-			page = BufferGetPage(buf);
-			maxoffno = PageGetMaxOffsetNumber(page);
-
-			for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
-			{
-				itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, offno));
-				datum = index_getattr(itup, 1, tupdesc, &isnull);
-				distance = DatumGetFloat8(FunctionCall2Coll(so->procinfo, so->collation, datum, value));
-
-				if (so->itemCount < so->maxItems)
-				{
-					scanitem = &so->items[so->itemCount];
-					scanitem->searchPage = searchPage;
-					scanitem->tid = itup->t_tid;
-					scanitem->distance = distance;
-					so->itemCount++;
-
-					/* Add to heap */
-					pairingheap_add(so->itemQueue, &scanitem->ph_node);
-
-					/* Calculate max distance */
-					if (so->itemCount == so->maxItems)
-					{
-						maxDistance = ((IvfflatScanItem *) pairingheap_first(so->itemQueue))->distance;
-						scanitem = &so->items[so->itemCount];
-					}
-				}
-				else if (distance <= maxDistance)
-				{
-					/* Reuse */
-					scanitem->searchPage = searchPage;
-					scanitem->tid = itup->t_tid;
-					scanitem->distance = distance;
-					pairingheap_add(so->itemQueue, &scanitem->ph_node);
-
-					/* Remove */
-					scanitem = (IvfflatScanItem *) pairingheap_remove_first(so->itemQueue);
-
-					/* Update max distance */
-					maxDistance = ((IvfflatScanItem *) pairingheap_first(so->itemQueue))->distance;
-				}
-			}
-
-			searchPage = IvfflatPageGetOpaque(page)->nextblkno;
-
-			UnlockReleaseBuffer(buf);
-		}
-	}
-
-	for (i = 0; i < so->itemCount; i++)
-		so->sortedItems[i] = (IvfflatScanItem *) pairingheap_remove_first(so->itemQueue);
-
-	Assert(pairingheap_is_empty(so->itemQueue));
-}
-
-/*
- * Initialize sort
- */
-static void
-InitSort(IvfflatScanOpaque so)
-{
-	AttrNumber	attNums[] = {1, 2};
-	Oid			sortOperators[] = {Float8LessOperator, TIDLessOperator};
-	Oid			sortCollations[] = {InvalidOid, InvalidOid};
-	bool		nullsFirstFlags[] = {false, false};
-
-	/* Create tuple description for sorting */
-#if PG_VERSION_NUM >= 120000
-	so->tupdesc = CreateTemplateTupleDesc(3);
-#else
-	so->tupdesc = CreateTemplateTupleDesc(3, false);
-#endif
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
-
-	/* Prep sort */
-#if PG_VERSION_NUM >= 110000
-	so->sortstate = tuplesort_begin_heap(so->tupdesc, sizeof(attNums) / sizeof(attNums[0]), attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
-#else
-	so->sortstate = tuplesort_begin_heap(so->tupdesc, sizeof(attNums) / sizeof(attNums[0]), attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, false);
-#endif
-
-#if PG_VERSION_NUM >= 120000
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
-#else
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc);
-#endif
+	if (so->probes > listCount)
+		so->probes = listCount;
 }
 
 /*
@@ -264,8 +95,8 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	OffsetNumber maxoffno;
 	Datum		datum;
 	bool		isnull;
-	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
 	int			i;
+	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
 
 #if PG_VERSION_NUM >= 120000
 	TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
@@ -280,10 +111,14 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	 */
 	BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
 
+	/* Set the max number of results */
+	if (ivfflat_bound > 0)
+		tuplesort_set_bound(so->sortstate, ivfflat_bound);
+
 	/* Search closest probes lists */
 	for (i = 0; i < so->probes; i++)
 	{
-		searchPage = so->sortedLists[i]->startPage;
+		searchPage = so->lists[i].startPage;
 
 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -307,10 +142,12 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				ExecClearTuple(slot);
 				slot->tts_values[0] = FunctionCall2Coll(so->procinfo, so->collation, datum, value);
 				slot->tts_isnull[0] = false;
-				slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
+				slot->tts_values[1] = Int32GetDatum((int) ItemPointerGetBlockNumberNoCheck(&itup->t_tid));
 				slot->tts_isnull[1] = false;
-				slot->tts_values[2] = Int32GetDatum((int) searchPage);
+				slot->tts_values[2] = Int32GetDatum((int) ItemPointerGetOffsetNumberNoCheck(&itup->t_tid));
 				slot->tts_isnull[2] = false;
+				slot->tts_values[3] = Int32GetDatum((int) searchPage);
+				slot->tts_isnull[3] = false;
 				ExecStoreVirtualTuple(slot);
 
 				tuplesort_puttupleslot(so->sortstate, slot);
@@ -323,7 +160,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	}
 
 	tuplesort_performsort(so->sortstate);
-	tuplesort_skiptuples(so->sortstate, so->maxItems, true);
 }
 
 /*
@@ -332,34 +168,49 @@ GetScanItems(IndexScanDesc scan, Datum value)
 IndexScanDesc
 ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 {
-	IndexScanDesc scan = RelationGetIndexScan(index, nkeys, norderbys);
-	int			lists = IvfflatGetLists(scan->indexRelation);
-	int			probes = ivfflat_probes;
+	IndexScanDesc scan;
 	IvfflatScanOpaque so;
+	int			lists;
+	AttrNumber	attNums[] = {1};
+	Oid			sortOperators[] = {Float8LessOperator};
+	Oid			sortCollations[] = {InvalidOid};
+	bool		nullsFirstFlags[] = {false};
 
-	if (probes > lists)
-		probes = lists;
+	scan = RelationGetIndexScan(index, nkeys, norderbys);
+	lists = IvfflatGetLists(scan->indexRelation);
 
-	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
+	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + lists * sizeof(IvfflatScanList));
 	so->buf = InvalidBuffer;
-	so->stage = 0;
-	so->probes = probes;
+	so->first = true;
 
 	/* Set support functions */
 	so->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC);
 	so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
 	so->collation = index->rd_indcollation[0];
 
-	so->listQueue = pairingheap_allocate(CompareLists, scan);
-	so->sortedLists = palloc(sizeof(IvfflatScanItem *) * probes);
+	/* Create tuple description for sorting */
+#if PG_VERSION_NUM >= 120000
+	so->tupdesc = CreateTemplateTupleDesc(4);
+#else
+	so->tupdesc = CreateTemplateTupleDesc(4, false);
+#endif
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "blkno", INT4OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "offset", INT4OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 4, "indexblkno", INT4OID, -1, 0);
 
-	so->maxItems = 1024;
-	so->itemCount = 0;
-	so->itemQueue = pairingheap_allocate(CompareItems, scan);
-	so->items = palloc(sizeof(IvfflatScanItem) * (so->maxItems + 1));
-	so->sortedItems = palloc(sizeof(IvfflatScanItem *) * so->maxItems);
+	/* Prep sort */
+#if PG_VERSION_NUM >= 110000
+	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
+#else
+	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, false);
+#endif
 
-	so->sortstate = NULL;
+#if PG_VERSION_NUM >= 120000
+	so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
+#else
+	so->slot = MakeSingleTupleTableSlot(so->tupdesc);
+#endif
 
 	scan->opaque = so;
 
@@ -375,14 +226,12 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
 #if PG_VERSION_NUM >= 130000
-	if (so->sortstate != NULL)
+	if (!so->first)
 		tuplesort_reset(so->sortstate);
 #endif
 
-	so->stage = 0;
-	pairingheap_reset(so->listQueue);
-	pairingheap_reset(so->itemQueue);
-	so->itemCount = 0;
+	so->first = true;
+	so->probes = ivfflat_probes;
 
 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -405,7 +254,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 	 */
 	Assert(ScanDirectionIsForward(dir));
 
-	if (so->stage == 0)
+	if (so->first)
 	{
 		Datum		value;
 
@@ -427,101 +276,43 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		}
 
 		IvfflatBench("GetScanLists", GetScanLists(scan, value));
-		IvfflatBench("GetScanItemsQuick", GetScanItemsQuick(scan, value));
-		so->heapFull = so->itemCount == so->maxItems;
-		so->stage++;
+		IvfflatBench("GetScanItems", GetScanItems(scan, value));
+		so->first = false;
 
 		/* Clean up if we allocated a new value */
 		if (value != scan->orderByData->sk_argument)
 			pfree(DatumGetPointer(value));
 	}
 
-	if (so->stage == 1)
-	{
-		if (so->itemCount > 0)
-		{
-			IvfflatScanItem *scanitem;
-
-			so->itemCount--;
-
-			scanitem = so->sortedItems[so->itemCount];
-
-#if PG_VERSION_NUM >= 120000
-			scan->xs_heaptid = scanitem->tid;
-#else
-			scan->xs_ctup.t_sef = scanitem->tid;
-#endif
-
-			if (BufferIsValid(so->buf))
-				ReleaseBuffer(so->buf);
-
-			/*
-			 * An index scan must maintain a pin on the index page holding the
-			 * item last returned by amgettuple
-			 *
-			 * https://www.postgresql.org/docs/current/index-locking.html
-			 */
-			so->buf = ReadBuffer(scan->indexRelation, scanitem->searchPage);
-
-			scan->xs_recheckorderby = false;
-			return true;
-		}
-		else if (so->heapFull)
-		{
-			Datum value = scan->orderByData->sk_argument;
-
-			if (so->normprocinfo != NULL)
-			{
-				/* No items will match if normalization fails */
-				if (!IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL))
-					return false;
-			}
-
-			if (so->sortstate == NULL)
-				InitSort(so);
-
-			IvfflatBench("GetScanItems", GetScanItems(scan, value));
-			so->stage++;
-
-			/* Clean up if we allocated a new value */
-			if (value != scan->orderByData->sk_argument)
-				pfree(DatumGetPointer(value));
-		}
-		else
-			so->stage = 3;
-	}
-
-	if (so->stage == 2)
-	{
 #if PG_VERSION_NUM >= 100000
-		if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
+	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
 #else
-		if (tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL))
+	if (tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL))
 #endif
-		{
-			ItemPointer tid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
-			BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+	{
+		BlockNumber blkno = DatumGetInt32(slot_getattr(so->slot, 2, &so->isnull));
+		OffsetNumber offset = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 4, &so->isnull));
 
 #if PG_VERSION_NUM >= 120000
-			scan->xs_heaptid = *tid;
+		ItemPointerSet(&scan->xs_heaptid, blkno, offset);
 #else
-			scan->xs_ctup.t_self = *tid;
+		ItemPointerSet(&scan->xs_ctup.t_self, blkno, offset);
 #endif
 
-			if (BufferIsValid(so->buf))
-				ReleaseBuffer(so->buf);
+		if (BufferIsValid(so->buf))
+			ReleaseBuffer(so->buf);
 
-			/*
-			 * An index scan must maintain a pin on the index page holding the
-			 * item last returned by amgettuple
-			 *
-			 * https://www.postgresql.org/docs/current/index-locking.html
-			 */
-			so->buf = ReadBuffer(scan->indexRelation, indexblkno);
+		/*
+		 * An index scan must maintain a pin on the index page holding the
+		 * item last returned by amgettuple
+		 *
+		 * https://www.postgresql.org/docs/current/index-locking.html
+		 */
+		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
 
-			scan->xs_recheckorderby = false;
-			return true;
-		}
+		scan->xs_recheckorderby = false;
+		return true;
 	}
 
 	return false;
@@ -539,15 +330,7 @@ ivfflatendscan(IndexScanDesc scan)
 	if (BufferIsValid(so->buf))
 		ReleaseBuffer(so->buf);
 
-	pairingheap_free(so->listQueue);
-	pfree(so->sortedLists);
-
-	if (so->sortstate != NULL)
-		tuplesort_end(so->sortstate);
-
-	pairingheap_free(so->itemQueue);
-	pfree(so->items);
-	pfree(so->sortedItems);
+	tuplesort_end(so->sortstate);
 
 	pfree(so);
 	scan->opaque = NULL;

test/t/003_recall.pl

@@ -2,16 +2,15 @@ use strict;
 use warnings;
 use PostgresNode;
 use TestLib;
-use Test::More tests => 9;
+use Test::More tests => 2;
 
 my $node;
 my @queries = ();
-my @expected;
-my $limit = 20;
+my @expected = ();
 
 sub test_recall
 {
-	my ($probes, $min, $operator) = @_;
+	my ($probes, $min) = @_;
 	my $correct = 0;
 	my $total = 0;
 
@@ -19,7 +18,7 @@ sub test_recall
 		my $actual = $node->safe_psql("postgres", qq(
 			SET enable_seqscan = off;
 			SET ivfflat.probes = $probes;
-			SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
+			SELECT i FROM tst ORDER BY v <-> '$queries[$i]' LIMIT 10;
 		));
 		my @actual_ids = split("\n", $actual);
 		my %actual_set = map { $_ => 1 } @actual_ids;
@@ -34,7 +33,7 @@ sub test_recall
 		}
 	}
 
-	cmp_ok($correct / $total, ">=", $min, $operator);
+	cmp_ok($correct / $total, ">=", $min);
 }
 
 # Initialize node
@@ -57,32 +56,17 @@ for (1..20) {
 	push(@queries, "[$r1,$r2,$r3]");
 }
 
-# Check each index type
-my @operators = ("<->", "<#>", "<=>");
-
-foreach (@operators) {
-	my $operator = $_;
-
-	# Get exact results
-	@expected = ();
-	foreach (@queries) {
-		my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
-		push(@expected, $res);
-	}
-
-	# Add index
-	my $opclass;
-	if ($operator == "<->") {
-		$opclass = "vector_l2_ops";
-	} elsif ($operator == "<#>") {
-		$opclass = "vector_ip_ops";
-	} else {
-		$opclass = "vector_cosine_ops";
-	}
-	$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
-
-	# Test approximate results
-	test_recall(1, 0.75, $operator);
-	test_recall(10, 0.95, $operator);
-	test_recall(100, 1.0, $operator);
+# Get exact results
+foreach (@queries) {
+	my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v <-> '$_' LIMIT 10;");
+	push(@expected, $res);
 }
+
+# Add index
+$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
+
+# Test approximate results
+test_recall(1, 0.8);
+
+# Test probes
+test_recall(100, 1.0);

test/t/005_query_recall.pl

@@ -1,45 +0,0 @@
-use strict;
-use warnings;
-use PostgresNode;
-use TestLib;
-use Test::More tests => 60;
-
-# Initialize node
-my $node = get_new_node('node');
-$node->init;
-$node->start;
-
-# Create table
-$node->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node->safe_psql("postgres", "CREATE TABLE tst (i int4 primary key, v vector(3));");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
-);
-
-# Check each index type
-my @operators = ("<->", "<#>", "<=>");
-foreach (@operators) {
-	my $operator = $_;
-
-	# Add index
-	my $opclass;
-	if ($operator == "<->") {
-		$opclass = "vector_l2_ops";
-	} elsif ($operator == "<#>") {
-		$opclass = "vector_ip_ops";
-	} else {
-		$opclass = "vector_cosine_ops";
-	}
-	$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
-
-	# Test 100% recall
-	for (1..20) {
-		my $i = int(rand() * 100000);
-		my $query = $node->safe_psql("postgres", "SELECT v FROM tst WHERE i = $i;");
-		my $res = $node->safe_psql("postgres", qq(
-			SET enable_seqscan = off;
-			SELECT v FROM tst ORDER BY v <-> '$query' LIMIT 1;
-		));
-		is($res, $query);
-	}
-}

test/t/006_lists.pl

@@ -1,31 +0,0 @@
-use strict;
-use warnings;
-use PostgresNode;
-use TestLib;
-use Test::More tests => 3;
-
-# Initialize node
-my $node = get_new_node('node');
-$node->init;
-$node->start;
-
-# Create table
-$node->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node->safe_psql("postgres", "CREATE TABLE tst (v vector(3));");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
-);
-
-$node->safe_psql("postgres", "CREATE INDEX lists50 ON tst USING ivfflat (v) WITH (lists = 50);");
-$node->safe_psql("postgres", "CREATE INDEX lists100 ON tst USING ivfflat (v) WITH (lists = 100);");
-
-# Test prefers more lists
-my $res = $node->safe_psql("postgres", "EXPLAIN SELECT v FROM tst ORDER BY v <-> '[0.5,0.5,0.5]' LIMIT 10;");
-like($res, qr/lists100/);
-unlike($res, qr/lists50/);
-
-# Test errors with too much memory
-my ($ret, $stdout, $stderr) = $node->psql("postgres",
-	"CREATE INDEX lists10000 ON tst USING ivfflat (v) WITH (lists = 10000);"
-);
-like($stderr, qr/memory required is/);

test/t/007_stages.pl

@@ -1,38 +0,0 @@
-use strict;
-use warnings;
-use PostgresNode;
-use TestLib;
-use Test::More tests => 2;
-
-# Initialize node
-my $node = get_new_node('node');
-$node->init;
-$node->start;
-
-# Create table
-$node->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i, ARRAY[i % 1000, i % 1000, i % 1000] FROM generate_series(1, 10000) i;"
-);
-
-my @limits = (128, 2048);
-my @expected = ();
-
-foreach (@limits) {
-	my $res = $node->safe_psql("postgres", "SELECT i, v FROM tst ORDER BY v <-> '[0,0,0]', i LIMIT $_;");
-	push(@expected, $res);
-}
-
-$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v) WITH (lists = 5);");
-
-for my $i (0 .. $#limits) {
-	my $res = $node->safe_psql("postgres", qq(
-		SET enable_seqscan = off;
-		SET ivfflat.probes = 5;
-		WITH tmp AS (
-			SELECT *, v <-> '[0,0,0]' AS d FROM tst ORDER BY v <-> '[0,0,0]' LIMIT $limits[$i]
-		) SELECT i, v FROM tmp ORDER BY d, i;
-	));
-	is($res, $expected[$i]);
-}