Add 8-bit scalar quantization support for IVF index.

Makefile

@@ -3,7 +3,7 @@ EXTVERSION = 0.4.4
 
 MODULE_big = vector
 DATA = $(wildcard sql/*--*.sql)
-OBJS = src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o
+OBJS = src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/common/ivf_list.o src/common/ivf_options.o src/common/metadata.o src/fixed_point/ivf_sq.o src/fixed_point/scalar_quantizer.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o
 
 TESTS = $(wildcard test/sql/*.sql)
 REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))

Makefile.win

@@ -1,7 +1,7 @@
 EXTENSION = vector
 EXTVERSION = 0.4.4
 
-OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
+OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\common\ivf_list.obj src\common\ivf_options.obj src\common\metadata.obj src\fixed_point\ivf_sq.obj src\fixed_point\scalar_quantizer.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
 
 REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
 REGRESS_OPTS = --inputdir=test --load-extension=vector

README.md

@@ -165,43 +165,56 @@ Three keys to achieving good recall are:
 2. Choose an appropriate number of lists - a good place to start is `rows / 1000` for up to 1M rows and `sqrt(rows)` for over 1M rows
 3. When querying, specify an appropriate number of [probes](#query-options) (higher is better for recall, lower is better for speed) - a good place to start is `sqrt(lists)`
 
-Add an index for each distance function you want to use.
+The IVF index supports both `flat` and `8-bit scalar quantized` storage. The `8-bit quantized` vectors map each vector from `float` space to `int8` space, resulting in up to ~2X faster scoring and ~4X lower storage cost with nearly no recall loss.
+
+Add an index for each distance function you want to use. Specify `ivfflat` to use the `flat` storage, or `ivf` with `quantizer = 'SQ8'` option to use the `8-bit quantized` IVF index.
 
 L2 distance
 
 ```sql
+CREATE INDEX ON items USING ivf (embedding vector_l2_ops) WITH (lists = 100, quantizer = 'SQ8');
 CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
 ```
 
 Inner product
 
 ```sql
+CREATE INDEX ON items USING ivf (embedding vector_ip_ops) WITH (lists = 100, quantizer='SQ8');
 CREATE INDEX ON items USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);
 ```
 
 Cosine distance
 
 ```sql
+CREATE INDEX ON items USING ivf (embedding vector_cosine_ops) WITH (lists = 100, quantizer='SQ8');
 CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
 ```
 
-Vectors with up to 2,000 dimensions can be indexed.
+Vectors with up to 8,000 dimensions can be indexed using `IVF` when `8-bit scalar quantization` is enabled, up to 2,000 dimensions can be indexed using `ivfflat` without quantization.
 
 ### Query Options
 
 Specify the number of probes (1 by default)
 
+#### ivfflat
+
 ```sql
 SET ivfflat.probes = 10;
 ```
 
+#### ivf
+
+```sql
+SET ivf.probes = 10;
+```
+
 A higher value provides better recall at the cost of speed, and it can be set to the number of lists for exact nearest neighbor search (at which point the planner won’t use the index)
 
 Use `SET LOCAL` inside a transaction to set it for a single query
 
 ```sql
 BEGIN;
-SET LOCAL ivfflat.probes = 10;
+SET LOCAL ivf[flat].probes = 10;
 SELECT ...
 COMMIT;
 ```
@@ -240,6 +253,9 @@ CREATE INDEX ON items (category_id);
 Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
 
 ```sql
+CREATE INDEX ON items USING ivf (embedding vector_l2_ops) WITH (lists = 100, quantizer='SQ8')
+    WHERE (category_id = 123);
+
 CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)
     WHERE (category_id = 123);
 ```
@@ -286,6 +302,8 @@ SELECT * FROM items ORDER BY embedding <#> '[3,1,2]' LIMIT 5;
 To speed up queries with an index, increase the number of inverted lists (at the expense of recall).
 
 ```sql
+CREATE INDEX ON items USING ivf (embedding vector_l2_ops) WITH (lists = 1000, quantizer='SQ8');
+
 CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 1000);
 ```
 
@@ -324,7 +342,7 @@ A non-partitioned table has a limit of 32 TB by default in Postgres. A partition
 
 Yes, pgvector uses the write-ahead log (WAL), which allows for replication and point-in-time recovery.
 
-#### What if I want to index vectors with more than 2,000 dimensions?
+#### What if I want to index vectors with more than 8,000 dimensions?
 
 You’ll need to use [dimensionality reduction](https://en.wikipedia.org/wiki/Dimensionality_reduction) at the moment.
 

sql/vector.sql

@@ -99,6 +99,12 @@ CREATE FUNCTION vector_avg(double precision[]) RETURNS vector
 CREATE FUNCTION vector_combine(double precision[], double precision[]) RETURNS double precision[]
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION inner_product_int8_float_batched(internal, internal, internal) RETURNS internal
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
+CREATE FUNCTION squared_l2_distance_int8_float_batched(internal, internal, internal) RETURNS internal
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 -- aggregates
 
 CREATE AGGREGATE avg(vector) (
@@ -236,6 +242,13 @@ CREATE ACCESS METHOD ivfflat TYPE INDEX HANDLER ivfflathandler;
 
 COMMENT ON ACCESS METHOD ivfflat IS 'ivfflat index access method';
 
+CREATE FUNCTION ivfhandler(internal) RETURNS index_am_handler
+       AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE ACCESS METHOD ivf TYPE INDEX HANDLER ivfhandler;
+
+COMMENT ON ACCESS METHOD ivf IS 'ivf index access method';
+
 CREATE FUNCTION hnswhandler(internal) RETURNS index_am_handler
 	AS 'MODULE_PATHNAME' LANGUAGE C;
 
@@ -290,3 +303,27 @@ CREATE OPERATOR CLASS vector_cosine_ops
 	OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
 	FUNCTION 1 vector_negative_inner_product(vector, vector),
 	FUNCTION 2 vector_norm(vector);
+
+CREATE OPERATOR CLASS vector_l2_ops
+	DEFAULT FOR TYPE vector USING ivf AS
+	OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops,
+	FUNCTION 1 vector_l2_squared_distance(vector, vector),
+	FUNCTION 3 l2_distance(vector, vector),
+	FUNCTION 5 squared_l2_distance_int8_float_batched(internal, internal, internal);
+
+CREATE OPERATOR CLASS vector_ip_ops
+	FOR TYPE vector USING ivf AS
+	OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
+	FUNCTION 1 vector_negative_inner_product(vector, vector),
+	FUNCTION 3 vector_spherical_distance(vector, vector),
+	FUNCTION 4 vector_norm(vector),
+	FUNCTION 5 inner_product_int8_float_batched(internal, internal, internal);
+
+CREATE OPERATOR CLASS vector_cosine_ops
+	FOR TYPE vector USING ivf AS
+	OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
+	FUNCTION 1 vector_negative_inner_product(vector, vector),
+	FUNCTION 2 vector_norm(vector),
+	FUNCTION 3 vector_spherical_distance(vector, vector),
+	FUNCTION 4 vector_norm(vector),
+	FUNCTION 5 inner_product_int8_float_batched(internal, internal, internal);

src/common/ivf_list.c

@@ -0,0 +1,36 @@
+#include "ivf_list.h"
+
+#include "metadata.h"
+
+/* PG C headers */
+#include "c.h"
+#include "common/relpath.h"
+#include "pg_config.h"
+#include "utils/relcache.h"
+
+IvfListV2 CreateIvfListV2(Relation rel, Metadata *metadata, bool external,
+						  ForkNumber fork_num, size_t *list_size)
+{
+	IvfListV2 list;
+	Metadata *list_metadata = metadata;
+	size_t metadata_size =
+		external ? EXTERNAL_METADATA_SIZE : VARSIZE_ANY(metadata);
+
+	Assert(!VARATT_IS_EXTERNAL(metadata));
+	*list_size = offsetof(IvfListDataV2, metadata) + metadata_size;
+
+	list = (IvfListV2)palloc0(*list_size);
+	list->startPage = InvalidBlockNumber;
+	list->insertPage = InvalidBlockNumber;
+	list->unused = 0UL;
+	if (external)
+	{
+		list_metadata = WriteMetadata(rel, metadata, fork_num);
+		Assert(VARATT_IS_EXTERNAL(list_metadata));
+		Assert(((ExternalMetadata *)list_metadata)->length ==
+			   (VARSIZE_ANY_EXHDR(metadata)));
+	}
+
+	memcpy(&list->metadata, list_metadata, metadata_size);
+	return list;
+}

src/common/ivf_list.h

@@ -0,0 +1,87 @@
+#ifndef PGVECTOR_SRC_COMMON_IVF_LIST_H_
+#define PGVECTOR_SRC_COMMON_IVF_LIST_H_
+
+#include "postgres.h"
+
+#include "storage/block.h"
+#include "common/relpath.h"
+#include "utils/relcache.h"
+
+#include "metadata.h"
+
+/*
+ * Partition ("list") metadata. It stores the location of its leaf pages
+ * (`start_page`), the location to insert new data into leaf pages
+ * (`insert_page`), together with some user defined metadata.
+ *
+ * STORAGE FORMAT
+ * --------------------------------------------------------
+ * | start_page(4) | insert_page(4) | metadata (varlena) |
+ * --------------------------------------------------------
+ */
+typedef struct IvfListDataV1
+{
+	BlockNumber startPage;
+	BlockNumber insertPage;
+	/* Inline or externalized storage. */
+	Metadata metadata;
+} IvfListDataV1;
+
+typedef IvfListDataV1 *IvfListV1;
+
+/*
+ * Similar to the above but added 8-bytes `unused` space for future use.
+ * Note that the newly built indexes are forced to use this format.
+ */
+typedef struct IvfListDataV2
+{
+	BlockNumber startPage;
+	BlockNumber insertPage;
+	uint64_t unused;
+	// Inline or externalized storage.
+	Metadata metadata;
+} IvfListDataV2;
+
+typedef IvfListDataV2 *IvfListV2;
+
+#define IVF_LIST_GET_START_PAGE(item, version)         \
+	(((version) == 1) ? ((IvfListV1)(item))->startPage \
+					  : ((IvfListV2)(item))->startPage)
+
+#define IVF_LIST_SET_START_PAGE(item, version, startPage) \
+	do                                                    \
+	{                                                     \
+		if ((version) == 1)                               \
+			((IvfListV1)(item))->startPage = startPage;   \
+		else                                              \
+			((IvfListV2)(item))->startPage = startPage;   \
+	} while (0);
+
+#define IVF_LIST_SET_INSERT_PAGE(item, version, insertPage) \
+	do                                                      \
+	{                                                       \
+		if ((version) == 1)                                 \
+			((IvfListV1)(item))->insertPage = insertPage;   \
+		else                                                \
+			((IvfListV2)(item))->insertPage = insertPage;   \
+	} while (0);
+
+#define IVF_LIST_GET_INSERT_PAGE(item, version)         \
+	(((version) == 1) ? ((IvfListV1)(item))->insertPage \
+					  : ((IvfListV2)(item))->insertPage)
+
+#define IVF_LIST_GET_METADATA(item, version)           \
+	(((version) == 1) ? &((IvfListV1)(item))->metadata \
+					  : &((IvfListV2)(item))->metadata)
+
+/*
+ * Creates `IvfListV2` structure from `metadata`. The input `metadata` should be inlined.
+ * When `external` is `true`, it writes the `metadata` into external pages and an
+ * `ExternalMetadata` is written into the structure, otherwise, the `metadata` is copied
+ * into the result.
+ *   RETURNS: `IvfListV2` structure on the heap with its size `list_size`.
+ */
+IvfListV2 CreateIvfListV2(Relation rel, Metadata *metadata, bool external,
+						  ForkNumber fork_num, size_t *list_size);
+
+#endif /* PGVECTOR_SRC_COMMON_IVF_LIST_H_ */