feat: support to set ef for search from python (#2569)

it's related to lancedb/lancedb#1428 --------- Signed-off-by: BubbleCal <[email protected]>
lancedb · Jul 9, 2024 · bf10c8d · bf10c8d
1 parent 49de38e
commit bf10c8d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
@@ -2166,6 +2166,7 @@ def nearest(
         nprobes: Optional[int] = None,
         refine_factor: Optional[int] = None,
         use_index: bool = True,
+        ef: Optional[int] = None,
     ) -> ScannerBuilder:
         q = _coerce_query_vector(q)
 
@@ -2192,6 +2193,10 @@ def nearest(
             raise ValueError(f"Nprobes must be > 0 but got {nprobes}")
         if refine_factor is not None and int(refine_factor) < 1:
             raise ValueError(f"Refine factor must be 1 or more got {refine_factor}")
+        if ef is not None and int(ef) <= 0:
+            # `ef` should be >= `k`, but `k` could be None so we can't check it here
+            # the rust code will check it
+            raise ValueError(f"ef must be > 0 but got {ef}")
         self._nearest = {
             "column": column,
             "q": q,
@@ -2200,6 +2205,7 @@ def nearest(
             "nprobes": nprobes,
             "refine_factor": refine_factor,
             "use_index": use_index,
+            "ef": ef,
         }
         return self
 

diff --git a/python/src/dataset.rs b/python/src/dataset.rs
@@ -579,6 +579,16 @@ impl Dataset {
                 true
             };
 
+            let ef: Option<usize> = if let Some(ef) = nearest.get_item("ef")? {
+                if ef.is_none() {
+                    None
+                } else {
+                    PyAny::downcast::<PyLong>(ef)?.extract()?
+                }
+            } else {
+                None
+            };
+
             scanner
                 .nearest(column.as_str(), &q, k)
                 .map(|s| {
@@ -589,6 +599,9 @@ impl Dataset {
                     if let Some(m) = metric_type {
                         s = s.distance_metric(m);
                     }
+                    if let Some(ef) = ef {
+                        s = s.ef(ef);
+                    }
                     s.use_index(use_index);
                     s
                 })