Skip to content

Commit

Permalink
Update test method
Browse files Browse the repository at this point in the history
Signed-off-by: Vijayan Balasubramanian <[email protected]>
  • Loading branch information
VijayanB committed Oct 30, 2024
1 parent 6cf4504 commit 73bd246
Show file tree
Hide file tree
Showing 6 changed files with 71 additions and 38 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,23 @@ public void testKNNIndexDefaultLegacyFieldMapping() throws Exception {
}
}

// ensure that index is created using legacy mapping in old cluster, and, then, add docs to both old and new cluster.
// when search is requested on new cluster it should return all docs irrespective of cluster.
public void testKNNIndexDefaultEngine() throws Exception {
waitForClusterHealthGreen(NODES_BWC_CLUSTER);
if (isRunningAgainstOldCluster()) {
createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS));
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, 5);
// Flush to ensure that index is not re-indexed when node comes back up
flush(testIndex, true);
} else {
validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, 5, 5);
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, 5, 5);
validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, 10, 10);
deleteKNNIndex(testIndex);
}
}

// Ensure that when segments created with old mapping are forcemerged in new cluster, they
// succeed
public void testKNNIndexDefaultLegacyFieldMappingForceMerge() throws Exception {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -500,9 +500,7 @@ private void resolveKNNMethodComponents(
.build()
);

// If the original parameters are from legacy, and it is created on or before 2_17_2 since default is changed to
// FAISS starting 2_18, which doesn't support accepting algo params from index settings
if (parserContext.indexVersionCreated().onOrBefore(Version.V_2_17_2) && builder.originalParameters.isLegacyMapping()) {
if (useKNNMethodContextFromLegacy(builder, parserContext)) {
// Then create KNNMethodContext to be used from the legacy index settings
builder.originalParameters.setResolvedKnnMethodContext(
createKNNMethodContextFromLegacy(parserContext.getSettings(), parserContext.indexVersionCreated(), resolvedSpaceType)
Expand Down Expand Up @@ -551,6 +549,12 @@ private void setEngine(final KNNMethodContext knnMethodContext, KNNEngine knnEng
}
}

static boolean useKNNMethodContextFromLegacy(Builder builder, Mapper.TypeParser.ParserContext parserContext) {
// If the original parameters are from legacy, and it is created on or before 2_17_2 since default is changed to
// FAISS starting 2_18, which doesn't support accepting algo params from index settings
return parserContext.indexVersionCreated().onOrBefore(Version.V_2_17_2) && builder.originalParameters.isLegacyMapping();
}

// We store the version of the index with the mapper as different version of Opensearch has different default
// values of KNN engine Algorithms hyperparameters.
protected Version indexCreatedVersion;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ private Map<Integer, Float> doANNSearch(
spaceType = modelMetadata.getSpaceType();
vectorDataType = modelMetadata.getVectorDataType();
} else {
String engineName = fieldInfo.attributes().getOrDefault(KNN_ENGINE, KNNEngine.NMSLIB.getName());
String engineName = fieldInfo.attributes().getOrDefault(KNN_ENGINE, KNNEngine.DEFAULT.getName());
knnEngine = KNNEngine.getEngine(engineName);
String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue());
spaceType = SpaceType.getSpace(spaceTypeName);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,10 +41,10 @@ public void testResolveEngine_whenModeAndCompressionAreFalse_thenDefault() {
);
}

public void testResolveEngine_whenModeSpecifiedAndCompressionIsNotSpecified_thenDefault() {
public void testResolveEngine_whenModeSpecifiedAndCompressionIsNotSpecified_thenNMSLIB() {
assertEquals(KNNEngine.DEFAULT, ENGINE_RESOLVER.resolveEngine(KNNMethodConfigContext.builder().build(), null, false));
assertEquals(
KNNEngine.DEFAULT,
KNNEngine.NMSLIB,
ENGINE_RESOLVER.resolveEngine(
KNNMethodConfigContext.builder().mode(Mode.IN_MEMORY).build(),
new KNNMethodContext(KNNEngine.DEFAULT, SpaceType.UNDEFINED, MethodComponentContext.EMPTY, false),
Expand All @@ -63,7 +63,7 @@ public void testResolveEngine_whenCompressionIs1x_thenEngineBasedOnMode() {
)
);
assertEquals(
KNNEngine.DEFAULT,
KNNEngine.NMSLIB,
ENGINE_RESOLVER.resolveEngine(KNNMethodConfigContext.builder().compressionLevel(CompressionLevel.x1).build(), null, false)
);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@
import java.util.Optional;
import java.util.stream.Collectors;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
Expand Down Expand Up @@ -284,36 +285,43 @@ public void testTypeParser_withDifferentSpaceTypeCombinations_thenSuccess() thro
);
assertTrue(knnVectorFieldMapper.fieldType().getKnnMappingConfig().getModelId().isEmpty());

// if space type is provided and legacy mappings is hit
xContentBuilder = XContentFactory.jsonBuilder()
.startObject()
.field(TYPE_FIELD_NAME, KNN_VECTOR_TYPE)
.field(DIMENSION_FIELD_NAME, TEST_DIMENSION)
.field(KNNConstants.TOP_LEVEL_PARAMETER_SPACE_TYPE, topLevelSpaceType.getValue())
.endObject();
builder = (KNNVectorFieldMapper.Builder) typeParser.parse(
"test-field-name-1",
xContentBuilderToMap(xContentBuilder),
buildParserContext("test", settings)
);
// mock useKNNMethodContextFromLegacy to simulate index ix created before 2_18
try (MockedStatic<KNNVectorFieldMapper> utilMockedStatic = Mockito.mockStatic(KNNVectorFieldMapper.class)) {
utilMockedStatic.when(() -> KNNVectorFieldMapper.useKNNMethodContextFromLegacy(any(), any())).thenReturn(true);
// if space type is provided and legacy mappings is hit
xContentBuilder = XContentFactory.jsonBuilder()
.startObject()
.field(TYPE_FIELD_NAME, KNN_VECTOR_TYPE)
.field(DIMENSION_FIELD_NAME, TEST_DIMENSION)
.field(KNNConstants.TOP_LEVEL_PARAMETER_SPACE_TYPE, topLevelSpaceType.getValue())
.endObject();
builder = (KNNVectorFieldMapper.Builder) typeParser.parse(
"test-field-name-1",
xContentBuilderToMap(xContentBuilder),
buildParserContext("test", settings)
);

builderContext = new Mapper.BuilderContext(settings, new ContentPath());
knnVectorFieldMapper = builder.build(builderContext);
assertTrue(knnVectorFieldMapper instanceof MethodFieldMapper);
assertTrue(knnVectorFieldMapper.fieldType().getKnnMappingConfig().getKnnMethodContext().isPresent());
assertEquals(topLevelSpaceType, knnVectorFieldMapper.fieldType().getKnnMappingConfig().getKnnMethodContext().get().getSpaceType());
// this check ensures that legacy mapping is hit, as in legacy mapping we pick M from index settings
assertEquals(
mForSetting,
knnVectorFieldMapper.fieldType()
.getKnnMappingConfig()
.getKnnMethodContext()
.get()
.getMethodComponentContext()
.getParameters()
.get(METHOD_PARAMETER_M)
);
assertTrue(knnVectorFieldMapper.fieldType().getKnnMappingConfig().getModelId().isEmpty());
builderContext = new Mapper.BuilderContext(settings, new ContentPath());
knnVectorFieldMapper = builder.build(builderContext);
assertTrue(knnVectorFieldMapper instanceof MethodFieldMapper);
assertTrue(knnVectorFieldMapper.fieldType().getKnnMappingConfig().getKnnMethodContext().isPresent());
assertEquals(
topLevelSpaceType,
knnVectorFieldMapper.fieldType().getKnnMappingConfig().getKnnMethodContext().get().getSpaceType()
);
// this check ensures that legacy mapping is hit, as in legacy mapping we pick M from index settings
assertEquals(
mForSetting,
knnVectorFieldMapper.fieldType()
.getKnnMappingConfig()
.getKnnMethodContext()
.get()
.getMethodComponentContext()
.getParameters()
.get(METHOD_PARAMETER_M)
);
assertTrue(knnVectorFieldMapper.fieldType().getKnnMappingConfig().getModelId().isEmpty());
}
}

public void testTypeParser_withSpaceTypeAndMode_thenSuccess() throws IOException {
Expand Down Expand Up @@ -1703,7 +1711,7 @@ public void testTypeParser_whenModeAndCompressionAreSet_thenHandle() throws IOEx
assertFalse(builder.getOriginalParameters().isLegacyMapping());
validateBuilderAfterParsing(
builder,
KNNEngine.DEFAULT,
KNNEngine.NMSLIB,
SpaceType.L2,
VectorDataType.FLOAT,
CompressionLevel.x1,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,11 @@ public void tearDownAfterTest() {
@SneakyThrows
public void testQueryResultScoreNmslib() {
for (SpaceType space : List.of(SpaceType.L2, SpaceType.L1, SpaceType.COSINESIMIL, SpaceType.INNER_PRODUCT, SpaceType.LINF)) {
testQueryScore(space::scoreTranslation, SEGMENT_FILES_NMSLIB, Map.of(SPACE_TYPE, space.getValue()));
testQueryScore(
space::scoreTranslation,
SEGMENT_FILES_NMSLIB,
Map.of(SPACE_TYPE, space.getValue(), KNN_ENGINE, KNNEngine.NMSLIB.getName())
);
}
}

Expand Down

0 comments on commit 73bd246

Please sign in to comment.