From 2af7d88ec7e974484e1be0d9f7bfba53732e0ef6 Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:23:21 -0400 Subject: [PATCH] fix(python): keep OpenaiReranker scores on the rows they rate (#4549) `OpenaiReranker._rerank` asks the chat model for a relevance score per document. The model can return the documents in any order, for example sorted by score. The reranker then replaces the `text` column with the documents in that order and appends the scores in that order. The other columns keep the input order. So after a rerank, the `text` and `_relevance_score` of a row belong to a different row. The `id`, the vector and every other column still describe the original row. If the model leaves out a document, `append_column` raises `ArrowInvalid`, because the new column has fewer rows than the table. ## Fix `_rerank` keeps the `text` column unchanged, so it also keeps its position and type. It maps each score back to its row by the document content. If the model leaves out a document, `_rerank` raises a `ValueError` that gives the number of documents with no score. ## Tests `python/python/tests/test_openai_reranker.py` replaces the chat client with a stub that returns a fixed reply. The tests do not need `openai` or an API key. - One test returns the documents sorted by score. It checks that each `id` keeps its own `text` and score. On `main`, the row with id 1 gets the text `cats purr`, which belongs to id 2. - One test returns 2 of 3 documents. It checks for the `ValueError` with the count of missing documents. On `main`, the call raises `ArrowInvalid: Added column's length must match table's length`. - Both tests fail on `main` and pass with the fix. - `test_rerankers.py`: 16 passed, 13 skipped, on the lancedb 0.40.0 wheel with and without the fix. The skipped tests need real API keys. - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed file copied over it. In that wheel, `rerankers/openai.py` and `rerankers/base.py` are byte-identical to `main`. The Node SDK has no OpenAI reranker, so there is no parity change. This change does not conflict with #4515, which edits `_client` in the same file. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/rerankers/openai.py | 21 ++++--- python/python/tests/test_openai_reranker.py | 66 +++++++++++++++++++++ 2 files changed, 78 insertions(+), 9 deletions(-) create mode 100644 python/python/tests/test_openai_reranker.py diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index 4cbbb0822..da0394d1c 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -70,15 +70,18 @@ class OpenaiReranker(Reranker): ], ) results = json.loads(response.choices[0].message.content)["documents"] - docs, scores = list( - zip(*[(result["content"], result["relevance_score"]) for result in results]) - ) # tuples - # replace the self.column column with the docs - result_set = result_set.drop(self.column) - result_set = result_set.append_column( - self.column, pa.array(docs, type=pa.string()) - ) - # add the scores + # The model returns the documents in its own order, so match each score + # back to its row by content. + scores_by_content = { + result["content"]: result["relevance_score"] for result in results + } + missing = sum(doc not in scores_by_content for doc in docs) + if missing: + raise ValueError( + f"The model returned no relevance_score for {missing} of " + f"{len(docs)} documents" + ) + scores = [scores_by_content[doc] for doc in docs] result_set = result_set.append_column( "_relevance_score", pa.array(scores, type=pa.float32()) ) diff --git a/python/python/tests/test_openai_reranker.py b/python/python/tests/test_openai_reranker.py new file mode 100644 index 000000000..43475103f --- /dev/null +++ b/python/python/tests/test_openai_reranker.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json +from types import SimpleNamespace + +import pyarrow as pa +import pytest + +from lancedb.rerankers import OpenaiReranker + + +def reranker_replying(documents): + """Return an OpenaiReranker whose chat client replies with these documents.""" + reply = json.dumps({"documents": documents}) + + def create(**kwargs): + message = SimpleNamespace(content=reply) + return SimpleNamespace(choices=[SimpleNamespace(message=message)]) + + reranker = OpenaiReranker(api_key="test-key") + reranker.__dict__["_client"] = SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ) + return reranker + + +def vector_results(): + return pa.table( + { + "id": [1, 2, 3], + "text": ["dogs bark", "cats purr", "birds sing"], + "_distance": [0.1, 0.2, 0.3], + } + ) + + +def test_scores_stay_with_their_rows_when_the_model_reorders(): + reranker = reranker_replying( + [ + {"content": "cats purr", "relevance_score": 0.9}, + {"content": "birds sing", "relevance_score": 0.5}, + {"content": "dogs bark", "relevance_score": 0.1}, + ] + ) + + result = reranker.rerank_vector("which animal purrs?", vector_results()) + + assert result.select(["id", "text"]).to_pylist() == [ + {"id": 2, "text": "cats purr"}, + {"id": 3, "text": "birds sing"}, + {"id": 1, "text": "dogs bark"}, + ] + assert result["_relevance_score"].to_pylist() == pytest.approx([0.9, 0.5, 0.1]) + + +def test_a_row_the_model_leaves_out_raises(): + reranker = reranker_replying( + [ + {"content": "cats purr", "relevance_score": 0.9}, + {"content": "dogs bark", "relevance_score": 0.1}, + ] + ) + + with pytest.raises(ValueError, match="1 of 3"): + reranker.rerank_vector("which animal purrs?", vector_results())