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 <noreply@anthropic.com>

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())
         )

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())