Skip more work

src/hnsw.h

@@ -110,13 +110,14 @@ typedef struct HnswCandidate
 {
 	HnswElement element;
 	float		distance;
+	bool		closer;
 }			HnswCandidate;
 
 typedef struct HnswNeighborArray
 {
 	int			length;
+	bool		closerSet;
 	HnswCandidate *items;
-	HnswElement firstPruned;
 }			HnswNeighborArray;
 
 typedef struct HnswPairingHeapNode

src/hnswutils.c

@@ -139,7 +139,7 @@ HnswInitNeighbors(HnswElement element, int m)
 		a = &element->neighbors[lc];
 		a->length = 0;
 		a->items = palloc(sizeof(HnswCandidate) * lm);
-		a->firstPruned = NULL;
+		a->closerSet = false;
 	}
 }
 
@@ -754,7 +754,8 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
 	List	   *r = NIL;
 	List	   *w = list_copy(c);
 	pairingheap *wd;
-	bool		mustCalculate = e2->neighbors[lc].firstPruned == NULL;
+	bool		mustCalculate = !e2->neighbors[lc].closerSet;
+	bool		foundNew = false;
 
 	if (list_length(w) <= m)
 		return w;
@@ -765,45 +766,39 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
 	{
 		/* Assumes w is already ordered desc */
 		HnswCandidate *e = llast(w);
-		bool		closer;
 
 		w = list_delete_last(w);
 
-		/*
-		 * r and wd will be the same as previous calls until the new
-		 * candidate, so can skip distance calculations for as many candidates
-		 * as there is state for
-		 */
+		/* Use previous state of r and wd to skip work when possible */
 		if (mustCalculate)
-			closer = CheckElementCloser(e, r, lc, procinfo, collation);
-		else if (e->element == e2->neighbors[lc].firstPruned)
+			e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
+		else if (foundNew)
 		{
-			closer = false;
+			/* If new or current candidate is not closer, no change in state */
+			if (newCandidate->closer && e->closer)
+			{
+				/* Only need to compare with new candidate */
+				float		distance = HnswGetDistance(e->element, newCandidate->element, lc, procinfo, collation);
 
-			/*
-			 * Could store multiple pruned and only calculate when exhausted
-			 * (or store full state) at the expense of memory
-			 */
-			mustCalculate = true;
+				e->closer = e->distance < distance;
+
+				if (!e->closer)
+					mustCalculate = true;
+			}
 		}
 		else if (e == newCandidate)
 		{
-			closer = CheckElementCloser(e, r, lc, procinfo, collation);
-			mustCalculate = true;
+			e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
+			foundNew = true;
 		}
-		else
-			closer = true;
 
-		if (closer)
+		if (e->closer)
 			r = lappend(r, e);
 		else
 			pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node));
 	}
 
-	/* Save first pruned */
-	/* OK to leave previous element if empty */
-	if (!pairingheap_is_empty(wd))
-		e2->neighbors[lc].firstPruned = ((HnswPairingHeapNode *) pairingheap_first(wd))->inner->element;
+	e2->neighbors[lc].closerSet = true;
 
 	/* Keep pruned connections */
 	while (!pairingheap_is_empty(wd) && list_length(r) < m)