Skip to content

Commit

Permalink
Add Lucene specific file extensions to core HybridFS (#721)
Browse files Browse the repository at this point in the history
* Add lucene vector specific file extensions for io with mmap

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
  • Loading branch information
martin-gaievski authored Jan 12, 2023
1 parent 4229216 commit 8a2aa04
Show file tree
Hide file tree
Showing 7 changed files with 71 additions and 0 deletions.
1 change: 1 addition & 0 deletions release-notes/opensearch-knn.release-notes-2.5.0.0.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ Compatible with OpenSearch 2.5.0
### Enhancements

* Extend SystemIndexPlugin for k-NN model system index ([#630](https://github.com/opensearch-project/k-NN/pull/630))
* Add Lucene specific file extensions to core HybridFS ([#721](https://github.com/opensearch-project/k-NN/pull/721))

### Bug Fixes

Expand Down
6 changes: 6 additions & 0 deletions src/main/java/org/opensearch/knn/index/util/KNNEngine.java
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.opensearch.knn.index.KNNMethodContext;
import org.opensearch.knn.index.SpaceType;

import java.util.List;
import java.util.Map;
import java.util.Set;

Expand Down Expand Up @@ -176,4 +177,9 @@ public Boolean isInitialized() {
public void setInitialized(Boolean isInitialized) {
knnLibrary.setInitialized(isInitialized);
}

@Override
public List<String> mmapFileExtensions() {
return knnLibrary.mmapFileExtensions();
}
}
11 changes: 11 additions & 0 deletions src/main/java/org/opensearch/knn/index/util/KNNLibrary.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
import org.opensearch.knn.index.KNNMethodContext;
import org.opensearch.knn.index.SpaceType;

import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
Expand Down Expand Up @@ -113,4 +115,13 @@ public interface KNNLibrary {
* @param isInitialized whether library has been initialized
*/
void setInitialized(Boolean isInitialized);

/**
* Getter for mmap file extensions
*
* @return list of file extensions that will be read/write with mmap
*/
default List<String> mmapFileExtensions() {
return Collections.EMPTY_LIST;
}
}
6 changes: 6 additions & 0 deletions src/main/java/org/opensearch/knn/index/util/Lucene.java
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import org.opensearch.knn.index.Parameter;
import org.opensearch.knn.index.SpaceType;

import java.util.List;
import java.util.Map;

import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW;
Expand Down Expand Up @@ -73,4 +74,9 @@ public float score(float rawScore, SpaceType spaceType) {
// score provided.
return rawScore;
}

@Override
public List<String> mmapFileExtensions() {
return List.of("vec", "vex");
}
}
24 changes: 24 additions & 0 deletions src/main/java/org/opensearch/knn/plugin/KNNPlugin.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.opensearch.knn.index.query.KNNWeight;
import org.opensearch.knn.index.codec.KNNCodecService;
import org.opensearch.knn.index.memory.NativeMemoryLoadStrategy;
import org.opensearch.knn.index.util.KNNEngine;
import org.opensearch.knn.indices.ModelGraveyard;
import org.opensearch.knn.indices.ModelCache;
import org.opensearch.knn.indices.ModelDao;
Expand Down Expand Up @@ -104,6 +105,8 @@
import java.util.Map;
import java.util.Optional;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static java.util.Collections.singletonList;
import static org.opensearch.knn.common.KNNConstants.KNN_THREAD_POOL_PREFIX;
Expand Down Expand Up @@ -338,4 +341,25 @@ public List<NamedXContentRegistry.Entry> getNamedXContent() {
public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings settings) {
return ImmutableList.of(new SystemIndexDescriptor(MODEL_INDEX_NAME, "Index for storing models used for k-NN indices"));
}

/**
* Plugin can provide additional node settings, that includes new settings or overrides for existing one from core.
*
* @return settings that are set by plugin
*/
@Override
public Settings additionalSettings() {
// We add engine specific extensions to the core list for HybridFS store type. We read existing values
// and append ours because in core setting will be replaced by override.
// Values are set as cluster defaults and are used at index creation time. Index specific overrides will take priority over values
// that are set here.
final List<String> engineSettings = Arrays.stream(KNNEngine.values())
.flatMap(engine -> engine.mmapFileExtensions().stream())
.collect(Collectors.toList());
final List<String> combinedSettings = Stream.concat(
IndexModule.INDEX_STORE_HYBRID_MMAP_EXTENSIONS.getDefault(Settings.EMPTY).stream(),
engineSettings.stream()
).collect(Collectors.toList());
return Settings.builder().putList(IndexModule.INDEX_STORE_HYBRID_MMAP_EXTENSIONS.getKey(), combinedSettings).build();
}
}
14 changes: 14 additions & 0 deletions src/test/java/org/opensearch/knn/index/util/KNNEngineTests.java
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,10 @@
import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.common.KNNConstants;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

public class KNNEngineTests extends KNNTestCase {
/**
* Check that version from engine and library match
Expand Down Expand Up @@ -47,4 +51,14 @@ public void testGetEngineFromPath() {
String invalidPath = "test.invalid";
expectThrows(IllegalArgumentException.class, () -> KNNEngine.getEngineNameFromPath(invalidPath));
}

public void testMmapFileExtensions() {
final List<String> mmapExtensions = Arrays.stream(KNNEngine.values())
.flatMap(engine -> engine.mmapFileExtensions().stream())
.collect(Collectors.toList());
assertNotNull(mmapExtensions);
final List<String> expectedSettings = List.of("vex", "vec");
assertTrue(expectedSettings.containsAll(mmapExtensions));
assertTrue(mmapExtensions.containsAll(expectedSettings));
}
}
9 changes: 9 additions & 0 deletions src/test/java/org/opensearch/knn/index/util/LuceneTests.java
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW;
Expand Down Expand Up @@ -114,4 +115,12 @@ public void testVersion() {
Lucene luceneInstance = Lucene.INSTANCE;
assertEquals(Version.LATEST.toString(), luceneInstance.getVersion());
}

public void testMmapFileExtensions() {
final List<String> luceneMmapExtensions = Lucene.INSTANCE.mmapFileExtensions();
assertNotNull(luceneMmapExtensions);
final List<String> expectedSettings = List.of("vex", "vec");
assertTrue(expectedSettings.containsAll(luceneMmapExtensions));
assertTrue(luceneMmapExtensions.containsAll(expectedSettings));
}
}

0 comments on commit 8a2aa04

Please sign in to comment.