CNDB-13583: Add vector ann and brute force metrics (#1683)

### What is the issue
riptano/cndb#13583

### What does this PR fix and why was it fixed
This PR adds comprehensive metrics for Storage Attached Indexes (SAI)
vector search operations, providing crucial insights into both ANN
(Approximate Nearest Neighbor) graph searches and brute force
operations.

New Vector Search Metrics:

Search Operation Counters:
- `ANNNodesVisited`: Total number of nodes visited during ANN searches
(this is equivalent to approximate similarity score computations)
- `ANNNodesReranked`: Number of nodes that underwent exact distance
computation for reranking (this is equivalent to exact similarity score
computations)
- `ANNNodesExpanded`: Total number of nodes whose edges were explored
- `ANNNodesExpandedBaseLayer`: Number of nodes expanded in the base
layer of the graph
- `ANNGraphSearches`: Count of new graph searches initiated
- `ANNGraphResumes`: Count of resumed graph searches (when a search
continues from previous results)
- `ANNGraphSearchLatency`: Timer measuring individual graph search
latency (Note: This measures per-graph search time, not total query time
which may involve multiple graphs)

Brute Force Operation Counters:
- `BruteForceNodesVisited`: Number of nodes visited during brute force
searches (approximate similarity comparisons)
- `BruteForceNodesReranked`: Number of nodes that underwent exact
similarity computation during brute force searches

Memory Usage Tracking:
- `quantizationMemoryBytes`: Current memory usage by the quantization
(PQ or BQ) data structures
- `ordinalsMapMemoryBytes`: Current memory usage by ordinals mapping
structures (only matters in some cases)
- `onDiskGraphsCount`: Number of currently loaded graph segments
- `onDiskGraphVectorsCount`: Total number of vectors in currently loaded
graphs

These metrics will help us:
1. Understand if we are performing more comparisons than expected
2. Get insight into number of graphs queried
3. Get insight into the brute force vs graph query path 
4. Understand current memory utilization

The counters provide operations/second metrics, allowing calculation of
per-query averages by dividing by the number of queries. The memory
tracking metrics help monitor resource usage across graph segments as
they are loaded and unloaded.

Author: michaeljmarshall

src/java/org/apache/cassandra/index/sai/IndexContext.java

@@ -198,7 +198,8 @@ public IndexContext(@Nonnull String keyspace,
             this.hasEuclideanSimilarityFunc = vectorSimilarityFunction == VectorSimilarityFunction.EUCLIDEAN;
 
             this.indexMetrics = new IndexMetrics(this);
-            this.columnQueryMetrics = isLiteral() ? new ColumnQueryMetrics.TrieIndexMetrics(keyspace, table, getIndexName())
+            this.columnQueryMetrics = isVector() ? new ColumnQueryMetrics.VectorIndexMetrics(keyspace, table, getIndexName()) :
+                                      isLiteral() ? new ColumnQueryMetrics.TrieIndexMetrics(keyspace, table, getIndexName())
                                                   : new ColumnQueryMetrics.BKDIndexMetrics(keyspace, table, getIndexName());
 
         }

src/java/org/apache/cassandra/index/sai/QueryContext.java

@@ -61,7 +61,8 @@ public class QueryContext
 
     private final LongAdder queryTimeouts = new LongAdder();
 
-    private final LongAdder annNodesVisited = new LongAdder();
+    private final LongAdder annGraphSearchLatency = new LongAdder();
+
     private float annRerankFloor = 0.0f; // only called from single-threaded setup code
 
     private final LongAdder shadowedPrimaryKeyCount = new LongAdder();
@@ -139,9 +140,10 @@ public void addQueryTimeouts(long val)
     {
         queryTimeouts.add(val);
     }
-    public void addAnnNodesVisited(long val)
+
+    public void addAnnGraphSearchLatency(long val)
     {
-        annNodesVisited.add(val);
+        annGraphSearchLatency.add(val);
     }
 
     public void setFilterSortOrder(FilterSortOrder filterSortOrder)
@@ -202,9 +204,9 @@ public long queryTimeouts()
     {
         return queryTimeouts.longValue();
     }
-    public long annNodesVisited()
+    public long annGraphSearchLatency()
     {
-        return annNodesVisited.longValue();
+        return annGraphSearchLatency.longValue();
     }
 
     public FilterSortOrder filterSortOrder()

src/java/org/apache/cassandra/index/sai/disk/v2/V2OnDiskOrdinalsMap.java

@@ -98,6 +98,12 @@ public Structure getStructure()
         return canFastMapOrdinalsView ? Structure.ONE_TO_ONE : Structure.ZERO_OR_ONE_TO_MANY;
     }
 
+    @Override
+    public long cachedBytesUsed()
+    {
+        return 0;
+    }
+
     @Override
     public RowIdsView getRowIdsView()
     {

src/java/org/apache/cassandra/index/sai/disk/v2/V2VectorIndexSearcher.java

@@ -58,6 +58,7 @@
 import org.apache.cassandra.index.sai.disk.vector.VectorCompression;
 import org.apache.cassandra.index.sai.disk.vector.VectorMemtableIndex;
 import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
+import org.apache.cassandra.index.sai.metrics.ColumnQueryMetrics;
 import org.apache.cassandra.index.sai.plan.Expression;
 import org.apache.cassandra.index.sai.plan.Orderer;
 import org.apache.cassandra.index.sai.plan.Plan.CostCoefficients;
@@ -103,6 +104,7 @@ public class V2VectorIndexSearcher extends IndexSearcher
     private final PrimaryKey.Factory keyFactory;
     private final PairedSlidingWindowReservoir expectedActualNodesVisited = new PairedSlidingWindowReservoir(20);
     private final ThreadLocal<SparseBits> cachedBits;
+    private final ColumnQueryMetrics.VectorIndexMetrics columnQueryMetrics;
 
     protected V2VectorIndexSearcher(PrimaryKeyMap.Factory primaryKeyMapFactory,
                                     PerIndexFiles perIndexFiles,
@@ -113,7 +115,8 @@ protected V2VectorIndexSearcher(PrimaryKeyMap.Factory primaryKeyMapFactory,
         super(primaryKeyMapFactory, perIndexFiles, segmentMetadata, indexContext);
         this.graph = graph;
         this.keyFactory = PrimaryKey.factory(indexContext.comparator(), indexContext.indexFeatureSet());
-        cachedBits = ThreadLocal.withInitial(SparseBits::new);
+        this.cachedBits = ThreadLocal.withInitial(SparseBits::new);
+        this.columnQueryMetrics = (ColumnQueryMetrics.VectorIndexMetrics) indexContext.getColumnQueryMetrics();
     }
 
     @Override
@@ -194,10 +197,7 @@ private CloseableIterator<RowIdWithScore> searchInternal(AbstractBounds<Partitio
             if (RangeUtil.coversFullRing(keyRange))
             {
                 var estimate = estimateCost(rerankK, graph.size());
-                return graph.search(queryVector, limit, rerankK, threshold, Bits.ALL, context, visited -> {
-                    estimate.updateStatistics(visited);
-                    context.addAnnNodesVisited(visited);
-                });
+                return graph.search(queryVector, limit, rerankK, threshold, Bits.ALL, context, estimate::updateStatistics);
             }
 
             PrimaryKey firstPrimaryKey = keyFactory.createTokenOnly(keyRange.left.getToken());
@@ -214,7 +214,7 @@ private CloseableIterator<RowIdWithScore> searchInternal(AbstractBounds<Partitio
 
             // if the range covers the entire segment, skip directly to an index search
             if (minSSTableRowId <= metadata.minSSTableRowId && maxSSTableRowId >= metadata.maxSSTableRowId)
-                return graph.search(queryVector, limit, rerankK, threshold, Bits.ALL, context, context::addAnnNodesVisited);
+                return graph.search(queryVector, limit, rerankK, threshold, Bits.ALL, context, visited -> {});
 
             minSSTableRowId = Math.max(minSSTableRowId, metadata.minSSTableRowId);
             maxSSTableRowId = min(maxSSTableRowId, metadata.maxSSTableRowId);
@@ -263,10 +263,7 @@ private CloseableIterator<RowIdWithScore> searchInternal(AbstractBounds<Partitio
             // the trouble to add it.
             var betterCostEstimate = estimateCost(rerankK, cardinality);
 
-            return graph.search(queryVector, limit, rerankK, threshold, bits, context, visited -> {
-                betterCostEstimate.updateStatistics(visited);
-                context.addAnnNodesVisited(visited);
-            });
+            return graph.search(queryVector, limit, rerankK, threshold, bits, context, betterCostEstimate::updateStatistics);
         }
     }
 
@@ -305,8 +302,9 @@ private CloseableIterator<RowIdWithScore> orderByBruteForce(CompressedVectors cv
         segmentOrdinalPairs.forEachIndexOrdinalPair((i, ordinal) -> {
             approximateScores.push(i, scoreFunction.similarityTo(ordinal));
         });
+        columnQueryMetrics.onBruteForceNodesVisited(segmentOrdinalPairs.size());
         var reranker = new CloseableReranker(similarityFunction, queryVector, graph.getView());
-        return new BruteForceRowIdIterator(approximateScores, segmentOrdinalPairs, reranker, limit, rerankK);
+        return new BruteForceRowIdIterator(approximateScores, segmentOrdinalPairs, reranker, limit, rerankK, columnQueryMetrics);
     }
 
     /**
@@ -325,6 +323,7 @@ private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> query
             segmentOrdinalPairs.forEachSegmentRowIdOrdinalPair((segmentRowId, ordinal) -> {
                 scoredRowIds.push(segmentRowId, esf.similarityTo(ordinal));
             });
+            columnQueryMetrics.onBruteForceNodesReranked(segmentOrdinalPairs.size());
             return new NodeQueueRowIdIterator(scoredRowIds);
         }
     }
@@ -348,6 +347,7 @@ private CloseableIterator<RowIdWithScore> filterByBruteForce(VectorFloat<?> quer
                 if (score >= threshold)
                     results.add(new RowIdWithScore(segmentRowId, score));
             });
+            columnQueryMetrics.onBruteForceNodesReranked(segmentOrdinalPairs.size());
         }
         return CloseableIterator.wrap(results.iterator());
     }

src/java/org/apache/cassandra/index/sai/disk/v5/V5OnDiskOrdinalsMap.java

@@ -354,4 +354,26 @@ public void close() {
             // no-op
         }
     }
+
+    @Override
+    public long cachedBytesUsed()
+    {
+        if (structure != V5VectorPostingsWriter.Structure.ONE_TO_MANY) {
+            return 0;
+        }
+
+        long bytes = 0;
+        if (extraRowIds != null) {
+            bytes += extraRowIds.length * 4L;
+        }
+        if (extraOrdinals != null) {
+            bytes += extraOrdinals.length * 4L;
+        }
+        if (extraRowsByOrdinal != null) {
+            for (int[] rowIds : extraRowsByOrdinal.values()) {
+                bytes += rowIds.length * 4L;
+            }
+        }
+        return bytes;
+    }
 }

src/java/org/apache/cassandra/index/sai/disk/vector/AutoResumingNodeScoreIterator.java

@@ -24,10 +24,13 @@
 
 import io.github.jbellis.jvector.graph.GraphSearcher;
 import io.github.jbellis.jvector.graph.SearchResult;
+import org.apache.cassandra.index.sai.QueryContext;
+import org.apache.cassandra.index.sai.metrics.ColumnQueryMetrics;
 import org.apache.cassandra.tracing.Tracing;
 import org.apache.cassandra.utils.AbstractIterator;
 
 import static java.lang.Math.max;
+import static org.apache.cassandra.utils.Clock.Global.nanoTime;
 
 /**
  * An iterator over {@link SearchResult.NodeScore} backed by a {@link SearchResult} that resumes search
@@ -41,6 +44,8 @@ public class AutoResumingNodeScoreIterator extends AbstractIterator<SearchResult
     private final int rerankK;
     private final boolean inMemory;
     private final String source;
+    private final QueryContext context;
+    private final ColumnQueryMetrics.VectorIndexMetrics columnQueryMetrics;
     private final IntConsumer nodesVisitedConsumer;
     private Iterator<SearchResult.NodeScore> nodeScores;
     private int cumulativeNodesVisited;
@@ -51,6 +56,8 @@ public class AutoResumingNodeScoreIterator extends AbstractIterator<SearchResult
      * no more results.
      * @param searcher the {@link GraphSearcher} to use to resume search.
      * @param result the first {@link SearchResult} to iterate over
+     * @param context the {@link QueryContext} to use to record metrics
+     * @param columnQueryMetrics object to record metrics
      * @param nodesVisitedConsumer a consumer that accepts the total number of nodes visited
      * @param limit the limit to pass to the {@link GraphSearcher} when resuming search
      * @param rerankK the rerankK to pass to the {@link GraphSearcher} when resuming search
@@ -60,6 +67,8 @@ public class AutoResumingNodeScoreIterator extends AbstractIterator<SearchResult
     public AutoResumingNodeScoreIterator(GraphSearcher searcher,
                                          GraphSearcherAccessManager accessManager,
                                          SearchResult result,
+                                         QueryContext context,
+                                         ColumnQueryMetrics.VectorIndexMetrics columnQueryMetrics,
                                          IntConsumer nodesVisitedConsumer,
                                          int limit,
                                          int rerankK,
@@ -69,7 +78,9 @@ public AutoResumingNodeScoreIterator(GraphSearcher searcher,
         this.searcher = searcher;
         this.accessManager = accessManager;
         this.nodeScores = Arrays.stream(result.getNodes()).iterator();
-        this.cumulativeNodesVisited = result.getVisitedCount();
+        this.context = context;
+        this.columnQueryMetrics = columnQueryMetrics;
+        this.cumulativeNodesVisited = 0;
         this.nodesVisitedConsumer = nodesVisitedConsumer;
         this.limit = max(1, limit / 2); // we shouldn't need as many results on resume
         this.rerankK = rerankK;
@@ -83,21 +94,29 @@ protected SearchResult.NodeScore computeNext()
         if (nodeScores.hasNext())
             return nodeScores.next();
 
+        long start = nanoTime();
+
+        // Search deeper into the graph
         var nextResult = searcher.resume(limit, rerankK);
-        maybeLogTrace(nextResult);
+
+        // Record metrics
+        long elapsed = nanoTime() - start;
+        columnQueryMetrics.onSearchResult(nextResult, elapsed, true);
+        context.addAnnGraphSearchLatency(elapsed);
         cumulativeNodesVisited += nextResult.getVisitedCount();
+
+        if (Tracing.isTracing())
+        {
+            String msg = inMemory ? "Memory based ANN resume for {}/{} visited {} nodes, reranked {} to return {} results from {}"
+                                  : "Disk based ANN resume for {}/{} visited {} nodes, reranked {} to return {} results from {}";
+            Tracing.trace(msg, limit, rerankK, nextResult.getVisitedCount(), nextResult.getRerankedCount(), nextResult.getNodes().length, source);
+        }
+
         // If the next result is empty, we are done searching.
         nodeScores = Arrays.stream(nextResult.getNodes()).iterator();
         return nodeScores.hasNext() ? nodeScores.next() : endOfData();
     }
 
-    private void maybeLogTrace(SearchResult result)
-    {
-        String msg = inMemory ? "ANN resume for {}/{} visited {} nodes, reranked {} to return {} results from {}"
-                              : "DiskANN resume for {}/{} visited {} nodes, reranked {} to return {} results from {}";
-        Tracing.trace(msg, limit, rerankK, result.getVisitedCount(), result.getRerankedCount(), result.getNodes().length, source);
-    }
-
     @Override
     public void close()
     {

src/java/org/apache/cassandra/index/sai/disk/vector/BruteForceRowIdIterator.java

@@ -20,6 +20,7 @@
 
 import io.github.jbellis.jvector.graph.NodeQueue;
 import io.github.jbellis.jvector.util.BoundedLongHeap;
+import org.apache.cassandra.index.sai.metrics.ColumnQueryMetrics;
 import org.apache.cassandra.index.sai.utils.SegmentRowIdOrdinalPairs;
 import org.apache.cassandra.index.sai.utils.RowIdWithMeta;
 import org.apache.cassandra.index.sai.utils.RowIdWithScore;
@@ -62,6 +63,7 @@ public class BruteForceRowIdIterator extends AbstractIterator<RowIdWithScore>
     private final CloseableReranker reranker;
     private final int topK;
     private final int limit;
+    private final ColumnQueryMetrics.VectorIndexMetrics columnQueryMetrics;
     private int rerankedCount;
 
     /**
@@ -70,12 +72,14 @@ public class BruteForceRowIdIterator extends AbstractIterator<RowIdWithScore>
      * @param reranker A function that takes a graph ordinal and returns the exact similarity score
      * @param limit The query limit
      * @param topK The number of vectors to resolve and score before returning results
+     * @param columnQueryMetrics object to record metrics
      */
     public BruteForceRowIdIterator(NodeQueue approximateScoreQueue,
                                    SegmentRowIdOrdinalPairs segmentOrdinalPairs,
                                    CloseableReranker reranker,
                                    int limit,
-                                   int topK)
+                                   int topK,
+                                   ColumnQueryMetrics.VectorIndexMetrics columnQueryMetrics)
     {
         this.approximateScoreQueue = approximateScoreQueue;
         this.segmentOrdinalPairs = segmentOrdinalPairs;
@@ -84,21 +88,25 @@ public BruteForceRowIdIterator(NodeQueue approximateScoreQueue,
         assert topK >= limit : "topK must be greater than or equal to limit. Found: " + topK + " < " + limit;
         this.limit = limit;
         this.topK = topK;
+        this.columnQueryMetrics = columnQueryMetrics;
         this.rerankedCount = topK; // placeholder to kick off computeNext
     }
 
     @Override
     protected RowIdWithScore computeNext() {
         int consumed = rerankedCount - exactScoreQueue.size();
         if (consumed >= limit) {
+            int exactComparisons = 0;
             // Refill the exactScoreQueue until it reaches topK exact scores, or the approximate score queue is empty
             while (approximateScoreQueue.size() > 0 && exactScoreQueue.size() < topK) {
                 int segmentOrdinalIndex = approximateScoreQueue.pop();
                 int rowId = segmentOrdinalPairs.getSegmentRowId(segmentOrdinalIndex);
                 int ordinal = segmentOrdinalPairs.getOrdinal(segmentOrdinalIndex);
                 float score = reranker.similarityTo(ordinal);
+                exactComparisons++;
                 exactScoreQueue.push(rowId, score);
             }
+            columnQueryMetrics.onBruteForceNodesReranked(exactComparisons);
             rerankedCount = exactScoreQueue.size();
         }
         if (exactScoreQueue.size() == 0)