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.
This commit is contained in:
Yingzi Zhang
2026-10-23 01:26:43 +08:00
committed by GitHub
parent 624defd109
commit 55897d7c3b
2 changed files with 112 additions and 1 deletions
+4 -1
View File
@@ -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):
+108
View File
@@ -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")