diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index fe8c33861..9d7610f72 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2575,10 +2575,6 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): with_row_ids: bool, offset: Optional[int] = None, ) -> pa.Table: - if norm == "rank": - vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance") - fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score") - # If both result sets are empty (e.g. after hard filtering), # return early to avoid errors in reranking or score restoration. if vector_results.num_rows == 0 and fts_results.num_rows == 0: @@ -2605,15 +2601,23 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): original_score_row_ids = None # normalize the scores to be between 0 and 1, 0 being most relevant # We check whether the results (vector and FTS) are empty, because when - # they are, they often are missing the _rowid column, which causes an error + # they are, they often are missing the _rowid column, which causes an error. + # The original values are captured before any rank conversion so that + # they can be restored after reranking. if vector_results.num_rows > 0: distance_i = vector_results.column_names.index("_distance") original_distances = vector_results.column(distance_i) original_distance_row_ids = vector_results.column("_rowid") + if norm == "rank": + vector_results = LanceHybridQueryBuilder._rank( + vector_results, "_distance" + ) vector_results = vector_results.set_column( distance_i, vector_results.field(distance_i), - LanceHybridQueryBuilder._normalize_scores(original_distances), + LanceHybridQueryBuilder._normalize_scores( + vector_results.column(distance_i) + ), ) # In fts higher scores represent relevance. Not inverting them here as @@ -2622,10 +2626,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): score_i = fts_results.column_names.index("_score") original_scores = fts_results.column(score_i) original_score_row_ids = fts_results.column("_rowid") + if norm == "rank": + fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score") fts_results = fts_results.set_column( score_i, fts_results.field(score_i), - LanceHybridQueryBuilder._normalize_scores(original_scores), + LanceHybridQueryBuilder._normalize_scores(fts_results.column(score_i)), ) results = reranker.rerank_hybrid(fts_query, vector_results, fts_results) diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index d5ee6b6a4..743669367 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -785,6 +785,46 @@ def test_empty_hybrid_result_reranker(): assert "_rowid" in result.column_names +@pytest.mark.parametrize("norm", ["score", "rank"]) +def test_hybrid_return_all_restores_original_scores(norm): + from lancedb.query import LanceHybridQueryBuilder + + vector_results = pa.table( + { + "text": ["a", "b", "c"], + "_rowid": pa.array([0, 1, 2], type=pa.uint64()), + "_distance": pa.array([0.5, 2.5, 10.0], type=pa.float32()), + } + ) + fts_results = pa.table( + { + "text": ["b", "a"], + "_rowid": pa.array([1, 0], type=pa.uint64()), + "_score": pa.array([3.25, 1.75], type=pa.float32()), + } + ) + + result = LanceHybridQueryBuilder._combine_hybrid_results( + fts_results=fts_results, + vector_results=vector_results, + norm=norm, + fts_query="query", + reranker=RRFReranker(return_score="all"), + limit=10, + with_row_ids=True, + ) + + # The raw distances and scores are returned, not the normalized values + # or ranks used for reranking. + rows = {row["_rowid"]: row for row in result.to_pylist()} + assert rows[0]["_distance"] == 0.5 + assert rows[1]["_distance"] == 2.5 + assert rows[2]["_distance"] == 10.0 + assert rows[0]["_score"] == 1.75 + assert rows[1]["_score"] == 3.25 + assert rows[2]["_score"] is None + + @pytest.mark.parametrize("return_score", ["relevance", "all"]) @pytest.mark.parametrize("first_row_id", [1, 2**63]) @pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()])