From 7649ff82e7c4df0bfd877b29929f02f4993b1db6 Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Wed, 9 Oct 2024 13:47:53 -0400 Subject: [PATCH 1/6] Adding new bbq index types behind a feature flag --- .../search.vectors/41_knn_search_bbq_hnsw.yml | 144 ++++++++++++++++ .../search.vectors/42_knn_search_bbq_flat.yml | 150 ++++++++++++++++ server/src/main/java/module-info.java | 7 +- .../index/mapper/MapperFeatures.java | 8 +- .../vectors/DenseVectorFieldMapper.java | 160 +++++++++++++++++- .../index/store/LuceneFilesExtensions.java | 4 +- .../vectors/DenseVectorFieldMapperTests.java | 66 +++++++- .../vectors/DenseVectorFieldTypeTests.java | 23 ++- 8 files changed, 543 insertions(+), 19 deletions(-) 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/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..d5d3d8f2ba95b --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml @@ -0,0 +1,144 @@ +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: 65 + index: true + similarity: l2_norm + index_options: + type: bbq_hnsw + another_vector: + type: dense_vector + dims: 65 + index: true + similarity: l2_norm + index_options: + type: bbq_hnsw + + - do: + index: + index: bbq_hnsw + id: "1" + body: + name: cow.jpg + vector: [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, 230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [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, 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: [-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, -0.5, 100.0, -13, 14.8, -156.0] + another_vector: [-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, -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: [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, 0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [-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, -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 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..1ef55d38b8d96 --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml @@ -0,0 +1,150 @@ +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: 65 + index: true + similarity: l2_norm + index_options: + type: bbq_flat + another_vector: + type: dense_vector + dims: 65 + index: true + similarity: l2_norm + index_options: + type: bbq_flat + + - do: + index: + index: bbq_flat + id: "1" + body: + name: cow.jpg + vector: [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, 230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [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, 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: [-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, -0.5, 100.0, -13, 14.8, -156.0] + another_vector: [-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, -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: [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, 0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [-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, -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 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 11965abf1dcd2..fc0c327e44a87 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. */ @@ -446,14 +445,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/mapper/MapperFeatures.java b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java index 4f90bd6e6f2c9..6b931e0f4adab 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; @@ -28,7 +29,7 @@ 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, @@ -54,6 +55,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 4adfe619ca4e1..04599c6f5e055 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; @@ -1167,7 +1173,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; } @@ -1347,6 +1353,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) { @@ -1712,6 +1762,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) @@ -2113,9 +2259,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/mapper/vectors/DenseVectorFieldMapperTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java index 8aede4940443c..d1b9465e9773a 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) @@ -1939,6 +1945,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 1f857c93cbe5483bc5815af0ba270e6ab781a214 Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Wed, 9 Oct 2024 16:28:11 -0400 Subject: [PATCH 2/6] adding tests, correcting scoring --- .../search.vectors/41_knn_search_bbq_hnsw.yml | 32 +++++++++--- .../search.vectors/42_knn_search_bbq_flat.yml | 31 +++++++++--- .../index/codec/vectors/BQVectorUtils.java | 8 +++ .../vectors/ES816BinaryFlatVectorsScorer.java | 49 ++++++------------- .../vectors/OffHeapBinarizedVectorValues.java | 22 +++++++++ ...RandomAccessBinarizedByteVectorValues.java | 14 ++++++ .../ES816BinaryFlatVectorsScorerTests.java | 2 +- 7 files changed, 107 insertions(+), 51 deletions(-) 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 index d5d3d8f2ba95b..188c155e4a836 100644 --- 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 @@ -15,14 +15,14 @@ setup: type: keyword vector: type: dense_vector - dims: 65 + dims: 64 index: true similarity: l2_norm index_options: type: bbq_hnsw another_vector: type: dense_vector - dims: 65 + dims: 64 index: true similarity: l2_norm index_options: @@ -34,8 +34,8 @@ setup: id: "1" body: name: cow.jpg - vector: [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, 230.0, 300.33, -34.8988, 15.555, -200.0] - another_vector: [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, 130.0, 115.0, -1.02, 15.555, -100.0] + 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: @@ -47,8 +47,8 @@ setup: id: "2" body: name: moose.jpg - vector: [-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, -0.5, 100.0, -13, 14.8, -156.0] - another_vector: [-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, -0.5, 50.0, -1, 1, 120] + 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: @@ -60,8 +60,8 @@ setup: id: "3" body: name: rabbit.jpg - vector: [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, 0.5, 111.3, -13.0, 14.8, -156.0] - another_vector: [-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, -0.5, 11.0, 0, 12, 111.0] + 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: @@ -72,6 +72,22 @@ setup: 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 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 index 1ef55d38b8d96..ed7a8dd5df65d 100644 --- 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 @@ -15,14 +15,14 @@ setup: type: keyword vector: type: dense_vector - dims: 65 + dims: 64 index: true similarity: l2_norm index_options: type: bbq_flat another_vector: type: dense_vector - dims: 65 + dims: 64 index: true similarity: l2_norm index_options: @@ -34,8 +34,8 @@ setup: id: "1" body: name: cow.jpg - vector: [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, 230.0, 300.33, -34.8988, 15.555, -200.0] - another_vector: [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, 130.0, 115.0, -1.02, 15.555, -100.0] + 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: @@ -47,8 +47,8 @@ setup: id: "2" body: name: moose.jpg - vector: [-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, -0.5, 100.0, -13, 14.8, -156.0] - another_vector: [-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, -0.5, 50.0, -1, 1, 120] + 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: @@ -60,8 +60,8 @@ setup: id: "3" body: name: rabbit.jpg - vector: [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, 0.5, 111.3, -13.0, 14.8, -156.0] - another_vector: [-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, -0.5, 11.0, 0, 12, 111.0] + 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: @@ -72,6 +72,21 @@ setup: 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 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/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); } } From 0c311aa93decbd326319caea936cc84398141355 Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Fri, 11 Oct 2024 10:51:38 -0400 Subject: [PATCH 3/6] adding docs --- .../mapping/types/dense-vector.asciidoc | 41 +++++++-- .../search-your-data/knn-search.asciidoc | 90 +++++++++++++++++++ 2 files changed, 126 insertions(+), 5 deletions(-) diff --git a/docs/reference/mapping/types/dense-vector.asciidoc b/docs/reference/mapping/types/dense-vector.asciidoc index 0cd9ee0578b70..12815b5cf6db0 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` - 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: `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 } -------------------------------------------------- +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 <>. +* `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`. +* `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..c9bf78984f0ab 100644 --- a/docs/reference/search/search-your-data/knn-search.asciidoc +++ b/docs/reference/search/search-your-data/knn-search.asciidoc @@ -1149,3 +1149,93 @@ 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 +<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. +<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. + +-------------------------------------------------- +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 30 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. From 2c6665300c2af7fd5abcbaa98eb6e8bb15a3d7af Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Fri, 11 Oct 2024 11:01:06 -0400 Subject: [PATCH 4/6] fixing docs --- docs/reference/search/search-your-data/knn-search.asciidoc | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/reference/search/search-your-data/knn-search.asciidoc b/docs/reference/search/search-your-data/knn-search.asciidoc index c9bf78984f0ab..e939fe565637b 100644 --- a/docs/reference/search/search-your-data/knn-search.asciidoc +++ b/docs/reference/search/search-your-data/knn-search.asciidoc @@ -1210,6 +1210,7 @@ will from all shards will be gathered before rescoring. 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 { From b33ee9b55c7d6fc3a6995ff57afd41085d7cbce8 Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Mon, 14 Oct 2024 12:35:10 -0400 Subject: [PATCH 5/6] addressing PR comments, adding experimental tags --- docs/reference/mapping/types/dense-vector.asciidoc | 10 +++++----- .../search/search-your-data/knn-search.asciidoc | 7 ++++--- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/docs/reference/mapping/types/dense-vector.asciidoc b/docs/reference/mapping/types/dense-vector.asciidoc index 12815b5cf6db0..44f90eded8632 100644 --- a/docs/reference/mapping/types/dense-vector.asciidoc +++ b/docs/reference/mapping/types/dense-vector.asciidoc @@ -121,7 +121,7 @@ The three following quantization strategies are supported: -- `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` - 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. +`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. -- When using a quantized format, you may want to oversample and rescore the results to improve accuracy. See <> for more information. @@ -134,7 +134,7 @@ This means disk usage will increase by ~25% for `int8`, ~12.5% for `int4`, and ~ NOTE: `int4` quantization requires an even number of vector dimensions. -NOTE: `bbq` quantization only supports vector dimensions that are greater than 64. +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: @@ -178,7 +178,7 @@ PUT my-byte-quantized-index } -------------------------------------------------- -Here is an example of how to create a binary quantized index: +experimental:[] Here is an example of how to create a binary quantized index: [source,console] -------------------------------------------------- @@ -327,7 +327,7 @@ by 4x at the cost of some accuracy. See <>. -* `bbq_hnsw` - This utilizes the https://arxiv.org/abs/1603.09320[HNSW algorithm] in addition to automatically binary +* 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. @@ -335,7 +335,7 @@ by 32x at the cost of accuracy. See < The number of results to return +<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. +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 @@ -1238,5 +1239,5 @@ POST /my-index/_search <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 30 candidates per shard to then be scored +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. From 70f3c44e35e7b41bd27ac01b856400435867a9a2 Mon Sep 17 00:00:00 2001 From: Benjamin Trent Date: Mon, 14 Oct 2024 14:57:31 -0400 Subject: [PATCH 6/6] Update docs/changelog/114439.yaml --- docs/changelog/114439.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 docs/changelog/114439.yaml 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: []