rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-05-19 04:50:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add maximum and minimum nprobes properties (#2430)

Browse Source 
This exposes the maximum_nprobes and minimum_nprobes feature that was
added in https://github.com/lancedb/lance/pull/3903

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added support for specifying minimum and maximum probe counts in
vector search queries, allowing finer control over search behavior.
- Users can now independently set minimum and maximum probes for vector
and hybrid queries via new methods and parameters in Python, Node.js,
and Rust APIs.

- **Bug Fixes**
- Improved parameter validation to ensure correct usage of minimum and
maximum probe values.

- **Tests**
- Expanded test coverage to validate correct handling, serialization,
and error cases for the new probe parameters.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
Weston Pace
2025-06-13 15:18:29 -07:00
committed by GitHub
parent fec8d58f06
commit 59b57e30ed
12 changed files with 505 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
rust/lancedb/src/query.rs
View File

@@ -796,8 +796,10 @@ pub struct VectorQueryRequest {
pub column: Option<String>,
/// The vector(s) to search for
pub query_vector: Vec<Arc<dyn Array>>,
/// The number of partitions to search
pub nprobes: usize,
/// The minimum number of partitions to search
pub minimum_nprobes: usize,
/// The maximum number of partitions to search
pub maximum_nprobes: Option<usize>,
/// The lower bound (inclusive) of the distance to search for.
pub lower_bound: Option<f32>,
/// The upper bound (exclusive) of the distance to search for.
@@ -819,7 +821,8 @@ impl Default for VectorQueryRequest {
base: QueryRequest::default(),
column: None,
query_vector: Vec::new(),
nprobes: 20,
minimum_nprobes: 20,
maximum_nprobes: Some(20),
lower_bound: None,
upper_bound: None,
ef: None,
@@ -925,11 +928,75 @@ impl VectorQuery {
/// For best results we recommend tuning this parameter with a benchmark against
/// your actual data to find the smallest possible value that will still give
/// you the desired recall.
///
/// This method sets both the minimum and maximum number of partitions to search.
/// For more fine-grained control see [`VectorQuery::minimum_nprobes`] and
/// [`VectorQuery::maximum_nprobes`].
pub fn nprobes(mut self, nprobes: usize) -> Self {
self.request.nprobes = nprobes;
self.request.minimum_nprobes = nprobes;
self.request.maximum_nprobes = Some(nprobes);
self
}
/// Set the minimum number of partitions to search
///
/// This argument is only used when the vector column has an IVF PQ index.
/// If there is no index then this value is ignored.
///
/// See [`VectorQuery::nprobes`] for more details.
///
/// These partitions will be searched on every indexed vector query.
///
/// Will return an error if the value is not greater than 0 or if maximum_nprobes
/// has been set and is less than the minimum_nprobes.
pub fn minimum_nprobes(mut self, minimum_nprobes: usize) -> Result<Self> {
if minimum_nprobes == 0 {
return Err(Error::InvalidInput {
message: "minimum_nprobes must be greater than 0".to_string(),
});
}
if let Some(maximum_nprobes) = self.request.maximum_nprobes {
if minimum_nprobes > maximum_nprobes {
return Err(Error::InvalidInput {
message: "minimum_nprobes must be less or equal to maximum_nprobes".to_string(),
});
}
}
self.request.minimum_nprobes = minimum_nprobes;
Ok(self)
}
/// Set the maximum number of partitions to search
///
/// This argument is only used when the vector column has an IVF PQ index.
/// If there is no index then this value is ignored.
///
/// See [`VectorQuery::nprobes`] for more details.
///
/// If this value is greater than minimum_nprobes then the excess partitions will
/// only be searched if the initial search does not return enough results.
///
/// This can be useful when there is a narrow filter to allow these queries to
/// spend more time searching and avoid potential false negatives.
///
/// Set to None to search all partitions, if needed, to satsify the limit
pub fn maximum_nprobes(mut self, maximum_nprobes: Option<usize>) -> Result<Self> {
if let Some(maximum_nprobes) = maximum_nprobes {
if maximum_nprobes == 0 {
return Err(Error::InvalidInput {
message: "maximum_nprobes must be greater than 0".to_string(),
});
}
if maximum_nprobes < self.request.minimum_nprobes {
return Err(Error::InvalidInput {
message: "maximum_nprobes must be greater than minimum_nprobes".to_string(),
});
}
}
self.request.maximum_nprobes = maximum_nprobes;
Ok(self)
}
/// Set the distance range for vector search,
/// only rows with distances in the range [lower_bound, upper_bound) will be returned
pub fn distance_range(mut self, lower_bound: Option<f32>, upper_bound: Option<f32>) -> Self {
@@ -1208,7 +1275,8 @@ mod tests {
);
assert_eq!(query.request.base.limit.unwrap(), 100);
assert_eq!(query.request.base.offset.unwrap(), 1);
assert_eq!(query.request.nprobes, 1000);
assert_eq!(query.request.minimum_nprobes, 1000);
assert_eq!(query.request.maximum_nprobes, Some(1000));
assert!(query.request.use_index);
assert_eq!(query.request.distance_type, Some(DistanceType::Cosine));
assert_eq!(query.request.refine_factor, Some(999));

18
rust/lancedb/src/remote/table.rs
View File

@@ -32,6 +32,7 @@ use lance::dataset::{ColumnAlteration, NewColumnTransform, Version};
use lance_datafusion::exec::{execute_plan, OneShotExec};
use reqwest::{RequestBuilder, Response};
use serde::{Deserialize, Serialize};
use serde_json::Number;
use std::collections::HashMap;
use std::io::Cursor;
use std::pin::Pin;
@@ -438,7 +439,18 @@ impl<S: HttpSend> RemoteTable<S> {
// Apply general parameters, before we dispatch based on number of query vectors.
body["distance_type"] = serde_json::json!(query.distance_type.unwrap_or_default());
body["nprobes"] = query.nprobes.into();
// In 0.23.1 we migrated from `nprobes` to `minimum_nprobes` and `maximum_nprobes`.
// Old client / new server: since minimum_nprobes is missing, fallback to nprobes
// New client / old server: old server will only see nprobes, make sure to set both
// nprobes and minimum_nprobes
// New client / new server: since minimum_nprobes is present, server can ignore nprobes
body["nprobes"] = query.minimum_nprobes.into();
body["minimum_nprobes"] = query.minimum_nprobes.into();
if let Some(maximum_nprobes) = query.maximum_nprobes {
body["maximum_nprobes"] = maximum_nprobes.into();
} else {
body["maximum_nprobes"] = serde_json::Value::Number(Number::from_u128(0).unwrap())
}
body["lower_bound"] = query.lower_bound.into();
body["upper_bound"] = query.upper_bound.into();
body["ef"] = query.ef.into();
@@ -2075,6 +2087,8 @@ mod tests {
"prefilter": true,
"distance_type": "l2",
"nprobes": 20,
"minimum_nprobes": 20,
"maximum_nprobes": 20,
"lower_bound": Option::<f32>::None,
"upper_bound": Option::<f32>::None,
"k": 10,
@@ -2175,6 +2189,8 @@ mod tests {
"bypass_vector_index": true,
"columns": ["a", "b"],
"nprobes": 12,
"minimum_nprobes": 12,
"maximum_nprobes": 12,
"lower_bound": Option::<f32>::None,
"upper_bound": Option::<f32>::None,
"ef": Option::<usize>::None,

5
rust/lancedb/src/table.rs
View File

@@ -2354,12 +2354,15 @@ impl BaseTable for NativeTable {
query.base.limit.unwrap_or(DEFAULT_TOP_K),
)?;
}
scanner.minimum_nprobes(query.minimum_nprobes);
if let Some(maximum_nprobes) = query.maximum_nprobes {
scanner.maximum_nprobes(maximum_nprobes);
}
}
scanner.limit(
query.base.limit.map(|limit| limit as i64),
query.base.offset.map(|offset| offset as i64),
)?;
scanner.nprobs(query.nprobes);
if let Some(ef) = query.ef {
scanner.ef(ef);
}