From 4f60c4c824f3779c05f55b7392c61506588382b0 Mon Sep 17 00:00:00 2001 From: Mingzhuo Yin Date: Wed, 27 Dec 2023 14:32:23 +0800 Subject: [PATCH] feat: allow user to set vbase range (#206) * feat: allow user to set vbase range Signed-off-by: silver-ymz * separate the enable_vbase and vbase_range settings Signed-off-by: silver-ymz --------- Signed-off-by: silver-ymz --- crates/service/src/index/mod.rs | 4 +--- crates/service/src/worker/instance.rs | 13 +++++++------ docs/searching.md | 3 ++- src/bgworker/mod.rs | 4 ++-- src/gucs.rs | 12 ++++++++++++ src/index/am_scan.rs | 3 ++- src/ipc/client/mod.rs | 8 ++++++-- src/ipc/packet/mod.rs | 2 +- src/ipc/server/mod.rs | 6 +++--- 9 files changed, 36 insertions(+), 19 deletions(-) diff --git a/crates/service/src/index/mod.rs b/crates/service/src/index/mod.rs index f9b049e05..52acd6913 100644 --- a/crates/service/src/index/mod.rs +++ b/crates/service/src/index/mod.rs @@ -330,11 +330,9 @@ impl IndexView { .map(|x| Pointer::from_u48(x.payload >> 16)) .collect() } - pub fn vbase(&self, vector: &[S::Scalar]) -> impl Iterator + '_ { + pub fn vbase(&self, vector: &[S::Scalar], range: usize) -> impl Iterator + '_ { assert_eq!(self.options.vector.dims as usize, vector.len()); - let range = 86; - struct Comparer<'a, S: G> { iter: ComparerIter<'a, S>, item: Option, diff --git a/crates/service/src/worker/instance.rs b/crates/service/src/worker/instance.rs index 16c2df1ce..8d10a1508 100644 --- a/crates/service/src/worker/instance.rs +++ b/crates/service/src/worker/instance.rs @@ -139,43 +139,44 @@ impl InstanceView { pub fn vbase( &self, vector: DynamicVector, + range: usize, ) -> Result + '_, FriendlyError> { match (self, vector) { (InstanceView::F32Cos(x), DynamicVector::F32(vector)) => { if x.options.vector.dims as usize != vector.len() { return Err(FriendlyError::Unmatched2); } - Ok(Box::new(x.vbase(&vector)) as Box>) + Ok(Box::new(x.vbase(&vector, range)) as Box>) } (InstanceView::F32Dot(x), DynamicVector::F32(vector)) => { if x.options.vector.dims as usize != vector.len() { return Err(FriendlyError::Unmatched2); } - Ok(Box::new(x.vbase(&vector))) + Ok(Box::new(x.vbase(&vector, range))) } (InstanceView::F32L2(x), DynamicVector::F32(vector)) => { if x.options.vector.dims as usize != vector.len() { return Err(FriendlyError::Unmatched2); } - Ok(Box::new(x.vbase(&vector))) + Ok(Box::new(x.vbase(&vector, range))) } (InstanceView::F16Cos(x), DynamicVector::F16(vector)) => { if x.options.vector.dims as usize != vector.len() { return Err(FriendlyError::Unmatched2); } - Ok(Box::new(x.vbase(&vector))) + Ok(Box::new(x.vbase(&vector, range))) } (InstanceView::F16Dot(x), DynamicVector::F16(vector)) => { if x.options.vector.dims as usize != vector.len() { return Err(FriendlyError::Unmatched2); } - Ok(Box::new(x.vbase(&vector))) + Ok(Box::new(x.vbase(&vector, range))) } (InstanceView::F16L2(x), DynamicVector::F16(vector)) => { if x.options.vector.dims as usize != vector.len() { return Err(FriendlyError::Unmatched2); } - Ok(Box::new(x.vbase(&vector))) + Ok(Box::new(x.vbase(&vector, range))) } _ => Err(FriendlyError::Unmatched2), } diff --git a/docs/searching.md b/docs/searching.md index 2220b6bd7..1441bd4b7 100644 --- a/docs/searching.md +++ b/docs/searching.md @@ -31,6 +31,7 @@ Search options are specified by PostgreSQL GUC. You can use `SET` command to app | 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`. | | vectors.enable_prefilter | boolean | Enable prefiltering or not. Default value is `off`. | | vectors.enable_vector_index | boolean | Enable vector indexes or not. This option is for debugging. Default value is `on`. | -| vectors.enable_vbase | boolean | Enable vbase optimization. Default value is `off`. | +| vectors.enable_vbase | boolean | Enable vbase optimization or not. Default value is `off`. | +| 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`. | Note: When `vectors.enable_vbase` is enabled, prefilter does not work. diff --git a/src/bgworker/mod.rs b/src/bgworker/mod.rs index 529a95e2e..16b011680 100644 --- a/src/bgworker/mod.rs +++ b/src/bgworker/mod.rs @@ -136,7 +136,7 @@ fn session(worker: Arc, mut handler: RpcHandler) -> Result<(), IpcError> let result = worker.call_stat(id); handler = x.leave(result)?; } - RpcHandle::Vbase { id, vector, x } => { + RpcHandle::Vbase { id, vbase, x } => { use crate::ipc::server::VbaseHandle::*; let instance = match worker.get_instance(id) { Ok(x) => x, @@ -146,7 +146,7 @@ fn session(worker: Arc, mut handler: RpcHandler) -> Result<(), IpcError> } }; let view = instance.view(); - let mut it = match view.vbase(vector) { + let mut it = match view.vbase(vbase.0, vbase.1) { Ok(x) => x, Err(e) => { x.error(Err(e))?; diff --git a/src/gucs.rs b/src/gucs.rs index e9c3e8189..96042bb11 100644 --- a/src/gucs.rs +++ b/src/gucs.rs @@ -25,6 +25,8 @@ pub static ENABLE_PREFILTER: GucSetting = GucSetting::::new(false); pub static ENABLE_VBASE: GucSetting = GucSetting::::new(false); +pub static VBASE_RANGE: GucSetting = GucSetting::::new(100); + pub static TRANSPORT: GucSetting = GucSetting::::new(Transport::default()); pub unsafe fn init() { @@ -70,6 +72,16 @@ pub unsafe fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_int_guc( + "vectors.vbase_range", + "The range of vbase.", + "The range size of vabse optimization.", + &VBASE_RANGE, + 1, + u16::MAX as _, + GucContext::Userset, + GucFlags::default(), + ); GucRegistry::define_enum_guc( "vectors.transport", "Transport for communicating with background worker.", diff --git a/src/index/am_scan.rs b/src/index/am_scan.rs index f437de764..9cca915c7 100644 --- a/src/index/am_scan.rs +++ b/src/index/am_scan.rs @@ -1,6 +1,7 @@ use crate::gucs::ENABLE_PREFILTER; use crate::gucs::ENABLE_VBASE; use crate::gucs::K; +use crate::gucs::VBASE_RANGE; use crate::index::utils::from_datum; use crate::ipc::client::ClientGuard; use crate::ipc::client::Vbase; @@ -96,7 +97,7 @@ pub unsafe fn next_scan(scan: pgrx::pg_sys::IndexScanDesc) -> bool { let mut rpc = crate::ipc::client::borrow_mut(); if ENABLE_VBASE.get() { - let vbase = rpc.vbase(id, vector.clone()); + let vbase = rpc.vbase(id, (vector.clone(), VBASE_RANGE.get() as _)); *scanner = Scanner::Vbase { node, vbase }; } else { let k = K.get() as _; diff --git a/src/ipc/client/mod.rs b/src/ipc/client/mod.rs index 38aa32b97..9d4633e75 100644 --- a/src/ipc/client/mod.rs +++ b/src/ipc/client/mod.rs @@ -120,8 +120,12 @@ impl Rpc { let stat::StatPacket::Leave { result } = self.socket.client_recv().friendly(); result.friendly() } - pub fn vbase(mut self: ClientGuard, id: Id, vector: DynamicVector) -> ClientGuard { - let packet = RpcPacket::Vbase { id, vector }; + pub fn vbase( + mut self: ClientGuard, + id: Id, + vbase: (DynamicVector, usize), + ) -> ClientGuard { + let packet = RpcPacket::Vbase { id, vbase }; self.socket.client_send(packet).friendly(); let vbase::VbaseErrorPacket { result } = self.socket.client_recv().friendly(); result.friendly(); diff --git a/src/ipc/packet/mod.rs b/src/ipc/packet/mod.rs index 2df63f0a6..3dbb1ed40 100644 --- a/src/ipc/packet/mod.rs +++ b/src/ipc/packet/mod.rs @@ -40,6 +40,6 @@ pub enum RpcPacket { }, Vbase { id: Id, - vector: DynamicVector, + vbase: (DynamicVector, usize), }, } diff --git a/src/ipc/server/mod.rs b/src/ipc/server/mod.rs index efafb6efd..9a60fe37d 100644 --- a/src/ipc/server/mod.rs +++ b/src/ipc/server/mod.rs @@ -65,9 +65,9 @@ impl RpcHandler { socket: self.socket, }, }, - RpcPacket::Vbase { id, vector } => RpcHandle::Vbase { + RpcPacket::Vbase { id, vbase } => RpcHandle::Vbase { id, - vector, + vbase, x: Vbase { socket: self.socket, }, @@ -111,7 +111,7 @@ pub enum RpcHandle { }, Vbase { id: Id, - vector: DynamicVector, + vbase: (DynamicVector, usize), x: Vbase, }, }