Skip to content

Commit

Permalink
Fixes for bwc and IT tests
Browse files Browse the repository at this point in the history 
Signed-off-by: AnnTian Shao <[email protected]>
  • Loading branch information
AnnTian Shao committed Jan 24, 2025
1 parent e9d4807 commit cd1b553
Show file tree
Hide file tree
Showing 8 changed files with 22 additions and 39 deletions.
4 changes: 2 additions & 2 deletions qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/FaissSQIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ public void testIVFSQFP16_onUpgradeWhenIndexedAndQueried_thenSucceed() throws Ex

// Add training data
createBasicKnnIndex(TRAIN_INDEX, TRAIN_TEST_FIELD, DIMENSION);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectors(TRAIN_INDEX, TRAIN_TEST_FIELD, trainingDataCount, DIMENSION);

XContentBuilder builder = XContentFactory.jsonBuilder()
Expand Down Expand Up @@ -278,7 +278,7 @@ public void testIVFSQFP16_onUpgradeWhenClipToFp16isTrueAndIndexedWithOutOfFP16Ra

// Add training data
createBasicKnnIndex(TRAIN_INDEX, TRAIN_TEST_FIELD, dimension);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectors(TRAIN_INDEX, TRAIN_TEST_FIELD, trainingDataCount, dimension);

XContentBuilder builder = XContentFactory.jsonBuilder()
Expand Down
16 changes: 8 additions & 8 deletions src/test/java/org/opensearch/knn/index/FaissIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@ public void testEndToEnd_whenDoRadiusSearch_whenDistanceThreshold_whenMethodIsHN

// training data needs to be at least equal to the number of centroids for PQ
// which is 2^8 = 256. 8 because that's the only valid code_size for HNSWPQ
int trainingDataCount = 256;
int trainingDataCount = 1100;

SpaceType spaceType = SpaceType.L2;

Expand Down Expand Up @@ -468,7 +468,7 @@ public void testEndToEnd_whenMethodIsHNSWPQ_thenSucceed() {

// training data needs to be at least equal to the number of centroids for PQ
// which is 2^8 = 256. 8 because thats the only valid code_size for HNSWPQ
int trainingDataCount = 256;
int trainingDataCount = 1100;

SpaceType spaceType = SpaceType.L2;

Expand Down Expand Up @@ -736,7 +736,7 @@ public void testIVFSQFP16_whenIndexedAndQueried_thenSucceed() {

// Add training data
createBasicKnnIndex(trainingIndexName, trainingFieldName, dimension);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectors(trainingIndexName, trainingFieldName, trainingDataCount, dimension);

XContentBuilder builder = XContentFactory.jsonBuilder()
Expand Down Expand Up @@ -960,7 +960,7 @@ public void testIVFSQFP16_whenIndexedWithOutOfFP16Range_thenThrowException() {

// Add training data
createBasicKnnIndex(trainingIndexName, trainingFieldName, dimension);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectors(trainingIndexName, trainingFieldName, trainingDataCount, dimension);

XContentBuilder builder = XContentFactory.jsonBuilder()
Expand Down Expand Up @@ -1064,7 +1064,7 @@ public void testIVFSQFP16_whenClipToFp16isTrueAndIndexedWithOutOfFP16Range_thenS

// Add training data
createBasicKnnIndex(trainingIndexName, trainingFieldName, dimension);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectors(trainingIndexName, trainingFieldName, trainingDataCount, dimension);

XContentBuilder builder = XContentFactory.jsonBuilder()
Expand Down Expand Up @@ -1144,7 +1144,7 @@ public void testEndToEnd_whenMethodIsHNSWPQAndHyperParametersNotSet_thenSucceed(

// training data needs to be at least equal to the number of centroids for PQ
// which is 2^8 = 256. 8 because thats the only valid code_size for HNSWPQ
int trainingDataCount = 256;
int trainingDataCount = 1100;

SpaceType spaceType = SpaceType.L2;

Expand Down Expand Up @@ -1412,7 +1412,7 @@ public void testKNNQuery_withModelDifferentCombination_thenSuccess() throws Exce

// Add training data
createBasicKnnIndex(trainingIndexName, trainingFieldName, dimension);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectors(trainingIndexName, trainingFieldName, trainingDataCount, dimension);

// Call train API - IVF with nlists = 1 is brute force, but will require training
Expand Down Expand Up @@ -1767,7 +1767,7 @@ public void testIVF_whenBinaryFormat_whenIVF_thenSuccess() {

createKnnIndex(trainingIndexName, trainIndexMapping);

int trainingDataCount = 40;
int trainingDataCount = 1100;
bulkIngestRandomBinaryVectors(trainingIndexName, trainingFieldName, trainingDataCount, dimension);

XContentBuilder trainModelXContentBuilder = XContentFactory.jsonBuilder()
Expand Down
2 changes: 1 addition & 1 deletion src/test/java/org/opensearch/knn/index/VectorDataTypeIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -603,7 +603,7 @@ public void testIVFByteVector_whenIndexedAndQueried_thenSucceed() {
.toString();
createKnnIndex(INDEX_NAME, trainIndexMapping);

int trainingDataCount = 100;
int trainingDataCount = 1100;
bulkIngestRandomByteVectors(INDEX_NAME, FIELD_NAME, trainingDataCount, dimension);

XContentBuilder trainModelXContentBuilder = XContentFactory.jsonBuilder()
Expand Down
2 changes: 1 addition & 1 deletion src/test/java/org/opensearch/knn/integ/KNNScriptScoringIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -620,7 +620,7 @@ public void testKNNScriptScoreOnModelBasedIndex() throws Exception {
int dimensions = randomIntBetween(2, 10);
String trainMapping = createKnnIndexMapping(TRAIN_FIELD_PARAMETER, dimensions);
createKnnIndex(TRAIN_INDEX_PARAMETER, trainMapping);
bulkIngestRandomVectors(TRAIN_INDEX_PARAMETER, TRAIN_FIELD_PARAMETER, dimensions * 3, dimensions);
bulkIngestRandomVectors(TRAIN_INDEX_PARAMETER, TRAIN_FIELD_PARAMETER, 1100, dimensions);

XContentBuilder methodBuilder = XContentFactory.jsonBuilder()
.startObject()
Expand Down
2 changes: 1 addition & 1 deletion src/test/java/org/opensearch/knn/integ/ModeAndCompressionIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ public class ModeAndCompressionIT extends KNNRestTestCase {

private static final String TRAINING_INDEX_NAME = "training_index";
private static final String TRAINING_FIELD_NAME = "training_field";
private static final int TRAINING_VECS = 20;
private static final int TRAINING_VECS = 1100;

private static final int DIMENSION = 16;
private static final int NUM_DOCS = 20;
Expand Down
2 changes: 1 addition & 1 deletion src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ public class RestKNNStatsHandlerIT extends KNNRestTestCase {
private static final String FIELD_LUCENE_NAME = "lucene_test_field";
private static final int DIMENSION = 4;
private static int DOC_ID = 0;
private static final int NUM_DOCS = 10;
private static final int NUM_DOCS = 1100;
private static final int DELAY_MILLI_SEC = 1000;
private static final int NUM_OF_ATTEMPTS = 30;

Expand Down
31 changes: 7 additions & 24 deletions src/test/java/org/opensearch/knn/plugin/action/RestTrainModelHandlerIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,28 +97,11 @@ public void testTrainModel_fail_notEnoughData() throws Exception {
.endObject();
Map<String, Object> method = xContentBuilderToMap(builder);

Response trainResponse = trainModel(null, trainingIndexName, trainingFieldName, dimension, method, "dummy description");

assertEquals(RestStatus.OK, RestStatus.fromCode(trainResponse.getStatusLine().getStatusCode()));

// Grab the model id from the response
String trainResponseBody = EntityUtils.toString(trainResponse.getEntity());
assertNotNull(trainResponseBody);

Map<String, Object> trainResponseMap = createParser(XContentType.JSON.xContent(), trainResponseBody).map();
String modelId = (String) trainResponseMap.get(MODEL_ID);
assertNotNull(modelId);

// Confirm that the model fails to create
Response getResponse = getModel(modelId, null);
String responseBody = EntityUtils.toString(getResponse.getEntity());
assertNotNull(responseBody);

Map<String, Object> responseMap = createParser(XContentType.JSON.xContent(), responseBody).map();

assertEquals(modelId, responseMap.get(MODEL_ID));

assertTrainingFails(modelId, 30, 1000);
ResponseException exception = expectThrows(
ResponseException.class,
() -> trainModel(null, trainingIndexName, trainingFieldName, dimension, method, "dummy description")
);
assertTrue(exception.getMessage().contains("Number of training points should be greater than"));
}

public void testTrainModel_fail_tooMuchData() throws Exception {
Expand All @@ -132,7 +115,7 @@ public void testTrainModel_fail_tooMuchData() throws Exception {

// Create a training index and randomly ingest data into it
createBasicKnnIndex(trainingIndexName, trainingFieldName, dimension);
int trainingDataCount = 20; // 20 * 16 * 4 ~= 10 kb
int trainingDataCount = 128;
bulkIngestRandomVectors(trainingIndexName, trainingFieldName, trainingDataCount, dimension);

// Call the train API with this definition:
Expand Down Expand Up @@ -491,7 +474,7 @@ public void testTrainModel_success_methodOverrideWithCompressionMode() throws Ex
// Create a training index and randomly ingest data into it
String mapping = createKnnIndexNestedMapping(dimension, nestedFieldPath);
createKnnIndex(trainingIndexName, mapping);
int trainingDataCount = 200;
int trainingDataCount = 1100;
bulkIngestRandomVectorsWithNestedField(trainingIndexName, nestedFieldPath, trainingDataCount, dimension);

// Call the train API with this definition:
Expand Down
2 changes: 1 addition & 1 deletion src/test/java/org/opensearch/knn/recall/RecallTestsIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ public class RecallTestsIT extends KNNRestTestCase {
private final static String TRAIN_FIELD_NAME = "train_field";
private final static String TEST_MODEL_ID = "test_model_id";
private final static int TEST_DIMENSION = 32;
private final static int DOC_COUNT = 500;
private final static int DOC_COUNT = 1100;
private final static int QUERY_COUNT = 100;
private final static int TEST_K = 100;
private final static double PERFECT_RECALL = 1.0;
Expand Down

0 comments on commit cd1b553

Please sign in to comment.