diff --git a/Cargo.lock b/Cargo.lock
index 57493495b47..1cf9dbd36f0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5032,9 +5032,9 @@ dependencies = [
[[package]]
name = "lance-namespace-reqwest-client"
-version = "0.12.0"
+version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d8d23e54b1634d5bbb434f8dd33dc3c05f6e58d876a9a27b3b4aef58ddbe11af"
+checksum = "828f291764feda98898452bea4ce3756b73e5d5c7c67b62edda27456b7cc72c1"
dependencies = [
"reqwest 0.12.28",
"serde",
diff --git a/Cargo.toml b/Cargo.toml
index ddafb660946..d13da27cb85 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -73,7 +73,7 @@ lance-io = { version = "=13.0.0-beta.2", path = "./rust/lance-io", default-featu
lance-linalg = { version = "=13.0.0-beta.2", path = "./rust/lance-linalg" }
lance-namespace = { version = "=13.0.0-beta.2", path = "./rust/lance-namespace" }
lance-namespace-impls = { version = "=13.0.0-beta.2", path = "./rust/lance-namespace-impls" }
-lance-namespace-reqwest-client = "0.12.0"
+lance-namespace-reqwest-client = "0.13.0"
lance-select = { version = "=13.0.0-beta.2", path = "./rust/lance-select" }
lance-tokenizer = { version = "=13.0.0-beta.2", path = "./rust/lance-tokenizer" }
lance-table = { version = "=13.0.0-beta.2", path = "./rust/lance-table" }
diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock
index 8c75878c028..77846500018 100644
--- a/java/lance-jni/Cargo.lock
+++ b/java/lance-jni/Cargo.lock
@@ -4288,9 +4288,9 @@ dependencies = [
[[package]]
name = "lance-namespace-reqwest-client"
-version = "0.12.0"
+version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d8d23e54b1634d5bbb434f8dd33dc3c05f6e58d876a9a27b3b4aef58ddbe11af"
+checksum = "828f291764feda98898452bea4ce3756b73e5d5c7c67b62edda27456b7cc72c1"
dependencies = [
"reqwest 0.12.28",
"serde",
diff --git a/java/src/main/java/org/lance/ipc/Query.java b/java/src/main/java/org/lance/ipc/Query.java
index 215865310df..76e931dcacb 100644
--- a/java/src/main/java/org/lance/ipc/Query.java
+++ b/java/src/main/java/org/lance/ipc/Query.java
@@ -176,16 +176,14 @@ public Builder setK(int k) {
}
/**
- * Sets the number of probes to load and search.
+ * Sets the maximum number of probes to load and search.
*
- * This is a convenience method that sets both the minimum and maximum number of probes to
- * the same value.
+ *
The minimum remains unchanged, so the search may stop before reaching this value.
*
* @param nprobes The number of probes.
* @return The Builder instance for method chaining.
*/
public Builder setNprobes(int nprobes) {
- this.minimumNprobes = nprobes;
this.maximumNprobes = Optional.of(nprobes);
return this;
}
diff --git a/java/src/test/java/org/lance/JNITest.java b/java/src/test/java/org/lance/JNITest.java
index 94db13d6dea..f313e574172 100644
--- a/java/src/test/java/org/lance/JNITest.java
+++ b/java/src/test/java/org/lance/JNITest.java
@@ -55,6 +55,15 @@ public void testQuery() {
new Query.Builder().setColumn("column").setKey(new float[] {1.0f, 2.0f, 3.0f}).build();
assertEquals(ApproxMode.NORMAL, defaultQuery.getApproxMode());
+ Query nprobesQuery =
+ new Query.Builder()
+ .setColumn("column")
+ .setKey(new float[] {1.0f, 2.0f, 3.0f})
+ .setNprobes(20)
+ .build();
+ assertEquals(1, nprobesQuery.getMinimumNprobes());
+ assertEquals(Optional.of(20), nprobesQuery.getMaximumNprobes());
+
JniTestHelper.parseQuery(
Optional.of(
new Query.Builder()
diff --git a/python/Cargo.lock b/python/Cargo.lock
index 77a416d9a10..1d7c6c48b81 100644
--- a/python/Cargo.lock
+++ b/python/Cargo.lock
@@ -4435,9 +4435,9 @@ dependencies = [
[[package]]
name = "lance-namespace-reqwest-client"
-version = "0.12.0"
+version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d8d23e54b1634d5bbb434f8dd33dc3c05f6e58d876a9a27b3b4aef58ddbe11af"
+checksum = "828f291764feda98898452bea4ce3756b73e5d5c7c67b62edda27456b7cc72c1"
dependencies = [
"reqwest 0.12.28",
"serde",
diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
index 662e9307f59..ec22d1c3a55 100644
--- a/python/python/lance/dataset.py
+++ b/python/python/lance/dataset.py
@@ -8368,8 +8368,7 @@ def _build_vector_search_query(
metric: str, optional
The distance metric to use (e.g., "L2", "cosine", "dot", "hamming").
nprobes: int, optional
- The number of partitions to search. Sets both minimum_nprobes and
- maximum_nprobes to the same value.
+ The maximum number of partitions to search. The minimum remains unchanged.
minimum_nprobes: int, optional
The minimum number of partitions to search.
maximum_nprobes: int, optional
@@ -8439,15 +8438,6 @@ def _build_vector_search_query(
if maximum_nprobes is not None and int(maximum_nprobes) < 0:
raise ValueError(f"Maximum nprobes must be >= 0 but got {maximum_nprobes}")
- if nprobes is not None:
- if minimum_nprobes is not None or maximum_nprobes is not None:
- raise ValueError(
- "nprobes cannot be set in combination with minimum_nprobes or "
- "maximum_nprobes"
- )
- else:
- minimum_nprobes = nprobes
- maximum_nprobes = nprobes
if (
minimum_nprobes is not None
and maximum_nprobes is not None
@@ -8483,6 +8473,7 @@ def _build_vector_search_query(
"q": q,
"k": k,
"metric": metric,
+ "nprobes": nprobes,
"minimum_nprobes": minimum_nprobes,
"maximum_nprobes": maximum_nprobes,
"refine_factor": refine_factor,
diff --git a/python/python/tests/test_row_addr_prefilter.py b/python/python/tests/test_row_addr_prefilter.py
index 7d3381ebdce..aef647e150f 100644
--- a/python/python/tests/test_row_addr_prefilter.py
+++ b/python/python/tests/test_row_addr_prefilter.py
@@ -46,8 +46,8 @@ def _write(tmp_path: Path, with_index: bool = False) -> lance.LanceDataset:
)
ds = lance.write_dataset(tbl, str(tmp_path / "t.lance"), mode="overwrite")
if with_index:
- # IVF_FLAT with nprobes == num_partitions is exact, so the masked result
- # can be compared against brute force without recall slack.
+ # IVF_FLAT with both probe bounds set to num_partitions is exact, so the
+ # masked result can be compared against brute force without recall slack.
ds.create_index("vector", index_type="IVF_FLAT", num_partitions=4, metric="l2")
return ds
@@ -136,7 +136,13 @@ def test_knn_topk_is_computed_over_masked_rows(
query = np.zeros(DIM, dtype=np.float32)
got = ds.scanner(
- nearest={"column": "vector", "q": query, "k": 5, "nprobes": 4},
+ nearest={
+ "column": "vector",
+ "q": query,
+ "k": 5,
+ "minimum_nprobes": 4,
+ "maximum_nprobes": 4,
+ },
with_row_id=True,
row_addr_allowlist=serialize_row_addrs(allowed),
).to_table()
diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py
index ab5f26379b6..2ae7e16b003 100644
--- a/python/python/tests/test_vector_index.py
+++ b/python/python/tests/test_vector_index.py
@@ -2526,9 +2526,22 @@ def test_vector_index_with_nprobes(indexed_dataset):
}
).explain_plan()
- assert "minimum_nprobes=7" in res
+ assert "minimum_nprobes=1" in res
assert "maximum_nprobes=Some(7)" in res
+ res = indexed_dataset.scanner(
+ nearest={
+ "column": "vector",
+ "q": np.random.randn(128),
+ "k": 10,
+ "nprobes": 10,
+ "minimum_nprobes": 7,
+ }
+ ).explain_plan()
+
+ assert "minimum_nprobes=7" in res
+ assert "maximum_nprobes=Some(10)" in res
+
res = indexed_dataset.scanner(
nearest={
"column": "vector",
diff --git a/python/src/dataset.rs b/python/src/dataset.rs
index 6b442fb3027..e56931b3669 100644
--- a/python/src/dataset.rs
+++ b/python/src/dataset.rs
@@ -5587,7 +5587,6 @@ fn vector_query_params_from_dict(
&& !nprobes.is_none()
{
let extracted: usize = nprobes.extract()?;
- minimum_nprobes = extracted;
maximum_nprobes = Some(extracted);
}
diff --git a/rust/lance-namespace-impls/src/dir.rs b/rust/lance-namespace-impls/src/dir.rs
index e7d91d11e3d..b4ce2f3da43 100644
--- a/rust/lance-namespace-impls/src/dir.rs
+++ b/rust/lance-namespace-impls/src/dir.rs
@@ -1013,6 +1013,29 @@ impl TransactionAlteration {
}
}
+fn apply_probe_bounds(
+ scanner: &mut Scanner,
+ nprobes: Option,
+ minimum_nprobes: Option,
+ maximum_nprobes: Option,
+) -> Result<()> {
+ let parse_probe_count = |name: &str, value: i32| {
+ usize::try_from(value)
+ .map_err(|_| Error::invalid_input(format!("{name} must be non-negative")))
+ };
+
+ if let Some(nprobes) = nprobes {
+ scanner.nprobes(parse_probe_count("nprobes", nprobes)?);
+ }
+ if let Some(minimum_nprobes) = minimum_nprobes {
+ scanner.minimum_nprobes(parse_probe_count("minimum_nprobes", minimum_nprobes)?);
+ }
+ if let Some(maximum_nprobes) = maximum_nprobes {
+ scanner.maximum_nprobes(parse_probe_count("maximum_nprobes", maximum_nprobes)?);
+ }
+ Ok(())
+}
+
impl DirectoryNamespace {
fn manifest_ns_for_read(&self) -> Option<&Arc> {
self.write_manifest_ns
@@ -3170,6 +3193,8 @@ impl DirectoryNamespace {
prefilter: Option,
bypass_vector_index: Option,
nprobes: Option,
+ minimum_nprobes: Option,
+ maximum_nprobes: Option,
ef: Option,
refine_factor: Option,
distance_type: Option<&str>,
@@ -3240,9 +3265,7 @@ impl DirectoryNamespace {
})?;
// ANN parameters — must be applied after nearest().
- if let Some(n) = nprobes {
- scanner.nprobes(n.max(1) as usize);
- }
+ apply_probe_bounds(scanner, nprobes, minimum_nprobes, maximum_nprobes)?;
if let Some(e) = ef {
scanner.ef(e.max(1) as usize);
}
@@ -5029,6 +5052,8 @@ impl LanceNamespace for DirectoryNamespace {
request.query.prefilter,
request.query.bypass_vector_index,
request.query.nprobes,
+ request.query.minimum_nprobes,
+ request.query.maximum_nprobes,
request.query.ef,
request.query.refine_factor,
request.query.distance_type.as_deref(),
@@ -5071,6 +5096,8 @@ impl LanceNamespace for DirectoryNamespace {
request.prefilter,
request.bypass_vector_index,
request.nprobes,
+ request.minimum_nprobes,
+ request.maximum_nprobes,
request.ef,
request.refine_factor,
request.distance_type.as_deref(),
@@ -5443,10 +5470,12 @@ impl LanceNamespace for DirectoryNamespace {
scanner.distance_metric(metric);
}
- // Apply nprobes if specified (maps to minimum_nprobes, matching lancedb behavior)
- if let Some(nprobes) = request.nprobes {
- scanner.minimum_nprobes(nprobes as usize);
- }
+ apply_probe_bounds(
+ &mut scanner,
+ request.nprobes,
+ request.minimum_nprobes,
+ request.maximum_nprobes,
+ )?;
// Apply ef (HNSW search effort) if specified
if let Some(ef) = request.ef {
@@ -14336,6 +14365,77 @@ mod tests {
assert_eq!(total_rows, 2);
}
+ #[tokio::test]
+ async fn test_explain_vector_probe_fields_are_applied_independently() {
+ use lance_namespace::models::ExplainTableQueryPlanRequest;
+
+ let (namespace, temp_dir, table_id) = create_ns_with_vector_table().await;
+ let table_uri = format!("{}/vector_table.lance", temp_dir.to_str().unwrap());
+ let mut dataset = Dataset::open(&table_uri).await.unwrap();
+ dataset
+ .create_index(
+ &["vector"],
+ IndexType::Vector,
+ Some("vector_idx".to_string()),
+ &VectorIndexParams::ivf_flat(1, MetricType::L2),
+ false,
+ )
+ .await
+ .unwrap();
+ let vector = || {
+ Box::new(lance_namespace::models::QueryTableRequestVector {
+ single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
+ multi_vector: None,
+ })
+ };
+
+ let query = QueryTableRequest {
+ id: None,
+ k: 2,
+ vector: vector(),
+ nprobes: Some(20),
+ minimum_nprobes: Some(3),
+ maximum_nprobes: Some(10),
+ ..Default::default()
+ };
+ let mut request = ExplainTableQueryPlanRequest::new(query);
+ request.id = Some(table_id.clone());
+
+ let plan = namespace.explain_table_query_plan(request).await.unwrap();
+ assert!(plan.contains("minimum_nprobes=3"), "{plan}");
+ assert!(plan.contains("maximum_nprobes=Some(10)"), "{plan}");
+
+ let query = QueryTableRequest {
+ id: None,
+ k: 2,
+ vector: vector(),
+ nprobes: Some(0),
+ ..Default::default()
+ };
+ let mut request = ExplainTableQueryPlanRequest::new(query);
+ request.id = Some(table_id.clone());
+
+ let plan = namespace.explain_table_query_plan(request).await.unwrap();
+ assert!(plan.contains("minimum_nprobes=1"), "{plan}");
+ assert!(plan.contains("maximum_nprobes=Some(0)"), "{plan}");
+
+ let query = QueryTableRequest {
+ id: None,
+ k: 2,
+ vector: vector(),
+ nprobes: Some(-1),
+ ..Default::default()
+ };
+ let mut request = ExplainTableQueryPlanRequest::new(query);
+ request.id = Some(table_id);
+
+ let err = namespace
+ .explain_table_query_plan(request)
+ .await
+ .unwrap_err();
+ assert!(err.to_string().contains("nprobes must be non-negative"));
+ }
+
#[tokio::test]
async fn test_namespace_id() {
let (namespace, _temp_dir) = create_test_namespace().await;
diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs
index 50f79260583..65da17b8c18 100644
--- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs
+++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs
@@ -633,11 +633,9 @@ impl MemTableScanner {
/// Set the number of probes for IVF search.
///
- /// This is a convenience method that sets both minimum and maximum nprobes
- /// to the same value, guaranteeing exactly `n` partitions will be searched.
+ /// The minimum remains unchanged, so fewer than `n` partitions may be searched.
pub fn nprobes(&mut self, n: usize) -> &mut Self {
if let Some(ref mut q) = self.nearest {
- q.nprobes = n;
q.maximum_nprobes = Some(n);
} else {
log::warn!("nprobes is not set because nearest has not been called yet");
diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs
index dbd5e9e4a79..71e34cd28f2 100644
--- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs
+++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs
@@ -30,6 +30,7 @@ use super::planner::LsmScanPlanner;
use super::point_lookup::LsmPointLookupPlanner;
use super::projection::validate_projection_names;
use super::sstable_cache::{DatasetCache, SsTableWarmer};
+use super::vector_search::ProbeBounds;
use crate::dataset::Dataset;
use crate::dataset::mem_wal::util::derived_store_params;
use crate::session::Session;
@@ -45,8 +46,8 @@ struct LsmVectorQuery {
key: Arc,
/// Number of nearest neighbors to fetch per source before the global merge.
k: usize,
- /// Number of IVF partitions to probe on the base arm.
- nprobes: usize,
+ /// IVF partition probe bounds for indexed arms.
+ probe_bounds: ProbeBounds,
/// Re-rank base candidates with exact distances when set (refine factor is
/// treated as a boolean; the LSM merge needs exact base distances).
refine: bool,
@@ -186,10 +187,11 @@ fn key_to_fsl(key: &dyn Array, dim: i32) -> Result {
/// ```
///
/// The query-building methods mirror [`crate::dataset::scanner::Scanner`]:
-/// [`Self::nearest`] (+ [`Self::nprobes`] / [`Self::refine`] /
-/// [`Self::distance_metric`]) for vector search and [`Self::full_text_search`]
-/// for FTS are state setters, and [`Self::create_plan`] dispatches to the right
-/// planner — so an LSM read reads like a normal scan.
+/// [`Self::nearest`] (+ [`Self::nprobes`] / [`Self::minimum_nprobes`] /
+/// [`Self::maximum_nprobes`] / [`Self::refine`] / [`Self::distance_metric`])
+/// for vector search and [`Self::full_text_search`] for FTS are state setters,
+/// and [`Self::create_plan`] dispatches to the right planner — so an LSM read
+/// reads like a normal scan.
pub struct LsmScanner {
// Data sources
base: BaseSource,
@@ -451,7 +453,8 @@ impl LsmScanner {
/// [`crate::dataset::scanner::Scanner::nearest`]; the LSM path supports a
/// single Float32 query vector. When combined with an offset, the LSM path
/// fetches `k + offset` per source before applying the final page. Tune with
- /// [`Self::nprobes`], [`Self::refine`], and [`Self::distance_metric`].
+ /// [`Self::nprobes`], [`Self::minimum_nprobes`], [`Self::maximum_nprobes`],
+ /// [`Self::refine`], and [`Self::distance_metric`].
pub fn nearest(mut self, column: &str, key: &dyn Array, k: usize) -> Result {
if k == 0 {
return Err(Error::invalid_input("k must be positive".to_string()));
@@ -465,18 +468,40 @@ impl LsmScanner {
column: column.to_string(),
key: key.slice(0, key.len()),
k,
- nprobes: 1,
+ probe_bounds: ProbeBounds::default(),
refine: false,
metric_type: None,
});
Ok(self)
}
- /// Number of IVF partitions to probe on the base arm (default 1). No-op
- /// unless [`Self::nearest`] was called.
+ /// Search up to `nprobes` IVF partitions on indexed arms. No-op unless
+ /// [`Self::nearest`] was called. The minimum remains unchanged.
pub fn nprobes(mut self, nprobes: usize) -> Self {
if let Some(q) = self.nearest.as_mut() {
- q.nprobes = nprobes;
+ q.probe_bounds.maximum_nprobes = Some(nprobes);
+ }
+ self
+ }
+
+ /// Set the minimum number of IVF partitions to search on indexed arms.
+ ///
+ /// When unset, the underlying Lance scanner supplies its default. No-op
+ /// unless [`Self::nearest`] was called.
+ pub fn minimum_nprobes(mut self, minimum_nprobes: usize) -> Self {
+ if let Some(q) = self.nearest.as_mut() {
+ q.probe_bounds.minimum_nprobes = Some(minimum_nprobes);
+ }
+ self
+ }
+
+ /// Set the maximum number of IVF partitions to search on indexed arms.
+ ///
+ /// When unset, all partitions may be searched if needed. No-op unless
+ /// [`Self::nearest`] was called.
+ pub fn maximum_nprobes(mut self, maximum_nprobes: usize) -> Self {
+ if let Some(q) = self.nearest.as_mut() {
+ q.probe_bounds.maximum_nprobes = Some(maximum_nprobes);
}
self
}
@@ -607,10 +632,10 @@ impl LsmScanner {
let per_source_k = nearest.k.saturating_add(self.offset.unwrap_or(0));
let overfetch_factor = self.overfetch_factor.unwrap_or(1.0);
let plan = planner
- .plan_search(
+ .plan_search_with_probe_bounds(
&query_fsl,
per_source_k,
- nearest.nprobes,
+ nearest.probe_bounds,
self.projection.as_deref(),
nearest.refine,
overfetch_factor,
@@ -1245,6 +1270,52 @@ mod tests {
]))
}
+ #[test]
+ fn vector_probe_settings_preserve_adaptive_defaults() {
+ use arrow_array::Float32Array;
+ use arrow_schema::{DataType, Field};
+
+ let schema = pk_schema_with(Field::new(
+ "vector",
+ DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4),
+ false,
+ ));
+ let new_scanner = || {
+ LsmScanner::without_base_table(
+ schema.clone(),
+ "memory://",
+ vec![],
+ vec!["id".to_string()],
+ )
+ .nearest(
+ "vector",
+ &Float32Array::from(vec![0.0f32, 1.0, 2.0, 3.0]),
+ 1,
+ )
+ .unwrap()
+ };
+
+ let scanner = new_scanner();
+ let query = scanner.nearest.as_ref().unwrap();
+ assert_eq!(query.probe_bounds.minimum_nprobes, None);
+ assert_eq!(query.probe_bounds.maximum_nprobes, None);
+
+ let scanner = new_scanner().nprobes(20);
+ let query = scanner.nearest.as_ref().unwrap();
+ assert_eq!(query.probe_bounds.minimum_nprobes, None);
+ assert_eq!(query.probe_bounds.maximum_nprobes, Some(20));
+
+ let scanner = new_scanner().minimum_nprobes(20);
+ let query = scanner.nearest.as_ref().unwrap();
+ assert_eq!(query.probe_bounds.minimum_nprobes, Some(20));
+ assert_eq!(query.probe_bounds.maximum_nprobes, None);
+
+ let scanner = new_scanner().maximum_nprobes(20);
+ let query = scanner.nearest.as_ref().unwrap();
+ assert_eq!(query.probe_bounds.minimum_nprobes, None);
+ assert_eq!(query.probe_bounds.maximum_nprobes, Some(20));
+ }
+
/// `LsmScanner::nearest(..).create_plan()` must route through the vector
/// planner (exercising the Scanner-aligned facade, `key_to_fsl`, and
/// `vector_dim`) and surface the in-memory match.
diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs
index c808a4e3a5a..9c3c97d6bc0 100644
--- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs
+++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs
@@ -36,6 +36,21 @@ use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable};
use crate::session::Session;
use lance_io::object_store::ObjectStoreParams;
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub(super) struct ProbeBounds {
+ pub minimum_nprobes: Option,
+ pub maximum_nprobes: Option,
+}
+
+impl ProbeBounds {
+ pub(super) fn maximum(nprobes: usize) -> Self {
+ Self {
+ minimum_nprobes: None,
+ maximum_nprobes: Some(nprobes),
+ }
+ }
+}
+
/// Plans vector search queries over LSM data.
///
/// Each source is independently newest-per-PK before the union — the active
@@ -217,7 +232,7 @@ impl LsmVectorSearchPlanner {
///
/// * `query_vector` - Query vector for KNN search
/// * `k` - Number of nearest neighbors to return
- /// * `nprobes` - Number of IVF partitions to search (for IVF-based indexes)
+ /// * `nprobes` - Maximum number of IVF partitions to search (for IVF-based indexes)
/// * `projection` - Columns to include in output (None = all columns)
/// * `refine_base_table` - When true, the base-table arm re-ranks its
/// candidates with exact distances (refine factor 1). Useful when the base
@@ -240,7 +255,6 @@ impl LsmVectorSearchPlanner {
///
/// An execution plan that returns the top-K nearest neighbors across all
/// LSM levels, with stale results filtered out.
- #[instrument(name = "lsm_vector_search", level = "info", skip_all, fields(k, nprobes, vector_column = %self.vector_column, distance_type = ?self.distance_type))]
pub async fn plan_search(
&self,
query_vector: &FixedSizeListArray,
@@ -249,12 +263,51 @@ impl LsmVectorSearchPlanner {
projection: Option<&[String]>,
refine_base_table: bool,
overfetch_factor: f64,
+ ) -> Result> {
+ if nprobes == 0 {
+ return Err(Error::invalid_input("nprobes must be positive".to_string()));
+ }
+ self.plan_search_with_probe_bounds(
+ query_vector,
+ k,
+ ProbeBounds::maximum(nprobes),
+ projection,
+ refine_base_table,
+ overfetch_factor,
+ )
+ .await
+ }
+
+ #[instrument(name = "lsm_vector_search", level = "info", skip_all, fields(k, minimum_nprobes = ?probe_bounds.minimum_nprobes, maximum_nprobes = ?probe_bounds.maximum_nprobes, vector_column = %self.vector_column, distance_type = ?self.distance_type))]
+ pub(super) async fn plan_search_with_probe_bounds(
+ &self,
+ query_vector: &FixedSizeListArray,
+ k: usize,
+ probe_bounds: ProbeBounds,
+ projection: Option<&[String]>,
+ refine_base_table: bool,
+ overfetch_factor: f64,
) -> Result> {
if k == 0 {
return Err(Error::invalid_input("k must be positive".to_string()));
}
- if nprobes == 0 {
- return Err(Error::invalid_input("nprobes must be positive".to_string()));
+ if probe_bounds.minimum_nprobes == Some(0) {
+ return Err(Error::invalid_input(
+ "minimum_nprobes must be positive".to_string(),
+ ));
+ }
+ if probe_bounds.maximum_nprobes == Some(0) {
+ return Err(Error::invalid_input(
+ "maximum_nprobes must be positive".to_string(),
+ ));
+ }
+ if let (Some(minimum_nprobes), Some(maximum_nprobes)) =
+ (probe_bounds.minimum_nprobes, probe_bounds.maximum_nprobes)
+ && minimum_nprobes > maximum_nprobes
+ {
+ return Err(Error::invalid_input(format!(
+ "minimum_nprobes ({minimum_nprobes}) must not exceed maximum_nprobes ({maximum_nprobes})"
+ )));
}
let sources = self.collector.collect()?;
@@ -319,7 +372,7 @@ impl LsmVectorSearchPlanner {
source,
query_vector,
*fetch_k,
- nprobes,
+ probe_bounds,
projection,
*is_base && refine_base,
))
@@ -442,7 +495,7 @@ impl LsmVectorSearchPlanner {
source: &LsmDataSource,
query_vector: &FixedSizeListArray,
k: usize,
- nprobes: usize,
+ probe_bounds: ProbeBounds,
projection: Option<&[String]>,
refine: bool,
) -> Result> {
@@ -471,7 +524,12 @@ impl LsmVectorSearchPlanner {
let query_arr = single_query_array(query_vector);
scanner.nearest(&self.vector_column, query_arr.as_ref(), k)?;
scanner.distance_range(self.distance_range.0, self.distance_range.1);
- scanner.nprobes(nprobes);
+ if let Some(minimum_nprobes) = probe_bounds.minimum_nprobes {
+ scanner.minimum_nprobes(minimum_nprobes);
+ }
+ if let Some(maximum_nprobes) = probe_bounds.maximum_nprobes {
+ scanner.maximum_nprobes(maximum_nprobes);
+ }
scanner.distance_metric(self.distance_type);
if let Some(ef) = self.ef {
scanner.ef(ef);
@@ -508,7 +566,12 @@ impl LsmVectorSearchPlanner {
let query_arr = single_query_array(query_vector);
scanner.nearest(&self.vector_column, query_arr.as_ref(), k)?;
scanner.distance_range(self.distance_range.0, self.distance_range.1);
- scanner.nprobes(nprobes);
+ if let Some(minimum_nprobes) = probe_bounds.minimum_nprobes {
+ scanner.minimum_nprobes(minimum_nprobes);
+ }
+ if let Some(maximum_nprobes) = probe_bounds.maximum_nprobes {
+ scanner.maximum_nprobes(maximum_nprobes);
+ }
scanner.distance_metric(self.distance_type);
if let Some(ef) = self.ef {
scanner.ef(ef);
@@ -541,7 +604,12 @@ impl LsmVectorSearchPlanner {
}
scanner.nearest(&self.vector_column, query_vector, k)?;
scanner.distance_range(self.distance_range.0, self.distance_range.1);
- scanner.nprobes(nprobes);
+ if let Some(minimum_nprobes) = probe_bounds.minimum_nprobes {
+ scanner.minimum_nprobes(minimum_nprobes);
+ }
+ if let Some(maximum_nprobes) = probe_bounds.maximum_nprobes {
+ scanner.maximum_nprobes(maximum_nprobes);
+ }
scanner.distance_metric(self.distance_type);
if let Some(ef) = self.ef {
scanner.ef(ef);
@@ -792,6 +860,45 @@ mod tests {
err.to_string().contains("nprobes must be positive"),
"expected nprobes validation error, got {err}"
);
+
+ let err = planner
+ .plan_search_with_probe_bounds(
+ &query,
+ 1,
+ ProbeBounds {
+ minimum_nprobes: None,
+ maximum_nprobes: Some(0),
+ },
+ None,
+ false,
+ 1.0,
+ )
+ .await
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("maximum_nprobes must be positive"),
+ "expected maximum_nprobes validation error, got {err}"
+ );
+
+ let err = planner
+ .plan_search_with_probe_bounds(
+ &query,
+ 1,
+ ProbeBounds {
+ minimum_nprobes: Some(2),
+ maximum_nprobes: Some(1),
+ },
+ None,
+ false,
+ 1.0,
+ )
+ .await
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("minimum_nprobes (2) must not exceed maximum_nprobes (1)"),
+ "expected probe-bound ordering error, got {err}"
+ );
}
#[tokio::test]
diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs
index 28f4d341470..2a6a1f8c75a 100644
--- a/rust/lance/src/dataset/scanner.rs
+++ b/rust/lance/src/dataset/scanner.rs
@@ -2047,13 +2047,12 @@ impl Scanner {
self
}
- /// Configures how many partititions will be searched in the vector index.
+ /// Configures the maximum number of partitions searched in the vector index.
///
- /// This method is a convenience method that sets both [Self::minimum_nprobes] and
- /// [Self::maximum_nprobes] to the same value.
+ /// The minimum remains unchanged, so the search may stop before reaching this
+ /// value when enough results have been found.
pub fn nprobes(&mut self, n: usize) -> &mut Self {
if let Some(q) = self.nearest.as_mut() {
- q.minimum_nprobes = n;
q.maximum_nprobes = Some(n);
} else {
log::warn!("nprobes is not set because nearest has not been called yet");
@@ -2061,14 +2060,10 @@ impl Scanner {
self
}
- /// Configures how many partititions will be searched in the vector index.
- ///
- /// This method is a convenience method that sets both [Self::minimum_nprobes] and
- /// [Self::maximum_nprobes] to the same value.
+ /// Configures the maximum number of partitions searched in the vector index.
#[deprecated(note = "Use nprobes instead")]
pub fn nprobs(&mut self, n: usize) -> &mut Self {
if let Some(q) = self.nearest.as_mut() {
- q.minimum_nprobes = n;
q.maximum_nprobes = Some(n);
} else {
log::warn!("nprobes is not set because nearest has not been called yet");
@@ -9324,7 +9319,7 @@ mod test {
k: usize,
use_index: bool,
distance_range: Option<(Option, Option)>,
- nprobes: Option,
+ fixed_nprobes: Option,
) {
let query_count = query_values.len() / 32;
assert_eq!(batch.num_rows(), query_count * k);
@@ -9338,8 +9333,8 @@ mod test {
// Pin nprobes to match the batch query: the single-query indexed path
// otherwise adaptively expands nprobes, which would make equivalence
// depend on data distribution rather than be guaranteed.
- if let Some(nprobes) = nprobes {
- scan.nprobes(nprobes);
+ if let Some(nprobes) = fixed_nprobes {
+ scan.minimum_nprobes(nprobes).maximum_nprobes(nprobes);
}
if let Some((lower, upper)) = distance_range {
scan.distance_range(lower, upper);
@@ -9799,7 +9794,7 @@ mod test {
// merged across multiple partitions and the batch result is
// deterministically equivalent to repeated single-query search (which
// would otherwise adaptively expand nprobes).
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
scan.project(&["i"]).unwrap();
let plan = scan.explain_plan(false).await.unwrap();
@@ -9821,7 +9816,7 @@ mod test {
// The batch node loads each probed partition once and scores every query
// that probes it, so it must report the *distinct* partitions read: with
- // 2 partitions and nprobes(2), both queries probe both partitions, so the
+ // 2 partitions and fixed bounds of 2, both queries probe both partitions, so the
// union is 2 -- not the per-query sum (2 queries x 2 = 4), and never 0
// (which is what a dropped metric would show). This guards the observed
// `partitions_searched` against silently regressing to either.
@@ -9851,7 +9846,8 @@ mod test {
.scan()
.nearest("vec", &queries, 2)
.unwrap()
- .nprobes(2)
+ .minimum_nprobes(2)
+ .maximum_nprobes(2)
.distance_range(Some(1.0), None)
.project(&["i"])
.unwrap()
@@ -9918,7 +9914,7 @@ mod test {
let k = 5;
let mut scan = dataset.scan();
scan.nearest("vec", &queries, k).unwrap();
- scan.nprobes(1);
+ scan.minimum_nprobes(1).maximum_nprobes(1);
scan.project(&["i"]).unwrap();
let plan = scan.explain_plan(false).await.unwrap();
@@ -10057,8 +10053,8 @@ mod test {
.await;
}
- /// `nprobes(0)` is not rejected by the query builder, so `minimum_nprobes ==
- /// maximum_nprobes == 0` slips past the fixed-nprobes gate. The single-query
+ /// `nprobes(0)` is not rejected by the query builder, so `maximum_nprobes == 0`
+ /// reaches the adaptive path. The single-query
/// path then probes nothing and returns an empty result, whereas the batch
/// node would clamp `nprobes` up to one partition — a silent divergence. The
/// scanner must fall back so the per-query loop defines the semantics of
@@ -10139,7 +10135,7 @@ mod test {
// batch-eligible, so the mask is the only thing that forces the fallback.
let mut unmasked = dataset.scan();
unmasked.nearest("vec", &queries, k).unwrap();
- unmasked.nprobes(2);
+ unmasked.minimum_nprobes(2).maximum_nprobes(2);
unmasked.project(&["i"]).unwrap();
let unmasked_plan = unmasked.explain_plan(false).await.unwrap();
assert!(
@@ -10157,7 +10153,7 @@ mod test {
let mut scan = dataset.scan();
scan.nearest("vec", &queries, k).unwrap();
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter(
allow.iter().copied(),
)));
@@ -10216,7 +10212,7 @@ mod test {
let mut scan = dataset.scan();
scan.nearest("vec", &queries, k).unwrap();
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
// Request both indexed fragments but select only the segment covering
// fragment 0; fragment 1 is covered only by the unselected segment.
scan.with_fragments(vec![fragments[0].clone(), fragments[1].clone()]);
@@ -10310,7 +10306,8 @@ mod test {
scan.nearest("vec", &queries, k).unwrap();
// Probe every partition so both paths are exact regardless of centroid
// proximity, and so the batch spans multiple streaming chunks.
- scan.nprobes(num_partitions);
+ scan.minimum_nprobes(num_partitions)
+ .maximum_nprobes(num_partitions);
scan.project(&["i"]).unwrap();
let plan = scan.explain_plan(false).await.unwrap();
@@ -10347,7 +10344,7 @@ mod test {
let mut scan = dataset.scan();
scan.nearest("vec", &queries, k).unwrap();
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
scan.project(&["i"]).unwrap();
let plan = scan.explain_plan(false).await.unwrap();
@@ -10393,7 +10390,7 @@ mod test {
let mut scan = dataset.scan();
scan.nearest("vec", &queries, 2).unwrap();
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
scan.project(&["i"]).unwrap();
let plan = scan.explain_plan(false).await.unwrap();
@@ -10423,7 +10420,7 @@ mod test {
let mut scan = dataset.scan();
scan.nearest("vec", &queries, k).unwrap();
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
scan.filter("i > 100").unwrap();
scan.prefilter(true);
scan.project(&["i"]).unwrap();
@@ -10455,7 +10452,8 @@ mod test {
.scan()
.nearest("vec", &query, k)
.unwrap()
- .nprobes(2)
+ .minimum_nprobes(2)
+ .maximum_nprobes(2)
.filter("i > 100")
.unwrap()
.prefilter(true)
@@ -10506,7 +10504,7 @@ mod test {
let k = 3;
let mut scan = dataset.scan();
scan.nearest("vec", &queries, k).unwrap();
- scan.nprobes(2);
+ scan.minimum_nprobes(2).maximum_nprobes(2);
scan.project(&["i"]).unwrap();
let plan = scan.explain_plan(false).await.unwrap();
@@ -16285,6 +16283,31 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\")
);
}
+ #[tokio::test]
+ async fn test_knn_probe_setters_preserve_independent_fields() {
+ let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false)
+ .await
+ .unwrap();
+ let query_vector = Float32Array::from(vec![0.0; 32]);
+ let mut scanner = test_ds.dataset.scan();
+ scanner.nearest("vec", &query_vector, 5).unwrap();
+
+ scanner.nprobes(20);
+ let query = scanner.nearest_mut().unwrap();
+ assert_eq!(query.minimum_nprobes, 1);
+ assert_eq!(query.maximum_nprobes, Some(20));
+
+ scanner.minimum_nprobes(5);
+ let query = scanner.nearest_mut().unwrap();
+ assert_eq!(query.minimum_nprobes, 5);
+ assert_eq!(query.maximum_nprobes, Some(20));
+
+ scanner.maximum_nprobes(10);
+ let query = scanner.nearest_mut().unwrap();
+ assert_eq!(query.minimum_nprobes, 5);
+ assert_eq!(query.maximum_nprobes, Some(10));
+ }
+
#[tokio::test]
async fn test_ivf_pq_query_parallelism_returns_same_results() {
let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false)
diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs
index 9e2d2078307..a10a2319d07 100644
--- a/rust/lance/src/dataset/tests/dataset_migrations.rs
+++ b/rust/lance/src/dataset/tests/dataset_migrations.rs
@@ -285,7 +285,8 @@ async fn test_v0_8_14_invalid_index_fragment_bitmap(
let batches = scan
.nearest("vector", &query_vec, 2000)
.unwrap()
- .nprobes(4)
+ .minimum_nprobes(4)
+ .maximum_nprobes(4)
.prefilter(true)
.try_into_stream()
.await
diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs
index 6695940ccfb..d4fb926cdce 100644
--- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs
+++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs
@@ -652,7 +652,8 @@ async fn test_vector_batch_falls_back_on_overlay(#[values(false, true)] stable_r
scanner
.nearest("vec", queries.as_ref(), 3)
.unwrap()
- .nprobes(1)
+ .minimum_nprobes(1)
+ .maximum_nprobes(1)
.project(&["id"])
.unwrap();
let plan = scanner.explain_plan(false).await.unwrap();
@@ -667,7 +668,8 @@ async fn test_vector_batch_falls_back_on_overlay(#[values(false, true)] stable_r
scanner
.nearest("vec", queries.as_ref(), 3)
.unwrap()
- .nprobes(1)
+ .minimum_nprobes(1)
+ .maximum_nprobes(1)
.project(&["id"])
.unwrap();
let plan = scanner.explain_plan(false).await.unwrap();
diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs
index 415958a438a..c1e07192bf6 100644
--- a/rust/lance/src/index.rs
+++ b/rust/lance/src/index.rs
@@ -11397,7 +11397,8 @@ mod tests {
.scan()
.nearest("vector", &Float32Array::from(query_vector), 10)
.unwrap()
- .nprobes(2)
+ .minimum_nprobes(2)
+ .maximum_nprobes(2)
.try_into_batch()
.await
.unwrap();
diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs
index 245c861de72..40172a40f5a 100644
--- a/rust/lance/src/index/append.rs
+++ b/rust/lance/src/index/append.rs
@@ -1502,7 +1502,8 @@ mod tests {
.unwrap()
.nearest("vector", query, 1)
.unwrap()
- .nprobes(num_probes)
+ .minimum_nprobes(num_probes)
+ .maximum_nprobes(num_probes)
.refine(1)
.try_into_batch()
.await
@@ -2541,7 +2542,8 @@ mod tests {
.unwrap()
.nearest("vector", array.value(0).as_primitive::(), 2)
.unwrap()
- .nprobes(2)
+ .minimum_nprobes(2)
+ .maximum_nprobes(2)
.refine(1);
let fanout_plan = fanout_scanner.explain_plan(true).await.unwrap();
assert!(
@@ -2563,7 +2565,8 @@ mod tests {
.unwrap()
.nearest("vector", array.value(0).as_primitive::(), 1)
.unwrap()
- .nprobes(2)
+ .minimum_nprobes(2)
+ .maximum_nprobes(2)
.refine(1)
.with_index_segments(vec![segment.uuid])
.unwrap();
@@ -2867,7 +2870,8 @@ mod tests {
.unwrap()
.nearest("vector", &query, 5)
.unwrap()
- .nprobes(1)
+ .minimum_nprobes(1)
+ .maximum_nprobes(1)
.try_into_batch()
.await
.unwrap();
diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs
index d4f3134c2d4..afecd76081d 100644
--- a/rust/lance/src/index/vector/ivf/v2.rs
+++ b/rust/lance/src/index/vector/ivf/v2.rs
@@ -5697,7 +5697,8 @@ mod tests {
.scan()
.nearest("vector", query.as_primitive::(), PQ_MATRIX_K)
.unwrap()
- .nprobes(nlist)
+ .minimum_nprobes(nlist)
+ .maximum_nprobes(nlist)
.with_row_id()
.try_into_batch()
.await
@@ -7226,7 +7227,8 @@ mod tests {
.scan()
.nearest(vector_column, query.as_primitive::(), k)
.unwrap()
- .nprobes(nlist)
+ .minimum_nprobes(nlist)
+ .maximum_nprobes(nlist)
.with_row_id()
.try_into_batch()
.await
@@ -8674,7 +8676,8 @@ mod tests {
.with_row_id()
.nearest("vector", &q, 10)
.unwrap()
- .nprobes(4)
+ .minimum_nprobes(4)
+ .maximum_nprobes(4)
.project(&["_rowid"])
.unwrap()
.try_into_batch()
@@ -8721,7 +8724,8 @@ mod tests {
.scan()
.nearest("vector", &q, 10)
.unwrap()
- .nprobes(4)
+ .minimum_nprobes(4)
+ .maximum_nprobes(4)
.project(&["_rowid"])
.unwrap()
.try_into_batch()