rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-18 12:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(node): preserve vector rerank query contracts

Browse Source
This commit is contained in:
Gatefixer
2026-08-06 00:16:39 +00:00
parent 48ca05c7d5
commit 3ee81fa556
11 changed files with 348 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/js/namespaces/rerankers/classes/RRFReranker.md
+16
View File
@@ -10,6 +10,22 @@ Reranks the results using the Reciprocal Rank Fusion (RRF) algorithm.
## Methods
### outputSchema()
```ts
outputSchema(inputSchema): Promise<Schema<any>>
```
#### Parameters
* **inputSchema**: `Schema`&lt;`any`&gt;
#### Returns
`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
***
### rerankHybrid()
```ts
docs/src/js/namespaces/rerankers/interfaces/Reranker.md
+21
View File
@@ -8,6 +8,27 @@
## Methods
### outputSchema()?
```ts
optional outputSchema(inputSchema): Promise<Schema<any>>
```
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`&lt;`any`&gt;
#### Returns
`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
***
### rerankHybrid()
```ts
nodejs/__test__/rerankers.test.ts
+7 -3
View File
@@ -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);
nodejs/lancedb/query.ts
+25 -13
View File
@@ -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<NativeVectorQuery> {
}
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;
}
nodejs/lancedb/rerankers/index.ts
+9 -1
View File
@@ -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<Schema>;
rerankHybrid(
query: string,
vecResults: RecordBatch,
nodejs/lancedb/rerankers/rrf.ts
+12 -1
View File
@@ -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<Schema> {
return new Schema(
[
...inputSchema.fields,
new Field("_relevance_score", new Float32(), false),
],
inputSchema.metadata,
);
}
async rerankHybrid(
query: string,
vecResults: RecordBatch,
nodejs/src/query.rs
+3 -2
View File
@@ -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<RerankHybridCallbackArgs, Promise<Buffer>>,
output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
) -> 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(())
}
nodejs/src/rerankers.rs
+51 -2
View File
@@ -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<Buffer>,
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<RerankOutputSchemaFn>,
}
impl Reranker {
pub fn new(
rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
) -> napi::Result<Self> {
let rerank_hybrid = rerank_hybrid
.build_threadsafe_function()
.weak::<true>()
.build()?;
Ok(Self { rerank_hybrid })
let output_schema = output_schema
.map(|output_schema| {
output_schema
.build_threadsafe_function()
.weak::<true>()
.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<arrow_schema::SchemaRef> {
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<Buffer> = 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<RecordBatch> {
let mut reader = ipc_file_to_batches(buffer.to_vec()).default_error()?;
reader
rust/lancedb/src/query.rs
+174 -4
View File
@@ -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<dyn Reranker>) -> 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<SchemaRef> {
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<SendableRecordBatchStream> {
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::<Vec<_>>().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<SchemaRef> {
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<AtomicBool>,
}
#[async_trait::async_trait]
impl Reranker for SlowReranker {
async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
RRFReranker::default().output_schema(input).await
}
async fn rerank_hybrid(
&self,
query: &str,
vector_results: RecordBatch,
fts_results: RecordBatch,
) -> Result<RecordBatch> {
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]
rust/lancedb/src/rerankers.rs
+14
View File
@@ -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<SchemaRef> {
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
rust/lancedb/src/rerankers/rrf.rs
+16 -9
View File
@@ -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<SchemaRef> {
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)
}