Add on-disk rescoring to disk BBQ

docs/changelog/135778.yaml

@@ -0,0 +1,5 @@
+pr: 135778
+summary: Add on-disk rescoring to disk BBQ
+area: Vector Search
+type: feature
+issues: []

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:

server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java

@@ -0,0 +1,175 @@
+/*
+ * 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 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;
+            }
+        };
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return List.of(InternalSettingsPlugin.class);
+    }
+
+    private String indexVectors(boolean directIO) {
+        String indexName = "test-vectors-" + directIO;
+        String type = randomFrom("bbq_disk");
+        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<String, Object>) response.fieldMappings(indexName, "fooVector").sourceAsMap().get("fooVector");
+        assertThat((String) ((Map<String, Object>) 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
+    }
+}

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<String, FlatVectorsFormat> supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat);
+    private static final Map<String, DirectIOCapableFlatVectorsFormat> 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,8 +119,9 @@ public ES920DiskBBQVectorsFormat() {
     @Override
     public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new ES920DiskBBQVectorsWriter(
-            rawVectorFormat.getName(),
             state,
+            rawVectorFormat.getName(),
+            useDirectIO,
             rawVectorFormat.fieldsWriter(state),
             vectorPerCluster,
             centroidsPerParentCluster
@@ -119,10 +130,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
 
     @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);
         });
     }
 

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<String, FlatVectorsReader> getFormatReader) throws IOException {
+    ES920DiskBBQVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException {
         super(state, getFormatReader);
     }