Merge branch 'master' into hqann2

Switched to macos-13 on CI [skip ci]
Use get_float8_infinity instead of INFINITY
2026-07-22 03:57:34 +08:00 · 2024-10-10 19:06:26 -07:00 · 2024-10-10 18:33:35 -07:00 · 2024-10-10 18:16:39 -07:00 · 2024-10-10 18:14:39 -07:00 · 2024-10-10 18:12:27 -07:00
24 changed files with 865 additions and 276 deletions
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -49,7 +49,7 @@ jobs:
          - postgres: 16
            os: macos-14
          - postgres: 14
-            os: macos-12
+            os: macos-13
    steps:
      - uses: actions/checkout@v4
      - uses: ankane/setup-postgres@v1
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## 0.8.0 (unreleased)

+- Added support for inline filtering with HNSW
 - Added support for iterative index scans
 - Added casts for arrays to `sparsevec`
 - Improved cost estimation
--- a/README.md
+++ b/README.md
@@ -439,6 +439,12 @@ Create an index on one [or more](https://www.postgresql.org/docs/current/indexes
 CREATE INDEX ON items (category_id);
 ```

+Or a composite HNSW index for approximate search (added in 0.8.0)
+
+```sql
+CREATE INDEX ON items USING hnsw (embedding vector_l2_ops, category_id);
+```
+
 Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search

 ```sql
@@ -451,77 +457,6 @@ Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html
 CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(category_id);
 ```

-## Iterative Search [unreleased]
-
-*Added in 0.8.0*
-
-With approximate indexes, you can end up with less results than expected due to filtering conditions in the query.
-
-Starting with 0.8.0, you can enable iterative search. If too few results from the initial index scan match the query filters, it will resume scanning until enough results are found. This can significantly improve recall (at the cost of speed).
-
-```tsql
-SET hnsw.streaming = on;
-- or
-SET ivfflat.streaming = on;
-```
-
-However, there are some important caveats.
-
-### Iterative Caveats
-
-With iterative search, it’s possible for rows to be slightly out of order by distance. For strict ordering, use:
-
-```sql
-WITH approx_order AS MATERIALIZED (
-    SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
-) SELECT * FROM approx_order ORDER BY distance;
-```
-
-For distance filters, use a CTE and place the filter outside it.
-
-```sql
-WITH approx_order AS MATERIALIZED (
-    SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
-) SELECT * FROM approx_order WHERE distance < 0.1 ORDER BY distance;
-```
-
-### Iterative Options
-
-Since scanning a large portion of the index is expensive, there are options to control when the scan ends.
-
-#### HNSW
-
-Specify the max number of additional tuples visited
-
-```sql
-SET hnsw.ef_stream = 10000;
-```
-
-The scan will also end if reaches `work_mem`. You can see when this happens by enabling debug messages.
-
-```sql
-SET client_min_messages = debug1;
-```
-
-```text
-DEBUG:  hnsw index scan exceeded work_mem after 10000 tuples
-HINT:  Increase work_mem to scan more tuples.
-```
-
-If the server has enough memory, you can adjust this with:
-
-```sql
-SET work_mem = '8MB';
-```
-
-#### IVFFlat
-
-Specify the max number of probes
-
-```sql
-SET ivfflat.max_probes = 100;
-```
-
 ## Half-Precision Vectors

 *Added in 0.7.0*
@@ -1260,6 +1195,7 @@ Thanks to:
 - [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
 - [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
 - [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf)
+- [HQANN: Efficient and Robust Similarity Search for Hybrid Queries with Structured and Unstructured Constraints](https://arxiv.org/pdf/2207.07940.pdf)

 ## History

--- a/sql/vector--0.7.4--0.8.0.sql
+++ b/sql/vector--0.7.4--0.8.0.sql
@@ -24,3 +24,11 @@ CREATE CAST (double precision[] AS sparsevec)

 CREATE CAST (numeric[] AS sparsevec)
 	WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT;
+
+CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
+	AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
+CREATE OPERATOR CLASS vector_integer_ops
+	DEFAULT FOR TYPE integer USING hnsw AS
+	OPERATOR 2 = (integer, integer),
+	FUNCTION 4 hnsw_attribute_distance(integer, integer);
--- a/sql/vector.sql
+++ b/sql/vector.sql
@@ -916,3 +916,13 @@ CREATE OPERATOR CLASS sparsevec_l1_ops
 	OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
 	FUNCTION 1 l1_distance(sparsevec, sparsevec),
 	FUNCTION 3 hnsw_sparsevec_support(internal);
+
+-- hnsw attributes
+
+CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
+	AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
+CREATE OPERATOR CLASS vector_integer_ops
+	DEFAULT FOR TYPE integer USING hnsw AS
+	OPERATOR 2 = (integer, integer),
+	FUNCTION 4 hnsw_attribute_distance(integer, integer);
--- a/src/hnsw.c
+++ b/src/hnsw.c
@@ -20,15 +20,13 @@

 static const struct config_enum_entry hnsw_iterative_search_options[] = {
 	{"off", HNSW_ITERATIVE_SEARCH_OFF, false},
-	{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
-	{"relaxed", HNSW_ITERATIVE_SEARCH_RELAXED, false},
-	/* TODO Change to strict before merging */
 	{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
+	{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
 	{NULL, 0, false}
 };

 int			hnsw_ef_search;
-int			hnsw_max_iterative_tuples;
+int			hnsw_iterative_search_max_tuples;
 int			hnsw_iterative_search;
 int			hnsw_lock_tranche_id;
 static relopt_kind hnsw_relopt_kind;
@@ -80,16 +78,14 @@ 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);

-	/* TODO Change name */
-	DefineCustomEnumVariable("hnsw.streaming", "Iterative search mode",
+	DefineCustomEnumVariable("hnsw.iterative_search", "Sets iterative search",
 							 NULL, &hnsw_iterative_search,
 							 HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);

-	/* TODO Change name */
 	/* TODO Ensure ivfflat.max_probes uses same value for "all" */
-	DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search",
-							"-1 means all", &hnsw_max_iterative_tuples,
-							HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, NULL, NULL, NULL);
+	DefineCustomIntVariable("hnsw.iterative_search_max_tuples", "Sets the max number of candidates to visit for iterative search",
+							"-1 means all", &hnsw_iterative_search_max_tuples,
+							-1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);

 	MarkGUCPrefixReserved("hnsw");
 }
@@ -249,13 +245,13 @@ hnswhandler(PG_FUNCTION_ARGS)
 	IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);

 	amroutine->amstrategies = 0;
-	amroutine->amsupport = 3;
+	amroutine->amsupport = 4;
 	amroutine->amoptsprocnum = 0;
 	amroutine->amcanorder = false;
 	amroutine->amcanorderbyop = true;
 	amroutine->amcanbackward = false;	/* can change direction mid-scan */
 	amroutine->amcanunique = false;
-	amroutine->amcanmulticol = false;
+	amroutine->amcanmulticol = true;
 	amroutine->amoptionalkey = true;
 	amroutine->amsearcharray = false;
 	amroutine->amsearchnulls = false;
@@ -307,3 +303,17 @@ hnswhandler(PG_FUNCTION_ARGS)

 	PG_RETURN_POINTER(amroutine);
 }
+
+/*
+ * Get the distance between two int4 attributes
+ */
+PGDLLEXPORT PG_FUNCTION_INFO_V1(hnsw_int4_attribute_distance);
+Datum
+hnsw_int4_attribute_distance(PG_FUNCTION_ARGS)
+{
+	int32		a = PG_GETARG_INT32(0);
+	int32		b = PG_GETARG_INT32(1);
+	double		distance = ((double) a) - ((double) b);
+
+	PG_RETURN_FLOAT8(distance);
+}
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -19,6 +19,7 @@
 #define HNSW_DISTANCE_PROC 1
 #define HNSW_NORM_PROC 2
 #define HNSW_TYPE_INFO_PROC 3
+#define HNSW_ATTRIBUTE_DISTANCE_PROC 4

 #define HNSW_VERSION	1
 #define HNSW_MAGIC_NUMBER 0xA953A953
@@ -42,9 +43,6 @@
 #define HNSW_DEFAULT_EF_SEARCH	40
 #define HNSW_MIN_EF_SEARCH		1
 #define HNSW_MAX_EF_SEARCH		1000
-#define HNSW_DEFAULT_EF_STREAM	-1
-#define HNSW_MIN_EF_STREAM		-1
-#define HNSW_MAX_EF_STREAM		INT_MAX

 /* Tuple types */
 #define HNSW_ELEMENT_TUPLE_TYPE  1
@@ -110,17 +108,19 @@
 #define HnswPtrPointer(hp) (hp).ptr
 #define HnswPtrOffset(hp) relptr_offset((hp).relptr)

+#define HnswUseIndexTuple(index) (IndexRelationGetNumberOfAttributes(index) > 1)
+
 /* Variables */
 extern int	hnsw_ef_search;
-extern int	hnsw_max_iterative_tuples;
 extern int	hnsw_iterative_search;
+extern int	hnsw_iterative_search_max_tuples;
 extern int	hnsw_lock_tranche_id;

 typedef enum HnswIterativeSearchType
 {
 	HNSW_ITERATIVE_SEARCH_OFF,
-	HNSW_ITERATIVE_SEARCH_STRICT,
-	HNSW_ITERATIVE_SEARCH_RELAXED
+	HNSW_ITERATIVE_SEARCH_RELAXED,
+	HNSW_ITERATIVE_SEARCH_STRICT
 }			HnswIterativeSearchType;

 typedef struct HnswElementData HnswElementData;
@@ -136,6 +136,7 @@ HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
 HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
 HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
 HnswPtrDeclare(char, DatumRelptr, DatumPtr);
+HnswPtrDeclare(IndexTupleData, IndexTupleRelptr, IndexTuplePtr);

 struct HnswElementData
 {
@@ -152,6 +153,7 @@ struct HnswElementData
 	OffsetNumber neighborOffno;
 	BlockNumber neighborPage;
 	DatumPtr	value;
+	IndexTuplePtr itup;
 	LWLock		lock;
 };

@@ -177,6 +179,7 @@ typedef struct HnswSearchCandidate
 	pairingheap_node w_node;
 	HnswElementPtr element;
 	double		distance;
+	bool		matches;
 }			HnswSearchCandidate;

 /* HNSW index options */
@@ -255,14 +258,16 @@ typedef struct HnswTypeInfo

 typedef struct HnswSupport
 {
-	FmgrInfo   *procinfo;
+	FmgrInfo   *procinfo[2];
 	FmgrInfo   *normprocinfo;
-	Oid			collation;
+	Oid		   *collation;
 }			HnswSupport;

 typedef struct HnswQuery
 {
 	Datum		value;
+	IndexTuple	itup;
+	ScanKeyData *keyData;
 }			HnswQuery;

 typedef struct HnswBuildState
@@ -291,6 +296,8 @@ typedef struct HnswBuildState
 	HnswGraph  *graph;
 	double		ml;
 	int			maxLevel;
+	bool		useIndexTuple;
+	TupleDesc	tupdesc;

 	/* Memory */
 	MemoryContext graphCtx;
@@ -418,30 +425,32 @@ bool		HnswCheckNorm(HnswSupport * support, Datum value);
 Buffer		HnswNewBuffer(Relation index, ForkNumber forkNum);
 void		HnswInitPage(Buffer buf, Page page);
 void		HnswInit(void);
-List	   *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
+List	   *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
 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, HnswSupport * support, int m, int efConstruction, bool existing);
-HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec);
+void		HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory);
+HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec, bool inMemory);
 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);
 HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
 void		HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
-bool		HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building);
+bool		HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc);
 void		HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building);
-void		HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
-void		HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
-bool		HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support);
-void		HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
+void		HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index);
+void		HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
+void		HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple);
 void		HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);
+bool		HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc);
 bool		HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
 void		HnswInitLockTranche(void);
 const		HnswTypeInfo *HnswGetTypeInfo(Relation index);
 PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
+Size		HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple);
+bool		HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc);

 /* Index access methods */
 IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);
--- a/src/hnswbuild.c
+++ b/src/hnswbuild.c
@@ -148,6 +148,7 @@ CreateGraphPages(HnswBuildState * buildstate)
 	Page		page;
 	HnswElementPtr iter = buildstate->graph->head;
 	char	   *base = buildstate->hnswarea;
+	bool		useIndexTuple = buildstate->useIndexTuple;

 	/* Calculate sizes */
 	maxSize = HNSW_MAX_SIZE;
@@ -167,7 +168,6 @@ CreateGraphPages(HnswBuildState * buildstate)
 		Size		etupSize;
 		Size		ntupSize;
 		Size		combinedSize;
-		Pointer		valuePtr = HnswPtrAccess(base, element->value);

 		/* Update iterator */
 		iter = element->next;
@@ -176,7 +176,7 @@ CreateGraphPages(HnswBuildState * buildstate)
 		MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE);

 		/* Calculate sizes */
-		etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
+		etupSize = HnswGetElementTupleSize(base, element, useIndexTuple);
 		ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
 		combinedSize = etupSize + ntupSize + sizeof(ItemIdData);

@@ -186,7 +186,7 @@ CreateGraphPages(HnswBuildState * buildstate)
 					(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
 					 errmsg("index tuple too large")));

-		HnswSetElementTuple(base, etup, element);
+		HnswSetElementTuple(base, etup, element, useIndexTuple);

 		/* Keep element and neighbors on the same page if possible */
 		if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -327,19 +327,18 @@ AddDuplicateInMemory(HnswElement element, HnswElement dup)
 * Find duplicate element
 */
 static bool
-FindDuplicateInMemory(char *base, HnswElement element)
+FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc)
 {
 	HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
-	Datum		value = HnswGetValue(base, element);
+	IndexTuple	itup = HnswPtrAccess(base, element->itup);

 	for (int i = 0; i < neighbors->length; i++)
 	{
 		HnswCandidate *neighbor = &neighbors->items[i];
 		HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
-		Datum		neighborValue = HnswGetValue(base, neighborElement);

 		/* Exit early since ordered by distance */
-		if (!datumIsEqual(value, neighborValue, false, -1))
+		if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
 			return false;

 		/* Check for space */
@@ -366,7 +365,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
 * Update neighbors
 */
 static void
-UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
+UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswElement e, int m)
 {
 	for (int lc = e->level; lc >= 0; lc--)
 	{
@@ -388,7 +387,7 @@ UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
 			Assert(neighborElement);

 			LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
-			HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, NULL, support);
+			HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, support);
 			LWLockRelease(&neighborElement->lock);
 		}
 	}
@@ -404,14 +403,14 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
 	char	   *base = buildstate->hnswarea;

 	/* Look for duplicate */
-	if (FindDuplicateInMemory(base, element))
+	if (FindDuplicateInMemory(base, element, buildstate->useIndexTuple, buildstate->tupdesc))
 		return;

 	/* Add element */
 	AddElementInMemory(base, graph, element);

 	/* Update neighbors */
-	UpdateNeighborsInMemory(base, support, element, m);
+	UpdateNeighborsInMemory(base, buildstate->index, support, element, m);

 	/* Update entry point if needed (already have lock) */
 	if (entryPoint == NULL || element->level > entryPoint->level)
@@ -424,6 +423,7 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
 static void
 InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
 {
+	Relation	index = buildstate->index;
 	HnswGraph  *graph = buildstate->graph;
 	HnswSupport *support = &buildstate->support;
 	HnswElement entryPoint;
@@ -457,7 +457,7 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
 	}

 	/* Find neighbors for element */
-	HnswFindElementNeighbors(base, element, entryPoint, NULL, support, m, efConstruction, false);
+	HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, true);

 	/* Update graph in memory */
 	UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate);
@@ -476,18 +476,20 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
 	HnswElement element;
 	HnswAllocator *allocator = &buildstate->allocator;
 	HnswSupport *support = &buildstate->support;
-	Size		valueSize;
-	Pointer		valuePtr;
 	LWLock	   *flushLock = &graph->flushLock;
 	char	   *base = buildstate->hnswarea;
-	Datum		value;
+	TupleDesc	tupdesc = buildstate->tupdesc;
+	IndexTuple	itup;
+	Size		itupSize;
+	IndexTuple	itupShared;
+	bool		unused;

-	/* Form index value */
-	if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, support))
+	/* Form index tuple */
+	if (!HnswFormIndexTuple(&itup, values, isnull, buildstate->typeInfo, support, tupdesc))
 		return false;

-	/* Get datum size */
-	valueSize = VARSIZE_ANY(DatumGetPointer(value));
+	/* Get tuple size */
+	itupSize = IndexTupleSize(itup);

 	/* Ensure graph not flushed when inserting */
 	LWLockAcquire(flushLock, LW_SHARED);
@@ -497,7 +499,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
 	{
 		LWLockRelease(flushLock);

-		return HnswInsertTupleOnDisk(index, support, value, heaptid, true);
+		return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
 	}

 	/*
@@ -529,12 +531,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn

 		LWLockRelease(flushLock);

-		return HnswInsertTupleOnDisk(index, support, value, heaptid, true);
+		return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
 	}

 	/* Ok, we can proceed to allocate the element */
 	element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
-	valuePtr = HnswAlloc(allocator, valueSize);
+	itupShared = HnswAlloc(allocator, itupSize);

 	/*
 	 * We have now allocated the space needed for the element, so we don't
@@ -543,9 +545,10 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
 	 */
 	LWLockRelease(&graph->allocatorLock);

-	/* Copy the datum */
-	memcpy(valuePtr, DatumGetPointer(value), valueSize);
-	HnswPtrStore(base, element->value, valuePtr);
+	/* Copy the tuple */
+	memcpy(itupShared, itup, itupSize);
+	HnswPtrStore(base, element->itup, itupShared);
+	HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupShared, 1, tupdesc, &unused)));

 	/* Create a lock for the element */
 	LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
@@ -672,6 +675,19 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("type not supported for hnsw index")));

+	/* TODO See if needed */
+	if (IndexRelationGetNumberOfKeyAttributes(index) > 2)
+		elog(ERROR, "index cannot have more than two columns");
+
+	if (!OidIsValid(index_getprocid(index, 1, HNSW_DISTANCE_PROC)))
+		elog(ERROR, "first column must be a vector");
+
+	for (int i = 1; i < IndexRelationGetNumberOfKeyAttributes(index); i++)
+	{
+		if (!OidIsValid(index_getprocid(index, i + 1, HNSW_ATTRIBUTE_DISTANCE_PROC)))
+			elog(ERROR, "column %d cannot be a vector", i + 1);
+	}
+
 	/* Require column to have dimensions to be indexed */
 	if (buildstate->dimensions < 0)
 		ereport(ERROR,
@@ -698,6 +714,8 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 	buildstate->graph = &buildstate->graphData;
 	buildstate->ml = HnswGetMl(buildstate->m);
 	buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
+	buildstate->useIndexTuple = HnswUseIndexTuple(index);
+	buildstate->tupdesc = RelationGetDescr(index);

 	buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
 												   "Hnsw build graph context",
--- a/src/hnswinsert.c
+++ b/src/hnswinsert.c
@@ -156,9 +156,10 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 	BlockNumber newInsertPage = InvalidBlockNumber;
 	uint8		tupleVersion;
 	char	   *base = NULL;
+	bool		useIndexTuple = HnswUseIndexTuple(index);

 	/* Calculate sizes */
-	etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value)));
+	etupSize = HnswGetElementTupleSize(base, e, useIndexTuple);
 	ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
 	combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
 	maxSize = HNSW_MAX_SIZE;
@@ -166,7 +167,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B

 	/* Prepare element tuple */
 	etup = palloc0(etupSize);
-	HnswSetElementTuple(base, etup, e);
+	HnswSetElementTuple(base, etup, e, useIndexTuple);

 	/* Prepare neighbor tuple */
 	ntup = palloc0(ntupSize);
@@ -383,8 +384,9 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, HnswQuery * q, int *idx, Re
 		HnswCandidate *hc = &neighbors->items[i];
 		HnswElement element = HnswPtrAccess(base, hc->element);
 		double		distance;
+		bool		matches;

-		HnswLoadElement(element, &distance, q, index, support, true, NULL);
+		HnswLoadElement(element, &distance, &matches, q, index, support, true, NULL);
 		hc->distance = distance;

 		/* Prune element if being deleted */
@@ -428,6 +430,8 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
 		HnswQuery	q;

 		q.value = HnswGetValue(base, element);
+		q.itup = HnswPtrAccess(base, element->itup);
+		q.keyData = NULL;

 		LoadElementsForInsert(neighbors, &q, &idx, index, support);

@@ -633,21 +637,30 @@ AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool bu
 * Find duplicate element
 */
 static bool
-FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
+FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDesc tupdesc)
 {
 	char	   *base = NULL;
 	HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
 	Datum		value = HnswGetValue(base, element);
+	IndexTuple	itup = HnswPtrAccess(base, element->itup);

 	for (int i = 0; i < neighbors->length; i++)
 	{
 		HnswCandidate *neighbor = &neighbors->items[i];
 		HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
-		Datum		neighborValue = HnswGetValue(base, neighborElement);

-		/* Exit early since ordered by distance */
-		if (!datumIsEqual(value, neighborValue, false, -1))
-			return false;
+		if (HnswUseIndexTuple(index))
+		{
+			/* Exit early since ordered by distance */
+			if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
+				return false;
+		}
+		else
+		{
+			/* Exit early since ordered by distance */
+			if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
+				return false;
+		}

 		if (AddDuplicateOnDisk(index, element, neighborElement, building))
 			return true;
@@ -660,12 +673,12 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
 * Update graph on disk
 */
 static void
-UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
+UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building, TupleDesc tupdesc)
 {
 	BlockNumber newInsertPage = InvalidBlockNumber;

 	/* Look for duplicate */
-	if (FindDuplicateOnDisk(index, element, building))
+	if (FindDuplicateOnDisk(index, element, building, tupdesc))
 		return;

 	/* Add element */
@@ -687,7 +700,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
 * Insert a tuple into the index
 */
 bool
-HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building)
+HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc)
 {
 	HnswElement entryPoint;
 	HnswElement element;
@@ -695,6 +708,7 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
 	int			efConstruction = HnswGetEfConstruction(index);
 	LOCKMODE	lockmode = ShareLock;
 	char	   *base = NULL;
+	bool		unused;

 	/*
 	 * Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -708,7 +722,8 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo

 	/* Create an element */
 	element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
-	HnswPtrStore(base, element->value, DatumGetPointer(value));
+	HnswPtrStore(base, element->itup, itup);
+	HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));

 	/* Prevent concurrent inserts when likely updating entry point */
 	if (entryPoint == NULL || element->level > entryPoint->level)
@@ -725,10 +740,10 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
 	}

 	/* Find neighbors for element */
-	HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false);
+	HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, false);

 	/* Update graph on disk */
-	UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building);
+	UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building, tupdesc);

 	/* Release lock */
 	UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -742,17 +757,18 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
 static void
 HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
 {
-	Datum		value;
+	IndexTuple	itup;
 	const		HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
+	TupleDesc	tupdesc = RelationGetDescr(index);
 	HnswSupport support;

 	HnswInitSupport(&support, index);

-	/* Form index value */
-	if (!HnswFormIndexValue(&value, values, isnull, typeInfo, &support))
+	/* Form index tuple */
+	if (!HnswFormIndexTuple(&itup, values, isnull, typeInfo, &support, tupdesc))
 		return;

-	HnswInsertTupleOnDisk(index, &support, value, heaptid, false);
+	HnswInsertTupleOnDisk(index, &support, itup, heaptid, false, tupdesc);
 }

 /*
--- a/src/hnswscan.c
+++ b/src/hnswscan.c
@@ -22,26 +22,30 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	int			m;
 	HnswElement entryPoint;
 	char	   *base = NULL;
+	bool		inMemory = false;
 	HnswQuery  *q = &so->q;

+	q->value = value;
+	q->itup = NULL;
+	q->keyData = scan->keyData;
+
 	/* Get m and entry point */
 	HnswGetMetaPageInfo(index, &m, &entryPoint);

-	q->value = value;
 	so->m = m;

 	if (entryPoint == NULL)
 		return NIL;

-	ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false));
+	ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false, inMemory));

 	for (int lc = entryPoint->level; lc >= 1; lc--)
 	{
-		w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL);
+		w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, inMemory, NULL, NULL, true, NULL);
 		ep = w;
 	}

-	return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
+	return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, inMemory, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
 }

 /*
@@ -72,7 +76,7 @@ ResumeScanItems(IndexScanDesc scan)
 		ep = lappend(ep, sc);
 	}

-	return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
+	return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, false, &so->v, &so->discarded, false, &so->tuples);
 }

 /*
@@ -96,7 +100,7 @@ GetScanValue(IndexScanDesc scan)

 		/* Normalize if needed */
 		if (so->support.normprocinfo != NULL)
-			value = HnswNormValue(so->typeInfo, so->support.collation, value);
+			value = HnswNormValue(so->typeInfo, so->support.collation[0], value);
 	}

 	return value;
@@ -225,7 +229,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 				break;

 			/* Reached max number of additional tuples */
-			if (hnsw_max_iterative_tuples != -1 && so->tuples >= hnsw_ef_search + hnsw_max_iterative_tuples)
+			if (hnsw_iterative_search_max_tuples != -1 && so->tuples >= hnsw_iterative_search_max_tuples)
 			{
 				if (pairingheap_is_empty(so->discarded))
 					break;
@@ -278,7 +282,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		element = HnswPtrAccess(base, sc->element);

 		/* Move to next element if no valid heap TIDs */
-		if (element->heaptidsLength == 0)
+		if (!sc->matches || element->heaptidsLength == 0)
 		{
 			so->w = list_delete_last(so->w);

--- a/src/hnswutils.c
+++ b/src/hnswutils.c
@@ -146,11 +146,39 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
 void
 HnswInitSupport(HnswSupport * support, Relation index)
 {
-	support->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
-	support->collation = index->rd_indcollation[0];
+	support->procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
+
+	if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
+		support->procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC);
+
+	support->collation = index->rd_indcollation;
 	support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
 }

+/*
+ * Get element tuple size
+ */
+Size
+HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple)
+{
+	Size		size;
+
+	if (useIndexTuple)
+	{
+		IndexTuple	itup = HnswPtrAccess(base, element->itup);
+
+		size = IndexTupleSize(itup);
+	}
+	else
+	{
+		Pointer		valuePtr = HnswPtrAccess(base, element->value);
+
+		size = VARSIZE_ANY(valuePtr);
+	}
+
+	return HNSW_ELEMENT_TUPLE_SIZE(size);
+}
+
 /*
 * Normalize value
 */
@@ -166,7 +194,38 @@ HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
 bool
 HnswCheckNorm(HnswSupport * support, Datum value)
 {
-	return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation, value)) > 0;
+	return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation[0], value)) > 0;
+}
+
+/*
+ * Check if index tuples are equal
+ */
+bool
+HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc)
+{
+	for (int i = 0; i < tupdesc->natts; i++)
+	{
+		bool		nullA;
+		bool		nullB;
+
+		Datum		datumA = index_getattr(a, i + 1, tupdesc, &nullA);
+		Datum		datumB = index_getattr(b, i + 1, tupdesc, &nullB);
+
+		if (nullA || nullB)
+		{
+			if (nullA != nullB)
+				return false;
+		}
+		else
+		{
+			Form_pg_attribute att = TupleDescAttr(tupdesc, i);
+
+			if (!datumIsEqual(datumA, datumB, att->attbyval, att->attlen))
+				return false;
+		}
+	}
+
+	return true;
 }

 /*
@@ -257,6 +316,7 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
 	HnswInitNeighbors(base, element, m, allocator);

 	HnswPtrStore(base, element->value, (Pointer) NULL);
+	HnswPtrStore(base, element->itup, (IndexTuple) NULL);

 	return element;
 }
@@ -283,6 +343,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
 	element->offno = offno;
 	HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
 	HnswPtrStore(base, element->value, (Pointer) NULL);
+	HnswPtrStore(base, element->itup, (IndexTuple) NULL);
 	return element;
 }

@@ -395,11 +456,13 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
 }

 /*
- * Form index value
+ * Form index tuple
 */
 bool
-HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support)
+HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc)
 {
+	Datum		newValues[2];
+
 	/* Detoast once for all calls */
 	Datum		value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));

@@ -413,10 +476,14 @@ HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo *
 		if (!HnswCheckNorm(support, value))
 			return false;

-		value = HnswNormValue(typeInfo, support->collation, value);
+		value = HnswNormValue(typeInfo, support->collation[0], value);
 	}

-	*out = value;
+	newValues[0] = value;
+	for (int i = 1; i < tupdesc->natts; i++)
+		newValues[i] = values[i];
+
+	*out = index_form_tuple(tupdesc, newValues, isnull);

 	return true;
 }
@@ -425,10 +492,8 @@ HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo *
 * Set element tuple, except for neighbor info
 */
 void
-HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
+HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple)
 {
-	Pointer		valuePtr = HnswPtrAccess(base, element->value);
-
 	etup->type = HNSW_ELEMENT_TUPLE_TYPE;
 	etup->level = element->level;
 	etup->deleted = 0;
@@ -440,7 +505,19 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
 		else
 			ItemPointerSetInvalid(&etup->heaptids[i]);
 	}
-	memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
+
+	if (useIndexTuple)
+	{
+		IndexTuple	itup = HnswPtrAccess(base, element->itup);
+
+		memcpy(&etup->data, itup, IndexTupleSize(itup));
+	}
+	else
+	{
+		Pointer		valuePtr = HnswPtrAccess(base, element->value);
+
+		memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
+	}
 }

 /*
@@ -482,7 +559,7 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
 * Load an element from a tuple
 */
 void
-HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec)
+HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index)
 {
 	element->level = etup->level;
 	element->deleted = etup->deleted;
@@ -506,26 +583,128 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
 	if (loadVec)
 	{
 		char	   *base = NULL;
-		Datum		value = datumCopy(PointerGetDatum(&etup->data), false, -1);

-		HnswPtrStore(base, element->value, DatumGetPointer(value));
+		if (HnswUseIndexTuple(index))
+		{
+			IndexTuple	itup = CopyIndexTuple((IndexTuple) &etup->data);
+			TupleDesc	tupdesc = RelationGetDescr(index);
+			bool		unused;
+
+			HnswPtrStore(base, element->itup, itup);
+			HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
+		}
+		else
+		{
+			Datum		value = datumCopy(PointerGetDatum(&etup->data), false, -1);
+
+			HnswPtrStore(base, element->value, DatumGetPointer(value));
+		}
 	}
 }

+/*
+ * Get the attribute distance
+ */
+static inline double
+AttributeDistance(double e)
+{
+	/* TODO Better bias */
+	/* must be >> max(w * g) + 1 / log10(2) */
+	double		bias = 4.32;
+
+	return e > 0 ? bias - 1.0 / log10(e + 1) : 0;
+}
+
 /*
 * Calculate the distance between values
 */
-static inline double
-HnswGetDistance(Datum a, Datum b, HnswSupport * support)
+static double
+HnswGetDistance(IndexTuple itup, Datum vec, HnswQuery * q, Relation index, HnswSupport * support, bool *matches)
 {
-	return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, a, b));
+	double		g;
+
+	if (DatumGetPointer(q->value) == NULL)
+		g = 0;
+	else
+		g = DatumGetFloat8(FunctionCall2Coll(support->procinfo[0], support->collation[0], q->value, vec));
+
+	Assert(PointerIsValid(matches));
+	*matches = true;
+
+	if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
+	{
+		double		w = 0.25;
+		double		e = 0.0;
+		TupleDesc	tupdesc = RelationGetDescr(index);
+
+		if (q->keyData)
+		{
+			/* TODO need to pass length of key data */
+			int			keyCount = 1;
+
+			for (int i = 0; i < keyCount; i++)
+			{
+				ScanKey		key = &q->keyData[i];
+				bool		isnull;
+				Datum		value = index_getattr(itup, key->sk_attno, tupdesc, &isnull);
+				bool		attnull = key->sk_flags & SK_ISNULL;
+
+				if (isnull || attnull)
+				{
+					if (isnull != attnull)
+					{
+						e += 1000;
+						*matches = false;
+					}
+				}
+				else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument)))
+				{
+					double		ei = fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[key->sk_attno - 1], support->collation[key->sk_attno - 1], value, key->sk_argument)));
+
+					if (ei > 0)
+						e += ei;
+					else
+						/* Distance is zero for inequality */
+						e += 1000;
+
+					*matches = false;
+				}
+			}
+
+			return w * g + AttributeDistance(e);
+		}
+		else if (q->itup)
+		{
+			int			keyCount = IndexRelationGetNumberOfKeyAttributes(index) - 1;
+
+			for (int i = 0; i < keyCount; i++)
+			{
+				bool		isnull;
+				bool		attnull;
+				Datum		value = index_getattr(itup, i + 2, tupdesc, &isnull);
+				Datum		value2 = index_getattr(q->itup, i + 2, tupdesc, &attnull);
+
+				if (isnull || attnull)
+				{
+					if (isnull != attnull)
+						e += 1000;
+				}
+				else
+					e += fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[i + 1], support->collation[i + 1], value, value2)));
+			}
+
+			return w * g + AttributeDistance(e);
+		}
+	}
+
+	return g;
 }

 /*
 * Load an element and optionally get its distance from q
 */
 static void
-HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
+HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
 {
 	Buffer		buf;
 	Page		page;
@@ -543,10 +722,23 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
 	/* Calculate distance */
 	if (distance != NULL)
 	{
-		if (DatumGetPointer(q->value) == NULL)
-			*distance = 0;
+		IndexTuple	itup = NULL;
+		Datum		value;
+
+		if (HnswUseIndexTuple(index))
+		{
+			TupleDesc	tupdesc = RelationGetDescr(index);
+			bool		unused;
+
+			itup = (IndexTuple) &etup->data;
+			value = index_getattr(itup, 1, tupdesc, &unused);
+		}
 		else
-			*distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support);
+		{
+			value = PointerGetDatum(&etup->data);
+		}
+
+		*distance = HnswGetDistance(itup, value, q, index, support, matches);
 	}

 	/* Load element */
@@ -555,7 +747,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
 		if (*element == NULL)
 			*element = HnswInitElementFromBlock(blkno, offno);

-		HnswLoadElementFromTuple(*element, etup, true, loadVec);
+		HnswLoadElementFromTuple(*element, etup, true, loadVec, index);
 	}

 	UnlockReleaseBuffer(buf);
@@ -565,36 +757,36 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
 * Load an element and optionally get its distance from q
 */
 void
-HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
+HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
 {
-	HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element);
+	HnswLoadElementImpl(element->blkno, element->offno, distance, matches, q, index, support, loadVec, maxDistance, &element);
 }

 /*
 * Get the distance for an element
 */
 static double
-GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport * support)
+GetElementDistance(char *base, HnswElement element, bool *matches, HnswQuery * q, Relation index, HnswSupport * support)
 {
 	Datum		value = HnswGetValue(base, element);
+	IndexTuple	itup = HnswPtrAccess(base, element->itup);

-	return HnswGetDistance(q->value, value, support);
+	return HnswGetDistance(itup, value, q, index, support, matches);
 }

 /*
 * Create a candidate for the entry point
 */
 HnswSearchCandidate *
-HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
+HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, bool inMemory)
 {
 	HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
-	bool		inMemory = index == NULL;

 	HnswPtrStore(base, sc->element, entryPoint);
 	if (inMemory)
-		sc->distance = GetElementDistance(base, entryPoint, q, support);
+		sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, index, support);
 	else
-		HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL);
+		HnswLoadElement(entryPoint, &sc->distance, &sc->matches, q, index, support, loadVec, NULL);
 	return sc;
 }

@@ -798,7 +990,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
 * Algorithm 2 from paper
 */
 List *
-HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
+HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
 {
 	List	   *w = NIL;
 	pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -811,7 +1003,8 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 	int			lm = HnswGetLayerM(m, lc);
 	HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
 	int			unvisitedLength;
-	bool		inMemory = index == NULL;
+	uint64		additional = 0;
+	uint64		maxAdditional = q->keyData && lc == 0 ? 10000 : 0;

 	if (v == NULL)
 	{
@@ -851,6 +1044,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 		pairingheap_add(C, &sc->c_node);
 		pairingheap_add(W, &sc->w_node);

+		/* Do not count elements that do not match filter towards ef */
+		if (!sc->matches && ++additional <= maxAdditional)
+			continue;
+
 		/*
 		 * Do not count elements being deleted towards ef when vacuuming. It
 		 * would be ideal to do this for inserts as well, but this could
@@ -884,6 +1081,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 			HnswElement eElement;
 			HnswSearchCandidate *e;
 			double		eDistance;
+			bool		eMatches;
 			bool		alwaysAdd = wlen < ef;

 			f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
@@ -891,7 +1089,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 			if (inMemory)
 			{
 				eElement = unvisited[i].element;
-				eDistance = GetElementDistance(base, eElement, q, support);
+				eDistance = GetElementDistance(base, eElement, &eMatches, q, index, support);
 			}
 			else
 			{
@@ -901,7 +1099,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in

 				/* Avoid any allocations if not adding */
 				eElement = NULL;
-				HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
+				HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);

 				if (eElement == NULL)
 					continue;
@@ -929,6 +1127,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 			e = palloc(sizeof(HnswSearchCandidate));
 			HnswPtrStore(base, e->element, eElement);
 			e->distance = eDistance;
+			e->matches = eMatches;
 			pairingheap_add(C, &e->c_node);
 			pairingheap_add(W, &e->w_node);

@@ -939,6 +1138,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 			 */
 			if (CountElement(skipElement, eElement))
 			{
+				/* Do not count elements that do not match filter towards ef */
+				if (!e->matches && ++additional <= maxAdditional)
+					continue;
+
 				wlen++;

 				/* No need to decrement wlen */
@@ -1016,18 +1219,24 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
 * Check if an element is closer to q than any element from R
 */
 static bool
-CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support)
+CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswSupport * support)
 {
 	HnswElement eElement = HnswPtrAccess(base, e->element);
-	Datum		eValue = HnswGetValue(base, eElement);
+	HnswQuery	q;
 	ListCell   *lc2;

+	q.value = HnswGetValue(base, eElement);
+	q.itup = HnswPtrAccess(base, eElement->itup);
+	q.keyData = NULL;
+
 	foreach(lc2, r)
 	{
 		HnswCandidate *ri = lfirst(lc2);
 		HnswElement riElement = HnswPtrAccess(base, ri->element);
 		Datum		riValue = HnswGetValue(base, riElement);
-		float		distance = HnswGetDistance(eValue, riValue, support);
+		IndexTuple	ritup = HnswPtrAccess(base, riElement->itup);
+		bool		matches;
+		float		distance = HnswGetDistance(ritup, riValue, &q, index, support, &matches);

 		if (distance <= e->distance)
 			return false;
@@ -1040,7 +1249,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support
 * Algorithm 4 from paper
 */
 static List *
-SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
+SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
 {
 	List	   *r = NIL;
 	List	   *w = list_copy(c);
@@ -1074,7 +1283,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer

 		/* Use previous state of r and wd to skip work when possible */
 		if (mustCalculate)
-			e->closer = CheckElementCloser(base, e, r, support);
+			e->closer = CheckElementCloser(base, e, r, index, support);
 		else if (list_length(added) > 0)
 		{
 			/* Keep Valgrind happy for in-memory, parallel builds */
@@ -1087,8 +1296,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
 			 */
 			if (e->closer)
 			{
-				e->closer = CheckElementCloser(base, e, added, support);
-
+				e->closer = CheckElementCloser(base, e, added, index, support);
 				if (!e->closer)
 					removedAny = true;
 			}
@@ -1100,7 +1308,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
 				 */
 				if (removedAny)
 				{
-					e->closer = CheckElementCloser(base, e, r, support);
+					e->closer = CheckElementCloser(base, e, r, index, support);
 					if (e->closer)
 						added = lappend(added, e);
 				}
@@ -1108,7 +1316,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
 		}
 		else if (e == newCandidate)
 		{
-			e->closer = CheckElementCloser(base, e, r, support);
+			e->closer = CheckElementCloser(base, e, r, index, support);
 			if (e->closer)
 				added = lappend(added, e);
 		}
@@ -1185,7 +1393,7 @@ HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newE
 			c = lappend(c, &neighbors->items[i]);
 		c = lappend(c, &newHc);

-		SelectNeighbors(base, c, lm, support, &neighbors->closerSet, &newHc, &pruned, true);
+		SelectNeighbors(base, c, lm, index, support, &neighbors->closerSet, &newHc, &pruned, true);

 		/* Should not happen */
 		if (pruned == NULL)
@@ -1256,17 +1464,19 @@ PrecomputeHash(char *base, HnswElement element)
 * Algorithm 1 from paper
 */
 void
-HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing)
+HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory)
 {
 	List	   *ep;
 	List	   *w;
 	int			level = element->level;
 	int			entryLevel;
 	HnswQuery	q;
+
 	HnswElement skipElement = existing ? element : NULL;
-	bool		inMemory = index == NULL;

 	q.value = HnswGetValue(base, element);
+	q.itup = HnswPtrAccess(base, element->itup);
+	q.keyData = NULL;

 	/* Precompute hash */
 	if (inMemory)
@@ -1277,13 +1487,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
 		return;

 	/* Get entry point and level */
-	ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true));
+	ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true, inMemory));
 	entryLevel = entryPoint->level;

 	/* 1st phase: greedy search to insert level */
 	for (int lc = entryLevel; lc >= level + 1; lc--)
 	{
-		w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
+		w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
 		ep = w;
 	}

@@ -1302,7 +1512,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
 		List	   *lw = NIL;
 		ListCell   *lc2;

-		w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
+		w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);

 		/* Convert search candidates to candidates */
 		foreach(lc2, w)
@@ -1326,7 +1536,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
 		 * sortCandidates to true for in-memory builds to enable closer
 		 * caching, but there does not seem to be a difference in performance.
 		 */
-		neighbors = SelectNeighbors(base, lw, lm, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);
+		neighbors = SelectNeighbors(base, lw, lm, index, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);

 		AddConnections(base, element, neighbors, lc);

--- a/src/hnswvacuum.c
+++ b/src/hnswvacuum.c
@@ -204,7 +204,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
 	element->heaptidsLength = 0;

 	/* Find neighbors for element, skipping itself */
-	HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true);
+	HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true, false);

 	/* Zero memory for each element */
 	MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
 		LockPage(index, HNSW_UPDATE_LOCK, ShareLock);

 		/* Load element */
-		HnswLoadElement(highestPoint, NULL, NULL, index, support, true, NULL);
+		HnswLoadElement(highestPoint, NULL, NULL, NULL, index, support, true, NULL);

 		/* Repair if needed */
 		if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
 			 * is outdated, this can remove connections at higher levels in
 			 * the graph until they are repaired, but this should be fine.
 			 */
-			HnswLoadElement(entryPoint, NULL, NULL, index, support, true, NULL);
+			HnswLoadElement(entryPoint, NULL, NULL, NULL, index, support, true, NULL);

 			if (NeedsUpdated(vacuumstate, entryPoint))
 			{
@@ -370,7 +370,7 @@ RepairGraph(HnswVacuumState * vacuumstate)

 			/* Create an element */
 			element = HnswInitElementFromBlock(blkno, offno);
-			HnswLoadElementFromTuple(element, etup, false, true);
+			HnswLoadElementFromTuple(element, etup, false, true, index);

 			elements = lappend(elements, element);
 		}
@@ -440,6 +440,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 	BlockNumber insertPage = InvalidBlockNumber;
 	Relation	index = vacuumstate->index;
 	BufferAccessStrategy bas = vacuumstate->bas;
+	bool		useIndexTuple = HnswUseIndexTuple(index);

 	/*
 	 * Wait for index scans to complete. Scans before this point may contain
@@ -521,7 +522,14 @@ MarkDeleted(HnswVacuumState * vacuumstate)

 			/* Overwrite element */
 			etup->deleted = 1;
-			MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
+			if (useIndexTuple)
+			{
+				IndexTuple	itup = (IndexTuple) &etup->data;
+
+				MemSet(itup, 0, IndexTupleSize(itup));
+			}
+			else
+				MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));

 			/* Overwrite neighbors */
 			for (int i = 0; i < ntup->count; i++)
--- a/src/ivfbuild.c
+++ b/src/ivfbuild.c
@@ -138,7 +138,7 @@ SampleRows(IvfflatBuildState * buildstate)
 * Add tuple to sort
 */
 static void
-AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState * buildstate)
+AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, IvfflatBuildState * buildstate)
 {
 	double		distance;
 	double		minDistance = DBL_MAX;
@@ -184,6 +184,11 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
 	slot->tts_isnull[1] = false;
 	slot->tts_values[2] = value;
 	slot->tts_isnull[2] = false;
+	for (int i = 1; i < buildstate->tupdesc->natts; i++)
+	{
+		slot->tts_values[2 + i] = values[i];
+		slot->tts_isnull[2 + i] = isnull[i];
+	}
 	ExecStoreVirtualTuple(slot);

 	/*
@@ -215,7 +220,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
 	oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);

 	/* Add tuple to sort */
-	AddTupleToSort(index, tid, values, buildstate);
+	AddTupleToSort(index, tid, values, isnull, buildstate);

 	/* Reset memory context */
 	MemoryContextSwitchTo(oldCtx);
@@ -226,19 +231,20 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
 * Get index tuple from sort state
 */
 static inline void
-GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list)
+GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, Datum *values, bool *isnull, IndexTuple *itup, int *list)
 {
-	Datum		value;
-	bool		isnull;
-
 	if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
 	{
-		*list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
-		value = slot_getattr(slot, 3, &isnull);
+		bool		unused;
+
+		*list = DatumGetInt32(slot_getattr(slot, 1, &unused));
+
+		for (int i = 0; i < tupdesc->natts; i++)
+			values[i] = slot_getattr(slot, 3 + i, &isnull[i]);

 		/* Form the index tuple */
-		*itup = index_form_tuple(tupdesc, &value, &isnull);
-		(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull)));
+		*itup = index_form_tuple(tupdesc, values, isnull);
+		(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &unused)));
 	}
 	else
 		*list = -1;
@@ -254,14 +260,16 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
 	IndexTuple	itup = NULL;	/* silence compiler warning */
 	int64		inserted = 0;

-	TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple);
-	TupleDesc	tupdesc = RelationGetDescr(index);
+	TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
+	TupleDesc	tupdesc = buildstate->tupdesc;
+	Datum	   *values = palloc(tupdesc->natts * sizeof(Datum));
+	bool	   *isnull = palloc(tupdesc->natts * sizeof(bool));

 	pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);

 	pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples);

-	GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
+	GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);

 	for (int i = 0; i < buildstate->centers->length; i++)
 	{
@@ -297,7 +305,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)

 			pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted);

-			GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
+			GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
 		}

 		insertPage = BufferGetBlockNumber(buf);
@@ -307,6 +315,9 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
 		/* Set the start and insert pages */
 		IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
 	}
+
+	pfree(values);
+	pfree(isnull);
 }

 /*
@@ -319,6 +330,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 	buildstate->index = index;
 	buildstate->indexInfo = indexInfo;
 	buildstate->typeInfo = IvfflatGetTypeInfo(index);
+	buildstate->tupdesc = RelationGetDescr(index);

 	buildstate->lists = IvfflatGetLists(index);
 	buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
@@ -356,12 +368,13 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 				 errmsg("dimensions must be greater than one for this opclass")));

 	/* Create tuple description for sorting */
-	buildstate->tupdesc = CreateTemplateTupleDesc(3);
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
+	buildstate->sortdesc = CreateTemplateTupleDesc(2 + buildstate->tupdesc->natts);
+	TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
+	TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
+	for (int i = 0; i < buildstate->tupdesc->natts; i++)
+		TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) (3 + i), NULL, buildstate->tupdesc->attrs[i].atttypid, -1, 0);

-	buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
+	buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);

 	buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions));
 	buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
@@ -633,7 +646,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
 	InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
 	memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
 	buildstate.centers->length = buildstate.centers->maxlen;
-	ivfspool->sortstate = InitBuildSortState(buildstate.tupdesc, sortmem, coordinate);
+	ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, sortmem, coordinate);
 	buildstate.sortstate = ivfspool->sortstate;
 	scan = table_beginscan_parallel(ivfspool->heap,
 									ParallelTableScanFromIvfflatShared(ivfshared));
@@ -950,7 +963,7 @@ AssignTuples(IvfflatBuildState * buildstate)
 	}

 	/* Begin serial/leader tuplesort */
-	buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate);
+	buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate);

 	/* Add tuples to sort */
 	if (buildstate->heap != NULL)
--- a/src/ivfflat.c
+++ b/src/ivfflat.c
@@ -17,8 +17,16 @@
 #endif

 int			ivfflat_probes;
+int			ivfflat_iterative_search;
+int			ivfflat_iterative_search_max_probes;
 static relopt_kind ivfflat_relopt_kind;

+static const struct config_enum_entry ivfflat_iterative_search_options[] = {
+	{"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false},
+	{"on", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false},
+	{NULL, 0, false}
+};
+
 /*
 * Initialize index options and variables
 */
@@ -33,6 +41,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);

+	DefineCustomEnumVariable("ivfflat.iterative_search", "Sets whether to use iterative search",
+							 NULL, &ivfflat_iterative_search,
+							 IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
+
+	DefineCustomIntVariable("ivfflat.iterative_search_max_probes", "Sets the max number of probes for iterative search",
+							"Zero sets to the number of lists", &ivfflat_iterative_search_max_probes,
+							0, 0, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
+
 	MarkGUCPrefixReserved("ivfflat");
 }

--- a/src/ivfflat.h
+++ b/src/ivfflat.h
@@ -80,6 +80,14 @@

 /* Variables */
 extern int	ivfflat_probes;
+extern int	ivfflat_iterative_search;
+extern int	ivfflat_iterative_search_max_probes;
+
+typedef enum IvfflatIterativeSearchType
+{
+	IVFFLAT_ITERATIVE_SEARCH_OFF,
+	IVFFLAT_ITERATIVE_SEARCH_RELAXED
+}			IvfflatIterativeSearchType;

 typedef struct VectorArrayData
 {
@@ -165,6 +173,7 @@ typedef struct IvfflatBuildState
 	Relation	index;
 	IndexInfo  *indexInfo;
 	const		IvfflatTypeInfo *typeInfo;
+	TupleDesc	tupdesc;

 	/* Settings */
 	int			dimensions;
@@ -198,7 +207,7 @@ typedef struct IvfflatBuildState

 	/* Sorting */
 	Tuplesortstate *sortstate;
-	TupleDesc	tupdesc;
+	TupleDesc	sortdesc;
 	TupleTableSlot *slot;

 	/* Memory */
@@ -247,8 +256,10 @@ typedef struct IvfflatScanOpaqueData
 {
 	const		IvfflatTypeInfo *typeInfo;
 	int			probes;
+	int			maxProbes;
 	int			dimensions;
 	bool		first;
+	Datum		value;

 	/* Sorting */
 	Tuplesortstate *sortstate;
@@ -265,6 +276,8 @@ typedef struct IvfflatScanOpaqueData

 	/* Lists */
 	pairingheap *listQueue;
+	BlockNumber *listPages;
+	int			listIndex;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;

--- a/src/ivfinsert.c
+++ b/src/ivfinsert.c
@@ -78,6 +78,8 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
 	BlockNumber insertPage = InvalidBlockNumber;
 	ListInfo	listInfo;
 	BlockNumber originalInsertPage;
+	TupleDesc	tupdesc = RelationGetDescr(index);
+	Datum	   *newValues = palloc(tupdesc->natts * sizeof(Datum));

 	/* Detoast once for all calls */
 	value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -98,12 +100,16 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
 	IvfflatGetMetaPageInfo(index, NULL, NULL);

 	/* Find the insert page - sets the page and list info */
-	FindInsertPage(index, values, &insertPage, &listInfo);
+	FindInsertPage(index, &value, &insertPage, &listInfo);
 	Assert(BlockNumberIsValid(insertPage));
 	originalInsertPage = insertPage;

+	newValues[0] = value;
+	for (int i = 1; i < tupdesc->natts; i++)
+		newValues[i] = values[i];
+
 	/* Form tuple */
-	itup = index_form_tuple(RelationGetDescr(index), &value, isnull);
+	itup = index_form_tuple(tupdesc, newValues, isnull);
 	itup->t_tid = *heap_tid;

 	/* Get tuple size */
--- a/src/ivfscan.c
+++ b/src/ivfscan.c
@@ -65,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;

@@ -78,7 +78,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
 				pairingheap_add(so->listQueue, &scanlist->ph_node);

 				/* Calculate max distance */
-				if (listCount == so->probes)
+				if (listCount == so->maxProbes)
 					maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
 			}
 			else if (distance < maxDistance)
@@ -102,6 +102,11 @@ GetScanLists(IndexScanDesc scan, Datum value)

 		UnlockReleaseBuffer(cbuf);
 	}
+
+	for (int i = listCount - 1; i >= 0; i--)
+		so->listPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
+
+	Assert(pairingheap_is_empty(so->listQueue));
 }

 /*
@@ -114,11 +119,14 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
 	double		tuples = 0;
 	TupleTableSlot *slot = so->vslot;
+	int			batchProbes = 0;
+
+	tuplesort_reset(so->sortstate);

 	/* Search closest probes lists */
-	while (!pairingheap_is_empty(so->listQueue))
+	while (so->listIndex < so->maxProbes && (++batchProbes) <= so->probes)
 	{
-		BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
+		BlockNumber searchPage = so->listPages[so->listIndex++];

 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -166,13 +174,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
 		}
 	}

-	if (tuples < 100)
+	if (tuples < 100 && ivfflat_iterative_search == IVFFLAT_ITERATIVE_SEARCH_OFF)
 		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
 }

 /*
@@ -240,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);

@@ -249,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_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF)
+	{
+		if (ivfflat_iterative_search_max_probes == 0)
+			maxProbes = lists;
+		else
+			maxProbes = Min(ivfflat_iterative_search_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 */
@@ -280,6 +304,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	so->bas = GetAccessStrategy(BAS_BULKREAD);

 	so->listQueue = pairingheap_allocate(CompareLists, scan);
+	so->listPages = palloc(maxProbes * sizeof(BlockNumber));
+	so->listIndex = 0;

 	scan->opaque = so;

@@ -294,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->listIndex = 0;

 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -314,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
@@ -341,28 +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)
-		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));
+		/* TODO clean up if we allocated a new value */
 	}

-	if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
+	while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
 	{
-		bool		isnull;
-		ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
+		if (so->listIndex == so->maxProbes)
+			return false;

-		scan->xs_heaptid = *heaptid;
-		scan->xs_recheck = false;
-		scan->xs_recheckorderby = false;
-		return true;
+		IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
 	}

-	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;
 }

 /*
@@ -374,6 +397,7 @@ ivfflatendscan(IndexScanDesc scan)
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;

 	pairingheap_free(so->listQueue);
+	pfree(so->listPages);
 	tuplesort_end(so->sortstate);
 	FreeAccessStrategy(so->bas);
 	FreeTupleDesc(so->tupdesc);
--- a/src/ivfvacuum.c
+++ b/src/ivfvacuum.c
@@ -26,7 +26,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
 		Page		cpage;
 		OffsetNumber coffno;
 		OffsetNumber cmaxoffno;
-		BlockNumber startPages[MaxOffsetNumber];
+		BlockNumber listPages[MaxOffsetNumber];
 		ListInfo	listInfo;

 		cbuf = ReadBuffer(index, blkno);
@@ -40,7 +40,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
 		{
 			IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));

-			startPages[coffno - FirstOffsetNumber] = list->startPage;
+			listPages[coffno - FirstOffsetNumber] = list->startPage;
 		}

 		listInfo.blkno = blkno;
@@ -50,7 +50,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,

 		for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
 		{
-			BlockNumber searchPage = startPages[coffno - FirstOffsetNumber];
+			BlockNumber searchPage = listPages[coffno - FirstOffsetNumber];
 			BlockNumber insertPage = InvalidBlockNumber;

 			/* Iterate over entry pages */
--- a/test/t/002_ivfflat_vacuum.pl
+++ b/test/t/002_ivfflat_vacuum.pl
@@ -6,13 +6,7 @@ use Test::More;

 my $dim = 3;

-my @r = ();
-for (1 .. $dim)
-{
-	my $v = int(rand(1000)) + 1;
-	push(@r, "i % $v");
-}
-my $array_sql = join(", ", @r);
+my $array_sql = join(",", ('random()') x $dim);

 # Initialize node
 my $node = PostgreSQL::Test::Cluster->new('node');
@@ -23,19 +17,20 @@ $node->start;
 $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 % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
+	"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);");

 # Get size
 my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");

+# Store values
+$node->safe_psql("postgres", "CREATE TABLE tmp AS SELECT * FROM tst;");
+
 # Delete all, vacuum, and insert same data
 $node->safe_psql("postgres", "DELETE FROM tst;");
 $node->safe_psql("postgres", "VACUUM tst;");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
-);
+$node->safe_psql("postgres", "INSERT INTO tst SELECT * FROM tmp;");

 # Check size
 my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
--- a/test/t/041_ivfflat_iterative_search.pl
+++ b/test/t/041_ivfflat_iterative_search.pl
@@ -0,0 +1,54 @@
+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 PRIMARY KEY, 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.iterative_search = 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);
+
+foreach ((30, 50, 70))
+{
+	my $max_probes = $_;
+	my $expected = $max_probes / 10;
+	my $sum = 0;
+
+	for my $i (1 .. 20)
+	{
+		$count = $node->safe_psql("postgres", qq(
+			SET enable_seqscan = off;
+			SET ivfflat.probes = 10;
+			SET ivfflat.iterative_search = on;
+			SET ivfflat.iterative_search_max_probes = $max_probes;
+			SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
+		));
+		$sum += $count;
+	}
+
+	my $avg = $sum / 20;
+	cmp_ok($avg, '>', $expected - 2);
+	cmp_ok($avg, '<', $expected + 2);
+}
+
+done_testing();
--- a/test/t/042_ivfflat_iterative_search_recall.pl
+++ b/test/t/042_ivfflat_iterative_search_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.iterative_search = 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.iterative_search = 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();
--- a/test/t/043_hnsw_iterative_search.pl
+++ b/test/t/043_hnsw_iterative_search.pl
@@ -26,7 +26,7 @@ $node->safe_psql("postgres", qq(

 my $count = $node->safe_psql("postgres", qq(
 	SET enable_seqscan = off;
-	SET hnsw.streaming = on;
+	SET hnsw.iterative_search = on;
 	SET work_mem = '8MB';
 	SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
 ));
@@ -34,16 +34,16 @@ is($count, 10);

 foreach ((30000, 50000, 70000))
 {
-	my $ef_stream = $_;
-	my $expected = $ef_stream / 10000;
+	my $max_tuples = $_;
+	my $expected = $max_tuples / 10000;
 	my $sum = 0;

 	for my $i (1 .. 20)
 	{
 		$count = $node->safe_psql("postgres", qq(
 			SET enable_seqscan = off;
-			SET hnsw.streaming = on;
-			SET hnsw.ef_stream = $ef_stream;
+			SET hnsw.iterative_search = on;
+			SET hnsw.iterative_search_max_tuples = $max_tuples;
 			SET work_mem = '8MB';
 			SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
 		));
@@ -57,7 +57,7 @@ foreach ((30000, 50000, 70000))

 my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
 	SET enable_seqscan = off;
-	SET hnsw.streaming = on;
+	SET hnsw.iterative_search = on;
 	SET client_min_messages = debug1;
 	SET work_mem = '2MB';
 	SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
--- a/test/t/044_hnsw_iterative_search_recall.pl
+++ b/test/t/044_hnsw_iterative_search_recall.pl
@@ -21,7 +21,7 @@ sub test_recall
 	my $explain = $node->safe_psql("postgres", qq(
 		SET enable_seqscan = off;
 		SET hnsw.ef_search = $ef_search;
-		SET hnsw.streaming = on;
+		SET hnsw.iterative_search = 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/);
@@ -31,7 +31,7 @@ sub test_recall
 		my $actual = $node->safe_psql("postgres", qq(
 			SET enable_seqscan = off;
 			SET hnsw.ef_search = $ef_search;
-			SET hnsw.streaming = on;
+			SET hnsw.iterative_search = on;
 			SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
 		));
 		my @actual_ids = split("\n", $actual);
--- a/test/t/045_hnsw_hqann.pl
+++ b/test/t/045_hnsw_hqann.pl
@@ -0,0 +1,113 @@
+use strict;
+use warnings FATAL => 'all';
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node;
+my @queries = ();
+my @cs = ();
+my @expected;
+my $limit = 20;
+my $dim = 3;
+my $array_sql = join(",", ('random()') x $dim);
+my $nc = 1000;
+
+sub test_recall
+{
+	my ($min, $operator) = @_;
+	my $correct = 0;
+	my $total = 0;
+
+	my $explain = $node->safe_psql("postgres", qq(
+		SET enable_seqscan = off;
+		EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $cs[0] ORDER BY v $operator '$queries[0]' LIMIT $limit;
+	));
+	like($explain, qr/Index Cond/);
+
+	for my $i (0 .. $#queries)
+	{
+		my $actual = $node->safe_psql("postgres", qq(
+			SET enable_seqscan = off;
+			SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v $operator '$queries[$i]' LIMIT $limit;
+		));
+		my @actual_ids = split("\n", $actual);
+		my %actual_set = map { $_ => 1 } @actual_ids;
+
+		is(scalar(@actual_ids), $limit);
+
+		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 = 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), c int4);");
+$node->safe_psql("postgres",
+	"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 20000) i;"
+);
+
+# Generate queries
+for (1 .. 20)
+{
+	my @r = ();
+	for (1 .. $dim)
+	{
+		push(@r, rand());
+	}
+	push(@queries, "[" . join(",", @r) . "]");
+	push(@cs, int(rand() * $nc));
+}
+
+# Get exact results
+@expected = ();
+for my $i (0 .. $#queries)
+{
+	my $res = $node->safe_psql("postgres", "SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <=> '$queries[$i]' LIMIT $limit;");
+	push(@expected, $res);
+}
+
+# Add index
+$node->safe_psql("postgres", qq(
+	SET maintenance_work_mem = '256MB';
+	SET max_parallel_maintenance_workers = 2;
+	CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, c);
+));
+
+# Test recall
+test_recall(0.99, '<=>');
+
+# Test vacuum
+$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
+$node->safe_psql("postgres", "VACUUM tst;");
+
+# Test columns
+my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);");
+like($stderr, qr/first column must be a vector/);
+
+($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_cosine_ops);");
+like($stderr, qr/first column must be a vector/);
+
+($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, c, c);");
+like($stderr, qr/index cannot have more than two columns/);
+
+($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, v vector_cosine_ops);");
+like($stderr, qr/column 2 cannot be a vector/);
+
+done_testing();
Author	SHA1	Message	Date
Andrew Kane	70785ff0fe	Merge branch 'master' into hqann2	2024-10-10 19:06:26 -07:00
Andrew Kane	35b252a3e3	Switched to macos-13 on CI [skip ci]	2024-10-10 18:33:35 -07:00
Andrew Kane	2832e746f0	Use get_float8_infinity instead of INFINITY Co-authored-by: "Jonathan S. Katz" <jkatz@amazon.com>	2024-10-10 18:16:39 -07:00
Andrew Kane	961cb17d80	Added iterative search for HNSW [skip ci]	2024-10-10 18:14:39 -07:00
Andrew Kane	c91ed7b2c3	Added iterative search for IVFFlat [skip ci]	2024-10-10 18:12:27 -07:00
Andrew Kane	48fe70c219	Fixed sortdesc for multiple attributes [skip ci]	2024-10-10 15:22:23 -07:00
Andrew Kane	29908405ab	List assignment does not change with scaling [skip ci]	2024-10-10 15:16:08 -07:00
Andrew Kane	08d0340655	Improved IVFFlat vacuum test [skip ci]	2024-10-10 14:24:26 -07:00
Andrew Kane	7d2eb49c2a	Fixed list assignment for cosine distance and non-normalized vectors for IVFFlat	2024-10-10 12:54:13 -07:00
Andrew Kane	772ab69de6	Updated IVFFlat to support multiple attributes (not enabled yet)	2024-10-10 12:34:03 -07:00
Andrew Kane	e13e9a9614	Improved scoping [skip ci]	2024-10-10 12:19:08 -07:00
Andrew Kane	e2fab306ac	Added tupdesc to buildstate [skip ci]	2024-10-10 12:18:41 -07:00
Andrew Kane	edc2126a4a	Renamed tupdesc to sortdesc [skip ci]	2024-10-10 12:16:57 -07:00
Andrew Kane	a73fe9e9e1	Use same TupleDesc [skip ci]	2024-10-10 02:01:23 -07:00
Andrew Kane	2ffbb2500c	Merge branch 'master' into hqann2	2024-10-09 23:59:11 -07:00
Andrew Kane	46209740ce	Merge branch 'master' into hqann2	2024-10-09 23:26:57 -07:00
Andrew Kane	d650d23da8	Merge branch 'master' into hqann2	2024-10-09 22:06:29 -07:00
Andrew Kane	ad755d19ca	Merge branch 'master' into hqann2	2024-10-09 21:52:45 -07:00
Andrew Kane	f39b515ce7	Merge branch 'master' into hqann2	2024-10-09 21:50:53 -07:00
Andrew Kane	ab88612c83	Always use tuple for duplicates [skip ci]	2024-10-09 21:30:10 -07:00
Andrew Kane	356f175359	Merge branch 'master' into hqann2	2024-10-09 21:25:30 -07:00
Andrew Kane	a49a2b9480	Improved test [skip ci]	2024-10-09 19:54:06 -07:00
Andrew Kane	3b61384367	Added HnswGetElementTupleSize method [skip ci]	2024-10-09 19:37:09 -07:00
Andrew Kane	67f9a3e61c	Init collation as well [skip ci]	2024-10-09 19:30:04 -07:00
Andrew Kane	3ccfab8f92	Added support for inline filtering with HNSW	2024-10-09 19:02:40 -07:00