Collect query count in index metrics

src/java/org/apache/cassandra/index/sai/metrics/IndexMetrics.java

@@ -40,6 +40,7 @@ public class IndexMetrics extends AbstractMetrics
     public final Counter compactionCount;
     public final Counter memtableIndexFlushErrors;
     public final Counter segmentFlushErrors;
+    public final Counter queriesCount;
 
     public final Histogram memtableFlushCellsPerSecond;
     public final Histogram segmentsPerCompaction;
@@ -60,6 +61,7 @@ public IndexMetrics(IndexContext context)
         compactionCount = Metrics.counter(createMetricName("CompactionCount"));
         memtableIndexFlushErrors = Metrics.counter(createMetricName("MemtableIndexFlushErrors"));
         segmentFlushErrors = Metrics.counter(createMetricName("CompactionSegmentFlushErrors"));
+        queriesCount = Metrics.counter(createMetricName("QueriesCount"));
         liveMemtableIndexWriteCount = Metrics.register(createMetricName("LiveMemtableIndexWriteCount"), context::liveMemtableWriteCount);
         memtableOnHeapIndexBytes = Metrics.register(createMetricName("MemtableOnHeapIndexBytes"), context::estimatedOnHeapMemIndexMemoryUsed);
         memtableOffHeapIndexBytes = Metrics.register(createMetricName("MemtableOffHeapIndexBytes"), context::estimatedOffHeapMemIndexMemoryUsed);

src/java/org/apache/cassandra/index/sai/plan/Plan.java

@@ -18,12 +18,7 @@
 
 package org.apache.cassandra.index.sai.plan;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Comparator;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Objects;
+import java.util.*;
 import java.util.function.DoubleSupplier;
 import java.util.function.Function;
 import javax.annotation.Nonnull;
@@ -328,6 +323,15 @@ protected String description()
         return "";
     }
 
+    /**
+     * Adds an index to the set of indexes used by the plan node.
+     * Needs to be overwritten by nodes that use indexes.
+     */
+    protected void countIndex(HashSet<IndexContext> indexes)
+    {
+        // No index to count by default
+    }
+
     /**
      * Returns an optimized plan.
      * <p>
@@ -888,6 +892,13 @@ public String getIndexName()
             assert predicate != null || ordering != null;
             return predicate != null ? predicate.getIndexName() : ordering.getIndexName();
         }
+
+        @Override
+        final protected void countIndex(HashSet<IndexContext> indexes)
+        {
+            assert predicate != null || ordering != null;
+            indexes.add(predicate != null ? predicate.context : ordering.context);
+        }
     }
     /**
      * Represents a scan over a numeric storage attached index.
@@ -1349,6 +1360,12 @@ protected double estimateSelectivity()
         {
             return 1.0;
         }
+
+        @Override
+        protected void countIndex(HashSet<IndexContext> indexes)
+        {
+            indexes.add(ordering.context);
+        }
     }
 
 

src/java/org/apache/cassandra/index/sai/plan/QueryController.java

@@ -18,13 +18,7 @@
 
 package org.apache.cassandra.index.sai.plan;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -344,6 +338,15 @@ public SinglePartitionReadCommand getPartitionReadCommand(PrimaryKey key, ReadEx
                                                  makeFilter(key));
     }
 
+    private void updateIndexMetricsQueriesCount(Plan plan) {
+        HashSet<IndexContext> queriedIndexes = new HashSet<>();
+        plan.forEach(node -> {
+            node.countIndex(queriedIndexes);
+            return Plan.ControlFlow.Continue;
+        });
+        queriedIndexes.stream().forEach(indexContext -> indexContext.getIndexMetrics().queriesCount.inc());
+    }
+
     Plan buildPlan()
     {
         Plan.KeysIteration keysIterationPlan = buildKeysIterationPlan();
@@ -372,6 +375,8 @@ Plan buildPlan()
         if (plan.contains(node -> node instanceof Plan.KeysSort))
             queryContext.setFilterSortOrder(QueryContext.FilterSortOrder.SEARCH_THEN_ORDER);
 
+        updateIndexMetricsQueriesCount(plan);
+
         if (logger.isTraceEnabled())
             logger.trace("Query execution plan:\n" + plan.toStringRecursive());
 

test/unit/org/apache/cassandra/index/sai/metrics/IndexMetricsTest.java

@@ -142,4 +142,51 @@ public void testMetricsThroughWriteLifecycle()
 
         waitForVerifyHistogram(objectName("CompactionSegmentCellsPerSecond", keyspace, TABLE, INDEX, "IndexMetrics"), 1);
     }
+
+    private void assertIndexQueryCount(String index, long expectedCount)
+    {
+        assertEquals(expectedCount,
+                     getMetricValue(objectName("QueriesCount", KEYSPACE, currentTable(), index, "IndexMetrics")));
+    }
+
+    @Test
+    public void testQueriesCount()
+    {
+        createTable("CREATE TABLE %s (id1 TEXT PRIMARY KEY, v1 INT, v2 TEXT, v3 VECTOR<FLOAT, 2>)");
+        String indexV1 = createIndex("CREATE CUSTOM INDEX IF NOT EXISTS ON %s (v1) USING 'StorageAttachedIndex'");
+
+        int rowCount = 10;
+        for (int i = 0; i < rowCount; i++)
+            execute("INSERT INTO %s (id1, v1, v2, v3) VALUES (?, ?, '0', [?, 0.0])", Integer.toString(i), i, i);
+
+        assertIndexQueryCount(indexV1, 0L);
+
+        ResultSet rows = executeNet("SELECT id1 FROM %s WHERE v1 >= 0");
+        assertEquals(rowCount, rows.all().size());
+        assertIndexQueryCount(indexV1, 1L);
+
+        executeNet("SELECT id1 FROM %s WHERE (v1 >= 0 OR v1 = 4) AND v2 = '2' ALLOW FILTERING");
+        assertIndexQueryCount(indexV1, 2L);
+
+        String indexV2 = createIndex("CREATE CUSTOM INDEX IF NOT EXISTS ON %s (v2) USING 'StorageAttachedIndex'");
+        executeNet("SELECT id1 FROM %s WHERE (v1 >= 0 OR v1 = 4)");
+        assertIndexQueryCount(indexV1, 3L);
+        assertIndexQueryCount(indexV2, 0L);
+
+        executeNet("SELECT id1 FROM %s WHERE v2 = '2'");
+        assertIndexQueryCount(indexV2, 1L);
+        executeNet("SELECT id1 FROM %s WHERE (v1 >= 0 OR v1 = 4) AND v2 = '2'");
+        assertIndexQueryCount(indexV1, 4L);
+        assertIndexQueryCount(indexV2, 1L);
+        executeNet("SELECT id1 FROM %s WHERE (v1 >= 0 OR v1 = 4) ORDER BY v2 LIMIT 10");
+        assertIndexQueryCount(indexV1, 4L);
+        assertIndexQueryCount(indexV2, 2L);
+
+        String indexV3 = createIndex("CREATE CUSTOM INDEX IF NOT EXISTS ON %s (v3) USING 'StorageAttachedIndex'");
+        assertIndexQueryCount(indexV3, 0L);
+        executeNet("SELECT id1 FROM %s WHERE v2 = '2' ORDER BY v3 ANN OF [5,0] LIMIT 10");
+        assertIndexQueryCount(indexV1, 4L);
+        assertIndexQueryCount(indexV2, 2L);
+        assertIndexQueryCount(indexV3, 1L);
+    }
 }