Skip to content
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

Add field parameter for exact search. #2475

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/main/java/org/opensearch/knn/common/KNNConstants.java
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ public class KNNConstants {
public static final String METHOD_PARAMETER_SPACE_TYPE = "space_type"; // used for mapping parameter
// used for defining toplevel parameter
public static final String TOP_LEVEL_PARAMETER_SPACE_TYPE = METHOD_PARAMETER_SPACE_TYPE;

public static final String SEARCH_MODE_PARAMETER = "search_mode";
public static final String COMPOUND_EXTENSION = "c";
public static final String MODEL = "model";
public static final String MODELS = "models";
Expand Down
44 changes: 44 additions & 0 deletions src/main/java/org/opensearch/knn/index/SearchMode.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.knn.index;

import java.util.Arrays;
import java.util.Locale;

public enum SearchMode {
ANN("ann") {

},
EXACT("exact") {

};

private final String value;

SearchMode(String value) {
this.value = value;
}

public String getValue() {
return value;
}

public static SearchMode getSearchMode(String searchModeName) {
for (SearchMode currentSearchMode : SearchMode.values()) {
if (currentSearchMode.getValue().equalsIgnoreCase(searchModeName)) {
return currentSearchMode;
}
}
throw new IllegalArgumentException(
String.format(
Locale.ROOT,
"Unable to find search mode: %s . Valid values are: %s",
searchModeName,
Arrays.toString(SearchMode.values())
)
);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@

import lombok.extern.log4j.Log4j2;
import org.opensearch.common.StopWatch;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.SearchMode;
import org.opensearch.knn.index.VectorDataType;
import org.opensearch.knn.index.engine.KNNEngine;
import org.opensearch.knn.index.vectorvalues.KNNVectorValues;
Expand Down Expand Up @@ -90,6 +92,11 @@ public void merge(MergeState mergeState) {
assert mergeState != null;
assert mergeState.mergeFieldInfos != null;
for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
if (fieldInfo.attributes().containsKey(KNNConstants.SEARCH_MODE_PARAMETER)
&& SearchMode.EXACT.equals(fieldInfo.attributes().get(KNNConstants.SEARCH_MODE_PARAMETER))) {
logger.debug("Field with exact search mode, skipping graph creation during merge.");
continue;
}
DocValuesType type = fieldInfo.getDocValuesType();
if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) {
StopWatch stopWatch = new StopWatch();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,10 @@
import org.opensearch.index.mapper.ParametrizedFieldMapper;
import org.opensearch.index.mapper.ParseContext;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.KNNSettings;
import org.opensearch.knn.index.*;
import org.opensearch.knn.index.engine.EngineResolver;
import org.opensearch.knn.index.engine.KNNMethodConfigContext;
import org.opensearch.knn.index.engine.KNNMethodContext;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.index.VectorDataType;
import org.opensearch.knn.index.VectorField;
import org.opensearch.knn.index.engine.KNNEngine;
import org.opensearch.knn.index.engine.ResolvedMethodContext;
import org.opensearch.knn.index.engine.SpaceTypeResolver;
Expand Down Expand Up @@ -184,6 +181,13 @@ public static class Builder extends ParametrizedFieldMapper.Builder {
SpaceType.UNDEFINED.getValue()
).setValidator(SpaceType::getSpace);

protected final Parameter<String> searchMode = Parameter.stringParam(
KNNConstants.SEARCH_MODE_PARAMETER,
false,
m -> toType(m).originalMappingParameters.getSearchMode(),
SearchMode.ANN.getValue()
).setValidator(SearchMode::getSearchMode);

protected final Parameter<Map<String, String>> meta = Parameter.metaParam();

protected ModelDao modelDao;
Expand Down Expand Up @@ -221,7 +225,8 @@ protected List<Parameter<?>> getParameters() {
modelId,
mode,
compressionLevel,
topLevelSpaceType
topLevelSpaceType,
searchMode
);
}

Expand Down Expand Up @@ -378,6 +383,9 @@ public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserCont
);
}

// TODO: validate and make sure this would work even if knn setting is disabled.
validateSearchMode(builder);

// Check for flat configuration and validate only if index is created after 2.17
if (isKNNDisabled(parserContext.getSettings()) && parserContext.indexVersionCreated().onOrAfter(Version.V_2_17_0)) {
validateFromFlat(builder);
Expand All @@ -404,6 +412,17 @@ public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserCont
return builder;
}

private void validateSearchMode(KNNVectorFieldMapper.Builder builder) {
final KNNMethodContext knnMethodContext = builder.knnMethodContext.get();
final SpaceType spaceType = SpaceType.getSpace(builder.topLevelSpaceType.get());
final SearchMode searchMode = SearchMode.getSearchMode(builder.searchMode.get());
if (SearchMode.EXACT.equals(searchMode) && (knnMethodContext != null || spaceType != SpaceType.UNDEFINED)) {
throw new MapperParsingException(
"knnMethodContext or space type is not expected to be passed if the" + " field have exact search mode."
);
}
}

private void validateSpaceType(KNNVectorFieldMapper.Builder builder) {
final KNNMethodContext knnMethodContext = builder.knnMethodContext.get();
// if context is defined
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public final class OriginalMappingParameters {
private final String compressionLevel;
private final String modelId;
private final String topLevelSpaceType;
private final String searchMode;

/**
* Initialize the parameters from the builder
Expand All @@ -60,6 +61,7 @@ public OriginalMappingParameters(KNNVectorFieldMapper.Builder builder) {
this.compressionLevel = builder.compressionLevel.get();
this.modelId = builder.modelId.get();
this.topLevelSpaceType = builder.topLevelSpaceType.get();
this.searchMode = builder.searchMode.get();
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,7 @@
import org.opensearch.index.mapper.ParseContext;
import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.KNNSettings;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.index.VectorDataType;
import org.opensearch.knn.index.VectorField;
import org.opensearch.knn.index.*;
import org.opensearch.knn.index.codec.util.KNNVectorSerializerFactory;
import org.opensearch.knn.index.engine.KNNEngine;
import org.opensearch.knn.index.engine.KNNMethodConfigContext;
Expand Down Expand Up @@ -123,7 +120,16 @@ public void testBuilder_getParameters() {
modelDao,
CURRENT,
null,
new OriginalMappingParameters(VectorDataType.DEFAULT, TEST_DIMENSION, null, null, null, null, SpaceType.UNDEFINED.getValue())
new OriginalMappingParameters(
VectorDataType.DEFAULT,
TEST_DIMENSION,
null,
null,
null,
null,
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
)
);

assertEquals(10, builder.getParameters().size());
Expand Down Expand Up @@ -1125,7 +1131,8 @@ public void testMethodFieldMapperParseCreateField_validInput_thenDifferentFieldT
Mode.NOT_CONFIGURED.getName(),
CompressionLevel.NOT_CONFIGURED.getName(),
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
);
originalMappingParameters.setResolvedKnnMethodContext(knnMethodContext);
MethodFieldMapper methodFieldMapper = MethodFieldMapper.createFieldMapper(
Expand Down Expand Up @@ -1233,7 +1240,8 @@ public void testModelFieldMapperParseCreateField_validInput_thenDifferentFieldTy
Mode.NOT_CONFIGURED.getName(),
CompressionLevel.NOT_CONFIGURED.getName(),
MODEL_ID,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
);

ModelFieldMapper modelFieldMapper = ModelFieldMapper.createFieldMapper(
Expand Down Expand Up @@ -1328,7 +1336,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withFloats() {
Mode.NOT_CONFIGURED.getName(),
CompressionLevel.NOT_CONFIGURED.getName(),
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
);
originalMappingParameters.setResolvedKnnMethodContext(originalMappingParameters.getKnnMethodContext());

Expand Down Expand Up @@ -1388,7 +1397,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withFloats() {
Mode.NOT_CONFIGURED.getName(),
CompressionLevel.NOT_CONFIGURED.getName(),
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
);
originalMappingParameters.setResolvedKnnMethodContext(originalMappingParameters.getKnnMethodContext());
luceneFieldMapper = LuceneFieldMapper.createFieldMapper(
Expand Down Expand Up @@ -1429,7 +1439,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withBytes() {
Mode.NOT_CONFIGURED.getName(),
CompressionLevel.NOT_CONFIGURED.getName(),
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
);
originalMappingParameters.setResolvedKnnMethodContext(originalMappingParameters.getKnnMethodContext());

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package org.opensearch.knn.index.mapper;

import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.index.SearchMode;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.index.VectorDataType;
import org.opensearch.knn.index.engine.KNNEngine;
Expand All @@ -18,12 +19,28 @@ public class OriginalMappingParametersTests extends KNNTestCase {

public void testIsLegacy() {
assertTrue(
new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, null, null, SpaceType.UNDEFINED.getValue())
.isLegacyMapping()
new OriginalMappingParameters(
VectorDataType.DEFAULT,
123,
null,
null,
null,
null,
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
).isLegacyMapping()
);
assertFalse(
new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, null, "model-id", SpaceType.UNDEFINED.getValue())
.isLegacyMapping()
new OriginalMappingParameters(
VectorDataType.DEFAULT,
123,
null,
null,
null,
"model-id",
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
).isLegacyMapping()
);
assertFalse(
new OriginalMappingParameters(
Expand All @@ -33,7 +50,8 @@ public void testIsLegacy() {
Mode.ON_DISK.getName(),
null,
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
).isLegacyMapping()
);
assertFalse(
Expand All @@ -44,7 +62,8 @@ public void testIsLegacy() {
null,
CompressionLevel.x2.getName(),
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
).isLegacyMapping()
);
assertFalse(
Expand All @@ -55,7 +74,8 @@ public void testIsLegacy() {
null,
null,
null,
SpaceType.UNDEFINED.getValue()
SpaceType.UNDEFINED.getValue(),
SearchMode.ANN.getValue()
).isLegacyMapping()
);
}
Expand Down