diff --git a/python/python/lancedb/rerankers/linear_combination.py b/python/python/lancedb/rerankers/linear_combination.py index 74f23ea61..682aac9e2 100644 --- a/python/python/lancedb/rerankers/linear_combination.py +++ b/python/python/lancedb/rerankers/linear_combination.py @@ -3,7 +3,6 @@ from collections import defaultdict -from numpy import nan import pyarrow as pa from .base import Reranker @@ -66,7 +65,7 @@ class LinearCombinationReranker(Reranker): elif self.score == "all": results = results.append_column( "_distance", - pa.array([nan] * len(fts_results), type=pa.float32()), + pa.nulls(len(fts_results), type=pa.float32()), ) return results @@ -87,7 +86,7 @@ class LinearCombinationReranker(Reranker): elif self.score == "all": results = results.append_column( "_score", - pa.array([nan] * len(vector_results), type=pa.float32()), + pa.nulls(len(vector_results), type=pa.float32()), ) return results results = defaultdict() diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index ac98efbfd..610f533f5 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -1150,6 +1150,52 @@ def test_linear_combination_missing_fts_is_penalised(): ) +@pytest.mark.parametrize("empty_side", ["vector", "fts"]) +def test_linear_combination_empty_side_scores_are_null(empty_side): + """With return_score="all", the missing side's score column is null, not NaN.""" + from lancedb.query import LanceHybridQueryBuilder + + vector_results = pa.table( + { + "_rowid": pa.array([1, 2], pa.uint64()), + "_distance": pa.array([0.1, 0.2], pa.float32()), + } + ) + fts_results = pa.table( + { + "_rowid": pa.array([1, 2], pa.uint64()), + "_score": pa.array([5.0, 4.0], pa.float32()), + } + ) + if empty_side == "vector": + expected_present = fts_results["_score"].to_pylist() + vector_results = vector_results.slice(0, 0) + missing = "_distance" + present = "_score" + else: + expected_present = vector_results["_distance"].to_pylist() + fts_results = fts_results.slice(0, 0) + missing = "_score" + present = "_distance" + + result = LanceHybridQueryBuilder._combine_hybrid_results( + fts_results=fts_results, + vector_results=vector_results, + norm="score", + fts_query="query", + reranker=LinearCombinationReranker(return_score="all"), + limit=10, + with_row_ids=True, + ) + + assert result[missing].type == pa.float32() + assert result[missing].null_count == len(result) == 2 + assert result[present].to_pylist() == expected_present + assert result["_relevance_score"].to_pylist() == sorted( + result["_relevance_score"].to_pylist(), reverse=True + ) + + @pytest.mark.skipif( os.environ.get("WATSONX_API_KEY") is None or (