Skip to content

Commit e530742

Browse files
cutecutecatVoVAllen
cutecutecat
and
VoVAllen
authored
feat: add nprobe into guc (#212)
Signed-off-by: cutecutecat <junyuchen@tensorchord.ai> Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com>
1 parent 2dc6b72 commit e530742

File tree

20 files changed

+153
-54
lines changed

20 files changed

+153
-54
lines changed

‎crates/service/src/algorithms/ivf/ivf_naive.rs‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,14 @@ impl<S: G> IvfNaive<S> {
5555
self.mmap.raw.payload(i)
5656
}
5757

58-
pub fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
59-
search(&self.mmap, k, vector, filter)
58+
pub fn search(
59+
&self,
60+
k: usize,
61+
vector: &[S::Scalar],
62+
nprobe: u32,
63+
filter: &mut impl Filter,
64+
) -> Heap {
65+
search(&self.mmap, k, vector, nprobe, filter)
6066
}
6167
}
6268

@@ -70,7 +76,6 @@ pub struct IvfRam<S: G> {
7076
dims: u16,
7177
// ----------------------
7278
nlist: u32,
73-
nprobe: u32,
7479
// ----------------------
7580
centroids: Vec2<S>,
7681
heads: Vec<AtomicU32>,
@@ -87,7 +92,6 @@ pub struct IvfMmap<S: G> {
8792
dims: u16,
8893
// ----------------------
8994
nlist: u32,
90-
nprobe: u32,
9195
// ----------------------
9296
centroids: MmapArray<S::Scalar>,
9397
heads: MmapArray<u32>,
@@ -116,7 +120,6 @@ pub fn make<S: G>(
116120
least_iterations,
117121
iterations,
118122
nlist,
119-
nprobe,
120123
nsample,
121124
quantization: quantization_opts,
122125
} = options.indexing.clone().unwrap_ivf();
@@ -186,7 +189,6 @@ pub fn make<S: G>(
186189
centroids,
187190
heads,
188191
nexts,
189-
nprobe,
190192
nlist,
191193
dims,
192194
}
@@ -212,7 +214,6 @@ pub fn save<S: G>(mut ram: IvfRam<S>, path: PathBuf) -> IvfMmap<S> {
212214
quantization: ram.quantization,
213215
dims: ram.dims,
214216
nlist: ram.nlist,
215-
nprobe: ram.nprobe,
216217
centroids,
217218
heads,
218219
nexts,
@@ -230,13 +231,12 @@ pub fn load<S: G>(path: PathBuf, options: IndexOptions) -> IvfMmap<S> {
230231
let centroids = MmapArray::open(path.join("centroids"));
231232
let heads = MmapArray::open(path.join("heads"));
232233
let nexts = MmapArray::open(path.join("nexts"));
233-
let IvfIndexingOptions { nlist, nprobe, .. } = options.indexing.unwrap_ivf();
234+
let IvfIndexingOptions { nlist, .. } = options.indexing.unwrap_ivf();
234235
IvfMmap {
235236
raw,
236237
quantization,
237238
dims: options.vector.dims,
238239
nlist,
239-
nprobe,
240240
centroids,
241241
heads,
242242
nexts,
@@ -247,11 +247,12 @@ pub fn search<S: G>(
247247
mmap: &IvfMmap<S>,
248248
k: usize,
249249
vector: &[S::Scalar],
250+
nprobe: u32,
250251
filter: &mut impl Filter,
251252
) -> Heap {
252253
let mut target = vector.to_vec();
253254
S::elkan_k_means_normalize(&mut target);
254-
let mut lists = Heap::new(mmap.nprobe as usize);
255+
let mut lists = Heap::new(nprobe as usize);
255256
for i in 0..mmap.nlist {
256257
let centroid = mmap.centroids(i);
257258
let distance = S::elkan_k_means_distance(&target, centroid);

‎crates/service/src/algorithms/ivf/ivf_pq.rs‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,14 @@ impl<S: G> IvfPq<S> {
5555
self.mmap.raw.payload(i)
5656
}
5757

58-
pub fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
59-
search(&self.mmap, k, vector, filter)
58+
pub fn search(
59+
&self,
60+
k: usize,
61+
vector: &[S::Scalar],
62+
nprobe: u32,
63+
filter: &mut impl Filter,
64+
) -> Heap {
65+
search(&self.mmap, k, vector, nprobe, filter)
6066
}
6167
}
6268

@@ -70,7 +76,6 @@ pub struct IvfRam<S: G> {
7076
dims: u16,
7177
// ----------------------
7278
nlist: u32,
73-
nprobe: u32,
7479
// ----------------------
7580
centroids: Vec2<S>,
7681
heads: Vec<AtomicU32>,
@@ -87,7 +92,6 @@ pub struct IvfMmap<S: G> {
8792
dims: u16,
8893
// ----------------------
8994
nlist: u32,
90-
nprobe: u32,
9195
// ----------------------
9296
centroids: MmapArray<S::Scalar>,
9397
heads: MmapArray<u32>,
@@ -116,7 +120,6 @@ pub fn make<S: G>(
116120
least_iterations,
117121
iterations,
118122
nlist,
119-
nprobe,
120123
nsample,
121124
quantization: quantization_opts,
122125
} = options.indexing.clone().unwrap_ivf();
@@ -189,7 +192,6 @@ pub fn make<S: G>(
189192
centroids,
190193
heads,
191194
nexts,
192-
nprobe,
193195
nlist,
194196
dims,
195197
}
@@ -215,7 +217,6 @@ pub fn save<S: G>(mut ram: IvfRam<S>, path: PathBuf) -> IvfMmap<S> {
215217
quantization: ram.quantization,
216218
dims: ram.dims,
217219
nlist: ram.nlist,
218-
nprobe: ram.nprobe,
219220
centroids,
220221
heads,
221222
nexts,
@@ -233,13 +234,12 @@ pub fn load<S: G>(path: PathBuf, options: IndexOptions) -> IvfMmap<S> {
233234
let centroids = MmapArray::open(path.join("centroids"));
234235
let heads = MmapArray::open(path.join("heads"));
235236
let nexts = MmapArray::open(path.join("nexts"));
236-
let IvfIndexingOptions { nlist, nprobe, .. } = options.indexing.unwrap_ivf();
237+
let IvfIndexingOptions { nlist, .. } = options.indexing.unwrap_ivf();
237238
IvfMmap {
238239
raw,
239240
quantization,
240241
dims: options.vector.dims,
241242
nlist,
242-
nprobe,
243243
centroids,
244244
heads,
245245
nexts,
@@ -250,11 +250,12 @@ pub fn search<S: G>(
250250
mmap: &IvfMmap<S>,
251251
k: usize,
252252
vector: &[S::Scalar],
253+
nprobe: u32,
253254
filter: &mut impl Filter,
254255
) -> Heap {
255256
let mut target = vector.to_vec();
256257
S::elkan_k_means_normalize(&mut target);
257-
let mut lists = Heap::new(mmap.nprobe as usize);
258+
let mut lists = Heap::new(nprobe as usize);
258259
for i in 0..mmap.nlist {
259260
let centroid = mmap.centroids(i);
260261
let distance = S::elkan_k_means_distance(&target, centroid);

‎crates/service/src/algorithms/ivf/mod.rs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,10 +70,16 @@ impl<S: G> Ivf<S> {
7070
}
7171
}
7272

73-
pub fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
73+
pub fn search(
74+
&self,
75+
k: usize,
76+
vector: &[S::Scalar],
77+
nprobe: u32,
78+
filter: &mut impl Filter,
79+
) -> Heap {
7480
match self {
75-
Ivf::Naive(x) => x.search(k, vector, filter),
76-
Ivf::Pq(x) => x.search(k, vector, filter),
81+
Ivf::Naive(x) => x.search(k, vector, nprobe, filter),
82+
Ivf::Pq(x) => x.search(k, vector, nprobe, filter),
7783
}
7884
}
7985
}

‎crates/service/src/index/indexing/flat.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use super::AbstractIndexing;
22
use crate::algorithms::quantization::QuantizationOptions;
33
use crate::index::segments::growing::GrowingSegment;
4+
use crate::index::segments::sealed::SealedSearchGucs;
45
use crate::index::IndexOptions;
56
use crate::prelude::*;
67
use crate::{algorithms::flat::Flat, index::segments::sealed::SealedSegment};
@@ -57,7 +58,13 @@ impl<S: G> AbstractIndexing<S> for FlatIndexing<S> {
5758
self.raw.payload(i)
5859
}
5960

60-
fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
61+
fn search(
62+
&self,
63+
k: usize,
64+
vector: &[S::Scalar],
65+
_gucs: SealedSearchGucs,
66+
filter: &mut impl Filter,
67+
) -> Heap {
6168
self.raw.search(k, vector, filter)
6269
}
6370
}

‎crates/service/src/index/indexing/hnsw.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use super::AbstractIndexing;
22
use crate::algorithms::hnsw::HnswIndexIter;
33
use crate::algorithms::quantization::QuantizationOptions;
44
use crate::index::segments::growing::GrowingSegment;
5+
use crate::index::segments::sealed::SealedSearchGucs;
56
use crate::index::IndexOptions;
67
use crate::prelude::*;
78
use crate::{algorithms::hnsw::Hnsw, index::segments::sealed::SealedSegment};
@@ -74,7 +75,13 @@ impl<S: G> AbstractIndexing<S> for HnswIndexing<S> {
7475
self.raw.payload(i)
7576
}
7677

77-
fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
78+
fn search(
79+
&self,
80+
k: usize,
81+
vector: &[S::Scalar],
82+
_gucs: SealedSearchGucs,
83+
filter: &mut impl Filter,
84+
) -> Heap {
7885
self.raw.search(k, vector, filter)
7986
}
8087
}

‎crates/service/src/index/indexing/ivf.rs‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use super::AbstractIndexing;
22
use crate::algorithms::ivf::Ivf;
33
use crate::algorithms::quantization::QuantizationOptions;
44
use crate::index::segments::growing::GrowingSegment;
5+
use crate::index::segments::sealed::SealedSearchGucs;
56
use crate::index::segments::sealed::SealedSegment;
67
use crate::index::IndexOptions;
78
use crate::prelude::*;
@@ -22,9 +23,6 @@ pub struct IvfIndexingOptions {
2223
#[serde(default = "IvfIndexingOptions::default_nlist")]
2324
#[validate(range(min = 1, max = 1_000_000))]
2425
pub nlist: u32,
25-
#[serde(default = "IvfIndexingOptions::default_nprobe")]
26-
#[validate(range(min = 1, max = 1_000_000))]
27-
pub nprobe: u32,
2826
#[serde(default = "IvfIndexingOptions::default_nsample")]
2927
#[validate(range(min = 1, max = 1_000_000))]
3028
pub nsample: u32,
@@ -43,9 +41,6 @@ impl IvfIndexingOptions {
4341
fn default_nlist() -> u32 {
4442
1000
4543
}
46-
fn default_nprobe() -> u32 {
47-
10
48-
}
4944
fn default_nsample() -> u32 {
5045
65536
5146
}
@@ -57,7 +52,6 @@ impl Default for IvfIndexingOptions {
5752
least_iterations: Self::default_least_iterations(),
5853
iterations: Self::default_iterations(),
5954
nlist: Self::default_nlist(),
60-
nprobe: Self::default_nprobe(),
6155
nsample: Self::default_nsample(),
6256
quantization: Default::default(),
6357
}
@@ -96,7 +90,13 @@ impl<S: G> AbstractIndexing<S> for IvfIndexing<S> {
9690
self.raw.payload(i)
9791
}
9892

99-
fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
100-
self.raw.search(k, vector, filter)
93+
fn search(
94+
&self,
95+
k: usize,
96+
vector: &[S::Scalar],
97+
gucs: SealedSearchGucs,
98+
filter: &mut impl Filter,
99+
) -> Heap {
100+
self.raw.search(k, vector, gucs.ivf_nprob, filter)
101101
}
102102
}

‎crates/service/src/index/indexing/mod.rs‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use super::segments::growing::GrowingSegment;
99
use super::segments::sealed::SealedSegment;
1010
use super::IndexOptions;
1111
use crate::algorithms::hnsw::HnswIndexIter;
12+
use crate::index::segments::sealed::SealedSearchGucs;
1213
use crate::prelude::*;
1314
use serde::{Deserialize, Serialize};
1415
use std::path::PathBuf;
@@ -72,7 +73,13 @@ pub trait AbstractIndexing<S: G>: Sized {
7273
fn len(&self) -> u32;
7374
fn vector(&self, i: u32) -> &[S::Scalar];
7475
fn payload(&self, i: u32) -> Payload;
75-
fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap;
76+
fn search(
77+
&self,
78+
k: usize,
79+
vector: &[S::Scalar],
80+
gucs: SealedSearchGucs,
81+
filter: &mut impl Filter,
82+
) -> Heap;
7683
}
7784

7885
pub enum DynamicIndexing<S: G> {
@@ -137,11 +144,17 @@ impl<S: G> DynamicIndexing<S> {
137144
}
138145
}
139146

140-
pub fn search(&self, k: usize, vector: &[S::Scalar], filter: &mut impl Filter) -> Heap {
147+
pub fn search(
148+
&self,
149+
k: usize,
150+
vector: &[S::Scalar],
151+
gucs: SealedSearchGucs,
152+
filter: &mut impl Filter,
153+
) -> Heap {
141154
match self {
142-
DynamicIndexing::Flat(x) => x.search(k, vector, filter),
143-
DynamicIndexing::Ivf(x) => x.search(k, vector, filter),
144-
DynamicIndexing::Hnsw(x) => x.search(k, vector, filter),
155+
DynamicIndexing::Flat(x) => x.search(k, vector, gucs, filter),
156+
DynamicIndexing::Ivf(x) => x.search(k, vector, gucs, filter),
157+
DynamicIndexing::Hnsw(x) => x.search(k, vector, gucs, filter),
145158
}
146159
}
147160

‎crates/service/src/index/mod.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use self::segments::SegmentsOptions;
1313
use crate::index::indexing::DynamicIndexIter;
1414
use crate::index::optimizing::indexing::OptimizerIndexing;
1515
use crate::index::optimizing::sealing::OptimizerSealing;
16+
use crate::index::segments::SearchGucs;
1617
use crate::prelude::*;
1718
use crate::utils::clean::clean;
1819
use crate::utils::dir_ops::sync_dir;
@@ -270,6 +271,7 @@ impl<S: G> IndexView<S> {
270271
&self,
271272
k: usize,
272273
vector: &[S::Scalar],
274+
gucs: SearchGucs,
273275
mut filter: F,
274276
) -> Vec<Pointer> {
275277
assert_eq!(self.options.vector.dims as usize, vector.len());
@@ -307,7 +309,9 @@ impl<S: G> IndexView<S> {
307309
let mut result = Heap::new(k);
308310
let mut heaps = BinaryHeap::with_capacity(1 + n);
309311
for (_, sealed) in self.sealed.iter() {
310-
let p = sealed.search(k, vector, &mut filter).into_reversed_heap();
312+
let p = sealed
313+
.search(k, vector, gucs.sealed, &mut filter)
314+
.into_reversed_heap();
311315
heaps.push(Comparer(p));
312316
}
313317
for (_, growing) in self.growing.iter() {

‎crates/service/src/index/segments/mod.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,18 @@ pub mod growing;
22
pub mod sealed;
33

44
use super::IndexTracker;
5+
use crate::index::segments::sealed::SealedSearchGucs;
56
use serde::{Deserialize, Serialize};
67
use std::path::PathBuf;
78
use std::sync::Arc;
89
use validator::Validate;
910
use validator::ValidationError;
1011

12+
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Validate)]
13+
pub struct SearchGucs {
14+
pub sealed: SealedSearchGucs,
15+
}
16+
1117
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
1218
#[serde(deny_unknown_fields)]
1319
#[validate(schema(function = "Self::validate_0"))]

0 commit comments

Comments
 (0)