diff --git a/python/python/lancedb/rerankers/base.py b/python/python/lancedb/rerankers/base.py index 4af938a1e..e1bc7f0dd 100644 --- a/python/python/lancedb/rerankers/base.py +++ b/python/python/lancedb/rerankers/base.py @@ -241,6 +241,7 @@ class Reranker(ABC): This op is slower than just keeping relevance score but can be useful for debugging. """ + combined_schema = pa.unify_schemas([vector_results.schema, fts_results.schema]) # add nulls to fts results for _distance if "_distance" not in fts_results.column_names: fts_results = fts_results.append_column( @@ -265,7 +266,9 @@ class Reranker(ABC): else: vector_results_dict[key] = value - combined = pa.Table.from_pylist(list(vector_results_dict.values())) + combined = pa.Table.from_pylist( + list(vector_results_dict.values()), schema=combined_schema + ) return combined def _keep_relevance_score(self, combined_results: pa.Table): diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 850b477a0..d5ee6b6a4 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -785,6 +785,114 @@ def test_empty_hybrid_result_reranker(): assert "_rowid" in result.column_names +@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()]) +def test_model_reranker_preserves_arrow_schema( + mocker, return_score, first_row_id, score_type +): + common_schema = pa.schema( + [ + pa.field("_rowid", pa.uint64()), + pa.field("text", pa.string()), + pa.field("vector", pa.list_(pa.float32(), 2)), + ], + metadata={b"source": b"reranker-test"}, + ) + vector_results = pa.Table.from_pylist( + [ + { + "_rowid": first_row_id, + "text": "vector", + "vector": [1, 0], + "_distance": 0.1, + }, + { + "_rowid": first_row_id + 1, + "text": "both", + "vector": [0, 1], + "_distance": 0.2, + }, + ], + schema=common_schema.append(pa.field("_distance", score_type)), + ) + fts_results = pa.Table.from_pylist( + [ + { + "_rowid": first_row_id + 1, + "text": "both", + "vector": [0, 1], + "_score": 0.9, + }, + { + "_rowid": first_row_id + 2, + "text": "fts", + "vector": [1, 1], + "_score": 0.8, + }, + ], + schema=common_schema.append(pa.field("_score", score_type)), + ) + reranker = CohereReranker(return_score=return_score) + reranker._client = mocker.Mock() + reranker._client.rerank.return_value.results = [ + mocker.Mock(index=index, relevance_score=score) + for index, score in [(2, 0.9), (1, 0.8), (0, 0.7)] + ] + + result = reranker.rerank_hybrid("query", vector_results, fts_results) + + assert result.select(common_schema.names).schema.equals( + common_schema, check_metadata=True + ) + assert result["_rowid"].to_pylist() == [ + first_row_id + 2, + first_row_id + 1, + first_row_id, + ] + assert result["text"].to_pylist() == ["fts", "both", "vector"] + if return_score == "all": + assert ( + result["_distance"] + .combine_chunks() + .equals(pa.array([None, 0.2, 0.1], type=score_type)) + ) + assert ( + result["_score"] + .combine_chunks() + .equals(pa.array([0.8, 0.9, None], type=score_type)) + ) + + +@pytest.mark.parametrize("return_score", ["relevance", "all"]) +def test_model_reranker_empty_results_preserve_schema(mocker, return_score): + common_schema = pa.schema( + [ + ("_rowid", pa.uint64()), + ("text", pa.string()), + ("vector", pa.list_(pa.float32(), 2)), + ], + metadata={b"source": b"reranker-test"}, + ) + vector_results = pa.Table.from_pylist( + [], schema=common_schema.append(pa.field("_distance", pa.float32())) + ) + fts_results = pa.Table.from_pylist( + [], schema=common_schema.append(pa.field("_score", pa.float32())) + ) + reranker = CohereReranker(return_score=return_score) + reranker._client = mocker.Mock() + + result = reranker.rerank_hybrid("query", vector_results, fts_results) + + assert len(result) == 0 + assert result.select(common_schema.names).schema.equals( + common_schema, check_metadata=True + ) + assert result.schema.field("_relevance_score").type == pa.float32() + reranker._client.rerank.assert_not_called() + + def test_cross_encoder_reranker_return_all(tmp_path): pytest.importorskip("sentence_transformers") reranker = CrossEncoderReranker(return_score="all")