feat: max_scan_tuples

Signed-off-by: usamoi <[email protected]>

src/gucs/executing.rs

@@ -3,6 +3,7 @@ use std::ffi::CStr;
 
 static PROBES: GucSetting<Option<&'static CStr>> = GucSetting::<Option<&CStr>>::new(Some(c"10"));
 static EPSILON: GucSetting<f64> = GucSetting::<f64>::new(1.9);
+static MAX_SCAN_TUPLES: GucSetting<i32> = GucSetting::<i32>::new(-1);
 
 pub unsafe fn init() {
     GucRegistry::define_string_guc(
@@ -23,6 +24,16 @@ pub unsafe fn init() {
         GucContext::Userset,
         GucFlags::default(),
     );
+    GucRegistry::define_int_guc(
+        "vchordrq.max_scan_tuples",
+        "`max_scan_tuples` argument of vchordrq.",
+        "`max_scan_tuples` argument of vchordrq.",
+        &MAX_SCAN_TUPLES,
+        -1,
+        u16::MAX as _,
+        GucContext::Userset,
+        GucFlags::default(),
+    );
 }
 
 pub fn probes() -> Vec<u32> {
@@ -54,3 +65,12 @@ pub fn probes() -> Vec<u32> {
 pub fn epsilon() -> f32 {
     EPSILON.get() as f32
 }
+
+pub fn max_scan_tuples() -> Option<u32> {
+    let x = MAX_SCAN_TUPLES.get();
+    if x < 0 {
+        None
+    } else {
+        Some(x as u32)
+    }
+}

src/index/am_scan.rs

@@ -1,6 +1,7 @@
 use super::am_options::Opfamily;
 use crate::algorithm::scan::scan;
 use crate::gucs::executing::epsilon;
+use crate::gucs::executing::max_scan_tuples;
 use crate::gucs::executing::probes;
 use crate::postgres::Relation;
 use base::distance::Distance;
@@ -86,7 +87,11 @@ pub fn scan_next(scanner: &mut Scanner, relation: Relation) -> Option<(Pointer,
                 epsilon(),
             );
             *scanner = Scanner::Vbase {
-                vbase: Box::new(vbase),
+                vbase: if let Some(max_scan_tuples) = max_scan_tuples() {
+                    Box::new(vbase.take(max_scan_tuples as usize))
+                } else {
+                    Box::new(vbase)
+                },
                 threshold: *threshold,
                 recheck: *recheck,
                 opfamily: *opfamily,