CNDB-13417: Reduce object alloc. on brute force hybrid ann path (#1643)

### What is the issue
Relates to riptano/cndb#13417 (the issue has
multiple PRs)

### What does this PR fix and why was it fixed

In hybrid ANN search resulting in brute force row sorting, we see many
object allocations per row materialized (so O(n) space complexity), and
this creates memory pressure leading to reduced performance. This PR
reduces object allocation by using jvector's `NodeQueue` data structure
that encodes an int and a float into a long and then sorts based on the
float. I also renamed the IntIntPairArray because I needed more
specialized methods.

- **Rename IntIntPairArray to SegmentRowIdOrdinalPairs**
- **CNDB-13417: Reduce object alloc. on brute force hybrid ann path**

Author: michaeljmarshall

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

@@ -22,18 +22,18 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.function.Consumer;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
-import com.google.common.util.concurrent.Runnables;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import io.github.jbellis.jvector.graph.NodeQueue;
 import io.github.jbellis.jvector.quantization.CompressedVectors;
 import io.github.jbellis.jvector.quantization.ProductQuantization;
 import io.github.jbellis.jvector.util.BitSet;
 import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.util.BoundedLongHeap;
 import io.github.jbellis.jvector.util.SparseBits;
 import io.github.jbellis.jvector.vector.VectorizationProvider;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
@@ -54,13 +54,14 @@
 import org.apache.cassandra.index.sai.disk.vector.BruteForceRowIdIterator;
 import org.apache.cassandra.index.sai.disk.vector.CassandraDiskAnn;
 import org.apache.cassandra.index.sai.disk.vector.CloseableReranker;
+import org.apache.cassandra.index.sai.disk.vector.NodeQueueRowIdIterator;
 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.plan.Expression;
 import org.apache.cassandra.index.sai.plan.Orderer;
 import org.apache.cassandra.index.sai.plan.Plan.CostCoefficients;
-import org.apache.cassandra.index.sai.utils.IntIntPairArray;
+import org.apache.cassandra.index.sai.utils.SegmentRowIdOrdinalPairs;
 import org.apache.cassandra.index.sai.utils.PrimaryKey;
 import org.apache.cassandra.index.sai.utils.PrimaryKeyWithSortKey;
 import org.apache.cassandra.index.sai.utils.RangeUtil;
@@ -72,7 +73,6 @@
 import org.apache.cassandra.metrics.QuickSlidingWindowReservoir;
 import org.apache.cassandra.tracing.Tracing;
 import org.apache.cassandra.utils.CloseableIterator;
-import org.apache.cassandra.utils.SortingIterator;
 
 import static java.lang.Math.ceil;
 import static java.lang.Math.min;
@@ -230,7 +230,7 @@ private CloseableIterator<RowIdWithScore> searchInternal(AbstractBounds<Partitio
             if (initialCostEstimate.shouldUseBruteForce())
             {
                 var maxSize = endSegmentRowId - startSegmentRowId + 1;
-                var segmentOrdinalPairs = new IntIntPairArray(maxSize);
+                var segmentOrdinalPairs = new SegmentRowIdOrdinalPairs(maxSize);
                 try (var ordinalsView = graph.getOrdinalsView())
                 {
                     ordinalsView.forEachOrdinalInRange(startSegmentRowId, endSegmentRowId, segmentOrdinalPairs::add);
@@ -270,7 +270,7 @@ private CloseableIterator<RowIdWithScore> searchInternal(AbstractBounds<Partitio
         }
     }
 
-    private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> queryVector, IntIntPairArray segmentOrdinalPairs, int limit, int rerankK) throws IOException
+    private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> queryVector, SegmentRowIdOrdinalPairs segmentOrdinalPairs, int limit, int rerankK) throws IOException
     {
         // If we use compressed vectors, we still have to order rerankK results using full resolution similarity
         // scores, so only use the compressed vectors when there are enough vectors to make it worthwhile.
@@ -289,33 +289,44 @@ private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> query
      */
     private CloseableIterator<RowIdWithScore> orderByBruteForce(CompressedVectors cv,
                                                                 VectorFloat<?> queryVector,
-                                                                IntIntPairArray segmentOrdinalPairs,
+                                                                SegmentRowIdOrdinalPairs segmentOrdinalPairs,
                                                                 int limit,
                                                                 int rerankK) throws IOException
     {
-        var approximateScores = new SortingIterator.Builder<BruteForceRowIdIterator.RowWithApproximateScore>(segmentOrdinalPairs.size());
+        // Use the jvector NodeQueue to avoid unnecessary object allocations since this part of the code operates on
+        // many rows.
+        var approximateScores = new NodeQueue(new BoundedLongHeap(segmentOrdinalPairs.size()), NodeQueue.Order.MAX_HEAP);
         var similarityFunction = indexContext.getIndexWriterConfig().getSimilarityFunction();
         var scoreFunction = cv.precomputedScoreFunctionFor(queryVector, similarityFunction);
 
-        segmentOrdinalPairs.forEachIntPair((segmentRowId, ordinal) -> {
-            var score = scoreFunction.similarityTo(ordinal);
-            approximateScores.add(new BruteForceRowIdIterator.RowWithApproximateScore(segmentRowId, ordinal, score));
+        // Store the index of the (rowId, ordinal) pair from the segmentOrdinalPairs in the NodeQueue so that we can
+        // retrieve both values with O(1) lookup when we need to resolve the full resolution score in the
+        // BruteForceRowIdIterator.
+        segmentOrdinalPairs.forEachIndexOrdinalPair((i, ordinal) -> {
+            approximateScores.push(i, scoreFunction.similarityTo(ordinal));
         });
-        var approximateScoresQueue = approximateScores.build(BruteForceRowIdIterator.RowWithApproximateScore::compare);
         var reranker = new CloseableReranker(similarityFunction, queryVector, graph.getView());
-        return new BruteForceRowIdIterator(approximateScoresQueue, reranker, limit, rerankK);
+        return new BruteForceRowIdIterator(approximateScores, segmentOrdinalPairs, reranker, limit, rerankK);
     }
 
     /**
      * Produces a correct ranking of the rows in the given segment. Because this graph does not have compressed
      * vectors, read all vectors and put them into a priority queue to rank them lazily. It is assumed that the whole
      * PQ will often not be needed.
      */
-    private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> queryVector, IntIntPairArray segmentOrdinalPairs) throws IOException
+    private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> queryVector, SegmentRowIdOrdinalPairs segmentOrdinalPairs) throws IOException
     {
-        var scoredRowIds = new SortingIterator.Builder<RowIdWithScore>(segmentOrdinalPairs.size());
-        addScoredRowIdsToCollector(queryVector, segmentOrdinalPairs, 0, scoredRowIds::add);
-        return scoredRowIds.closeable(RowIdWithScore::compare, Runnables.doNothing());
+        var scoredRowIds = new NodeQueue(new BoundedLongHeap(segmentOrdinalPairs.size()), NodeQueue.Order.MAX_HEAP);
+        try (var vectorsView = graph.getView())
+        {
+            var similarityFunction = indexContext.getIndexWriterConfig().getSimilarityFunction();
+            var esf = vectorsView.rerankerFor(queryVector, similarityFunction);
+            // Because the scores are exact, we only store the rowid, score pair.
+            segmentOrdinalPairs.forEachSegmentRowIdOrdinalPair((segmentRowId, ordinal) -> {
+                scoredRowIds.push(segmentRowId, esf.similarityTo(ordinal));
+            });
+            return new NodeQueueRowIdIterator(scoredRowIds);
+        }
     }
 
     /**
@@ -324,29 +335,21 @@ private CloseableIterator<RowIdWithScore> orderByBruteForce(VectorFloat<?> query
      * NOTE: because the threshold is not used for ordering, the result is returned in PK order, not score order.
      */
     private CloseableIterator<RowIdWithScore> filterByBruteForce(VectorFloat<?> queryVector,
-                                                                 IntIntPairArray segmentOrdinalPairs,
+                                                                 SegmentRowIdOrdinalPairs segmentOrdinalPairs,
                                                                  float threshold) throws IOException
     {
         var results = new ArrayList<RowIdWithScore>(segmentOrdinalPairs.size());
-        addScoredRowIdsToCollector(queryVector, segmentOrdinalPairs, threshold, results::add);
-        return CloseableIterator.wrap(results.iterator());
-    }
-
-    private void addScoredRowIdsToCollector(VectorFloat<?> queryVector,
-                                            IntIntPairArray segmentOrdinalPairs,
-                                            float threshold,
-                                            Consumer<RowIdWithScore> collector) throws IOException
-    {
-        var similarityFunction = indexContext.getIndexWriterConfig().getSimilarityFunction();
         try (var vectorsView = graph.getView())
         {
+            var similarityFunction = indexContext.getIndexWriterConfig().getSimilarityFunction();
             var esf = vectorsView.rerankerFor(queryVector, similarityFunction);
-            segmentOrdinalPairs.forEachIntPair((segmentRowId, ordinal) -> {
+            segmentOrdinalPairs.forEachSegmentRowIdOrdinalPair((segmentRowId, ordinal) -> {
                 var score = esf.similarityTo(ordinal);
                 if (score >= threshold)
-                    collector.accept(new RowIdWithScore(segmentRowId, score));
+                    results.add(new RowIdWithScore(segmentRowId, score));
             });
         }
+        return CloseableIterator.wrap(results.iterator());
     }
 
     private long getMaxSSTableRowId(PrimaryKeyMap primaryKeyMap, PartitionPosition right)
@@ -489,7 +492,7 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderResultsBy(SSTableReader rea
         }
         // Create bits from the mapping
         var bits = bitSetForSearch();
-        segmentOrdinalPairs.forEachRightInt(bits::set);
+        segmentOrdinalPairs.forEachOrdinal(bits::set);
         // else ask the index to perform a search limited to the bits we created
         var queryVector = vts.createFloatVector(orderer.getVectorTerm());
         var results = graph.search(queryVector, limit, rerankK, 0, bits, context, cost::updateStatistics);
@@ -504,9 +507,9 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderResultsBy(SSTableReader rea
      * @return a mapping of segment row id to ordinal
      * @throws IOException
      */
-    private IntIntPairArray flatmapPrimaryKeysToBitsAndRows(List<PrimaryKey> keysInRange) throws IOException
+    private SegmentRowIdOrdinalPairs flatmapPrimaryKeysToBitsAndRows(List<PrimaryKey> keysInRange) throws IOException
     {
-        var segmentOrdinalPairs = new IntIntPairArray(keysInRange.size());
+        var segmentOrdinalPairs = new SegmentRowIdOrdinalPairs(keysInRange.size());
         int lastSegmentRowId = -1;
         try (var primaryKeyMap = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap();
              var ordinalsView = graph.getOrdinalsView())

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

@@ -18,11 +18,13 @@
 
 package org.apache.cassandra.index.sai.disk.vector;
 
+import io.github.jbellis.jvector.graph.NodeQueue;
+import io.github.jbellis.jvector.util.BoundedLongHeap;
+import org.apache.cassandra.index.sai.utils.SegmentRowIdOrdinalPairs;
 import org.apache.cassandra.index.sai.utils.RowIdWithMeta;
 import org.apache.cassandra.index.sai.utils.RowIdWithScore;
 import org.apache.cassandra.io.util.FileUtils;
 import org.apache.cassandra.utils.AbstractIterator;
-import org.apache.cassandra.utils.LucenePriorityQueue;
 import org.apache.cassandra.utils.SortingIterator;
 
 
@@ -45,55 +47,39 @@
  * is consumed. We do this because we expect that most often the first limit-many will pass the final verification
  * and only query more if some didn't (e.g. because the vector was deleted in a newer sstable).
  * <p>
- * As an implementation detail, we use a PriorityQueue to maintain state rather than a List and sorting.
+ * As an implementation detail, we use a heap to maintain state rather than a List and sorting.
  */
 public class BruteForceRowIdIterator extends AbstractIterator<RowIdWithScore>
 {
-    public static class RowWithApproximateScore
-    {
-        private final int rowId;
-        private final int ordinal;
-        private final float appoximateScore;
-
-        public RowWithApproximateScore(int rowId, int ordinal, float appoximateScore)
-        {
-            this.rowId = rowId;
-            this.ordinal = ordinal;
-            this.appoximateScore = appoximateScore;
-        }
-
-        public static int compare(RowWithApproximateScore l, RowWithApproximateScore r)
-        {
-            // Inverted comparison to sort in descending order
-            return Float.compare(r.appoximateScore, l.appoximateScore);
-        }
-    }
-
-    // We use two binary heaps (a SortingIterator and LucenePriorityQueue) because we do not need an eager ordering of
+    // We use two binary heaps (NodeQueue) because we do not need an eager ordering of
     // these results. Depending on how many sstables the query hits and the relative scores of vectors from those
     // sstables, we may not need to return more than the first handful of scores.
-    // Priority queue with compressed vector scores
-    private final SortingIterator<RowWithApproximateScore> approximateScoreQueue;
-    // Priority queue with full resolution scores
-    private final LucenePriorityQueue<RowIdWithScore> exactScoreQueue;
+    // Heap with compressed vector scores
+    private final NodeQueue approximateScoreQueue;
+    private final SegmentRowIdOrdinalPairs segmentOrdinalPairs;
+    // Use the jvector NodeQueue to avoid unnecessary object allocations
+    private final NodeQueue exactScoreQueue;
     private final CloseableReranker reranker;
     private final int topK;
     private final int limit;
     private int rerankedCount;
 
     /**
-     * @param approximateScoreQueue A priority queue of rows and their ordinal ordered by their approximate similarity scores
+     * @param approximateScoreQueue A heap of indexes ordered by their approximate similarity scores
+     * @param segmentOrdinalPairs A mapping from the index in the approximateScoreQueue to the node's rowId and ordinal
      * @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
      */
-    public BruteForceRowIdIterator(SortingIterator<RowWithApproximateScore> approximateScoreQueue,
+    public BruteForceRowIdIterator(NodeQueue approximateScoreQueue,
+                                   SegmentRowIdOrdinalPairs segmentOrdinalPairs,
                                    CloseableReranker reranker,
                                    int limit,
                                    int topK)
     {
         this.approximateScoreQueue = approximateScoreQueue;
-        this.exactScoreQueue = new LucenePriorityQueue<>(topK, RowIdWithScore::compare);
+        this.segmentOrdinalPairs = segmentOrdinalPairs;
+        this.exactScoreQueue = new NodeQueue(new BoundedLongHeap(topK), NodeQueue.Order.MAX_HEAP);
         this.reranker = reranker;
         assert topK >= limit : "topK must be greater than or equal to limit. Found: " + topK + " < " + limit;
         this.limit = limit;
@@ -106,15 +92,21 @@ protected RowIdWithScore computeNext() {
         int consumed = rerankedCount - exactScoreQueue.size();
         if (consumed >= limit) {
             // Refill the exactScoreQueue until it reaches topK exact scores, or the approximate score queue is empty
-            while (approximateScoreQueue.hasNext() && exactScoreQueue.size() < topK) {
-                RowWithApproximateScore rowOrdinalScore = approximateScoreQueue.next();
-                float score = reranker.similarityTo(rowOrdinalScore.ordinal);
-                exactScoreQueue.add(new RowIdWithScore(rowOrdinalScore.rowId, score));
+            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);
+                exactScoreQueue.push(rowId, score);
             }
             rerankedCount = exactScoreQueue.size();
         }
-        RowIdWithScore top = exactScoreQueue.pop();
-        return top == null ? endOfData() : top;
+        if (exactScoreQueue.size() == 0)
+            return endOfData();
+
+        float score = exactScoreQueue.topScore();
+        int rowId = exactScoreQueue.pop();
+        return new RowIdWithScore(rowId, score);
     }
 
     @Override

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.index.sai.disk.vector;
+
+import io.github.jbellis.jvector.graph.NodeQueue;
+import org.apache.cassandra.index.sai.utils.RowIdWithScore;
+import org.apache.cassandra.utils.AbstractIterator;
+
+/**
+ * An iterator over {@link RowIdWithScore} that lazily consumes a {@link NodeQueue}.
+ */
+public class NodeQueueRowIdIterator extends AbstractIterator<RowIdWithScore>
+{
+    private final NodeQueue scoreQueue;
+
+    public NodeQueueRowIdIterator(NodeQueue scoreQueue)
+    {
+        this.scoreQueue = scoreQueue;
+    }
+
+    @Override
+    protected RowIdWithScore computeNext()
+    {
+        if (scoreQueue.size() == 0)
+            return endOfData();
+        float score = scoreQueue.topScore();
+        int rowId = scoreQueue.pop();
+        return new RowIdWithScore(rowId, score);
+    }
+}