Skip to content

Commit

Permalink
Add bwc test for model id
Browse files Browse the repository at this point in the history 
This is applicable for index created before 2.17

Signed-off-by: Vijayan Balasubramanian <[email protected]>
  • Loading branch information
@VijayanB
VijayanB committed Dec 10, 2024
1 parent 5ede283 commit d1857aa
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 18 deletions.
13 changes: 13 additions & 0 deletions qa/restart-upgrade/build.gradle
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,12 @@ testClusters {
excludeTestsMatching "org.opensearch.knn.bwc.IndexingIT.testKNNIndexLuceneForceMerge"
}
}
if (!(knn_bwc_version.startsWith("2.15.") ||
knn_bwc_version.startsWith("2.16."))){
filter {
excludeTestsMatching "org.opensearch.knn.bwc.ModelIT.testNonKNNIndex_withModelId"
}
}

nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
Expand Down Expand Up @@ -170,6 +176,13 @@ testClusters {
}
}

if (!(knn_bwc_version.startsWith("2.15.") ||
knn_bwc_version.startsWith("2.16."))){
filter {
excludeTestsMatching "org.opensearch.knn.bwc.ModelIT.testNonKNNIndex_withModelId"
}
}

nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
systemProperty 'tests.security.manager', 'false'
Expand Down
59 changes: 57 additions & 2 deletions qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ModelIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,26 +41,32 @@
public class ModelIT extends AbstractRestartUpgradeTestCase {
private static final String TEST_MODEL_INDEX = KNN_BWC_PREFIX + "test-model-index";
private static final String TEST_MODEL_INDEX_DEFAULT = KNN_BWC_PREFIX + "test-model-index-default";
private static final String TEST_MODEL_INDEX_FOR_NON_KNN_INDEX = KNN_BWC_PREFIX + "test-model-index-for-non-knn-index";
private static final String TRAINING_INDEX = KNN_BWC_PREFIX + "train-index";
private static final String TRAINING_INDEX_DEFAULT = KNN_BWC_PREFIX + "train-index-default";
private static final String TRAINING_INDEX_FOR_NON_KNN_INDEX = KNN_BWC_PREFIX + "train-index-for-non-knn-index";
private static final String TRAINING_FIELD = "train-field";
private static final String TEST_FIELD = "test-field";
private static final int DIMENSIONS = 5;
private static int DOC_ID = 0;
private static int DOC_ID_TEST_MODEL_INDEX = 0;
private static int DOC_ID_TEST_MODEL_INDEX_DEFAULT = 0;
private static int DOC_ID_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX = 0;
private static final int DELAY_MILLI_SEC = 1000;
private static final int EXP_NUM_OF_MODELS = 2;
private static final int MIN_NUM_OF_MODELS = 2;
private static final int K = 5;
private static final int NUM_DOCS = 10;
private static final int NUM_DOCS_TEST_MODEL_INDEX = 100;
private static final int NUM_DOCS_TEST_MODEL_INDEX_DEFAULT = 100;
private static final int NUM_DOCS_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX = 100;
private static final int NUM_OF_ATTEMPTS = 30;
private static int QUERY_COUNT = 0;
private static int QUERY_COUNT_TEST_MODEL_INDEX = 0;
private static int QUERY_COUNT_TEST_MODEL_INDEX_DEFAULT = 0;
private static int QUERY_COUNT_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX = 0;
private static final String TEST_MODEL_ID = "test-model-id";
private static final String TEST_MODEL_ID_DEFAULT = "test-model-id-default";
private static final String TEST_MODEL_ID_FOR_NON_KNN_INDEX = "test-model-id-for-non-knn-index";
private static final String MODEL_DESCRIPTION = "Description for train model test";

// KNN model test
Expand Down Expand Up @@ -135,6 +141,55 @@ public void testKNNModelDefault() throws Exception {
}
}

public void testNonKNNIndex_withModelId() throws Exception {
if (isRunningAgainstOldCluster()) {

// Create a training index and randomly ingest data into it
createBasicKnnIndex(TRAINING_INDEX_FOR_NON_KNN_INDEX, TRAINING_FIELD, DIMENSIONS);
bulkIngestRandomVectors(TRAINING_INDEX_FOR_NON_KNN_INDEX, TRAINING_FIELD, NUM_DOCS, DIMENSIONS);

trainKNNModel(TEST_MODEL_ID_FOR_NON_KNN_INDEX, TRAINING_INDEX_FOR_NON_KNN_INDEX, TRAINING_FIELD, DIMENSIONS, MODEL_DESCRIPTION);
validateModelCreated(TEST_MODEL_ID_FOR_NON_KNN_INDEX);

createKnnIndex(
testIndex,
createKNNDefaultScriptScoreSettings(),
modelIndexMapping(TEST_FIELD, TEST_MODEL_ID_FOR_NON_KNN_INDEX)
);
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
} else {
Thread.sleep(1000);
DOC_ID = NUM_DOCS;
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
QUERY_COUNT = 2 * NUM_DOCS;
validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, QUERY_COUNT, K);

searchKNNModel(TEST_MODEL_ID_FOR_NON_KNN_INDEX);

createKnnIndex(TEST_MODEL_INDEX_FOR_NON_KNN_INDEX, modelIndexMapping(TEST_FIELD, TEST_MODEL_ID_FOR_NON_KNN_INDEX));
addKNNDocs(
TEST_MODEL_INDEX_FOR_NON_KNN_INDEX,
TEST_FIELD,
DIMENSIONS,
DOC_ID_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX,
NUM_DOCS_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX
);
QUERY_COUNT_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX = NUM_DOCS_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX;
validateKNNSearch(
TEST_MODEL_INDEX_FOR_NON_KNN_INDEX,
TEST_FIELD,
DIMENSIONS,
QUERY_COUNT_TEST_MODEL_INDEX_FOR_NON_KNN_INDEX,
K
);

deleteKNNIndex(testIndex);
deleteKNNIndex(TRAINING_INDEX_FOR_NON_KNN_INDEX);
deleteKNNIndex(TEST_MODEL_INDEX_FOR_NON_KNN_INDEX);
deleteKNNModel(TEST_MODEL_INDEX_FOR_NON_KNN_INDEX);
}
}

// Delete Models and ".opensearch-knn-models" index to clear cluster metadata
@AfterClass
public static void wipeAllModels() throws IOException {
Expand Down Expand Up @@ -168,7 +223,7 @@ public void searchKNNModel(String testModelID) throws Exception {
XContentParser parser = createParser(MediaTypeRegistry.getDefaultMediaType().xContent(), responseBody);
SearchResponse searchResponse = SearchResponse.fromXContent(parser);
assertNotNull(searchResponse);
assertEquals(EXP_NUM_OF_MODELS, searchResponse.getHits().getHits().length);
assertTrue(MIN_NUM_OF_MODELS <= searchResponse.getHits().getHits().length);

for (SearchHit hit : searchResponse.getHits().getHits()) {
assertTrue(hit.getId().startsWith(testModelID));
Expand Down
16 changes: 0 additions & 16 deletions src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -396,22 +396,6 @@ protected String createKnnIndexMapping(String fieldName, Integer dimensions) thr
.toString();
}

/**
* Utility to create a Knn Index Mapping with model id
*/
protected String createKnnIndexMapping(String fieldName, String modelId) throws IOException {
return XContentFactory.jsonBuilder()
.startObject()
.startObject("properties")
.startObject(fieldName)
.field("type", "knn_vector")
.field("model_id", modelId)
.endObject()
.endObject()
.endObject()
.toString();
}

protected String createKnnIndexMapping(final String fieldName, final Integer dimensions, final VectorDataType vectorDataType)
throws IOException {
return XContentFactory.jsonBuilder()
Expand Down

0 comments on commit d1857aa

Please sign in to comment.