mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
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:
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user