From 73e3cc0110fa6e05e9ea04f117facd188ff7d270 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Wed, 9 Jul 2025 10:50:44 +0100 Subject: [PATCH 1/3] Add a direct IO option to rescore_vector for bbq_hnsw --- .../test/knn/KnnIndexTester.java | 2 +- .../DirectIOLucene99FlatVectorsFormat.java | 1 - .../ES818BinaryQuantizedVectorsFormat.java | 17 ++-- ...ES818HnswBinaryQuantizedVectorsFormat.java | 28 +++++- .../vectors/DenseVectorFieldMapper.java | 56 ++++++------ ...S818BinaryQuantizedVectorsFormatTests.java | 57 ------------ ...ctIOBinaryQuantizedVectorsFormatTests.java | 85 ++++++++++++++++++ ...HnswBinaryQuantizedVectorsFormatTests.java | 90 +++++++++++++++++++ ...HnswBinaryQuantizedVectorsFormatTests.java | 61 +------------ .../vectors/DenseVectorFieldTypeTests.java | 9 +- .../mapper/SemanticTextFieldMapper.java | 3 +- .../mapper/SemanticTextFieldMapperTests.java | 4 +- 12 files changed, 252 insertions(+), 161 deletions(-) create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOBinaryQuantizedVectorsFormatTests.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOHnswBinaryQuantizedVectorsFormatTests.java diff --git a/qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java b/qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java index afa155dfb2d21..3bdb6dea70d86 100644 --- a/qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java +++ b/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) { diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java index 8e328b5c500ad..2a542f6b77864 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java +++ b/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); } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java index 146164b55f00a..4baa18b580173 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java @@ -87,8 +87,6 @@ */ public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat { - public static final boolean USE_DIRECT_IO = Boolean.parseBoolean(System.getProperty("vector.rescoring.directio", "false")); - public static final String BINARIZED_VECTOR_COMPONENT = "BVEC"; public static final String NAME = "ES818BinaryQuantizedVectorsFormat"; @@ -100,17 +98,24 @@ public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat { static final String VECTOR_DATA_EXTENSION = "veb"; static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; - 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 diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormat.java index 56942017c3cef..b6c76ced89d56 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormat.java +++ b/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 { 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 4d1c4fc41526c..91a86c527cd4e 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 @@ -387,7 +387,7 @@ private DenseVectorIndexOptions defaultIndexOptions(boolean defaultInt8Hnsw, boo return new BBQHnswIndexOptions( Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN, Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH, - new RescoreVector(DEFAULT_OVERSAMPLE) + null ); } else if (defaultInt8Hnsw) { return new Int8HnswIndexOptions( @@ -1632,9 +1632,6 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map indexOptionsMap, IndexVersion indexVersion) { Object rescoreVectorNode = indexOptionsMap.remove(NAME); @@ -2352,26 +2345,35 @@ static RescoreVector fromIndexOptions(Map indexOptionsMap, IndexVersi return null; } Map mappedNode = XContentMapValues.nodeMapValue(rescoreVectorNode, NAME); + + Float oversampleValue = null; Object oversampleNode = mappedNode.get(OVERSAMPLE); - if (oversampleNode == null) { - throw new IllegalArgumentException("Invalid rescore_vector value. Missing required field " + OVERSAMPLE); - } - float oversampleValue = (float) XContentMapValues.nodeDoubleValue(oversampleNode); - if (oversampleValue == 0 && allowsZeroRescore(indexVersion) == false) { - throw new IllegalArgumentException("oversample must be greater than 1"); - } - if (oversampleValue < 1 && oversampleValue != 0) { - throw new IllegalArgumentException("oversample must be greater than 1 or exactly 0"); - } else if (oversampleValue > 10) { - throw new IllegalArgumentException("oversample must be less than or equal to 10"); + if (oversampleNode != null) { + oversampleValue = (float) XContentMapValues.nodeDoubleValue(oversampleNode); + if (oversampleValue == 0 && allowsZeroRescore(indexVersion) == false) { + throw new IllegalArgumentException("oversample must be greater than 1"); + } + if (oversampleValue < 1 && oversampleValue != 0) { + throw new IllegalArgumentException("oversample must be greater than 1 or exactly 0"); + } else if (oversampleValue > 10) { + throw new IllegalArgumentException("oversample must be less than or equal to 10"); + } } - return new RescoreVector(oversampleValue); + + Boolean directIO = (Boolean) mappedNode.get(DIRECT_IO); + + return new RescoreVector(oversampleValue, directIO); } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(NAME); - builder.field(OVERSAMPLE, oversample); + if (oversample != null) { + builder.field(OVERSAMPLE, oversample); + } + if (useDirectIO != null) { + builder.field(DIRECT_IO, useDirectIO); + } builder.endObject(); return builder; } @@ -2710,6 +2712,10 @@ && isNotUnitVector(squaredMagnitude)) { && quantizedIndexOptions.rescoreVector != null) { oversample = quantizedIndexOptions.rescoreVector.oversample; } + if (oversample == null) { + oversample = DEFAULT_OVERSAMPLE; + } + boolean rescore = needsRescore(oversample); if (rescore) { // Will get k * oversample for rescoring, and get the top k diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormatTests.java index 4a82a7e3f13e6..d0a5ddcf8a12f 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormatTests.java @@ -38,7 +38,6 @@ import org.apache.lucene.index.SoftDeletesRetentionMergePolicy; import org.apache.lucene.index.Term; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.misc.store.DirectIODirectory; import org.apache.lucene.search.FieldExistsQuery; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; @@ -52,32 +51,19 @@ import org.apache.lucene.search.join.DiversifyingChildrenFloatKnnVectorQuery; import org.apache.lucene.search.join.QueryBitSetProducer; 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.store.MMapDirectory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.store.MockDirectoryWrapper; import org.apache.lucene.tests.util.TestUtil; import org.elasticsearch.common.logging.LogConfigurator; -import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.index.IndexModule; -import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.codec.vectors.BQVectorUtils; import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer; -import org.elasticsearch.index.shard.ShardId; -import org.elasticsearch.index.shard.ShardPath; -import org.elasticsearch.index.store.FsDirectoryFactory; -import org.elasticsearch.test.IndexSettingsModule; import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Locale; -import java.util.OptionalLong; import static java.lang.String.format; import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; @@ -268,14 +254,6 @@ public void testSimpleOffHeapSize() throws IOException { } } - public void testSimpleOffHeapSizeFSDir() throws IOException { - checkDirectIOSupported(); - var config = newIndexWriterConfig().setUseCompoundFile(false); // avoid compound files to allow directIO - try (Directory dir = newFSDirectory()) { - testSimpleOffHeapSizeImpl(dir, config, false); - } - } - public void testSimpleOffHeapSizeMMapDir() throws IOException { try (Directory dir = newMMapDirectory()) { testSimpleOffHeapSizeImpl(dir, newIndexWriterConfig(), true); @@ -315,39 +293,4 @@ static Directory newMMapDirectory() throws IOException { } return dir; } - - private Directory newFSDirectory() throws IOException { - Settings settings = Settings.builder() - .put(IndexModule.INDEX_STORE_TYPE_SETTING.getKey(), IndexModule.Type.HYBRIDFS.name().toLowerCase(Locale.ROOT)) - .build(); - IndexSettings idxSettings = IndexSettingsModule.newIndexSettings("foo", settings); - Path tempDir = createTempDir().resolve(idxSettings.getUUID()).resolve("0"); - Files.createDirectories(tempDir); - ShardPath path = new ShardPath(false, tempDir, tempDir, new ShardId(idxSettings.getIndex(), 0)); - Directory dir = (new FsDirectoryFactory()).newDirectory(idxSettings, path); - if (random().nextBoolean()) { - dir = new MockDirectoryWrapper(random(), dir); - } - return dir; - } - - static void checkDirectIOSupported() { - assumeTrue("Direct IO is not enabled", ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO); - - Path path = createTempDir("directIOProbe"); - try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { - out.writeString("test"); - } catch (IOException e) { - assumeNoException("test requires a filesystem that supports Direct IO", e); - } - } - - 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; - } - }; - } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOBinaryQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOBinaryQuantizedVectorsFormatTests.java new file mode 100644 index 0000000000000..4f27815a576df --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOBinaryQuantizedVectorsFormatTests.java @@ -0,0 +1,85 @@ +/* + * 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.codec.vectors.es818; + +import org.apache.lucene.codecs.Codec; +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.store.MockDirectoryWrapper; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.IndexModule; +import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.index.shard.ShardId; +import org.elasticsearch.index.shard.ShardPath; +import org.elasticsearch.index.store.FsDirectoryFactory; +import org.elasticsearch.test.IndexSettingsModule; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Locale; +import java.util.OptionalLong; + +public class ES818DirectIOBinaryQuantizedVectorsFormatTests extends ES818BinaryQuantizedVectorsFormatTests { + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new ES818BinaryQuantizedVectorsFormat(true)); + + @Override + protected Codec getCodec() { + return codec; + } + + @BeforeClass + public static void checkDirectIOSupport() { + Path path = createTempDir("directIOProbe"); + try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { + out.writeString("test"); + } catch (IOException e) { + assumeNoException("test requires a filesystem that supports Direct IO", e); + } + } + + 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 + public void testSimpleOffHeapSize() throws IOException { + var config = newIndexWriterConfig().setUseCompoundFile(false); // avoid compound files to allow directIO + try (Directory dir = newFSDirectory()) { + testSimpleOffHeapSizeImpl(dir, config, false); + } + } + + private Directory newFSDirectory() throws IOException { + Settings settings = Settings.builder() + .put(IndexModule.INDEX_STORE_TYPE_SETTING.getKey(), IndexModule.Type.HYBRIDFS.name().toLowerCase(Locale.ROOT)) + .build(); + IndexSettings idxSettings = IndexSettingsModule.newIndexSettings("foo", settings); + Path tempDir = createTempDir().resolve(idxSettings.getUUID()).resolve("0"); + Files.createDirectories(tempDir); + ShardPath path = new ShardPath(false, tempDir, tempDir, new ShardId(idxSettings.getIndex(), 0)); + Directory dir = (new FsDirectoryFactory()).newDirectory(idxSettings, path); + if (random().nextBoolean()) { + dir = new MockDirectoryWrapper(random(), dir); + } + return dir; + } +} diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOHnswBinaryQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOHnswBinaryQuantizedVectorsFormatTests.java new file mode 100644 index 0000000000000..7140947049bb8 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOHnswBinaryQuantizedVectorsFormatTests.java @@ -0,0 +1,90 @@ +/* + * 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.codec.vectors.es818; + +import org.apache.lucene.codecs.Codec; +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.store.MockDirectoryWrapper; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.IndexModule; +import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.index.shard.ShardId; +import org.elasticsearch.index.shard.ShardPath; +import org.elasticsearch.index.store.FsDirectoryFactory; +import org.elasticsearch.test.IndexSettingsModule; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Locale; +import java.util.OptionalLong; + +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; + +public class ES818DirectIOHnswBinaryQuantizedVectorsFormatTests extends ES818HnswBinaryQuantizedVectorsFormatTests { + + static final Codec codec = TestUtil.alwaysKnnVectorsFormat( + new ES818HnswBinaryQuantizedVectorsFormat(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, true) + ); + + @Override + protected Codec getCodec() { + return codec; + } + + @BeforeClass + public static void checkDirectIOSupport() { + Path path = createTempDir("directIOProbe"); + try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { + out.writeString("test"); + } catch (IOException e) { + assumeNoException("test requires a filesystem that supports Direct IO", e); + } + } + + 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 + public void testSimpleOffHeapSize() throws IOException { + var config = newIndexWriterConfig().setUseCompoundFile(false); // avoid compound files to allow directIO + try (Directory dir = newFSDirectory()) { + testSimpleOffHeapSizeImpl(dir, config, false); + } + } + + private Directory newFSDirectory() throws IOException { + Settings settings = Settings.builder() + .put(IndexModule.INDEX_STORE_TYPE_SETTING.getKey(), IndexModule.Type.HYBRIDFS.name().toLowerCase(Locale.ROOT)) + .build(); + IndexSettings idxSettings = IndexSettingsModule.newIndexSettings("foo", settings); + Path tempDir = createTempDir().resolve(idxSettings.getUUID()).resolve("0"); + Files.createDirectories(tempDir); + ShardPath path = new ShardPath(false, tempDir, tempDir, new ShardId(idxSettings.getIndex(), 0)); + Directory dir = (new FsDirectoryFactory()).newDirectory(idxSettings, path); + if (random().nextBoolean()) { + dir = new MockDirectoryWrapper(random(), dir); + } + return dir; + } +} diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormatTests.java index 35bac97013487..8197423b742c4 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818HnswBinaryQuantizedVectorsFormatTests.java @@ -36,32 +36,18 @@ import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.misc.store.DirectIODirectory; import org.apache.lucene.search.TopDocs; 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.store.MMapDirectory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.store.MockDirectoryWrapper; import org.apache.lucene.tests.util.TestUtil; import org.apache.lucene.util.SameThreadExecutorService; import org.elasticsearch.common.logging.LogConfigurator; -import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.index.IndexModule; -import org.elasticsearch.index.IndexSettings; -import org.elasticsearch.index.shard.ShardId; -import org.elasticsearch.index.shard.ShardPath; -import org.elasticsearch.index.store.FsDirectoryFactory; -import org.elasticsearch.test.IndexSettingsModule; import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; import java.util.Arrays; import java.util.Locale; -import java.util.OptionalLong; import static java.lang.String.format; import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; @@ -87,7 +73,7 @@ public void testToString() { FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) { @Override public KnnVectorsFormat knnVectorsFormat() { - return new ES818HnswBinaryQuantizedVectorsFormat(10, 20, 1, null); + return new ES818HnswBinaryQuantizedVectorsFormat(10, 20, 1, false, null); } }; String expectedPattern = @@ -137,7 +123,7 @@ public void testLimits() { expectThrows(IllegalArgumentException.class, () -> new ES818HnswBinaryQuantizedVectorsFormat(20, 3201)); expectThrows( IllegalArgumentException.class, - () -> new ES818HnswBinaryQuantizedVectorsFormat(20, 100, 1, new SameThreadExecutorService()) + () -> new ES818HnswBinaryQuantizedVectorsFormat(20, 100, 1, false, new SameThreadExecutorService()) ); } @@ -155,14 +141,6 @@ public void testSimpleOffHeapSize() throws IOException { } } - public void testSimpleOffHeapSizeFSDir() throws IOException { - checkDirectIOSupported(); - var config = newIndexWriterConfig().setUseCompoundFile(false); // avoid compound files to allow directIO - try (Directory dir = newFSDirectory()) { - testSimpleOffHeapSizeImpl(dir, config, false); - } - } - public void testSimpleOffHeapSizeMMapDir() throws IOException { try (Directory dir = newMMapDirectory()) { testSimpleOffHeapSizeImpl(dir, newIndexWriterConfig(), true); @@ -203,39 +181,4 @@ static Directory newMMapDirectory() throws IOException { } return dir; } - - private Directory newFSDirectory() throws IOException { - Settings settings = Settings.builder() - .put(IndexModule.INDEX_STORE_TYPE_SETTING.getKey(), IndexModule.Type.HYBRIDFS.name().toLowerCase(Locale.ROOT)) - .build(); - IndexSettings idxSettings = IndexSettingsModule.newIndexSettings("foo", settings); - Path tempDir = createTempDir().resolve(idxSettings.getUUID()).resolve("0"); - Files.createDirectories(tempDir); - ShardPath path = new ShardPath(false, tempDir, tempDir, new ShardId(idxSettings.getIndex(), 0)); - Directory dir = (new FsDirectoryFactory()).newDirectory(idxSettings, path); - if (random().nextBoolean()) { - dir = new MockDirectoryWrapper(random(), dir); - } - return dir; - } - - static void checkDirectIOSupported() { - assumeTrue("Direct IO is not enabled", ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO); - - Path path = createTempDir("directIOProbe"); - try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { - out.writeString("test"); - } catch (IOException e) { - assumeNoException("test requires a filesystem that supports Direct IO", e); - } - } - - 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; - } - }; - } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java index 2524422ed8f90..8e67b7f414023 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java @@ -55,7 +55,10 @@ public DenseVectorFieldTypeTests() { } private static DenseVectorFieldMapper.RescoreVector randomRescoreVector() { - return new DenseVectorFieldMapper.RescoreVector(randomBoolean() ? 0 : randomFloatBetween(1.0F, 10.0F, false)); + return new DenseVectorFieldMapper.RescoreVector( + randomBoolean() ? 0 : randomFloatBetween(1.0F, 10.0F, false), + randomOptionalBoolean() + ); } private DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsNonQuantized() { @@ -663,7 +666,7 @@ public void testRescoreOversampleQueryOverrides() { 3, true, VectorSimilarity.COSINE, - randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(randomFloatBetween(1.1f, 9.9f, false))), + randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(randomFloatBetween(1.1f, 9.9f, false), null)), Collections.emptyMap(), false ); @@ -692,7 +695,7 @@ public void testRescoreOversampleQueryOverrides() { 3, true, VectorSimilarity.COSINE, - randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(0)), + randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(0f, null)), Collections.emptyMap(), false ); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java index 9972fa9e5ae0b..4cf784edd260e 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java @@ -1254,8 +1254,7 @@ static boolean indexVersionDefaultsToBbqHnsw(IndexVersion indexVersion) { public static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() { int m = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; int efConstruction = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; - DenseVectorFieldMapper.RescoreVector rescoreVector = new DenseVectorFieldMapper.RescoreVector(DEFAULT_RESCORE_OVERSAMPLE); - return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, rescoreVector); + return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, null); } static SemanticTextIndexOptions defaultIndexOptions(IndexVersion indexVersionCreated, MinimalServiceSettings modelSettings) { diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java index cc87edf59e9d3..a573bcc863de3 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java @@ -105,7 +105,6 @@ import static org.elasticsearch.xpack.inference.mapper.SemanticTextField.getChunksFieldName; import static org.elasticsearch.xpack.inference.mapper.SemanticTextField.getEmbeddingsFieldName; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper.DEFAULT_ELSER_2_INFERENCE_ID; -import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper.DEFAULT_RESCORE_OVERSAMPLE; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper.INDEX_OPTIONS_FIELD; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldTests.generateRandomChunkingSettings; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldTests.generateRandomChunkingSettingsOtherThan; @@ -1185,8 +1184,7 @@ private static SemanticTextIndexOptions defaultDenseVectorSemanticIndexOptions() private static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() { int m = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; int efConstruction = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; - DenseVectorFieldMapper.RescoreVector rescoreVector = new DenseVectorFieldMapper.RescoreVector(DEFAULT_RESCORE_OVERSAMPLE); - return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, rescoreVector); + return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, null); } private static SemanticTextIndexOptions defaultBbqHnswSemanticTextIndexOptions() { From 3776a01e46692c10e7fea3360e1e721b66b23d9c Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Mon, 14 Jul 2025 14:24:45 +0100 Subject: [PATCH 2/3] Use a separate option --- .../ES818BinaryQuantizedVectorsFormat.java | 1 - .../vectors/DenseVectorFieldMapper.java | 79 +++++++++++-------- .../vectors/DenseVectorFieldTypeTests.java | 19 +++-- .../mapper/SemanticTextFieldMapper.java | 3 +- .../mapper/SemanticTextFieldMapperTests.java | 4 +- 5 files changed, 60 insertions(+), 46 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java index cb68772ae1f9e..4baa18b580173 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java +++ b/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; 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 4f1373866485a..2c7c31857f13f 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 @@ -388,7 +388,8 @@ private DenseVectorIndexOptions defaultIndexOptions(boolean defaultInt8Hnsw, boo return new BBQHnswIndexOptions( Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN, Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH, - null + 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 indexOptionsMap, IndexVersion indexVersion) { Object mNode = indexOptionsMap.remove("m"); Object efConstructionNode = indexOptionsMap.remove("ef_construction"); + Object useDirectIONode = indexOptionsMap.remove("use_direct_io"); + if (mNode == null) { mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; } @@ -1630,12 +1633,19 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map indexOptionsMap, IndexVersion indexVersion) { Object rescoreVectorNode = indexOptionsMap.remove(NAME); @@ -2346,35 +2368,26 @@ static RescoreVector fromIndexOptions(Map indexOptionsMap, IndexVersi return null; } Map mappedNode = XContentMapValues.nodeMapValue(rescoreVectorNode, NAME); - - Float oversampleValue = null; Object oversampleNode = mappedNode.get(OVERSAMPLE); - if (oversampleNode != null) { - oversampleValue = (float) XContentMapValues.nodeDoubleValue(oversampleNode); - if (oversampleValue == 0 && allowsZeroRescore(indexVersion) == false) { - throw new IllegalArgumentException("oversample must be greater than 1"); - } - if (oversampleValue < 1 && oversampleValue != 0) { - throw new IllegalArgumentException("oversample must be greater than 1 or exactly 0"); - } else if (oversampleValue > 10) { - throw new IllegalArgumentException("oversample must be less than or equal to 10"); - } + if (oversampleNode == null) { + throw new IllegalArgumentException("Invalid rescore_vector value. Missing required field " + OVERSAMPLE); } - - Boolean directIO = (Boolean) mappedNode.get(DIRECT_IO); - - return new RescoreVector(oversampleValue, directIO); + float oversampleValue = (float) XContentMapValues.nodeDoubleValue(oversampleNode); + if (oversampleValue == 0 && allowsZeroRescore(indexVersion) == false) { + throw new IllegalArgumentException("oversample must be greater than 1"); + } + if (oversampleValue < 1 && oversampleValue != 0) { + throw new IllegalArgumentException("oversample must be greater than 1 or exactly 0"); + } else if (oversampleValue > 10) { + throw new IllegalArgumentException("oversample must be less than or equal to 10"); + } + return new RescoreVector(oversampleValue); } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(NAME); - if (oversample != null) { - builder.field(OVERSAMPLE, oversample); - } - if (useDirectIO != null) { - builder.field(DIRECT_IO, useDirectIO); - } + builder.field(OVERSAMPLE, oversample); builder.endObject(); return builder; } @@ -2716,10 +2729,6 @@ && isNotUnitVector(squaredMagnitude)) { && quantizedIndexOptions.rescoreVector != null) { oversample = quantizedIndexOptions.rescoreVector.oversample; } - if (oversample == null) { - oversample = DEFAULT_OVERSAMPLE; - } - boolean rescore = needsRescore(oversample); if (rescore) { // Will get k * oversample for rescoring, and get the top k diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java index 8e67b7f414023..f0753203a4da4 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java @@ -55,10 +55,7 @@ public DenseVectorFieldTypeTests() { } private static DenseVectorFieldMapper.RescoreVector randomRescoreVector() { - return new DenseVectorFieldMapper.RescoreVector( - randomBoolean() ? 0 : randomFloatBetween(1.0F, 10.0F, false), - randomOptionalBoolean() - ); + return new DenseVectorFieldMapper.RescoreVector(randomBoolean() ? 0 : randomFloatBetween(1.0F, 10.0F, false)); } private DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsNonQuantized() { @@ -95,7 +92,8 @@ public static DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsA new DenseVectorFieldMapper.BBQHnswIndexOptions( randomIntBetween(1, 100), randomIntBetween(1, 10_000), - randomFrom((DenseVectorFieldMapper.RescoreVector) null, randomRescoreVector()) + randomFrom((DenseVectorFieldMapper.RescoreVector) null, randomRescoreVector()), + randomBoolean() ), new DenseVectorFieldMapper.BBQFlatIndexOptions(randomFrom((DenseVectorFieldMapper.RescoreVector) null, randomRescoreVector())) ); @@ -121,7 +119,12 @@ private DenseVectorFieldMapper.DenseVectorIndexOptions randomIndexOptionsHnswQua randomFrom((Float) null, 0f, (float) randomDoubleBetween(0.9, 1.0, true)), rescoreVector ), - new DenseVectorFieldMapper.BBQHnswIndexOptions(randomIntBetween(1, 100), randomIntBetween(1, 10_000), rescoreVector) + new DenseVectorFieldMapper.BBQHnswIndexOptions( + randomIntBetween(1, 100), + randomIntBetween(1, 10_000), + rescoreVector, + randomBoolean() + ) ); } @@ -666,7 +669,7 @@ public void testRescoreOversampleQueryOverrides() { 3, true, VectorSimilarity.COSINE, - randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(randomFloatBetween(1.1f, 9.9f, false), null)), + randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(randomFloatBetween(1.1f, 9.9f, false))), Collections.emptyMap(), false ); @@ -695,7 +698,7 @@ public void testRescoreOversampleQueryOverrides() { 3, true, VectorSimilarity.COSINE, - randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(0f, null)), + randomIndexOptionsHnswQuantized(new DenseVectorFieldMapper.RescoreVector(0)), Collections.emptyMap(), false ); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java index 4cf784edd260e..6f3b36d25f31f 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java @@ -1254,7 +1254,8 @@ static boolean indexVersionDefaultsToBbqHnsw(IndexVersion indexVersion) { public static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() { int m = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; int efConstruction = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; - return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, null); + DenseVectorFieldMapper.RescoreVector rescoreVector = new DenseVectorFieldMapper.RescoreVector(DEFAULT_RESCORE_OVERSAMPLE); + return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, rescoreVector, false); } static SemanticTextIndexOptions defaultIndexOptions(IndexVersion indexVersionCreated, MinimalServiceSettings modelSettings) { diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java index a573bcc863de3..3c7a62ad62411 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapperTests.java @@ -105,6 +105,7 @@ import static org.elasticsearch.xpack.inference.mapper.SemanticTextField.getChunksFieldName; import static org.elasticsearch.xpack.inference.mapper.SemanticTextField.getEmbeddingsFieldName; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper.DEFAULT_ELSER_2_INFERENCE_ID; +import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper.DEFAULT_RESCORE_OVERSAMPLE; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper.INDEX_OPTIONS_FIELD; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldTests.generateRandomChunkingSettings; import static org.elasticsearch.xpack.inference.mapper.SemanticTextFieldTests.generateRandomChunkingSettingsOtherThan; @@ -1184,7 +1185,8 @@ private static SemanticTextIndexOptions defaultDenseVectorSemanticIndexOptions() private static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() { int m = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; int efConstruction = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; - return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, null); + DenseVectorFieldMapper.RescoreVector rescoreVector = new DenseVectorFieldMapper.RescoreVector(DEFAULT_RESCORE_OVERSAMPLE); + return new DenseVectorFieldMapper.BBQHnswIndexOptions(m, efConstruction, rescoreVector, false); } private static SemanticTextIndexOptions defaultBbqHnswSemanticTextIndexOptions() { From 31acb0148a6596a58fd192677f1a78bd8b23d117 Mon Sep 17 00:00:00 2001 From: Simon Cooper Date: Mon, 14 Jul 2025 17:11:34 +0100 Subject: [PATCH 3/3] Rename option, add basic tests --- .../upgrades/VectorSearchIT.java | 59 +++++++++++++++++ .../search.vectors/41_knn_search_bbq_hnsw.yml | 63 +++++++++++++++++++ .../vectors/DenseVectorFieldMapper.java | 22 +++---- .../elasticsearch/search/SearchFeatures.java | 4 +- 4 files changed, 136 insertions(+), 12 deletions(-) diff --git a/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/upgrades/VectorSearchIT.java b/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/upgrades/VectorSearchIT.java index afee17cd82e2d..e79b8eb0eda65 100644 --- a/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/upgrades/VectorSearchIT.java +++ b/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 response = search(searchRequest); + assertThat(extractValue(response, "hits.total.value"), equalTo(2)); + List> 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()) { diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml index e3c1155ed2000..7ccdf3131158d 100644 --- a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml +++ b/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"] 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 2c7c31857f13f..6ad4bc50be242 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 @@ -1623,7 +1623,7 @@ public boolean supportsDimension(int dims) { public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map indexOptionsMap, IndexVersion indexVersion) { Object mNode = indexOptionsMap.remove("m"); Object efConstructionNode = indexOptionsMap.remove("ef_construction"); - Object useDirectIONode = indexOptionsMap.remove("use_direct_io"); + Object disableOffheapCacheRescoringNode = indexOptionsMap.remove("disable_offheap_cache_rescoring"); if (mNode == null) { mNode = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; @@ -1642,10 +1642,10 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map 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 getTestFeatures() { @@ -43,7 +44,8 @@ public Set 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 ); } }