Add a direct IO option to rescore_vector for bbq_hnsw

qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/upgrades/VectorSearchIT.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
+import org.elasticsearch.search.SearchFeatures;
 
 import java.io.IOException;
 import java.util.List;
@@ -35,6 +36,7 @@ public VectorSearchIT(@Name("upgradedNodes") int upgradedNodes) {
     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 BBQ_INDEX_NAME_RESCORE = "bbq_vector_index_rescore";
     private static final String FLAT_QUANTIZED_INDEX_NAME = "flat_quantized_vector_index";
     private static final String FLAT_BBQ_INDEX_NAME = "flat_bbq_vector_index";
 
@@ -507,6 +509,63 @@ public void testBBQVectorSearch() throws Exception {
         );
     }
 
+    public void testBBQVectorSearchOffheapRescoring() throws Exception {
+        assumeTrue("Disabling off-heap rescoring is not supported", oldClusterHasFeature(SearchFeatures.BBQ_OFFHEAP_RESCORING));
+        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,
+                        "disable_offheap_cache_rescoring": true
+                      }
+                    }
+                  }
+                }
+                """;
+            // create index and index 10 random floating point vectors
+            createIndex(
+                BBQ_INDEX_NAME_RESCORE,
+                Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0).build(),
+                mapping
+            );
+            index64DimVectors(BBQ_INDEX_NAME_RESCORE);
+            // force merge the index
+            client().performRequest(new Request("POST", "/" + BBQ_INDEX_NAME_RESCORE + "/_forcemerge?max_num_segments=1"));
+        }
+        Request searchRequest = new Request("POST", "/" + BBQ_INDEX_NAME_RESCORE + "/_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,
+                  "rescore_vector": {
+                    "oversample": 2.0
+                  }
+                }
+            }
+            """);
+        Map<String, Object> response = search(searchRequest);
+        assertThat(extractValue(response, "hits.total.value"), equalTo(2));
+        List<Map<String, Object>> hits = extractValue(response, "hits.hits");
+        assertThat("expected: 0 received" + hits.get(0).get("_id") + " hits: " + response, hits.get(0).get("_id"), equalTo("0"));
+        assertThat(
+            "expected_near: 0.99 received" + hits.get(0).get("_score") + "hits: " + response,
+            (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", oldClusterHasFeature(BBQ_VECTOR_SEARCH_TEST_FEATURE));
         if (isOldCluster()) {

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -95,7 +95,7 @@ static Codec createCodec(CmdLineArgs args) {
                 if (args.indexType() == IndexType.FLAT) {
                     format = new ES818BinaryQuantizedVectorsFormat();
                 } else {
-                    format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
+                    format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, false, null);
                 }
             } else if (args.quantizeBits() < 32) {
                 if (args.indexType() == IndexType.FLAT) {

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml

@@ -338,6 +338,69 @@ setup:
   - match: { hits.hits.1._score: $rescore_score1 }
   - match: { hits.hits.2._score: $rescore_score2 }
 ---
+"Test index configured rescore vector with no off-heap scoring":
+  - requires:
+      cluster_features: ["search.vectors.bbq_offheap_rescoring"]
+      reason: Needs bbq_offheap_rescoring feature
+  - skip:
+      features: "headers"
+  - do:
+      indices.create:
+        index: bbq_rescore_hnsw
+        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_hnsw
+                  disable_offheap_cache_rescoring: true
+                  rescore_vector:
+                    oversample: 1.5
+
+  - do:
+      bulk:
+        index: bbq_rescore_hnsw
+        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_hnsw
+        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 }
+---
 "Test index configured rescore vector updateable and settable to 0":
   - requires:
       cluster_features: ["mapper.dense_vector.rescore_zero_vector"]

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java

@@ -68,7 +68,6 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
     }
 
     static boolean shouldUseDirectIO(SegmentReadState state) {
-        assert ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO;
         return FsDirectoryFactory.isHybridFs(state.directory);
     }
 

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java

@@ -26,7 +26,6 @@
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
-import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 
 import java.io.IOException;
@@ -88,8 +87,6 @@
  */
 public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
 
-    public static final boolean USE_DIRECT_IO = getUseDirectIO();
-
     public static final String BINARIZED_VECTOR_COMPONENT = "BVEC";
     public static final String NAME = "ES818BinaryQuantizedVectorsFormat";
 
@@ -101,24 +98,24 @@ public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
     static final String VECTOR_DATA_EXTENSION = "veb";
     static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
-    @SuppressForbidden(
-        reason = "TODO Deprecate any lenient usage of Boolean#parseBoolean https://github.com/elastic/elasticsearch/issues/128993"
-    )
-    private static boolean getUseDirectIO() {
-        return Boolean.parseBoolean(System.getProperty("vector.rescoring.directio", "false"));
-    }
-
-    private static final FlatVectorsFormat rawVectorFormat = USE_DIRECT_IO
-        ? new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
-        : new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
-
     private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
+    private final FlatVectorsFormat rawVectorFormat;
+
     /** Creates a new instance with the default number of vectors per cluster. */
     public ES818BinaryQuantizedVectorsFormat() {
+        this(false);
+    }
+
+    /** Creates a new instance with the default number of vectors per cluster,
+     * and whether direct IO should be used to access raw vectors. */
+    public ES818BinaryQuantizedVectorsFormat(boolean useDirectIO) {
         super(NAME);
+        rawVectorFormat = useDirectIO
+            ? new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
+            : new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormat.java

@@ -62,14 +62,14 @@ public class ES818HnswBinaryQuantizedVectorsFormat extends KnnVectorsFormat {
     private final int beamWidth;
 
     /** The format for storing, reading, merging vectors on disk */
-    private static final FlatVectorsFormat flatVectorsFormat = new ES818BinaryQuantizedVectorsFormat();
+    private final FlatVectorsFormat flatVectorsFormat;
 
     private final int numMergeWorkers;
     private final TaskExecutor mergeExec;
 
     /** Constructs a format using default graph construction parameters */
     public ES818HnswBinaryQuantizedVectorsFormat() {
-        this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null);
+        this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, false, null);
     }
 
     /**
@@ -79,7 +79,18 @@ public ES818HnswBinaryQuantizedVectorsFormat() {
      * @param beamWidth the size of the queue maintained during graph construction.
      */
     public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
-        this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null);
+        this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, false, null);
+    }
+
+    /**
+     * Constructs a format using the given graph construction parameters.
+     *
+     * @param maxConn the maximum number of connections to a node in the HNSW graph
+     * @param beamWidth the size of the queue maintained during graph construction.
+     * @param useDirectIO whether direct IO should be used to access raw vectors
+     */
+    public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, boolean useDirectIO) {
+        this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, useDirectIO, null);
     }
 
     /**
@@ -92,7 +103,13 @@ public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
      * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
      *     generated by this format to do the merge
      */
-    public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
+    public ES818HnswBinaryQuantizedVectorsFormat(
+        int maxConn,
+        int beamWidth,
+        int numMergeWorkers,
+        boolean useDirectIO,
+        ExecutorService mergeExec
+    ) {
         super(NAME);
         if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
             throw new IllegalArgumentException(
@@ -110,6 +127,9 @@ public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, int num
             throw new IllegalArgumentException("No executor service is needed as we'll use single thread to merge");
         }
         this.numMergeWorkers = numMergeWorkers;
+
+        flatVectorsFormat = new ES818BinaryQuantizedVectorsFormat(useDirectIO);
+
         if (mergeExec != null) {
             this.mergeExec = new TaskExecutor(mergeExec);
         } else {

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -388,7 +388,8 @@ private DenseVectorIndexOptions defaultIndexOptions(boolean defaultInt8Hnsw, boo
                 return new BBQHnswIndexOptions(
                     Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN,
                     Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH,
-                    new RescoreVector(DEFAULT_OVERSAMPLE)
+                    new RescoreVector(DEFAULT_OVERSAMPLE),
+                    false
                 );
             } else if (defaultInt8Hnsw) {
                 return new Int8HnswIndexOptions(
@@ -1622,6 +1623,8 @@ public boolean supportsDimension(int dims) {
             public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map<String, ?> indexOptionsMap, IndexVersion indexVersion) {
                 Object mNode = indexOptionsMap.remove("m");
                 Object efConstructionNode = indexOptionsMap.remove("ef_construction");
+                Object disableOffheapCacheRescoringNode = indexOptionsMap.remove("disable_offheap_cache_rescoring");
+
                 if (mNode == null) {
                     mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
                 }
@@ -1630,15 +1633,19 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map<String, ?
                 }
                 int m = XContentMapValues.nodeIntegerValue(mNode);
                 int efConstruction = XContentMapValues.nodeIntegerValue(efConstructionNode);
+
                 RescoreVector rescoreVector = null;
                 if (hasRescoreIndexVersion(indexVersion)) {
                     rescoreVector = RescoreVector.fromIndexOptions(indexOptionsMap, indexVersion);
                     if (rescoreVector == null && defaultOversampleForBBQ(indexVersion)) {
                         rescoreVector = new RescoreVector(DEFAULT_OVERSAMPLE);
                     }
                 }
+
+                boolean disableOffheapCacheRescoring = XContentMapValues.nodeBooleanValue(disableOffheapCacheRescoringNode, false);
+
                 MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
-                return new BBQHnswIndexOptions(m, efConstruction, rescoreVector);
+                return new BBQHnswIndexOptions(m, efConstruction, rescoreVector, disableOffheapCacheRescoring);
             }
 
             @Override
@@ -2174,17 +2181,19 @@ public String toString() {
     public static class BBQHnswIndexOptions extends QuantizedIndexOptions {
         private final int m;
         private final int efConstruction;
+        private final boolean disableOffheapCacheRescoring;
 
-        public BBQHnswIndexOptions(int m, int efConstruction, RescoreVector rescoreVector) {
+        public BBQHnswIndexOptions(int m, int efConstruction, RescoreVector rescoreVector, boolean disableOffheapCacheRescoring) {
             super(VectorIndexType.BBQ_HNSW, rescoreVector);
             this.m = m;
             this.efConstruction = efConstruction;
+            this.disableOffheapCacheRescoring = disableOffheapCacheRescoring;
         }
 
         @Override
         KnnVectorsFormat getVectorsFormat(ElementType elementType) {
             assert elementType == ElementType.FLOAT;
-            return new ES818HnswBinaryQuantizedVectorsFormat(m, efConstruction);
+            return new ES818HnswBinaryQuantizedVectorsFormat(m, efConstruction, disableOffheapCacheRescoring);
         }
 
         @Override
@@ -2195,12 +2204,15 @@ public boolean updatableTo(DenseVectorIndexOptions update) {
         @Override
         boolean doEquals(DenseVectorIndexOptions other) {
             BBQHnswIndexOptions that = (BBQHnswIndexOptions) other;
-            return m == that.m && efConstruction == that.efConstruction && Objects.equals(rescoreVector, that.rescoreVector);
+            return m == that.m
+                && efConstruction == that.efConstruction
+                && Objects.equals(rescoreVector, that.rescoreVector)
+                && disableOffheapCacheRescoring == that.disableOffheapCacheRescoring;
         }
 
         @Override
         int doHashCode() {
-            return Objects.hash(m, efConstruction, rescoreVector);
+            return Objects.hash(m, efConstruction, rescoreVector, disableOffheapCacheRescoring);
         }
 
         @Override
@@ -2214,6 +2226,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field("type", type);
             builder.field("m", m);
             builder.field("ef_construction", efConstruction);
+            if (disableOffheapCacheRescoring) {
+                builder.field("disable_offheap_cache_rescoring", true);
+            }
             if (rescoreVector != null) {
                 rescoreVector.toXContent(builder, params);
             }

server/src/main/java/org/elasticsearch/search/SearchFeatures.java

@@ -33,6 +33,7 @@ public Set<NodeFeature> getFeatures() {
     static final NodeFeature MULTI_MATCH_CHECKS_POSITIONS = new NodeFeature("search.multi.match.checks.positions");
     public static final NodeFeature BBQ_HNSW_DEFAULT_INDEXING = new NodeFeature("search.vectors.mappers.default_bbq_hnsw");
     public static final NodeFeature SEARCH_WITH_NO_DIMENSIONS_BUGFIX = new NodeFeature("search.vectors.no_dimensions_bugfix");
+    public static final NodeFeature BBQ_OFFHEAP_RESCORING = new NodeFeature("search.vectors.bbq_offheap_rescoring");
 
     @Override
     public Set<NodeFeature> getTestFeatures() {
@@ -43,7 +44,8 @@ public Set<NodeFeature> getTestFeatures() {
             INT_SORT_FOR_INT_SHORT_BYTE_FIELDS,
             MULTI_MATCH_CHECKS_POSITIONS,
             BBQ_HNSW_DEFAULT_INDEXING,
-            SEARCH_WITH_NO_DIMENSIONS_BUGFIX
+            SEARCH_WITH_NO_DIMENSIONS_BUGFIX,
+            BBQ_OFFHEAP_RESCORING
         );
     }
 }