Skip to content

perf!: dynamic pruning for vector search #4773

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BubbleCal merged 17 commits into from
Nov 10, 2025
Merged

perf!: dynamic pruning for vector search #4773

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
a44159c
return dists
BubbleCal Sep 18, 2025
c7e7335
perf: dynamic pruning for vector search
BubbleCal Sep 19, 2025
4510b56
Merge branch 'main' of https://github.com/lancedb/lance into dynamic-…
BubbleCal Sep 19, 2025
7326200
fmt
BubbleCal Sep 19, 2025
db5f098
fix ut
BubbleCal Sep 19, 2025
63a5736
Merge branch 'main' of https://github.com/lancedb/lance into dynamic-…
BubbleCal Nov 5, 2025
c9c3d22
fix
BubbleCal Nov 5, 2025
4aa97e9
fix
BubbleCal Nov 5, 2025
22f33d8
Merge branch 'main' of https://github.com/lancedb/lance into dynamic-…
BubbleCal Nov 6, 2025
b25aa20
fmt
BubbleCal Nov 6, 2025
6d832ef
fix
BubbleCal Nov 6, 2025
8ee0ba7
fix ut
BubbleCal Nov 6, 2025
2d55da3
fix ut
BubbleCal Nov 6, 2025
5376315
Merge branch 'main' of https://github.com/lancedb/lance into dynamic-…
BubbleCal Nov 7, 2025
2398244
fix
BubbleCal Nov 7, 2025
f492fd3
more comments
BubbleCal Nov 10, 2025
d767012
Merge branch 'main' of https://github.com/lancedb/lance into dynamic-…
BubbleCal Nov 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion java/src/main/java/com/lancedb/lance/ipc/Query.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ public static class Builder {
private String column;
private float[] key;
private int k = 10;
private int minimumNprobes = 20;
private int minimumNprobes = 1;
private Optional<Integer> maximumNprobes = Optional.empty();
private Optional<Integer> ef = Optional.empty();
private Optional<Integer> refineFactor = Optional.empty();
Expand Down
4 changes: 2 additions & 2 deletions python/python/lance/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -717,7 +717,7 @@ def scanner(
"column": <embedding col name>,
"q": <query vector as pa.Float32Array>,
"k": 10,
"minimum_nprobes": 20,
"minimum_nprobes": 1,
"maximum_nprobes": 50,
"refine_factor": 1
}
Expand Down Expand Up @@ -980,7 +980,7 @@ def to_table(
"q": <query vector as pa.Float32Array>,
"k": 10,
"metric": "cosine",
"minimum_nprobes": 20,
"minimum_nprobes": 1,
"maximum_nprobes": 50,
"refine_factor": 1
}
Expand Down
2 changes: 1 addition & 1 deletion python/python/tests/test_vector_index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1092,7 +1092,7 @@ def query_index(ds, ntimes, q=None):
nearest={
"column": "vector",
"q": q if q is not None else rng.standard_normal(ndim),
"minimum_nprobes": 1,
"nprobes": 20,
},
)

Expand Down
23 changes: 14 additions & 9 deletions python/src/dataset.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ pub mod commit;
pub mod optimize;
pub mod stats;

const DEFAULT_NPROBES: usize = 20;
const DEFAULT_NPROBES: usize = 1;
const LANCE_COMMIT_MESSAGE_KEY: &str = "__lance_commit_message";

fn convert_reader(reader: &Bound<PyAny>) -> PyResult<Box<dyn RecordBatchReader + Send>> {
Expand Down Expand Up @@ -986,8 +986,9 @@ impl Dataset {

if let Some(nprobes) = nearest.get_item("nprobes")? {
if !nprobes.is_none() {
minimum_nprobes = nprobes.extract()?;
maximum_nprobes = Some(minimum_nprobes);
let extracted: usize = nprobes.extract()?;
minimum_nprobes = extracted;
maximum_nprobes = Some(extracted);
}
}

Expand All @@ -1003,18 +1004,22 @@ impl Dataset {
}
}

if minimum_nprobes > maximum_nprobes.unwrap_or(usize::MAX) {
return Err(PyValueError::new_err(
"minimum_nprobes must be <= maximum_nprobes",
));
if let Some(maximum_nprobes) = maximum_nprobes {
if minimum_nprobes > maximum_nprobes {
return Err(PyValueError::new_err(
"minimum_nprobes must be <= maximum_nprobes",
));
}
}

if minimum_nprobes < 1 {
return Err(PyValueError::new_err("minimum_nprobes must be >= 1"));
}

if maximum_nprobes.unwrap_or(usize::MAX) < 1 {
return Err(PyValueError::new_err("maximum_nprobes must be >= 1"));
if let Some(maximum_nprobes) = maximum_nprobes {
if maximum_nprobes < 1 {
return Err(PyValueError::new_err("maximum_nprobes must be >= 1"));
}
}

let metric_type: Option<MetricType> =
Expand Down
5 changes: 4 additions & 1 deletion rust/lance-index/src/vector.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,10 @@ pub struct Query {
pub upper_bound: Option<f32>,

/// The minimum number of probes to load and search. More partitions
/// will only be loaded if we have not found k results.
/// will only be loaded if we have not found k results, or the the algorithm
/// determines more partitions are needed to satisfy recall requirements.
///
/// The planner will always search at least this many partitions. Defaults to 1.
Comment on lines +91 to +92
Copy link
Copy Markdown
Contributor

@wjones127 wjones127 Nov 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why this change? It seems like it would change the recall, wouldn't it?

Copy link
Copy Markdown
Contributor Author

@BubbleCal BubbleCal Nov 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, it will set min_nprobes = max(min_probes, pruned_nprobes) where pruned_nprobes is a value that the number of partitions that are close to the closest one.

Will add more comments here.

The figure I posted in PR desc contains recalls, it doesn't affect recall, because the dynamic pruning only excludes partitions that are too far

pub minimum_nprobes: usize,

/// The maximum number of probes to load and search. If not set then
Expand Down
38 changes: 20 additions & 18 deletions rust/lance/src/dataset/scanner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1047,7 +1047,7 @@ impl Scanner {
k,
lower_bound: None,
upper_bound: None,
minimum_nprobes: 20,
minimum_nprobes: 1,
maximum_nprobes: None,
ef: None,
refine_factor: None,
Expand Down Expand Up @@ -1110,6 +1110,8 @@ impl Scanner {
/// If we have found k matching results after searching this many partitions then
/// the search will stop. Increasing this number can increase recall but will increase
/// latency on all queries.
///
/// The default value is 1.
pub fn minimum_nprobes(&mut self, n: usize) -> &mut Self {
if let Some(q) = self.nearest.as_mut() {
q.minimum_nprobes = n;
Expand Down Expand Up @@ -6848,7 +6850,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=42), expr=...
ANNSubIndex: name=..., k=42, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1";
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1";
assert_plan_equals(
&dataset.dataset,
|scan| scan.nearest("vec", &q, 42),
Expand All @@ -6868,7 +6870,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=40), expr=...
ANNSubIndex: name=..., k=40, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1";
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1";
assert_plan_equals(
&dataset.dataset,
|scan| Ok(scan.nearest("vec", &q, 10)?.refine(4)),
Expand Down Expand Up @@ -6912,7 +6914,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=17), expr=...
ANNSubIndex: name=..., k=17, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1";
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1";
assert_plan_equals(
&dataset.dataset,
|scan| {
Expand All @@ -6933,7 +6935,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=17), expr=...
ANNSubIndex: name=..., k=17, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
FilterExec: i@0 > 10
LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None"
} else {
Expand All @@ -6942,7 +6944,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=17), expr=...
ANNSubIndex: name=..., k=17, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
LanceRead: uri=..., projection=[], num_fragments=2, range_before=None, range_after=None, \
row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)
"
Expand Down Expand Up @@ -6978,7 +6980,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=6), expr=...
ANNSubIndex: name=..., k=6, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1";
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1";
assert_plan_equals(
&dataset.dataset,
|scan| scan.nearest("vec", &q, 6),
Expand Down Expand Up @@ -7010,7 +7012,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=15), expr=...
ANNSubIndex: name=..., k=15, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1";
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1";
assert_plan_equals(
&dataset.dataset,
|scan| scan.nearest("vec", &q, 15)?.filter("i > 10"),
Expand Down Expand Up @@ -7039,7 +7041,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=5), expr=...
ANNSubIndex: name=..., k=5, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
FilterExec: i@0 > 10
LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None"
} else {
Expand All @@ -7061,7 +7063,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=5), expr=...
ANNSubIndex: name=..., k=5, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
LanceRead: uri=..., projection=[], num_fragments=2, range_before=None, range_after=None, \
row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)"
};
Expand Down Expand Up @@ -7092,7 +7094,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=5), expr=...
ANNSubIndex: name=..., k=5, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
ScalarIndexQuery: query=[i > 10]@i_idx";
assert_plan_equals(
&dataset.dataset,
Expand All @@ -7113,7 +7115,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=5), expr=...
ANNSubIndex: name=..., k=5, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
FilterExec: i@0 > 10
LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None"
} else {
Expand All @@ -7122,7 +7124,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=5), expr=...
ANNSubIndex: name=..., k=5, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
LanceRead: uri=..., projection=[], num_fragments=3, range_before=None, \
range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)"
};
Expand Down Expand Up @@ -7160,7 +7162,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=8), expr=...
ANNSubIndex: name=..., k=8, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
ScalarIndexQuery: query=[i > 10]@i_idx";
assert_plan_equals(
&dataset.dataset,
Expand Down Expand Up @@ -7196,7 +7198,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=11), expr=...
ANNSubIndex: name=..., k=11, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1
ScalarIndexQuery: query=[i > 10]@i_idx";
dataset.make_scalar_index().await?;
assert_plan_equals(
Expand Down Expand Up @@ -7536,7 +7538,7 @@ mod test {
},
"SortExec: TopK(fetch=32), expr=[_distance@0 ASC NULLS LAST, _rowid@1 ASC NULLS LAST]...
ANNSubIndex: name=idx, k=32, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1",
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1",
)
.await
.unwrap();
Expand All @@ -7551,7 +7553,7 @@ mod test {
},
"SortExec: TopK(fetch=33), expr=[_distance@0 ASC NULLS LAST, _rowid@1 ASC NULLS LAST]...
ANNSubIndex: name=idx, k=33, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1",
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1",
)
.await
.unwrap();
Expand Down Expand Up @@ -7579,7 +7581,7 @@ mod test {
CoalesceBatchesExec: target_batch_size=8192
SortExec: TopK(fetch=34), expr=[_distance@0 ASC NULLS LAST, _rowid@1 ASC NULLS LAST]...
ANNSubIndex: name=idx, k=34, deltas=1
ANNIvfPartition: uuid=..., minimum_nprobes=20, maximum_nprobes=None, deltas=1",
ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1",
)
.await
.unwrap();
Expand Down
Loading
Loading