feat(index): add layered RaBitQ storage and precision queries

protos/ann.proto

@@ -61,6 +61,10 @@ message VectorQueryProto {
    * indexes, such as IVF_RQ. Other index types ignore this setting.
    */
   VectorApproxMode approx_mode = 14;
+  // Empty selects full precision; otherwise sign, high or full.
+  string rq_precision = 17;
+  // Absent scans full codes; present reranks k times this positive factor.
+  optional uint32 rq_cascade_factor = 18;
   /* Which covering columns the index must materialize for this query. Absent means
    * "not computed" -- see CoveringProjection. Carried across the wire so a remote
    * executor declares the same search output schema the planner did; without it the

rust/lance-index/Cargo.toml

@@ -79,6 +79,8 @@ async-channel = "2.3.1"
 rand_distr.workspace = true
 rangemap.workspace = true
 
+uuid.workspace = true
+
 [dev-dependencies]
 approx.workspace = true
 criterion.workspace = true
@@ -92,7 +94,6 @@ test-log.workspace = true
 rstest.workspace = true
 serial_test.workspace = true
 chrono.workspace = true
-uuid.workspace = true
 
 [features]
 geo = ["dep:lance-geo", "lance-geo/geo", "dep:geoarrow-array", "dep:geoarrow-schema", "dep:geo-types"]

rust/lance-index/benches/rq.rs

@@ -252,6 +252,7 @@ fn ex_code_storage_load(c: &mut Criterion) {
         code_dim: LOAD_DIM as u32,
         num_bits: NUM_BITS,
         packed: true,
+        layered: false,
         query_estimator: RabitQueryEstimator::RawQuery,
     };
     let code_len = LOAD_DIM / 8;
@@ -465,7 +466,10 @@ fn ex_bulk_paths(c: &mut Criterion) {
                 0.0,
                 None,
                 &mut f32_scratch,
-                DistanceCalculatorOptions { approx_mode },
+                DistanceCalculatorOptions {
+                    rq_precision: Default::default(),
+                    approx_mode,
+                },
             );
             let mut dists = Vec::new();
             let mut u16_scratch = Vec::new();
@@ -613,7 +617,10 @@ fn heap_topk(c: &mut Criterion) {
             1.0,
             None,
             &mut f32_scratch,
-            DistanceCalculatorOptions { approx_mode },
+            DistanceCalculatorOptions {
+                rq_precision: Default::default(),
+                approx_mode,
+            },
         );
         let mut heap = BinaryHeap::with_capacity(TOPK_K + 1);
         let mut dists = Vec::new();
@@ -670,6 +677,7 @@ fn build_transform(c: &mut Criterion) {
                     code_dim: dim as u32,
                     num_bits,
                     packed: false,
+                    layered: false,
                     query_estimator: RabitQueryEstimator::RawQuery,
                 };
                 let rq = RabitQuantizer::try_from(

rust/lance-index/protos-cache/cache.proto

@@ -217,4 +217,5 @@ message RabitPartitionHeader {
   optional bytes fast_rotation_signs = 5;
   // Estimator the RabitQ storage uses at query time (residual vs raw query).
   RabitQueryEstimator query_estimator = 6;
+  bool layered = 7;
 }

rust/lance-index/src/vector.rs

@@ -219,6 +219,10 @@ pub struct Query {
     /// This currently only affects RQ-quantized vector indexes, such as IVF_RQ.
     /// Other index types ignore this setting.
     pub approx_mode: ApproxMode,
+    /// Precision of a layered RaBitQ query; defaults to full.
+    pub rq_precision: bq::layered::RQPrecision,
+    /// Candidate oversampling for a high-to-full cascade. None scans full codes.
+    pub rq_cascade_factor: Option<u32>,
 }
 
 impl Query {

rust/lance-index/src/vector/bq.rs

@@ -22,6 +22,7 @@ pub(crate) mod dist_table_quant;
 pub mod ex_dot;
 pub mod layered;
 pub(crate) mod pairwise;
+mod plane_cache;
 pub mod prune;
 pub mod rotation;
 pub mod storage;
@@ -124,6 +125,8 @@ pub struct RQBuildParams {
     /// rotates vectors identically. This is transient build-time state and is never
     /// persisted to the `RabitQuantization` params proto.
     pub rotation: Option<RabitQuantizationMetadata>,
+    /// Store prefix code planes; supported for 3 through 9 bits. Default off.
+    pub layered: bool,
 }
 
 pub fn validate_rq_num_bits(num_bits: u8) -> Result<()> {
@@ -160,11 +163,18 @@ pub fn rabit_ex_code_bytes(rotated_dim: usize, ex_bits: u8) -> Result<usize> {
 }
 
 impl RQBuildParams {
+    /// Opt into the layered format. Requires a reader with layered RQ support.
+    pub fn with_layered(mut self, layered: bool) -> Self {
+        self.layered = layered;
+        self
+    }
+
     pub fn new(num_bits: u8) -> Self {
         Self {
             num_bits,
             rotation_type: RQRotationType::default(),
             rotation: None,
+            layered: false,
         }
     }
 
@@ -173,6 +183,7 @@ impl RQBuildParams {
             num_bits,
             rotation_type,
             rotation: None,
+            layered: false,
         }
     }
 }
@@ -182,7 +193,7 @@ impl From<&RQBuildParams> for RabitQuantization {
         use crate::pb::vector_index_details::rabit_quantization::{RotationType, RowLayout};
         Self {
             num_bits: value.num_bits as u32,
-            layered: false,
+            layered: value.layered,
             row_layout: RowLayout::Columns as i32,
             rotation_type: match value.rotation_type {
                 RQRotationType::Fast => RotationType::Fast as i32,
@@ -204,6 +215,7 @@ impl Default for RQBuildParams {
             num_bits: RABIT_DEFAULT_NUM_BITS,
             rotation_type: RQRotationType::default(),
             rotation: None,
+            layered: false,
         }
     }
 }

rust/lance-index/src/vector/bq/builder.rs

@@ -63,11 +63,8 @@ pub struct RabitQuantizer {
 pub(crate) struct RabitQuantizedBatch {
     pub binary_codes: ArrayRef,
     pub ex_codes: Option<ArrayRef>,
-    #[cfg(test)]
     pub ex_res_dot_dists: Option<Vec<f32>>,
-    #[cfg(test)]
     pub rotated_residuals: Option<Vec<f32>>,
-    #[cfg(test)]
     pub ex_code_values: Option<Vec<u8>>,
 }
 
@@ -542,11 +539,8 @@ impl RabitQuantizer {
             } else {
                 None
             },
-            #[cfg(test)]
             rotated_residuals: None,
-            #[cfg(test)]
             ex_code_values: None,
-            #[cfg(test)]
             ex_res_dot_dists: None,
         })
     }
@@ -583,6 +577,7 @@ impl RabitQuantizer {
                     code_dim: code_dim as u32,
                     num_bits,
                     packed: false,
+                    layered: false,
                     query_estimator: RabitQueryEstimator::RawQuery,
                 }
             }
@@ -594,6 +589,7 @@ impl RabitQuantizer {
                 code_dim: code_dim as u32,
                 num_bits,
                 packed: false,
+                layered: false,
                 query_estimator: RabitQueryEstimator::RawQuery,
             },
         };
@@ -1035,11 +1031,8 @@ impl RabitQuantizer {
                         .map(|array| Arc::new(array) as ArrayRef)
                 })
                 .transpose()?,
-            #[cfg(test)]
             ex_res_dot_dists,
-            #[cfg(test)]
             rotated_residuals: Some(rotated_residuals),
-            #[cfg(test)]
             ex_code_values,
         })
     }
@@ -1056,18 +1049,22 @@ impl Quantization for RabitQuantizer {
         params: &Self::BuildParams,
     ) -> Result<Self> {
         validate_rq_num_bits(params.num_bits)?;
+        if params.layered {
+            super::layered::RQLayout::try_new(params.num_bits)?;
+        }
 
         let dim = data.as_fixed_size_list().value_length() as usize;
         if !dim.is_multiple_of(u8::BITS as usize) {
             return Err(Error::invalid_input(
                 "vector dimension must be divisible by 8 for IVF_RQ",
             ));
         }
-        if let Some(q) = Self::from_supplied_rotation(params, dim)? {
+        if let Some(mut q) = Self::from_supplied_rotation(params, dim)? {
+            q.metadata.layered = params.layered;
             return Ok(q);
         }
 
-        let q = match data.as_fixed_size_list().value_type() {
+        let mut q = match data.as_fixed_size_list().value_type() {
             DataType::Float16 => Self::new_with_rotation::<Float16Type>(
                 params.num_bits,
                 data.as_fixed_size_list().value_length(),
@@ -1090,6 +1087,7 @@ impl Quantization for RabitQuantizer {
                 )));
             }
         };
+        q.metadata.layered = params.layered;
         Ok(q)
     }
 
@@ -1152,6 +1150,14 @@ impl Quantization for RabitQuantizer {
         _: lance_linalg::distance::DistanceType,
     ) -> Result<Quantizer> {
         validate_rq_num_bits(metadata.num_bits)?;
+        if metadata.layered {
+            super::layered::RQLayout::try_new(metadata.num_bits)?;
+            if metadata.query_estimator != RabitQueryEstimator::RawQuery {
+                return Err(Error::invalid_input(
+                    "layered IVF_RQ requires the raw-query estimator",
+                ));
+            }
+        }
         Ok(Quantizer::Rabit(Self {
             metadata: metadata.clone(),
         }))
@@ -1166,6 +1172,10 @@ impl Quantization for RabitQuantizer {
         if self.metadata.query_estimator == RabitQueryEstimator::RawQuery {
             fields.push(ERROR_FACTORS_FIELD.clone());
         }
+        if self.metadata.layered {
+            return super::layered::storage_fields(self.code_dim(), self.metadata.num_bits, fields)
+                .expect("layered layout validated at build");
+        }
         if let Some(ex_code_field) = rabit_ex_code_field(self.code_dim(), self.metadata.num_bits)
             .expect("RabitQ num_bits should be validated")
         {