fix(node): rerank vector search results

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";

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,

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<dyn Reranker>) -> 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<SendableRecordBatchStream> {
+        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::<Vec<_>>().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
     }
 

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,