From a0946282ca57a58e18b8cea9a50be39195b3323a Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:06 +0100 Subject: [PATCH] fix(python): return original scores from hybrid search with normalize="rank" (#4343) ## Summary With `rerank(..., normalize="rank")`, a hybrid query returns ranks (1.0, 2.0, 3.0, ...) in the `_distance` and `_score` columns instead of the real vector distances and BM25 scores. `_combine_hybrid_results` converted both result sets to ranks first, and only then captured `original_distances` / `original_scores`. So the step that restores the original values after reranking put the ranks back. With `normalize="score"`, the real values are returned as expected. Repro (RRFReranker(return_score="all"), vector distances 0 / 25 / 100): | normalize | `_distance` returned | `_score` returned | | --- | --- | --- | | `"score"` | 0.0, 25.0, 100.0 | 0.47, 0.67, None | | `"rank"` (before) | 1.0, 2.0, 3.0 | 1.0, 2.0, None | | `"rank"` (after) | 0.0, 25.0, 100.0 | 0.47, 0.67, None | ## Changes - Capture the original `_distance` / `_score` columns before converting them to ranks, then normalize the ranked column. The reranker input is unchanged; only the values restored afterwards are fixed. `AsyncHybridQuery` uses the same function, so both sync and async are covered. ## Test plan - [x] New `test_hybrid_return_all_restores_original_scores` in `test_rerankers.py`, parametrized over `score` / `rank`: the `rank` case fails before the fix and passes after - [x] `pytest python/tests/test_rerankers.py python/tests/test_hybrid_query.py` - [x] `ruff format --check` / `ruff check` I have no local Rust toolchain, so I ran the tests against the released 0.39.0 native module with the changed Python files overlaid. The changed functions are identical between 0.39.0 and main. Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 20 +++++++++----- python/python/tests/test_rerankers.py | 40 +++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 7 deletions(-) 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()])