From 3ee81fa5560275c44987224e2079a879dd71d223 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:16:39 +0000 Subject: [PATCH] fix(node): preserve vector rerank query contracts --- .../rerankers/classes/RRFReranker.md | 16 ++ .../rerankers/interfaces/Reranker.md | 21 +++ nodejs/__test__/rerankers.test.ts | 10 +- nodejs/lancedb/query.ts | 38 ++-- nodejs/lancedb/rerankers/index.ts | 10 +- nodejs/lancedb/rerankers/rrf.ts | 13 +- nodejs/src/query.rs | 5 +- nodejs/src/rerankers.rs | 53 +++++- rust/lancedb/src/query.rs | 178 +++++++++++++++++- rust/lancedb/src/rerankers.rs | 14 ++ rust/lancedb/src/rerankers/rrf.rs | 25 ++- 11 files changed, 348 insertions(+), 35 deletions(-) diff --git a/docs/src/js/namespaces/rerankers/classes/RRFReranker.md b/docs/src/js/namespaces/rerankers/classes/RRFReranker.md index d054b7f0b..34c25f2fa 100644 --- a/docs/src/js/namespaces/rerankers/classes/RRFReranker.md +++ b/docs/src/js/namespaces/rerankers/classes/RRFReranker.md @@ -10,6 +10,22 @@ Reranks the results using the Reciprocal Rank Fusion (RRF) algorithm. ## Methods +### outputSchema() + +```ts +outputSchema(inputSchema): Promise> +``` + +#### Parameters + +* **inputSchema**: `Schema`<`any`> + +#### Returns + +`Promise`<`Schema`<`any`>> + +*** + ### rerankHybrid() ```ts diff --git a/docs/src/js/namespaces/rerankers/interfaces/Reranker.md b/docs/src/js/namespaces/rerankers/interfaces/Reranker.md index 1d0566730..1fd7def36 100644 --- a/docs/src/js/namespaces/rerankers/interfaces/Reranker.md +++ b/docs/src/js/namespaces/rerankers/interfaces/Reranker.md @@ -8,6 +8,27 @@ ## Methods +### outputSchema()? + +```ts +optional outputSchema(inputSchema): Promise> +``` + +Declare the schema returned when reranking a vector-only query. + +This is required for vector-only reranking so query schema introspection +and execution agree. Hybrid-only rerankers may omit it. + +#### Parameters + +* **inputSchema**: `Schema`<`any`> + +#### Returns + +`Promise`<`Schema`<`any`>> + +*** + ### rerankHybrid() ```ts diff --git a/nodejs/__test__/rerankers.test.ts b/nodejs/__test__/rerankers.test.ts index 3e5628f82..b0fe5d1ee 100644 --- a/nodejs/__test__/rerankers.test.ts +++ b/nodejs/__test__/rerankers.test.ts @@ -80,12 +80,16 @@ describe("rerankers", function () { }); it("returns relevance scores when reranking a vector search", async function () { - const result = await table + const query = table .vectorSearch([0.1, 0.1]) .limit(2) - .rerank(await RRFReranker.create()) - .toArray(); + .rerank(await RRFReranker.create()); + const schema = await query.outputSchema(); + const result = await query.toArray(); + expect(schema.fields.map((field) => field.name)).toContain( + "_relevance_score", + ); expect(result).toHaveLength(2); expect(result[0]._relevance_score).toBeCloseTo(1 / 60); expect(result[1]._relevance_score).toBeCloseTo(1 / 61); diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 843a1276f..d3712161d 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -5,9 +5,11 @@ import { Table as ArrowTable, type IntoVector, RecordBatch, + createEmptyTable, extractVectorBuffer, fromBufferToRecordBatch, fromRecordBatchToBuffer, + fromTableToBuffer, tableFromIPC, } from "./arrow"; import { type IvfPqOptions } from "./indices"; @@ -744,20 +746,30 @@ export class VectorQuery extends StandardQueryBase { } rerank(reranker: Reranker): VectorQuery { - super.doCall((inner) => - inner.rerank(async (args) => { - const vecResults = await fromBufferToRecordBatch(args.vecResults); - const ftsResults = await fromBufferToRecordBatch(args.ftsResults); - const result = await reranker.rerankHybrid( - args.query, - vecResults as RecordBatch, - ftsResults as RecordBatch, - ); + super.doCall((inner) => { + const outputSchema = reranker.outputSchema?.bind(reranker); + inner.rerank( + async (args) => { + const vecResults = await fromBufferToRecordBatch(args.vecResults); + const ftsResults = await fromBufferToRecordBatch(args.ftsResults); + const result = await reranker.rerankHybrid( + args.query, + vecResults as RecordBatch, + ftsResults as RecordBatch, + ); - const buffer = fromRecordBatchToBuffer(result); - return buffer; - }), - ); + const buffer = fromRecordBatchToBuffer(result); + return buffer; + }, + outputSchema + ? async (args) => { + const inputSchema = tableFromIPC(args.inputSchema).schema; + const result = await outputSchema(inputSchema); + return fromTableToBuffer(createEmptyTable(result)); + } + : undefined, + ); + }); return this; } diff --git a/nodejs/lancedb/rerankers/index.ts b/nodejs/lancedb/rerankers/index.ts index 97999f8a2..2e04b2a45 100644 --- a/nodejs/lancedb/rerankers/index.ts +++ b/nodejs/lancedb/rerankers/index.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -import { RecordBatch } from "apache-arrow"; +import { RecordBatch, Schema } from "apache-arrow"; export * from "./rrf"; @@ -9,6 +9,14 @@ export * from "./rrf"; // search results. For vector-only searches, query is empty and ftsResults is an // empty batch with the same schema as vecResults. export interface Reranker { + /** + * Declare the schema returned when reranking a vector-only query. + * + * This is required for vector-only reranking so query schema introspection + * and execution agree. Hybrid-only rerankers may omit it. + */ + outputSchema?(inputSchema: Schema): Promise; + rerankHybrid( query: string, vecResults: RecordBatch, diff --git a/nodejs/lancedb/rerankers/rrf.ts b/nodejs/lancedb/rerankers/rrf.ts index 654e61301..ce4e5d877 100644 --- a/nodejs/lancedb/rerankers/rrf.ts +++ b/nodejs/lancedb/rerankers/rrf.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -import { RecordBatch } from "apache-arrow"; +import { Field, Float32, RecordBatch, Schema } from "apache-arrow"; import { fromBufferToRecordBatch, fromRecordBatchToBuffer } from "../arrow"; import { RrfReranker as NativeRRFReranker } from "../native"; @@ -24,6 +24,17 @@ export class RRFReranker { ); } + /** Declare the RRF output schema for vector-only query execution. */ + async outputSchema(inputSchema: Schema): Promise { + return new Schema( + [ + ...inputSchema.fields, + new Field("_relevance_score", new Float32(), false), + ], + inputSchema.metadata, + ); + } + async rerankHybrid( query: string, vecResults: RecordBatch, diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index 3828023a9..42af9ae94 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use crate::error::NapiErrorExt; use crate::error::convert_error; use crate::iterator::RecordBatchIterator; -use crate::rerankers::RerankHybridCallbackArgs; use crate::rerankers::Reranker; +use crate::rerankers::{RerankHybridCallbackArgs, RerankOutputSchemaCallbackArgs}; use crate::util::{parse_distance_type, schema_to_buffer}; use arrow_array::{ Array, Float16Array as ArrowFloat16Array, Float32Array as ArrowFloat32Array, @@ -388,8 +388,9 @@ impl VectorQuery { pub fn rerank( &mut self, rerank_hybrid: Function>, + output_schema: Option>>, ) -> napi::Result<()> { - let reranker = Reranker::new(rerank_hybrid)?; + let reranker = Reranker::new(rerank_hybrid, output_schema)?; self.inner = self.inner.clone().rerank(Arc::new(reranker)); Ok(()) } diff --git a/nodejs/src/rerankers.rs b/nodejs/src/rerankers.rs index 7a4d71f93..b09d707a2 100644 --- a/nodejs/src/rerankers.rs +++ b/nodejs/src/rerankers.rs @@ -6,7 +6,7 @@ use async_trait::async_trait; use napi::{bindgen_prelude::*, threadsafe_function::ThreadsafeFunction}; use napi_derive::napi; -use lancedb::ipc::batches_to_ipc_file; +use lancedb::ipc::{batches_to_ipc_file, ipc_file_to_schema, schema_to_ipc_file}; use lancedb::rerankers::Reranker as LanceDBReranker; use lancedb::{error::Error, ipc::ipc_file_to_batches}; @@ -21,28 +21,72 @@ type RerankHybridFn = ThreadsafeFunction< true, >; +type RerankOutputSchemaFn = ThreadsafeFunction< + RerankOutputSchemaCallbackArgs, + Promise, + RerankOutputSchemaCallbackArgs, + Status, + false, + true, +>; + /// Reranker implementation that "wraps" a NodeJS Reranker implementation. /// This contains references to the callbacks that can be used to invoke the /// reranking methods on the NodeJS implementation and handles serializing the /// record batches to Arrow IPC buffers. pub struct Reranker { rerank_hybrid: RerankHybridFn, + output_schema: Option, } impl Reranker { pub fn new( rerank_hybrid: Function>, + output_schema: Option>>, ) -> napi::Result { let rerank_hybrid = rerank_hybrid .build_threadsafe_function() .weak::() .build()?; - Ok(Self { rerank_hybrid }) + let output_schema = output_schema + .map(|output_schema| { + output_schema + .build_threadsafe_function() + .weak::() + .build() + }) + .transpose()?; + Ok(Self { + rerank_hybrid, + output_schema, + }) } } #[async_trait] impl lancedb::rerankers::Reranker for Reranker { + async fn output_schema( + &self, + input: &arrow_schema::SchemaRef, + ) -> lancedb::error::Result { + let output_schema = self.output_schema.as_ref().ok_or(Error::NotSupported { + message: "vector rerankers must declare their output schema".to_string(), + })?; + let callback_args = RerankOutputSchemaCallbackArgs { + input_schema: Buffer::from(schema_to_ipc_file(input.as_ref())?), + }; + let promised_buffer: Promise = output_schema + .call_async(callback_args) + .await + .map_err(|e| Error::Runtime { + message: format!("napi error status={}, reason={}", e.status, e.reason), + })?; + let buffer = promised_buffer.await.map_err(|e| Error::Runtime { + message: format!("napi error status={}, reason={}", e.status, e.reason), + })?; + ipc_file_to_schema(buffer.to_vec()) + } + async fn rerank_hybrid( &self, query: &str, @@ -86,6 +130,11 @@ pub struct RerankHybridCallbackArgs { pub fts_results: Buffer, } +#[napi(object)] +pub struct RerankOutputSchemaCallbackArgs { + pub input_schema: Buffer, +} + fn buffer_to_record_batch(buffer: Buffer) -> Result { let mut reader = ipc_file_to_batches(buffer.to_vec()).default_error()?; reader diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 2797721c8..5f336986a 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -512,7 +512,8 @@ pub trait QueryBase { /// Rerank the results using the specified reranker. /// /// For vector-only searches, the reranker receives the vector results and an - /// empty full-text result set. + /// empty full-text result set and must declare its output schema. Reranking + /// multiple query vectors in one query is not supported. fn rerank(self, reranker: Arc) -> Self; /// The method to normalize the scores. Can be "rank" or "Score". If "Rank", @@ -1139,6 +1140,44 @@ pub struct VectorQuery { } impl VectorQuery { + fn check_vector_rerank_supported(&self) -> Result<()> { + if self.request.query_vector.len() > 1 { + return Err(Error::NotSupported { + message: "reranking multiple query vectors is not supported; execute one query per vector" + .to_string(), + }); + } + Ok(()) + } + + async fn vector_rerank_output_schema(&self) -> Result { + self.check_vector_rerank_supported()?; + + // Rerankers receive row IDs internally. Apply their schema transform to + // that exact input and then hide the row ID from the declared public + // schema unless it was explicitly requested. + let vector_query = self.clone().with_row_id(); + let plan = vector_query + .create_plan(QueryExecutionOptions::default()) + .await?; + let reranker = self + .request + .base + .reranker + .as_ref() + .expect("vector_rerank_output_schema requires a reranker"); + let input_schema = plan.schema(); + let output_schema = reranker.output_schema(&input_schema).await?; + + if self.request.base.with_row_id { + Ok(output_schema) + } else { + Ok(RecordBatch::new_empty(output_schema) + .drop_column(ROW_ID)? + .schema()) + } + } + fn new(base: Query) -> Self { Self { parent: base.parent, @@ -1448,6 +1487,8 @@ impl VectorQuery { &self, options: QueryExecutionOptions, ) -> Result { + self.check_vector_rerank_supported()?; + let max_batch_length = options.max_batch_length as usize; let internal_options = options.without_output_batch_length_limit(); @@ -1460,7 +1501,8 @@ impl VectorQuery { 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 vector_schema = vector_results.schema(); + let fts_results = RecordBatch::new_empty(vector_schema.clone()); let reranker = self .request @@ -1468,11 +1510,21 @@ impl VectorQuery { .reranker .as_ref() .expect("execute_vector_rerank requires a reranker"); + let expected_schema = reranker.output_schema(&vector_schema).await?; let mut results = reranker .rerank_hybrid("", vector_results, fts_results) .await?; check_reranker_result(&results)?; + if results.schema() != expected_schema { + return Err(Error::Schema { + message: format!( + "reranker returned schema {:?}, but declared {:?}", + results.schema(), + expected_schema + ), + }); + } let limit = self.request.base.limit.unwrap_or(DEFAULT_TOP_K); if results.num_rows() > limit { @@ -1539,7 +1591,20 @@ impl ExecutableQuery for VectorQuery { } if self.request.base.reranker.is_some() { - return self.execute_vector_rerank(options).await; + let timeout = options.timeout; + let mut rerank_options = options; + // A single outer deadline covers planning, candidate collection, + // schema declaration, and the complete reranker callback. + rerank_options.timeout = None; + let execution = self.execute_vector_rerank(rerank_options); + return match timeout { + Some(timeout) => tokio::time::timeout(timeout, execution) + .await + .map_err(|_| Error::Timeout { + message: format!("Query timeout after {} ms", timeout.as_millis()), + })?, + None => execution.await, + }; } self.inner_execute_with_options(options).await @@ -1554,6 +1619,15 @@ impl ExecutableQuery for VectorQuery { let query = AnyQuery::VectorQuery(self.request.clone()); self.parent.analyze_plan(&query, options).await } + + async fn output_schema(&self) -> Result { + if self.request.base.full_text_search.is_none() && self.request.base.reranker.is_some() { + self.vector_rerank_output_schema().await + } else { + let plan = self.create_plan(QueryExecutionOptions::default()).await?; + Ok(plan.schema()) + } + } } impl HasQuery for VectorQuery { @@ -1690,7 +1764,13 @@ impl ExecutableQuery for TakeQuery { #[cfg(test)] mod tests { - use std::{collections::HashSet, sync::Arc}; + use std::{ + collections::HashSet, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + }; use super::*; use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type}; @@ -1706,6 +1786,31 @@ mod tests { use crate::{Table, connect, database::CreateTableMode, index::Index}; + #[derive(Debug)] + struct SlowReranker { + invoked: Arc, + } + + #[async_trait::async_trait] + impl Reranker for SlowReranker { + async fn output_schema(&self, input: &SchemaRef) -> Result { + RRFReranker::default().output_schema(input).await + } + + async fn rerank_hybrid( + &self, + query: &str, + vector_results: RecordBatch, + fts_results: RecordBatch, + ) -> Result { + self.invoked.store(true, Ordering::SeqCst); + tokio::time::sleep(Duration::from_secs(2)).await; + RRFReranker::default() + .rerank_hybrid(query, vector_results, fts_results) + .await + } + } + #[tokio::test] async fn test_setters_getters() { // TODO: Switch back to memory://foo after https://github.com/lancedb/lancedb/issues/1051 @@ -2417,6 +2522,71 @@ mod tests { // We don't guarantee order. assert!(query_index.values().contains(&0)); assert!(query_index.values().contains(&1)); + + let reranked = query.rerank(Arc::new(RRFReranker::default())); + let Err(execute_error) = reranked.execute().await else { + panic!("multi-vector reranking should be rejected"); + }; + assert!( + execute_error + .to_string() + .contains("reranking multiple query vectors is not supported") + ); + let schema_error = reranked.output_schema().await.unwrap_err(); + assert!( + schema_error + .to_string() + .contains("reranking multiple query vectors is not supported") + ); + } + + #[tokio::test] + async fn test_vector_rerank_timeout_covers_reranker() { + let tmp_dir = tempdir().unwrap(); + let table = make_test_table(&tmp_dir).await; + let invoked = Arc::new(AtomicBool::new(false)); + let reranker = SlowReranker { + invoked: invoked.clone(), + }; + + let result = table + .vector_search(&[0.1, 0.2, 0.3, 0.4]) + .unwrap() + .limit(1) + .rerank(Arc::new(reranker)) + .execute_with_options(QueryExecutionOptions { + timeout: Some(Duration::from_secs(1)), + ..Default::default() + }) + .await; + + assert!(invoked.load(Ordering::SeqCst)); + assert!(matches!(result, Err(Error::Timeout { .. }))); + } + + #[tokio::test] + async fn test_vector_rerank_output_schema_matches_execution() { + let tmp_dir = tempdir().unwrap(); + let table = make_test_table(&tmp_dir).await; + let query = table + .vector_search(&[0.1, 0.2, 0.3, 0.4]) + .unwrap() + .limit(1) + .rerank(Arc::new(RRFReranker::default())); + + let promised = query.output_schema().await.unwrap(); + let actual = query + .execute() + .await + .unwrap() + .next() + .await + .unwrap() + .unwrap() + .schema(); + + assert_eq!(promised, actual); + assert!(promised.column_with_name("_relevance_score").is_some()); } #[tokio::test] diff --git a/rust/lancedb/src/rerankers.rs b/rust/lancedb/src/rerankers.rs index 9e2c71f7b..cea5d3130 100644 --- a/rust/lancedb/src/rerankers.rs +++ b/rust/lancedb/src/rerankers.rs @@ -8,6 +8,7 @@ use arrow::{ compute::{concat_batches, filter_record_batch}, }; use arrow_array::{BooleanArray, RecordBatch, UInt64Array}; +use arrow_schema::SchemaRef; use async_trait::async_trait; use lance::dataset::ROW_ID; @@ -52,6 +53,19 @@ impl std::fmt::Display for NormalizeMethod { /// methods or assigning a relevance score to vector search results. #[async_trait] pub trait Reranker: std::fmt::Debug + Sync + Send { + /// Declare the schema returned by [`Self::rerank_hybrid`] for a vector-only + /// query. + /// + /// Vector reranking validates the returned batch against this schema so + /// [`crate::query::ExecutableQuery::output_schema`] and execution cannot + /// disagree. Rerankers that only support hybrid search do not need to + /// implement this method. + async fn output_schema(&self, _input: &SchemaRef) -> Result { + Err(Error::NotSupported { + message: "vector rerankers must declare their output schema".to_string(), + }) + } + /// 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. For a vector-only search, `query` is empty and diff --git a/rust/lancedb/src/rerankers/rrf.rs b/rust/lancedb/src/rerankers/rrf.rs index 4da1927be..9515a6eca 100644 --- a/rust/lancedb/src/rerankers/rrf.rs +++ b/rust/lancedb/src/rerankers/rrf.rs @@ -9,7 +9,7 @@ use arrow::{ compute::{sort_to_indices, take}, }; use arrow_array::{Float32Array, RecordBatch, UInt64Array}; -use arrow_schema::{DataType, Field, Schema, SortOptions}; +use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions}; use async_trait::async_trait; use lance::dataset::ROW_ID; @@ -44,6 +44,19 @@ impl Default for RRFReranker { #[async_trait] impl Reranker for RRFReranker { + async fn output_schema(&self, input: &SchemaRef) -> Result { + let mut fields = input.fields().to_vec(); + fields.push(Arc::new(Field::new( + RELEVANCE_SCORE, + DataType::Float32, + false, + ))); + Ok(Arc::new(Schema::new_with_metadata( + fields, + input.metadata().clone(), + ))) + } + async fn rerank_hybrid( &self, _query: &str, @@ -135,15 +148,9 @@ impl Reranker for RRFReranker { .collect(); // add relevance score to schema - let mut fields = combined_results.schema().fields().to_vec(); - fields.push(Arc::new(Field::new( - RELEVANCE_SCORE, - DataType::Float32, - false, - ))); - let schema = Schema::new(fields); + let schema = self.output_schema(&combined_results.schema()).await?; - let combined_results = RecordBatch::try_new(Arc::new(schema), columns)?; + let combined_results = RecordBatch::try_new(schema, columns)?; Ok(combined_results) }