diff --git a/nodejs/__test__/rerankers.test.ts b/nodejs/__test__/rerankers.test.ts index 256ef686c..3e5628f82 100644 --- a/nodejs/__test__/rerankers.test.ts +++ b/nodejs/__test__/rerankers.test.ts @@ -79,6 +79,18 @@ describe("rerankers", function () { expect(result).toHaveLength(2); }); + it("returns relevance scores when reranking a vector search", async function () { + const result = await table + .vectorSearch([0.1, 0.1]) + .limit(2) + .rerank(await RRFReranker.create()) + .toArray(); + + expect(result).toHaveLength(2); + expect(result[0]._relevance_score).toBeCloseTo(1 / 60); + expect(result[1]._relevance_score).toBeCloseTo(1 / 61); + }); + it("does not keep process alive after rerank query", async function () { const script = ` import * as lancedb from "./dist/index.js"; diff --git a/nodejs/lancedb/rerankers/index.ts b/nodejs/lancedb/rerankers/index.ts index 653499e32..97999f8a2 100644 --- a/nodejs/lancedb/rerankers/index.ts +++ b/nodejs/lancedb/rerankers/index.ts @@ -5,9 +5,9 @@ import { RecordBatch } from "apache-arrow"; export * from "./rrf"; -// Interface for a reranker. A reranker is used to rerank the results from a -// vector and FTS search. This is useful for combining the results from both -// search methods. +// Interface for a reranker. A reranker is used to rerank vector and hybrid +// search results. For vector-only searches, query is empty and ftsResults is an +// empty batch with the same schema as vecResults. export interface Reranker { rerankHybrid( query: string, diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index b76865043..2797721c8 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -511,7 +511,8 @@ pub trait QueryBase { /// Rerank the results using the specified reranker. /// - /// This is currently only supported for Hybrid Search. + /// For vector-only searches, the reranker receives the vector results and an + /// empty full-text result set. fn rerank(self, reranker: Arc) -> Self; /// The method to normalize the scores. Can be "rank" or "Score". If "Rank", @@ -1443,6 +1444,48 @@ impl VectorQuery { Ok(single_batch_stream(results, max_batch_length)) } + async fn execute_vector_rerank( + &self, + options: QueryExecutionOptions, + ) -> Result { + let max_batch_length = options.max_batch_length as usize; + let internal_options = options.without_output_batch_length_limit(); + + // RRF needs row IDs to assign and preserve scores. Keep them internal unless + // the caller explicitly requested them. + let vector_query = self.clone().with_row_id(); + let vector_results = vector_query + .inner_execute_with_options(internal_options) + .await?; + let schema = vector_results.schema(); + let vector_results = vector_results.try_collect::>().await?; + let vector_results = concat_batches(&schema, vector_results.iter())?; + let fts_results = RecordBatch::new_empty(vector_results.schema()); + + let reranker = self + .request + .base + .reranker + .as_ref() + .expect("execute_vector_rerank requires a reranker"); + let mut results = reranker + .rerank_hybrid("", vector_results, fts_results) + .await?; + + check_reranker_result(&results)?; + + let limit = self.request.base.limit.unwrap_or(DEFAULT_TOP_K); + if results.num_rows() > limit { + results = results.slice(0, limit); + } + + if !self.request.base.with_row_id { + results = results.drop_column(ROW_ID)?; + } + + Ok(single_batch_stream(results, max_batch_length)) + } + async fn inner_execute_with_options( &self, options: QueryExecutionOptions, @@ -1495,6 +1538,10 @@ impl ExecutableQuery for VectorQuery { return Ok(hybrid_result); } + if self.request.base.reranker.is_some() { + return self.execute_vector_rerank(options).await; + } + self.inner_execute_with_options(options).await } diff --git a/rust/lancedb/src/rerankers.rs b/rust/lancedb/src/rerankers.rs index f439c28a5..9e2c71f7b 100644 --- a/rust/lancedb/src/rerankers.rs +++ b/rust/lancedb/src/rerankers.rs @@ -47,16 +47,15 @@ impl std::fmt::Display for NormalizeMethod { } } -/// Interface for a reranker. A reranker is used to rerank the results from a -/// vector and FTS search. This is useful for combining the results from both -/// search methods. +/// Interface for a reranker. A reranker is used to rerank vector and hybrid +/// search results. This is useful for combining results from multiple search +/// methods or assigning a relevance score to vector search results. #[async_trait] pub trait Reranker: std::fmt::Debug + Sync + Send { - // TODO support vector reranking and FTS reranking. Currently only hybrid reranking is supported. - /// Rerank function receives the individual results from the vector and FTS search /// results. You can choose to use any of the results to generate the final results, - /// allowing maximum flexibility. + /// allowing maximum flexibility. For a vector-only search, `query` is empty and + /// `fts_results` is an empty batch with the same schema as `vector_results`. async fn rerank_hybrid( &self, query: &str,