chore: add ef as query parameter

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

rust/lance-index/src/vector.rs

@@ -55,6 +55,10 @@ pub struct Query {
     /// The number of probes to load and search.
     pub nprobes: usize,
 
+    /// The number of candidates to reserve while searching.
+    /// this is an optional parameter for HNSW related index types.
+    pub ef: Option<usize>,
+
     /// If presented, apply a refine step.
     /// TODO: should we support fraction / float number here?
     pub refine_factor: Option<u32>,

rust/lance/src/dataset/scanner.rs

@@ -475,6 +475,7 @@ impl Scanner {
             key: key.into(),
             k,
             nprobes: 1,
+            ef: None,
             refine_factor: None,
             metric_type: MetricType::L2,
             use_index: true,

rust/lance/src/index/vector/fixture_test.rs

@@ -198,6 +198,7 @@ mod test {
                 key: Arc::new(Float32Array::from(query)),
                 k: 1,
                 nprobes: 1,
+                ef: None,
                 refine_factor: None,
                 metric_type: metric,
                 use_index: true,

rust/lance/src/index/vector/hnsw.rs

@@ -163,8 +163,16 @@ impl<Q: Quantizer + Send + Sync + 'static> VectorIndex for HNSWIndex<Q> {
             )
         };
 
-        let k = query.k * query.refine_factor.unwrap_or(1) as usize;
-        let ef = k + k / 2;
+        let refine_factor = query.refine_factor.unwrap_or(1) as usize;
+        let k = query.k * refine_factor;
+        let ef = query.ef.unwrap_or(k + k / 2);
+        if ef < k {
+            return Err(Error::Index {
+                message: "ef must be greater than or equal to k".to_string(),
+                location: location!(),
+            });
+        }
+
         let results = self.hnsw.search(
             query.key.as_primitive::<Float32Type>().as_slice(),
             k,

rust/lance/src/index/vector/ivf.rs

@@ -1771,6 +1771,7 @@ mod tests {
                     key: Arc::new(row),
                     k: 5,
                     nprobes: 1,
+                    ef: None,
                     refine_factor: None,
                     metric_type: MetricType::L2,
                     use_index: true,