Adding stat for query with filter

Signed-off-by: Martin Gaievski <[email protected]>
opensearch-project · Oct 20, 2022 · 6f7821f · 6f7821f
1 parent 44f10de
commit 6f7821f
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java b/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java
@@ -158,6 +158,7 @@ public static KNNQueryBuilder fromXContent(XContentParser parser) throws IOExcep
                         String tokenName = parser.currentName();
                         if (FILTER_FIELD.getPreferredName().equals(tokenName)) {
                             log.debug(String.format("Start parsing filter for field [%s]", fieldName));
+                            KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.increment();
                             if (isClusterOnOrAfterMinRequiredVersion()) {
                                 filter = parseInnerQueryBuilder(parser);
                             } else {

diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java
@@ -21,7 +21,8 @@ public enum KNNCounter {
     SCRIPT_QUERY_REQUESTS("script_query_requests"),
     SCRIPT_QUERY_ERRORS("script_query_errors"),
     TRAINING_REQUESTS("training_requests"),
-    TRAINING_ERRORS("training_errors");
+    TRAINING_ERRORS("training_errors"),
+    KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests");
 
     private String name;
     private AtomicLong count;

diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java
@@ -55,6 +55,10 @@ public class KNNStatsConfig {
         .put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new KNNStat<>(true, new KNNCircuitBreakerSupplier()))
         .put(StatNames.MODEL_INDEX_STATUS.getName(), new KNNStat<>(true, new ModelIndexStatusSupplier<>(ModelDao::getHealthStatus)))
         .put(StatNames.KNN_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS)))
+        .put(
+            StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName(),
+            new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS))
+        )
         .put(StatNames.SCRIPT_COMPILATIONS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATIONS)))
         .put(
             StatNames.SCRIPT_COMPILATION_ERRORS.getName(),

diff --git a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java
@@ -40,7 +40,8 @@ public enum StatNames {
     TRAINING_ERRORS(KNNCounter.TRAINING_ERRORS.getName()),
     TRAINING_MEMORY_USAGE("training_memory_usage"),
     TRAINING_MEMORY_USAGE_PERCENTAGE("training_memory_usage_percentage"),
-    SCRIPT_QUERY_ERRORS(KNNCounter.SCRIPT_QUERY_ERRORS.getName());
+    SCRIPT_QUERY_ERRORS(KNNCounter.SCRIPT_QUERY_ERRORS.getName()),
+    KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName());
 
     private String name;
 

diff --git a/src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java b/src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java
@@ -22,6 +22,7 @@
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.knn.KNNRestTestCase;
 import org.opensearch.knn.index.query.KNNQueryBuilder;
 import org.opensearch.knn.index.SpaceType;
@@ -67,6 +68,7 @@ public class RestKNNStatsHandlerIT extends KNNRestTestCase {
     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 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;
@@ -106,6 +108,7 @@ public void testStatsValueCheck() throws Exception {
         Map<String, Object> nodeStats0 = parseNodeStatsResponse(responseBody).get(0);
         Integer hitCount0 = (Integer) nodeStats0.get(StatNames.HIT_COUNT.getName());
         Integer missCount0 = (Integer) nodeStats0.get(StatNames.MISS_COUNT.getName());
+        Integer knnQueryWithFilterCount0 = (Integer) nodeStats0.get(StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName());
 
         // Setup index
         createKnnIndex(INDEX_NAME, createKnnIndexMapping(FIELD_NAME, 2));
@@ -124,9 +127,11 @@ public void testStatsValueCheck() throws Exception {
         Map<String, Object> nodeStats1 = parseNodeStatsResponse(responseBody).get(0);
         Integer hitCount1 = (Integer) nodeStats1.get(StatNames.HIT_COUNT.getName());
         Integer missCount1 = (Integer) nodeStats1.get(StatNames.MISS_COUNT.getName());
+        Integer knnQueryWithFilterCount1 = (Integer) nodeStats1.get(StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName());
 
         assertEquals(hitCount0, hitCount1);
         assertEquals((Integer) (missCount0 + 1), missCount1);
+        assertEquals(knnQueryWithFilterCount0, knnQueryWithFilterCount1);
 
         // Second search: Ensure that hits=1
         searchKNNIndex(INDEX_NAME, new KNNQueryBuilder(FIELD_NAME, qvector, 1), 1);
@@ -137,9 +142,24 @@ public void testStatsValueCheck() throws Exception {
         Map<String, Object> nodeStats2 = parseNodeStatsResponse(responseBody).get(0);
         Integer hitCount2 = (Integer) nodeStats2.get(StatNames.HIT_COUNT.getName());
         Integer missCount2 = (Integer) nodeStats2.get(StatNames.MISS_COUNT.getName());
+        Integer knnQueryWithFilterCount2 = (Integer) nodeStats2.get(StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName());
 
         assertEquals(missCount1, missCount2);
         assertEquals((Integer) (hitCount1 + 1), hitCount2);
+        assertEquals(knnQueryWithFilterCount0, knnQueryWithFilterCount2);
+
+        putMappingRequest(INDEX_NAME, createKnnIndexMapping(FIELD_LUCENE_NAME, 2, METHOD_HNSW, LUCENE_NAME));
+        addKnnDoc(INDEX_NAME, "2", FIELD_LUCENE_NAME, vector);
+
+        searchKNNIndex(INDEX_NAME, new KNNQueryBuilder(FIELD_LUCENE_NAME, qvector, 1, QueryBuilders.termQuery("_id", "1")), 1);
+
+        response = getKnnStats(Collections.emptyList(), Collections.emptyList());
+        responseBody = EntityUtils.toString(response.getEntity());
+
+        Map<String, Object> nodeStats3 = parseNodeStatsResponse(responseBody).get(0);
+        Integer knnQueryWithFilterCount3 = (Integer) nodeStats3.get(StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName());
+
+        assertEquals((Integer) (knnQueryWithFilterCount0 + 1), knnQueryWithFilterCount3);
     }
 
     /**