Skip to content

Commit efcc3fc

Browse files
retareta
committed
Introduce extensible KNNEngine API and BuiltinKNNEngine enumeration
Signed-off-by: Andriy Redko <drreta@gmail.com>
1 parent 86ad566 commit efcc3fc

167 files changed

Lines changed: 2004 additions & 1373 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎CHANGELOG.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
99
* Add debug mode to MMR rerank that injects per-hit scoring details (original_score, max_similarity_to_selected, mmr_score, mmr_formula) into _source via the `debug` flag in the mmr search extension [#3254](https://github.com/opensearch-project/k-NN/pull/3254)
1010
* Support derived source for knn with other fields [#3260](https://github.com/opensearch-project/k-NN/pull/3260)
1111
* Added support for 1 bit SQ with remote build [#3270](https://github.com/opensearch-project/k-NN/pull/3270)
12-
12+
* Introduce extensible KNNEngine API and BuiltinKNNEngine enumeration [#3288](https://github.com/opensearch-project/k-NN/pull/3288)
1313
### Maintenance
1414

1515

‎qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/FaissSQIT.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import org.opensearch.knn.common.KNNConstants;
2525
import org.opensearch.knn.index.SpaceType;
2626
import org.opensearch.knn.index.query.KNNQueryBuilder;
27-
import org.opensearch.knn.index.engine.KNNEngine;
27+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
2828

2929
import java.io.IOException;
3030
import java.util.Arrays;
@@ -101,7 +101,7 @@ public void testHNSWSQFP16_onUpgradeWhenIndexedAndQueried_thenSucceed() throws E
101101
.startObject(KNNConstants.KNN_METHOD)
102102
.field(KNNConstants.NAME, KNNConstants.METHOD_HNSW)
103103
.field(KNNConstants.METHOD_PARAMETER_SPACE_TYPE, spaceType.getValue())
104-
.field(KNNConstants.KNN_ENGINE, KNNEngine.FAISS.getName())
104+
.field(KNNConstants.KNN_ENGINE, BuiltinKNNEngine.FAISS.getName())
105105
.startObject(KNNConstants.PARAMETERS)
106106
.field(KNNConstants.METHOD_PARAMETER_M, mValues.get(random().nextInt(mValues.size())))
107107
.field(
@@ -178,7 +178,7 @@ public void testHNSWSQFP16_onUpgradeWhenClipToFp16isTrueAndIndexedWithOutOfFP16R
178178
.startObject(KNNConstants.KNN_METHOD)
179179
.field(KNNConstants.NAME, KNNConstants.METHOD_HNSW)
180180
.field(KNNConstants.METHOD_PARAMETER_SPACE_TYPE, SpaceType.L2.getValue())
181-
.field(KNNConstants.KNN_ENGINE, KNNEngine.FAISS.getName())
181+
.field(KNNConstants.KNN_ENGINE, BuiltinKNNEngine.FAISS.getName())
182182
.startObject(PARAMETERS)
183183
.field(KNNConstants.METHOD_PARAMETER_M, mValues.get(random().nextInt(mValues.size())))
184184
.field(

‎qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import org.opensearch.knn.index.KNNSettings;
1717
import org.opensearch.knn.index.SpaceType;
1818
import org.opensearch.knn.index.VectorDataType;
19-
import org.opensearch.knn.index.engine.KNNEngine;
19+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
2020
import org.opensearch.knn.index.engine.faiss.QFrameBitEncoder;
2121
import org.opensearch.knn.index.mapper.CompressionLevel;
2222
import org.opensearch.knn.index.mapper.Mode;
@@ -272,7 +272,7 @@ public void testKNNIndexBinaryForceMerge() throws Exception {
272272
TEST_FIELD,
273273
dimension,
274274
METHOD_HNSW,
275-
KNNEngine.FAISS.getName(),
275+
BuiltinKNNEngine.FAISS.getName(),
276276
SpaceType.HAMMING.getValue(),
277277
true,
278278
VectorDataType.BINARY
@@ -456,7 +456,7 @@ public void testBlockNMSLIBIndexCreationPost3_0_0() throws Exception {
456456
createKnnIndex(
457457
testIndex,
458458
getKNNDefaultIndexSettings(),
459-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, METHOD_HNSW, KNNEngine.NMSLIB.getName())
459+
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, METHOD_HNSW, BuiltinKNNEngine.NMSLIB.getName())
460460
);
461461
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
462462
// Flush to ensure the index persists after upgrade
@@ -468,7 +468,7 @@ public void testBlockNMSLIBIndexCreationPost3_0_0() throws Exception {
468468
() -> createKnnIndex(
469469
testIndex + "_new",
470470
getKNNDefaultIndexSettings(),
471-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, METHOD_HNSW, KNNEngine.NMSLIB.getName())
471+
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, METHOD_HNSW, BuiltinKNNEngine.NMSLIB.getName())
472472
)
473473
);
474474
deleteKNNIndex(testIndex);
@@ -482,7 +482,7 @@ public void testBlockModeAndCompressionBefore2_17_0() throws IOException {
482482
createKnnIndex(
483483
testIndex,
484484
getKNNDefaultIndexSettings(),
485-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, METHOD_HNSW, KNNEngine.NMSLIB.getName())
485+
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, METHOD_HNSW, BuiltinKNNEngine.NMSLIB.getName())
486486
);
487487
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
488488
flush(testIndex, true);

‎qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/PainlessScriptScoringIT.java‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
package org.opensearch.knn.bwc;
77

88
import org.opensearch.knn.index.SpaceType;
9-
import org.opensearch.knn.index.engine.KNNEngine;
9+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
1010

1111
public class PainlessScriptScoringIT extends AbstractRestartUpgradeTestCase {
1212
private static final String TEST_FIELD = "test-field";
@@ -61,7 +61,7 @@ public void testNonKNNIndex_withMethodParams_withFaissEngine() throws Exception
6161
createKnnIndex(
6262
testIndex,
6363
createKNNDefaultScriptScoreSettings(),
64-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", KNNEngine.FAISS.getName(), SpaceType.DEFAULT.getValue(), false)
64+
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", BuiltinKNNEngine.FAISS.getName(), SpaceType.DEFAULT.getValue(), false)
6565
);
6666
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
6767
} else {
@@ -84,7 +84,14 @@ public void testNonKNNIndex_withMethodParams_withNmslibEngine() throws Exception
8484
createKnnIndex(
8585
testIndex,
8686
createKNNDefaultScriptScoreSettings(),
87-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", KNNEngine.NMSLIB.getName(), SpaceType.DEFAULT.getValue(), false)
87+
createKnnIndexMapping(
88+
TEST_FIELD,
89+
DIMENSIONS,
90+
"hnsw",
91+
BuiltinKNNEngine.NMSLIB.getName(),
92+
SpaceType.DEFAULT.getValue(),
93+
false
94+
)
8895
);
8996
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
9097
} else {
@@ -107,7 +114,14 @@ public void testNonKNNIndex_withMethodParams_withLuceneEngine() throws Exception
107114
createKnnIndex(
108115
testIndex,
109116
createKNNDefaultScriptScoreSettings(),
110-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", KNNEngine.LUCENE.getName(), SpaceType.DEFAULT.getValue(), false)
117+
createKnnIndexMapping(
118+
TEST_FIELD,
119+
DIMENSIONS,
120+
"hnsw",
121+
BuiltinKNNEngine.LUCENE.getName(),
122+
SpaceType.DEFAULT.getValue(),
123+
false
124+
)
111125
);
112126
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
113127
} else {

‎qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import org.opensearch.knn.KNNResult;
1515
import org.opensearch.knn.index.SpaceType;
1616
import org.opensearch.core.rest.RestStatus;
17-
import org.opensearch.knn.index.engine.KNNEngine;
17+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
1818

1919
import java.util.Collections;
2020
import java.util.HashMap;
@@ -86,7 +86,7 @@ public void testNonKNNIndex_withMethodParams_withFaissEngine() throws Exception
8686
createKnnIndex(
8787
testIndex,
8888
createKNNDefaultScriptScoreSettings(),
89-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", KNNEngine.FAISS.getName(), SpaceType.DEFAULT.getValue(), false)
89+
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", BuiltinKNNEngine.FAISS.getName(), SpaceType.DEFAULT.getValue(), false)
9090
);
9191
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
9292
} else {
@@ -105,7 +105,14 @@ public void testNonKNNIndex_withMethodParams_withNmslibEngine() throws Exception
105105
createKnnIndex(
106106
testIndex,
107107
createKNNDefaultScriptScoreSettings(),
108-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", KNNEngine.NMSLIB.getName(), SpaceType.DEFAULT.getValue(), false)
108+
createKnnIndexMapping(
109+
TEST_FIELD,
110+
DIMENSIONS,
111+
"hnsw",
112+
BuiltinKNNEngine.NMSLIB.getName(),
113+
SpaceType.DEFAULT.getValue(),
114+
false
115+
)
109116
);
110117
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
111118
} else {
@@ -124,7 +131,14 @@ public void testNonKNNIndex_withMethodParams_withLuceneEngine() throws Exception
124131
createKnnIndex(
125132
testIndex,
126133
createKNNDefaultScriptScoreSettings(),
127-
createKnnIndexMapping(TEST_FIELD, DIMENSIONS, "hnsw", KNNEngine.LUCENE.getName(), SpaceType.DEFAULT.getValue(), false)
134+
createKnnIndexMapping(
135+
TEST_FIELD,
136+
DIMENSIONS,
137+
"hnsw",
138+
BuiltinKNNEngine.LUCENE.getName(),
139+
SpaceType.DEFAULT.getValue(),
140+
false
141+
)
128142
);
129143
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
130144
} else {

‎src/main/java/org/opensearch/knn/common/FieldInfoExtractor.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.opensearch.index.mapper.MapperService;
1616
import org.opensearch.knn.index.SpaceType;
1717
import org.opensearch.knn.index.VectorDataType;
18+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
1819
import org.opensearch.knn.index.engine.KNNEngine;
1920
import org.opensearch.knn.index.engine.MemoryOptimizedSearchSupportSpec;
2021
import org.opensearch.knn.index.engine.faiss.SQConfig;
@@ -53,8 +54,8 @@ public static KNNEngine extractKNNEngine(final FieldInfo field) {
5354
if (modelMetadata != null) {
5455
return modelMetadata.getKnnEngine();
5556
}
56-
final String engineName = field.attributes().getOrDefault(KNNConstants.KNN_ENGINE, KNNEngine.DEFAULT.getName());
57-
return KNNEngine.getEngine(engineName);
57+
final String engineName = field.attributes().getOrDefault(KNNConstants.KNN_ENGINE, BuiltinKNNEngine.DEFAULT.getName());
58+
return BuiltinKNNEngine.getEngine(engineName);
5859
}
5960

6061
/**

‎src/main/java/org/opensearch/knn/index/KNNIndexShard.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.opensearch.index.shard.IndexShard;
2222
import org.opensearch.knn.common.FieldInfoExtractor;
2323
import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper;
24+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
2425
import org.opensearch.knn.index.engine.KNNEngine;
2526
import org.opensearch.knn.index.engine.qframe.QuantizationConfig;
2627
import org.opensearch.knn.index.mapper.KNNVectorFieldMapper;
@@ -134,7 +135,7 @@ private void warmUpOffHeapIndex(final List<EngineFileContext> engineFileContexts
134135
NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
135136
getParametersAtLoading(
136137
engineFileContext.getSpaceType(),
137-
KNNEngine.getEngineNameFromPath(engineFileContext.getVectorFileName()),
138+
BuiltinKNNEngine.getEngineNameFromPath(engineFileContext.getVectorFileName()),
138139
getIndexName(),
139140
engineFileContext.getVectorDataType(),
140141
engineFileContext.getSegmentLevelQuantizationInfo()
@@ -195,7 +196,7 @@ public void clearCache() {
195196
List<EngineFileContext> getAllEngineFileContexts(final Set<String> loadedFieldNames, final LeafReaderContext leafReaderContext)
196197
throws IOException {
197198
List<EngineFileContext> engineFiles = new ArrayList<>();
198-
for (KNNEngine knnEngine : KNNEngine.getEnginesThatCreateCustomSegmentFiles()) {
199+
for (KNNEngine knnEngine : BuiltinKNNEngine.getEnginesThatCreateCustomSegmentFiles()) {
199200
engineFiles.addAll(getEngineFileContexts(loadedFieldNames, leafReaderContext, knnEngine));
200201
}
201202
return engineFiles;

‎src/main/java/org/opensearch/knn/index/codec/KNN1040BasePerFieldKnnVectorsFormat.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
1111
import org.opensearch.index.mapper.MapperService;
1212
import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
13+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
1314
import org.opensearch.knn.index.engine.CodecFormatResolver;
1415
import org.opensearch.knn.index.engine.KNNEngine;
1516
import org.opensearch.knn.index.engine.KNNMethodContext;
@@ -91,7 +92,7 @@ public KnnVectorsFormat getKnnVectorsFormatForField(final String field) {
9192
final KNNEngine engine = knnMethodContext.getKnnEngine();
9293
final Map<String, Object> params = knnMethodContext.getMethodComponentContext().getParameters();
9394

94-
if (engine == KNNEngine.LUCENE) {
95+
if (engine == BuiltinKNNEngine.LUCENE) {
9596
return luceneFormatResolver.resolve(field, knnMethodContext, params, defaultMaxConnections, defaultBeamWidth);
9697
}
9798

‎src/main/java/org/opensearch/knn/index/codec/KNN1040Codec/Faiss1040ScalarQuantizedKnnVectorsFormat.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import org.apache.lucene.index.SegmentReadState;
1515
import org.apache.lucene.index.SegmentWriteState;
1616
import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
17-
import org.opensearch.knn.index.engine.KNNEngine;
17+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
1818

1919
import java.io.IOException;
2020

@@ -84,7 +84,7 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException
8484
*/
8585
@Override
8686
public int getMaxDimensions(String fieldName) {
87-
return KNNEngine.getMaxDimensionByEngine(KNNEngine.FAISS);
87+
return BuiltinKNNEngine.getMaxDimensionByEngine(BuiltinKNNEngine.FAISS);
8888
}
8989

9090
@Override

‎src/main/java/org/opensearch/knn/index/codec/KNN1040Codec/KNN1040HnswScalarQuantizedVectorsFormat.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import org.apache.lucene.index.SegmentReadState;
1515
import org.apache.lucene.index.SegmentWriteState;
1616
import org.apache.lucene.search.TaskExecutor;
17-
import org.opensearch.knn.index.engine.KNNEngine;
17+
import org.opensearch.knn.index.engine.BuiltinKNNEngine;
1818

1919
import java.io.IOException;
2020
import java.util.concurrent.ExecutorService;
@@ -89,7 +89,7 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException
8989

9090
@Override
9191
public int getMaxDimensions(String fieldName) {
92-
return KNNEngine.getMaxDimensionByEngine(KNNEngine.LUCENE);
92+
return BuiltinKNNEngine.getMaxDimensionByEngine(BuiltinKNNEngine.LUCENE);
9393
}
9494

9595
@Override

0 commit comments

Comments
 (0)