Use pairing heap [skip ci]

README.md

@@ -312,10 +312,6 @@ Two things you can try are:
 1. use dimensionality reduction
 2. compile Postgres with a larger block size (`./configure --with-blocksize=32`) and edit the limit in `src/ivfflat.h`
 
-#### Why am I seeing less results after adding an index?
-
-The index was likely created with too little data for the number of lists. Drop the index until the table has more data.
-
 ## Reference
 
 ### Vector Type

src/ivfflat.h

@@ -187,6 +187,14 @@ typedef struct IvfflatScanList
 	double		distance;
 }			IvfflatScanList;
 
+typedef struct IvfflatScanItem
+{
+	pairingheap_node ph_node;
+	BlockNumber searchPage;
+	double		distance;
+	ItemPointerData tid;
+}			IvfflatScanItem;
+
 typedef struct IvfflatScanOpaqueData
 {
 	int			probes;
@@ -204,9 +212,15 @@ typedef struct IvfflatScanOpaqueData
 	FmgrInfo   *normprocinfo;
 	Oid			collation;
 
+	/* Items */
+	int			maxItems;
+	int			itemCount;
+	pairingheap *itemQueue;
+	IvfflatScanItem *items;
+	IvfflatScanItem **sortedItems;
+
 	/* Lists */
 	pairingheap *listQueue;
-	double		minDistance;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;
 

src/ivfscan.c

@@ -26,6 +26,21 @@ CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
 	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 0;
+}
+
 /*
  * Get lists and sort by distance
  */
@@ -60,9 +75,6 @@ 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 (distance < so->minDistance)
-				so->minDistance = distance;
-
 			if (listCount < so->probes)
 			{
 				scanlist = &so->lists[listCount];
@@ -114,13 +126,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	Datum		datum;
 	bool		isnull;
 	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
-	double		tuples = 0;
+	int			i;
-
+	double		distance;
-#if PG_VERSION_NUM >= 120000
+	IvfflatScanItem *scanitem;
-	TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
+	double		maxDistance = DBL_MAX;
-#else
-	TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc);
-#endif
 
 	/*
 	 * Reuse same set of shared buffers for scan
@@ -132,13 +141,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	/* Search closest probes lists */
 	while (!pairingheap_is_empty(so->listQueue))
 	{
-		IvfflatScanList *scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
+		searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
-
-		/* Query-aware dynamic pruning */
-		if (fabs(scanlist->distance) > 1.5 * fabs(so->minDistance))
-			continue;
-
-		searchPage = scanlist->startPage;
 
 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -152,25 +155,40 @@ GetScanItems(IndexScanDesc scan, Datum value)
 			{
 				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)
-				 * Add virtual tuple
+				{
-				 *
+					scanitem = &so->items[so->itemCount];
-				 * Use procinfo from the index instead of scan key for
+					scanitem->searchPage = searchPage;
-				 * performance
+					scanitem->tid = itup->t_tid;
-				 */
+					scanitem->distance = distance;
-				ExecClearTuple(slot);
+					so->itemCount++;
-				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_isnull[1] = false;
-				slot->tts_values[2] = Int32GetDatum((int) searchPage);
-				slot->tts_isnull[2] = false;
-				ExecStoreVirtualTuple(slot);
 
-				tuplesort_puttupleslot(so->sortstate, slot);
+					/* Add to heap */
+					pairingheap_add(so->itemQueue, &scanitem->ph_node);
 
-				tuples++;
+					/* 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;
@@ -179,14 +197,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
 		}
 	}
 
-	/* TODO Scan more lists */
+	for (i = 0; i < so->itemCount; i++)
-	if (tuples < 100)
+		so->sortedItems[i] = (IvfflatScanItem *) pairingheap_remove_first(so->itemQueue);
-		ereport(DEBUG1,
-				(errmsg("index scan found few tuples"),
-				 errdetail("index may have been created without data or lists is too high"),
-				 errhint("recreate the index and possibly decrease lists")));
 
-	tuplesort_performsort(so->sortstate);
+	Assert(pairingheap_is_empty(so->itemQueue));
 }
 
 /*
@@ -198,10 +212,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	IndexScanDesc scan;
 	IvfflatScanOpaque so;
 	int			lists;
-	AttrNumber	attNums[] = {1};
-	Oid			sortOperators[] = {Float8LessOperator};
-	Oid			sortCollations[] = {InvalidOid};
-	bool		nullsFirstFlags[] = {false};
 	int			probes = ivfflat_probes;
 
 	scan = RelationGetIndexScan(index, nkeys, norderbys);
@@ -220,27 +230,14 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
 	so->collation = index->rd_indcollation[0];
 
-	/* 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 */
-	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
-
-#if PG_VERSION_NUM >= 120000
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
-#else
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc);
-#endif
-
 	so->listQueue = pairingheap_allocate(CompareLists, scan);
 
+	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);
+
 	scan->opaque = so;
 
 	return scan;
@@ -254,14 +251,10 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 {
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
-#if PG_VERSION_NUM >= 130000
-	if (!so->first)
-		tuplesort_reset(so->sortstate);
-#endif
-
 	so->first = true;
 	pairingheap_reset(so->listQueue);
-	so->minDistance = DBL_MAX;
+	pairingheap_reset(so->itemQueue);
+	so->itemCount = 0;
 
 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -321,15 +314,18 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 			pfree(DatumGetPointer(value));
 	}
 
-	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
+	if (so->itemCount > 0)
 	{
-		ItemPointer tid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
+		IvfflatScanItem *scanitem;
-		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+
+		so->itemCount--;
+
+		scanitem = so->sortedItems[so->itemCount];
 
 #if PG_VERSION_NUM >= 120000
-		scan->xs_heaptid = *tid;
+		scan->xs_heaptid = scanitem->tid;
 #else
-		scan->xs_ctup.t_self = *tid;
+		scan->xs_ctup.t_self = scanitem->tid;
 #endif
 
 		if (BufferIsValid(so->buf))
@@ -341,7 +337,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		 *
 		 * https://www.postgresql.org/docs/current/index-locking.html
 		 */
-		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
+		so->buf = ReadBuffer(scan->indexRelation, scanitem->searchPage);
 
 		scan->xs_recheckorderby = false;
 		return true;
@@ -363,7 +359,10 @@ ivfflatendscan(IndexScanDesc scan)
 		ReleaseBuffer(so->buf);
 
 	pairingheap_free(so->listQueue);
-	tuplesort_end(so->sortstate);
+
+	pairingheap_free(so->itemQueue);
+	pfree(so->items);
+	pfree(so->sortedItems);
 
 	pfree(so);
 	scan->opaque = NULL;