feat!: add minimum probes and maximum probes to IVF search (#3903)

Currently IVF search takes an `nprobes` parameter that selects how many
partitions are searched. This can be a problem with highly selective
prefilters because the partitions closest to the query vector might not
have any results. Setting a larger `nprobes` value will lead to more
results being returned but all queries will be more expensive.

I think we will eventually want to work around this problem through a
combination of column statistics, selectivity estimation, and partition
bloom filters. However, in the short term, this PR splits `nprobes` into
`minimum_nprobes` and `maximum_nprobes`. At first we will search
`minimum_nprobes`, like we do today. If we don't find enough results
after searching those partitions then we will search up to
`maximum_nprobes` partitions until we do.

This PR also changes the default nprobes to 20. The previous default was
1 which was unlikely to give good recall and could set misleading
expectations. The new default of 20 is the same as the default in
lancedb.

We also consider the prefilter max_results, which will allow us to stop
early if we have found all results matching the prefilter.

Also, if the number of prefilter results is very small, we will skip the
late search entirely and just return all prefiltered row ids to be
sorted in the refine stage.

The approach is slightly complex but one of my primary goals is to avoid
impacting our current search performance which relies on our ability to
select partitions and execute the prefilter in parallel (so, for
example, we won't know the actual selectivity until after we've already
loaded several partitions).

BREAKING CHANGE: the default for nprobes is now 20 instead of 1

Author: westonpace

java/core/lance-jni/src/blocking_scanner.rs

@@ -190,8 +190,13 @@ fn inner_create_scanner<'local>(
         let k = env.get_int_as_usize_from_method(&java_obj, "getK")?;
         let _ = scanner.nearest(&column, &key, k);
 
-        let nprobes = env.get_int_as_usize_from_method(&java_obj, "getNprobes")?;
-        scanner.nprobs(nprobes);
+        let minimum_nprobes = env.get_int_as_usize_from_method(&java_obj, "getMinimumNprobes")?;
+        scanner.minimum_nprobes(minimum_nprobes);
+
+        let maximum_nprobes = env.get_optional_usize_from_method(&java_obj, "getMaximumNprobes")?;
+        if let Some(maximum_nprobes) = maximum_nprobes {
+            scanner.maximum_nprobes(maximum_nprobes);
+        }
 
         if let Some(ef) = env.get_optional_usize_from_method(&java_obj, "getEf")? {
             scanner.ef(ef);

java/core/lance-jni/src/utils.rs

@@ -99,7 +99,8 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result<Option<Query>>
         let key = Arc::new(Float32Array::from(key_array));
 
         let k = env.get_int_as_usize_from_method(&java_obj, "getK")?;
-        let nprobes = env.get_int_as_usize_from_method(&java_obj, "getNprobes")?;
+        let minimum_nprobes = env.get_int_as_usize_from_method(&java_obj, "getMinimumNprobes")?;
+        let maximum_nprobes = env.get_optional_usize_from_method(&java_obj, "getMaximumNprobes")?;
 
         let ef = env.get_optional_usize_from_method(&java_obj, "getEf")?;
 
@@ -120,7 +121,8 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result<Option<Query>>
             k,
             lower_bound: None,
             upper_bound: None,
-            nprobes,
+            minimum_nprobes,
+            maximum_nprobes,
             ef,
             refine_factor,
             metric_type: distance_type,

java/core/src/main/java/com/lancedb/lance/ipc/Query.java

@@ -25,7 +25,8 @@ public class Query {
   private final String column;
   private final float[] key;
   private final int k;
-  private final int nprobes;
+  private final int minimumNprobes;
+  private final Optional<Integer> maximumNprobes;
   private final Optional<Integer> ef;
   private final Optional<Integer> refineFactor;
   private final DistanceType distanceType;
@@ -36,9 +37,15 @@ private Query(Builder builder) {
     Preconditions.checkArgument(!builder.column.isEmpty(), "Column must not be empty");
     this.key = Preconditions.checkNotNull(builder.key, "Key must be set");
     Preconditions.checkArgument(builder.k > 0, "K must be greater than 0");
-    Preconditions.checkArgument(builder.nprobes > 0, "Nprobes must be greater than 0");
+    Preconditions.checkArgument(
+        builder.minimumNprobes > 0, "Minimum Nprobes must be greater than 0");
+    Preconditions.checkArgument(
+        !builder.maximumNprobes.isPresent()
+            || builder.maximumNprobes.get() >= builder.minimumNprobes,
+        "Maximum Nprobes must be greater than minimum Nprobes");
     this.k = builder.k;
-    this.nprobes = builder.nprobes;
+    this.minimumNprobes = builder.minimumNprobes;
+    this.maximumNprobes = builder.maximumNprobes;
     this.ef = builder.ef;
     this.refineFactor = builder.refineFactor;
     this.distanceType = Preconditions.checkNotNull(builder.distanceType, "Metric type must be set");
@@ -57,8 +64,12 @@ public int getK() {
     return k;
   }
 
-  public int getNprobes() {
-    return nprobes;
+  public int getMinimumNprobes() {
+    return minimumNprobes;
+  }
+
+  public Optional<Integer> getMaximumNprobes() {
+    return maximumNprobes;
   }
 
   public Optional<Integer> getEf() {
@@ -83,7 +94,8 @@ public String toString() {
         .append("column", column)
         .append("key", key)
         .append("k", k)
-        .append("nprobes", nprobes)
+        .append("minimumNprobes", minimumNprobes)
+        .append("maximumNprobes", maximumNprobes.orElse(null))
         .append("ef", ef.orElse(null))
         .append("refineFactor", refineFactor.orElse(null))
         .append("distanceType", distanceType)
@@ -95,7 +107,8 @@ public static class Builder {
     private String column;
     private float[] key;
     private int k = 10;
-    private int nprobes = 1;
+    private int minimumNprobes = 20;
+    private Optional<Integer> maximumNprobes = Optional.empty();
     private Optional<Integer> ef = Optional.empty();
     private Optional<Integer> refineFactor = Optional.empty();
     private DistanceType distanceType = DistanceType.L2;
@@ -137,11 +150,46 @@ public Builder setK(int k) {
     /**
      * Sets the number of probes to load and search.
      *
+     * <p>This is a convenience method that sets both the minimum and maximum number of probes to
+     * the same value.
+     *
      * @param nprobes The number of probes.
      * @return The Builder instance for method chaining.
      */
     public Builder setNprobes(int nprobes) {
-      this.nprobes = nprobes;
+      this.minimumNprobes = nprobes;
+      this.maximumNprobes = Optional.of(nprobes);
+      return this;
+    }
+
+    /**
+     * Sets the minimum number of partitions to search.
+     *
+     * <p>This many partitions will always be loaded and searched on the query. Increasing this
+     * number can improve recall at the cost of latency.
+     *
+     * @param minimumNprobes The minimum number of partitions to search.
+     * @return The Builder instance for method chaining.
+     */
+    public Builder setMinimumNprobes(int minimumNprobes) {
+      this.minimumNprobes = minimumNprobes;
+      return this;
+    }
+
+    /**
+     * Sets the maximum number of partitions to search.
+     *
+     * <p>These partitions will only be loaded and searched if we have not found the desired number
+     * of results after searching the minimum number of partitions. Increasing this number can avoid
+     * false negatives on queries with a highly selective prefilter. This setting does not affect
+     * the recall of the query and will only affect the latency if the prefilter is highly
+     * selective.
+     *
+     * @param maximumNprobes The maximum number of partitions to search.
+     * @return The Builder instance for method chaining.
+     */
+    public Builder setMaximumNprobes(int maximumNprobes) {
+      this.maximumNprobes = Optional.of(maximumNprobes);
       return this;
     }
 

python/python/lance/dataset.py

@@ -404,7 +404,8 @@ def scanner(
                     "column": <embedding col name>,
                     "q": <query vector as pa.Float32Array>,
                     "k": 10,
-                    "nprobes": 1,
+                    "minimum_nprobes": 20,
+                    "maximum_nprobes": 50,
                     "refine_factor": 1
                 }
 
@@ -643,7 +644,8 @@ def to_table(
                     "q": <query vector as pa.Float32Array>,
                     "k": 10,
                     "metric": "cosine",
-                    "nprobes": 1,
+                    "minimum_nprobes": 20,
+                    "maximum_nprobes": 50,
                     "refine_factor": 1
                 }
 
@@ -3402,6 +3404,8 @@ def nearest(
         k: Optional[int] = None,
         metric: Optional[str] = None,
         nprobes: Optional[int] = None,
+        minimum_nprobes: Optional[int] = None,
+        maximum_nprobes: Optional[int] = None,
         refine_factor: Optional[int] = None,
         use_index: bool = True,
         ef: Optional[int] = None,
@@ -3435,6 +3439,26 @@ def nearest(
             raise ValueError(f"Nearest-K must be > 0 but got {k}")
         if nprobes is not None and int(nprobes) <= 0:
             raise ValueError(f"Nprobes must be > 0 but got {nprobes}")
+        if minimum_nprobes is not None and int(minimum_nprobes) < 0:
+            raise ValueError(f"Minimum nprobes must be >= 0 but got {minimum_nprobes}")
+        if maximum_nprobes is not None and int(maximum_nprobes) < 0:
+            raise ValueError(f"Maximum nprobes must be >= 0 but got {maximum_nprobes}")
+
+        if nprobes is not None:
+            if minimum_nprobes is not None or maximum_nprobes is not None:
+                raise ValueError(
+                    "nprobes cannot be set in combination with minimum_nprobes or "
+                    "maximum_nprobes"
+                )
+            else:
+                minimum_nprobes = nprobes
+                maximum_nprobes = nprobes
+        if (
+            minimum_nprobes is not None
+            and maximum_nprobes is not None
+            and minimum_nprobes > maximum_nprobes
+        ):
+            raise ValueError("minimum_nprobes must be <= maximum_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:
@@ -3446,7 +3470,8 @@ def nearest(
             "q": q,
             "k": k,
             "metric": metric,
-            "nprobes": nprobes,
+            "minimum_nprobes": minimum_nprobes,
+            "maximum_nprobes": maximum_nprobes,
             "refine_factor": refine_factor,
             "use_index": use_index,
             "ef": ef,

python/python/tests/test_vector_index.py

@@ -14,8 +14,6 @@
 import pytest
 from lance import LanceFragment
 from lance.dataset import VectorIndexReader
-
-torch = pytest.importorskip("torch")
 from lance.util import validate_vector_index  # noqa: E402
 from lance.vector import vec_to_table  # noqa: E402
 
@@ -255,6 +253,8 @@ def test_index_with_nans(tmp_path):
 
 
 def test_torch_index_with_nans(tmp_path):
+    torch = pytest.importorskip("torch")
+
     # 1024 rows, the entire table should be sampled
     tbl = create_table(nvec=1000, nans=24)
 
@@ -271,6 +271,8 @@ def test_torch_index_with_nans(tmp_path):
 
 
 def test_index_with_no_centroid_movement(tmp_path):
+    torch = pytest.importorskip("torch")
+
     # this test makes the centroids essentially [1..]
     # this makes sure the early stop condition in the index building code
     # doesn't do divide by zero
@@ -343,6 +345,10 @@ def test_create_index_using_cuda(tmp_path, nullify):
 
 
 def test_create_index_unsupported_accelerator(tmp_path):
+    # Even attempting to use an accelerator will trigger torch import
+    # so make sure it's available
+    pytest.importorskip("torch")
+
     tbl = create_table()
     dataset = lance.write_dataset(tbl, tmp_path)
     with pytest.raises(ValueError):
@@ -896,6 +902,7 @@ def query_index(ds, ntimes, q=None):
                 nearest={
                     "column": "vector",
                     "q": q if q is not None else rng.standard_normal(ndim),
+                    "minimum_nprobes": 1,
                 },
             )
 
@@ -1033,6 +1040,8 @@ def test_dynamic_projection_with_vectors_index(tmp_path: Path):
 
 
 def test_index_cast_centroids(tmp_path):
+    torch = pytest.importorskip("torch")
+
     tbl = create_table(nvec=1000)
 
     dataset = lance.write_dataset(tbl, tmp_path)
@@ -1260,3 +1269,53 @@ def test_vector_index_with_prefilter_and_scalar_index(indexed_dataset):
         prefilter=True,
     )
     assert len(res) == 10
+
+
+def test_vector_index_with_nprobes(indexed_dataset):
+    res = indexed_dataset.scanner(
+        nearest={
+            "column": "vector",
+            "q": np.random.randn(128),
+            "k": 10,
+            "nprobes": 7,
+        }
+    ).explain_plan()
+
+    assert "minimum_nprobes=7" in res
+    assert "maximum_nprobes=Some(7)" in res
+
+    res = indexed_dataset.scanner(
+        nearest={
+            "column": "vector",
+            "q": np.random.randn(128),
+            "k": 10,
+            "minimum_nprobes": 7,
+        }
+    ).explain_plan()
+
+    assert "minimum_nprobes=7" in res
+    assert "maximum_nprobes=None" in res
+
+    res = indexed_dataset.scanner(
+        nearest={
+            "column": "vector",
+            "q": np.random.randn(128),
+            "k": 10,
+            "minimum_nprobes": 7,
+            "maximum_nprobes": 10,
+        }
+    ).explain_plan()
+
+    assert "minimum_nprobes=7" in res
+    assert "maximum_nprobes=Some(10)" in res
+
+    res = indexed_dataset.scanner(
+        nearest={
+            "column": "vector",
+            "q": np.random.randn(128),
+            "k": 10,
+            "maximum_nprobes": 30,
+        }
+    ).analyze_plan()
+
+    print(res)

python/src/dataset.rs

@@ -92,7 +92,7 @@ pub mod commit;
 pub mod optimize;
 pub mod stats;
 
-const DEFAULT_NPROBS: usize = 1;
+const DEFAULT_NPROBS: usize = 20;
 const DEFAULT_INDEX_CACHE_SIZE: usize = 256;
 const DEFAULT_METADATA_CACHE_SIZE: usize = 256;
 
@@ -746,15 +746,41 @@ impl Dataset {
                 10
             };
 
-            let nprobes: usize = if let Some(nprobes) = nearest.get_item("nprobes")? {
-                if nprobes.is_none() {
-                    DEFAULT_NPROBS
-                } else {
-                    nprobes.extract()?
+            let mut minimum_nprobes = DEFAULT_NPROBS;
+            let mut maximum_nprobes = None;
+
+            if let Some(nprobes) = nearest.get_item("nprobes")? {
+                if !nprobes.is_none() {
+                    minimum_nprobes = nprobes.extract()?;
+                    maximum_nprobes = Some(minimum_nprobes);
                 }
-            } else {
-                DEFAULT_NPROBS
-            };
+            }
+
+            if let Some(min_nprobes) = nearest.get_item("minimum_nprobes")? {
+                if !min_nprobes.is_none() {
+                    minimum_nprobes = min_nprobes.extract()?;
+                }
+            }
+
+            if let Some(max_nprobes) = nearest.get_item("maximum_nprobes")? {
+                if !max_nprobes.is_none() {
+                    maximum_nprobes = Some(max_nprobes.extract()?);
+                }
+            }
+
+            if minimum_nprobes > maximum_nprobes.unwrap_or(usize::MAX) {
+                return Err(PyValueError::new_err(
+                    "minimum_nprobes must be <= maximum_nprobes",
+                ));
+            }
+
+            if minimum_nprobes < 1 {
+                return Err(PyValueError::new_err("minimum_nprobes must be >= 1"));
+            }
+
+            if maximum_nprobes.unwrap_or(usize::MAX) < 1 {
+                return Err(PyValueError::new_err("maximum_nprobes must be >= 1"));
+            }
 
             let metric_type: Option<MetricType> =
                 if let Some(metric) = nearest.get_item("metric")? {
@@ -813,7 +839,10 @@ impl Dataset {
             };
             scanner
                 .map(|s| {
-                    let mut s = s.nprobs(nprobes);
+                    let mut s = s.minimum_nprobes(minimum_nprobes);
+                    if let Some(maximum_nprobes) = maximum_nprobes {
+                        s = s.maximum_nprobes(maximum_nprobes);
+                    }
                     if let Some(factor) = refine_factor {
                         s = s.refine(factor);
                     }