Limit max probes [skip ci]

src/hnsw.c

@@ -18,7 +18,6 @@
 #endif
 
 int			hnsw_ef_search;
-int			hnsw_time_budget;
 int			hnsw_lock_tranche_id;
 static relopt_kind hnsw_relopt_kind;
 
@@ -69,10 +68,6 @@ HnswInit(void)
 							"Valid range is 1..1000.", &hnsw_ef_search,
 							HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
 
-	DefineCustomIntVariable("hnsw.time_budget", "Sets the time budget for search",
-							NULL, &hnsw_time_budget,
-							-1, -1, INT_MAX, PGC_USERSET, GUC_UNIT_MS, NULL, NULL, NULL);
-
 	MarkGUCPrefixReserved("hnsw");
 }
 

src/hnsw.h

@@ -106,7 +106,6 @@
 
 /* Variables */
 extern int	hnsw_ef_search;
-extern int	hnsw_time_budget;
 extern int	hnsw_lock_tranche_id;
 
 typedef struct HnswElementData HnswElementData;
@@ -156,12 +155,13 @@ struct HnswNeighborArray
 	HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
 };
 
-typedef struct HnswPairingHeapNode
+typedef struct HnswSearchCandidate
 {
-	HnswCandidate *inner;
 	pairingheap_node c_node;
 	pairingheap_node w_node;
-}			HnswPairingHeapNode;
+	HnswElementPtr element;
+	float		distance;
+}			HnswSearchCandidate;
 
 /* HNSW index options */
 typedef struct HnswOptions
@@ -331,7 +331,6 @@ typedef struct HnswScanOpaqueData
 	bool		first;
 	List	   *w;
 	MemoryContext tmpCtx;
-	instr_time	start;
 
 	/* Support functions */
 	FmgrInfo   *procinfo;
@@ -376,14 +375,14 @@ bool		HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
 Buffer		HnswNewBuffer(Relation index, ForkNumber forkNum);
 void		HnswInitPage(Buffer buf, Page page);
 void		HnswInit(void);
-List	   *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, instr_time *start);
+List	   *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
 HnswElement HnswGetEntryPoint(Relation index);
 void		HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
 void	   *HnswAlloc(HnswAllocator * allocator, Size size);
 HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
 HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
 void		HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
-HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
+HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
 void		HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
 void		HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
 void		HnswAddHeapTid(HnswElement element, ItemPointer heaptid);

src/hnswscan.c

@@ -3,7 +3,6 @@
 #include "access/relscan.h"
 #include "hnsw.h"
 #include "pgstat.h"
-#include "portability/instr_time.h"
 #include "storage/bufmgr.h"
 #include "storage/lmgr.h"
 #include "utils/memutils.h"
@@ -23,7 +22,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
 	int			m;
 	HnswElement entryPoint;
 	char	   *base = NULL;
-	instr_time *start = hnsw_time_budget == -1 ? NULL : &so->start;
 
 	/* Get m and entry point */
 	HnswGetMetaPageInfo(index, &m, &entryPoint);
@@ -35,11 +33,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
 
 	for (int lc = entryPoint->level; lc >= 1; lc--)
 	{
-		w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, start);
+		w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
 		ep = w;
 	}
 
-	return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, start);
+	return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
 }
 
 /*
@@ -92,9 +90,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 	so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
 	so->collation = index->rd_indcollation[0];
 
-	if (hnsw_time_budget != -1)
-		INSTR_TIME_SET_CURRENT(so->start);
-
 	scan->opaque = so;
 
 	return scan;
@@ -166,14 +161,14 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		so->first = false;
 
 #if defined(HNSW_MEMORY)
-		elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(so->tmpCtx, false) / (1024 * 1024));
+		elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
 #endif
 	}
 
 	while (list_length(so->w) > 0)
 	{
 		char	   *base = NULL;
-		HnswCandidate *hc = llast(so->w);
+		HnswSearchCandidate *hc = llast(so->w);
 		HnswElement element = HnswPtrAccess(base, hc->element);
 		ItemPointer heaptid;
 

src/hnswutils.c

@@ -9,7 +9,6 @@
 #include "fmgr.h"
 #include "hnsw.h"
 #include "lib/pairingheap.h"
-#include "portability/instr_time.h"
 #include "sparsevec.h"
 #include "storage/bufmgr.h"
 #include "utils/datum.h"
@@ -609,10 +608,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
 /*
  * Create a candidate for the entry point
  */
-HnswCandidate *
+HnswSearchCandidate *
 HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
 {
-	HnswCandidate *hc = palloc(sizeof(HnswCandidate));
+	HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate));
 
 	HnswPtrStore(base, hc->element, entryPoint);
 	if (index == NULL)
@@ -622,8 +621,8 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
 	return hc;
 }
 
-#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner)
+#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
-#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner)
+#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
 
 /*
  * Compare candidate distances
@@ -631,10 +630,10 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
 static int
 CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
 {
-	if (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance)
+	if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance)
 		return 1;
 
-	if (HnswGetPairingHeapCandidateConst(c_node, a)->distance > HnswGetPairingHeapCandidateConst(c_node, b)->distance)
+	if (HnswGetSearchCandidateConst(c_node, a)->distance > HnswGetSearchCandidateConst(c_node, b)->distance)
 		return -1;
 
 	return 0;
@@ -646,27 +645,15 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
 static int
 CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
 {
-	if (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance)
+	if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
 		return -1;
 
-	if (HnswGetPairingHeapCandidateConst(w_node, a)->distance > HnswGetPairingHeapCandidateConst(w_node, b)->distance)
+	if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
 		return 1;
 
 	return 0;
 }
 
-/*
- * Create a pairing heap node for a candidate
- */
-static HnswPairingHeapNode *
-CreatePairingHeapNode(HnswCandidate * c)
-{
-	HnswPairingHeapNode *node = palloc(sizeof(HnswPairingHeapNode));
-
-	node->inner = c;
-	return node;
-}
-
 /*
  * Init visited
  */
@@ -800,7 +787,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
  * Algorithm 2 from paper
  */
 List *
-HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, instr_time *start)
+HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
 {
 	List	   *w = NIL;
 	pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -826,15 +813,13 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
 	/* Add entry points to v, C, and W */
 	foreach(lc2, ep)
 	{
-		HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
+		HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2);
 		bool		found;
-		HnswPairingHeapNode *node;
 
 		AddToVisited(base, &v, hc->element, index, &found);
 
-		node = CreatePairingHeapNode(hc);
+		pairingheap_add(C, &hc->c_node);
-		pairingheap_add(C, &node->c_node);
+		pairingheap_add(W, &hc->w_node);
-		pairingheap_add(W, &node->w_node);
 
 		/*
 		 * Do not count elements being deleted towards ef when vacuuming. It
@@ -847,21 +832,10 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
 
 	while (!pairingheap_is_empty(C))
 	{
-		HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C));
+		HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C));
-		HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
+		HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
 		HnswElement cElement;
 
-		/* Check time budget */
-		if (start != NULL)
-		{
-			instr_time	duration;
-
-			INSTR_TIME_SET_CURRENT(duration);
-			INSTR_TIME_SUBTRACT(duration, *start);
-			if (INSTR_TIME_GET_MILLISEC(duration) >= hnsw_time_budget)
-				break;
-		}
-
 		if (c->distance > f->distance)
 			break;
 
@@ -875,12 +849,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
 		for (int i = 0; i < unvisitedLength; i++)
 		{
 			HnswElement eElement;
-			HnswCandidate *e;
+			HnswSearchCandidate *e;
-			HnswPairingHeapNode *node;
 			float		eDistance;
 			bool		alwaysAdd = wlen < ef;
 
-			f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
+			f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
 
 			if (index == NULL)
 			{
@@ -911,13 +884,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
 				continue;
 
 			/* Create a new candidate */
-			e = palloc(sizeof(HnswCandidate));
+			e = palloc(sizeof(HnswSearchCandidate));
 			HnswPtrStore(base, e->element, eElement);
 			e->distance = eDistance;
-
+			pairingheap_add(C, &e->c_node);
-			node = CreatePairingHeapNode(e);
+			pairingheap_add(W, &e->w_node);
-			pairingheap_add(C, &node->c_node);
-			pairingheap_add(W, &node->w_node);
 
 			/*
 			 * Do not count elements being deleted towards ef when vacuuming.
@@ -938,7 +909,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
 	/* Add each element of W to w */
 	while (!pairingheap_is_empty(W))
 	{
-		HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W));
+		HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
 
 		w = lappend(w, hc);
 	}
@@ -1303,7 +1274,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
 	/* 1st phase: greedy search to insert level */
 	for (int lc = entryLevel; lc >= level + 1; lc--)
 	{
-		w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL);
+		w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
 		ep = w;
 	}
 
@@ -1319,16 +1290,27 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
 	{
 		int			lm = HnswGetLayerM(m, lc);
 		List	   *neighbors;
-		List	   *lw;
+		List	   *lw = NIL;
+		ListCell   *lc2;
 
-		w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL);
+		w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement);
+
+		/* Convert search candidates to candidates */
+		foreach(lc2, w)
+		{
+			HnswSearchCandidate *sc = lfirst(lc2);
+			HnswCandidate *hc = palloc(sizeof(HnswCandidate));
+
+			hc->element = sc->element;
+			hc->distance = sc->distance;
+
+			lw = lappend(lw, hc);
+		}
 
 		/* Elements being deleted or skipped can help with search */
 		/* but should be removed before selecting neighbors */
 		if (index != NULL)
-			lw = RemoveElements(base, w, skipElement);
+			lw = RemoveElements(base, lw, skipElement);
-		else
-			lw = w;
 
 		/*
 		 * Candidates are sorted, but not deterministically. Could set

src/ivfflat.c

@@ -17,6 +17,8 @@
 #endif
 
 int			ivfflat_probes;
+int			ivfflat_max_probes;
+bool		ivfflat_streaming;
 static relopt_kind ivfflat_relopt_kind;
 
 /*
@@ -33,6 +35,14 @@ IvfflatInit(void)
 							"Valid range is 1..lists.", &ivfflat_probes,
 							IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
 
+	DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative scans",
+							NULL, &ivfflat_max_probes,
+							-1, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
+
+	DefineCustomBoolVariable("ivfflat.streaming", "Use streaming mode",
+							 NULL, &ivfflat_streaming,
+							 IVFFLAT_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL);
+
 	MarkGUCPrefixReserved("ivfflat");
 }
 
@@ -57,6 +67,35 @@ ivfflatbuildphasename(int64 phasenum)
 	}
 }
 
+/*
+ * Estimate the number of probes for iterative scans
+ */
+static int
+EstimateProbes(PlannerInfo *root, IndexPath *path, int lists)
+{
+	double		selectivity = 1;
+	ListCell   *lc;
+	double		tuplesPerList;
+
+	/* Cannot estimate without limit */
+	/* limit_tuples includes offset */
+	if (root->limit_tuples < 0)
+		return 0;
+
+	/* Get the selectivity of non-index conditions */
+	foreach(lc, path->indexinfo->indrestrictinfo)
+	{
+		RestrictInfo *rinfo = lfirst(lc);
+
+		/* Skip DEFAULT_INEQ_SEL since it may be a distance filter */
+		if (rinfo->norm_selec >= 0 && rinfo->norm_selec <= 1 && rinfo->norm_selec != (Selectivity) DEFAULT_INEQ_SEL)
+			selectivity *= rinfo->norm_selec;
+	}
+
+	tuplesPerList = path->indexinfo->tuples / (double) lists;
+	return root->limit_tuples / (tuplesPerList * selectivity);
+}
+
 /*
  * Estimate the cost of an index scan
  */
@@ -68,6 +107,7 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 {
 	GenericCosts costs;
 	int			lists;
+	int			probes;
 	double		ratio;
 	double		spc_seq_page_cost;
 	Relation	index;
@@ -89,8 +129,17 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 	IvfflatGetMetaPageInfo(index, &lists, NULL);
 	index_close(index, NoLock);
 
+	probes = ivfflat_probes;
+	if (ivfflat_streaming)
+	{
+		probes = Max(probes, EstimateProbes(root, path, lists));
+
+		if (ivfflat_max_probes != -1)
+			probes = Min(probes, ivfflat_max_probes);
+	}
+
 	/* Get the ratio of lists that we need to visit */
-	ratio = ((double) ivfflat_probes) / lists;
+	ratio = ((double) probes) / lists;
 	if (ratio > 1.0)
 		ratio = 1.0;
 

src/ivfflat.h

@@ -43,6 +43,7 @@
 #define IVFFLAT_MIN_LISTS		1
 #define IVFFLAT_MAX_LISTS		32768
 #define IVFFLAT_DEFAULT_PROBES	1
+#define IVFFLAT_DEFAULT_STREAMING	false
 
 /* Build phases */
 /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
@@ -80,6 +81,8 @@
 
 /* Variables */
 extern int	ivfflat_probes;
+extern int	ivfflat_max_probes;
+extern bool ivfflat_streaming;
 
 typedef struct VectorArrayData
 {
@@ -247,14 +250,17 @@ typedef struct IvfflatScanOpaqueData
 {
 	const		IvfflatTypeInfo *typeInfo;
 	int			probes;
+	int			maxProbes;
 	int			dimensions;
 	bool		first;
+	Datum		value;
 
 	/* Sorting */
 	Tuplesortstate *sortstate;
 	TupleDesc	tupdesc;
-	TupleTableSlot *slot;
+	TupleTableSlot *vslot;
-	bool		isnull;
+	TupleTableSlot *mslot;
+	BufferAccessStrategy bas;
 
 	/* Support functions */
 	FmgrInfo   *procinfo;
@@ -264,6 +270,8 @@ typedef struct IvfflatScanOpaqueData
 
 	/* Lists */
 	pairingheap *listQueue;
+	BlockNumber *startPages;
+	int			currentIndex;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;
 

src/ivfscan.c

@@ -15,16 +15,19 @@
 #include "utils/memutils.h"
 #endif
 
+#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
+#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
+
 /*
  * Compare list distances
  */
 static int
 CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
 {
-	if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance)
+	if (GetScanListConst(a)->distance > GetScanListConst(b)->distance)
 		return 1;
 
-	if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
+	if (GetScanListConst(a)->distance < GetScanListConst(b)->distance)
 		return -1;
 
 	return 0;
@@ -62,7 +65,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
 			/* Use procinfo from the index instead of scan key for performance */
 			distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
 
-			if (listCount < so->probes)
+			if (listCount < so->maxProbes)
 			{
 				IvfflatScanList *scanlist;
 
@@ -76,14 +79,14 @@ GetScanLists(IndexScanDesc scan, Datum value)
 
 				/* Calculate max distance */
 				if (listCount == so->probes)
-					maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
+					maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
 			}
 			else if (distance < maxDistance)
 			{
 				IvfflatScanList *scanlist;
 
 				/* Remove */
-				scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
+				scanlist = GetScanList(pairingheap_remove_first(so->listQueue));
 
 				/* Reuse */
 				scanlist->startPage = list->startPage;
@@ -91,7 +94,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
 				pairingheap_add(so->listQueue, &scanlist->ph_node);
 
 				/* Update max distance */
-				maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
+				maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
 			}
 		}
 
@@ -99,6 +102,11 @@ GetScanLists(IndexScanDesc scan, Datum value)
 
 		UnlockReleaseBuffer(cbuf);
 	}
+
+	for (int i = listCount - 1; i >= 0; i--)
+		so->startPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
+
+	Assert(pairingheap_is_empty(so->listQueue));
 }
 
 /*
@@ -110,19 +118,15 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
 	double		tuples = 0;
-	TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
+	TupleTableSlot *slot = so->vslot;
+	int			batchProbes = 0;
 
-	/*
+	tuplesort_reset(so->sortstate);
-	 * 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 */
-	while (!pairingheap_is_empty(so->listQueue))
+	while (so->currentIndex < so->maxProbes && (++batchProbes) <= so->probes)
 	{
-		BlockNumber searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
+		BlockNumber searchPage = so->startPages[so->currentIndex++];
 
 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -131,7 +135,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 			Page		page;
 			OffsetNumber maxoffno;
 
-			buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
+			buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, so->bas);
 			LockBuffer(buf, BUFFER_LOCK_SHARE);
 			page = BufferGetPage(buf);
 			maxoffno = PageGetMaxOffsetNumber(page);
@@ -170,15 +174,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
 		}
 	}
 
-	FreeAccessStrategy(bas);
+	if (tuples < 100 && !ivfflat_streaming)
-
-	if (tuples < 100)
 		ereport(DEBUG1,
 				(errmsg("index scan found few tuples"),
 				 errdetail("Index may have been created with little data."),
 				 errhint("Recreate the index and possibly decrease lists.")));
 
 	tuplesort_performsort(so->sortstate);
+
+#if defined(IVFFLAT_MEMORY)
+	elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
+#endif
 }
 
 /*
@@ -246,6 +252,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	int			lists;
 	int			dimensions;
 	int			probes = ivfflat_probes;
+	int			maxProbes;
 
 	scan = RelationGetIndexScan(index, nkeys, norderbys);
 
@@ -255,10 +262,21 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	if (probes > lists)
 		probes = lists;
 
-	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
+	if (ivfflat_streaming)
+	{
+		if (ivfflat_max_probes == -1)
+			maxProbes = lists;
+		else
+			maxProbes = Min(ivfflat_max_probes, lists);
+	}
+	else
+		maxProbes = probes;
+
+	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList));
 	so->typeInfo = IvfflatGetTypeInfo(index);
 	so->first = true;
 	so->probes = probes;
+	so->maxProbes = maxProbes;
 	so->dimensions = dimensions;
 
 	/* Set support functions */
@@ -274,9 +292,20 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	/* Prep sort */
 	so->sortstate = InitScanSortState(so->tupdesc);
 
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
+	/* Need separate slots for puttuple and gettuple */
+	so->vslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
+	so->mslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
+
+	/*
+	 * Reuse same set of shared buffers for scan
+	 *
+	 * See postgres/src/backend/storage/buffer/README for description
+	 */
+	so->bas = GetAccessStrategy(BAS_BULKREAD);
 
 	so->listQueue = pairingheap_allocate(CompareLists, scan);
+	so->startPages = palloc(maxProbes * sizeof(BlockNumber));
+	so->currentIndex = 0;
 
 	scan->opaque = so;
 
@@ -291,11 +320,9 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 {
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
-	if (!so->first)
-		tuplesort_reset(so->sortstate);
-
 	so->first = true;
 	pairingheap_reset(so->listQueue);
+	so->currentIndex = 0;
 
 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -311,6 +338,8 @@ bool
 ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 {
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
+	ItemPointer heaptid;
+	bool		isnull;
 
 	/*
 	 * Index can be used to scan backward, but Postgres doesn't support
@@ -338,27 +367,25 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		IvfflatBench("GetScanLists", GetScanLists(scan, value));
 		IvfflatBench("GetScanItems", GetScanItems(scan, value));
 		so->first = false;
+		so->value = value;
 
-#if defined(IVFFLAT_MEMORY)
+		/* TODO clean up if we allocated a new value */
-		elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
-#endif
-
-		/* Clean up if we allocated a new value */
-		if (value != scan->orderByData->sk_argument)
-			pfree(DatumGetPointer(value));
 	}
 
-	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
+	while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
 	{
-		ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
+		if (so->currentIndex == so->maxProbes)
+			return false;
 
-		scan->xs_heaptid = *heaptid;
+		IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
-		scan->xs_recheck = false;
-		scan->xs_recheckorderby = false;
-		return true;
 	}
 
-	return false;
+	heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
+
+	scan->xs_heaptid = *heaptid;
+	scan->xs_recheck = false;
+	scan->xs_recheckorderby = false;
+	return true;
 }
 
 /*
@@ -370,7 +397,12 @@ ivfflatendscan(IndexScanDesc scan)
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
 	pairingheap_free(so->listQueue);
+	pfree(so->startPages);
 	tuplesort_end(so->sortstate);
+	FreeAccessStrategy(so->bas);
+	FreeTupleDesc(so->tupdesc);
+
+	/* TODO Free vslot and mslot without freeing TupleDesc */
 
 	pfree(so);
 	scan->opaque = NULL;

test/t/039_ivfflat_streaming.pl

@@ -0,0 +1,31 @@
+use strict;
+use warnings FATAL => 'all';
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $dim = 3;
+my $array_sql = join(",", ('random()') x $dim);
+
+# Initialize node
+my $node = PostgreSQL::Test::Cluster->new('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($dim));");
+$node->safe_psql("postgres",
+	"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
+);
+$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
+
+my $count = $node->safe_psql("postgres", qq(
+	SET enable_seqscan = off;
+	SET ivfflat.probes = 10;
+	SET ivfflat.streaming = on;
+	SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
+));
+is($count, 10);
+
+done_testing();

test/t/040_ivfflat_streaming_recall.pl

@@ -0,0 +1,125 @@
+use strict;
+use warnings FATAL => 'all';
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node;
+my @queries = ();
+my @expected;
+my $limit = 20;
+my @cs = (100, 1000);
+
+sub test_recall
+{
+	my ($c, $probes, $min, $operator) = @_;
+	my $correct = 0;
+	my $total = 0;
+
+	my $explain = $node->safe_psql("postgres", qq(
+		SET enable_seqscan = off;
+		SET ivfflat.probes = $probes;
+		SET ivfflat.streaming = on;
+		EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
+	));
+	like($explain, qr/Index Scan using idx on tst/);
+
+	for my $i (0 .. $#queries)
+	{
+		my $actual = $node->safe_psql("postgres", qq(
+			SET enable_seqscan = off;
+			SET ivfflat.probes = $probes;
+			SET ivfflat.streaming = on;
+			SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
+		));
+		my @actual_ids = split("\n", $actual);
+
+		my @expected_ids = split("\n", $expected[$i]);
+		my %expected_set = map { $_ => 1 } @expected_ids;
+
+		foreach (@actual_ids)
+		{
+			if (exists($expected_set{$_}))
+			{
+				$correct++;
+			}
+		}
+
+		$total += $limit;
+	}
+
+	cmp_ok($correct / $total, ">=", $min, $operator);
+}
+
+# Initialize node
+$node = PostgreSQL::Test::Cluster->new('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[random(), random(), random()] FROM generate_series(1, 100000) i;"
+);
+
+# Generate queries
+for (1 .. 20)
+{
+	my $r1 = rand();
+	my $r2 = rand();
+	my $r3 = rand();
+	push(@queries, "[$r1,$r2,$r3]");
+}
+
+# Check each index type
+my @operators = ("<->", "<=>");
+my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
+
+for my $i (0 .. $#operators)
+{
+	my $operator = $operators[$i];
+	my $opclass = $opclasses[$i];
+
+	$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v $opclass);");
+
+	foreach (@cs)
+	{
+		my $c = $_;
+
+		# Get exact results
+		@expected = ();
+		foreach (@queries)
+		{
+			my $res = $node->safe_psql("postgres", qq(
+				SET enable_indexscan = off;
+				WITH top AS (
+					SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
+				)
+				SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
+			));
+			push(@expected, $res);
+		}
+
+		if ($c == 100)
+		{
+			test_recall($c, 1, 0.58, $operator);
+			test_recall($c, 10, 0.98, $operator);
+		}
+		else
+		{
+			if ($operator eq "<->")
+			{
+				test_recall($c, 1, 0.80, $operator);
+			}
+			else
+			{
+				test_recall($c, 1, 0.88, $operator);
+			}
+		}
+	}
+
+	$node->safe_psql("postgres", "DROP INDEX idx;");
+}
+
+done_testing();