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]>
tensorchord · Dec 27, 2023 · 4f60c4c · 4f60c4c
1 parent 7a2eb0a
commit 4f60c4c
Show file tree

Hide file tree

Showing 9 changed files with 36 additions and 19 deletions.
diff --git a/crates/service/src/index/mod.rs b/crates/service/src/index/mod.rs
@@ -330,11 +330,9 @@ impl<S: G> IndexView<S> {
             .map(|x| Pointer::from_u48(x.payload >> 16))
             .collect()
     }
-    pub fn vbase(&self, vector: &[S::Scalar]) -> impl Iterator<Item = Pointer> + '_ {
+    pub fn vbase(&self, vector: &[S::Scalar], range: usize) -> impl Iterator<Item = Pointer> + '_ {
         assert_eq!(self.options.vector.dims as usize, vector.len());
 
-        let range = 86;
-
         struct Comparer<'a, S: G> {
             iter: ComparerIter<'a, S>,
             item: Option<HeapElement>,

diff --git 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<impl Iterator<Item = Pointer> + '_, 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<dyn Iterator<Item = Pointer>>)
+                Ok(Box::new(x.vbase(&vector, range)) as Box<dyn Iterator<Item = Pointer>>)
             }
             (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
@@ -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
@@ -136,7 +136,7 @@ fn session(worker: Arc<Worker>, 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<Worker>, 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
@@ -25,6 +25,8 @@ pub static ENABLE_PREFILTER: GucSetting<bool> = GucSetting::<bool>::new(false);
 
 pub static ENABLE_VBASE: GucSetting<bool> = GucSetting::<bool>::new(false);
 
+pub static VBASE_RANGE: GucSetting<i32> = GucSetting::<i32>::new(100);
+
 pub static TRANSPORT: GucSetting<Transport> = GucSetting::<Transport>::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
@@ -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
@@ -120,8 +120,12 @@ impl Rpc {
         let stat::StatPacket::Leave { result } = self.socket.client_recv().friendly();
         result.friendly()
     }
-    pub fn vbase(mut self: ClientGuard<Self>, id: Id, vector: DynamicVector) -> ClientGuard<Vbase> {
-        let packet = RpcPacket::Vbase { id, vector };
+    pub fn vbase(
+        mut self: ClientGuard<Self>,
+        id: Id,
+        vbase: (DynamicVector, usize),
+    ) -> ClientGuard<Vbase> {
+        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
@@ -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
@@ -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,
     },
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -40,6 +40,6 @@ pub enum RpcPacket { @@
         },
         Vbase {
             id: Id,
-            vector: DynamicVector,
+            vbase: (DynamicVector, usize),
         },
     }