diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/StatsIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/StatsIT.java index df9babd77..cca6a45c7 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/StatsIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/StatsIT.java @@ -20,14 +20,16 @@ public class StatsIT extends AbstractRollingUpgradeTestCase { private KNNStats knnStats = new KNNStats(KNN_STATS); - // Validate if all the KNN Stats metrics are returned + // Validate if all the KNN Stats metrics from old version are present in new version public void testAllMetricStatsReturned() throws IOException { Response response = getKnnStats(Collections.emptyList(), Collections.emptyList()); String responseBody = EntityUtils.toString(response.getEntity()); Map clusterStats = parseClusterStatsResponse(responseBody); - assertEquals(knnStats.getClusterStats().keySet(), clusterStats.keySet()); + assertNotNull(clusterStats); + assertTrue(knnStats.getClusterStats().keySet().containsAll(clusterStats.keySet())); List> nodeStats = parseNodeStatsResponse(responseBody); - assertEquals(knnStats.getNodeStats().keySet(), nodeStats.get(0).keySet()); + assertNotNull(nodeStats.get(0)); + assertTrue(knnStats.getNodeStats().keySet().containsAll(nodeStats.get(0).keySet())); } // Verify if it returns failure for invalid metric 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 96fea1fbc..3964dd98f 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java @@ -197,6 +197,7 @@ public KNNVectorFieldMapper build(BuilderContext context) { final CopyTo copyToBuilder = copyTo.build(); final Explicit ignoreMalformed = ignoreMalformed(context); final Map metaValue = meta.getValue(); + if (knnMethodContext != null) { final KNNVectorFieldType mappedFieldType = new KNNVectorFieldType( buildFullName(context), @@ -407,6 +408,7 @@ public KNNVectorFieldMapper( this.stored = stored; this.hasDocValues = hasDocValues; this.dimension = mappedFieldType.getDimension(); + updateEngineStats(); } public KNNVectorFieldMapper clone() { @@ -538,6 +540,11 @@ protected void doXContentBody(XContentBuilder builder, boolean includeDefaults, } } + /** + * Overwrite at child level in case specific stat needs to be updated + */ + void updateEngineStats() {} + public static class Names { public static final String IGNORE_MALFORMED = "ignore_malformed"; } diff --git a/src/main/java/org/opensearch/knn/index/mapper/LuceneFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/LuceneFieldMapper.java index 5b6c65486..5dcb09318 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/LuceneFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/LuceneFieldMapper.java @@ -105,6 +105,11 @@ protected void parseCreateField(ParseContext context, int dimension) throws IOEx context.path().remove(); } + @Override + void updateEngineStats() { + KNNEngine.LUCENE.setInitialized(true); + } + @AllArgsConstructor @lombok.Builder @Getter diff --git a/src/main/java/org/opensearch/knn/index/util/JVMLibrary.java b/src/main/java/org/opensearch/knn/index/util/JVMLibrary.java index f6d4dc283..e1d48cb0a 100644 --- a/src/main/java/org/opensearch/knn/index/util/JVMLibrary.java +++ b/src/main/java/org/opensearch/knn/index/util/JVMLibrary.java @@ -15,6 +15,8 @@ */ public abstract class JVMLibrary extends AbstractKNNLibrary { + boolean initialized; + /** * Constructor * @@ -29,4 +31,14 @@ public abstract class JVMLibrary extends AbstractKNNLibrary { public int estimateOverheadInKB(KNNMethodContext knnMethodContext, int dimension) { throw new UnsupportedOperationException("Estimating overhead is not supported for JVM based libraries."); } + + @Override + public Boolean isInitialized() { + return initialized; + } + + @Override + public void setInitialized(Boolean isInitialized) { + initialized = isInitialized; + } } diff --git a/src/main/java/org/opensearch/knn/index/util/Lucene.java b/src/main/java/org/opensearch/knn/index/util/Lucene.java index df1353106..e1996353f 100644 --- a/src/main/java/org/opensearch/knn/index/util/Lucene.java +++ b/src/main/java/org/opensearch/knn/index/util/Lucene.java @@ -73,14 +73,4 @@ public float score(float rawScore, SpaceType spaceType) { // score provided. return rawScore; } - - @Override - public Boolean isInitialized() { - return true; - } - - @Override - public void setInitialized(Boolean isInitialized) { - throw new UnsupportedOperationException("Setting Lucene as initialized is not supported"); - } } diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java index 56089ed84..c41170b32 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java @@ -75,6 +75,7 @@ public class KNNStatsConfig { ) .put(StatNames.FAISS_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.FAISS))) .put(StatNames.NMSLIB_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.NMSLIB))) + .put(StatNames.LUCENE_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.LUCENE))) .put(StatNames.TRAINING_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.TRAINING_REQUESTS))) .put(StatNames.TRAINING_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.TRAINING_ERRORS))) .put( diff --git a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java index f807c3eaa..ffe5882bb 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java @@ -26,6 +26,7 @@ public enum StatNames { MODEL_INDEX_STATUS("model_index_status"), FAISS_LOADED("faiss_initialized"), NMSLIB_LOADED("nmslib_initialized"), + LUCENE_LOADED("lucene_initialized"), INDEXING_FROM_MODEL_DEGRADED("indexing_from_model_degraded"), GRAPH_QUERY_ERRORS(KNNCounter.GRAPH_QUERY_ERRORS.getName()), GRAPH_QUERY_REQUESTS(KNNCounter.GRAPH_QUERY_REQUESTS.getName()), 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 13e2ab9a3..3c311f86b 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java @@ -193,6 +193,8 @@ public void testBuilder_parse_fromKnnMethodContext_luceneEngine() throws IOExcep ModelDao modelDao = mock(ModelDao.class); KNNVectorFieldMapper.TypeParser typeParser = new KNNVectorFieldMapper.TypeParser(() -> modelDao); + KNNEngine.LUCENE.setInitialized(false); + int efConstruction = 321; int m = 12; int dimension = 133; @@ -215,12 +217,15 @@ public void testBuilder_parse_fromKnnMethodContext_luceneEngine() throws IOExcep xContentBuilderToMap(xContentBuilder), buildParserContext(indexName, settings) ); + Mapper.BuilderContext builderContext = new Mapper.BuilderContext(settings, new ContentPath()); + builder.build(builderContext); assertEquals(METHOD_HNSW, builder.knnMethodContext.get().getMethodComponent().getName()); assertEquals( efConstruction, builder.knnMethodContext.get().getMethodComponent().getParameters().get(METHOD_PARAMETER_EF_CONSTRUCTION) ); + assertTrue(KNNEngine.LUCENE.isInitialized()); XContentBuilder xContentBuilderEmptyParams = XContentFactory.jsonBuilder() .startObject() diff --git a/src/test/java/org/opensearch/knn/index/util/LuceneTests.java b/src/test/java/org/opensearch/knn/index/util/LuceneTests.java index d5fffc1bb..57e02e173 100644 --- a/src/test/java/org/opensearch/knn/index/util/LuceneTests.java +++ b/src/test/java/org/opensearch/knn/index/util/LuceneTests.java @@ -100,11 +100,12 @@ public void testScore() { public void testIsInitialized() { Lucene luceneLibrary = new Lucene(Collections.emptyMap(), ""); - assertTrue(luceneLibrary.isInitialized()); + assertFalse(luceneLibrary.isInitialized()); } public void testSetInitialized() { Lucene luceneLibrary = new Lucene(Collections.emptyMap(), ""); - expectThrows(UnsupportedOperationException.class, () -> luceneLibrary.setInitialized(true)); + luceneLibrary.setInitialized(true); + assertTrue(luceneLibrary.isInitialized()); } } diff --git a/src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java b/src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java index 10fa47c4c..2b6bb757e 100644 --- a/src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java +++ b/src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java @@ -14,8 +14,11 @@ import org.opensearch.client.Response; import org.opensearch.client.ResponseException; import org.opensearch.cluster.health.ClusterHealthStatus; +import org.opensearch.common.Strings; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.xcontent.XContentBuilder; +import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.common.xcontent.XContentType; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.QueryBuilder; @@ -33,6 +36,20 @@ import java.util.List; import java.util.Map; +import static org.opensearch.knn.TestUtils.KNN_VECTOR; +import static org.opensearch.knn.TestUtils.PROPERTIES; +import static org.opensearch.knn.TestUtils.VECTOR_TYPE; +import static org.opensearch.knn.common.KNNConstants.FAISS_NAME; +import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE; +import static org.opensearch.knn.common.KNNConstants.LUCENE_NAME; +import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; +import static org.opensearch.knn.common.KNNConstants.METHOD_IVF; +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NLIST; +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_SPACE_TYPE; +import static org.opensearch.knn.common.KNNConstants.MODEL_ID; +import static org.opensearch.knn.common.KNNConstants.NAME; +import static org.opensearch.knn.common.KNNConstants.NMSLIB_NAME; +import static org.opensearch.knn.common.KNNConstants.PARAMETERS; import static org.opensearch.knn.plugin.stats.KNNStatsConfig.KNN_STATS; /** @@ -41,8 +58,20 @@ public class RestKNNStatsHandlerIT extends KNNRestTestCase { private static final Logger logger = LogManager.getLogger(RestKNNStatsHandlerIT.class); + private static final String TRAINING_INDEX = "training-index"; + private static final String TRAINING_FIELD = "training-field"; + private static final String TEST_MODEL_ID = "model-id"; + private static final String TEST_INDEX = "test-index"; + private static final String MODEL_DESCRIPTION = "Description for train model test"; private boolean isDebuggingTest = new DisableOnDebug(null).isDebugging(); private boolean isDebuggingRemoteCluster = System.getProperty("cluster.debug", "false").equals("true"); + private static final String FIELD_NAME_2 = "test_field_two"; + private static final String FIELD_NAME_3 = "test_field_three"; + private static final int DIMENSION = 4; + private static int DOC_ID = 0; + private static final int NUM_DOCS = 10; + private static final int DELAY_MILLI_SEC = 1000; + private static final int NUM_OF_ATTEMPTS = 30; private KNNStats knnStats; @@ -333,6 +362,98 @@ public void testModelIndexingDegradedMetricsStats() throws IOException { assertEquals(false, nodeStats.get(statName)); } + /** + * Test checks that handler correctly returns value for field per engine stats + * + * @throws IOException throws IOException + */ + public void testFieldByEngineStats() throws Exception { + createKnnIndex(INDEX_NAME, createKnnIndexMapping(FIELD_NAME, 2, METHOD_HNSW, NMSLIB_NAME)); + putMappingRequest(INDEX_NAME, createKnnIndexMapping(FIELD_NAME_2, 3, METHOD_HNSW, LUCENE_NAME)); + putMappingRequest(INDEX_NAME, createKnnIndexMapping(FIELD_NAME_3, 3, METHOD_HNSW, FAISS_NAME)); + + Response response = getKnnStats(Collections.emptyList(), Collections.emptyList()); + + String responseBody = EntityUtils.toString(response.getEntity()); + + Map nodeStats0 = parseNodeStatsResponse(responseBody).get(0); + boolean faissField = (Boolean) nodeStats0.get(StatNames.FAISS_LOADED.getName()); + boolean luceneField = (Boolean) nodeStats0.get(StatNames.LUCENE_LOADED.getName()); + boolean nmslibField = (Boolean) nodeStats0.get(StatNames.NMSLIB_LOADED.getName()); + + assertTrue(faissField); + assertTrue(luceneField); + assertTrue(nmslibField); + } + + public void testFieldsByEngineModelTraining() throws Exception { + createBasicKnnIndex(TRAINING_INDEX, TRAINING_FIELD, DIMENSION); + bulkIngestRandomVectors(TRAINING_INDEX, TRAINING_FIELD, NUM_DOCS, DIMENSION); + trainKnnModel(TEST_MODEL_ID, TRAINING_INDEX, TRAINING_FIELD, DIMENSION, MODEL_DESCRIPTION); + + validateModelCreated(TEST_MODEL_ID); + + createKnnIndex(TEST_INDEX, modelIndexMapping(FIELD_NAME, TEST_MODEL_ID)); + + addKNNDocs(TEST_INDEX, FIELD_NAME, DIMENSION, DOC_ID, NUM_DOCS); + + final Response response = getKnnStats(Collections.emptyList(), Collections.emptyList()); + final String responseBody = EntityUtils.toString(response.getEntity()); + + final Map nodeStats0 = parseNodeStatsResponse(responseBody).get(0); + + boolean faissField = (Boolean) nodeStats0.get(StatNames.FAISS_LOADED.getName()); + + assertTrue(faissField); + } + + public void trainKnnModel(String modelId, String trainingIndexName, String trainingFieldName, int dimension, String description) + throws IOException { + XContentBuilder builder = XContentFactory.jsonBuilder() + .startObject() + .field(NAME, METHOD_IVF) + .field(KNN_ENGINE, FAISS_NAME) + .field(METHOD_PARAMETER_SPACE_TYPE, SpaceType.L2.getValue()) + .startObject(PARAMETERS) + .field(METHOD_PARAMETER_NLIST, 1) + .endObject() + .endObject(); + Map method = xContentBuilderToMap(builder); + + Response trainResponse = trainModel(modelId, trainingIndexName, trainingFieldName, dimension, method, description); + assertEquals(RestStatus.OK, RestStatus.fromCode(trainResponse.getStatusLine().getStatusCode())); + } + + public void validateModelCreated(String modelId) throws IOException, InterruptedException { + Response getResponse = getModel(modelId, null); + String responseBody = EntityUtils.toString(getResponse.getEntity()); + assertNotNull(responseBody); + + Map responseMap = createParser(XContentType.JSON.xContent(), responseBody).map(); + assertEquals(modelId, responseMap.get(MODEL_ID)); + assertTrainingSucceeds(modelId, NUM_OF_ATTEMPTS, DELAY_MILLI_SEC); + } + + // mapping to create index from model + public String modelIndexMapping(String fieldName, String modelId) throws IOException { + return Strings.toString( + XContentFactory.jsonBuilder() + .startObject() + .startObject(PROPERTIES) + .startObject(fieldName) + .field(VECTOR_TYPE, KNN_VECTOR) + .field(MODEL_ID, modelId) + .endObject() + .endObject() + .endObject() + ); + } + + @Override + protected boolean preserveClusterUponCompletion() { + return false; + } + // Useful settings when debugging to prevent timeouts @Override protected Settings restClientSettings() { diff --git a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java index 3edf508ab..797b038d9 100644 --- a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java +++ b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java @@ -297,6 +297,27 @@ protected String createKnnIndexMapping(String fieldName, Integer dimensions) thr ); } + /** + * Utility to create a Knn Index Mapping with specific algorithm and engine + */ + protected String createKnnIndexMapping(String fieldName, Integer dimensions, String algoName, String knnEngine) throws IOException { + return Strings.toString( + XContentFactory.jsonBuilder() + .startObject() + .startObject("properties") + .startObject(fieldName) + .field("type", "knn_vector") + .field("dimension", dimensions.toString()) + .startObject("method") + .field("name", algoName) + .field("engine", knnEngine) + .endObject() + .endObject() + .endObject() + .endObject() + ); + } + /** * Utility to create a Knn Index Mapping with multiple k-NN fields */