Top-K & Pagination
Your BM25 scoring is perfect. Your inverted index is blazing fast.
But what if you have 2 million documents in your index and a user searches for "software" which appears in 1 million documents.
Your current code does this:
- Finds 1,000,000 matching Document IDs.
- Calculates the BM25 score for all 1,000,000 documents.
- Pushes all 1,000,000 documents into a massive Array in RAM.
- Sorts the entire 1,000,000-element Array
O(N log N). - Slices the first
limit=10elements and returns them.
The problem is steps 3 and 4. Storing and sorting 1 million documents in RAM for every single concurrent search request will instantly OOM (Out-of-Memory) crash your server.
We only need the top k results. We don’t care about sorting the 999,990 losers.
The Solution: A Min-Heap
Instead of sorting the entire array at the end, we can maintain a Priority Queue (specifically, a Min-Heap) of size k (where k is the limit parameter).
As you iterate through the matching documents and calculate their scores:
- If the heap has fewer than
kelements, just push the document in. - If the heap is full (size
k), compare the current document’s score to the smallest score in the heap (which sits at the root of a Min-Heap,O(1)). - If the current document’s score is larger than the root, pop the root (
O(log k)) and push the new document (O(log k)). - If it’s smaller, just ignore the document entirely.
This reduces our memory footprint from O(N) to O(k), and our time complexity drops significantly!
Your Task
Update your GET /search endpoint. Stop collecting all results into a massive array.
Initialize a Priority Queue / Min-Heap sized to the limit={k} query parameter. Feed the documents into the heap as you score them. Once you’ve evaluated all matches, the heap will contain exactly the top k highest-scoring documents.
Pop them out, sort them descending (since the heap pops smallest first), and return them.
Note
Some languages call this a
PriorityQueue,MinHeap, orBinaryHeap. Make sure you configure it so that the element with the lowest score is at the top (popped first).
Important
Run the below docker command to test your solution.
docker run \ --rm \ --add-host host.docker.internal:host-gateway \ codeberg.org/level0/buildit/search-engine:latest \ --addr host.docker.internal:8080 \ --until top_k