diff --git a/release-notes/opensearch-knn.release-notes-2.17.0.0.md b/release-notes/opensearch-knn.release-notes-2.17.0.0.md index 4892d8b69..9876b7b38 100644 --- a/release-notes/opensearch-knn.release-notes-2.17.0.0.md +++ b/release-notes/opensearch-knn.release-notes-2.17.0.0.md @@ -8,6 +8,7 @@ Compatible with OpenSearch 2.17.0 * Add support for byte vector with Faiss Engine HNSW algorithm [#1823](https://github.com/opensearch-project/k-NN/pull/1823) * Add support for byte vector with Faiss Engine IVF algorithm [#2002](https://github.com/opensearch-project/k-NN/pull/2002) * Add mode/compression configuration support for disk-based vector search [#2034](https://github.com/opensearch-project/k-NN/pull/2034) +* Add spaceType as a top level optional parameter while creating vector field. [#2044](https://github.com/opensearch-project/k-NN/pull/2044) ### Enhancements * 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) ### Bug Fixes @@ -32,4 +33,4 @@ Compatible with OpenSearch 2.17.0 * Added Quantization Framework and implemented 1Bit and multibit quantizer[#1889](https://github.com/opensearch-project/k-NN/issues/1889) * Encapsulate dimension, vector data type validation/processing inside Library [#1957](https://github.com/opensearch-project/k-NN/pull/1957) * Add quantization state cache [#1960](https://github.com/opensearch-project/k-NN/pull/1960) -* Add quantization state reader and writer [#1997](https://github.com/opensearch-project/k-NN/pull/1997) \ No newline at end of file +* Add quantization state reader and writer [#1997](https://github.com/opensearch-project/k-NN/pull/1997) diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java index 441b4ea1c..15994e9d7 100644 --- a/src/main/java/org/opensearch/knn/common/KNNConstants.java +++ b/src/main/java/org/opensearch/knn/common/KNNConstants.java @@ -33,6 +33,8 @@ public class KNNConstants { public static final String METHOD_IVF = "ivf"; public static final String METHOD_PARAMETER_NLIST = "nlist"; 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 COMPOUND_EXTENSION = "c"; public static final String MODEL = "model"; public static final String MODELS = "models"; @@ -72,6 +74,7 @@ public class KNNConstants { public static final String MODEL_VECTOR_DATA_TYPE_KEY = VECTOR_DATA_TYPE_FIELD; public static final VectorDataType DEFAULT_VECTOR_DATA_TYPE_FIELD = VectorDataType.FLOAT; public static final String MINIMAL_MODE_AND_COMPRESSION_FEATURE = "mode_and_compression_feature"; + public static final String TOP_LEVEL_SPACE_TYPE_FEATURE = "top_level_space_type_feature"; public static final String RADIAL_SEARCH_KEY = "radial_search"; public static final String QUANTIZATION_STATE_FILE_SUFFIX = "osknnqstate"; diff --git a/src/main/java/org/opensearch/knn/index/SpaceType.java b/src/main/java/org/opensearch/knn/index/SpaceType.java index 43ff45e1d..abe265a01 100644 --- a/src/main/java/org/opensearch/knn/index/SpaceType.java +++ b/src/main/java/org/opensearch/knn/index/SpaceType.java @@ -11,10 +11,12 @@ package org.opensearch.knn.index; +import java.util.Arrays; import java.util.Locale; import java.util.HashSet; import java.util.Set; +import java.util.stream.Collectors; import static org.opensearch.knn.common.KNNVectorUtil.isZeroVector; @@ -149,6 +151,12 @@ public KNNVectorSimilarityFunction getKnnVectorSimilarityFunction() { public static SpaceType DEFAULT = L2; public static SpaceType DEFAULT_BINARY = HAMMING; + private static final String[] VALID_VALUES = Arrays.stream(SpaceType.values()) + .filter(space -> space != SpaceType.UNDEFINED) + .map(SpaceType::getValue) + .collect(Collectors.toList()) + .toArray(new String[0]); + private final String value; SpaceType(String value) { @@ -221,7 +229,9 @@ public static SpaceType getSpace(String spaceTypeName) { return currentSpaceType; } } - throw new IllegalArgumentException("Unable to find space: " + spaceTypeName); + throw new IllegalArgumentException( + String.format(Locale.ROOT, "Unable to find space: %s . Valid values are: %s", spaceTypeName, Arrays.toString(VALID_VALUES)) + ); } /** diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java index d2bb8e41a..10ba0d94d 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java @@ -161,6 +161,14 @@ public static class Builder extends ParametrizedFieldMapper.Builder { CompressionLevel.NAMES_ARRAY ).acceptsNull(); + // A top level space Type field. + protected final Parameter<String> topLevelSpaceType = Parameter.stringParam( + KNNConstants.TOP_LEVEL_PARAMETER_SPACE_TYPE, + false, + m -> toType(m).originalMappingParameters.getTopLevelSpaceType(), + SpaceType.UNDEFINED.getValue() + ).setValidator(SpaceType::getSpace); + protected final Parameter<Map<String, String>> meta = Parameter.metaParam(); protected ModelDao modelDao; @@ -187,7 +195,18 @@ public Builder( @Override protected List<Parameter<?>> getParameters() { - return Arrays.asList(stored, hasDocValues, dimension, vectorDataType, meta, knnMethodContext, modelId, mode, compressionLevel); + return Arrays.asList( + stored, + hasDocValues, + dimension, + vectorDataType, + meta, + knnMethodContext, + modelId, + mode, + compressionLevel, + topLevelSpaceType + ); } protected Explicit<Boolean> ignoreMalformed(BuilderContext context) { @@ -346,6 +365,7 @@ public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserCont validateFromModel(builder); } else { validateMode(builder); + validateSpaceType(builder); resolveKNNMethodComponents(builder, parserContext); validateFromKNNMethod(builder); } @@ -353,6 +373,23 @@ public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserCont return builder; } + private void validateSpaceType(KNNVectorFieldMapper.Builder builder) { + final KNNMethodContext knnMethodContext = builder.knnMethodContext.get(); + // if context is defined + if (knnMethodContext != null) { + // now ensure both space types are same. + final SpaceType knnMethodContextSpaceType = knnMethodContext.getSpaceType(); + final SpaceType topLevelSpaceType = SpaceType.getSpace(builder.topLevelSpaceType.get()); + if (topLevelSpaceType != SpaceType.UNDEFINED + && topLevelSpaceType != knnMethodContextSpaceType + && knnMethodContextSpaceType != SpaceType.UNDEFINED) { + throw new MapperParsingException( + "Space type in \"method\" and top level space type should be same or one of them should be defined" + ); + } + } + } + private void validateMode(KNNVectorFieldMapper.Builder builder) { boolean isKNNMethodContextConfigured = builder.originalParameters.getKnnMethodContext() != null; boolean isModeConfigured = builder.mode.isConfigured() || builder.compressionLevel.isConfigured(); @@ -386,6 +423,11 @@ private void validateFromModel(KNNVectorFieldMapper.Builder builder) { if (builder.dimension.getValue() == UNSET_MODEL_DIMENSION_IDENTIFIER && builder.modelId.get() == null) { throw new IllegalArgumentException(String.format(Locale.ROOT, "Dimension value missing for vector: %s", builder.name())); } + // ensure model and top level spaceType is not defined + if (builder.modelId.get() != null && SpaceType.getSpace(builder.topLevelSpaceType.get()) != SpaceType.UNDEFINED) { + throw new IllegalArgumentException("TopLevel Space type and model can not be both specified in the " + "mapping"); + } + validateCompressionAndModeNotSet(builder, builder.name(), "model"); } @@ -439,19 +481,33 @@ private void resolveKNNMethodComponents(KNNVectorFieldMapper.Builder builder, Pa // Configure method from map or legacy if (builder.originalParameters.isLegacyMapping()) { builder.originalParameters.setResolvedKnnMethodContext( - createKNNMethodContextFromLegacy(parserContext.getSettings(), parserContext.indexVersionCreated()) + createKNNMethodContextFromLegacy( + parserContext.getSettings(), + parserContext.indexVersionCreated(), + SpaceType.getSpace(builder.topLevelSpaceType.get()) + ) ); } else if (Mode.isConfigured(Mode.fromName(builder.mode.get())) || CompressionLevel.isConfigured(CompressionLevel.fromName(builder.compressionLevel.get()))) { + // we need don't need to resolve the space type, whatever default we are using will be passed down to + // while resolving KNNMethodContext for the mode and compression. and then when we resolve the spaceType + // we will set the correct spaceType. builder.originalParameters.setResolvedKnnMethodContext( ModeBasedResolver.INSTANCE.resolveKNNMethodContext( builder.knnMethodConfigContext.getMode(), builder.knnMethodConfigContext.getCompressionLevel(), - false + false, + SpaceType.getSpace(builder.originalParameters.getTopLevelSpaceType()) ) ); } - setDefaultSpaceType(builder.originalParameters.getResolvedKnnMethodContext(), builder.originalParameters.getVectorDataType()); + // this function should now correct the space type for the above resolved context too, if spaceType was + // not provided. + setSpaceType( + builder.originalParameters.getResolvedKnnMethodContext(), + builder.originalParameters.getVectorDataType(), + builder.topLevelSpaceType.get() + ); } private boolean isKNNDisabled(Settings settings) { @@ -459,16 +515,30 @@ private boolean isKNNDisabled(Settings settings) { return !isSettingPresent || !KNNSettings.IS_KNN_INDEX_SETTING.get(settings); } - private void setDefaultSpaceType(final KNNMethodContext knnMethodContext, final VectorDataType vectorDataType) { + private void setSpaceType( + final KNNMethodContext knnMethodContext, + final VectorDataType vectorDataType, + final String topLevelSpaceType + ) { + // Now KNNMethodContext should never be null. Because only case it could be null is flatMapper which is + // already handled if (knnMethodContext == null) { - return; + throw new IllegalArgumentException("KNNMethodContext cannot be null"); } - + final SpaceType topLevelSpaceTypeEnum = SpaceType.getSpace(topLevelSpaceType); + // Now set the spaceSpaceType for KNNMethodContext if (SpaceType.UNDEFINED == knnMethodContext.getSpaceType()) { - if (VectorDataType.BINARY == vectorDataType) { - knnMethodContext.setSpaceType(SpaceType.DEFAULT_BINARY); + // We are handling the case when top level space type is defined but method level spaceType is not + // defined. + if (topLevelSpaceTypeEnum != SpaceType.UNDEFINED) { + knnMethodContext.setSpaceType(topLevelSpaceTypeEnum); } else { - knnMethodContext.setSpaceType(SpaceType.DEFAULT); + // If both spaceTypes are undefined then put the default spaceType based on datatype + if (VectorDataType.BINARY == vectorDataType) { + knnMethodContext.setSpaceType(SpaceType.DEFAULT_BINARY); + } else { + knnMethodContext.setSpaceType(SpaceType.DEFAULT); + } } } } diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java index 551905793..b3727f2ef 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java @@ -193,10 +193,18 @@ private static int getEfConstruction(Settings indexSettings, Version indexVersio return Integer.parseInt(efConstruction); } - static KNNMethodContext createKNNMethodContextFromLegacy(Settings indexSettings, Version indexCreatedVersion) { + static KNNMethodContext createKNNMethodContextFromLegacy( + Settings indexSettings, + Version indexCreatedVersion, + SpaceType topLevelSpaceType + ) { + // If top level spaceType is set then use that spaceType otherwise default to spaceType from index-settings + final SpaceType finalSpaceToSet = topLevelSpaceType != SpaceType.UNDEFINED + ? topLevelSpaceType + : KNNVectorFieldMapperUtil.getSpaceType(indexSettings); return new KNNMethodContext( KNNEngine.NMSLIB, - KNNVectorFieldMapperUtil.getSpaceType(indexSettings), + finalSpaceToSet, new MethodComponentContext( METHOD_HNSW, Map.of( diff --git a/src/main/java/org/opensearch/knn/index/mapper/ModeBasedResolver.java b/src/main/java/org/opensearch/knn/index/mapper/ModeBasedResolver.java index 06b34fcd3..2a0c8ef46 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/ModeBasedResolver.java +++ b/src/main/java/org/opensearch/knn/index/mapper/ModeBasedResolver.java @@ -59,15 +59,19 @@ private ModeBasedResolver() {} * @param requiresTraining whether config requires trianing * @return {@link KNNMethodContext} */ - public KNNMethodContext resolveKNNMethodContext(Mode mode, CompressionLevel compressionLevel, boolean requiresTraining) { + public KNNMethodContext resolveKNNMethodContext( + Mode mode, + CompressionLevel compressionLevel, + boolean requiresTraining, + SpaceType spaceType + ) { if (requiresTraining) { - return resolveWithTraining(mode, compressionLevel); + return resolveWithTraining(mode, compressionLevel, spaceType); } - - return resolveWithoutTraining(mode, compressionLevel); + return resolveWithoutTraining(mode, compressionLevel, spaceType); } - private KNNMethodContext resolveWithoutTraining(Mode mode, CompressionLevel compressionLevel) { + private KNNMethodContext resolveWithoutTraining(Mode mode, CompressionLevel compressionLevel, final SpaceType spaceType) { CompressionLevel resolvedCompressionLevel = resolveCompressionLevel(mode, compressionLevel); MethodComponentContext encoderContext = resolveEncoder(resolvedCompressionLevel); @@ -76,7 +80,7 @@ private KNNMethodContext resolveWithoutTraining(Mode mode, CompressionLevel comp if (encoderContext != null) { return new KNNMethodContext( knnEngine, - SpaceType.DEFAULT, + spaceType, new MethodComponentContext( METHOD_HNSW, Map.of( @@ -96,7 +100,7 @@ private KNNMethodContext resolveWithoutTraining(Mode mode, CompressionLevel comp if (knnEngine == KNNEngine.FAISS) { return new KNNMethodContext( knnEngine, - SpaceType.DEFAULT, + spaceType, new MethodComponentContext( METHOD_HNSW, Map.of( @@ -113,7 +117,7 @@ private KNNMethodContext resolveWithoutTraining(Mode mode, CompressionLevel comp return new KNNMethodContext( knnEngine, - SpaceType.DEFAULT, + spaceType, new MethodComponentContext( METHOD_HNSW, Map.of( @@ -126,13 +130,13 @@ private KNNMethodContext resolveWithoutTraining(Mode mode, CompressionLevel comp ); } - private KNNMethodContext resolveWithTraining(Mode mode, CompressionLevel compressionLevel) { + private KNNMethodContext resolveWithTraining(Mode mode, CompressionLevel compressionLevel, SpaceType spaceType) { CompressionLevel resolvedCompressionLevel = resolveCompressionLevel(mode, compressionLevel); MethodComponentContext encoderContext = resolveEncoder(resolvedCompressionLevel); if (encoderContext != null) { return new KNNMethodContext( KNNEngine.FAISS, - SpaceType.DEFAULT, + spaceType, new MethodComponentContext( METHOD_IVF, Map.of( @@ -149,7 +153,7 @@ private KNNMethodContext resolveWithTraining(Mode mode, CompressionLevel compres return new KNNMethodContext( KNNEngine.FAISS, - SpaceType.DEFAULT, + spaceType, new MethodComponentContext( METHOD_IVF, Map.of(METHOD_PARAMETER_NLIST, METHOD_PARAMETER_NLIST_DEFAULT, METHOD_PARAMETER_NPROBES, METHOD_PARAMETER_NPROBES_DEFAULT) diff --git a/src/main/java/org/opensearch/knn/index/mapper/OriginalMappingParameters.java b/src/main/java/org/opensearch/knn/index/mapper/OriginalMappingParameters.java index f7235620f..77bf07a90 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/OriginalMappingParameters.java +++ b/src/main/java/org/opensearch/knn/index/mapper/OriginalMappingParameters.java @@ -42,6 +42,7 @@ public final class OriginalMappingParameters { private final String mode; private final String compressionLevel; private final String modelId; + private final String topLevelSpaceType; /** * Initialize the parameters from the builder @@ -56,6 +57,7 @@ public OriginalMappingParameters(KNNVectorFieldMapper.Builder builder) { this.mode = builder.mode.get(); this.compressionLevel = builder.compressionLevel.get(); this.modelId = builder.modelId.get(); + this.topLevelSpaceType = builder.topLevelSpaceType.get(); } /** diff --git a/src/main/java/org/opensearch/knn/index/util/IndexUtil.java b/src/main/java/org/opensearch/knn/index/util/IndexUtil.java index 70d17da81..b279f1181 100644 --- a/src/main/java/org/opensearch/knn/index/util/IndexUtil.java +++ b/src/main/java/org/opensearch/knn/index/util/IndexUtil.java @@ -52,6 +52,7 @@ public class IndexUtil { private static final Version MINIMAL_SUPPORTED_VERSION_FOR_MODEL_VECTOR_DATA_TYPE = Version.V_2_16_0; private static final Version MINIMAL_RESCORE_FEATURE = Version.V_2_17_0; private static final Version MINIMAL_MODE_AND_COMPRESSION_FEATURE = Version.V_2_17_0; + private static final Version MINIMAL_TOP_LEVEL_SPACE_TYPE_FEATURE = Version.V_2_17_0; // public so neural search can access it public static final Map<String, Version> minimalRequiredVersionMap = initializeMinimalRequiredVersionMap(); public static final Set<VectorDataType> VECTOR_DATA_TYPES_NOT_SUPPORTING_ENCODERS = Set.of(VectorDataType.BINARY, VectorDataType.BYTE); @@ -390,6 +391,7 @@ private static Map<String, Version> initializeMinimalRequiredVersionMap() { put(KNNConstants.MODEL_VECTOR_DATA_TYPE_KEY, MINIMAL_SUPPORTED_VERSION_FOR_MODEL_VECTOR_DATA_TYPE); put(RESCORE_PARAMETER, MINIMAL_RESCORE_FEATURE); put(KNNConstants.MINIMAL_MODE_AND_COMPRESSION_FEATURE, MINIMAL_MODE_AND_COMPRESSION_FEATURE); + put(KNNConstants.TOP_LEVEL_SPACE_TYPE_FEATURE, MINIMAL_TOP_LEVEL_SPACE_TYPE_FEATURE); } }; diff --git a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java index 8770449eb..0f7e8523b 100644 --- a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java +++ b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java @@ -95,6 +95,7 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr int dimension = DEFAULT_NOT_SET_INT_VALUE; int maximumVectorCount = DEFAULT_NOT_SET_INT_VALUE; int searchSize = DEFAULT_NOT_SET_INT_VALUE; + SpaceType topLevelSpaceType = SpaceType.UNDEFINED; String compressionLevel = null; String mode = null; @@ -109,9 +110,6 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr trainingField = parser.textOrNull(); } else if (KNN_METHOD.equals(fieldName) && ensureNotSet(fieldName, knnMethodContext)) { knnMethodContext = KNNMethodContext.parse(parser.map()); - if (SpaceType.UNDEFINED == knnMethodContext.getSpaceType()) { - knnMethodContext.setSpaceType(SpaceType.L2); - } } else if (DIMENSION.equals(fieldName) && ensureNotSet(fieldName, dimension)) { dimension = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false); } else if (MAX_VECTOR_COUNT_PARAMETER.equals(fieldName) && ensureNotSet(fieldName, maximumVectorCount)) { @@ -127,6 +125,8 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr mode = parser.text(); } else if (KNNConstants.COMPRESSION_LEVEL_PARAMETER.equals(fieldName) && ensureNotSet(fieldName, compressionLevel)) { compressionLevel = parser.text(); + } else if (KNNConstants.SPACE_TYPE.equals(fieldName) && ensureSpaceTypeNotSet(topLevelSpaceType)) { + topLevelSpaceType = SpaceType.getSpace(parser.text()); } else { throw new IllegalArgumentException("Unable to parse token. \"" + fieldName + "\" is not a valid " + "parameter."); } @@ -160,7 +160,11 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr vectorDataType, VectorDataType.FLOAT.getValue() ); - + resolveSpaceTypeAndSetInKNNMethodContext(topLevelSpaceType, knnMethodContext); + // if KNNMethodContext was not null then spaceTypes we should fix the space type if it is not set. + if (knnMethodContext == null && topLevelSpaceType == SpaceType.UNDEFINED) { + topLevelSpaceType = SpaceType.DEFAULT; + } TrainingModelRequest trainingModelRequest = new TrainingModelRequest( modelId, knnMethodContext, @@ -171,7 +175,8 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr description, vectorDataType, Mode.fromName(mode), - CompressionLevel.fromName(compressionLevel) + CompressionLevel.fromName(compressionLevel), + topLevelSpaceType ); if (maximumVectorCount != DEFAULT_NOT_SET_INT_VALUE) { @@ -205,6 +210,35 @@ private void ensureMutualExclusion(String fieldNameA, Object valueA, String fiel } } + private boolean ensureSpaceTypeNotSet(SpaceType spaceType) { + if (spaceType != SpaceType.UNDEFINED) { + throw new IllegalArgumentException("Unable to parse SpaceType as it is duplicated."); + } + return true; + } + + private void resolveSpaceTypeAndSetInKNNMethodContext(SpaceType topLevelSpaceType, KNNMethodContext knnMethodContext) { + // First check if KNNMethodContext is not null as it can be null + if (knnMethodContext != null) { + // if space type is not provided by user then it will undefined + if (knnMethodContext.getSpaceType() == SpaceType.UNDEFINED) { + // fix the top level spaceType if it is undefined + if (topLevelSpaceType == SpaceType.UNDEFINED) { + topLevelSpaceType = SpaceType.DEFAULT; + } + // set the space type now in KNNMethodContext + knnMethodContext.setSpaceType(topLevelSpaceType); + } else { + // if spaceType is set at 2 places lets ensure that we validate those cases and throw error + if (topLevelSpaceType != SpaceType.UNDEFINED) { + throw new IllegalArgumentException( + "Top Level spaceType and space type in method both are set. Set space type at 1 place." + ); + } + } + } + } + private void ensureIfSetThenEquals( String fieldNameA, Object valueA, diff --git a/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java b/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java index 82669d4a8..df24baf0e 100644 --- a/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java +++ b/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java @@ -21,6 +21,7 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.engine.KNNMethodConfigContext; import org.opensearch.knn.index.mapper.CompressionLevel; import org.opensearch.knn.index.mapper.Mode; @@ -56,6 +57,33 @@ public class TrainingModelRequest extends ActionRequest { private final Mode mode; private final CompressionLevel compressionLevel; + TrainingModelRequest( + String modelId, + KNNMethodContext knnMethodContext, + int dimension, + String trainingIndex, + String trainingField, + String preferredNodeId, + String description, + VectorDataType vectorDataType, + Mode mode, + CompressionLevel compressionLevel + ) { + this( + modelId, + knnMethodContext, + dimension, + trainingIndex, + trainingField, + preferredNodeId, + description, + vectorDataType, + mode, + compressionLevel, + SpaceType.DEFAULT + ); + } + /** * Constructor. * @@ -77,7 +105,8 @@ public TrainingModelRequest( String description, VectorDataType vectorDataType, Mode mode, - CompressionLevel compressionLevel + CompressionLevel compressionLevel, + SpaceType spaceType ) { super(); this.modelId = modelId; @@ -107,7 +136,7 @@ public TrainingModelRequest( .build(); if (knnMethodContext == null && (Mode.isConfigured(mode) || CompressionLevel.isConfigured(compressionLevel))) { - this.knnMethodContext = ModeBasedResolver.INSTANCE.resolveKNNMethodContext(mode, compressionLevel, true); + this.knnMethodContext = ModeBasedResolver.INSTANCE.resolveKNNMethodContext(mode, compressionLevel, true, spaceType); } else { this.knnMethodContext = knnMethodContext; } diff --git a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java index 6d0a3d5df..abc4f563e 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java @@ -36,6 +36,7 @@ import org.opensearch.index.mapper.MapperService; 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; @@ -117,10 +118,10 @@ public void testBuilder_getParameters() { modelDao, CURRENT, null, - new OriginalMappingParameters(VectorDataType.DEFAULT, TEST_DIMENSION, null, null, null, null) + new OriginalMappingParameters(VectorDataType.DEFAULT, TEST_DIMENSION, null, null, null, null, SpaceType.UNDEFINED.getValue()) ); - assertEquals(9, builder.getParameters().size()); + assertEquals(10, builder.getParameters().size()); List<String> actualParams = builder.getParameters().stream().map(a -> a.name).collect(Collectors.toList()); List<String> expectedParams = Arrays.asList( "store", @@ -131,7 +132,8 @@ public void testBuilder_getParameters() { KNN_METHOD, MODEL_ID, MODE_PARAMETER, - COMPRESSION_LEVEL_PARAMETER + COMPRESSION_LEVEL_PARAMETER, + KNNConstants.TOP_LEVEL_PARAMETER_SPACE_TYPE ); assertEquals(expectedParams, actualParams); } @@ -899,7 +901,8 @@ public void testMethodFieldMapperParseCreateField_validInput_thenDifferentFieldT knnMethodContext, Mode.NOT_CONFIGURED.getName(), CompressionLevel.NOT_CONFIGURED.getName(), - null + null, + SpaceType.UNDEFINED.getValue() ); originalMappingParameters.setResolvedKnnMethodContext(knnMethodContext); MethodFieldMapper methodFieldMapper = MethodFieldMapper.createFieldMapper( @@ -1000,7 +1003,8 @@ public void testModelFieldMapperParseCreateField_validInput_thenDifferentFieldTy null, Mode.NOT_CONFIGURED.getName(), CompressionLevel.NOT_CONFIGURED.getName(), - MODEL_ID + MODEL_ID, + SpaceType.UNDEFINED.getValue() ); ModelFieldMapper modelFieldMapper = ModelFieldMapper.createFieldMapper( @@ -1092,7 +1096,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withFloats() { getDefaultKNNMethodContext(), Mode.NOT_CONFIGURED.getName(), CompressionLevel.NOT_CONFIGURED.getName(), - null + null, + SpaceType.UNDEFINED.getValue() ); originalMappingParameters.setResolvedKnnMethodContext(originalMappingParameters.getKnnMethodContext()); @@ -1151,7 +1156,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withFloats() { knnMethodContext, Mode.NOT_CONFIGURED.getName(), CompressionLevel.NOT_CONFIGURED.getName(), - null + null, + SpaceType.UNDEFINED.getValue() ); originalMappingParameters.setResolvedKnnMethodContext(originalMappingParameters.getKnnMethodContext()); luceneFieldMapper = LuceneFieldMapper.createFieldMapper( @@ -1191,7 +1197,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withBytes() { getDefaultByteKNNMethodContext(), Mode.NOT_CONFIGURED.getName(), CompressionLevel.NOT_CONFIGURED.getName(), - null + null, + SpaceType.UNDEFINED.getValue() ); originalMappingParameters.setResolvedKnnMethodContext(originalMappingParameters.getKnnMethodContext()); diff --git a/src/test/java/org/opensearch/knn/index/mapper/OriginalMappingParametersTests.java b/src/test/java/org/opensearch/knn/index/mapper/OriginalMappingParametersTests.java index 2822a882e..4b089b149 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/OriginalMappingParametersTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/OriginalMappingParametersTests.java @@ -17,11 +17,35 @@ public class OriginalMappingParametersTests extends KNNTestCase { public void testIsLegacy() { - assertTrue(new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, null, null).isLegacyMapping()); - assertFalse(new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, null, "model-id").isLegacyMapping()); - assertFalse(new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, Mode.ON_DISK.getName(), null, null).isLegacyMapping()); + assertTrue( + new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, null, null, SpaceType.UNDEFINED.getValue()) + .isLegacyMapping() + ); + assertFalse( + new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, null, "model-id", SpaceType.UNDEFINED.getValue()) + .isLegacyMapping() + ); + assertFalse( + new OriginalMappingParameters( + VectorDataType.DEFAULT, + 123, + null, + Mode.ON_DISK.getName(), + null, + null, + SpaceType.UNDEFINED.getValue() + ).isLegacyMapping() + ); assertFalse( - new OriginalMappingParameters(VectorDataType.DEFAULT, 123, null, null, CompressionLevel.x2.getName(), null).isLegacyMapping() + new OriginalMappingParameters( + VectorDataType.DEFAULT, + 123, + null, + null, + CompressionLevel.x2.getName(), + null, + SpaceType.UNDEFINED.getValue() + ).isLegacyMapping() ); assertFalse( new OriginalMappingParameters( @@ -30,7 +54,8 @@ public void testIsLegacy() { new KNNMethodContext(KNNEngine.DEFAULT, SpaceType.L2, new MethodComponentContext(null, Collections.emptyMap())), null, null, - null + null, + SpaceType.UNDEFINED.getValue() ).isLegacyMapping() ); }