Added comment [skip ci]

CHANGELOG.md

@@ -1,13 +1,13 @@
 ## 0.5.0 (2023-08-28)
 
 - Added HNSW index type
-- Added support for parallel index builds
+- Added support for parallel index builds for IVFFlat
 - Added `l1_distance` function
 - Added element-wise multiplication for vectors
 - Added `sum` aggregate
 - Improved performance of distance functions
 - Fixed out of range results for cosine distance
-- Fixed results for NULL and NaN distances
+- Fixed results for NULL and NaN distances for IVFFlat
 
 ## 0.4.4 (2023-06-12)
 

README.md

@@ -162,7 +162,7 @@ You can add an index to use approximate nearest neighbor search, which trades so
 Supported index types are:
 
 - [IVFFlat](#ivfflat)
-- [HNSW](#hnsw) - *added in 0.5.0*
+- [HNSW](#hnsw) - added in 0.5.0
 
 ## IVFFlat
 
@@ -282,8 +282,8 @@ SELECT phase, tuples_done, tuples_total FROM pg_stat_progress_create_index;
 The phases are:
 
 1. `initializing`
-2. `performing k-means` (IVFFlat only)
-3. `assigning tuples` (IVFFlat only)
+2. `performing k-means` - IVFFlat only
+3. `assigning tuples` - IVFFlat only
 4. `loading tuples`
 
 Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase

src/hnsw.h

@@ -219,6 +219,9 @@ typedef struct HnswScanOpaqueData
 {
 	bool		first;
 	Buffer		buf;
+	ItemPointerData heaptid;
+	OffsetNumber offno;
+	int			removedCount;
 	List	   *w;
 	MemoryContext tmpCtx;
 

src/hnswscan.c

@@ -58,6 +58,75 @@ GetDimensions(Relation index)
 	return dimensions;
 }
 
+/*
+ * Remove deleted heap TID
+ */
+static void
+RemoveHeapTid(IndexScanDesc scan)
+{
+	HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
+	Relation	index = scan->indexRelation;
+	Buffer		buf = so->buf;
+	Page		page;
+	GenericXLogState *state;
+	ItemId		itemid;
+	HnswElementTuple etup;
+	Size		etupSize;
+	int			idx = -1;
+
+	/* Safety check */
+	if (!BufferIsValid(buf) || !OffsetNumberIsValid(so->offno) || !ItemPointerIsValid(&so->heaptid))
+		return;
+
+	/* Use WAL rather than hint */
+	LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
+	state = GenericXLogStart(index);
+	page = GenericXLogRegisterBuffer(state, buf, 0);
+	itemid = PageGetItemId(page, so->offno);
+	etup = (HnswElementTuple) PageGetItem(page, itemid);
+	etupSize = ItemIdGetLength(itemid);
+
+	Assert(HnswIsElementTuple(etup));
+
+	/* Find index */
+	for (int i = 0; i < HNSW_HEAPTIDS; i++)
+	{
+		if (!ItemPointerIsValid(&etup->heaptids[i]))
+			break;
+
+		if (ItemPointerEquals(&etup->heaptids[i], &so->heaptid))
+		{
+			idx = i;
+			break;
+		}
+	}
+
+	if (idx == -1)
+		GenericXLogAbort(state);
+	else
+	{
+		/* Move pointers forward */
+		for (int i = idx; i < HNSW_HEAPTIDS; i++)
+		{
+			if (i + 1 == HNSW_HEAPTIDS || !ItemPointerIsValid(&etup->heaptids[i + 1]))
+				ItemPointerSetInvalid(&etup->heaptids[i]);
+			else
+				ItemPointerCopy(&etup->heaptids[i + 1], &etup->heaptids[i]);
+		}
+
+		/* Overwrite tuple */
+		if (!PageIndexTupleOverwrite(page, so->offno, (Item) etup, etupSize))
+			elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
+
+		/* Commit */
+		MarkBufferDirty(buf);
+		GenericXLogFinish(state);
+	}
+
+	/* Unlock buffer */
+	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
+}
+
 /*
  * Prepare for an index scan
  */
@@ -71,6 +140,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 
 	so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
 	so->buf = InvalidBuffer;
+	ItemPointerSetInvalid(&so->heaptid);
+	so->offno = InvalidOffsetNumber;
+	so->removedCount = 0;
 	so->first = true;
 	so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
 									   "Hnsw scan temporary context",
@@ -95,6 +167,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
 	HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
 
 	so->first = true;
+	ItemPointerSetInvalid(&so->heaptid);
 	MemoryContextReset(so->tmpCtx);
 
 	if (keys && scan->numberOfKeys > 0)
@@ -158,12 +231,25 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 		so->first = false;
 	}
+	else
+	{
+		/*
+		 * Remove dead tuples. kill_prior_tuple will only be true if not in
+		 * recovery. Limit the number removed per scan for performance.
+		 */
+		if (scan->kill_prior_tuple && so->removedCount < 3)
+		{
+			RemoveHeapTid(scan);
+			so->removedCount++;
+		}
+	}
 
 	while (list_length(so->w) > 0)
 	{
 		HnswCandidate *hc = llast(so->w);
 		ItemPointer tid;
 		BlockNumber indexblkno;
+		OffsetNumber indexoffno;
 
 		/* Move to next element if no valid heap tids */
 		if (list_length(hc->element->heaptids) == 0)
@@ -174,6 +260,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 		tid = llast(hc->element->heaptids);
 		indexblkno = hc->element->blkno;
+		indexoffno = hc->element->offno;
 
 		hc->element->heaptids = list_delete_last(hc->element->heaptids);
 
@@ -185,6 +272,11 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *tid;
 #endif
 
+		/* Keep track of info needed to remove dead tuples */
+		so->heaptid = *tid;
+		so->offno = indexoffno;
+
+		/* Unpin buffer */
 		if (BufferIsValid(so->buf))
 			ReleaseBuffer(so->buf);
 

src/hnswutils.c

@@ -551,6 +551,8 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 	pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
 	pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
 	int			wlen = 0;
+	uint64		dead = 0;
+	uint64		maxAdditional = skipElement == NULL ? ef : PG_UINT64_MAX;
 	HASHCTL		hash_ctl;
 	HTAB	   *v;
 
@@ -579,13 +581,14 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 		pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
 		pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
 
-		/*
-		 * Do not count elements being deleted towards ef when vacuuming. It
-		 * would be ideal to do this for inserts as well, but this could
-		 * affect insert performance.
-		 */
-		if (skipElement == NULL || list_length(hc->element->heaptids) != 0)
-			wlen++;
+		/* Do not count certain number of dead elements towards ef */
+		if (list_length(hc->element->heaptids) == 0)
+		{
+			if ((++dead) <= maxAdditional)
+				continue;
+		}
+
+		wlen++;
 	}
 
 	while (!pairingheap_is_empty(C))
@@ -638,19 +641,18 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 					pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node));
 					pairingheap_add(W, &(CreatePairingHeapNode(ec)->ph_node));
 
-					/*
-					 * Do not count elements being deleted towards ef when
-					 * vacuuming. It would be ideal to do this for inserts as
-					 * well, but this could affect insert performance.
-					 */
-					if (skipElement == NULL || list_length(e->element->heaptids) != 0)
+					/* Do not count certain number of dead elements towards ef */
+					if (list_length(e->element->heaptids) == 0)
 					{
-						wlen++;
-
-						/* No need to decrement wlen */
-						if (wlen > ef)
-							pairingheap_remove_first(W);
+						if ((++dead) <= maxAdditional)
+							continue;
 					}
+
+					wlen++;
+
+					/* No need to decrement wlen */
+					if (wlen > ef)
+						pairingheap_remove_first(W);
 				}
 			}
 		}

src/ivfbuild.c

@@ -10,8 +10,8 @@
 #include "ivfflat.h"
 #include "miscadmin.h"
 #include "storage/bufmgr.h"
-#include "utils/memutils.h"
 #include "tcop/tcopprot.h"
+#include "utils/memutils.h"
 
 #if PG_VERSION_NUM >= 140000
 #include "utils/backend_progress.h"

src/ivfscan.c

@@ -343,6 +343,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *tid;
 #endif
 
+		/* Unpin buffer */
 		if (BufferIsValid(so->buf))
 			ReleaseBuffer(so->buf);
 

test/t/003_ivfflat_build_recall.pl

@@ -94,7 +94,7 @@ for my $i (0 .. $#operators)
 	# Test approximate results
 	if ($operator ne "<#>")
 	{
-		# TODO fix test
+		# TODO Fix test (uniform random vectors all have similar inner product)
 		test_recall(1, 0.71, $operator);
 		test_recall(10, 0.95, $operator);
 	}
@@ -115,7 +115,7 @@ for my $i (0 .. $#operators)
 	# Test approximate results
 	if ($operator ne "<#>")
 	{
-		# TODO fix test
+		# TODO Fix test (uniform random vectors all have similar inner product)
 		test_recall(1, 0.71, $operator);
 		test_recall(10, 0.95, $operator);
 	}

test/t/010_hnsw_wal.pl

@@ -38,8 +38,9 @@ sub test_index_replay
 	);
 
 	# Run test queries and compare their result
-	my $primary_result = $node_primary->safe_psql("postgres", $queries);
+	# Query replica first since index scan on primary can generate WAL removing tuples
 	my $replica_result = $node_replica->safe_psql("postgres", $queries);
+	my $primary_result = $node_primary->safe_psql("postgres", $queries);
 
 	is($primary_result, $replica_result, "$test_name: query result matches");
 	return;

test/t/015_hnsw_duplicates.pl

@@ -23,25 +23,27 @@ sub insert_vectors
 
 sub test_duplicates
 {
+	my ($exp) = @_;
+
 	my $res = $node->safe_psql("postgres", qq(
 		SET enable_seqscan = off;
 		SET hnsw.ef_search = 1;
 		SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t;
 	));
-	is($res, 10);
+	is($res, $exp);
 }
 
 # Test duplicates with build
 insert_vectors();
 $node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
-test_duplicates();
+test_duplicates(10);
 
 # Reset
 $node->safe_psql("postgres", "TRUNCATE tst;");
 
 # Test duplicates with inserts
 insert_vectors();
-test_duplicates();
+test_duplicates(10);
 
 # Test fallback path for inserts
 $node->pgbench(
@@ -55,4 +57,15 @@ $node->pgbench(
 	}
 );
 
+# Reset
+$node->safe_psql("postgres", "TRUNCATE tst;");
+
+# Test deletes with index scan
+$node->safe_psql("postgres", "INSERT INTO tst SELECT '[1,1,1]' FROM generate_series(1, 10) i;");
+$node->safe_psql("postgres", "DELETE FROM tst WHERE ctid IN (SELECT ctid FROM tst ORDER BY random() LIMIT 5);");
+for (1 .. 3)
+{
+	test_duplicates(5);
+}
+
 done_testing();

test/t/016_hnsw_vacuum_recall.pl

@@ -89,7 +89,7 @@ foreach (@queries)
 test_recall(0.20, $limit, "before vacuum");
 test_recall(0.95, 100, "before vacuum");
 
-# TODO test concurrent inserts with vacuum
+# TODO Test concurrent inserts with vacuum
 $node->safe_psql("postgres", "VACUUM tst;");
 
 test_recall(0.95, $limit, "after vacuum");

test/t/017_ivfflat_insert_recall.pl

@@ -0,0 +1,117 @@
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More;
+
+my $node;
+my @queries = ();
+my @expected;
+my $limit = 20;
+
+sub test_recall
+{
+	my ($probes, $min, $operator) = @_;
+	my $correct = 0;
+	my $total = 0;
+
+	my $explain = $node->safe_psql("postgres", qq(
+		SET enable_seqscan = off;
+		SET ivfflat.probes = $probes;
+		EXPLAIN ANALYZE SELECT i FROM tst 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;
+			SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
+		));
+		my @actual_ids = split("\n", $actual);
+		my %actual_set = map { $_ => 1 } @actual_ids;
+
+		my @expected_ids = split("\n", $expected[$i]);
+
+		foreach (@expected_ids)
+		{
+			if (exists($actual_set{$_}))
+			{
+				$correct++;
+			}
+			$total++;
+		}
+	}
+
+	cmp_ok($correct / $total, ">=", $min, $operator);
+}
+
+# Initialize node
+$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 serial, v vector(3));");
+
+# 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_ip_ops", "vector_cosine_ops");
+
+for my $i (0 .. $#operators)
+{
+	my $operator = $operators[$i];
+	my $opclass = $opclasses[$i];
+
+	# Add index
+	$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v $opclass);");
+
+	# Use concurrent inserts
+	$node->pgbench(
+		"--no-vacuum --client=10 --transactions=1000",
+		0,
+		[qr{actually processed}],
+		[qr{^$}],
+		"concurrent INSERTs",
+		{
+			"017_ivfflat_insert_recall_$opclass" => "INSERT INTO tst (v) SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 10) i;"
+		}
+	);
+
+	# Get exact results
+	@expected = ();
+	foreach (@queries)
+	{
+		my $res = $node->safe_psql("postgres", qq(
+			SET enable_indexscan = off;
+			SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
+		));
+		push(@expected, $res);
+	}
+
+	# Test approximate results
+	if ($operator ne "<#>")
+	{
+		# TODO Fix test (uniform random vectors all have similar inner product)
+		test_recall(1, 0.71, $operator);
+		test_recall(10, 0.95, $operator);
+	}
+	# Account for equal distances
+	test_recall(100, 0.9925, $operator);
+
+	$node->safe_psql("postgres", "DROP INDEX idx;");
+	$node->safe_psql("postgres", "TRUNCATE tst;");
+}
+
+done_testing();