From 9a57ee935e92c234552d93b884ec83221836c61c Mon Sep 17 00:00:00 2001 From: Benjamin Trent Date: Mon, 14 Oct 2024 20:13:27 -0400 Subject: [PATCH 1/2] Adding new bbq index types behind a feature flag (#114439) new index types of bbq_hnsw and bbq_flat which utilize the better binary quantization formats. A 32x reduction in memory, with nice recall properties. (cherry picked from commit 6c752abc231598f852ff47b4cde79bd97b9a1f5f) --- docs/changelog/114439.yaml | 5 + .../mapping/types/dense-vector.asciidoc | 41 ++++- .../search-your-data/knn-search.asciidoc | 92 ++++++++++ .../search.vectors/41_knn_search_bbq_hnsw.yml | 160 +++++++++++++++++ .../search.vectors/42_knn_search_bbq_flat.yml | 165 ++++++++++++++++++ server/src/main/java/module-info.java | 7 +- .../index/codec/vectors/BQVectorUtils.java | 8 + .../vectors/ES816BinaryFlatVectorsScorer.java | 49 ++---- .../vectors/OffHeapBinarizedVectorValues.java | 22 +++ ...RandomAccessBinarizedByteVectorValues.java | 14 ++ .../index/mapper/MapperFeatures.java | 9 +- .../vectors/DenseVectorFieldMapper.java | 160 ++++++++++++++++- .../index/store/LuceneFilesExtensions.java | 4 +- .../ES816BinaryFlatVectorsScorerTests.java | 2 +- .../vectors/DenseVectorFieldMapperTests.java | 66 ++++++- .../vectors/DenseVectorFieldTypeTests.java | 23 ++- 16 files changed, 768 insertions(+), 59 deletions(-) create mode 100644 docs/changelog/114439.yaml create mode 100644 rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml create mode 100644 rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml diff --git a/docs/changelog/114439.yaml b/docs/changelog/114439.yaml new file mode 100644 index 0000000000000..fd097d02f885f --- /dev/null +++ b/docs/changelog/114439.yaml @@ -0,0 +1,5 @@ +pr: 114439 +summary: Adding new bbq index types behind a feature flag +area: Vector Search +type: feature +issues: [] diff --git a/docs/reference/mapping/types/dense-vector.asciidoc b/docs/reference/mapping/types/dense-vector.asciidoc index 0cd9ee0578b70..44f90eded8632 100644 --- a/docs/reference/mapping/types/dense-vector.asciidoc +++ b/docs/reference/mapping/types/dense-vector.asciidoc @@ -115,22 +115,27 @@ that sacrifices result accuracy for improved speed. ==== Automatically quantize vectors for kNN search The `dense_vector` type supports quantization to reduce the memory footprint required when <> `float` vectors. -The two following quantization strategies are supported: +The three following quantization strategies are supported: + -- -`int8` - Quantizes each dimension of the vector to 1-byte integers. This can reduce the memory footprint by 75% at the cost of some accuracy. -`int4` - Quantizes each dimension of the vector to half-byte integers. This can reduce the memory footprint by 87% at the cost of some accuracy. +`int8` - Quantizes each dimension of the vector to 1-byte integers. This reduces the memory footprint by 75% (or 4x) at the cost of some accuracy. +`int4` - Quantizes each dimension of the vector to half-byte integers. This reduces the memory footprint by 87% (or 8x) at the cost of accuracy. +`bbq` - experimental:[] Better binary quantization which reduces each dimension to a single bit precision. This reduces the memory footprint by 96% (or 32x) at a larger cost of accuracy. Generally, oversampling during query time and reranking can help mitigate the accuracy loss. -- -To use a quantized index, you can set your index type to `int8_hnsw` or `int4_hnsw`. When indexing `float` vectors, the current default +When using a quantized format, you may want to oversample and rescore the results to improve accuracy. See <> for more information. + +To use a quantized index, you can set your index type to `int8_hnsw`, `int4_hnsw`, or `bbq_hnsw`. When indexing `float` vectors, the current default index type is `int8_hnsw`. NOTE: Quantization will continue to keep the raw float vector values on disk for reranking, reindexing, and quantization improvements over the lifetime of the data. -This means disk usage will increase by ~25% for `int8` and ~12.5% for `int4` due to the overhead of storing the quantized and raw vectors. +This means disk usage will increase by ~25% for `int8`, ~12.5% for `int4`, and ~3.1% for `bbq` due to the overhead of storing the quantized and raw vectors. NOTE: `int4` quantization requires an even number of vector dimensions. +NOTE: experimental:[] `bbq` quantization only supports vector dimensions that are greater than 64. + Here is an example of how to create a byte-quantized index: [source,console] @@ -173,6 +178,27 @@ PUT my-byte-quantized-index } -------------------------------------------------- +experimental:[] Here is an example of how to create a binary quantized index: + +[source,console] +-------------------------------------------------- +PUT my-byte-quantized-index +{ + "mappings": { + "properties": { + "my_vector": { + "type": "dense_vector", + "dims": 64, + "index": true, + "index_options": { + "type": "bbq_hnsw" + } + } + } + } +} +-------------------------------------------------- + [role="child_attributes"] [[dense-vector-params]] ==== Parameters for dense vector fields @@ -301,11 +327,16 @@ by 4x at the cost of some accuracy. See <>. +* experimental:[] `bbq_hnsw` - This utilizes the https://arxiv.org/abs/1603.09320[HNSW algorithm] in addition to automatically binary +quantization for scalable approximate kNN search with `element_type` of `float`. This can reduce the memory footprint +by 32x at the cost of accuracy. See <>. * `flat` - This utilizes a brute-force search algorithm for exact kNN search. This supports all `element_type` values. * `int8_flat` - This utilizes a brute-force search algorithm in addition to automatically scalar quantization. Only supports `element_type` of `float`. * `int4_flat` - This utilizes a brute-force search algorithm in addition to automatically half-byte scalar quantization. Only supports `element_type` of `float`. +* experimental:[] `bbq_flat` - This utilizes a brute-force search algorithm in addition to automatically binary quantization. Only supports +`element_type` of `float`. -- `m`::: (Optional, integer) diff --git a/docs/reference/search/search-your-data/knn-search.asciidoc b/docs/reference/search/search-your-data/knn-search.asciidoc index 70cf9eec121d7..6fb7f1747051f 100644 --- a/docs/reference/search/search-your-data/knn-search.asciidoc +++ b/docs/reference/search/search-your-data/knn-search.asciidoc @@ -1149,3 +1149,95 @@ POST product-index/_search ---- //TEST[continued] +[discrete] +[[dense-vector-knn-search-reranking]] +==== Oversampling and rescoring for quantized vectors + +All forms of quantization will result in some accuracy loss and as the quantization level increases the accuracy loss will also increase. +Generally, we have found that: +- `int8` requires minimal if any rescoring +- `int4` requires some rescoring for higher accuracy and larger recall scenarios. Generally, oversampling by 1.5x-2x recovers most of the accuracy loss. +- `bbq` requires rescoring except on exceptionally large indices or models specifically designed for quantization. We have found that between 3x-5x oversampling is generally sufficient. But for fewer dimensions or vectors that do not quantize well, higher oversampling may be required. + +There are two main ways to oversample and rescore. The first is to utilize the <> in the `_search` request. + +Here is an example using the top level `knn` search with oversampling and using `rescore` to rerank the results: + +[source,console] +-------------------------------------------------- +POST /my-index/_search +{ + "size": 10, <1> + "knn": { + "query_vector": [0.04283529, 0.85670587, -0.51402352, 0], + "field": "my_int4_vector", + "k": 20, <2> + "num_candidates": 50 + }, + "rescore": { + "window_size": 20, <3> + "query": { + "rescore_query": { + "script_score": { + "query": { + "match_all": {} + }, + "script": { + "source": "(dotProduct(params.queryVector, 'my_int4_vector') + 1.0)", <4> + "params": { + "queryVector": [0.04283529, 0.85670587, -0.51402352, 0] + } + } + } + }, + "query_weight": 0, <5> + "rescore_query_weight": 1 <6> + } + } +} +-------------------------------------------------- +// TEST[skip: setup not provided] +<1> The number of results to return, note its only 10 and we will oversample by 2x, gathering 20 nearest neighbors. +<2> The number of results to return from the KNN search. This will do an approximate KNN search with 50 candidates +per HNSW graph and use the quantized vectors, returning the 20 most similar vectors +according to the quantized score. Additionally, since this is the top-level `knn` object, the global top 20 results +will from all shards will be gathered before rescoring. Combining with `rescore`, this is oversampling by `2x`, meaning +gathering 20 nearest neighbors according to quantized scoring and rescoring with higher fidelity float vectors. +<3> The number of results to rescore, if you want to rescore all results, set this to the same value as `k` +<4> The script to rescore the results. Script score will interact directly with the originally provided float32 vector. +<5> The weight of the original query, here we simply throw away the original score +<6> The weight of the rescore query, here we only use the rescore query + +The second way is to score per shard with the <> and <>. Generally, this means that there will be more rescoring per shard, but this +can increase overall recall at the cost of compute. + +[source,console] +-------------------------------------------------- +POST /my-index/_search +{ + "size": 10, <1> + "query": { + "script_score": { + "query": { + "knn": { <2> + "query_vector": [0.04283529, 0.85670587, -0.51402352, 0], + "field": "my_int4_vector", + "num_candidates": 20 <3> + } + }, + "script": { + "source": "(dotProduct(params.queryVector, 'my_int4_vector') + 1.0)", <4> + "params": { + "queryVector": [0.04283529, 0.85670587, -0.51402352, 0] + } + } + } + } +} +-------------------------------------------------- +// TEST[skip: setup not provided] +<1> The number of results to return +<2> The `knn` query to perform the initial search, this is executed per-shard +<3> The number of candidates to use for the initial approximate `knn` search. This will search using the quantized vectors +and return the top 20 candidates per shard to then be scored +<4> The script to score the results. Script score will interact directly with the originally provided float32 vector. diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml new file mode 100644 index 0000000000000..188c155e4a836 --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml @@ -0,0 +1,160 @@ +setup: + - requires: + cluster_features: "mapper.vectors.bbq" + reason: 'kNN float to better-binary quantization is required' + - do: + indices.create: + index: bbq_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + dims: 64 + index: true + similarity: l2_norm + index_options: + type: bbq_hnsw + another_vector: + type: dense_vector + dims: 64 + index: true + similarity: l2_norm + index_options: + type: bbq_hnsw + + - do: + index: + index: bbq_hnsw + id: "1" + body: + name: cow.jpg + vector: [300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_hnsw + + - do: + index: + index: bbq_hnsw + id: "2" + body: + name: moose.jpg + vector: [100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0] + another_vector: [50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_hnsw + + - do: + index: + index: bbq_hnsw + id: "3" + body: + name: rabbit.jpg + vector: [111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_hnsw + + - do: + indices.forcemerge: + index: bbq_hnsw + max_num_segments: 1 +--- +"Test knn search": + - do: + search: + index: bbq_hnsw + body: + knn: + field: vector + query_vector: [ 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0] + k: 3 + num_candidates: 3 + + # Depending on how things are distributed, docs 2 and 3 might be swapped + # here we verify that are last hit is always the worst one + - match: { hits.hits.2._id: "1" } + +--- +"Test bad quantization parameters": + - do: + catch: bad_request + indices.create: + index: bad_bbq_hnsw + body: + mappings: + properties: + vector: + type: dense_vector + dims: 64 + element_type: byte + index: true + index_options: + type: bbq_hnsw + + - do: + catch: bad_request + indices.create: + index: bad_bbq_hnsw + body: + mappings: + properties: + vector: + type: dense_vector + dims: 64 + index: false + index_options: + type: bbq_hnsw +--- +"Test few dimensions fail indexing": + - do: + catch: bad_request + indices.create: + index: bad_bbq_hnsw + body: + mappings: + properties: + vector: + type: dense_vector + dims: 42 + index: true + index_options: + type: bbq_hnsw + + - do: + indices.create: + index: dynamic_dim_bbq_hnsw + body: + mappings: + properties: + vector: + type: dense_vector + index: true + similarity: l2_norm + index_options: + type: bbq_hnsw + + - do: + catch: bad_request + index: + index: dynamic_dim_bbq_hnsw + body: + vector: [1.0, 2.0, 3.0, 4.0, 5.0] + + - do: + index: + index: dynamic_dim_bbq_hnsw + body: + vector: [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0] diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml new file mode 100644 index 0000000000000..ed7a8dd5df65d --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml @@ -0,0 +1,165 @@ +setup: + - requires: + cluster_features: "mapper.vectors.bbq" + reason: 'kNN float to better-binary quantization is required' + - do: + indices.create: + index: bbq_flat + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + dims: 64 + index: true + similarity: l2_norm + index_options: + type: bbq_flat + another_vector: + type: dense_vector + dims: 64 + index: true + similarity: l2_norm + index_options: + type: bbq_flat + + - do: + index: + index: bbq_flat + id: "1" + body: + name: cow.jpg + vector: [300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0, 230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0, 130.0, 115.0, -1.02, 15.555, -100.0] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_flat + + - do: + index: + index: bbq_flat + id: "2" + body: + name: moose.jpg + vector: [100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0, -0.5, 100.0, -13, 14.8, -156.0] + another_vector: [50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120, -0.5, 50.0, -1, 1, 120] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_flat + + - do: + index: + index: bbq_flat + id: "3" + body: + name: rabbit.jpg + vector: [111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0, 0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0, -0.5, 11.0, 0, 12, 111.0] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_flat + + - do: + indices.forcemerge: + index: bbq_flat + max_num_segments: 1 +--- +"Test knn search": + - do: + search: + index: bbq_flat + body: + knn: + field: vector + query_vector: [ 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0, -0.5, 90.0, -10, 14.8, -156.0] + k: 3 + num_candidates: 3 + + # Depending on how things are distributed, docs 2 and 3 might be swapped + # here we verify that are last hit is always the worst one + - match: { hits.hits.2._id: "1" } +--- +"Test bad parameters": + - do: + catch: bad_request + indices.create: + index: bad_bbq_flat + body: + mappings: + properties: + vector: + type: dense_vector + dims: 64 + index: true + index_options: + type: bbq_flat + m: 42 + + - do: + catch: bad_request + indices.create: + index: bad_bbq_flat + body: + mappings: + properties: + vector: + type: dense_vector + dims: 64 + element_type: byte + index: true + index_options: + type: bbq_flat +--- +"Test few dimensions fail indexing": + # verify index creation fails + - do: + catch: bad_request + indices.create: + index: bad_bbq_flat + body: + mappings: + properties: + vector: + type: dense_vector + dims: 42 + index: true + similarity: l2_norm + index_options: + type: bbq_flat + + # verify dynamic dimension fails + - do: + indices.create: + index: dynamic_dim_bbq_flat + body: + mappings: + properties: + vector: + type: dense_vector + index: true + similarity: l2_norm + index_options: + type: bbq_flat + + # verify index fails for odd dim vector + - do: + catch: bad_request + index: + index: dynamic_dim_bbq_flat + body: + vector: [1.0, 2.0, 3.0, 4.0, 5.0] + + # verify that we can index an even dim vector after the odd dim vector failure + - do: + index: + index: dynamic_dim_bbq_flat + body: + vector: [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0] diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 2f03fda72ee2c..6971fa0e2f294 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -7,7 +7,6 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import org.elasticsearch.index.codec.tsdb.ES87TSDBDocValuesFormat; import org.elasticsearch.plugins.internal.RestExtension; /** The Elasticsearch Server Module. */ @@ -445,14 +444,16 @@ org.elasticsearch.index.codec.bloomfilter.ES85BloomFilterPostingsFormat, org.elasticsearch.index.codec.bloomfilter.ES87BloomFilterPostingsFormat, org.elasticsearch.index.codec.postings.ES812PostingsFormat; - provides org.apache.lucene.codecs.DocValuesFormat with ES87TSDBDocValuesFormat; + provides org.apache.lucene.codecs.DocValuesFormat with org.elasticsearch.index.codec.tsdb.ES87TSDBDocValuesFormat; provides org.apache.lucene.codecs.KnnVectorsFormat with org.elasticsearch.index.codec.vectors.ES813FlatVectorFormat, org.elasticsearch.index.codec.vectors.ES813Int8FlatVectorFormat, org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.ES815HnswBitVectorsFormat, - org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat; + org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat, + org.elasticsearch.index.codec.vectors.ES816BinaryQuantizedVectorsFormat, + org.elasticsearch.index.codec.vectors.ES816HnswBinaryQuantizedVectorsFormat; provides org.apache.lucene.codecs.Codec with diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java index 3d2acb533e26d..5201e57179cc7 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java @@ -27,6 +27,14 @@ public class BQVectorUtils { private static final float EPSILON = 1e-4f; + public static double sqrtNewtonRaphson(double x, double curr, double prev) { + return (curr == prev) ? curr : sqrtNewtonRaphson(x, 0.5 * (curr + x / curr), curr); + } + + public static double constSqrt(double x) { + return x >= 0 && Double.isInfinite(x) == false ? sqrtNewtonRaphson(x, x, 0) : Double.NaN; + } + public static boolean isUnitVector(float[] v) { double l1norm = VectorUtil.dotProduct(v, v); return Math.abs(l1norm - 1.0d) <= EPSILON; diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorer.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorer.java index 78fa282709098..f4d22edc6dfdb 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorer.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorer.java @@ -153,6 +153,7 @@ public static class BinarizedRandomVectorScorer extends RandomVectorScorer.Abstr private final VectorSimilarityFunction similarityFunction; private final float sqrtDimensions; + private final float maxX1; public BinarizedRandomVectorScorer( BinaryQueryVector queryVectors, @@ -164,24 +165,12 @@ public BinarizedRandomVectorScorer( this.targetVectors = targetVectors; this.similarityFunction = similarityFunction; // FIXME: precompute this once? - this.sqrtDimensions = (float) Utils.constSqrt(targetVectors.dimension()); - } - - // FIXME: utils class; pull this out - private static class Utils { - public static double sqrtNewtonRaphson(double x, double curr, double prev) { - return (curr == prev) ? curr : sqrtNewtonRaphson(x, 0.5 * (curr + x / curr), curr); - } - - public static double constSqrt(double x) { - return x >= 0 && Double.isInfinite(x) == false ? sqrtNewtonRaphson(x, x, 0) : Double.NaN; - } + this.sqrtDimensions = targetVectors.sqrtDimensions(); + this.maxX1 = targetVectors.maxX1(); } @Override public float score(int targetOrd) throws IOException { - // FIXME: implement fastscan in the future? - byte[] quantizedQuery = queryVector.vector(); int quantizedSum = queryVector.factors().quantizedSum(); float lower = queryVector.factors().lower(); @@ -218,17 +207,13 @@ public float score(int targetOrd) throws IOException { } assert Float.isFinite(dist); - // TODO: this is useful for mandatory rescoring by accounting for bias - // However, for just oversampling & rescoring, it isn't strictly useful. - // We should consider utilizing this bias in the future to determine which vectors need to - // be rescored - // float ooqSqr = (float) Math.pow(ooq, 2); - // float errorBound = (float) (normVmC * normOC * (maxX1 * Math.sqrt((1 - ooqSqr) / ooqSqr))); - // float score = dist - errorBound; + float ooqSqr = (float) Math.pow(ooq, 2); + float errorBound = (float) (vmC * normOC * (maxX1 * Math.sqrt((1 - ooqSqr) / ooqSqr))); + float score = Float.isFinite(errorBound) ? dist - errorBound : dist; if (similarityFunction == MAXIMUM_INNER_PRODUCT) { - return VectorUtil.scaleMaxInnerProductScore(dist); + return VectorUtil.scaleMaxInnerProductScore(score); } - return Math.max((1f + dist) / 2f, 0); + return Math.max((1f + score) / 2f, 0); } private float euclideanScore( @@ -256,17 +241,13 @@ private float euclideanScore( long qcDist = ESVectorUtil.ipByteBinByte(quantizedQuery, binaryCode); float score = sqrX + distanceToCentroid + factorPPC * lower + (qcDist * 2 - quantizedSum) * factorIP * width; - // TODO: this is useful for mandatory rescoring by accounting for bias - // However, for just oversampling & rescoring, it isn't strictly useful. - // We should consider utilizing this bias in the future to determine which vectors need to - // be rescored - // float projectionDist = (float) Math.sqrt(xX0 * xX0 - targetDistToC * targetDistToC); - // float error = 2.0f * maxX1 * projectionDist; - // float y = (float) Math.sqrt(distanceToCentroid); - // float errorBound = y * error; - // if (Float.isFinite(errorBound)) { - // score = dist + errorBound; - // } + float projectionDist = (float) Math.sqrt(xX0 * xX0 - targetDistToC * targetDistToC); + float error = 2.0f * maxX1 * projectionDist; + float y = (float) Math.sqrt(distanceToCentroid); + float errorBound = y * error; + if (Float.isFinite(errorBound)) { + score = score + errorBound; + } return Math.max(1 / (1f + score), 0); } } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapBinarizedVectorValues.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapBinarizedVectorValues.java index 2a3c3aca60e54..628480e273b34 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapBinarizedVectorValues.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/OffHeapBinarizedVectorValues.java @@ -34,6 +34,7 @@ import java.nio.ByteBuffer; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.elasticsearch.index.codec.vectors.BQVectorUtils.constSqrt; /** Binarized vector values loaded from off-heap */ public abstract class OffHeapBinarizedVectorValues extends BinarizedByteVectorValues implements RandomAccessBinarizedByteVectorValues { @@ -53,6 +54,9 @@ public abstract class OffHeapBinarizedVectorValues extends BinarizedByteVectorVa protected final BinaryQuantizer binaryQuantizer; protected final float[] centroid; protected final float centroidDp; + private final int discretizedDimensions; + private final float maxX1; + private final float sqrtDimensions; private final int correctionsCount; OffHeapBinarizedVectorValues( @@ -79,6 +83,9 @@ public abstract class OffHeapBinarizedVectorValues extends BinarizedByteVectorVa this.byteBuffer = ByteBuffer.allocate(numBytes); this.binaryValue = byteBuffer.array(); this.binaryQuantizer = quantizer; + this.discretizedDimensions = BQVectorUtils.discretize(dimension, 64); + this.sqrtDimensions = (float) constSqrt(dimension); + this.maxX1 = (float) (1.9 / constSqrt(discretizedDimensions - 1.0)); } @Override @@ -103,6 +110,21 @@ public byte[] vectorValue(int targetOrd) throws IOException { return binaryValue; } + @Override + public int discretizedDimensions() { + return discretizedDimensions; + } + + @Override + public float sqrtDimensions() { + return sqrtDimensions; + } + + @Override + public float maxX1() { + return maxX1; + } + @Override public float getCentroidDP() { return centroidDp; diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/RandomAccessBinarizedByteVectorValues.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/RandomAccessBinarizedByteVectorValues.java index 2417353373ba5..5163baf617c29 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/RandomAccessBinarizedByteVectorValues.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/RandomAccessBinarizedByteVectorValues.java @@ -24,6 +24,8 @@ import java.io.IOException; +import static org.elasticsearch.index.codec.vectors.BQVectorUtils.constSqrt; + /** * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10 */ @@ -54,6 +56,18 @@ public interface RandomAccessBinarizedByteVectorValues extends RandomAccessVecto */ BinaryQuantizer getQuantizer(); + default int discretizedDimensions() { + return BQVectorUtils.discretize(dimension(), 64); + } + + default float sqrtDimensions() { + return (float) constSqrt(dimension()); + } + + default float maxX1() { + return (float) (1.9 / constSqrt(discretizedDimensions() - 1.0)); + } + /** * @return coarse grained centroids for the vectors */ diff --git a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java index 7d17079c20d89..866bb3eaacfdf 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java @@ -9,6 +9,7 @@ package org.elasticsearch.index.mapper; +import org.elasticsearch.common.util.set.Sets; import org.elasticsearch.features.FeatureSpecification; import org.elasticsearch.features.NodeFeature; import org.elasticsearch.index.IndexSettings; @@ -23,7 +24,8 @@ public class MapperFeatures implements FeatureSpecification { @Override public Set getFeatures() { - return Set.of( + Set features = Set.of( + BWC_WORKAROUND_9_0, IgnoredSourceFieldMapper.TRACK_IGNORED_SOURCE, PassThroughObjectMapper.PASS_THROUGH_PRIORITY, RangeFieldMapper.NULL_VALUES_OFF_BY_ONE_FIX, @@ -48,6 +50,11 @@ public Set getFeatures() { TimeSeriesRoutingHashFieldMapper.TS_ROUTING_HASH_FIELD_PARSES_BYTES_REF, FlattenedFieldMapper.IGNORE_ABOVE_WITH_ARRAYS_SUPPORT ); + // BBQ is currently behind a feature flag for testing + if (DenseVectorFieldMapper.BBQ_FEATURE_FLAG.isEnabled()) { + return Sets.union(features, Set.of(DenseVectorFieldMapper.BBQ_FORMAT)); + } + return features; } @Override diff --git a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java index c3959bd442a1a..52ff7a3014d1d 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java @@ -36,6 +36,7 @@ import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.VectorUtil; import org.elasticsearch.common.ParsingException; +import org.elasticsearch.common.util.FeatureFlag; import org.elasticsearch.common.xcontent.support.XContentMapValues; import org.elasticsearch.features.NodeFeature; import org.elasticsearch.index.IndexVersion; @@ -45,6 +46,8 @@ import org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat; import org.elasticsearch.index.codec.vectors.ES815HnswBitVectorsFormat; +import org.elasticsearch.index.codec.vectors.ES816BinaryQuantizedVectorsFormat; +import org.elasticsearch.index.codec.vectors.ES816HnswBinaryQuantizedVectorsFormat; import org.elasticsearch.index.fielddata.FieldDataContext; import org.elasticsearch.index.fielddata.IndexFieldData; import org.elasticsearch.index.mapper.ArraySourceValueFetcher; @@ -98,6 +101,7 @@ public class DenseVectorFieldMapper extends FieldMapper { public static final String COSINE_MAGNITUDE_FIELD_SUFFIX = "._magnitude"; private static final float EPS = 1e-3f; + static final int BBQ_MIN_DIMS = 64; public static boolean isNotUnitVector(float magnitude) { return Math.abs(magnitude - 1.0f) > EPS; @@ -105,6 +109,8 @@ public static boolean isNotUnitVector(float magnitude) { public static final NodeFeature INT4_QUANTIZATION = new NodeFeature("mapper.vectors.int4_quantization"); public static final NodeFeature BIT_VECTORS = new NodeFeature("mapper.vectors.bit_vectors"); + public static final NodeFeature BBQ_FORMAT = new NodeFeature("mapper.vectors.bbq"); + public static final FeatureFlag BBQ_FEATURE_FLAG = new FeatureFlag("bbq_index_format"); public static final IndexVersion MAGNITUDE_STORED_INDEX_VERSION = IndexVersions.V_7_5_0; public static final IndexVersion INDEXED_BY_DEFAULT_INDEX_VERSION = IndexVersions.FIRST_DETACHED_INDEX_VERSION; @@ -1162,7 +1168,7 @@ final void validateElementType(ElementType elementType) { abstract boolean updatableTo(IndexOptions update); - public final void validateDimension(int dim) { + public void validateDimension(int dim) { if (type.supportsDimension(dim)) { return; } @@ -1342,6 +1348,50 @@ public boolean supportsElementType(ElementType elementType) { public boolean supportsDimension(int dims) { return dims % 2 == 0; } + }, + BBQ_HNSW("bbq_hnsw") { + @Override + public IndexOptions parseIndexOptions(String fieldName, Map indexOptionsMap) { + Object mNode = indexOptionsMap.remove("m"); + Object efConstructionNode = indexOptionsMap.remove("ef_construction"); + if (mNode == null) { + mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; + } + if (efConstructionNode == null) { + efConstructionNode = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; + } + int m = XContentMapValues.nodeIntegerValue(mNode); + int efConstruction = XContentMapValues.nodeIntegerValue(efConstructionNode); + MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap); + return new BBQHnswIndexOptions(m, efConstruction); + } + + @Override + public boolean supportsElementType(ElementType elementType) { + return elementType == ElementType.FLOAT; + } + + @Override + public boolean supportsDimension(int dims) { + return dims >= BBQ_MIN_DIMS; + } + }, + BBQ_FLAT("bbq_flat") { + @Override + public IndexOptions parseIndexOptions(String fieldName, Map indexOptionsMap) { + MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap); + return new BBQFlatIndexOptions(); + } + + @Override + public boolean supportsElementType(ElementType elementType) { + return elementType == ElementType.FLOAT; + } + + @Override + public boolean supportsDimension(int dims) { + return dims >= BBQ_MIN_DIMS; + } }; static Optional fromString(String type) { @@ -1707,6 +1757,102 @@ public String toString() { } } + static class BBQHnswIndexOptions extends IndexOptions { + private final int m; + private final int efConstruction; + + BBQHnswIndexOptions(int m, int efConstruction) { + super(VectorIndexType.BBQ_HNSW); + this.m = m; + this.efConstruction = efConstruction; + } + + @Override + KnnVectorsFormat getVectorsFormat(ElementType elementType) { + assert elementType == ElementType.FLOAT; + return new ES816HnswBinaryQuantizedVectorsFormat(m, efConstruction); + } + + @Override + boolean updatableTo(IndexOptions update) { + return update.type.equals(this.type); + } + + @Override + boolean doEquals(IndexOptions other) { + BBQHnswIndexOptions that = (BBQHnswIndexOptions) other; + return m == that.m && efConstruction == that.efConstruction; + } + + @Override + int doHashCode() { + return Objects.hash(m, efConstruction); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field("type", type); + builder.field("m", m); + builder.field("ef_construction", efConstruction); + builder.endObject(); + return builder; + } + + @Override + public void validateDimension(int dim) { + if (type.supportsDimension(dim)) { + return; + } + throw new IllegalArgumentException(type.name + " does not support dimensions fewer than " + BBQ_MIN_DIMS + "; provided=" + dim); + } + } + + static class BBQFlatIndexOptions extends IndexOptions { + private final int CLASS_NAME_HASH = this.getClass().getName().hashCode(); + + BBQFlatIndexOptions() { + super(VectorIndexType.BBQ_FLAT); + } + + @Override + KnnVectorsFormat getVectorsFormat(ElementType elementType) { + assert elementType == ElementType.FLOAT; + return new ES816BinaryQuantizedVectorsFormat(); + } + + @Override + boolean updatableTo(IndexOptions update) { + return update.type.equals(this.type); + } + + @Override + boolean doEquals(IndexOptions other) { + return other instanceof BBQFlatIndexOptions; + } + + @Override + int doHashCode() { + return CLASS_NAME_HASH; + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field("type", type); + builder.endObject(); + return builder; + } + + @Override + public void validateDimension(int dim) { + if (type.supportsDimension(dim)) { + return; + } + throw new IllegalArgumentException(type.name + " does not support dimensions fewer than " + BBQ_MIN_DIMS + "; provided=" + dim); + } + } + public static final TypeParser PARSER = new TypeParser( (n, c) -> new Builder(n, c.indexVersionCreated()), notInMultiFields(CONTENT_TYPE) @@ -2108,9 +2254,15 @@ private static IndexOptions parseIndexOptions(String fieldName, Object propNode) throw new MapperParsingException("[index_options] requires field [type] to be configured"); } String type = XContentMapValues.nodeStringValue(typeNode); - return VectorIndexType.fromString(type) - .orElseThrow(() -> new MapperParsingException("Unknown vector index options type [" + type + "] for field [" + fieldName + "]")) - .parseIndexOptions(fieldName, indexOptionsMap); + Optional vectorIndexType = VectorIndexType.fromString(type); + if (vectorIndexType.isEmpty()) { + throw new MapperParsingException("Unknown vector index options type [" + type + "] for field [" + fieldName + "]"); + } + VectorIndexType parsedType = vectorIndexType.get(); + if ((parsedType == VectorIndexType.BBQ_FLAT || parsedType == VectorIndexType.BBQ_HNSW) && BBQ_FEATURE_FLAG.isEnabled() == false) { + throw new MapperParsingException("Unknown vector index options type [" + type + "] for field [" + fieldName + "]"); + } + return parsedType.parseIndexOptions(fieldName, indexOptionsMap); } /** diff --git a/server/src/main/java/org/elasticsearch/index/store/LuceneFilesExtensions.java b/server/src/main/java/org/elasticsearch/index/store/LuceneFilesExtensions.java index 186aff230b8d0..387385ea2d6a4 100644 --- a/server/src/main/java/org/elasticsearch/index/store/LuceneFilesExtensions.java +++ b/server/src/main/java/org/elasticsearch/index/store/LuceneFilesExtensions.java @@ -81,7 +81,9 @@ public enum LuceneFilesExtensions { VEM("vem", "Vector Metadata", true, false), VEMF("vemf", "Flat Vector Metadata", true, false), VEMQ("vemq", "Scalar Quantized Vector Metadata", true, false), - VEQ("veq", "Scalar Quantized Vector Data", false, true); + VEQ("veq", "Scalar Quantized Vector Data", false, true), + VEMB("vemb", "Binarized Vector Metadata", true, false), + VEB("veb", "Binarized Vector Data", false, true); /** * Allow plugin developers of custom codecs to opt out of the assertion in {@link #fromExtension} diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorerTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorerTests.java index 4ac66a9f63a3f..04d4ef2079b99 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorerTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/ES816BinaryFlatVectorsScorerTests.java @@ -1741,6 +1741,6 @@ public int dimension() { similarityFunction ); - assertEquals(132.30249f, scorer.score(0), 0.0001f); + assertEquals(129.64046f, scorer.score(0), 0.0001f); } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java index 0b11129f0e519..442b1ae5e9e5f 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java @@ -63,6 +63,7 @@ import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.BBQ_FEATURE_FLAG; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; @@ -1227,13 +1228,18 @@ public void testInvalidParameters() { e.getMessage(), containsString("Failed to parse mapping: Mapping definition for [field] has unsupported parameters: [foo : {}]") ); - for (String quantizationKind : new String[] { "int4_hnsw", "int8_hnsw", "int8_flat", "int4_flat" }) { + List floatOnlyQuantizations = new ArrayList<>(Arrays.asList("int4_hnsw", "int8_hnsw", "int8_flat", "int4_flat")); + if (BBQ_FEATURE_FLAG.isEnabled()) { + floatOnlyQuantizations.add("bbq_hnsw"); + floatOnlyQuantizations.add("bbq_flat"); + } + for (String quantizationKind : floatOnlyQuantizations) { e = expectThrows( MapperParsingException.class, () -> createDocumentMapper( fieldMapping( b -> b.field("type", "dense_vector") - .field("dims", dims) + .field("dims", 64) .field("element_type", "byte") .field("similarity", "l2_norm") .field("index", true) @@ -1957,6 +1963,62 @@ public void testKnnQuantizedHNSWVectorsFormat() throws IOException { assertEquals(expectedString, knnVectorsFormat.toString()); } + public void testKnnBBQHNSWVectorsFormat() throws IOException { + assumeTrue("BBQ vectors are not supported in the current version", BBQ_FEATURE_FLAG.isEnabled()); + final int m = randomIntBetween(1, DEFAULT_MAX_CONN + 10); + final int efConstruction = randomIntBetween(1, DEFAULT_BEAM_WIDTH + 10); + final int dims = randomIntBetween(64, 4096); + MapperService mapperService = createMapperService(fieldMapping(b -> { + b.field("type", "dense_vector"); + b.field("dims", dims); + b.field("index", true); + b.field("similarity", "dot_product"); + b.startObject("index_options"); + b.field("type", "bbq_hnsw"); + b.field("m", m); + b.field("ef_construction", efConstruction); + b.endObject(); + })); + CodecService codecService = new CodecService(mapperService, BigArrays.NON_RECYCLING_INSTANCE); + Codec codec = codecService.codec("default"); + KnnVectorsFormat knnVectorsFormat; + if (CodecService.ZSTD_STORED_FIELDS_FEATURE_FLAG.isEnabled()) { + assertThat(codec, instanceOf(PerFieldMapperCodec.class)); + knnVectorsFormat = ((PerFieldMapperCodec) codec).getKnnVectorsFormatForField("field"); + } else { + if (codec instanceof CodecService.DeduplicateFieldInfosCodec deduplicateFieldInfosCodec) { + codec = deduplicateFieldInfosCodec.delegate(); + } + assertThat(codec, instanceOf(LegacyPerFieldMapperCodec.class)); + knnVectorsFormat = ((LegacyPerFieldMapperCodec) codec).getKnnVectorsFormatForField("field"); + } + String expectedString = "ES816HnswBinaryQuantizedVectorsFormat(name=ES816HnswBinaryQuantizedVectorsFormat, maxConn=" + + m + + ", beamWidth=" + + efConstruction + + ", flatVectorFormat=ES816BinaryQuantizedVectorsFormat(" + + "name=ES816BinaryQuantizedVectorsFormat, " + + "flatVectorScorer=ES816BinaryFlatVectorsScorer(nonQuantizedDelegate=DefaultFlatVectorScorer())))"; + assertEquals(expectedString, knnVectorsFormat.toString()); + } + + public void testInvalidVectorDimensionsBBQ() { + assumeTrue("BBQ vectors are not supported in the current version", BBQ_FEATURE_FLAG.isEnabled()); + for (String quantizedFlatFormat : new String[] { "bbq_hnsw", "bbq_flat" }) { + MapperParsingException e = expectThrows(MapperParsingException.class, () -> createDocumentMapper(fieldMapping(b -> { + b.field("type", "dense_vector"); + b.field("dims", randomIntBetween(1, 63)); + b.field("element_type", "float"); + b.field("index", true); + b.field("similarity", "dot_product"); + b.startObject("index_options"); + b.field("type", quantizedFlatFormat); + b.endObject(); + }))); + assertThat(e.getMessage(), containsString("does not support dimensions fewer than 64")); + } + } + public void testKnnHalfByteQuantizedHNSWVectorsFormat() throws IOException { final int m = randomIntBetween(1, DEFAULT_MAX_CONN + 10); final int efConstruction = randomIntBetween(1, DEFAULT_BEAM_WIDTH + 10); diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java index 23864777db961..6433cf2f1c0d4 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Set; +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.BBQ_MIN_DIMS; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.instanceOf; @@ -61,7 +62,9 @@ private DenseVectorFieldMapper.IndexOptions randomIndexOptionsAll() { ), new DenseVectorFieldMapper.FlatIndexOptions(), new DenseVectorFieldMapper.Int8FlatIndexOptions(randomFrom((Float) null, 0f, (float) randomDoubleBetween(0.9, 1.0, true))), - new DenseVectorFieldMapper.Int4FlatIndexOptions(randomFrom((Float) null, 0f, (float) randomDoubleBetween(0.9, 1.0, true))) + new DenseVectorFieldMapper.Int4FlatIndexOptions(randomFrom((Float) null, 0f, (float) randomDoubleBetween(0.9, 1.0, true))), + new DenseVectorFieldMapper.BBQHnswIndexOptions(randomIntBetween(1, 100), randomIntBetween(1, 10_000)), + new DenseVectorFieldMapper.BBQFlatIndexOptions() ); } @@ -70,7 +73,7 @@ private DenseVectorFieldType createFloatFieldType() { "f", IndexVersion.current(), DenseVectorFieldMapper.ElementType.FLOAT, - 6, + BBQ_MIN_DIMS, indexed, VectorSimilarity.COSINE, indexed ? randomIndexOptionsAll() : null, @@ -147,7 +150,7 @@ public void testFetchSourceValue() throws IOException { public void testCreateNestedKnnQuery() { BitSetProducer producer = context -> null; - int dims = randomIntBetween(2, 2048); + int dims = randomIntBetween(BBQ_MIN_DIMS, 2048); if (dims % 2 != 0) { dims++; } @@ -197,7 +200,7 @@ public void testCreateNestedKnnQuery() { } public void testExactKnnQuery() { - int dims = randomIntBetween(2, 2048); + int dims = randomIntBetween(BBQ_MIN_DIMS, 2048); if (dims % 2 != 0) { dims++; } @@ -260,15 +263,19 @@ public void testFloatCreateKnnQuery() { "f", IndexVersion.current(), DenseVectorFieldMapper.ElementType.FLOAT, - 4, + BBQ_MIN_DIMS, true, VectorSimilarity.DOT_PRODUCT, randomIndexOptionsAll(), Collections.emptyMap() ); + float[] queryVector = new float[BBQ_MIN_DIMS]; + for (int i = 0; i < BBQ_MIN_DIMS; i++) { + queryVector[i] = i; + } e = expectThrows( IllegalArgumentException.class, - () -> dotProductField.createKnnQuery(VectorData.fromFloats(new float[] { 0.3f, 0.1f, 1.0f, 0.0f }), 10, 10, null, null, null) + () -> dotProductField.createKnnQuery(VectorData.fromFloats(queryVector), 10, 10, null, null, null) ); assertThat(e.getMessage(), containsString("The [dot_product] similarity can only be used with unit-length vectors.")); @@ -276,7 +283,7 @@ public void testFloatCreateKnnQuery() { "f", IndexVersion.current(), DenseVectorFieldMapper.ElementType.FLOAT, - 4, + BBQ_MIN_DIMS, true, VectorSimilarity.COSINE, randomIndexOptionsAll(), @@ -284,7 +291,7 @@ public void testFloatCreateKnnQuery() { ); e = expectThrows( IllegalArgumentException.class, - () -> cosineField.createKnnQuery(VectorData.fromFloats(new float[] { 0.0f, 0.0f, 0.0f, 0.0f }), 10, 10, null, null, null) + () -> cosineField.createKnnQuery(VectorData.fromFloats(new float[BBQ_MIN_DIMS]), 10, 10, null, null, null) ); assertThat(e.getMessage(), containsString("The [cosine] similarity does not support vectors with zero magnitude.")); } From 8574546747760d7bc778abf26c0c7f07f4e194fe Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Mon, 14 Oct 2024 20:43:24 -0400 Subject: [PATCH 2/2] spotless --- .../main/java/org/elasticsearch/index/mapper/MapperFeatures.java | 1 - 1 file changed, 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java index 866bb3eaacfdf..c461242b9753b 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java @@ -25,7 +25,6 @@ public class MapperFeatures implements FeatureSpecification { @Override public Set getFeatures() { Set features = Set.of( - BWC_WORKAROUND_9_0, IgnoredSourceFieldMapper.TRACK_IGNORED_SOURCE, PassThroughObjectMapper.PASS_THROUGH_PRIORITY, RangeFieldMapper.NULL_VALUES_OFF_BY_ONE_FIX,