diff --git a/docs/changelog/135778.yaml b/docs/changelog/135778.yaml new file mode 100644 index 0000000000000..6e4fb9c3f37c5 --- /dev/null +++ b/docs/changelog/135778.yaml @@ -0,0 +1,5 @@ +pr: 135778 +summary: "Add `on_disk_rescore: true` option to disk BBQ to rescore vectors on-disk without loading into memory" +area: Vector Search +type: feature +issues: [] diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/180_update_dense_vector_type.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/180_update_dense_vector_type.yml index 23ecdbe95e22f..06173aab39f6e 100644 --- a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/180_update_dense_vector_type.yml +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/180_update_dense_vector_type.yml @@ -1682,6 +1682,109 @@ setup: - match: { test_index.mappings.properties.embedding.index_options.confidence_interval: 0.3 } --- +"Allowed dense vector updates on same type but different other index_options, bbq_disk": + - requires: + cluster_features: "mapper.vectors.diskbbq_on_disk_rescoring" + reason: 'diskbbq needs to support on-disk rescoring' + - requires: + test_runner_features: [ contains ] + - do: + indices.create: + index: test_index + + - do: + indices.put_mapping: + index: test_index + body: + properties: + embedding: + type: dense_vector + dims: 64 + index_options: + type: bbq_disk + + - do: + indices.get_mapping: + index: test_index + + - match: { test_index.mappings.properties.embedding.type: dense_vector } + - match: { test_index.mappings.properties.embedding.index_options.type: bbq_disk } + + - do: + index: + index: test_index + id: "1" + body: + embedding: [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] + - do: + indices.flush: + index: test_index + + - do: + indices.put_mapping: + index: test_index + body: + properties: + embedding: + type: dense_vector + dims: 64 + index_options: + type: bbq_disk + on_disk_rescore: true + + - do: + indices.get_mapping: + index: test_index + + - match: { test_index.mappings.properties.embedding.type: dense_vector } + - match: { test_index.mappings.properties.embedding.index_options.type: bbq_disk } + - match: { test_index.mappings.properties.embedding.index_options.on_disk_rescore: true } + + - do: + index: + index: test_index + id: "2" + body: + embedding: [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: + indices.flush: + index: test_index + - do: + indices.refresh: { } + - do: + search: + index: test_index + body: + knn: + field: embedding + 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: 2 + num_candidates: 2 + + - match: { hits.hits.0._id: "1" } + - match: { hits.hits.1._id: "2" } +--- "Test create and update dense vector mapping to int4 with per-doc indexing and flush": - requires: cluster_features: "gte_v8.16.0" 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 index 3fbbbc3de1f48..507b46090e52c 100644 --- 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 @@ -306,6 +306,98 @@ setup: - 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 with on-disk rescore": + - requires: + cluster_features: [ "mapper.vectors.diskbbq_on_disk_rescoring" ] + reason: Needs on_disk_rescoring feature for DiskBBQ + - skip: + features: "headers" + - do: + indices.create: + index: bbq_on_disk_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_disk + on_disk_rescore: true + rescore_vector: + oversample: 1.5 + + - do: + bulk: + index: bbq_on_disk_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_on_disk_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_on_disk_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: diff --git a/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java b/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java new file mode 100644 index 0000000000000..efbc19b30079c --- /dev/null +++ b/server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java @@ -0,0 +1,187 @@ +/* + * 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.store; + +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.apache.logging.log4j.Level; +import org.apache.lucene.misc.store.DirectIODirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.core.Strings; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.search.vectors.KnnSearchBuilder; +import org.elasticsearch.search.vectors.VectorData; +import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.test.InternalSettingsPlugin; +import org.elasticsearch.test.MockLog; +import org.elasticsearch.test.junit.annotations.TestLogging; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.OptionalLong; +import java.util.stream.IntStream; + +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; + +@LuceneTestCase.SuppressCodecs("*") // only use our own codecs +@ESTestCase.WithoutEntitlements // requires entitlement delegation ES-10920 +public class DirectIOIT extends ESIntegTestCase { + + private static boolean SUPPORTED; + + @BeforeClass + public static void checkSupported() { + Path path = createTempDir("directIOProbe"); + try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { + out.writeString("test"); + SUPPORTED = true; + } catch (IOException e) { + SUPPORTED = false; + } + } + + static DirectIODirectory open(Path path) throws IOException { + return new DirectIODirectory(FSDirectory.open(path)) { + @Override + protected boolean useDirectIO(String name, IOContext context, OptionalLong fileLength) { + return true; + } + }; + } + + private final String type; + + @ParametersFactory + public static Iterable parameters() { + return List.of(new Object[] { "bbq_disk" }); + } + + public DirectIOIT(String type) { + this.type = type; + } + + @Override + protected Collection> nodePlugins() { + return List.of(InternalSettingsPlugin.class); + } + + private String indexVectors(boolean directIO) { + String indexName = "test-vectors-" + directIO; + assertAcked( + prepareCreate(indexName).setSettings(Settings.builder().put(InternalSettingsPlugin.USE_COMPOUND_FILE.getKey(), false)) + .setMapping(Strings.format(""" + { + "properties": { + "fooVector": { + "type": "dense_vector", + "dims": 64, + "element_type": "float", + "index": true, + "similarity": "l2_norm", + "index_options": { + "type": "%s", + "on_disk_rescore": %s + } + } + } + } + """, type, directIO)) + ); + ensureGreen(indexName); + + for (int i = 0; i < 1000; i++) { + indexDoc(indexName, Integer.toString(i), "fooVector", IntStream.range(0, 64).mapToDouble(d -> randomFloat()).toArray()); + } + refresh(); + assertBBQIndexType(indexName, type); // test assertion to ensure that the correct index type is being used + return indexName; + } + + @SuppressWarnings("unchecked") + static void assertBBQIndexType(String indexName, String type) { + var response = indicesAdmin().prepareGetFieldMappings(indexName).setFields("fooVector").get(); + var map = (Map) response.fieldMappings(indexName, "fooVector").sourceAsMap().get("fooVector"); + assertThat((String) ((Map) map.get("index_options")).get("type"), is(equalTo(type))); + } + + @TestLogging(value = "org.elasticsearch.index.store.FsDirectoryFactory:DEBUG", reason = "to capture trace logging for direct IO") + public void testDirectIOUsed() { + try (MockLog mockLog = MockLog.capture(FsDirectoryFactory.class)) { + // we're just looking for some evidence direct IO is used (or not) + MockLog.LoggingExpectation expectation = SUPPORTED + ? new MockLog.PatternSeenEventExpectation( + "Direct IO used", + FsDirectoryFactory.class.getCanonicalName(), + Level.DEBUG, + "Opening .*\\.vec with direct IO" + ) + : new MockLog.PatternSeenEventExpectation( + "Direct IO not used", + FsDirectoryFactory.class.getCanonicalName(), + Level.DEBUG, + "Could not open .*\\.vec with direct IO" + ); + mockLog.addExpectation(expectation); + + String indexName = indexVectors(true); + + // do a search + var knn = List.of(new KnnSearchBuilder("fooVector", new VectorData(null, new byte[64]), 10, 20, 10f, null, null)); + assertHitCount(prepareSearch(indexName).setKnnSearch(knn), 10); + mockLog.assertAllExpectationsMatched(); + } + } + + @TestLogging(value = "org.elasticsearch.index.store.FsDirectoryFactory:DEBUG", reason = "to capture trace logging for direct IO") + public void testDirectIONotUsed() { + try (MockLog mockLog = MockLog.capture(FsDirectoryFactory.class)) { + // nothing about direct IO should be logged at all + MockLog.LoggingExpectation expectation = SUPPORTED + ? new MockLog.PatternNotSeenEventExpectation( + "Direct IO used", + FsDirectoryFactory.class.getCanonicalName(), + Level.DEBUG, + "Opening .*\\.vec with direct IO" + ) + : new MockLog.PatternNotSeenEventExpectation( + "Direct IO not used", + FsDirectoryFactory.class.getCanonicalName(), + Level.DEBUG, + "Could not open .*\\.vec with direct IO" + ); + mockLog.addExpectation(expectation); + + String indexName = indexVectors(false); + + // do a search + var knn = List.of(new KnnSearchBuilder("fooVector", new VectorData(null, new byte[64]), 10, 20, 10f, null, null)); + assertHitCount(prepareSearch(indexName).setKnnSearch(knn), 10); + mockLog.assertAllExpectationsMatched(); + } + } + + @Override + protected boolean addMockFSIndexStore() { + return false; // we require to always use the "real" hybrid directory + } +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java index 9a458485d838c..64796d10662e4 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java @@ -13,11 +13,11 @@ 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.lucene99.Lucene99FlatVectorsFormat; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; +import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat; import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer; +import org.elasticsearch.index.codec.vectors.es93.DirectIOCapableLucene99FlatVectorsFormat; import java.io.IOException; import java.util.Map; @@ -55,12 +55,16 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat { static final String IVF_META_EXTENSION = "mivf"; public static final int VERSION_START = 0; - public static final int VERSION_CURRENT = VERSION_START; + public static final int VERSION_DIRECT_IO = 1; + public static final int VERSION_CURRENT = VERSION_DIRECT_IO; - private static final Lucene99FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat( + private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat( FlatVectorScorerUtil.getLucene99FlatVectorsScorer() ); - private static final Map supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat); + private static final Map supportedFormats = Map.of( + rawVectorFormat.getName(), + rawVectorFormat + ); // 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 ratio. @@ -74,8 +78,13 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat { private final int vectorPerCluster; private final int centroidsPerParentCluster; + private final boolean useDirectIO; public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) { + this(vectorPerCluster, centroidsPerParentCluster, false); + } + + public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO) { super(NAME); if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) { throw new IllegalArgumentException( @@ -99,6 +108,7 @@ public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentClu } this.vectorPerCluster = vectorPerCluster; this.centroidsPerParentCluster = centroidsPerParentCluster; + this.useDirectIO = useDirectIO; } /** Constructs a format using the given graph construction parameters and scalar quantization. */ @@ -109,20 +119,34 @@ public ES920DiskBBQVectorsFormat() { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { return new ES920DiskBBQVectorsWriter( - rawVectorFormat.getName(), state, + rawVectorFormat.getName(), + useDirectIO, rawVectorFormat.fieldsWriter(state), vectorPerCluster, centroidsPerParentCluster ); } + // for testing + KnnVectorsWriter version0FieldsWriter(SegmentWriteState state) throws IOException { + return new ES920DiskBBQVectorsWriter( + state, + rawVectorFormat.getName(), + null, + rawVectorFormat.fieldsWriter(state), + vectorPerCluster, + centroidsPerParentCluster, + VERSION_START + ); + } + @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new ES920DiskBBQVectorsReader(state, f -> { + return new ES920DiskBBQVectorsReader(state, (f, dio) -> { var format = supportedFormats.get(f); if (format == null) return null; - return format.fieldsReader(state); + return format.fieldsReader(state, dio); }); } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java index aa01fa822eed3..7d8a563c99f7f 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsReader.java @@ -9,14 +9,12 @@ package org.elasticsearch.index.codec.vectors.diskbbq; -import org.apache.lucene.codecs.hnsw.FlatVectorsReader; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.KnnCollector; import org.apache.lucene.store.IndexInput; import org.apache.lucene.util.Bits; -import org.apache.lucene.util.IOFunction; import org.apache.lucene.util.VectorUtil; import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer; import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue; @@ -40,7 +38,7 @@ */ public class ES920DiskBBQVectorsReader extends IVFVectorsReader { - public ES920DiskBBQVectorsReader(SegmentReadState state, IOFunction getFormatReader) throws IOException { + ES920DiskBBQVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException { super(state, getFormatReader); } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java index d312eb373eed7..7f77f0de22bc7 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsWriter.java @@ -52,13 +52,34 @@ public class ES920DiskBBQVectorsWriter extends IVFVectorsWriter { private final int centroidsPerParentCluster; public ES920DiskBBQVectorsWriter( - String rawVectorFormatName, SegmentWriteState state, + String rawVectorFormatName, + boolean useDirectIOReads, FlatVectorsWriter rawVectorDelegate, int vectorPerCluster, int centroidsPerParentCluster ) throws IOException { - super(state, rawVectorFormatName, rawVectorDelegate); + this( + state, + rawVectorFormatName, + useDirectIOReads, + rawVectorDelegate, + vectorPerCluster, + centroidsPerParentCluster, + ES920DiskBBQVectorsFormat.VERSION_CURRENT + ); + } + + ES920DiskBBQVectorsWriter( + SegmentWriteState state, + String rawVectorFormatName, + Boolean useDirectIOReads, + FlatVectorsWriter rawVectorDelegate, + int vectorPerCluster, + int centroidsPerParentCluster, + int writeVersion + ) throws IOException { + super(state, rawVectorFormatName, useDirectIOReads, rawVectorDelegate, writeVersion); this.vectorPerCluster = vectorPerCluster; this.centroidsPerParentCluster = centroidsPerParentCluster; } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java index 68aafc9e9fd69..81c1e144effdb 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsReader.java @@ -29,7 +29,6 @@ import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.util.Bits; -import org.apache.lucene.util.IOFunction; import org.elasticsearch.core.IOUtils; import org.elasticsearch.search.vectors.IVFKnnSearchStrategy; @@ -43,20 +42,37 @@ import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS; import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.DYNAMIC_VISIT_RATIO; +import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO; /** * Reader for IVF vectors. This reader is used to read the IVF vectors from the index. */ public abstract class IVFVectorsReader extends KnnVectorsReader { + private record FlatVectorsReaderKey(String formatName, boolean useDirectIO) { + private FlatVectorsReaderKey(FieldEntry entry) { + this(entry.rawVectorFormatName, entry.useDirectIOReads); + } + + @Override + public String toString() { + return formatName + (useDirectIO ? " with Direct IO" : ""); + } + } + private final IndexInput ivfCentroids, ivfClusters; private final SegmentReadState state; private final FieldInfos fieldInfos; protected final IntObjectHashMap fields; - private final Map rawVectorReaders; + private final Map rawVectorReaders; + + @FunctionalInterface + public interface GetFormatReader { + FlatVectorsReader getReader(String formatName, boolean useDirectIO) throws IOException; + } @SuppressWarnings("this-escape") - protected IVFVectorsReader(SegmentReadState state, IOFunction getFormatReader) throws IOException { + protected IVFVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException { this.state = state; this.fieldInfos = state.fieldInfos; this.fields = new IntObjectHashMap<>(); @@ -70,7 +86,7 @@ protected IVFVectorsReader(SegmentReadState state, IOFunction readers = null; + Map readers = null; try { versionMeta = CodecUtil.checkIndexHeader( ivfMeta, @@ -80,7 +96,7 @@ protected IVFVectorsReader(SegmentReadState state, IOFunction readFields(ChecksumIndexInput meta, IOFunction loadReader) + private Map readFields(ChecksumIndexInput meta, GetFormatReader loadReader, int versionMeta) throws IOException { - Map readers = new HashMap<>(); + Map readers = new HashMap<>(); for (int fieldNumber = meta.readInt(); fieldNumber != -1; fieldNumber = meta.readInt()) { final FieldInfo info = fieldInfos.fieldInfo(fieldNumber); if (info == null) { throw new CorruptIndexException("Invalid field number: " + fieldNumber, meta); } - FieldEntry fieldEntry = readField(meta, info); + FieldEntry fieldEntry = readField(meta, info, versionMeta); + FlatVectorsReaderKey key = new FlatVectorsReaderKey(fieldEntry); - FlatVectorsReader reader = readers.get(fieldEntry.rawVectorFormatName); + FlatVectorsReader reader = readers.get(key); if (reader == null) { - reader = loadReader.apply(fieldEntry.rawVectorFormatName); + reader = loadReader.getReader(fieldEntry.rawVectorFormatName, fieldEntry.useDirectIOReads); if (reader == null) { throw new IllegalStateException("Cannot find flat vector format: " + fieldEntry.rawVectorFormatName); } - readers.put(fieldEntry.rawVectorFormatName, reader); + readers.put(key, reader); } fields.put(info.number, fieldEntry); @@ -178,8 +195,9 @@ private Map readFields(ChecksumIndexInput meta, IOFun return readers; } - private FieldEntry readField(IndexInput input, FieldInfo info) throws IOException { + private FieldEntry readField(IndexInput input, FieldInfo info, int versionMeta) throws IOException { final String rawVectorFormat = input.readString(); + final boolean useDirectIOReads = versionMeta >= VERSION_DIRECT_IO && input.readByte() == 1; final VectorEncoding vectorEncoding = readVectorEncoding(input); final VectorSimilarityFunction similarityFunction = readSimilarityFunction(input); if (similarityFunction != info.getVectorSimilarityFunction()) { @@ -207,6 +225,7 @@ private FieldEntry readField(IndexInput input, FieldInfo info) throws IOExceptio } return new FieldEntry( rawVectorFormat, + useDirectIOReads, similarityFunction, vectorEncoding, numCentroids, @@ -254,10 +273,10 @@ private FieldEntry getFieldEntryOrThrow(String field) { } private FlatVectorsReader getReaderForField(String field) { - var formatName = getFieldEntryOrThrow(field).rawVectorFormatName; - FlatVectorsReader reader = rawVectorReaders.get(formatName); + var readerKey = new FlatVectorsReaderKey(getFieldEntryOrThrow(field)); + FlatVectorsReader reader = rawVectorReaders.get(readerKey); if (reader == null) throw new IllegalArgumentException( - "Could not find raw vector format [" + formatName + "] for field [" + field + "]" + "Could not find raw vector format [" + readerKey + "] for field [" + field + "]" ); return reader; } @@ -387,6 +406,7 @@ public void close() throws IOException { protected record FieldEntry( String rawVectorFormatName, + boolean useDirectIOReads, VectorSimilarityFunction similarityFunction, VectorEncoding vectorEncoding, int numCentroids, diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsWriter.java index a14849c382571..b24255ef3baa8 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsWriter.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsWriter.java @@ -53,12 +53,26 @@ public abstract class IVFVectorsWriter extends KnnVectorsWriter { private final IndexOutput ivfCentroids, ivfClusters; private final IndexOutput ivfMeta; private final String rawVectorFormatName; + private final int writeVersion; + private final Boolean useDirectIOReads; private final FlatVectorsWriter rawVectorDelegate; @SuppressWarnings("this-escape") - protected IVFVectorsWriter(SegmentWriteState state, String rawVectorFormatName, FlatVectorsWriter rawVectorDelegate) - throws IOException { + protected IVFVectorsWriter( + SegmentWriteState state, + String rawVectorFormatName, + Boolean useDirectIOReads, + FlatVectorsWriter rawVectorDelegate, + int writeVersion + ) throws IOException { + // if version >= VERSION_DIRECT_IO, useDirectIOReads should have a value + if ((writeVersion >= ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO) == (useDirectIOReads == null)) throw new IllegalArgumentException( + "Write version " + writeVersion + " does not match direct IO value " + useDirectIOReads + ); + this.rawVectorFormatName = rawVectorFormatName; + this.writeVersion = writeVersion; + this.useDirectIOReads = useDirectIOReads; this.rawVectorDelegate = rawVectorDelegate; final String metaFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, @@ -82,7 +96,7 @@ protected IVFVectorsWriter(SegmentWriteState state, String rawVectorFormatName, CodecUtil.writeIndexHeader( ivfMeta, ES920DiskBBQVectorsFormat.NAME, - ES920DiskBBQVectorsFormat.VERSION_CURRENT, + writeVersion, state.segmentInfo.getId(), state.segmentSuffix ); @@ -90,7 +104,7 @@ protected IVFVectorsWriter(SegmentWriteState state, String rawVectorFormatName, CodecUtil.writeIndexHeader( ivfCentroids, ES920DiskBBQVectorsFormat.NAME, - ES920DiskBBQVectorsFormat.VERSION_CURRENT, + writeVersion, state.segmentInfo.getId(), state.segmentSuffix ); @@ -98,7 +112,7 @@ protected IVFVectorsWriter(SegmentWriteState state, String rawVectorFormatName, CodecUtil.writeIndexHeader( ivfClusters, ES920DiskBBQVectorsFormat.NAME, - ES920DiskBBQVectorsFormat.VERSION_CURRENT, + writeVersion, state.segmentInfo.getId(), state.segmentSuffix ); @@ -497,6 +511,9 @@ private void writeMeta( ) throws IOException { ivfMeta.writeInt(field.number); ivfMeta.writeString(rawVectorFormatName); + if (writeVersion >= ES920DiskBBQVectorsFormat.VERSION_DIRECT_IO) { + ivfMeta.writeByte(useDirectIOReads ? (byte) 1 : 0); + } ivfMeta.writeInt(field.getVectorEncoding().ordinal()); ivfMeta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction())); ivfMeta.writeInt(numCentroids); diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsFormat.java index 9a7d6e1db1009..c5900d83f7997 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsFormat.java @@ -13,11 +13,11 @@ 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.lucene99.Lucene99FlatVectorsFormat; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; +import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat; import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer; +import org.elasticsearch.index.codec.vectors.es93.DirectIOCapableLucene99FlatVectorsFormat; import java.io.IOException; import java.util.Map; @@ -54,13 +54,16 @@ public class ESNextDiskBBQVectorsFormat extends KnnVectorsFormat { public static final String CLUSTER_EXTENSION = "clivf"; static final String IVF_META_EXTENSION = "mivf"; - public static final int VERSION_START = 0; + public static final int VERSION_START = 1; public static final int VERSION_CURRENT = VERSION_START; - private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat( + private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat( FlatVectorScorerUtil.getLucene99FlatVectorsScorer() ); - private static final Map supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat); + private static final Map supportedFormats = Map.of( + rawVectorFormat.getName(), + rawVectorFormat + ); // 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 ratio. @@ -74,8 +77,13 @@ public class ESNextDiskBBQVectorsFormat extends KnnVectorsFormat { private final int vectorPerCluster; private final int centroidsPerParentCluster; + private final boolean useDirectIO; public ESNextDiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) { + this(vectorPerCluster, centroidsPerParentCluster, false); + } + + public ESNextDiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO) { super(NAME); if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) { throw new IllegalArgumentException( @@ -99,6 +107,7 @@ public ESNextDiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCl } this.vectorPerCluster = vectorPerCluster; this.centroidsPerParentCluster = centroidsPerParentCluster; + this.useDirectIO = useDirectIO; } /** Constructs a format using the given graph construction parameters and scalar quantization. */ @@ -109,8 +118,9 @@ public ESNextDiskBBQVectorsFormat() { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { return new ESNextDiskBBQVectorsWriter( - rawVectorFormat.getName(), state, + rawVectorFormat.getName(), + useDirectIO, rawVectorFormat.fieldsWriter(state), vectorPerCluster, centroidsPerParentCluster @@ -119,10 +129,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { - return new ESNextDiskBBQVectorsReader(state, f -> { + return new ESNextDiskBBQVectorsReader(state, (f, dio) -> { var format = supportedFormats.get(f); if (format == null) return null; - return format.fieldsReader(state); + return format.fieldsReader(state, dio); }); } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsReader.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsReader.java index b21a1b9055876..13e343e5273b9 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsReader.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsReader.java @@ -9,14 +9,12 @@ package org.elasticsearch.index.codec.vectors.diskbbq.next; -import org.apache.lucene.codecs.hnsw.FlatVectorsReader; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.KnnCollector; import org.apache.lucene.store.IndexInput; import org.apache.lucene.util.Bits; -import org.apache.lucene.util.IOFunction; import org.apache.lucene.util.VectorUtil; import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer; import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue; @@ -42,7 +40,7 @@ */ public class ESNextDiskBBQVectorsReader extends IVFVectorsReader { - public ESNextDiskBBQVectorsReader(SegmentReadState state, IOFunction getFormatReader) throws IOException { + public ESNextDiskBBQVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException { super(state, getFormatReader); } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsWriter.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsWriter.java index 6799743d9dc9f..13d6fb4cd93da 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsWriter.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQVectorsWriter.java @@ -60,13 +60,14 @@ public class ESNextDiskBBQVectorsWriter extends IVFVectorsWriter { private final int centroidsPerParentCluster; public ESNextDiskBBQVectorsWriter( - String rawVectorFormatName, SegmentWriteState state, + String rawVectorFormatName, + boolean useDirectIOReads, FlatVectorsWriter rawVectorDelegate, int vectorPerCluster, int centroidsPerParentCluster ) throws IOException { - super(state, rawVectorFormatName, rawVectorDelegate); + super(state, rawVectorFormatName, useDirectIOReads, rawVectorDelegate, ESNextDiskBBQVectorsFormat.VERSION_CURRENT); this.vectorPerCluster = vectorPerCluster; this.centroidsPerParentCluster = centroidsPerParentCluster; } 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 4fe064b1af067..914acd791e0ee 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/MapperFeatures.java @@ -52,6 +52,7 @@ public class MapperFeatures implements FeatureSpecification { static final NodeFeature IGNORED_SOURCE_FIELDS_PER_ENTRY = new NodeFeature("mapper.ignored_source_fields_per_entry"); public static final NodeFeature MULTI_FIELD_UNICODE_OPTIMISATION_FIX = new NodeFeature("mapper.multi_field.unicode_optimisation_fix"); static final NodeFeature PATTERN_TEXT_RENAME = new NodeFeature("mapper.pattern_text_rename"); + public static final NodeFeature DISKBBQ_ON_DISK_RESCORING = new NodeFeature("mapper.vectors.diskbbq_on_disk_rescoring"); @Override public Set getTestFeatures() { @@ -89,7 +90,8 @@ public Set getTestFeatures() { IGNORED_SOURCE_FIELDS_PER_ENTRY, MULTI_FIELD_UNICODE_OPTIMISATION_FIX, MATCH_ONLY_TEXT_BLOCK_LOADER_FIX, - PATTERN_TEXT_RENAME + PATTERN_TEXT_RENAME, + DISKBBQ_ON_DISK_RESCORING ); } } 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 8f92f1d0628ba..5e4012a5e53e1 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 @@ -1458,6 +1458,7 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map