diff --git a/rest-api-spec/src/yamlRestTest/java/org/elasticsearch/test/rest/ClientYamlTestSuiteIT.java b/rest-api-spec/src/yamlRestTest/java/org/elasticsearch/test/rest/ClientYamlTestSuiteIT.java index a4499a07898a0..ee7cd9059e50d 100644 --- a/rest-api-spec/src/yamlRestTest/java/org/elasticsearch/test/rest/ClientYamlTestSuiteIT.java +++ b/rest-api-spec/src/yamlRestTest/java/org/elasticsearch/test/rest/ClientYamlTestSuiteIT.java @@ -38,6 +38,7 @@ public class ClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase { .feature(FeatureFlag.SUB_OBJECTS_AUTO_ENABLED) .feature(FeatureFlag.DOC_VALUES_SKIPPER) .feature(FeatureFlag.USE_LUCENE101_POSTINGS_FORMAT) + .feature(FeatureFlag.IVF_FORMAT) .build(); public ClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) { diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/46_knn_search_bbq_ivf.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/46_knn_search_bbq_ivf.yml new file mode 100644 index 0000000000000..86eda64701404 --- /dev/null +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/46_knn_search_bbq_ivf.yml @@ -0,0 +1,555 @@ +setup: + - requires: + cluster_features: ["mapper.ivf_format_cluster_feature"] + reason: Needs mapper.ivf_format_cluster_feature feature + - skip: + features: "headers" + - do: + indices.create: + index: bbq_ivf + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: bbq_ivf + + - do: + index: + index: bbq_ivf + id: "1" + body: + 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] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_ivf + + - do: + index: + index: bbq_ivf + id: "2" + body: + 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] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_ivf + + - do: + index: + index: bbq_ivf + id: "3" + body: + name: rabbit.jpg + vector: [0.139, 0.178, -0.117, 0.399, 0.014, -0.139, 0.347, -0.33 , + 0.139, 0.34 , -0.052, -0.052, -0.249, 0.327, -0.288, 0.049, + 0.464, 0.338, 0.516, 0.247, -0.104, 0.259, -0.209, -0.246, + -0.11 , 0.323, 0.091, 0.442, -0.254, 0.195, -0.109, -0.058, + -0.279, 0.402, -0.107, 0.308, -0.273, 0.019, 0.082, 0.399, + -0.658, -0.03 , 0.276, 0.041, 0.187, -0.331, 0.165, 0.017, + 0.171, -0.203, -0.198, 0.115, -0.007, 0.337, -0.444, 0.615, + -0.683, 1.331, 0.194, -0.157, -0.1 , -0.279, -0.098, -0.176] + # Flush in order to provoke a merge later + - do: + indices.flush: + index: bbq_ivf + + - do: + indices.forcemerge: + index: bbq_ivf + max_num_segments: 1 + + - do: + indices.refresh: { } +--- +"Test knn search": + - do: + search: + index: bbq_ivf + 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.hits.0._id: "1" } + - match: { hits.hits.1._id: "3" } + - match: { hits.hits.2._id: "2" } +--- +"Vector rescoring has same scoring as exact search for kNN section": + - skip: + features: "headers" + + # Rescore + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: bbq_ivf + 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 + 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: "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 } + +--- +"Test bad quantization parameters": + - do: + catch: bad_request + indices.create: + index: bad_bbq_ivf + body: + mappings: + properties: + vector: + type: dense_vector + dims: 64 + element_type: byte + index: true + index_options: + type: bbq_ivf + + - do: + catch: bad_request + indices.create: + index: bad_bbq_ivf + body: + mappings: + properties: + vector: + type: dense_vector + dims: 64 + index: false + index_options: + type: bbq_ivf +--- +"Test few dimensions fail indexing": + - do: + catch: bad_request + indices.create: + index: bad_bbq_ivf + body: + mappings: + properties: + vector: + type: dense_vector + dims: 42 + index: true + index_options: + type: bbq_ivf + + - do: + indices.create: + index: dynamic_dim_bbq_ivf + body: + mappings: + properties: + vector: + type: dense_vector + index: true + similarity: l2_norm + index_options: + type: bbq_ivf + + - do: + catch: bad_request + index: + index: dynamic_dim_bbq_ivf + body: + vector: [1.0, 2.0, 3.0, 4.0, 5.0] + + - do: + index: + index: dynamic_dim_bbq_ivf + body: + vector: [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0] +--- +"Test index configured rescore vector": + - skip: + features: "headers" + - do: + indices.create: + index: bbq_rescore_ivf + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: bbq_ivf + rescore_vector: + oversample: 1.5 + + - do: + bulk: + index: bbq_rescore_ivf + 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: bbq_rescore_ivf + 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: bbq_rescore_ivf + 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 } +--- +"Test index configured rescore vector updateable and settable to 0": + - do: + indices.create: + index: bbq_rescore_0_ivf + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + index_options: + type: bbq_ivf + rescore_vector: + oversample: 0 + + - do: + indices.create: + index: bbq_rescore_update_ivf + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + index_options: + type: bbq_ivf + rescore_vector: + oversample: 1 + + - do: + indices.put_mapping: + index: bbq_rescore_update_ivf + body: + properties: + vector: + type: dense_vector + index_options: + type: bbq_ivf + rescore_vector: + oversample: 0 + + - do: + indices.get_mapping: + index: bbq_rescore_update_ivf + + - match: { .bbq_rescore_update_ivf.mappings.properties.vector.index_options.rescore_vector.oversample: 0 } +--- +"Test index configured rescore vector score consistency": + - skip: + features: "headers" + - do: + indices.create: + index: bbq_rescore_zero_ivf + body: + settings: + index: + number_of_shards: 1 + mappings: + properties: + vector: + type: dense_vector + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: bbq_ivf + rescore_vector: + oversample: 0 + + - do: + bulk: + index: bbq_rescore_zero_ivf + 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: bbq_rescore_zero_ivf + 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 } + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: bbq_rescore_zero_ivf + 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 + 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: bbq_rescore_zero_ivf + body: + properties: + vector: + type: dense_vector + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: bbq_ivf + rescore_vector: + oversample: 2 + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: bbq_rescore_zero_ivf + 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: default_rescore0 } + - set: { hits.hits.1._score: default_rescore1 } + - set: { hits.hits.2._score: default_rescore2 } + + - do: + indices.put_mapping: + index: bbq_rescore_zero_ivf + body: + properties: + vector: + type: dense_vector + dims: 64 + index: true + similarity: max_inner_product + index_options: + type: bbq_ivf + rescore_vector: + oversample: 0 + + - do: + headers: + Content-Type: application/json + search: + rest_total_hits_as_int: true + index: bbq_rescore_zero_ivf + 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: $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 } + +--- +"default oversample value": + - do: + indices.get_mapping: + index: bbq_ivf + + - match: { bbq_ivf.mappings.properties.vector.index_options.rescore_vector.oversample: 3.0 } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormat.java index f124e978116e2..03e4a38b169de 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormat.java @@ -61,14 +61,26 @@ public class IVFVectorsFormat extends KnnVectorsFormat { FlatVectorScorerUtil.getLucene99FlatVectorsScorer() ); - private static final int DEFAULT_VECTORS_PER_CLUSTER = 1000; + // This dynamically sets the cluster probe based on the `k` requested and the number of clusters. + // useful when searching with 'efSearch' type parameters instead of requiring a specific nprobe. + public static final int DYNAMIC_NPROBE = -1; + public static final int DEFAULT_VECTORS_PER_CLUSTER = 384; + public static final int MIN_VECTORS_PER_CLUSTER = 64; + public static final int MAX_VECTORS_PER_CLUSTER = 1 << 16; // 65536 private final int vectorPerCluster; public IVFVectorsFormat(int vectorPerCluster) { super(NAME); - if (vectorPerCluster <= 0) { - throw new IllegalArgumentException("vectorPerCluster must be > 0"); + if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) { + throw new IllegalArgumentException( + "vectorsPerCluster must be between " + + MIN_VECTORS_PER_CLUSTER + + " and " + + MAX_VECTORS_PER_CLUSTER + + ", got: " + + vectorPerCluster + ); } this.vectorPerCluster = vectorPerCluster; } @@ -90,12 +102,12 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException @Override public int getMaxDimensions(String fieldName) { - return 1024; + return 4096; } @Override public String toString() { - return "IVFVectorFormat"; + return "IVFVectorsFormat(" + "vectorPerCluster=" + vectorPerCluster + ')'; } static IVFVectorsReader getIVFReader(KnnVectorsReader vectorsReader, String fieldName) { diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java index d8d68569d6159..12726836719be 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java @@ -38,6 +38,7 @@ import java.util.function.IntPredicate; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; +import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.DYNAMIC_NPROBE; /** * Reader for IVF vectors. This reader is used to read the IVF vectors from the index. @@ -226,17 +227,6 @@ public final ByteVectorValues getByteVectorValues(String field) throws IOExcepti return rawVectorsReader.getByteVectorValues(field); } - protected float[] getGlobalCentroid(FieldInfo info) { - if (info == null || info.getVectorEncoding().equals(VectorEncoding.BYTE)) { - return null; - } - FieldEntry entry = fields.get(info.number); - if (entry == null) { - return null; - } - return entry.globalCentroid(); - } - @Override public final void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException { final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field); @@ -261,12 +251,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector } return visitedDocs.getAndSet(docId) == false; }; - final int nProbe; + int nProbe = DYNAMIC_NPROBE; if (knnCollector.getSearchStrategy() instanceof IVFKnnSearchStrategy ivfSearchStrategy) { nProbe = ivfSearchStrategy.getNProbe(); - } else { - // TODO calculate nProbe given the number of centroids vs. number of vectors for given `k` - nProbe = 10; } FieldEntry entry = fields.get(fieldInfo.number); @@ -277,17 +264,27 @@ public final void search(String field, float[] target, KnnCollector knnCollector target, ivfClusters ); + if (nProbe == DYNAMIC_NPROBE) { + // empirically based, and a good dynamic to get decent recall while scaling a la "efSearch" + // scaling by the number of centroids vs. the nearest neighbors requested + // not perfect, but a comparative heuristic. + // we might want to utilize the total vector count as well, but this is a good start + nProbe = (int) Math.round(Math.log10(centroidQueryScorer.size()) * Math.sqrt(knnCollector.k())); + // clip to be between 1 and the number of centroids + nProbe = Math.max(Math.min(nProbe, centroidQueryScorer.size()), 1); + } final NeighborQueue centroidQueue = scorePostingLists(fieldInfo, knnCollector, centroidQueryScorer, nProbe); PostingVisitor scorer = getPostingVisitor(fieldInfo, ivfClusters, target, needsScoring); int centroidsVisited = 0; long expectedDocs = 0; long actualDocs = 0; // initially we visit only the "centroids to search" - while (centroidQueue.size() > 0 && centroidsVisited < nProbe) { + while (centroidQueue.size() > 0 && centroidsVisited < nProbe && actualDocs < knnCollector.k()) { ++centroidsVisited; // todo do we actually need to know the score??? int centroidOrdinal = centroidQueue.pop(); - // todo do we need direct access to the raw centroid??? + // todo do we need direct access to the raw centroid???, this is used for quantizing, maybe hydrating and quantizing + // is enough? expectedDocs += scorer.resetPostingsScorer(centroidOrdinal, centroidQueryScorer.centroid(centroidOrdinal)); actualDocs += scorer.visit(knnCollector); } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java index 18ae1fa802df6..70b6790c4dc1d 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java @@ -41,6 +41,7 @@ public class MapperFeatures implements FeatureSpecification { "mapper.unknown_field_mapping_update_error_message" ); static final NodeFeature NPE_ON_DIMS_UPDATE_FIX = new NodeFeature("mapper.npe_on_dims_update_fix"); + static final NodeFeature IVF_FORMAT_CLUSTER_FEATURE = new NodeFeature("mapper.ivf_format_cluster_feature"); @Override public Set getTestFeatures() { @@ -68,7 +69,8 @@ public Set getTestFeatures() { DateFieldMapper.INVALID_DATE_FIX, NPE_ON_DIMS_UPDATE_FIX, RESCORE_ZERO_VECTOR_QUANTIZED_VECTOR_MAPPING, - USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ + USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ, + IVF_FORMAT_CLUSTER_FEATURE ); } } 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 fd38cebb58e3d..0d6970fba1927 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 @@ -39,6 +39,7 @@ import org.apache.lucene.util.VectorUtil; import org.elasticsearch.common.ParsingException; import org.elasticsearch.common.settings.Setting; +import org.elasticsearch.common.util.FeatureFlag; import org.elasticsearch.common.xcontent.support.XContentMapValues; import org.elasticsearch.features.NodeFeature; import org.elasticsearch.index.IndexVersion; @@ -48,6 +49,7 @@ import org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat; import org.elasticsearch.index.codec.vectors.ES815HnswBitVectorsFormat; +import org.elasticsearch.index.codec.vectors.IVFVectorsFormat; import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat; import org.elasticsearch.index.fielddata.FieldDataContext; @@ -62,6 +64,7 @@ import org.elasticsearch.index.mapper.Mapper; import org.elasticsearch.index.mapper.MapperBuilderContext; import org.elasticsearch.index.mapper.MapperParsingException; +import org.elasticsearch.index.mapper.MappingLookup; import org.elasticsearch.index.mapper.MappingParser; import org.elasticsearch.index.mapper.NumberFieldMapper; import org.elasticsearch.index.mapper.SimpleMappedFieldType; @@ -78,6 +81,7 @@ import org.elasticsearch.search.vectors.ESDiversifyingChildrenFloatKnnVectorQuery; import org.elasticsearch.search.vectors.ESKnnByteVectorQuery; import org.elasticsearch.search.vectors.ESKnnFloatVectorQuery; +import org.elasticsearch.search.vectors.IVFKnnFloatVectorQuery; import org.elasticsearch.search.vectors.RescoreKnnVectorQuery; import org.elasticsearch.search.vectors.VectorData; import org.elasticsearch.search.vectors.VectorSimilarityQuery; @@ -106,6 +110,8 @@ import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_INDEX_VERSION_CREATED; import static org.elasticsearch.common.Strings.format; import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken; +import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MAX_VECTORS_PER_CLUSTER; +import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MIN_VECTORS_PER_CLUSTER; /** * A {@link FieldMapper} for indexing a dense vector of floats. @@ -115,6 +121,8 @@ public class DenseVectorFieldMapper extends FieldMapper { private static final float EPS = 1e-3f; public static final int BBQ_MIN_DIMS = 64; + public static final FeatureFlag IVF_FORMAT = new FeatureFlag("ivf_format"); + public static boolean isNotUnitVector(float magnitude) { return Math.abs(magnitude - 1.0f) > EPS; } @@ -1594,6 +1602,52 @@ public boolean supportsElementType(ElementType elementType) { return elementType == ElementType.FLOAT; } + @Override + public boolean supportsDimension(int dims) { + return dims >= BBQ_MIN_DIMS; + } + }, + BBQ_IVF("bbq_ivf", true) { + @Override + public IndexOptions parseIndexOptions(String fieldName, Map indexOptionsMap, IndexVersion indexVersion) { + Object clusterSizeNode = indexOptionsMap.remove("cluster_size"); + int clusterSize = IVFVectorsFormat.DEFAULT_VECTORS_PER_CLUSTER; + if (clusterSizeNode != null) { + clusterSize = XContentMapValues.nodeIntegerValue(clusterSizeNode); + if (clusterSize < MIN_VECTORS_PER_CLUSTER || clusterSize > MAX_VECTORS_PER_CLUSTER) { + throw new IllegalArgumentException( + "cluster_size must be between " + + MIN_VECTORS_PER_CLUSTER + + " and " + + MAX_VECTORS_PER_CLUSTER + + ", got: " + + clusterSize + ); + } + } + RescoreVector rescoreVector = RescoreVector.fromIndexOptions(indexOptionsMap, indexVersion); + if (rescoreVector == null) { + rescoreVector = new RescoreVector(DEFAULT_OVERSAMPLE); + } + Object nProbeNode = indexOptionsMap.remove("default_n_probe"); + int nProbe = -1; + if (nProbeNode != null) { + nProbe = XContentMapValues.nodeIntegerValue(nProbeNode); + if (nProbe < 1 && nProbe != -1) { + throw new IllegalArgumentException( + "default_n_probe must be at least 1 or exactly -1, got: " + nProbe + " for field [" + fieldName + "]" + ); + } + } + MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap); + return new BBQIVFIndexOptions(clusterSize, nProbe, rescoreVector); + } + + @Override + public boolean supportsElementType(ElementType elementType) { + return elementType == ElementType.FLOAT; + } + @Override public boolean supportsDimension(int dims) { return dims >= BBQ_MIN_DIMS; @@ -1601,7 +1655,10 @@ public boolean supportsDimension(int dims) { }; static Optional fromString(String type) { - return Stream.of(VectorIndexType.values()).filter(vectorIndexType -> vectorIndexType.name.equals(type)).findFirst(); + return Stream.of(VectorIndexType.values()) + .filter(vectorIndexType -> vectorIndexType != VectorIndexType.BBQ_IVF || IVF_FORMAT.isEnabled()) + .filter(vectorIndexType -> vectorIndexType.name.equals(type)) + .findFirst(); } private final String name; @@ -2100,6 +2157,54 @@ public boolean validateDimension(int dim, boolean throwOnError) { } + static class BBQIVFIndexOptions extends QuantizedIndexOptions { + final int clusterSize; + final int defaultNProbe; + + BBQIVFIndexOptions(int clusterSize, int defaultNProbe, RescoreVector rescoreVector) { + super(VectorIndexType.BBQ_IVF, rescoreVector); + this.clusterSize = clusterSize; + this.defaultNProbe = defaultNProbe; + } + + @Override + KnnVectorsFormat getVectorsFormat(ElementType elementType) { + assert elementType == ElementType.FLOAT; + return new IVFVectorsFormat(clusterSize); + } + + @Override + boolean updatableTo(IndexOptions update) { + return update.type.equals(this.type); + } + + @Override + boolean doEquals(IndexOptions other) { + BBQIVFIndexOptions that = (BBQIVFIndexOptions) other; + return clusterSize == that.clusterSize + && defaultNProbe == that.defaultNProbe + && Objects.equals(rescoreVector, that.rescoreVector); + } + + @Override + int doHashCode() { + return Objects.hash(clusterSize, defaultNProbe, rescoreVector); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field("type", type); + builder.field("cluster_size", clusterSize); + builder.field("default_n_probe", defaultNProbe); + if (rescoreVector != null) { + rescoreVector.toXContent(builder, params); + } + builder.endObject(); + return builder; + } + } + public record RescoreVector(float oversample) implements ToXContentObject { static final String NAME = "rescore_vector"; static final String OVERSAMPLE = "oversample"; @@ -2411,17 +2516,25 @@ && isNotUnitVector(squaredMagnitude)) { adjustedK = Math.min((int) Math.ceil(k * oversample), OVERSAMPLE_LIMIT); numCands = Math.max(adjustedK, numCands); } - Query knnQuery = parentFilter != null - ? new ESDiversifyingChildrenFloatKnnVectorQuery( - name(), - queryVector, - filter, - adjustedK, - numCands, - parentFilter, - knnSearchStrategy - ) - : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, knnSearchStrategy); + if (parentFilter != null && indexOptions instanceof BBQIVFIndexOptions) { + throw new IllegalArgumentException("IVF index does not support nested queries"); + } + Query knnQuery; + if (indexOptions instanceof BBQIVFIndexOptions bbqIndexOptions) { + knnQuery = new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe); + } else { + knnQuery = parentFilter != null + ? new ESDiversifyingChildrenFloatKnnVectorQuery( + name(), + queryVector, + filter, + adjustedK, + numCands, + parentFilter, + knnSearchStrategy + ) + : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, knnSearchStrategy); + } if (rescore) { knnQuery = new RescoreKnnVectorQuery( name(), @@ -2651,6 +2764,19 @@ public FieldMapper.Builder getMergeBuilder() { return new Builder(leafName(), indexCreatedVersion).init(this); } + @Override + public void doValidate(MappingLookup mappers) { + if (indexOptions instanceof BBQIVFIndexOptions && mappers.nestedLookup().getNestedParent(fullPath()) != null) { + throw new IllegalArgumentException( + "[" + + CONTENT_TYPE + + "] fields with index type [" + + indexOptions.type + + "] cannot be indexed if they're within [nested] mappings" + ); + } + } + private static IndexOptions parseIndexOptions(String fieldName, Object propNode, IndexVersion indexVersion) { @SuppressWarnings("unchecked") Map indexOptionsMap = (Map) propNode; diff --git a/server/src/main/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQuery.java b/server/src/main/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQuery.java index 31461d6ab1238..66446a0dc43bd 100644 --- a/server/src/main/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQuery.java +++ b/server/src/main/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQuery.java @@ -50,15 +50,26 @@ abstract class AbstractIVFKnnVectorQuery extends Query implements QueryProfilerP protected final String field; protected final int nProbe; protected final int k; + protected final int numCands; protected final Query filter; protected final KnnSearchStrategy searchStrategy; protected int vectorOpsCount; - protected AbstractIVFKnnVectorQuery(String field, int nProbe, int k, Query filter) { + protected AbstractIVFKnnVectorQuery(String field, int nProbe, int k, int numCands, Query filter) { + if (k < 1) { + throw new IllegalArgumentException("k must be at least 1, got: " + k); + } + if (nProbe < 1 && nProbe != -1) { + throw new IllegalArgumentException("nProbe must be at least 1 or exactly -1, got: " + nProbe); + } + if (numCands < k) { + throw new IllegalArgumentException("numCands must be at least k, got: " + numCands); + } this.field = field; this.nProbe = nProbe; this.k = k; this.filter = filter; + this.numCands = numCands; this.searchStrategy = new IVFKnnSearchStrategy(nProbe); } @@ -103,7 +114,8 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { } else { filterWeight = null; } - KnnCollectorManager knnCollectorManager = getKnnCollectorManager(k, indexSearcher); + // we request numCands as we are using it as an approximation measure + KnnCollectorManager knnCollectorManager = getKnnCollectorManager(numCands, indexSearcher); TaskExecutor taskExecutor = indexSearcher.getTaskExecutor(); List leafReaderContexts = reader.leaves(); List> tasks = new ArrayList<>(leafReaderContexts.size()); diff --git a/server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java b/server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java index 84793df2ae894..1b4cb44eb0c0a 100644 --- a/server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java +++ b/server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java @@ -30,17 +30,12 @@ public class IVFKnnFloatVectorQuery extends AbstractIVFKnnVectorQuery { * @param field the field to search * @param query the query vector * @param k the number of nearest neighbors to return + * @param numCands the number of nearest neighbors to gather per shard * @param filter the filter to apply to the results * @param nProbe the number of probes to use for the IVF search strategy */ - public IVFKnnFloatVectorQuery(String field, float[] query, int k, Query filter, int nProbe) { - super(field, nProbe, k, filter); - if (k < 1) { - throw new IllegalArgumentException("k must be at least 1, got: " + k); - } - if (nProbe < 1) { - throw new IllegalArgumentException("nProbe must be at least 1, got: " + nProbe); - } + public IVFKnnFloatVectorQuery(String field, float[] query, int k, int numCands, Query filter, int nProbe) { + super(field, nProbe, k, numCands, filter); this.query = query; } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormatTests.java index c822d71a358ff..f7eb4cf5241ce 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/IVFVectorsFormatTests.java @@ -11,6 +11,7 @@ import com.carrotsearch.randomizedtesting.generators.RandomPicks; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.index.VectorSimilarityFunction; @@ -20,6 +21,13 @@ import org.junit.Before; import java.util.List; +import java.util.Locale; + +import static java.lang.String.format; +import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MAX_VECTORS_PER_CLUSTER; +import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.MIN_VECTORS_PER_CLUSTER; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.oneOf; public class IVFVectorsFormatTests extends BaseKnnVectorsFormatTestCase { @@ -32,7 +40,7 @@ public class IVFVectorsFormatTests extends BaseKnnVectorsFormatTestCase { @Before @Override public void setUp() throws Exception { - format = new IVFVectorsFormat(random().nextInt(10, 1000)); + format = new IVFVectorsFormat(random().nextInt(MIN_VECTORS_PER_CLUSTER, IVFVectorsFormat.MAX_VECTORS_PER_CLUSTER)); super.setUp(); } @@ -62,4 +70,28 @@ public void testSearchWithVisitedLimit() { protected Codec getCodec() { return TestUtil.alwaysKnnVectorsFormat(format); } + + @Override + public void testAdvance() throws Exception { + // TODO re-enable with hierarchical IVF, clustering as it is is flaky + } + + public void testToString() { + FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) { + @Override + public KnnVectorsFormat knnVectorsFormat() { + return new IVFVectorsFormat(128); + } + }; + String expectedPattern = "IVFVectorsFormat(vectorPerCluster=128)"; + + var defaultScorer = format(Locale.ROOT, expectedPattern, "DefaultFlatVectorScorer"); + var memSegScorer = format(Locale.ROOT, expectedPattern, "Lucene99MemorySegmentFlatVectorsScorer"); + assertThat(customCodec.knnVectorsFormat().toString(), is(oneOf(defaultScorer, memSegScorer))); + } + + public void testLimits() { + expectThrows(IllegalArgumentException.class, () -> new IVFVectorsFormat(MIN_VECTORS_PER_CLUSTER - 1)); + expectThrows(IllegalArgumentException.class, () -> new IVFVectorsFormat(MAX_VECTORS_PER_CLUSTER + 1)); + } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java index 1ce916533a2c9..9db39b669c558 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java @@ -31,6 +31,7 @@ import org.elasticsearch.index.codec.CodecService; import org.elasticsearch.index.codec.LegacyPerFieldMapperCodec; import org.elasticsearch.index.codec.PerFieldMapperCodec; +import org.elasticsearch.index.codec.vectors.IVFVectorsFormat; import org.elasticsearch.index.mapper.DocumentMapper; import org.elasticsearch.index.mapper.DocumentParsingException; import org.elasticsearch.index.mapper.LuceneDocument; @@ -64,6 +65,8 @@ import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; +import static org.elasticsearch.index.codec.vectors.IVFVectorsFormat.DYNAMIC_NPROBE; +import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; @@ -884,6 +887,69 @@ protected void assertExistsQuery(MappedFieldType fieldType, Query query, LuceneD @Override public void testAggregatableConsistency() {} + public void testIVFParsing() throws IOException { + assumeTrue("feature flag [ivf_format] must be enabled", IVF_FORMAT.isEnabled()); + { + DocumentMapper mapperService = createDocumentMapper(fieldMapping(b -> { + b.field("type", "dense_vector"); + b.field("dims", 128); + b.field("index", true); + b.field("similarity", "dot_product"); + b.startObject("index_options"); + b.field("type", "bbq_ivf"); + b.endObject(); + })); + + DenseVectorFieldMapper denseVectorFieldMapper = (DenseVectorFieldMapper) mapperService.mappers().getMapper("field"); + DenseVectorFieldMapper.BBQIVFIndexOptions indexOptions = (DenseVectorFieldMapper.BBQIVFIndexOptions) denseVectorFieldMapper + .fieldType() + .getIndexOptions(); + assertEquals(3.0F, indexOptions.rescoreVector.oversample(), 0.0F); + assertEquals(IVFVectorsFormat.DEFAULT_VECTORS_PER_CLUSTER, indexOptions.clusterSize); + assertEquals(DYNAMIC_NPROBE, indexOptions.defaultNProbe); + } + { + DocumentMapper mapperService = createDocumentMapper(fieldMapping(b -> { + b.field("type", "dense_vector"); + b.field("dims", 128); + b.field("index", true); + b.field("similarity", "dot_product"); + b.startObject("index_options"); + b.field("type", "bbq_ivf"); + b.field("cluster_size", 1000); + b.field("default_n_probe", 10); + b.field(DenseVectorFieldMapper.RescoreVector.NAME, Map.of("oversample", 2.0f)); + b.endObject(); + })); + + DenseVectorFieldMapper denseVectorFieldMapper = (DenseVectorFieldMapper) mapperService.mappers().getMapper("field"); + DenseVectorFieldMapper.BBQIVFIndexOptions indexOptions = (DenseVectorFieldMapper.BBQIVFIndexOptions) denseVectorFieldMapper + .fieldType() + .getIndexOptions(); + assertEquals(2F, indexOptions.rescoreVector.oversample(), 0.0F); + assertEquals(1000, indexOptions.clusterSize); + assertEquals(10, indexOptions.defaultNProbe); + } + } + + public void testIVFParsingFailureInRelease() { + assumeFalse("feature flag [ivf_format] must be disabled", IVF_FORMAT.isEnabled()); + + Exception e = expectThrows( + MapperParsingException.class, + () -> createDocumentMapper( + fieldMapping( + b -> b.field("type", "dense_vector") + .field("dims", dims) + .startObject("index_options") + .field("type", "bbq_ivf") + .endObject() + ) + ) + ); + assertThat(e.getMessage(), containsString("Unknown vector index options")); + } + public void testRescoreVectorForNonQuantized() { for (String indexType : List.of("hnsw", "flat")) { Exception e = expectThrows( @@ -2250,6 +2316,57 @@ public void testKnnBBQHNSWVectorsFormat() throws IOException { assertEquals(expectedString, knnVectorsFormat.toString()); } + public void testBBQIVFVectorsFormatDisallowsNested() throws IOException { + assumeTrue("feature flag [ivf_format] must be enabled", IVF_FORMAT.isEnabled()); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> createDocumentMapper(fieldMapping(b -> { + b.field("type", "nested"); + b.startObject("properties"); + b.startObject("field"); + b.field("type", "dense_vector"); + b.field("dims", randomIntBetween(64, 4096)); + b.field("index", true); + b.field("similarity", "dot_product"); + b.startObject("index_options"); + b.field("type", "bbq_ivf"); + b.endObject(); + b.endObject(); + b.endObject(); + }))); + assertThat( + e.getMessage(), + containsString("fields with index type [bbq_ivf] cannot be indexed if they're within [nested] mappings") + ); + } + + public void testKnnBBQIVFVectorsFormat() throws IOException { + assumeTrue("feature flag [ivf_format] must be enabled", IVF_FORMAT.isEnabled()); + final int dims = randomIntBetween(64, 4096); + MapperService mapperService = createMapperService(fieldMapping(b -> { + b.field("type", "dense_vector"); + b.field("dims", dims); + b.field("index", true); + b.field("similarity", "dot_product"); + b.startObject("index_options"); + b.field("type", "bbq_ivf"); + b.endObject(); + })); + CodecService codecService = new CodecService(mapperService, BigArrays.NON_RECYCLING_INSTANCE); + Codec codec = codecService.codec("default"); + KnnVectorsFormat knnVectorsFormat; + if (CodecService.ZSTD_STORED_FIELDS_FEATURE_FLAG) { + assertThat(codec, instanceOf(PerFieldMapperCodec.class)); + knnVectorsFormat = ((PerFieldMapperCodec) codec).getKnnVectorsFormatForField("field"); + } else { + if (codec instanceof CodecService.DeduplicateFieldInfosCodec deduplicateFieldInfosCodec) { + codec = deduplicateFieldInfosCodec.delegate(); + } + assertThat(codec, instanceOf(LegacyPerFieldMapperCodec.class)); + knnVectorsFormat = ((LegacyPerFieldMapperCodec) codec).getKnnVectorsFormatForField("field"); + } + String expectedString = "IVFVectorsFormat(vectorPerCluster=384)"; + assertEquals(expectedString, knnVectorsFormat.toString()); + } + public void testInvalidVectorDimensionsBBQ() { for (String quantizedFlatFormat : new String[] { "bbq_hnsw", "bbq_flat" }) { MapperParsingException e = expectThrows(MapperParsingException.class, () -> createDocumentMapper(fieldMapping(b -> { diff --git a/server/src/test/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQueryTests.java b/server/src/test/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQueryTests.java index d15d981e3997a..2c57b6958f9ca 100644 --- a/server/src/test/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQueryTests.java +++ b/server/src/test/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQueryTests.java @@ -27,7 +27,7 @@ public class IVFKnnFloatVectorQueryTests extends AbstractIVFKnnVectorQueryTestCa @Override IVFKnnFloatVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter, int nProbe) { - return new IVFKnnFloatVectorQuery(field, query, k, queryFilter, nProbe); + return new IVFKnnFloatVectorQuery(field, query, k, k, queryFilter, nProbe); } @Override diff --git a/test/test-clusters/src/main/java/org/elasticsearch/test/cluster/FeatureFlag.java b/test/test-clusters/src/main/java/org/elasticsearch/test/cluster/FeatureFlag.java index 53e4a971add7d..830adfee64e49 100644 --- a/test/test-clusters/src/main/java/org/elasticsearch/test/cluster/FeatureFlag.java +++ b/test/test-clusters/src/main/java/org/elasticsearch/test/cluster/FeatureFlag.java @@ -20,7 +20,8 @@ public enum FeatureFlag { SUB_OBJECTS_AUTO_ENABLED("es.sub_objects_auto_feature_flag_enabled=true", Version.fromString("8.16.0"), null), DOC_VALUES_SKIPPER("es.doc_values_skipper_feature_flag_enabled=true", Version.fromString("8.18.1"), null), USE_LUCENE101_POSTINGS_FORMAT("es.use_lucene101_postings_format_feature_flag_enabled=true", Version.fromString("9.1.0"), null), - INFERENCE_CUSTOM_SERVICE_ENABLED("es.inference_custom_service_feature_flag_enabled=true", Version.fromString("8.19.0"), null); + INFERENCE_CUSTOM_SERVICE_ENABLED("es.inference_custom_service_feature_flag_enabled=true", Version.fromString("8.19.0"), null), + IVF_FORMAT("es.ivf_format_feature_flag_enabled=true", Version.fromString("9.1.0"), null); public final String systemProperty; public final Version from;