diff --git a/python/python/lancedb/rerankers/mrr.py b/python/python/lancedb/rerankers/mrr.py index 81a03c9ca..38f742bb8 100644 --- a/python/python/lancedb/rerankers/mrr.py +++ b/python/python/lancedb/rerankers/mrr.py @@ -97,7 +97,12 @@ class MRRReranker(Reranker): weighted_mrr = self.weight_vector * vector_rr + self.weight_fts * fts_rr final_mrr_scores[result_id] = weighted_mrr - combined_results = self.merge_results(vector_results, fts_results) + # With return_score="all", a row found by both searches must keep its + # `_distance` and its `_score`. + if self.score == "all": + combined_results = self._merge_and_keep_scores(vector_results, fts_results) + else: + combined_results = self.merge_results(vector_results, fts_results) combined_row_ids = combined_results["_rowid"].to_pylist() relevance_scores = [final_mrr_scores[row_id] for row_id in combined_row_ids] combined_results = combined_results.append_column( diff --git a/python/python/lancedb/rerankers/rrf.py b/python/python/lancedb/rerankers/rrf.py index 9c835a3ea..c9b54a5a7 100644 --- a/python/python/lancedb/rerankers/rrf.py +++ b/python/python/lancedb/rerankers/rrf.py @@ -54,8 +54,12 @@ class RRFReranker(Reranker): for i, result_id in enumerate(ids, 1): rrf_score_map[result_id] += 1 / (i + self.K) - # Sort the results based on RRF score - combined_results = self.merge_results(vector_results, fts_results) + # Sort the results based on RRF score. With return_score="all", a row found + # by both searches must keep its `_distance` and its `_score`. + if self.score == "all": + combined_results = self._merge_and_keep_scores(vector_results, fts_results) + else: + combined_results = self.merge_results(vector_results, fts_results) combined_row_ids = combined_results["_rowid"].to_pylist() relevance_scores = [rrf_score_map[row_id] for row_id in combined_row_ids] combined_results = combined_results.append_column( diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index dd6cb4ab5..78b64613e 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -428,6 +428,38 @@ def test_rrf_multivector_ranks_each_result_list(): assert result["_rowid"].to_pylist() == [2, 1, 3] +@pytest.mark.parametrize("reranker_cls", [RRFReranker, MRRReranker]) +def test_rank_fusion_return_score_all_keeps_both_scores(reranker_cls): + # A row found by both searches must keep its vector `_distance` and its FTS + # `_score` when the reranker is asked to return all scores. + vector_results = pa.table( + { + "_rowid": pa.array([1, 2], type=pa.uint64()), + "text": ["a", "b"], + "_distance": pa.array([0.1, 0.2], type=pa.float32()), + } + ) + fts_results = pa.table( + { + "_rowid": pa.array([2, 3], type=pa.uint64()), + "text": ["b", "c"], + "_score": pa.array([5.0, 4.0], type=pa.float32()), + } + ) + + result = reranker_cls(return_score="all").rerank_hybrid( + "query", vector_results, fts_results + ) + rows = {row["_rowid"]: row for row in result.to_pylist()} + + assert rows[1]["_distance"] == pytest.approx(0.1) + assert rows[1]["_score"] is None + assert rows[2]["_distance"] == pytest.approx(0.2) + assert rows[2]["_score"] == pytest.approx(5.0) + assert rows[3]["_distance"] is None + assert rows[3]["_score"] == pytest.approx(4.0) + + def test_rrf_reranker_distance(): data = pa.table( {