Adding k-NN engine stat (opensearch-project#523)

* Adding field by engine stat

Signed-off-by: Martin Gaievski <[email protected]>

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<String, Object> clusterStats = parseClusterStatsResponse(responseBody);
-        assertEquals(knnStats.getClusterStats().keySet(), clusterStats.keySet());
+        assertNotNull(clusterStats);
+        assertTrue(knnStats.getClusterStats().keySet().containsAll(clusterStats.keySet()));
         List<Map<String, Object>> 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

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<Boolean> ignoreMalformed = ignoreMalformed(context);
             final Map<String, String> 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";
     }

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

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;
+    }
 }

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");
-    }
 }

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(

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()),

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()

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());
     }
 }

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<String, Object> 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<String, Object> 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<String, Object> 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<String, Object> 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() {

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
      */