|
@@ -33,11 +33,14 @@ public class VectorSearchIT extends AbstractRollingUpgradeTestCase {
|
|
|
private static final String SCRIPT_BYTE_INDEX_NAME = "script_byte_vector_index";
|
|
|
private static final String BYTE_INDEX_NAME = "byte_vector_index";
|
|
|
private static final String QUANTIZED_INDEX_NAME = "quantized_vector_index";
|
|
|
+ private static final String BBQ_INDEX_NAME = "bbq_vector_index";
|
|
|
private static final String FLAT_QUANTIZED_INDEX_NAME = "flat_quantized_vector_index";
|
|
|
+ private static final String FLAT_BBQ_INDEX_NAME = "flat_bbq_vector_index";
|
|
|
private static final String FLOAT_VECTOR_SEARCH_VERSION = "8.4.0";
|
|
|
private static final String BYTE_VECTOR_SEARCH_VERSION = "8.6.0";
|
|
|
private static final String QUANTIZED_VECTOR_SEARCH_VERSION = "8.12.1";
|
|
|
private static final String FLAT_QUANTIZED_VECTOR_SEARCH_VERSION = "8.13.0";
|
|
|
+ private static final String BBQ_VECTOR_SEARCH_VERSION = "8.18.0";
|
|
|
|
|
|
public void testScriptByteVectorSearch() throws Exception {
|
|
|
assumeTrue("byte vector search is not supported on this version", getOldClusterTestVersion().onOrAfter(BYTE_VECTOR_SEARCH_VERSION));
|
|
@@ -429,6 +432,182 @@ public class VectorSearchIT extends AbstractRollingUpgradeTestCase {
|
|
|
assertThat((double) hits.get(0).get("_score"), closeTo(0.9934857, 0.005));
|
|
|
}
|
|
|
|
|
|
+ public void testBBQVectorSearch() throws Exception {
|
|
|
+ assumeTrue(
|
|
|
+ "Quantized vector search is not supported on this version",
|
|
|
+ getOldClusterTestVersion().onOrAfter(BBQ_VECTOR_SEARCH_VERSION)
|
|
|
+ );
|
|
|
+ if (isOldCluster()) {
|
|
|
+ String mapping = """
|
|
|
+ {
|
|
|
+ "properties": {
|
|
|
+ "vector": {
|
|
|
+ "type": "dense_vector",
|
|
|
+ "dims": 64,
|
|
|
+ "index": true,
|
|
|
+ "similarity": "cosine",
|
|
|
+ "index_options": {
|
|
|
+ "type": "bbq_hnsw",
|
|
|
+ "ef_construction": 100,
|
|
|
+ "m": 16
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ """;
|
|
|
+ // create index and index 10 random floating point vectors
|
|
|
+ createIndex(BBQ_INDEX_NAME, Settings.EMPTY, mapping);
|
|
|
+ index64DimVectors(BBQ_INDEX_NAME);
|
|
|
+ // force merge the index
|
|
|
+ client().performRequest(new Request("POST", "/" + BBQ_INDEX_NAME + "/_forcemerge?max_num_segments=1"));
|
|
|
+ }
|
|
|
+ Request searchRequest = new Request("POST", "/" + BBQ_INDEX_NAME + "/_search");
|
|
|
+ searchRequest.setJsonEntity("""
|
|
|
+ {
|
|
|
+ "query": {
|
|
|
+ "script_score": {
|
|
|
+ "query": {
|
|
|
+ "exists": {
|
|
|
+ "field": "vector"
|
|
|
+ }
|
|
|
+ },
|
|
|
+ "script": {
|
|
|
+ "source": "cosineSimilarity(params.query, 'vector') + 1.0",
|
|
|
+ "params": {
|
|
|
+ "query": [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
|
|
|
+ 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ """);
|
|
|
+ Map<String, Object> response = search(searchRequest);
|
|
|
+ assertThat(extractValue(response, "hits.total.value"), equalTo(7));
|
|
|
+ List<Map<String, Object>> hits = extractValue(response, "hits.hits");
|
|
|
+ assertThat(hits.get(0).get("_id"), equalTo("0"));
|
|
|
+ assertThat((double) hits.get(0).get("_score"), closeTo(1.9869276, 0.0001));
|
|
|
+
|
|
|
+ // search with knn
|
|
|
+ searchRequest = new Request("POST", "/" + BBQ_INDEX_NAME + "/_search");
|
|
|
+ searchRequest.setJsonEntity("""
|
|
|
+ {
|
|
|
+ "knn": {
|
|
|
+ "field": "vector",
|
|
|
+ "query_vector": [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
|
|
|
+ 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
|
|
|
+ "k": 2,
|
|
|
+ "num_candidates": 5
|
|
|
+ }
|
|
|
+ }
|
|
|
+ """);
|
|
|
+ response = search(searchRequest);
|
|
|
+ assertThat(extractValue(response, "hits.total.value"), equalTo(2));
|
|
|
+ hits = extractValue(response, "hits.hits");
|
|
|
+ assertThat(hits.get(0).get("_id"), equalTo("0"));
|
|
|
+ assertThat((double) hits.get(0).get("_score"), closeTo(0.9934857, 0.005));
|
|
|
+ }
|
|
|
+
|
|
|
+ public void testFlatBBQVectorSearch() throws Exception {
|
|
|
+ assumeTrue(
|
|
|
+ "Quantized vector search is not supported on this version",
|
|
|
+ getOldClusterTestVersion().onOrAfter(BBQ_VECTOR_SEARCH_VERSION)
|
|
|
+ );
|
|
|
+ if (isOldCluster()) {
|
|
|
+ String mapping = """
|
|
|
+ {
|
|
|
+ "properties": {
|
|
|
+ "vector": {
|
|
|
+ "type": "dense_vector",
|
|
|
+ "dims": 64,
|
|
|
+ "index": true,
|
|
|
+ "similarity": "cosine",
|
|
|
+ "index_options": {
|
|
|
+ "type": "bbq_flat"
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ """;
|
|
|
+ // create index and index 10 random floating point vectors
|
|
|
+ createIndex(FLAT_BBQ_INDEX_NAME, Settings.EMPTY, mapping);
|
|
|
+ index64DimVectors(FLAT_BBQ_INDEX_NAME);
|
|
|
+ // force merge the index
|
|
|
+ client().performRequest(new Request("POST", "/" + FLAT_BBQ_INDEX_NAME + "/_forcemerge?max_num_segments=1"));
|
|
|
+ }
|
|
|
+ Request searchRequest = new Request("POST", "/" + FLAT_BBQ_INDEX_NAME + "/_search");
|
|
|
+ searchRequest.setJsonEntity("""
|
|
|
+ {
|
|
|
+ "query": {
|
|
|
+ "script_score": {
|
|
|
+ "query": {
|
|
|
+ "exists": {
|
|
|
+ "field": "vector"
|
|
|
+ }
|
|
|
+ },
|
|
|
+ "script": {
|
|
|
+ "source": "cosineSimilarity(params.query, 'vector') + 1.0",
|
|
|
+ "params": {
|
|
|
+ "query": [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
|
|
|
+ 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ """);
|
|
|
+ Map<String, Object> response = search(searchRequest);
|
|
|
+ assertThat(extractValue(response, "hits.total.value"), equalTo(7));
|
|
|
+ List<Map<String, Object>> hits = extractValue(response, "hits.hits");
|
|
|
+ assertThat(hits.get(0).get("_id"), equalTo("0"));
|
|
|
+ assertThat((double) hits.get(0).get("_score"), closeTo(1.9869276, 0.0001));
|
|
|
+
|
|
|
+ // search with knn
|
|
|
+ searchRequest = new Request("POST", "/" + FLAT_BBQ_INDEX_NAME + "/_search");
|
|
|
+ searchRequest.setJsonEntity("""
|
|
|
+ {
|
|
|
+ "knn": {
|
|
|
+ "field": "vector",
|
|
|
+ "query_vector": [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
|
|
|
+ 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6],
|
|
|
+ "k": 2,
|
|
|
+ "num_candidates": 5
|
|
|
+ }
|
|
|
+ }
|
|
|
+ """);
|
|
|
+ response = search(searchRequest);
|
|
|
+ assertThat(extractValue(response, "hits.total.value"), equalTo(2));
|
|
|
+ hits = extractValue(response, "hits.hits");
|
|
|
+ assertThat(hits.get(0).get("_id"), equalTo("0"));
|
|
|
+ assertThat((double) hits.get(0).get("_score"), closeTo(0.9934857, 0.005));
|
|
|
+ }
|
|
|
+
|
|
|
+ private void index64DimVectors(String indexName) throws Exception {
|
|
|
+ String[] vectors = new String[] {
|
|
|
+ "{\"vector\":[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "
|
|
|
+ + "1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}",
|
|
|
+ "{\"vector\":[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "
|
|
|
+ + "1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]}",
|
|
|
+ "{\"vector\":[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "
|
|
|
+ + "1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]}",
|
|
|
+ "{\"vector\":[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, "
|
|
|
+ + "2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}",
|
|
|
+ "{\"vector\":[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, "
|
|
|
+ + "3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}",
|
|
|
+ "{\"vector\":[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "
|
|
|
+ + "1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}",
|
|
|
+ "{\"vector\":[3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "
|
|
|
+ + "1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}",
|
|
|
+ "{}" };
|
|
|
+ for (int i = 0; i < vectors.length; i++) {
|
|
|
+ Request indexRequest = new Request("PUT", "/" + indexName + "/_doc/" + i);
|
|
|
+ indexRequest.setJsonEntity(vectors[i]);
|
|
|
+ assertOK(client().performRequest(indexRequest));
|
|
|
+ }
|
|
|
+ // always refresh to ensure the data is visible
|
|
|
+ refresh(indexName);
|
|
|
+ }
|
|
|
+
|
|
|
private void indexVectors(String indexName) throws Exception {
|
|
|
String[] vectors = new String[] {
|
|
|
"{\"vector\":[1, 1, 1]}",
|