feat(index): add opt-in quantized vector refinement

java/lance-jni/src/blocking_scanner.rs

@@ -403,6 +403,12 @@ pub(crate) fn build_scanner_with_options<'a>(
             scanner.refine(refine_factor);
         }
 
+        if let Some(factor) =
+            env.get_optional_u32_from_method(&java_obj, "getQuantizedRefineFactor")?
+        {
+            scanner.quantized_refine(factor);
+        }
+
         if let Some(distance_type_str) =
             env.get_optional_string_from_method(&java_obj, "getDistanceTypeString")?
         {

java/lance-jni/src/utils.rs

@@ -300,6 +300,10 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result<Option<Query>>
         let ef = env.get_optional_usize_from_method(&java_obj, "getEf")?;
 
         let refine_factor = env.get_optional_u32_from_method(&java_obj, "getRefineFactor")?;
+        // Query filters cannot represent scanner execution stages, but read the
+        // property here so JNI contract tests cover the Java method signature.
+        let _quantized_refine_factor =
+            env.get_optional_u32_from_method(&java_obj, "getQuantizedRefineFactor")?;
 
         let distance_type = if let Some(distance_type_str) =
             env.get_optional_string_from_method(&java_obj, "getDistanceTypeString")?

java/src/main/java/org/lance/ipc/Query.java

@@ -29,6 +29,7 @@ public class Query {
   private final Optional<Integer> maximumNprobes;
   private final Optional<Integer> ef;
   private final Optional<Integer> refineFactor;
+  private final Optional<Integer> quantizedRefineFactor;
   private final Optional<DistanceType> distanceType;
   private final boolean useIndex;
   private final int queryParallelism;
@@ -50,6 +51,10 @@ private Query(Builder builder) {
     this.maximumNprobes = builder.maximumNprobes;
     this.ef = builder.ef;
     this.refineFactor = builder.refineFactor;
+    Preconditions.checkArgument(
+        !builder.quantizedRefineFactor.isPresent() || builder.quantizedRefineFactor.get() > 0,
+        "Quantized refine factor must be greater than 0");
+    this.quantizedRefineFactor = builder.quantizedRefineFactor;
     this.distanceType = builder.distanceType;
     this.useIndex = builder.useIndex;
     this.queryParallelism = builder.queryParallelism;
@@ -84,6 +89,10 @@ public Optional<Integer> getRefineFactor() {
     return refineFactor;
   }
 
+  public Optional<Integer> getQuantizedRefineFactor() {
+    return quantizedRefineFactor;
+  }
+
   public Optional<DistanceType> getDistanceType() {
     return distanceType;
   }
@@ -118,6 +127,7 @@ public String toString() {
         .add("maximumNprobes", maximumNprobes.orElse(null))
         .add("ef", ef.orElse(null))
         .add("refineFactor", refineFactor.orElse(null))
+        .add("quantizedRefineFactor", quantizedRefineFactor.orElse(null))
         .add("distanceType", distanceType.orElse(null))
         .add("useIndex", useIndex)
         .add("queryParallelism", queryParallelism)
@@ -133,6 +143,7 @@ public static class Builder {
     private Optional<Integer> maximumNprobes = Optional.empty();
     private Optional<Integer> ef = Optional.empty();
     private Optional<Integer> refineFactor = Optional.empty();
+    private Optional<Integer> quantizedRefineFactor = Optional.empty();
     private Optional<DistanceType> distanceType = Optional.empty();
     private boolean useIndex = true;
     private int queryParallelism = 0;
@@ -244,6 +255,20 @@ public Builder setRefineFactor(int refineFactor) {
       return this;
     }
 
+    /**
+     * Reranks overfetched IVF_RQ candidates with the index's most accurate stored RQ data.
+     *
+     * <p>This stage does not read original vectors. It can be combined with {@link
+     * #setRefineFactor(int)}, in which case exact refinement remains the final stage.
+     *
+     * @param quantizedRefineFactor One-bit candidate overfetch factor; must be greater than zero.
+     * @return The Builder instance for method chaining.
+     */
+    public Builder setQuantizedRefineFactor(int quantizedRefineFactor) {
+      this.quantizedRefineFactor = Optional.of(quantizedRefineFactor);
+      return this;
+    }
+
     /**
      * Sets the distance metric type.
      *

java/src/test/java/org/lance/JNITest.java

@@ -31,6 +31,7 @@
 import java.util.Optional;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class JNITest {
@@ -54,21 +55,33 @@ public void testQuery() {
     Query defaultQuery =
         new Query.Builder().setColumn("column").setKey(new float[] {1.0f, 2.0f, 3.0f}).build();
     assertEquals(ApproxMode.NORMAL, defaultQuery.getApproxMode());
+    assertFalse(defaultQuery.getQuantizedRefineFactor().isPresent());
+
+    Query configuredQuery =
+        new Query.Builder()
+            .setColumn("column")
+            .setKey(new float[] {1.0f, 2.0f, 3.0f})
+            .setK(10)
+            .setNprobes(20)
+            .setEf(30)
+            .setRefineFactor(40)
+            .setQuantizedRefineFactor(4)
+            .setDistanceType(DistanceType.L2)
+            .setUseIndex(true)
+            .setQueryParallelism(-1)
+            .setApproxMode(ApproxMode.ACCURATE)
+            .build();
+    assertEquals(Optional.of(4), configuredQuery.getQuantizedRefineFactor());
+    JniTestHelper.parseQuery(Optional.of(configuredQuery));
 
-    JniTestHelper.parseQuery(
-        Optional.of(
+    assertThrows(
+        IllegalArgumentException.class,
+        () ->
             new Query.Builder()
                 .setColumn("column")
-                .setKey(new float[] {1.0f, 2.0f, 3.0f})
-                .setK(10)
-                .setNprobes(20)
-                .setEf(30)
-                .setRefineFactor(40)
-                .setDistanceType(DistanceType.L2)
-                .setUseIndex(true)
-                .setQueryParallelism(-1)
-                .setApproxMode(ApproxMode.ACCURATE)
-                .build()));
+                .setKey(new float[] {1.0f})
+                .setQuantizedRefineFactor(0)
+                .build());
   }
 
   @Test

python/python/lance/dataset.py

@@ -7111,6 +7111,7 @@ def nearest(
         minimum_nprobes: Optional[int] = None,
         maximum_nprobes: Optional[int] = None,
         refine_factor: Optional[int] = None,
+        quantized_refine_factor: Optional[int] = None,
         use_index: bool = True,
         ef: Optional[int] = None,
         query_parallelism: Optional[int] = None,
@@ -7145,6 +7146,12 @@ def nearest(
             setting. ``fast`` favors lower latency and may reduce recall,
             ``normal`` uses the default balance, and ``accurate`` favors higher
             recall and may increase latency.
+        quantized_refine_factor: int, optional
+            For an IVF_RQ flat index, overfetch this many one-bit candidates per
+            requested result and rerank them with the index's most accurate stored
+            RQ representation. Unlike ``refine_factor``, this does not read the
+            original vectors. The two stages can be combined; exact refinement
+            remains the final stage.
         """
         self._nearest = _build_vector_search_query(
             column,
@@ -7156,6 +7163,7 @@ def nearest(
             minimum_nprobes=minimum_nprobes,
             maximum_nprobes=maximum_nprobes,
             refine_factor=refine_factor,
+            quantized_refine_factor=quantized_refine_factor,
             use_index=use_index,
             ef=ef,
             query_parallelism=query_parallelism,
@@ -8364,6 +8372,7 @@ def _build_vector_search_query(
     minimum_nprobes: Optional[int] = None,
     maximum_nprobes: Optional[int] = None,
     refine_factor: Optional[int] = None,
+    quantized_refine_factor: Optional[int] = None,
     use_index: bool = True,
     ef: Optional[int] = None,
     query_parallelism: Optional[int] = None,
@@ -8399,6 +8408,10 @@ def _build_vector_search_query(
         The maximum number of partitions to search.
     refine_factor: int, optional
         The refine factor for the search.
+    quantized_refine_factor: int, optional
+        For an IVF_RQ flat index, overfetch this many one-bit candidates per
+        requested result and rerank them using the most accurate stored RQ data.
+        This does not read original vectors.
     use_index: bool, default True
         Whether to use the index for the search.
     ef: int, optional
@@ -8479,6 +8492,10 @@ def _build_vector_search_query(
         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 quantized_refine_factor is not None and int(quantized_refine_factor) < 1:
+        raise ValueError(
+            f"Quantized refine factor must be 1 or more got {quantized_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
@@ -8509,6 +8526,7 @@ def _build_vector_search_query(
         "minimum_nprobes": minimum_nprobes,
         "maximum_nprobes": maximum_nprobes,
         "refine_factor": refine_factor,
+        "quantized_refine_factor": quantized_refine_factor,
         "use_index": use_index,
         "ef": ef,
         "query_parallelism": query_parallelism,

python/python/tests/test_vector_index.py

@@ -1341,6 +1341,17 @@ def test_create_ivf_rq_multi_bit_searches_l2_and_cosine():
         )
         assert result.num_rows == 10
 
+    quantized_refined = ds.to_table(
+        nearest={
+            "column": "vector",
+            "q": mat[0],
+            "k": 10,
+            "quantized_refine_factor": 4,
+        },
+        columns=["id"],
+    )
+    assert quantized_refined.num_rows == 10
+
     cosine_ds = lance.write_dataset(tbl, "memory://")
     cosine_ds = _assert_recall_at_least(cosine_ds, mat[1], metric="cosine")
     cosine_stats = cosine_ds.stats.index_stats("vector_idx")
@@ -2669,6 +2680,18 @@ def test_vector_index_invalid_approx_mode(indexed_dataset):
         )
 
 
+def test_vector_index_invalid_quantized_refine_factor(indexed_dataset):
+    with pytest.raises(ValueError, match="Quantized refine factor"):
+        indexed_dataset.scanner(
+            nearest={
+                "column": "vector",
+                "q": np.random.randn(128),
+                "k": 10,
+                "quantized_refine_factor": 0,
+            }
+        )
+
+
 def test_knn_deleted_rows(tmp_path):
     data = create_table()
     ds = lance.write_dataset(data, tmp_path)

python/src/dataset.rs

@@ -1534,6 +1534,7 @@ impl Dataset {
                 maximum_nprobes,
                 metric_type,
                 refine_factor,
+                quantized_refine_factor,
                 use_index,
                 ef,
                 query_parallelism,
@@ -1597,6 +1598,9 @@ impl Dataset {
                     if let Some(factor) = refine_factor {
                         s = s.refine(factor);
                     }
+                    if let Some(factor) = quantized_refine_factor {
+                        s = s.quantized_refine(factor);
+                    }
                     if let Some(m) = metric_type {
                         s = s.distance_metric(m);
                     }
@@ -5512,6 +5516,7 @@ type VectorQueryParams = (
     Option<usize>,
     Option<MetricType>,
     Option<u32>,
+    Option<u32>,
     bool,
     Option<usize>,
     i32,
@@ -5643,6 +5648,23 @@ fn vector_query_params_from_dict(
         None
     };
 
+    let quantized_refine_factor: Option<u32> =
+        if let Some(factor) = dict.get_item("quantized_refine_factor")? {
+            if factor.is_none() {
+                None
+            } else {
+                let factor: u32 = factor.extract()?;
+                if factor == 0 {
+                    return Err(PyValueError::new_err(
+                        "quantized_refine_factor must be greater than zero",
+                    ));
+                }
+                Some(factor)
+            }
+        } else {
+            None
+        };
+
     let use_index: bool = if let Some(idx) = dict.get_item("use_index")? {
         idx.extract()?
     } else {
@@ -5670,6 +5692,7 @@ fn vector_query_params_from_dict(
         maximum_nprobes,
         metric_type,
         refine_factor,
+        quantized_refine_factor,
         use_index,
         ef,
         query_parallelism,
@@ -5707,12 +5730,19 @@ impl PySearchFilter {
             maximum_nprobes,
             metric_type_opt,
             refine_factor,
+            quantized_refine_factor,
             use_index,
             ef,
             query_parallelism,
             approx_mode,
         ) = vector_query_params_from_dict(query, default_k)?;
 
+        if quantized_refine_factor.is_some() {
+            return Err(PyValueError::new_err(
+                "quantized_refine_factor is only supported on dataset scanners",
+            ));
+        }
+
         let metric_type = Some(metric_type_opt.unwrap_or(MetricType::L2));
 
         let vector_query = VectorQuery {

rust/lance-index/src/vector.rs

@@ -168,6 +168,25 @@ pub struct Query {
     pub approx_mode: ApproxMode,
 }
 
+/// Stable partition-local identity for one candidate emitted by an IVF search.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, DeepSizeOf)]
+pub struct PartitionSearchCandidate {
+    /// IVF partition containing the candidate.
+    pub partition_id: u32,
+    /// Zero-based offset inside the partition's vector storage.
+    pub offset_in_partition: u32,
+}
+
+/// Result of a partition search that also exposes identities suitable for
+/// selected-candidate re-scoring.
+#[derive(Debug, Clone)]
+pub struct PartitionSearchResult {
+    /// Ordinary `(distance, row id)` result batch.
+    pub batch: RecordBatch,
+    /// Candidate identities in the same order as rows in `batch`.
+    pub candidates: Vec<PartitionSearchCandidate>,
+}
+
 impl From<pb::VectorMetricType> for DistanceType {
     fn from(proto: pb::VectorMetricType) -> Self {
         match proto {
@@ -248,6 +267,37 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index {
         metrics: &dyn MetricsCollector,
     ) -> Result<RecordBatch>;
 
+    /// Search one partition and preserve the local offsets of emitted rows.
+    ///
+    /// Implementations may reject this operation. Quantized refinement uses it
+    /// only for flat IVF_RQ sub-indices.
+    async fn search_in_partition_with_candidates(
+        &self,
+        _partition_id: usize,
+        _query: &Query,
+        _pre_filter: Arc<dyn PreFilter>,
+        _metrics: &dyn MetricsCollector,
+        _candidate_limit: usize,
+    ) -> Result<PartitionSearchResult> {
+        Err(Error::not_supported("search_in_partition_with_candidates"))
+    }
+
+    /// Score partition-local candidate offsets with the index's most accurate
+    /// native storage representation.
+    ///
+    /// The returned rows preserve the order of `offsets_in_partition`. This
+    /// operation does not apply a pre-filter or distance bounds; callers must
+    /// validate candidate identity and apply query result shaping around it.
+    async fn score_partition_candidates(
+        &self,
+        _partition_id: usize,
+        _query: &Query,
+        _offsets_in_partition: &[u32],
+        _metrics: &dyn MetricsCollector,
+    ) -> Result<RecordBatch> {
+        Err(Error::not_supported("score_partition_candidates"))
+    }
+
     /// Asynchronously prepare a single-partition search so the CPU-heavy portion
     /// can be executed separately.
     async fn prepare_partition_search(