remove reranking from orderByBruteForce

datastax · Dec 30, 2024 · 83263e4 · 83263e4
1 parent b7f6789
commit 83263e4
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/src/java/org/apache/cassandra/index/sai/disk/v2/V2VectorIndexSearcher.java b/src/java/org/apache/cassandra/index/sai/disk/v2/V2VectorIndexSearcher.java
@@ -21,7 +21,9 @@
 import java.lang.invoke.MethodHandles;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.function.Consumer;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -291,7 +293,7 @@ private CloseableIterator<RowIdWithScore> orderByBruteForce(CompressedVectors cv
                                                                 VectorFloat<?> queryVector,
                                                                 IntIntPairArray segmentOrdinalPairs,
                                                                 int limit,
-                                                                int rerankK) throws IOException
+                                                                int rerankK)
     {
         var approximateScores = new SortingIterator.Builder<BruteForceRowIdIterator.RowWithApproximateScore>(segmentOrdinalPairs.size());
         var similarityFunction = indexContext.getIndexWriterConfig().getSimilarityFunction();
@@ -302,8 +304,27 @@ private CloseableIterator<RowIdWithScore> orderByBruteForce(CompressedVectors cv
             approximateScores.add(new BruteForceRowIdIterator.RowWithApproximateScore(segmentRowId, ordinal, score));
         });
         var approximateScoresQueue = approximateScores.build(BruteForceRowIdIterator.RowWithApproximateScore::compare);
-        var reranker = new CloseableReranker(similarityFunction, queryVector, graph.getView());
-        return new BruteForceRowIdIterator(approximateScoresQueue, reranker, limit, rerankK);
+        var transformed = new Iterator<RowIdWithScore>() {
+            int consumed = 0;
+
+            @Override
+            public boolean hasNext()
+            {
+                if (consumed++ >= limit)
+                    return false;
+                return approximateScoresQueue.hasNext();
+            }
+
+            @Override
+            public RowIdWithScore next()
+            {
+                if (!hasNext())
+                    throw new NoSuchElementException();
+                var approximated = approximateScoresQueue.next();
+                return new RowIdWithScore(approximated.rowId, approximated.appoximateScore);
+            }
+        };
+        return CloseableIterator.wrap(transformed);
     }
 
     /**

diff --git a/src/java/org/apache/cassandra/index/sai/disk/vector/BruteForceRowIdIterator.java b/src/java/org/apache/cassandra/index/sai/disk/vector/BruteForceRowIdIterator.java
@@ -51,9 +51,9 @@ public class BruteForceRowIdIterator extends AbstractIterator<RowIdWithScore>
 {
     public static class RowWithApproximateScore
     {
-        private final int rowId;
-        private final int ordinal;
-        private final float appoximateScore;
+        public final int rowId;
+        public final int ordinal;
+        public final float appoximateScore;
 
         public RowWithApproximateScore(int rowId, int ordinal, float appoximateScore)
         {