From 18760f74cd43b24e2e8026d88baf9ace1facd3a3 Mon Sep 17 00:00:00 2001 From: Andrew Chen Date: Sat, 25 Jul 2026 06:02:23 +0800 Subject: [PATCH] fix: crash in AnswerdotaiRerankers/ColbertReranker for return_score="all" (#3671) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What `AnswerdotaiRerankers(return_score="all").rerank_hybrid(...)` (and `ColbertReranker`, which subclasses it without overriding `rerank_hybrid`) raises: ``` pyarrow.lib.ArrowInvalid: Invalid sort key column: No match for FieldRef.Name(_relevance_score) in _rowid: int64 ... ``` ## Why ```python combined_results = self.merge_results(vector_results, fts_results) combined_results = self._rerank(combined_results, query) if self.score == "relevance": combined_results = self._keep_relevance_score(combined_results) elif self.score == "all": combined_results = self._merge_and_keep_scores(vector_results, fts_results) ``` When `score == "all"`, `combined_results` is unconditionally overwritten by `_merge_and_keep_scores(vector_results, fts_results)` **after** `_rerank()` already computed and appended `_relevance_score` — discarding it. The following `sort_by("_relevance_score", ...)` then has nothing to sort on. Every sibling reranker that supports `return_score="all"` (`cross_encoder`, `openai`, `cohere`, `jinaai`, `voyageai`, `watsonx`) instead calls `_merge_and_keep_scores()` **before** `_rerank()`. This file is the one place the ordering got inverted when `"all"` support was added (#2509) — a copy/paste inconsistency across the six files that PR touched. Fix mirrors the pattern already used (and tested) by the other five rerankers. Also drops the now-stale `"Only 'relevance' is supported for now"` docstring line on both classes, left over from before `"all"` support existed. ## Testing Added `test_answerdotai_reranker_return_all`, mirroring the existing `test_cross_encoder_reranker_return_all`. Verified locally with the real built Rust extension: red (reproduces the exact `ArrowInvalid` above) → green, using the actual `rerank_hybrid`/`_rerank`/`base.py` code path with the model call mocked out — my local environment's `rerankers==0.10.0` fails to load the real ColBERT model against the available `transformers` version (`AttributeError: 'ColBERTModel' object has no attribute 'all_tied_weights_keys'`), which I confirmed also breaks the **pre-existing**, unmodified `test_colbert_reranker`/`test_answerdotai_reranker` baseline tests identically — an unrelated local dependency-version issue, not a regression from this change. `ruff check`/`ruff format` clean; full `test_rerankers.py` run: 9 passed / 8 skipped / 3 failed (the 3 failures are exactly those two pre-existing tests plus my new one, all failing at model-loading time for the same unrelated reason before reaching the changed code). --- Disclosure: this PR was drafted with AI assistance (Claude); I reviewed, tested, and take responsibility for the change. Co-authored-by: Claude Opus 4.8 --- python/python/lancedb/rerankers/answerdotai.py | 9 +++++---- python/python/lancedb/rerankers/colbert.py | 2 +- python/python/tests/test_rerankers.py | 15 +++++++++++++++ 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/python/python/lancedb/rerankers/answerdotai.py b/python/python/lancedb/rerankers/answerdotai.py index 0b44569e3..0e2d9c6cb 100644 --- a/python/python/lancedb/rerankers/answerdotai.py +++ b/python/python/lancedb/rerankers/answerdotai.py @@ -23,7 +23,7 @@ class AnswerdotaiRerankers(Reranker): column : str, default "text" The name of the column to use as input to the cross encoder model. return_score : str, default "relevance" - options are "relevance" or "all". Only "relevance" is supported for now. + options are "relevance" or "all". **kwargs Additional keyword arguments to pass to the model. For example, 'device'. See AnswerDotAI/rerankers for more information. @@ -77,12 +77,13 @@ class AnswerdotaiRerankers(Reranker): vector_results: pa.Table, fts_results: pa.Table, ): - combined_results = self.merge_results(vector_results, fts_results) + if self.score == "all": + combined_results = self._merge_and_keep_scores(vector_results, fts_results) + else: + combined_results = self.merge_results(vector_results, fts_results) combined_results = self._rerank(combined_results, query) if self.score == "relevance": combined_results = self._keep_relevance_score(combined_results) - elif self.score == "all": - combined_results = self._merge_and_keep_scores(vector_results, fts_results) combined_results = combined_results.sort_by( [("_relevance_score", "descending")] ) diff --git a/python/python/lancedb/rerankers/colbert.py b/python/python/lancedb/rerankers/colbert.py index a7a98b361..ff3bc8c4a 100644 --- a/python/python/lancedb/rerankers/colbert.py +++ b/python/python/lancedb/rerankers/colbert.py @@ -16,7 +16,7 @@ class ColbertReranker(AnswerdotaiRerankers): column : str, default "text" The name of the column to use as input to the cross encoder model. return_score : str, default "relevance" - options are "relevance" or "all". Only "relevance" is supported for now. + options are "relevance" or "all". **kwargs Additional keyword arguments to pass to the model, for example, 'device'. See AnswerDotAI/rerankers for more information. diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 7c2b52ffb..372a6b0f7 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -644,6 +644,21 @@ def test_cross_encoder_reranker_return_all(tmp_path): assert "_distance" in result.column_names +def test_answerdotai_reranker_return_all(tmp_path): + pytest.importorskip("rerankers") + reranker = AnswerdotaiRerankers(return_score="all") + table, schema = get_test_table(tmp_path) + query = "single player experience" + result = ( + table.search(query, query_type="hybrid", vector_column_name="vector") + .rerank(reranker=reranker) + .to_arrow() + ) + assert "_relevance_score" in result.column_names + assert "_score" in result.column_names + assert "_distance" in result.column_names + + # --------------------------------------------------------------------------- # Regression tests for LinearCombinationReranker scoring bugs (issue #3154) # ---------------------------------------------------------------------------