From 55897d7c3b03c6b94b461161abc1b28d34903fde Mon Sep 17 00:00:00 2001 From: Yingzi Zhang <164733623+LindseyZ1205@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:26:43 -0400 Subject: [PATCH] fix(python): preserve Arrow schema when merging reranker scores (#4332) Fixes #4331 With `return_score="all"`, model rerankers reconstruct merged rows without their Arrow schema. This changes fixed-size float32 vectors into variable-length float64 lists, loses metadata, and can overflow valid uint64 row IDs. Unify the original input schemas before inserting placeholder scores, then use the unified schema when reconstructing the merged table. Row merging and reranking behavior stay the same. Capturing the schema first also preserves float64 scores without conflicting with the float32 null placeholders. Validation: - New regression cases: 5 failures and 5 passing controls before the fix; all 10 pass afterward. - Reranker, hybrid-query and TypeSafe reranker tests: 150 passed, 17 skipped. Optional live model/API tests were not run. - The fake-client tests that use random metadata passed again with seeds 7 and 41. - Repository-wide Ruff lint/format checks, configured Pyright check and applicable pre-commit hooks pass. - Tests used the locally built Rust/Python extension and the repository's locked uv environment. No Rust code changed. AI disclosure: I used Codex for investigation, implementation, tests, local checks, and drafting this description. --- python/python/lancedb/rerankers/base.py | 5 +- python/python/tests/test_rerankers.py | 108 ++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) 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")