Skip to content

Reciprocal Rank Fusion (RRF) in TopDocs #13470

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jpountz merged 15 commits into from
Feb 27, 2025
Merged

Reciprocal Rank Fusion (RRF) in TopDocs #13470

Changes from 7 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
c6da395
add function rrf signature
harenlin May 18, 2024
9a7b57b
update TopDocs
harenlin Jun 7, 2024
57327aa
fix TopDocs error
harenlin Jun 7, 2024
ce55148
Float -> float in TopDocs.java RRF
harenlin Jun 7, 2024
4a6e10f
fix bug in rrf
harenlin Jun 7, 2024
ccb940f
Merge branch 'apache:main' into TopDocs-RRF
harenlin Jun 7, 2024
8e21bbc
style: gradlew tidy
hack4chang Jun 7, 2024
c018e50
feat: add javadocs on rrf method
hack4chang Jun 8, 2024
9579816
fix: gradlew format
hack4chang Jun 12, 2024
8d8f760
Merge branch 'main' into TopDocs-RRF
jpountz Feb 14, 2025
68b041c
Fixes + tests.
jpountz Feb 14, 2025
99c616c
Merge branch 'main' into TopDocs-RRF
jpountz Feb 26, 2025
3164006
Remove usage of break + label
jpountz Feb 26, 2025
67cc61f
CHANGES
jpountz Feb 26, 2025
2af86e3
Merge branch 'main' into TopDocs-RRF
jpountz Feb 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions lucene/core/src/java/org/apache/lucene/search/TopDocs.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,11 @@
*/
package org.apache.lucene.search;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.lucene.util.PriorityQueue;

/** Represents hits returned by {@link IndexSearcher#search(Query,int)}. */
Expand Down Expand Up @@ -350,4 +354,38 @@ private static TopDocs mergeAux(
return new TopFieldDocs(totalHits, hits, sort.getSort());
}
}

/** Reciprocal Rank Fusion method. */
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Users could use more details in these javadocs, e.g. what are k and topN?

Copy link
Contributor

@hack4chang hack4chang Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have to check if the k value is greater than or equal to 1 in the code? And maybe mention it in the javadocs?

Copy link
Contributor

@jpountz jpountz Jun 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 to validating parameters

public static TopDocs rrf(int TopN, int k, TopDocs[] hits) {
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: function arguments should use lower camel case

Suggested change
public static TopDocs rrf(int TopN, int k, TopDocs[] hits) {
public static TopDocs rrf(int topN, int k, TopDocs[] hits) {

Copy link
Contributor

@hack4chang hack4chang Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oops, my bad.

Map<Integer, Float> rrfScore = new HashMap<>();
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that we should identify documents not only by their doc ID, but by the combination of ScoreDoc.shardIndex and ScoreDoc.doc, as there could be multiple documents that have the same doc ID but come from different shards.

long minHits = Long.MAX_VALUE;
for (TopDocs topDoc : hits) {
minHits = Math.min(minHits, topDoc.totalHits.value);
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if it should be a max rather than a min? Presumably, hits from either top hits are considered as hits globally, and we are just using RRF to boost hits that are ranked high in multiple top hits?

Copy link
Contributor

@hack4chang hack4chang Jun 8, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The totalHits was a tricky part that we didn't know what value to assign to. IIUC, The totalHits means all the matched Document in a query, and we couldn't really calculate the union of the totalHits for all the TopDocs. So for this min totalHits, I just wanted to assign a min totalHits temporarily, to match the totalHits relation "greater than or equal to". And I want to ask for your opinion on this.

Copy link
Contributor

@jpountz jpountz Jun 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with using GREATER_THAN_OR_EQUAL_TO all the time, but I would still take the max, because a document is a match of it is a match to either query. For instance, imagine combining top hits from two queries where one query didn't match any docs, the total hit count should be the hit count of the other query, not zero?

Copy link
Contributor

@hack4chang hack4chang Jun 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make sense to me, thank you for the detailed explanation!!

Map<Integer, Float> scoreMap = new HashMap<>();
for (ScoreDoc scoreDoc : topDoc.scoreDocs) {
scoreMap.put(scoreDoc.doc, scoreDoc.score);
}

List<Map.Entry<Integer, Float>> scoreList = new ArrayList<>(scoreMap.entrySet());
scoreList.sort(Map.Entry.comparingByValue());
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't seem to be using these scoreMap and scoreList anywhere?

Copy link
Contributor Author

@harenlin harenlin Jun 9, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oops! My bad. I think we got something wrong right here. The for loop traversal for (ScoreDoc scoreDoc : topDoc.scoreDocs) is wrong, we should actually traverse the sorted results, i.e., scoreList, to add the ranking result to rrfScore.

int rank = 1;
for (Map.Entry<Integer, Float> entry : scoreList.entrySet()) {
    rrfScore.put(entry.getKey(), 
                 rrfScore.getOrDefault(entry.getKey(), 0.0f) + 1 / (rank + k));
    rank++;
}

P.S. For this part, however, I think we should determine the implementation of combining ScoreDoc.doc and ScoreDoc.shardIndex together first.


int rank = 1;
for (ScoreDoc scoreDoc : topDoc.scoreDocs) {
rrfScore.put(scoreDoc.doc, rrfScore.getOrDefault(scoreDoc.doc, 0.0f) + 1.0f / (rank + k));
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use Map#compute instead of getOrDefault + put?

Copy link
Contributor Author

@harenlin harenlin Jun 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good! I don't know Map#compute before tho. It should be like below:

int rank = 1;
for (Map.Entry<Integer, Float> entry : scoreList.entrySet()) {
    rrfScore.compute(entry.getKey(), (key, value) -> 
        (value == null ? 0.0f : value) + 1 / (float)(rank + k)
    );
    rank++;
}

rank++;
}
}

List<Map.Entry<Integer, Float>> rrfScoreRank = new ArrayList<>(rrfScore.entrySet());
rrfScoreRank.sort(
Map.Entry.<Integer, Float>comparingByValue().reversed()); // Sort in descending order

ScoreDoc[] rrfScoreDocs = new ScoreDoc[Math.min(TopN, rrfScoreRank.size())];
for (int i = 0; i < rrfScoreDocs.length; i++) {
rrfScoreDocs[i] = new ScoreDoc(rrfScoreRank.get(i).getKey(), rrfScoreRank.get(i).getValue());
Copy link
Contributor

@jpountz jpountz Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: we should preserve the original shardIndex that is configured on the ScoreDoc object that identifies the shard that it is coming from.

Copy link
Contributor

@hack4chang hack4chang Jun 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So this was also a tricky part for us. For my understanding, the RRF would combine search result based on the different ranks of a documents in different results. We supposed to combine the ranks for all individual doucments, but a document come from different shards should be treated as different documents?

Copy link
Contributor

@jpountz jpountz Jun 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is correct. When working with shards, hits should first be combined on a per-query basis using TopDocs#merge, which will set the shardIndex field. And then global top hits can be merged across queries with RRF.

Copy link
Contributor

@hack4chang hack4chang Jun 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we do the rrf by setting the unique key as docid and shardIndex, what would be the difference between TopDocs#rrf and TopDocs#merge? I think giving an example could express better. Suppose that we have two Shards, and we want to retrieve top 3 results from each shards and do rrf on top of them. There's three documents A, B and C. In Shard1, the top 3 is A -> B -> C. In Shard2, it's B -> C -> A. The original rrf method would calculate the rank by aggregating the docid, assume the constant k is 1. Top 3 results would be B (1/(k+2) + 1/(k+1)) - A (1/(k+1) + 1/(k+3)) - C (1/(k+3) + 1/(k+2)).
If we are going to consider the shardIndex as a unique key as well, how should the rrf rank to be presented.

Copy link
Contributor

@jpountz jpountz Jun 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When working with shards, shards managed non-overlapping subsets of the data, so you could not have documents A, B and C in both shards.

}

return new TopDocs(
new TotalHits(minHits, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO), rrfScoreDocs);
}
}