Use bounded top-k selection instead of full result sorting

Author: rudrankriyam

Sources/VecturaKit/SearchEngine/BM25Index.swift

@@ -205,6 +205,10 @@ public actor BM25Index {
   ///   - topK: Maximum number of results to return
   /// - Returns: Array of tuples containing lightweight documents and their BM25 scores
   public func search(query: String, topK: Int = 10) -> [(document: BM25Document, score: Float)] {
+    guard topK > 0 else {
+      return []
+    }
+
     let queryTerms = tokenize(query)
     guard !queryTerms.isEmpty else {
       return []
@@ -223,8 +227,10 @@ public actor BM25Index {
       queryIDFs[term] = log(max(idfArgument, 1e-9))
     }
 
-    var scores: [(BM25Document, Float)] = []
-    scores.reserveCapacity(documents.count)
+    var topResults = TopKSelector<(BM25Document, Float)>(
+      maxCount: topK,
+      isHigherRanked: { $0.1 > $1.1 }
+    )
 
     for document in documents.values {
       let docLength = Float(documentLengths[document.id] ?? 0)
@@ -246,15 +252,12 @@ public actor BM25Index {
         score += Float(queryTermCount) * idf * (numerator / denominator)
       }
 
-      scores.append((document, score))
+      if score > 0 {
+        topResults.insert((document, score))
+      }
     }
 
-    return Array(
-      scores
-        .sorted { $0.1 > $1.1 }
-        .filter { $0.1 > 0 }
-        .prefix(topK)
-    )
+    return topResults.sortedElements()
   }
 
   /// Add a new document to the index incrementally

Sources/VecturaKit/SearchEngine/HybridSearchEngine.swift

@@ -220,10 +220,17 @@ public struct HybridSearchEngine: VecturaSearchEngine {
       )
     }
 
-    // Sort and return top K
-    return combinedResults.values
-      .sorted(by: { $0.score > $1.score })
-      .prefix(topK)
-      .map { $0 }
+    guard topK > 0 else {
+      return []
+    }
+
+    var topResults = TopKSelector<VecturaSearchResult>(
+      maxCount: topK,
+      isHigherRanked: { $0.score > $1.score }
+    )
+    for result in combinedResults.values {
+      topResults.insert(result)
+    }
+    return topResults.sortedElements()
   }
 }

Sources/VecturaKit/SearchEngine/VectorSearchEngine.swift

@@ -114,9 +114,10 @@ public struct VectorSearchEngine: VecturaSearchEngine {
       1
     )
 
-    // Build results
-    var results = [VecturaSearchResult]()
-    results.reserveCapacity(docsCount)
+    var topResults = TopKSelector<VecturaSearchResult>(
+      maxCount: options.numResults,
+      isHigherRanked: { $0.score > $1.score }
+    )
 
     for (i, similarity) in similarities.enumerated() {
       if let threshold = options.threshold, similarity < threshold {
@@ -125,7 +126,7 @@ public struct VectorSearchEngine: VecturaSearchEngine {
 
       // docIds and documents are built in parallel, so indices correspond
       let doc = documents[i]
-      results.append(
+      topResults.insert(
         VecturaSearchResult(
           id: doc.id,
           text: doc.text,
@@ -135,11 +136,7 @@ public struct VectorSearchEngine: VecturaSearchEngine {
       )
     }
 
-    results.sort { $0.score > $1.score }
-    if results.count > options.numResults {
-      results.removeSubrange(options.numResults..<results.count)
-    }
-    return results
+    return topResults.sortedElements()
   }
 
   private func searchWithIndexedStorage(
@@ -223,17 +220,18 @@ public struct VectorSearchEngine: VecturaSearchEngine {
       1
     )
 
-    // Build results
-    var results = [VecturaSearchResult]()
-    results.reserveCapacity(candidatesCount)
+    var topResults = TopKSelector<VecturaSearchResult>(
+      maxCount: options.numResults,
+      isHigherRanked: { $0.score > $1.score }
+    )
 
     for (i, similarity) in similarities.enumerated() {
       if let threshold = options.threshold, similarity < threshold {
         continue
       }
 
       let doc = candidateDocs[i]
-      results.append(
+      topResults.insert(
         VecturaSearchResult(
           id: doc.id,
           text: doc.text,
@@ -243,11 +241,7 @@ public struct VectorSearchEngine: VecturaSearchEngine {
       )
     }
 
-    results.sort { $0.score > $1.score }
-    if results.count > options.numResults {
-      results.removeSubrange(options.numResults..<results.count)
-    }
-    return results
+    return topResults.sortedElements()
   }
 
   // MARK: - Helper Methods

Sources/VecturaKit/Utilities/TopKSelector.swift

@@ -0,0 +1,85 @@
+import Foundation
+
+/// Keeps only the top K elements using a bounded min-heap.
+///
+/// The heap root stores the lowest-ranked element currently retained, so inserts
+/// are `O(log k)` and the full pass is `O(n log k)`.
+struct TopKSelector<Element> {
+  private let maxCount: Int
+  private let isHigherRanked: (Element, Element) -> Bool
+  private var heap: [Element] = []
+
+  init(
+    maxCount: Int,
+    isHigherRanked: @escaping (Element, Element) -> Bool
+  ) {
+    precondition(maxCount > 0, "maxCount must be greater than zero")
+    self.maxCount = maxCount
+    self.isHigherRanked = isHigherRanked
+    self.heap.reserveCapacity(maxCount)
+  }
+
+  mutating func insert(_ element: Element) {
+    if heap.count < maxCount {
+      heap.append(element)
+      siftUp(from: heap.count - 1)
+      return
+    }
+
+    guard let currentLowest = heap.first, isHigherRanked(element, currentLowest) else {
+      return
+    }
+
+    heap[0] = element
+    siftDown(from: 0)
+  }
+
+  func sortedElements() -> [Element] {
+    heap.sorted(by: isHigherRanked)
+  }
+
+  private func isLowerRanked(_ lhs: Element, than rhs: Element) -> Bool {
+    isHigherRanked(rhs, lhs)
+  }
+
+  private mutating func siftUp(from index: Int) {
+    var childIndex = index
+
+    while childIndex > 0 {
+      let parentIndex = (childIndex - 1) / 2
+      guard isLowerRanked(heap[childIndex], than: heap[parentIndex]) else {
+        break
+      }
+
+      heap.swapAt(childIndex, parentIndex)
+      childIndex = parentIndex
+    }
+  }
+
+  private mutating func siftDown(from index: Int) {
+    var parentIndex = index
+
+    while true {
+      let leftChildIndex = 2 * parentIndex + 1
+      let rightChildIndex = leftChildIndex + 1
+      var candidateIndex = parentIndex
+
+      if leftChildIndex < heap.count,
+        isLowerRanked(heap[leftChildIndex], than: heap[candidateIndex]) {
+        candidateIndex = leftChildIndex
+      }
+
+      if rightChildIndex < heap.count,
+        isLowerRanked(heap[rightChildIndex], than: heap[candidateIndex]) {
+        candidateIndex = rightChildIndex
+      }
+
+      guard candidateIndex != parentIndex else {
+        return
+      }
+
+      heap.swapAt(parentIndex, candidateIndex)
+      parentIndex = candidateIndex
+    }
+  }
+}