-
Notifications
You must be signed in to change notification settings - Fork 25.5k
Add on-disk rescoring to disk BBQ #135778
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 4 commits
Commits
Show all changes
12 commits
Select commit
Hold shift + click to select a range
cf67fc8
Add Direct IO support to DiskBBQ
thecoop 1ec7452
Update docs/changelog/135778.yaml
thecoop 19b1d40
Use parameters instead of random
thecoop a92fed8
Add bwc format test
thecoop b115fc1
Add updateable index test
thecoop 3bb425b
Merge branch 'main' into diskbbq-disk-rescoring
thecoop cb914d8
Update test
thecoop 8a8d9b9
Update the base version used by the next format
thecoop 5d6a3b1
Merge branch 'main' into diskbbq-disk-rescoring
thecoop 450a8a6
Merge branch 'main' into diskbbq-disk-rescoring
thecoop 50beb8c
Check the version, not the directIO value
thecoop f9a879c
Merge branch 'main' into diskbbq-disk-rescoring
thecoop File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| pr: 135778 | ||
| summary: "Add `on_disk_rescore: true` option to disk BBQ to rescore vectors on-disk without loading into memory" | ||
| area: Vector Search | ||
| type: feature | ||
| issues: [] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
187 changes: 187 additions & 0 deletions
187
server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,187 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the "Elastic License | ||
| * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side | ||
| * Public License v 1"; you may not use this file except in compliance with, at | ||
| * your election, the "Elastic License 2.0", the "GNU Affero General Public | ||
| * License v3.0 only", or the "Server Side Public License, v 1". | ||
| */ | ||
|
|
||
| package org.elasticsearch.index.store; | ||
|
|
||
| import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; | ||
|
|
||
| import org.apache.logging.log4j.Level; | ||
| import org.apache.lucene.misc.store.DirectIODirectory; | ||
| import org.apache.lucene.store.Directory; | ||
| import org.apache.lucene.store.FSDirectory; | ||
| import org.apache.lucene.store.IOContext; | ||
| import org.apache.lucene.store.IndexOutput; | ||
| import org.apache.lucene.tests.util.LuceneTestCase; | ||
| import org.elasticsearch.common.settings.Settings; | ||
| import org.elasticsearch.core.Strings; | ||
| import org.elasticsearch.plugins.Plugin; | ||
| import org.elasticsearch.search.vectors.KnnSearchBuilder; | ||
| import org.elasticsearch.search.vectors.VectorData; | ||
| import org.elasticsearch.test.ESIntegTestCase; | ||
| import org.elasticsearch.test.ESTestCase; | ||
| import org.elasticsearch.test.InternalSettingsPlugin; | ||
| import org.elasticsearch.test.MockLog; | ||
| import org.elasticsearch.test.junit.annotations.TestLogging; | ||
| import org.junit.BeforeClass; | ||
|
|
||
| import java.io.IOException; | ||
| import java.nio.file.Path; | ||
| import java.util.Collection; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.OptionalLong; | ||
| import java.util.stream.IntStream; | ||
|
|
||
| import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; | ||
| import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount; | ||
| import static org.hamcrest.Matchers.equalTo; | ||
| import static org.hamcrest.Matchers.is; | ||
|
|
||
| @LuceneTestCase.SuppressCodecs("*") // only use our own codecs | ||
| @ESTestCase.WithoutEntitlements // requires entitlement delegation ES-10920 | ||
| public class DirectIOIT extends ESIntegTestCase { | ||
|
|
||
| private static boolean SUPPORTED; | ||
|
|
||
| @BeforeClass | ||
| public static void checkSupported() { | ||
| Path path = createTempDir("directIOProbe"); | ||
| try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { | ||
| out.writeString("test"); | ||
| SUPPORTED = true; | ||
| } catch (IOException e) { | ||
| SUPPORTED = false; | ||
| } | ||
| } | ||
|
|
||
| static DirectIODirectory open(Path path) throws IOException { | ||
| return new DirectIODirectory(FSDirectory.open(path)) { | ||
| @Override | ||
| protected boolean useDirectIO(String name, IOContext context, OptionalLong fileLength) { | ||
| return true; | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| private final String type; | ||
|
|
||
| @ParametersFactory | ||
| public static Iterable<Object[]> parameters() { | ||
| return List.<Object[]>of(new Object[] { "bbq_disk" }); | ||
| } | ||
|
|
||
| public DirectIOIT(String type) { | ||
| this.type = type; | ||
| } | ||
|
|
||
| @Override | ||
| protected Collection<Class<? extends Plugin>> nodePlugins() { | ||
| return List.of(InternalSettingsPlugin.class); | ||
| } | ||
|
|
||
| private String indexVectors(boolean directIO) { | ||
| String indexName = "test-vectors-" + directIO; | ||
| assertAcked( | ||
| prepareCreate(indexName).setSettings(Settings.builder().put(InternalSettingsPlugin.USE_COMPOUND_FILE.getKey(), false)) | ||
| .setMapping(Strings.format(""" | ||
| { | ||
| "properties": { | ||
| "fooVector": { | ||
| "type": "dense_vector", | ||
| "dims": 64, | ||
| "element_type": "float", | ||
| "index": true, | ||
| "similarity": "l2_norm", | ||
| "index_options": { | ||
| "type": "%s", | ||
| "on_disk_rescore": %s | ||
| } | ||
| } | ||
| } | ||
| } | ||
| """, type, directIO)) | ||
| ); | ||
| ensureGreen(indexName); | ||
|
|
||
| for (int i = 0; i < 1000; i++) { | ||
| indexDoc(indexName, Integer.toString(i), "fooVector", IntStream.range(0, 64).mapToDouble(d -> randomFloat()).toArray()); | ||
| } | ||
| refresh(); | ||
| assertBBQIndexType(indexName, type); // test assertion to ensure that the correct index type is being used | ||
| return indexName; | ||
| } | ||
|
|
||
| @SuppressWarnings("unchecked") | ||
| static void assertBBQIndexType(String indexName, String type) { | ||
| var response = indicesAdmin().prepareGetFieldMappings(indexName).setFields("fooVector").get(); | ||
| var map = (Map<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 | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
awesome, thank you!