feat: allow user to set vbase range (#206)

silver-ymz · web-flow · commit 4f60c4c824f3 · 2023-12-27T06:32:23.000Z
* feat: allow user to set vbase range

Signed-off-by: silver-ymz &lt;yinmingzhuo@gmail.com&gt;

* separate the enable_vbase and vbase_range settings

Signed-off-by: silver-ymz &lt;yinmingzhuo@gmail.com&gt;

---------

Signed-off-by: silver-ymz &lt;yinmingzhuo@gmail.com&gt;
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 number	Diff line change
`@@ -139,43 +139,44 @@ impl InstanceView {`
`139`	`139`	`pub fn vbase(`
`140`	`140`	`&self,`
`141`	`141`	`vector: DynamicVector,`
	`142`	`+ range: usize,`
`142`	`143`	`) -> Result<impl Iterator<Item = Pointer> + '_, FriendlyError> {`
`143`	`144`	`match (self, vector) {`
`144`	`145`	`(InstanceView::F32Cos(x), DynamicVector::F32(vector)) => {`
`145`	`146`	`if x.options.vector.dims as usize != vector.len() {`
`146`	`147`	`return Err(FriendlyError::Unmatched2);`
`147`	`148`	`}`
`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>>)`
`149`	`150`	`}`
`150`	`151`	`(InstanceView::F32Dot(x), DynamicVector::F32(vector)) => {`
`151`	`152`	`if x.options.vector.dims as usize != vector.len() {`
`152`	`153`	`return Err(FriendlyError::Unmatched2);`
`153`	`154`	`}`
`154`		`- Ok(Box::new(x.vbase(&vector)))`
	`155`	`+ Ok(Box::new(x.vbase(&vector, range)))`
`155`	`156`	`}`
`156`	`157`	`(InstanceView::F32L2(x), DynamicVector::F32(vector)) => {`
`157`	`158`	`if x.options.vector.dims as usize != vector.len() {`
`158`	`159`	`return Err(FriendlyError::Unmatched2);`
`159`	`160`	`}`
`160`		`- Ok(Box::new(x.vbase(&vector)))`
	`161`	`+ Ok(Box::new(x.vbase(&vector, range)))`
`161`	`162`	`}`
`162`	`163`	`(InstanceView::F16Cos(x), DynamicVector::F16(vector)) => {`
`163`	`164`	`if x.options.vector.dims as usize != vector.len() {`
`164`	`165`	`return Err(FriendlyError::Unmatched2);`
`165`	`166`	`}`
`166`		`- Ok(Box::new(x.vbase(&vector)))`
	`167`	`+ Ok(Box::new(x.vbase(&vector, range)))`
`167`	`168`	`}`
`168`	`169`	`(InstanceView::F16Dot(x), DynamicVector::F16(vector)) => {`
`169`	`170`	`if x.options.vector.dims as usize != vector.len() {`
`170`	`171`	`return Err(FriendlyError::Unmatched2);`
`171`	`172`	`}`
`172`		`- Ok(Box::new(x.vbase(&vector)))`
	`173`	`+ Ok(Box::new(x.vbase(&vector, range)))`
`173`	`174`	`}`
`174`	`175`	`(InstanceView::F16L2(x), DynamicVector::F16(vector)) => {`
`175`	`176`	`if x.options.vector.dims as usize != vector.len() {`
`176`	`177`	`return Err(FriendlyError::Unmatched2);`
`177`	`178`	`}`
`178`		`- Ok(Box::new(x.vbase(&vector)))`
	`179`	`+ Ok(Box::new(x.vbase(&vector, range)))`
`179`	`180`	`}`
`180`	`181`	`_ => Err(FriendlyError::Unmatched2),`
`181`	`182`	`}`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,6 @@ pub enum RpcPacket {`
`40`	`40`	`},`
`41`	`41`	`Vbase {`
`42`	`42`	`id: Id,`
`43`		`- vector: DynamicVector,`
	`43`	`+ vbase: (DynamicVector, usize),`
`44`	`44`	`},`
`45`	`45`	`}`