Skip to content

Commit 4f60c4c

Browse files
silver-ymzsilver-ymz
authored
feat: allow user to set vbase range (#206)
* feat: allow user to set vbase range Signed-off-by: silver-ymz <[email protected]> * separate the enable_vbase and vbase_range settings Signed-off-by: silver-ymz <[email protected]> --------- Signed-off-by: silver-ymz <[email protected]>
1 parent 7a2eb0a commit 4f60c4c

File tree

9 files changed

+36
-19
lines changed

9 files changed

+36
-19
lines changed

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -330,11 +330,9 @@ impl<S: G> IndexView<S> {
330330
.map(|x| Pointer::from_u48(x.payload >> 16))
331331
.collect()
332332
}
333-
pub fn vbase(&self, vector: &[S::Scalar]) -> impl Iterator<Item = Pointer> + '_ {
333+
pub fn vbase(&self, vector: &[S::Scalar], range: usize) -> impl Iterator<Item = Pointer> + '_ {
334334
assert_eq!(self.options.vector.dims as usize, vector.len());
335335

336-
let range = 86;
337-
338336
struct Comparer<'a, S: G> {
339337
iter: ComparerIter<'a, S>,
340338
item: Option<HeapElement>,

‎crates/service/src/worker/instance.rs

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -139,43 +139,44 @@ impl InstanceView {
139139
pub fn vbase(
140140
&self,
141141
vector: DynamicVector,
142+
range: usize,
142143
) -> Result<impl Iterator<Item = Pointer> + '_, FriendlyError> {
143144
match (self, vector) {
144145
(InstanceView::F32Cos(x), DynamicVector::F32(vector)) => {
145146
if x.options.vector.dims as usize != vector.len() {
146147
return Err(FriendlyError::Unmatched2);
147148
}
148-
Ok(Box::new(x.vbase(&vector)) as Box<dyn Iterator<Item = Pointer>>)
149+
Ok(Box::new(x.vbase(&vector, range)) as Box<dyn Iterator<Item = Pointer>>)
149150
}
150151
(InstanceView::F32Dot(x), DynamicVector::F32(vector)) => {
151152
if x.options.vector.dims as usize != vector.len() {
152153
return Err(FriendlyError::Unmatched2);
153154
}
154-
Ok(Box::new(x.vbase(&vector)))
155+
Ok(Box::new(x.vbase(&vector, range)))
155156
}
156157
(InstanceView::F32L2(x), DynamicVector::F32(vector)) => {
157158
if x.options.vector.dims as usize != vector.len() {
158159
return Err(FriendlyError::Unmatched2);
159160
}
160-
Ok(Box::new(x.vbase(&vector)))
161+
Ok(Box::new(x.vbase(&vector, range)))
161162
}
162163
(InstanceView::F16Cos(x), DynamicVector::F16(vector)) => {
163164
if x.options.vector.dims as usize != vector.len() {
164165
return Err(FriendlyError::Unmatched2);
165166
}
166-
Ok(Box::new(x.vbase(&vector)))
167+
Ok(Box::new(x.vbase(&vector, range)))
167168
}
168169
(InstanceView::F16Dot(x), DynamicVector::F16(vector)) => {
169170
if x.options.vector.dims as usize != vector.len() {
170171
return Err(FriendlyError::Unmatched2);
171172
}
172-
Ok(Box::new(x.vbase(&vector)))
173+
Ok(Box::new(x.vbase(&vector, range)))
173174
}
174175
(InstanceView::F16L2(x), DynamicVector::F16(vector)) => {
175176
if x.options.vector.dims as usize != vector.len() {
176177
return Err(FriendlyError::Unmatched2);
177178
}
178-
Ok(Box::new(x.vbase(&vector)))
179+
Ok(Box::new(x.vbase(&vector, range)))
179180
}
180181
_ => Err(FriendlyError::Unmatched2),
181182
}

‎docs/searching.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ Search options are specified by PostgreSQL GUC. You can use `SET` command to app
3131
| vectors.k | integer | Expected number of candidates returned by index. The parameter will influence the recall if you use HNSW or quantization for indexing. Default value is `64`. |
3232
| vectors.enable_prefilter | boolean | Enable prefiltering or not. Default value is `off`. |
3333
| vectors.enable_vector_index | boolean | Enable vector indexes or not. This option is for debugging. Default value is `on`. |
34-
| vectors.enable_vbase | boolean | Enable vbase optimization. Default value is `off`. |
34+
| vectors.enable_vbase | boolean | Enable vbase optimization or not. Default value is `off`. |
35+
| vectors.vbase_range | integer | Range size of vbase optimization. The larger the range, the more accurate the result, but the slower the speed. Default value is `100`. |
3536

3637
Note: When `vectors.enable_vbase` is enabled, prefilter does not work.

‎src/bgworker/mod.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ fn session(worker: Arc<Worker>, mut handler: RpcHandler) -> Result<(), IpcError>
136136
let result = worker.call_stat(id);
137137
handler = x.leave(result)?;
138138
}
139-
RpcHandle::Vbase { id, vector, x } => {
139+
RpcHandle::Vbase { id, vbase, x } => {
140140
use crate::ipc::server::VbaseHandle::*;
141141
let instance = match worker.get_instance(id) {
142142
Ok(x) => x,
@@ -146,7 +146,7 @@ fn session(worker: Arc<Worker>, mut handler: RpcHandler) -> Result<(), IpcError>
146146
}
147147
};
148148
let view = instance.view();
149-
let mut it = match view.vbase(vector) {
149+
let mut it = match view.vbase(vbase.0, vbase.1) {
150150
Ok(x) => x,
151151
Err(e) => {
152152
x.error(Err(e))?;

‎src/gucs.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ pub static ENABLE_PREFILTER: GucSetting<bool> = GucSetting::<bool>::new(false);
2525

2626
pub static ENABLE_VBASE: GucSetting<bool> = GucSetting::<bool>::new(false);
2727

28+
pub static VBASE_RANGE: GucSetting<i32> = GucSetting::<i32>::new(100);
29+
2830
pub static TRANSPORT: GucSetting<Transport> = GucSetting::<Transport>::new(Transport::default());
2931

3032
pub unsafe fn init() {
@@ -70,6 +72,16 @@ pub unsafe fn init() {
7072
GucContext::Userset,
7173
GucFlags::default(),
7274
);
75+
GucRegistry::define_int_guc(
76+
"vectors.vbase_range",
77+
"The range of vbase.",
78+
"The range size of vabse optimization.",
79+
&VBASE_RANGE,
80+
1,
81+
u16::MAX as _,
82+
GucContext::Userset,
83+
GucFlags::default(),
84+
);
7385
GucRegistry::define_enum_guc(
7486
"vectors.transport",
7587
"Transport for communicating with background worker.",

‎src/index/am_scan.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crate::gucs::ENABLE_PREFILTER;
22
use crate::gucs::ENABLE_VBASE;
33
use crate::gucs::K;
4+
use crate::gucs::VBASE_RANGE;
45
use crate::index::utils::from_datum;
56
use crate::ipc::client::ClientGuard;
67
use crate::ipc::client::Vbase;
@@ -96,7 +97,7 @@ pub unsafe fn next_scan(scan: pgrx::pg_sys::IndexScanDesc) -> bool {
9697
let mut rpc = crate::ipc::client::borrow_mut();
9798

9899
if ENABLE_VBASE.get() {
99-
let vbase = rpc.vbase(id, vector.clone());
100+
let vbase = rpc.vbase(id, (vector.clone(), VBASE_RANGE.get() as _));
100101
*scanner = Scanner::Vbase { node, vbase };
101102
} else {
102103
let k = K.get() as _;

‎src/ipc/client/mod.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,12 @@ impl Rpc {
120120
let stat::StatPacket::Leave { result } = self.socket.client_recv().friendly();
121121
result.friendly()
122122
}
123-
pub fn vbase(mut self: ClientGuard<Self>, id: Id, vector: DynamicVector) -> ClientGuard<Vbase> {
124-
let packet = RpcPacket::Vbase { id, vector };
123+
pub fn vbase(
124+
mut self: ClientGuard<Self>,
125+
id: Id,
126+
vbase: (DynamicVector, usize),
127+
) -> ClientGuard<Vbase> {
128+
let packet = RpcPacket::Vbase { id, vbase };
125129
self.socket.client_send(packet).friendly();
126130
let vbase::VbaseErrorPacket { result } = self.socket.client_recv().friendly();
127131
result.friendly();

‎src/ipc/packet/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,6 @@ pub enum RpcPacket {
4040
},
4141
Vbase {
4242
id: Id,
43-
vector: DynamicVector,
43+
vbase: (DynamicVector, usize),
4444
},
4545
}

‎src/ipc/server/mod.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,9 +65,9 @@ impl RpcHandler {
6565
socket: self.socket,
6666
},
6767
},
68-
RpcPacket::Vbase { id, vector } => RpcHandle::Vbase {
68+
RpcPacket::Vbase { id, vbase } => RpcHandle::Vbase {
6969
id,
70-
vector,
70+
vbase,
7171
x: Vbase {
7272
socket: self.socket,
7373
},
@@ -111,7 +111,7 @@ pub enum RpcHandle {
111111
},
112112
Vbase {
113113
id: Id,
114-
vector: DynamicVector,
114+
vbase: (DynamicVector, usize),
115115
x: Vbase,
116116
},
117117
}

0 commit comments

Comments
 (0)