diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_byte_quantized_bfloat16.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_byte_quantized_bfloat16.yml new file mode 100644 index 0000000000000..92141ba8e220a --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_byte_quantized_bfloat16.yml @@ -0,0 +1,896 @@ +setup: + - requires: + cluster_features: [ "mapper.vectors.generic_vector_format" ] + reason: Needs generic vector support + - do: + indices.create: + index: hnsw_byte_quantized + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: int8_hnsw + another_vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: int8_hnsw + + - do: + index: + index: hnsw_byte_quantized + id: "1" + body: + name: cow.jpg + vector: [230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [130.0, 115.0, -1.02, 15.555, -100.0] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: hnsw_byte_quantized + + - do: + index: + index: hnsw_byte_quantized + id: "2" + body: + name: moose.jpg + vector: [-0.5, 100.0, -13, 14.8, -156.0] + another_vector: [-0.5, 50.0, -1, 1, 120] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: hnsw_byte_quantized + + - do: + index: + index: hnsw_byte_quantized + id: "3" + body: + name: rabbit.jpg + vector: [0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [-0.5, 11.0, 0, 12, 111.0] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: hnsw_byte_quantized + + - do: + indices.forcemerge: + index: hnsw_byte_quantized + max_num_segments: 1 + + - do: + indices.refresh: { } +--- +"kNN search only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.1.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8, -156.0], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12, 111.0], k: 2, num_candidates: 3} + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} +--- +"kNN search plus query": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + query: + term: + name: cow.jpg + + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.0.fields.name.0: "cow.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} + + - match: {hits.hits.2._id: "3"} + - match: {hits.hits.2.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search with query": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8, -156.0], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12, 111.0], k: 2, num_candidates: 3} + query: + term: + name: cow.jpg + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "1"} + - match: {hits.hits.1.fields.name.0: "cow.jpg"} + + - match: {hits.hits.2._id: "2"} + - match: {hits.hits.2.fields.name.0: "moose.jpg"} +--- +"kNN search with filter": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + filter: + term: + name: "rabbit.jpg" + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + filter: + - term: + name: "rabbit.jpg" + - term: + _id: 2 + + - match: {hits.total.value: 0} + +--- +"KNN Vector similarity search only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 10.3 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} +--- +"Vector similarity with filter only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 11 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + filter: {"term": {"name": "moose.jpg"}} + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + filter: {"term": {"name": "cow.jpg"}} + + - length: {hits.hits: 0} +--- +"Knn search with mip": + - do: + indices.create: + index: mip + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: max_inner_product + index_options: + type: int8_hnsw + + - do: + index: + index: mip + id: "1" + body: + name: cow.jpg + vector: [1, 2, 3, 4, 5] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: mip + id: "2" + body: + name: moose.jpg + vector: [1, 1, 1, 1, 1] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: mip + id: "3" + body: + name: rabbit.jpg + vector: [1, 2, 2, 2, 2] + + # We force merge into a single segment to make sure scores are more uniform + # Each segment can have a different quantization error, which can affect scores and mip is especially sensitive to this + - do: + indices.forcemerge: + index: mip + max_num_segments: 1 + + - do: + indices.refresh: {} + + - do: + search: + index: mip + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + query_vector: [1, 2, 3, 4, 5] + + + - length: {hits.hits: 3} + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.2._id: "2"} + + - do: + search: + index: mip + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + query_vector: [1, 2, 3, 4, 5] + filter: { "term": { "name": "moose.jpg" } } + + + + - length: {hits.hits: 1} + - match: {hits.hits.0._id: "2"} +--- +"Cosine similarity with indexed vector": + - skip: + features: "headers" + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "cosineSimilarity(params.query_vector, 'vector')" + params: + query_vector: [0.5, 111.3, -13.0, 14.8, -156.0] + + - match: {hits.total: 3} + + - match: {hits.hits.0._id: "3"} + - gte: {hits.hits.0._score: 0.999} + - lte: {hits.hits.0._score: 1.001} + + - match: {hits.hits.1._id: "2"} + - gte: {hits.hits.1._score: 0.998} + - lte: {hits.hits.1._score: 1.0} + + - match: {hits.hits.2._id: "1"} + - gte: {hits.hits.2._score: 0.78} + - lte: {hits.hits.2._score: 0.791} + +--- +# Won't be true for larger datasets, but this helps checking kNN vs rescoring vs exact search +"Vector rescoring has the same scoring as exact search for kNN section": + - requires: + reason: 'Quantized vector rescoring is required' + test_runner_features: [capabilities] + capabilities: + - method: GET + path: /_search + capabilities: [knn_quantized_vector_rescore_oversample] + - skip: + features: "headers" + + # Rescore + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: hnsw_byte_quantized + body: + size: 3 + query: + knn: + k: 3 + num_candidates: 3 + field: vector + query_vector: [0.5, 111.3, -13.0, 14.8, -156.0] + rescore_vector: + oversample: 1.5 + + # Get rescoring scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + # Exact knn via script score + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "1.0 / (1.0 + Math.pow(l2norm(params.query_vector, 'vector'), 2.0))" + params: + query_vector: [0.5, 111.3, -13.0, 14.8, -156.0] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } + +--- +"Test bad quantization parameters": + - do: + catch: bad_request + indices.create: + index: bad_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + dims: 5 + element_type: byte + index: true + index_options: + type: int8_hnsw + + - do: + catch: bad_request + indices.create: + index: bad_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + dims: 5 + index: false + index_options: + type: int8_hnsw +--- +"Test create, merge, and search cosine": + - do: + indices.create: + index: hnsw_byte_quantized_merge_cosine + body: + settings: + index: + number_of_shards: 1 + + - do: + indices.put_mapping: + index: hnsw_byte_quantized_merge_cosine + body: + properties: + embedding: + type: dense_vector + element_type: bfloat16 + similarity: cosine + index_options: + type: int8_hnsw + + - do: + index: + index: hnsw_byte_quantized_merge_cosine + id: "1" + body: + embedding: [1.0, 1.0, 1.0, 1.0] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_cosine + id: "2" + body: + embedding: [1.0, 1.0, 1.0, 2.0] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_cosine + id: "3" + body: + embedding: [1.0, 1.0, 1.0, 3.0] + + - do: + indices.forcemerge: + index: hnsw_byte_quantized_merge_cosine + max_num_segments: 1 + + - do: + indices.refresh: {} + + - do: + search: + index: hnsw_byte_quantized_merge_cosine + body: + size: 3 + query: + knn: + field: embedding + query_vector: [1.0, 1.0, 1.0, 1.0] + num_candidates: 10 + + - length: { hits.hits: 3 } + - match: { hits.hits.0._id: "1"} + - match: { hits.hits.1._id: "2"} + - match: { hits.hits.2._id: "3"} +--- +"Test create, merge, and search dot_product": + - requires: + cluster_features: "gte_v8.12.0" + reason: 'kNN float to byte quantization added in 8.12' + - do: + indices.create: + index: hnsw_byte_quantized_merge_dot_product + body: + settings: + index: + number_of_shards: 1 + + - do: + indices.put_mapping: + index: hnsw_byte_quantized_merge_dot_product + body: + properties: + embedding: + type: dense_vector + element_type: bfloat16 + similarity: dot_product + index_options: + type: int8_hnsw + + - do: + index: + index: hnsw_byte_quantized_merge_dot_product + id: "1" + body: + embedding: [0.6, 0.8] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_dot_product + id: "2" + body: + embedding: [0.8, 0.6] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_dot_product + id: "3" + body: + embedding: [-0.6, -0.8] + + - do: + indices.forcemerge: + index: hnsw_byte_quantized_merge_dot_product + max_num_segments: 1 + + - do: + indices.refresh: {} + + - do: + search: + index: hnsw_byte_quantized_merge_dot_product + body: + size: 3 + query: + knn: + field: embedding + query_vector: [0.6, 0.8] + num_candidates: 10 + + - length: { hits.hits: 3 } + - match: { hits.hits.0._id: "1"} + - match: { hits.hits.1._id: "2"} + - match: { hits.hits.2._id: "3"} +--- +"Test index configured rescore vector": + - requires: + cluster_features: ["mapper.dense_vector.rescore_vector"] + reason: Needs rescore_vector feature + - skip: + features: "headers" + - do: + indices.create: + index: int8_rescore_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + similarity: max_inner_product + index_options: + type: int8_hnsw + rescore_vector: + oversample: 1.5 + + - do: + bulk: + index: int8_rescore_hnsw + refresh: true + body: | + { "index": {"_id": "1"}} + { "vector": [230.0, 300.33, -34.8988, 15.555] } + { "index": {"_id": "2"}} + { "vector": [-0.5, 100.0, -13, 14.8] } + { "index": {"_id": "3"}} + { "vector": [0.5, 111.3, -13.0, 14.8] } + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_hnsw + body: + query: + script_score: + query: {match_all: {} } + script: + source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1" + params: + query_vector: [-0.5, 90.0, -10, 14.8] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } +--- +"Test index configured rescore vector updateable and settable to 0": + - requires: + cluster_features: ["mapper.dense_vector.rescore_zero_vector"] + reason: Needs rescore_zero_vector feature + + - do: + indices.create: + index: int8_rescore_0_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index_options: + type: int8_hnsw + rescore_vector: + oversample: 0 + + - do: + indices.create: + index: int8_rescore_update_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index_options: + type: int8_hnsw + rescore_vector: + oversample: 1 + + - do: + indices.put_mapping: + index: int8_rescore_update_hnsw + body: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index_options: + type: int8_hnsw + rescore_vector: + oversample: 0 + + - do: + indices.get_mapping: + index: int8_rescore_update_hnsw + + - match: { .int8_rescore_update_hnsw.mappings.properties.vector.index_options.rescore_vector.oversample: 0 } +--- +"Test index configured rescore vector score consistency": + - requires: + cluster_features: ["mapper.dense_vector.rescore_zero_vector"] + reason: Needs rescore_zero_vector feature + - skip: + features: "headers" + - do: + indices.create: + index: int8_rescore_zero_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + similarity: max_inner_product + dims: 4 + index_options: + type: int8_hnsw + rescore_vector: + oversample: 0 + + - do: + bulk: + index: int8_rescore_zero_hnsw + refresh: true + body: | + { "index": {"_id": "1"}} + { "vector": [230.0, 300.33, -34.8988, 15.555] } + { "index": {"_id": "2"}} + { "vector": [-0.5, 100.0, -13, 14.8] } + { "index": {"_id": "3"}} + { "vector": [0.5, 111.3, -13.0, 14.8] } + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_zero_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: raw_score0 } + - set: { hits.hits.1._score: raw_score1 } + - set: { hits.hits.2._score: raw_score2 } + + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_zero_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + rescore_vector: + oversample: 2 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: override_score0 } + - set: { hits.hits.1._score: override_score1 } + - set: { hits.hits.2._score: override_score2 } + + - do: + indices.put_mapping: + index: int8_rescore_zero_hnsw + body: + properties: + vector: + type: dense_vector + element_type: bfloat16 + similarity: max_inner_product + dims: 4 + index_options: + type: int8_hnsw + rescore_vector: + oversample: 2 + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_zero_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: default_rescore0 } + - set: { hits.hits.1._score: default_rescore1 } + - set: { hits.hits.2._score: default_rescore2 } + + - do: + indices.put_mapping: + index: int8_rescore_zero_hnsw + body: + properties: + vector: + type: dense_vector + element_type: bfloat16 + similarity: max_inner_product + dims: 4 + index_options: + type: int8_hnsw + rescore_vector: + oversample: 0 + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_zero_hnsw + body: + query: + script_score: + query: {match_all: {} } + script: + source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1" + params: + query_vector: [-0.5, 90.0, -10, 14.8] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $override_score0 } + - match: { hits.hits.0._score: $default_rescore0 } + - match: { hits.hits.1._score: $override_score1 } + - match: { hits.hits.1._score: $default_rescore1 } + - match: { hits.hits.2._score: $override_score2 } + - match: { hits.hits.2._score: $default_rescore2 } diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_half_byte_quantized_bfloat16.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_half_byte_quantized_bfloat16.yml new file mode 100644 index 0000000000000..8b9467f0d97e1 --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_half_byte_quantized_bfloat16.yml @@ -0,0 +1,933 @@ +setup: + - requires: + cluster_features: [ "mapper.vectors.generic_vector_format" ] + reason: Needs generic vector support + - do: + indices.create: + index: hnsw_byte_quantized + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + index: true + similarity: l2_norm + index_options: + type: int4_hnsw + another_vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + index: true + similarity: l2_norm + index_options: + type: int4_hnsw + + - do: + index: + index: hnsw_byte_quantized + id: "1" + body: + name: cow.jpg + vector: [230.0, 300.33, -34.8988, 15.555] + another_vector: [130.0, 115.0, -1.02, 15.555] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized + id: "2" + body: + name: moose.jpg + vector: [-0.5, 100.0, -13, 14.8] + another_vector: [-0.5, 50.0, -1, 1] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized + id: "3" + body: + name: rabbit.jpg + vector: [0.5, 111.3, -13.0, 14.8] + another_vector: [-0.5, 11.0, 0, 12] + + - do: + indices.forcemerge: + index: hnsw_byte_quantized + max_num_segments: 1 + + - do: + indices.refresh: {} + +--- +"kNN search only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.1.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12], k: 2, num_candidates: 3} + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} +--- +"kNN search plus query": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + query: + term: + name: + value: cow.jpg + boost: 1.5 + + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.0.fields.name.0: "cow.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} + + - match: {hits.hits.2._id: "3"} + - match: {hits.hits.2.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search with query": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12], k: 2, num_candidates: 3, boost: 2.0} + query: + term: + name: + value: cow.jpg + boost: 2.0 + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "1"} + - match: {hits.hits.1.fields.name.0: "cow.jpg"} + + - match: {hits.hits.2._id: "2"} + - match: {hits.hits.2.fields.name.0: "moose.jpg"} +--- +"kNN search with filter": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + filter: + term: + name: "rabbit.jpg" + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + filter: + - term: + name: "rabbit.jpg" + - term: + _id: 2 + + - match: {hits.total.value: 0} + +--- +"KNN Vector similarity search only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 17 + query_vector: [-0.5, 90.0, -10, 14.8] + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} +--- +"Vector similarity with filter only": + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 17 + query_vector: [-0.5, 90.0, -10, 14.8] + filter: {"term": {"name": "moose.jpg"}} + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - do: + search: + index: hnsw_byte_quantized + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8] + filter: {"term": {"name": "cow.jpg"}} + + - length: {hits.hits: 0} +--- +"Knn search with mip": + - do: + indices.create: + index: mip + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 6 + index: true + similarity: max_inner_product + index_options: + type: int4_hnsw + + - do: + index: + index: mip + id: "1" + body: + name: cow.jpg + vector: [1, 2, 3, 4, 5, 0] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: mip + id: "2" + body: + name: moose.jpg + vector: [1, 1, 1, 1, 1, 0] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: mip + id: "3" + body: + name: rabbit.jpg + vector: [1, 2, 2, 2, 2, 0] + + # We force merge into a single segment to make sure scores are more uniform + # Each segment can have a different quantization error, which can affect scores and mip is especially sensitive to this + - do: + indices.forcemerge: + index: mip + max_num_segments: 1 + + - do: + indices.refresh: {} + + - do: + search: + index: mip + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + query_vector: [1, 2, 3, 4, 5, 0] + + + - length: {hits.hits: 3} + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.2._id: "2"} + + - do: + search: + index: mip + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + query_vector: [1, 2, 3, 4, 5, 0] + filter: { "term": { "name": "moose.jpg" } } + + + + - length: {hits.hits: 1} + - match: {hits.hits.0._id: "2"} +--- +"Cosine similarity with indexed vector": + - skip: + features: "headers" + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "cosineSimilarity(params.query_vector, 'vector')" + params: + query_vector: [0.5, 111.3, -13.0, 14.8] + + - match: {hits.total: 3} + + - match: {hits.hits.0._id: "3"} + - gte: {hits.hits.0._score: 0.999} + - lte: {hits.hits.0._score: 1.001} + + - match: {hits.hits.1._id: "2"} + - gte: {hits.hits.1._score: 0.998} + - lte: {hits.hits.1._score: 1.0} + + - match: {hits.hits.2._id: "1"} + - gte: {hits.hits.2._score: 0.78} + - lte: {hits.hits.2._score: 0.80} +--- +"Test bad quantization parameters": + - do: + catch: bad_request + indices.create: + index: bad_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + dims: 6 + element_type: byte + index: true + index_options: + type: int4_hnsw + + - do: + catch: bad_request + indices.create: + index: bad_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + dims: 6 + index: false + index_options: + type: int4_hnsw +--- +"Test create, merge, and search cosine": + - do: + indices.create: + index: hnsw_byte_quantized_merge_cosine + body: + settings: + index: + number_of_shards: 1 + - do: + indices.put_mapping: + index: hnsw_byte_quantized_merge_cosine + body: + properties: + embedding: + type: dense_vector + element_type: bfloat16 + similarity: cosine + index_options: + type: int4_hnsw + + - do: + index: + index: hnsw_byte_quantized_merge_cosine + id: "1" + body: + embedding: [0.5, 0.5, 0.5, 0.5, 0.5, 1.0] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_cosine + id: "2" + body: + embedding: [0.0, 0.0, 0.0, 1.0, 1.0, 0.5] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_cosine + id: "3" + body: + embedding: [0.0, 0.0, 0.0, 0.0, 0.0, 10.5] + + - do: + indices.forcemerge: + index: hnsw_byte_quantized_merge_cosine + max_num_segments: 1 + + - do: + indices.refresh: {} + + - do: + search: + index: hnsw_byte_quantized_merge_cosine + body: + size: 3 + query: + knn: + field: embedding + query_vector: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + num_candidates: 10 + + - length: { hits.hits: 3 } + - match: { hits.hits.0._id: "1"} + - match: { hits.hits.1._id: "2"} + - match: { hits.hits.2._id: "3"} +--- +"Test create, merge, and search dot_product": + - do: + indices.create: + index: hnsw_byte_quantized_merge_dot_product + body: + settings: + index: + number_of_shards: 1 + - do: + indices.put_mapping: + index: hnsw_byte_quantized_merge_dot_product + body: + properties: + embedding: + type: dense_vector + element_type: bfloat16 + similarity: dot_product + index_options: + type: int4_hnsw + + - do: + index: + index: hnsw_byte_quantized_merge_dot_product + id: "1" + body: + embedding: [0.6, 0.8] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_dot_product + id: "2" + body: + embedding: [0.8, 0.6] + + # Flush in order to provoke a merge later + - do: + indices.flush: { } + + - do: + index: + index: hnsw_byte_quantized_merge_dot_product + id: "3" + body: + embedding: [-0.6, -0.8] + + - do: + indices.forcemerge: + index: hnsw_byte_quantized_merge_dot_product + max_num_segments: 1 + + - do: + indices.refresh: {} + + - do: + search: + index: hnsw_byte_quantized_merge_dot_product + body: + size: 3 + query: + knn: + field: embedding + query_vector: [0.6, 0.8] + num_candidates: 10 + + - length: { hits.hits: 3 } + - match: { hits.hits.0._id: "1"} + - match: { hits.hits.1._id: "2"} + - match: { hits.hits.2._id: "3"} +--- +"Vector rescoring has the same scoring as exact search for kNN section": + - requires: + reason: 'Quantized vector rescoring is required' + test_runner_features: [capabilities] + capabilities: + - method: GET + path: /_search + capabilities: [knn_quantized_vector_rescore_oversample] + - skip: + features: "headers" + + # Rescore + - do: + headers: + Content-Type: application/json + search: + index: hnsw_byte_quantized + rest_total_hits_as_int: true + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + rescore_vector: + oversample: 1.5 + + # Get rescoring scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + # Exact knn via script score + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "1.0 / (1.0 + Math.pow(l2norm(params.query_vector, 'vector'), 2.0))" + params: + query_vector: [-0.5, 90.0, -10, 14.8] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } + +--- +"Test odd dimensions fail indexing": + - do: + catch: bad_request + indices.create: + index: bad_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + index_options: + type: int4_hnsw + + - do: + indices.create: + index: dynamic_dim_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index: true + similarity: l2_norm + index_options: + type: int4_hnsw + + - do: + catch: bad_request + index: + index: dynamic_dim_hnsw_quantized + body: + vector: [1.0, 2.0, 3.0, 4.0, 5.0] + + - do: + index: + index: dynamic_dim_hnsw_quantized + body: + vector: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] +--- +"Test index configured rescore vector": + - requires: + cluster_features: ["mapper.dense_vector.rescore_vector"] + reason: Needs rescore_vector feature + - skip: + features: "headers" + - do: + indices.create: + index: int4_rescore_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + similarity: max_inner_product + index_options: + type: int4_hnsw + rescore_vector: + oversample: 1.5 + + - do: + bulk: + index: int4_rescore_hnsw + refresh: true + body: | + { "index": {"_id": "1"}} + { "vector": [230.0, 300.33, -34.8988, 15.555] } + { "index": {"_id": "2"}} + { "vector": [-0.5, 100.0, -13, 14.8] } + { "index": {"_id": "3"}} + { "vector": [0.5, 111.3, -13.0, 14.8] } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_hnsw + body: + query: + script_score: + query: {match_all: {} } + script: + source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1" + params: + query_vector: [-0.5, 90.0, -10, 14.8] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } +--- +"Test index configured rescore vector updateable and settable to 0": + - requires: + cluster_features: ["mapper.dense_vector.rescore_zero_vector"] + reason: Needs rescore_zero_vector feature + + - do: + indices.create: + index: int4_rescore_0_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index_options: + type: int4_hnsw + rescore_vector: + oversample: 0 + + - do: + indices.create: + index: int4_rescore_update_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index_options: + type: int4_hnsw + rescore_vector: + oversample: 1 + + - do: + indices.put_mapping: + index: int4_rescore_update_hnsw + body: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index_options: + type: int4_hnsw + rescore_vector: + oversample: 0 + + - do: + indices.get_mapping: + index: int4_rescore_update_hnsw + + - match: { .int4_rescore_update_hnsw.mappings.properties.vector.index_options.rescore_vector.oversample: 0 } +--- +"Test index configured rescore vector score consistency": + - requires: + cluster_features: ["mapper.dense_vector.rescore_zero_vector"] + reason: Needs rescore_zero_vector feature + - skip: + features: "headers" + - do: + indices.create: + index: int4_rescore_zero_hnsw + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + similarity: max_inner_product + dims: 4 + index_options: + type: int4_hnsw + rescore_vector: + oversample: 0 + + - do: + bulk: + index: int4_rescore_zero_hnsw + refresh: true + body: | + { "index": {"_id": "1"}} + { "vector": [230.0, 300.33, -34.8988, 15.555] } + { "index": {"_id": "2"}} + { "vector": [-0.5, 100.0, -13, 14.8] } + { "index": {"_id": "3"}} + { "vector": [0.5, 111.3, -13.0, 14.8] } + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_zero_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: raw_score0 } + - set: { hits.hits.1._score: raw_score1 } + - set: { hits.hits.2._score: raw_score2 } + + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_zero_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + rescore_vector: + oversample: 2 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: override_score0 } + - set: { hits.hits.1._score: override_score1 } + - set: { hits.hits.2._score: override_score2 } + + - do: + indices.put_mapping: + index: int4_rescore_zero_hnsw + body: + properties: + vector: + type: dense_vector + element_type: bfloat16 + similarity: max_inner_product + dims: 4 + index_options: + type: int4_hnsw + rescore_vector: + oversample: 2 + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_zero_hnsw + body: + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: default_rescore0 } + - set: { hits.hits.1._score: default_rescore1 } + - set: { hits.hits.2._score: default_rescore2 } + + - do: + indices.put_mapping: + index: int4_rescore_zero_hnsw + body: + properties: + vector: + type: dense_vector + element_type: bfloat16 + similarity: max_inner_product + dims: 4 + index_options: + type: int4_hnsw + rescore_vector: + oversample: 0 + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_zero_hnsw + body: + query: + script_score: + query: {match_all: {} } + script: + source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1" + params: + query_vector: [-0.5, 90.0, -10, 14.8] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $override_score0 } + - match: { hits.hits.0._score: $default_rescore0 } + - match: { hits.hits.1._score: $override_score1 } + - match: { hits.hits.1._score: $default_rescore1 } + - match: { hits.hits.2._score: $override_score2 } + - match: { hits.hits.2._score: $default_rescore2 } diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_flat_bfloat16.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_flat_bfloat16.yml new file mode 100644 index 0000000000000..4d4d20f9dcb4e --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_flat_bfloat16.yml @@ -0,0 +1,432 @@ +setup: + - requires: + cluster_features: [ "mapper.vectors.generic_vector_format" ] + reason: Needs generic vector support + - do: + indices.create: + index: flat + body: + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: flat + another_vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: flat + + - do: + index: + index: flat + id: "1" + body: + name: cow.jpg + vector: [230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [130.0, 115.0, -1.02, 15.555, -100.0] + + - do: + index: + index: flat + id: "2" + body: + name: moose.jpg + vector: [-0.5, 100.0, -13, 14.8, -156.0] + another_vector: [-0.5, 50.0, -1, 1, 120] + + - do: + index: + index: flat + id: "3" + body: + name: rabbit.jpg + vector: [0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [-0.5, 11.0, 0, 12, 111.0] + + - do: + indices.refresh: {} + +--- +"kNN search only": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.1.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search only": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8, -156.0], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12, 111.0], k: 2, num_candidates: 3} + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} +--- +"kNN search plus query": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + query: + term: + name: cow.jpg + + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.0.fields.name.0: "cow.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} + + - match: {hits.hits.2._id: "3"} + - match: {hits.hits.2.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search with query": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8, -156.0], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12, 111.0], k: 2, num_candidates: 3} + query: + term: + name: cow.jpg + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "1"} + - match: {hits.hits.1.fields.name.0: "cow.jpg"} + + - match: {hits.hits.2._id: "2"} + - match: {hits.hits.2.fields.name.0: "moose.jpg"} +--- +"kNN search with filter": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + filter: + term: + name: "rabbit.jpg" + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + filter: + - term: + name: "rabbit.jpg" + - term: + _id: 2 + + - match: {hits.total.value: 0} + +--- +"KNN Vector similarity search only": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 11 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} +--- +"Vector similarity with filter only": + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 11 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + filter: {"term": {"name": "moose.jpg"}} + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - do: + search: + index: flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + filter: {"term": {"name": "cow.jpg"}} + + - length: {hits.hits: 0} +--- +"Cosine similarity with indexed vector": + - skip: + features: "headers" + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "cosineSimilarity(params.query_vector, 'vector')" + params: + query_vector: [0.5, 111.3, -13.0, 14.8, -156.0] + + - match: {hits.total: 3} + + - match: {hits.hits.0._id: "3"} + - gte: {hits.hits.0._score: 0.999} + - lte: {hits.hits.0._score: 1.001} + + - match: {hits.hits.1._id: "2"} + - gte: {hits.hits.1._score: 0.998} + - lte: {hits.hits.1._score: 1.0} + + - match: {hits.hits.2._id: "1"} + - gte: {hits.hits.2._score: 0.78} + - lte: {hits.hits.2._score: 0.791} +--- +"Vector rescoring has no effect for non-quantized vectors and provides same results as non-rescored knn": + - requires: + reason: 'Quantized vector rescoring is required' + test_runner_features: [capabilities] + capabilities: + - method: GET + path: /_search + capabilities: [knn_quantized_vector_rescore_oversample] + - skip: + features: "headers" + + # Non-rescored knn + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 3 + num_candidates: 3 + + # Get scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - set: { hits.hits.0._score: knn_score0 } + - set: { hits.hits.1._score: knn_score1 } + - set: { hits.hits.2._score: knn_score2 } + + # Rescored knn + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 3 + num_candidates: 3 + rescore_vector: + oversample: 1.5 + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $knn_score0 } + - match: { hits.hits.1._score: $knn_score1 } + - match: { hits.hits.2._score: $knn_score2 } +--- +"Test bad parameters": + - do: + catch: bad_request + indices.create: + index: bad_flat + body: + mappings: + properties: + vector: + type: dense_vector + dims: 5 + index: true + index_options: + type: flat + m: 42 +--- +"Nested flat search": + - do: + indices.create: + index: flat_nested + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + nested: + type: nested + properties: + paragraph_id: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: flat + - do: + index: + index: flat_nested + id: "1" + body: + name: cow.jpg + nested: + - paragraph_id: 0 + vector: [230.0, 300.33, -34.8988, 15.555, -200.0] + - paragraph_id: 1 + vector: [240.0, 300, -3, 1, -20] + + - do: + index: + index: flat_nested + id: "2" + body: + name: moose.jpg + nested: + - paragraph_id: 0 + vector: [-0.5, 100.0, -13, 14.8, -156.0] + - paragraph_id: 2 + vector: [0, 100.0, 0, 14.8, -156.0] + - paragraph_id: 3 + vector: [0, 1.0, 0, 1.8, -15.0] + + - do: + index: + index: flat_nested + id: "3" + body: + name: rabbit.jpg + nested: + - paragraph_id: 0 + vector: [0.5, 111.3, -13.0, 14.8, -156.0] + + - do: + indices.refresh: {} + - do: + search: + index: flat_nested + body: + fields: [ "name" ] + query: + nested: + path: nested + query: + knn: + field: nested.vector + query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ] + num_candidates: 3 + k: 2 + + - match: { hits.hits.0._id: "2" } + - match: { hits.hits.0.fields.name.0: "moose.jpg" } + + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.1.fields.name.0: "rabbit.jpg" } + + - do: + search: + index: flat_nested + body: + fields: [ "name" ] + knn: + field: nested.vector + query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ] + num_candidates: 3 + k: 2 + + - match: { hits.hits.0._id: "2" } + - match: { hits.hits.0.fields.name.0: "moose.jpg" } + + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.1.fields.name.0: "rabbit.jpg" } diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int4_flat_bfloat16.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int4_flat_bfloat16.yml new file mode 100644 index 0000000000000..e60ea690664e5 --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int4_flat_bfloat16.yml @@ -0,0 +1,619 @@ +setup: + - requires: + cluster_features: [ "mapper.vectors.generic_vector_format" ] + reason: Needs generic vector support + - do: + indices.create: + index: int4_flat + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + index: true + similarity: l2_norm + index_options: + type: int4_flat + another_vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + index: true + similarity: l2_norm + index_options: + type: int4_flat + + - do: + index: + index: int4_flat + id: "1" + body: + name: cow.jpg + vector: [230.0, 300.33, -34.8988, 15.555] + another_vector: [130.0, 115.0, -1.02, 15.555] + # Flush in order to provoke a merge later & ensure replicas have same doc order + - do: + indices.flush: { } + - do: + index: + index: int4_flat + id: "2" + body: + name: moose.jpg + vector: [-0.5, 100.0, -13, 14.8] + another_vector: [-0.5, 50.0, -1, 1] + # Flush in order to provoke a merge later & ensure replicas have same doc order + - do: + indices.flush: { } + - do: + index: + index: int4_flat + id: "3" + body: + name: rabbit.jpg + vector: [0.5, 111.3, -13.0, 14.8] + another_vector: [-0.5, 11.0, 0, 12] + + - do: + indices.flush: { } + + # For added test reliability, pending the resolution of https://github.com/elastic/elasticsearch/issues/109416. + - do: + indices.forcemerge: + max_num_segments: 1 + index: int4_flat + - do: + indices.refresh: {} +--- +"kNN search only": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.1.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search only": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12], k: 2, num_candidates: 3} + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} +--- +"kNN search plus query": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + query: + term: + name: + value: cow.jpg + boost: 1.5 + + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.0.fields.name.0: "cow.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} + + - match: {hits.hits.2._id: "3"} + - match: {hits.hits.2.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search with query": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12], k: 2, num_candidates: 3, boost: 2.0} + query: + term: + name: + value: cow.jpg + boost: 2.0 + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "1"} + - match: {hits.hits.1.fields.name.0: "cow.jpg"} + + - match: {hits.hits.2._id: "2"} + - match: {hits.hits.2.fields.name.0: "moose.jpg"} +--- +"kNN search with filter": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + filter: + term: + name: "rabbit.jpg" + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 2 + num_candidates: 3 + filter: + - term: + name: "rabbit.jpg" + - term: + _id: 2 + + - match: {hits.total.value: 0} + +--- +"KNN Vector similarity search only": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + # Set high allowed similarity, reduce once we can update underlying quantization algo + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8] + + - is_true: hits.hits.0 + + #- match: {hits.hits.0._id: "2"} + #- match: {hits.hits.0.fields.name.0: "moose.jpg"} +--- +"Vector similarity with filter only": + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + # Set high allowed similarity, reduce once we can update underlying quantization algo + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8] + filter: {"term": {"name": "moose.jpg"}} + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - do: + search: + index: int4_flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8] + filter: {"term": {"name": "cow.jpg"}} + + - length: {hits.hits: 0} +--- +"Cosine similarity with indexed vector": + - skip: + features: "headers" + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "cosineSimilarity(params.query_vector, 'vector')" + params: + query_vector: [0.5, 111.3, -13.0, 14.8] + + - match: {hits.total: 3} + + - match: {hits.hits.0._id: "3"} + - gte: {hits.hits.0._score: 0.999} + - lte: {hits.hits.0._score: 1.001} + + - match: {hits.hits.1._id: "2"} + - gte: {hits.hits.1._score: 0.998} + - lte: {hits.hits.1._score: 1.0} + + - match: {hits.hits.2._id: "1"} + - gte: {hits.hits.2._score: 0.78} + - lte: {hits.hits.2._score: 0.80} +--- +"Test bad parameters": + - do: + catch: bad_request + indices.create: + index: bad_int4_flat + body: + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 6 + index: true + index_options: + type: int4_flat + m: 42 + + - do: + catch: bad_request + indices.create: + index: bad_int4_flat + body: + mappings: + properties: + vector: + type: dense_vector + dims: 6 + element_type: byte + index: true + index_options: + type: int4_flat +--- +"Test odd dimensions fail indexing": + # verify index creation fails + - do: + catch: bad_request + indices.create: + index: bad_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: int4_flat + + # verify dynamic dimension fails + - do: + indices.create: + index: dynamic_dim_hnsw_quantized + body: + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + index: true + similarity: l2_norm + index_options: + type: int4_hnsw + + # verify index fails for odd dim vector + - do: + catch: bad_request + index: + index: dynamic_dim_hnsw_quantized + 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_hnsw_quantized + body: + vector: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] +--- +"Vector rescoring has the same scoring as exact search for kNN section": + - requires: + reason: 'Quantized vector rescoring is required' + test_runner_features: [capabilities] + capabilities: + - method: GET + path: /_search + capabilities: [knn_quantized_vector_rescore_oversample] + - skip: + features: "headers" + + # Rescore + - do: + headers: + Content-Type: application/json + search: + index: int4_flat + rest_total_hits_as_int: true + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8] + k: 3 + num_candidates: 3 + rescore_vector: + oversample: 1.5 + + # Get rescoring scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + # Exact knn via script score + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "1.0 / (1.0 + Math.pow(l2norm(params.query_vector, 'vector'), 2.0))" + params: + query_vector: [-0.5, 90.0, -10, 14.8] + + # Get rescoring scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } +--- +"Test index configured rescore vector": + - requires: + cluster_features: ["mapper.dense_vector.rescore_vector"] + reason: Needs rescore_vector feature + - skip: + features: "headers" + - do: + indices.create: + index: int4_rescore_flat + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: int4_flat + rescore_vector: + oversample: 1.5 + + - do: + bulk: + index: int4_rescore_flat + refresh: true + body: | + { "index": {"_id": "1"}} + { "vector": [0.077, 0.32 , -0.205, 0.63 , 0.032, 0.201, 0.167, -0.313, 0.176, 0.531, -0.375, 0.334, -0.046, 0.078, -0.349, 0.272, 0.307, -0.083, 0.504, 0.255, -0.404, 0.289, -0.226, -0.132, -0.216, 0.49 , 0.039, 0.507, -0.307, 0.107, 0.09 , -0.265, -0.285, 0.336, -0.272, 0.369, -0.282, 0.086, -0.132, 0.475, -0.224, 0.203, 0.439, 0.064, 0.246, -0.396, 0.297, 0.242, -0.028, 0.321, -0.022, -0.009, -0.001 , 0.031, -0.533, 0.45, -0.683, 1.331, 0.194, -0.157, -0.1 , -0.279, -0.098, -0.176] } + { "index": {"_id": "2"}} + { "vector": [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348, -0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048, 0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438, -0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138, -0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429, -0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166, 0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569, -0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013] } + { "index": {"_id": "3"}} + { "vector": [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348, -0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048, 0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438, -0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138, -0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429, -0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166, 0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569, -0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013] } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_flat + body: + knn: + field: vector + query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393, + 0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015, + 0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259, + -0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 , + -0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232, + -0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034, + -0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582, + -0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int4_rescore_flat + body: + query: + script_score: + query: {match_all: {} } + script: + source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1" + params: + query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393, + 0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015, + 0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259, + -0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 , + -0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232, + -0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034, + -0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582, + -0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } + +--- +"no default oversample value": + - requires: + cluster_features: ["mapper.dense_vector.default_oversample_value_for_bbq"] + reason: "Needs default_oversample_value_for_bbq feature" + - do: + indices.get_mapping: + index: int4_flat + + - not_exists: int4_flat.mappings.properties.vector.index_options.rescore_vector +--- +"Nested flat search": + - do: + indices.create: + index: int4_flat_nested + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + nested: + type: nested + properties: + paragraph_id: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 4 + index: true + similarity: l2_norm + index_options: + type: int4_flat + - do: + index: + index: int4_flat_nested + id: "1" + body: + name: cow.jpg + nested: + - paragraph_id: 0 + vector: [230.0, 300.33, -34.8988, 15.555 ] + - paragraph_id: 1 + vector: [240.0, 300, -3, 1 ] + + - do: + index: + index: int4_flat_nested + id: "2" + body: + name: moose.jpg + nested: + - paragraph_id: 0 + vector: [-0.5, 100.0, -13, 14.8] + - paragraph_id: 2 + vector: [0, 100.0, 0, 14.8] + - paragraph_id: 3 + vector: [0, 1.0, 0, 1.8] + + - do: + index: + index: int4_flat_nested + id: "3" + body: + name: rabbit.jpg + nested: + - paragraph_id: 0 + vector: [0.5, 111.3, -13.0, 14.8] + + - do: + indices.refresh: {} + - do: + search: + index: int4_flat_nested + body: + fields: [ "name" ] + query: + nested: + path: nested + query: + knn: + field: nested.vector + query_vector: [ -0.5, 90.0, -10, 14.8] + num_candidates: 3 + k: 2 + + - match: { hits.hits.0._id: "2" } + - match: { hits.hits.0.fields.name.0: "moose.jpg" } + + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.1.fields.name.0: "rabbit.jpg" } + + - do: + search: + index: int4_flat_nested + body: + fields: [ "name" ] + knn: + field: nested.vector + query_vector: [ -0.5, 90.0, -10, 14.8] + num_candidates: 3 + k: 2 + + - match: { hits.hits.0._id: "2" } + - match: { hits.hits.0.fields.name.0: "moose.jpg" } + + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.1.fields.name.0: "rabbit.jpg" } + diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int8_flat_bfloat16.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int8_flat_bfloat16.yml new file mode 100644 index 0000000000000..065593282b751 --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int8_flat_bfloat16.yml @@ -0,0 +1,555 @@ +setup: + - requires: + cluster_features: [ "mapper.vectors.generic_vector_format" ] + reason: Needs generic vector support + - do: + indices.create: + index: int8_flat + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: int8_flat + confidence_interval: 0.9 + another_vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: int8_flat + confidence_interval: 0.9 + + - do: + index: + index: int8_flat + id: "1" + body: + name: cow.jpg + vector: [230.0, 300.33, -34.8988, 15.555, -200.0] + another_vector: [130.0, 115.0, -1.02, 15.555, -100.0] + + - do: + index: + index: int8_flat + id: "2" + body: + name: moose.jpg + vector: [-0.5, 100.0, -13, 14.8, -156.0] + another_vector: [-0.5, 50.0, -1, 1, 120] + + - do: + index: + index: int8_flat + id: "3" + body: + name: rabbit.jpg + vector: [0.5, 111.3, -13.0, 14.8, -156.0] + another_vector: [-0.5, 11.0, 0, 12, 111.0] + + - do: + indices.refresh: {} + +--- +"kNN search only": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - match: {hits.hits.1._id: "3"} + - match: {hits.hits.1.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search only": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8, -156.0], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12, 111.0], k: 2, num_candidates: 3} + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} +--- +"kNN search plus query": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + query: + term: + name: cow.jpg + + - match: {hits.hits.0._id: "1"} + - match: {hits.hits.0.fields.name.0: "cow.jpg"} + + - match: {hits.hits.1._id: "2"} + - match: {hits.hits.1.fields.name.0: "moose.jpg"} + + - match: {hits.hits.2._id: "3"} + - match: {hits.hits.2.fields.name.0: "rabbit.jpg"} +--- +"kNN multi-field search with query": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + - {field: vector, query_vector: [-0.5, 90.0, -10, 14.8, -156.0], k: 2, num_candidates: 3} + - {field: another_vector, query_vector: [-0.5, 11.0, 0, 12, 111.0], k: 2, num_candidates: 3} + query: + term: + name: cow.jpg + + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - match: {hits.hits.1._id: "1"} + - match: {hits.hits.1.fields.name.0: "cow.jpg"} + + - match: {hits.hits.2._id: "2"} + - match: {hits.hits.2.fields.name.0: "moose.jpg"} +--- +"kNN search with filter": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + filter: + term: + name: "rabbit.jpg" + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "3"} + - match: {hits.hits.0.fields.name.0: "rabbit.jpg"} + + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 2 + num_candidates: 3 + filter: + - term: + name: "rabbit.jpg" + - term: + _id: 2 + + - match: {hits.total.value: 0} + +--- +"KNN Vector similarity search only": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 10.3 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} +--- +"Vector similarity with filter only": + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 11 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + filter: {"term": {"name": "moose.jpg"}} + + - length: {hits.hits: 1} + + - match: {hits.hits.0._id: "2"} + - match: {hits.hits.0.fields.name.0: "moose.jpg"} + + - do: + search: + index: int8_flat + body: + fields: [ "name" ] + knn: + num_candidates: 3 + k: 3 + field: vector + similarity: 110 + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + filter: {"term": {"name": "cow.jpg"}} + + - length: {hits.hits: 0} +--- +"Cosine similarity with indexed vector": + - skip: + features: "headers" + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "cosineSimilarity(params.query_vector, 'vector')" + params: + query_vector: [0.5, 111.3, -13.0, 14.8, -156.0] + + - match: {hits.total: 3} + + - match: {hits.hits.0._id: "3"} + - gte: {hits.hits.0._score: 0.999} + - lte: {hits.hits.0._score: 1.001} + + - match: {hits.hits.1._id: "2"} + - gte: {hits.hits.1._score: 0.998} + - lte: {hits.hits.1._score: 1.0} + + - match: {hits.hits.2._id: "1"} + - gte: {hits.hits.2._score: 0.78} + - lte: {hits.hits.2._score: 0.791} +--- +"Vector rescoring has the same scoring as exact search for kNN section": + - requires: + reason: 'Quantized vector rescoring is required' + test_runner_features: [capabilities] + capabilities: + - method: GET + path: /_search + capabilities: [knn_quantized_vector_rescore_oversample] + - skip: + features: "headers" + + # Rescore + - do: + headers: + Content-Type: application/json + search: + index: int8_flat + rest_total_hits_as_int: true + body: + fields: [ "name" ] + knn: + field: vector + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + k: 3 + num_candidates: 3 + rescore_vector: + oversample: 1.5 + + # Get rescoring scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + body: + query: + script_score: + query: {match_all: {} } + script: + source: "1.0 / (1.0 + Math.pow(l2norm(params.query_vector, 'vector'), 2.0))" + params: + query_vector: [-0.5, 90.0, -10, 14.8, -156.0] + + # Get rescoring scores - hit ordering may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } +--- +"Test bad parameters": + - do: + catch: bad_request + indices.create: + index: bad_int8_flat + body: + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + index_options: + type: int8_flat + m: 42 + + - do: + catch: bad_request + indices.create: + index: bad_int8_flat + body: + mappings: + properties: + vector: + type: dense_vector + dims: 5 + element_type: byte + index: true + index_options: + type: int8_flat +--- +"Test index configured rescore vector": + - requires: + cluster_features: ["mapper.dense_vector.rescore_vector"] + reason: Needs rescore_vector feature + - skip: + features: "headers" + - do: + indices.create: + index: int8_rescore_flat + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + element_type: bfloat16 + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: int8_flat + rescore_vector: + oversample: 1.5 + + - do: + bulk: + index: int8_rescore_flat + refresh: true + body: | + { "index": {"_id": "1"}} + { "vector": [0.077, 0.32 , -0.205, 0.63 , 0.032, 0.201, 0.167, -0.313, 0.176, 0.531, -0.375, 0.334, -0.046, 0.078, -0.349, 0.272, 0.307, -0.083, 0.504, 0.255, -0.404, 0.289, -0.226, -0.132, -0.216, 0.49 , 0.039, 0.507, -0.307, 0.107, 0.09 , -0.265, -0.285, 0.336, -0.272, 0.369, -0.282, 0.086, -0.132, 0.475, -0.224, 0.203, 0.439, 0.064, 0.246, -0.396, 0.297, 0.242, -0.028, 0.321, -0.022, -0.009, -0.001 , 0.031, -0.533, 0.45, -0.683, 1.331, 0.194, -0.157, -0.1 , -0.279, -0.098, -0.176] } + { "index": {"_id": "2"}} + { "vector": [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348, -0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048, 0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438, -0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138, -0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429, -0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166, 0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569, -0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013] } + { "index": {"_id": "3"}} + { "vector": [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348, -0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048, 0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438, -0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138, -0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429, -0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166, 0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569, -0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013] } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_flat + body: + knn: + field: vector + query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393, + 0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015, + 0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259, + -0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 , + -0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232, + -0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034, + -0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582, + -0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158] + k: 3 + num_candidates: 3 + + - match: { hits.total: 3 } + - set: { hits.hits.0._score: rescore_score0 } + - set: { hits.hits.1._score: rescore_score1 } + - set: { hits.hits.2._score: rescore_score2 } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: int8_rescore_flat + body: + query: + script_score: + query: {match_all: {} } + script: + source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1" + params: + query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393, + 0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015, + 0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259, + -0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 , + -0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232, + -0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034, + -0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582, + -0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158] + + # Compare scores as hit IDs may change depending on how things are distributed + - match: { hits.total: 3 } + - match: { hits.hits.0._score: $rescore_score0 } + - match: { hits.hits.1._score: $rescore_score1 } + - match: { hits.hits.2._score: $rescore_score2 } + +--- +"no default oversample value": + - requires: + cluster_features: ["mapper.dense_vector.default_oversample_value_for_bbq"] + reason: "Needs default_oversample_value_for_bbq feature" + - do: + indices.get_mapping: + index: int8_flat + + - not_exists: int8_flat.mappings.properties.vector.index_options.rescore_vector +--- +"Nested flat search": + - do: + indices.create: + index: int8_flat_nested + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + name: + type: keyword + nested: + type: nested + properties: + paragraph_id: + type: keyword + vector: + type: dense_vector + element_type: bfloat16 + dims: 5 + index: true + similarity: l2_norm + index_options: + type: int8_flat + - do: + index: + index: int8_flat_nested + id: "1" + body: + name: cow.jpg + nested: + - paragraph_id: 0 + vector: [230.0, 300.33, -34.8988, 15.555, -200.0] + - paragraph_id: 1 + vector: [240.0, 300, -3, 1, -20] + + - do: + index: + index: int8_flat_nested + id: "2" + body: + name: moose.jpg + nested: + - paragraph_id: 0 + vector: [-0.5, 100.0, -13, 14.8, -156.0] + - paragraph_id: 2 + vector: [0, 100.0, 0, 14.8, -156.0] + - paragraph_id: 3 + vector: [0, 1.0, 0, 1.8, -15.0] + + - do: + index: + index: int8_flat_nested + id: "3" + body: + name: rabbit.jpg + nested: + - paragraph_id: 0 + vector: [0.5, 111.3, -13.0, 14.8, -156.0] + + - do: + indices.refresh: {} + - do: + search: + index: int8_flat_nested + body: + fields: [ "name" ] + query: + nested: + path: nested + query: + knn: + field: nested.vector + query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ] + num_candidates: 3 + k: 2 + + - match: { hits.hits.0._id: "2" } + - match: { hits.hits.0.fields.name.0: "moose.jpg" } + + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.1.fields.name.0: "rabbit.jpg" } + + - do: + search: + index: int8_flat_nested + body: + fields: [ "name" ] + knn: + field: nested.vector + query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ] + num_candidates: 3 + k: 2 + + - match: { hits.hits.0._id: "2" } + - match: { hits.hits.0.fields.name.0: "moose.jpg" } + + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.1.fields.name.0: "rabbit.jpg" } + diff --git a/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java b/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java index 578182b7e5aec..a85a435408402 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java @@ -38,6 +38,7 @@ import java.util.Map; import java.util.OptionalLong; import java.util.stream.IntStream; +import java.util.stream.Stream; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount; @@ -76,7 +77,7 @@ protected boolean useDirectIO(String name, IOContext context, OptionalLong fileL @ParametersFactory public static Iterable parameters() { - return List.of(new Object[] { "bbq_hnsw" }, new Object[] { "bbq_disk" }); + return Stream.of("int4_hnsw", "int8_hnsw", "bbq_hnsw", "bbq_disk").map(s -> new Object[] { s }).toList(); } public DirectIOIT(String type) { diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 9c5d11e1cf9e1..1abeb15209ce7 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -465,8 +465,11 @@ org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat, org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat, - org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat, + org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat, org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; provides org.apache.lucene.codecs.Codec diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java new file mode 100644 index 0000000000000..172ad11be497b --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93FlatVectorFormat.java @@ -0,0 +1,125 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; + +import java.io.IOException; +import java.util.Map; + +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; + +public class ES93FlatVectorFormat extends KnnVectorsFormat { + + static final String NAME = "ES93FlatVectorFormat"; + + private final FlatVectorsFormat format; + + /** + * Sole constructor + */ + public ES93FlatVectorFormat() { + super(NAME); + format = new ES93GenericFlatVectorsFormat(); + } + + public ES93FlatVectorFormat(DenseVectorFieldMapper.ElementType elementType) { + super(NAME); + assert elementType != DenseVectorFieldMapper.ElementType.BIT : "ES815BitFlatVectorFormat should be used for bits"; + format = new ES93GenericFlatVectorsFormat(elementType, false); + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return format.fieldsWriter(state); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new ES93FlatVectorReader(format.fieldsReader(state)); + } + + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMS_COUNT; + } + + static class ES93FlatVectorReader extends KnnVectorsReader { + + private final FlatVectorsReader reader; + + ES93FlatVectorReader(FlatVectorsReader reader) { + this.reader = reader; + } + + @Override + public void checkIntegrity() throws IOException { + reader.checkIntegrity(); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return reader.getFloatVectorValues(field); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return reader.getByteVectorValues(field); + } + + @Override + public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) + throws IOException { + OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); + Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); + for (int i = 0; i < scorer.maxOrd(); i++) { + if (acceptedOrds == null || acceptedOrds.get(i)) { + collector.collect(i, scorer.score(i)); + collector.incVisitedCount(1); + } + } + assert collector.earlyTerminated() == false; + } + + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + return reader.getOffHeapByteSize(fieldInfo); + } + + @Override + public void close() throws IOException { + reader.close(); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java new file mode 100644 index 0000000000000..cd861954ddc02 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswScalarQuantizedVectorsFormat.java @@ -0,0 +1,168 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.elasticsearch.index.codec.vectors.AbstractHnswVectorsFormat; +import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; + +import java.io.IOException; +import java.util.concurrent.ExecutorService; + +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER; +import static org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat.DYNAMIC_CONFIDENCE_INTERVAL; + +public class ES93HnswScalarQuantizedVectorsFormat extends AbstractHnswVectorsFormat { + + static final String NAME = "ES93HnswScalarQuantizedVectorsFormat"; + private static final int ALLOWED_BITS = (1 << 7) | (1 << 4); + + /** The minimum confidence interval */ + private static final float MINIMUM_CONFIDENCE_INTERVAL = 0.9f; + + /** The maximum confidence interval */ + private static final float MAXIMUM_CONFIDENCE_INTERVAL = 1f; + + static final FlatVectorsScorer flatVectorScorer = new ES93ScalarQuantizedVectorsFormat.ESQuantizedFlatVectorsScorer( + new ScalarQuantizedVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()) + ); + + private final FlatVectorsFormat rawVectorFormat; + + /** + * Controls the confidence interval used to scalar quantize the vectors the default value is + * calculated as `1-1/(vector_dimensions + 1)` + */ + public final Float confidenceInterval; + + private final byte bits; + private final boolean compress; + + public ES93HnswScalarQuantizedVectorsFormat() { + super(NAME); + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(DenseVectorFieldMapper.ElementType.FLOAT, false); + this.confidenceInterval = null; + this.bits = 7; + this.compress = false; + } + + public ES93HnswScalarQuantizedVectorsFormat( + int maxConn, + int beamWidth, + DenseVectorFieldMapper.ElementType elementType, + Float confidenceInterval, + int bits, + boolean compress, + boolean useDirectIO + ) { + this(maxConn, beamWidth, elementType, confidenceInterval, bits, compress, useDirectIO, DEFAULT_NUM_MERGE_WORKER, null); + } + + public ES93HnswScalarQuantizedVectorsFormat( + int maxConn, + int beamWidth, + DenseVectorFieldMapper.ElementType elementType, + Float confidenceInterval, + int bits, + boolean compress, + boolean useDirectIO, + int numMergeWorkers, + ExecutorService mergeExec + ) { + super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec); + + if (confidenceInterval != null + && confidenceInterval != DYNAMIC_CONFIDENCE_INTERVAL + && (confidenceInterval < MINIMUM_CONFIDENCE_INTERVAL || confidenceInterval > MAXIMUM_CONFIDENCE_INTERVAL)) { + throw new IllegalArgumentException( + "confidenceInterval must be between " + + MINIMUM_CONFIDENCE_INTERVAL + + " and " + + MAXIMUM_CONFIDENCE_INTERVAL + + "; confidenceInterval=" + + confidenceInterval + ); + } + if (bits < 1 || bits > 8 || (ALLOWED_BITS & (1 << bits)) == 0) { + throw new IllegalArgumentException("bits must be one of: 4, 7; bits=" + bits); + } + assert elementType != DenseVectorFieldMapper.ElementType.BIT : "BIT should not be used with scalar quantization"; + + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, useDirectIO); + this.confidenceInterval = confidenceInterval; + this.bits = (byte) bits; + this.compress = compress; + } + + @Override + protected FlatVectorsFormat flatVectorsFormat() { + return rawVectorFormat; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return new Lucene99HnswVectorsWriter( + state, + maxConn, + beamWidth, + new Lucene99ScalarQuantizedVectorsWriter( + state, + confidenceInterval, + bits, + compress, + rawVectorFormat.fieldsWriter(state), + flatVectorScorer + ), + numMergeWorkers, + mergeExec + ); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new Lucene99HnswVectorsReader( + state, + new Lucene99ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer) + ); + } + + @Override + public String toString() { + return NAME + + "(name=" + + NAME + + ", maxConn=" + + maxConn + + ", beamWidth=" + + beamWidth + + ", confidenceInterval=" + + confidenceInterval + + ", bits=" + + bits + + ", compressed=" + + compress + + ", flatVectorScorer=" + + flatVectorScorer + + ", flatVectorFormat=" + + rawVectorFormat + + ")"; + } +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java new file mode 100644 index 0000000000000..e69cbbd331872 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93ScalarQuantizedVectorsFormat.java @@ -0,0 +1,271 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors.es93; + +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.quantization.QuantizedByteVectorValues; +import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; +import org.elasticsearch.simdvec.VectorScorerFactory; +import org.elasticsearch.simdvec.VectorSimilarityType; + +import java.io.IOException; +import java.util.Map; + +import static org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat.DYNAMIC_CONFIDENCE_INTERVAL; +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; + +public class ES93ScalarQuantizedVectorsFormat extends KnnVectorsFormat { + + static final String NAME = "ES93ScalarQuantizedVectorsFormat"; + private static final int ALLOWED_BITS = (1 << 7) | (1 << 4); + + static final FlatVectorsScorer flatVectorScorer = new ESQuantizedFlatVectorsScorer( + new ScalarQuantizedVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()) + ); + + /** The minimum confidence interval */ + private static final float MINIMUM_CONFIDENCE_INTERVAL = 0.9f; + + /** The maximum confidence interval */ + private static final float MAXIMUM_CONFIDENCE_INTERVAL = 1f; + + private final FlatVectorsFormat rawVectorFormat; + + /** + * Controls the confidence interval used to scalar quantize the vectors the default value is + * calculated as `1-1/(vector_dimensions + 1)` + */ + public final Float confidenceInterval; + + private final byte bits; + private final boolean compress; + + public ES93ScalarQuantizedVectorsFormat() { + this(DenseVectorFieldMapper.ElementType.FLOAT, null, 7, false); + } + + public ES93ScalarQuantizedVectorsFormat(DenseVectorFieldMapper.ElementType elementType) { + this(elementType, null, 7, false); + } + + public ES93ScalarQuantizedVectorsFormat( + DenseVectorFieldMapper.ElementType elementType, + Float confidenceInterval, + int bits, + boolean compress + ) { + super(NAME); + if (confidenceInterval != null + && confidenceInterval != DYNAMIC_CONFIDENCE_INTERVAL + && (confidenceInterval < MINIMUM_CONFIDENCE_INTERVAL || confidenceInterval > MAXIMUM_CONFIDENCE_INTERVAL)) { + throw new IllegalArgumentException( + "confidenceInterval must be between " + + MINIMUM_CONFIDENCE_INTERVAL + + " and " + + MAXIMUM_CONFIDENCE_INTERVAL + + "; confidenceInterval=" + + confidenceInterval + ); + } + if (bits < 1 || bits > 8 || (ALLOWED_BITS & (1 << bits)) == 0) { + throw new IllegalArgumentException("bits must be one of: 4, 7; bits=" + bits); + } + assert elementType != DenseVectorFieldMapper.ElementType.BIT : "BIT should not be used with scalar quantization"; + + this.rawVectorFormat = new ES93GenericFlatVectorsFormat(elementType, false); + this.confidenceInterval = confidenceInterval; + this.bits = (byte) bits; + this.compress = compress; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return new Lucene99ScalarQuantizedVectorsWriter( + state, + confidenceInterval, + bits, + compress, + rawVectorFormat.fieldsWriter(state), + flatVectorScorer + ); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new ES93FlatVectorReader( + new Lucene99ScalarQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), flatVectorScorer) + ); + } + + @Override + public int getMaxDimensions(String fieldName) { + return MAX_DIMS_COUNT; + } + + @Override + public String toString() { + return NAME + + "(name=" + + NAME + + ", confidenceInterval=" + + confidenceInterval + + ", bits=" + + bits + + ", compressed=" + + compress + + ", flatVectorScorer=" + + flatVectorScorer + + ", rawVectorFormat=" + + rawVectorFormat + + ")"; + } + + static class ES93FlatVectorReader extends KnnVectorsReader { + + private final FlatVectorsReader reader; + + ES93FlatVectorReader(FlatVectorsReader reader) { + this.reader = reader; + } + + @Override + public void checkIntegrity() throws IOException { + reader.checkIntegrity(); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return reader.getFloatVectorValues(field); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return reader.getByteVectorValues(field); + } + + @Override + public void search(String field, float[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + private void collectAllMatchingDocs(KnnCollector knnCollector, AcceptDocs acceptDocs, RandomVectorScorer scorer) + throws IOException { + OrdinalTranslatedKnnCollector collector = new OrdinalTranslatedKnnCollector(knnCollector, scorer::ordToDoc); + Bits acceptedOrds = scorer.getAcceptOrds(acceptDocs.bits()); + for (int i = 0; i < scorer.maxOrd(); i++) { + if (acceptedOrds == null || acceptedOrds.get(i)) { + collector.collect(i, scorer.score(i)); + collector.incVisitedCount(1); + } + } + assert collector.earlyTerminated() == false; + } + + @Override + public void search(String field, byte[] target, KnnCollector knnCollector, AcceptDocs acceptDocs) throws IOException { + collectAllMatchingDocs(knnCollector, acceptDocs, reader.getRandomVectorScorer(field, target)); + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + return reader.getOffHeapByteSize(fieldInfo); + } + + @Override + public void close() throws IOException { + reader.close(); + } + } + + static final class ESQuantizedFlatVectorsScorer implements FlatVectorsScorer { + + final FlatVectorsScorer delegate; + final VectorScorerFactory factory; + + ESQuantizedFlatVectorsScorer(FlatVectorsScorer delegate) { + this.delegate = delegate; + factory = VectorScorerFactory.instance().orElse(null); + } + + @Override + public String toString() { + return "ESQuantizedFlatVectorsScorer(" + "delegate=" + delegate + ", factory=" + factory + ')'; + } + + @Override + public RandomVectorScorerSupplier getRandomVectorScorerSupplier(VectorSimilarityFunction sim, KnnVectorValues values) + throws IOException { + if (values instanceof QuantizedByteVectorValues qValues && qValues.getSlice() != null) { + // TODO: optimize int4 quantization + if (qValues.getScalarQuantizer().getBits() != 7) { + return delegate.getRandomVectorScorerSupplier(sim, values); + } + if (factory != null) { + var scorer = factory.getInt7SQVectorScorerSupplier( + VectorSimilarityType.of(sim), + qValues.getSlice(), + qValues, + qValues.getScalarQuantizer().getConstantMultiplier() + ); + if (scorer.isPresent()) { + return scorer.get(); + } + } + } + return delegate.getRandomVectorScorerSupplier(sim, values); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(VectorSimilarityFunction sim, KnnVectorValues values, float[] query) + throws IOException { + if (values instanceof QuantizedByteVectorValues qValues && qValues.getSlice() != null) { + // TODO: optimize int4 quantization + if (qValues.getScalarQuantizer().getBits() != 7) { + return delegate.getRandomVectorScorer(sim, values, query); + } + if (factory != null) { + var scorer = factory.getInt7SQVectorScorer(sim, qValues, query); + if (scorer.isPresent()) { + return scorer.get(); + } + } + } + return delegate.getRandomVectorScorer(sim, values, query); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(VectorSimilarityFunction sim, KnnVectorValues values, byte[] query) + throws IOException { + return delegate.getRandomVectorScorer(sim, values, query); + } + } +} 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 aa02950e460b9..112ca683ef078 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 @@ -59,9 +59,12 @@ import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat; +import org.elasticsearch.index.codec.vectors.es93.ES93FlatVectorFormat; import org.elasticsearch.index.codec.vectors.es93.ES93GenericFlatVectorsFormat; import org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; +import org.elasticsearch.index.codec.vectors.es93.ES93HnswScalarQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.es93.ES93HnswVectorsFormat; +import org.elasticsearch.index.codec.vectors.es93.ES93ScalarQuantizedVectorsFormat; import org.elasticsearch.index.fielddata.FieldDataContext; import org.elasticsearch.index.fielddata.IndexFieldData; import org.elasticsearch.index.mapper.BlockLoader; @@ -406,6 +409,7 @@ private DenseVectorIndexOptions defaultIndexOptions(boolean defaultInt8Hnsw, boo Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN, Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH, null, + false, null ); } @@ -1512,14 +1516,14 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map