Skip to content

Commit bb651a7

Browse files
jmazanec15jmazanec15
committed
Introduce mode and compression param resolution
Adds mode and compression based parameter resolution. With this, if a user specifies the mode and/or compression params, we will create a default configuration with the aim of meeting those hints. Currently, it does not contain support for overriding any of the parameters. This will be taken in a future commit. Signed-off-by: John Mazanec <[email protected]>
1 parent a58a9dc commit bb651a7

20 files changed

Lines changed: 783 additions & 336 deletions

‎release-notes/opensearch-knn.release-notes-2.17.0.0.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ Compatible with OpenSearch 2.17.0
77
* k-NN query rescore support for native engines [#1984](https://github.com/opensearch-project/k-NN/pull/1984)
88
* Add support for byte vector with Faiss Engine HNSW algorithm [#1823](https://github.com/opensearch-project/k-NN/pull/1823)
99
* Add support for byte vector with Faiss Engine IVF algorithm [#2002](https://github.com/opensearch-project/k-NN/pull/2002)
10+
* Add mode/compression configuration support for disk-based vector search [#2034](https://github.com/opensearch-project/k-NN/pull/2034)
1011
### Enhancements
1112
* Adds iterative graph build capability into a faiss index to improve the memory footprint during indexing and Integrates KNNVectorsFormat for native engines[#1950](https://github.com/opensearch-project/k-NN/pull/1950)
1213
### Bug Fixes

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import org.opensearch.common.lucene.Lucene;
1919
import org.opensearch.index.engine.Engine;
2020
import org.opensearch.index.shard.IndexShard;
21+
import org.opensearch.knn.common.FieldInfoExtractor;
22+
import org.opensearch.knn.index.engine.qframe.QuantizationConfig;
2123
import org.opensearch.knn.index.mapper.KNNVectorFieldMapper;
2224
import org.opensearch.knn.index.memory.NativeMemoryAllocation;
2325
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
@@ -182,7 +184,11 @@ List<EngineFileContext> getEngineFileContexts(IndexReader indexReader, KNNEngine
182184
shardPath,
183185
spaceType,
184186
modelId,
185-
VectorDataType.get(fieldInfo.attributes().getOrDefault(VECTOR_DATA_TYPE_FIELD, VectorDataType.FLOAT.getValue()))
187+
FieldInfoExtractor.extractQuantizationConfig(fieldInfo) == QuantizationConfig.EMPTY
188+
? VectorDataType.get(
189+
fieldInfo.attributes().getOrDefault(VECTOR_DATA_TYPE_FIELD, VectorDataType.FLOAT.getValue())
190+
)
191+
: VectorDataType.BINARY
186192
)
187193
);
188194
}

‎src/main/java/org/opensearch/knn/index/engine/KNNMethodConfigContext.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
import lombok.Setter;
1313
import org.opensearch.Version;
1414
import org.opensearch.knn.index.VectorDataType;
15+
import org.opensearch.knn.index.mapper.CompressionLevel;
16+
import org.opensearch.knn.index.mapper.Mode;
1517

1618
/**
1719
* This object provides additional context that the user does not provide when {@link KNNMethodContext} is
@@ -27,5 +29,10 @@ public final class KNNMethodConfigContext {
2729
private VectorDataType vectorDataType;
2830
private Integer dimension;
2931
private Version versionCreated;
32+
@Builder.Default
33+
private Mode mode = Mode.NOT_CONFIGURED;
34+
@Builder.Default
35+
private CompressionLevel compressionLevel = CompressionLevel.NOT_CONFIGURED;
36+
3037
public static final KNNMethodConfigContext EMPTY = KNNMethodConfigContext.builder().build();
3138
}

‎src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,9 @@
88
import lombok.AllArgsConstructor;
99
import lombok.Getter;
1010
import org.opensearch.core.common.Strings;
11+
import org.opensearch.knn.index.query.rescore.RescoreContext;
1112

12-
import java.util.Arrays;
1313
import java.util.Locale;
14-
import java.util.stream.Collectors;
1514

1615
/**
1716
* Enum representing the compression level for float vectors. Compression in this sense refers to compressing a
@@ -20,20 +19,23 @@
2019
*/
2120
@AllArgsConstructor
2221
public enum CompressionLevel {
23-
NOT_CONFIGURED(-1, ""),
24-
x1(1, "1x"),
25-
x2(2, "2x"),
26-
x4(4, "4x"),
27-
x8(8, "8x"),
28-
x16(16, "16x"),
29-
x32(32, "32x");
22+
NOT_CONFIGURED(-1, "", null),
23+
x1(1, "1x", null),
24+
x2(2, "2x", null),
25+
x4(4, "4x", new RescoreContext(1.0f)),
26+
x8(8, "8x", new RescoreContext(1.5f)),
27+
x16(16, "16x", new RescoreContext(2.0f)),
28+
x32(32, "32x", new RescoreContext(2.0f));
3029

3130
// Internally, an empty string is easier to deal with them null. However, from the mapping,
3231
// we do not want users to pass in the empty string and instead want null. So we make the conversion herex
33-
static final String[] NAMES_ARRAY = Arrays.stream(CompressionLevel.values())
34-
.map(compressionLevel -> compressionLevel == NOT_CONFIGURED ? null : compressionLevel.getName())
35-
.collect(Collectors.toList())
36-
.toArray(new String[0]);
32+
public static final String[] NAMES_ARRAY = new String[] {
33+
NOT_CONFIGURED.getName(),
34+
x1.getName(),
35+
x2.getName(),
36+
x8.getName(),
37+
x16.getName(),
38+
x32.getName() };
3739

3840
/**
3941
* Default is set to 1x and is a noop
@@ -62,6 +64,8 @@ public static CompressionLevel fromName(String name) {
6264
private final int compressionLevel;
6365
@Getter
6466
private final String name;
67+
@Getter
68+
private final RescoreContext defaultRescoreContext;
6569

6670
/**
6771
* Gets the number of bits used to represent a float in order to achieve this compression. For instance, for

‎src/main/java/org/opensearch/knn/index/mapper/KNNMappingConfig.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,24 @@ default Optional<KNNMethodContext> getKnnMethodContext() {
3030
return Optional.empty();
3131
}
3232

33+
/**
34+
* Return the mode to be used for this field
35+
*
36+
* @return {@link Mode}
37+
*/
38+
default Mode getMode() {
39+
return Mode.NOT_CONFIGURED;
40+
}
41+
42+
/**
43+
* Return compression level to be used for this field
44+
*
45+
* @return {@link CompressionLevel}
46+
*/
47+
default CompressionLevel getCompressionLevel() {
48+
return CompressionLevel.NOT_CONFIGURED;
49+
}
50+
3351
/**
3452
*
3553
* @return the dimension of the index; for model based indices, it will be null

‎src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java‎

Lines changed: 45 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -145,16 +145,7 @@ public static class Builder extends ParametrizedFieldMapper.Builder {
145145
b.startObject(n);
146146
v.toXContent(b, ToXContent.EMPTY_PARAMS);
147147
b.endObject();
148-
}), m -> m.getMethodComponentContext().getName()).setValidator(v -> {
149-
if (v == null) return;
150-
151-
ValidationException validationException;
152-
if (v.isTrainingRequired()) {
153-
validationException = new ValidationException();
154-
validationException.addValidationError(String.format(Locale.ROOT, "\"%s\" requires training.", KNN_METHOD));
155-
throw validationException;
156-
}
157-
});
148+
}), m -> m.getMethodComponentContext().getName());
158149

159150
protected final Parameter<String> mode = Parameter.restrictedStringParam(
160151
KNNConstants.MODE_PARAMETER,
@@ -354,13 +345,34 @@ public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserCont
354345
} else if (builder.modelId.get() != null) {
355346
validateFromModel(builder);
356347
} else {
348+
validateMode(builder);
357349
resolveKNNMethodComponents(builder, parserContext);
358350
validateFromKNNMethod(builder);
359351
}
360352

361353
return builder;
362354
}
363355

356+
private void validateMode(KNNVectorFieldMapper.Builder builder) {
357+
boolean isKNNMethodContextConfigured = builder.originalParameters.getKnnMethodContext() != null;
358+
boolean isModeConfigured = builder.mode.isConfigured() || builder.compressionLevel.isConfigured();
359+
if (isModeConfigured && isKNNMethodContextConfigured) {
360+
throw new MapperParsingException(
361+
String.format(
362+
Locale.ROOT,
363+
"Compression and mode can not be specified in a \"method\" mapping configuration for field: %s",
364+
builder.name
365+
)
366+
);
367+
}
368+
369+
if (isModeConfigured && builder.vectorDataType.getValue() != VectorDataType.FLOAT) {
370+
throw new MapperParsingException(
371+
String.format(Locale.ROOT, "Compression and mode cannot be used for non-float32 data type for field %s", builder.name)
372+
);
373+
}
374+
}
375+
364376
private void validateFromFlat(KNNVectorFieldMapper.Builder builder) {
365377
if (builder.modelId.get() != null || builder.knnMethodContext.get() != null) {
366378
throw new IllegalArgumentException("Cannot set modelId or method parameters when index.knn setting is false");
@@ -378,9 +390,15 @@ private void validateFromModel(KNNVectorFieldMapper.Builder builder) {
378390
}
379391

380392
private void validateFromKNNMethod(KNNVectorFieldMapper.Builder builder) {
393+
ValidationException validationException;
394+
if (builder.originalParameters.getResolvedKnnMethodContext().isTrainingRequired()) {
395+
validationException = new ValidationException();
396+
validationException.addValidationError(String.format(Locale.ROOT, "\"%s\" requires training.", KNN_METHOD));
397+
throw validationException;
398+
}
399+
381400
if (builder.originalParameters.getResolvedKnnMethodContext() != null) {
382-
ValidationException validationException = builder.originalParameters.getResolvedKnnMethodContext()
383-
.validate(builder.knnMethodConfigContext);
401+
validationException = builder.originalParameters.getResolvedKnnMethodContext().validate(builder.knnMethodConfigContext);
384402
if (validationException != null) {
385403
throw validationException;
386404
}
@@ -410,9 +428,11 @@ private void validateCompressionAndModeNotSet(KNNVectorFieldMapper.Builder build
410428
private void resolveKNNMethodComponents(KNNVectorFieldMapper.Builder builder, ParserContext parserContext) {
411429
builder.setKnnMethodConfigContext(
412430
KNNMethodConfigContext.builder()
413-
.vectorDataType(builder.vectorDataType.getValue())
431+
.vectorDataType(builder.originalParameters.getVectorDataType())
414432
.versionCreated(parserContext.indexVersionCreated())
415-
.dimension(builder.dimension.getValue())
433+
.dimension(builder.originalParameters.getDimension())
434+
.mode(Mode.fromName(builder.originalParameters.getMode()))
435+
.compressionLevel(CompressionLevel.fromName(builder.originalParameters.getCompressionLevel()))
416436
.build()
417437
);
418438

@@ -421,8 +441,17 @@ private void resolveKNNMethodComponents(KNNVectorFieldMapper.Builder builder, Pa
421441
builder.originalParameters.setResolvedKnnMethodContext(
422442
createKNNMethodContextFromLegacy(parserContext.getSettings(), parserContext.indexVersionCreated())
423443
);
424-
}
425-
setDefaultSpaceType(builder.originalParameters.getResolvedKnnMethodContext(), builder.vectorDataType.getValue());
444+
} else if (Mode.isConfigured(Mode.fromName(builder.mode.get()))
445+
|| CompressionLevel.isConfigured(CompressionLevel.fromName(builder.compressionLevel.get()))) {
446+
builder.originalParameters.setResolvedKnnMethodContext(
447+
ModeBasedResolver.INSTANCE.resolveKNNMethodContext(
448+
builder.knnMethodConfigContext.getMode(),
449+
builder.knnMethodConfigContext.getCompressionLevel(),
450+
false
451+
)
452+
);
453+
}
454+
setDefaultSpaceType(builder.originalParameters.getResolvedKnnMethodContext(), builder.originalParameters.getVectorDataType());
426455
}
427456

428457
private boolean isKNNDisabled(Settings settings) {

‎src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.opensearch.index.query.QueryShardException;
1818
import org.opensearch.knn.index.KNNVectorIndexFieldData;
1919
import org.opensearch.knn.index.VectorDataType;
20+
import org.opensearch.knn.index.query.rescore.RescoreContext;
2021
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
2122
import org.opensearch.search.lookup.SearchLookup;
2223

@@ -81,4 +82,20 @@ public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName, S
8182
public Object valueForDisplay(Object value) {
8283
return deserializeStoredVector((BytesRef) value, vectorDataType);
8384
}
85+
86+
/**
87+
* Resolve the rescore context provided for a user based on the field configuration
88+
*
89+
* @param userProvidedContext {@link RescoreContext} user passed; if null, the default should be configured
90+
* @return resolved {@link RescoreContext}
91+
*/
92+
public RescoreContext resolveRescoreContext(RescoreContext userProvidedContext) {
93+
if (userProvidedContext != null) {
94+
return userProvidedContext;
95+
}
96+
return ModeBasedResolver.INSTANCE.resolveRescoreContext(
97+
getKnnMappingConfig().getMode(),
98+
getKnnMappingConfig().getCompressionLevel()
99+
);
100+
}
84101
}

‎src/main/java/org/opensearch/knn/index/mapper/MethodFieldMapper.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,16 @@ public Optional<KNNMethodContext> getKnnMethodContext() {
6565
public int getDimension() {
6666
return knnMethodConfigContext.getDimension();
6767
}
68+
69+
@Override
70+
public Mode getMode() {
71+
return knnMethodConfigContext.getMode();
72+
}
73+
74+
@Override
75+
public CompressionLevel getCompressionLevel() {
76+
return knnMethodConfigContext.getCompressionLevel();
77+
}
6878
}
6979
);
7080
return new MethodFieldMapper(

‎src/main/java/org/opensearch/knn/index/mapper/Mode.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ public enum Mode {
2626

2727
// Internally, an empty string is easier to deal with them null. However, from the mapping,
2828
// we do not want users to pass in the empty string and instead want null. So we make the conversion herex
29-
static final String[] NAMES_ARRAY = Arrays.stream(Mode.values())
29+
public static final String[] NAMES_ARRAY = Arrays.stream(Mode.values())
3030
.map(mode -> mode == NOT_CONFIGURED ? null : mode.getName())
3131
.collect(Collectors.toList())
3232
.toArray(new String[0]);

0 commit comments

Comments
 (0)