feat(index): read and write optional RaBitQ plane-row files

Cargo.toml

@@ -215,6 +215,7 @@ serde = { version = "^1" }
 serde_json = { version = "1" }
 semver = "1.0"
 serial_test = "3"
+sha2 = "0.10"
 snafu = "0.9"
 syn = { version = "2.0.37", features = ["full"] }
 lindera = { version = "3.0.7" }

rust/lance-index-core/src/lib.rs

@@ -224,6 +224,11 @@ impl IndexType {
     /// Indices which higher version than these will be ignored for compatibility,
     /// This would happen when creating index in a newer version of Lance,
     /// but then opening the index in older version of Lance
+    ///
+    /// For [`Self::IvfRq`] this is only the highest version this build reads:
+    /// an IVF_RQ writer commits the version of the row layout it writes
+    /// (`lance_index::ivf_rq_index_version`), so a column-layout index keeps
+    /// the version older readers accept.
     pub fn version(&self) -> i32 {
         match self {
             Self::Scalar => 0,
@@ -253,7 +258,8 @@ impl IndexType {
             | Self::IvfHnswSq
             | Self::IvfHnswPq
             | Self::IvfHnswFlat => 1,
-            Self::IvfRq => 2,
+            // 2: the column layout; 3: the plane-row layout.
+            Self::IvfRq => 3,
         }
     }
 

rust/lance-index/benches/rq.rs

@@ -253,6 +253,7 @@ fn ex_code_storage_load(c: &mut Criterion) {
         num_bits: NUM_BITS,
         packed: true,
         layered: false,
+        row_layout: RQRowLayout::Columns,
         query_estimator: RabitQueryEstimator::RawQuery,
     };
     let code_len = LOAD_DIM / 8;
@@ -678,6 +679,7 @@ fn build_transform(c: &mut Criterion) {
                     num_bits,
                     packed: false,
                     layered: false,
+                    row_layout: RQRowLayout::Columns,
                     query_estimator: RabitQueryEstimator::RawQuery,
                 };
                 let rq = RabitQuantizer::try_from(

rust/lance-index/src/lib.rs

@@ -43,8 +43,23 @@ pub const INDEX_METADATA_SCHEMA_KEY: &str = "lance:index";
 /// Most vector indices should use this version unless they need to bump for a
 /// format change.
 pub const VECTOR_INDEX_VERSION: u32 = 1;
-/// Version for IVF_RQ indices.
-pub const IVF_RQ_INDEX_VERSION: u32 = 2;
+/// Version of IVF_RQ indices in the plane-row layout, and the highest IVF_RQ
+/// version this build reads.
+pub const IVF_RQ_INDEX_VERSION: u32 = 3;
+/// Version of IVF_RQ indices in the column layout, which readers that predate
+/// the plane-row layout accept.
+pub const IVF_RQ_COLUMNS_INDEX_VERSION: u32 = 2;
+
+/// The index version an IVF_RQ segment whose auxiliary file is in `row_layout`
+/// is committed with. Every writer (build, append, optimize, conversion) takes
+/// it from the layout it writes, so a column-layout index never takes the
+/// plane-row version, which older readers leave out of their usable indexes.
+pub fn ivf_rq_index_version(row_layout: vector::bq::RQRowLayout) -> i32 {
+    match row_layout {
+        vector::bq::RQRowLayout::Columns => IVF_RQ_COLUMNS_INDEX_VERSION as i32,
+        vector::bq::RQRowLayout::PlaneRows => IVF_RQ_INDEX_VERSION as i32,
+    }
+}
 
 /// The factor of threshold to trigger split / join for vector index.
 ///
@@ -108,6 +123,27 @@ mod tests {
         assert_eq!(IndexType::max_vector_version(), IVF_RQ_INDEX_VERSION);
     }
 
+    #[test]
+    fn test_ivf_rq_index_version_follows_row_layout() {
+        use crate::vector::bq::RQRowLayout;
+        assert_eq!(
+            ivf_rq_index_version(RQRowLayout::Columns),
+            IVF_RQ_COLUMNS_INDEX_VERSION as i32
+        );
+        assert_eq!(
+            ivf_rq_index_version(RQRowLayout::PlaneRows),
+            IVF_RQ_INDEX_VERSION as i32
+        );
+        // Plane rows are the read ceiling; columns keep the version readers
+        // without plane-row support accept, which still exceeds IVF_PQ's.
+        assert_eq!(
+            ivf_rq_index_version(RQRowLayout::PlaneRows),
+            IndexType::IvfRq.version()
+        );
+        assert!(ivf_rq_index_version(RQRowLayout::Columns) < IndexType::IvfRq.version());
+        assert!(ivf_rq_index_version(RQRowLayout::Columns) > IndexType::IvfPq.version());
+    }
+
     #[test]
     fn test_ivf_rq_target_partition_size() {
         assert_eq!(IndexType::IvfRq.target_partition_size(), 4096);

rust/lance-index/src/vector.rs

@@ -28,6 +28,7 @@ use v3::subindex::SubIndexType;
 
 pub mod bq;
 pub mod distributed;
+pub(crate) mod exact_buffers;
 pub mod flat;
 pub mod graph;
 pub mod hnsw;

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

@@ -14,6 +14,7 @@ use lance_core::{Error, Result};
 use num_traits::Float;
 use serde::{Deserialize, Serialize};
 
+pub use crate::vector::bq::storage::RQRowLayout;
 use crate::vector::bq::storage::RabitQuantizationMetadata;
 use crate::vector::quantizer::QuantizerBuildParams;
 
@@ -27,6 +28,7 @@ mod lazy_tests;
 pub(crate) mod pairwise;
 pub mod partition_codes;
 mod plane_cache;
+pub mod plane_rows;
 pub mod prune;
 pub mod raw_body;
 pub mod resident;
@@ -133,6 +135,10 @@ pub struct RQBuildParams {
     pub rotation: Option<RabitQuantizationMetadata>,
     /// Store prefix code planes; supported for 3 through 9 bits. Default off.
     pub layered: bool,
+    /// How the auxiliary file stores each row's fields. Default columns; the
+    /// plane-row layout is committed as index version 3, which readers
+    /// without plane-row support leave out of their usable indexes.
+    pub row_layout: RQRowLayout,
 }
 
 pub fn validate_rq_num_bits(num_bits: u8) -> Result<()> {
@@ -175,12 +181,20 @@ impl RQBuildParams {
         self
     }
 
+    /// Write the auxiliary file in `row_layout`. Plane rows require a reader
+    /// with plane-row support.
+    pub fn with_row_layout(mut self, row_layout: RQRowLayout) -> Self {
+        self.row_layout = row_layout;
+        self
+    }
+
     pub fn new(num_bits: u8) -> Self {
         Self {
             num_bits,
             rotation_type: RQRotationType::default(),
             rotation: None,
             layered: false,
+            row_layout: RQRowLayout::Columns,
         }
     }
 
@@ -190,6 +204,7 @@ impl RQBuildParams {
             rotation_type,
             rotation: None,
             layered: false,
+            row_layout: RQRowLayout::Columns,
         }
     }
 }
@@ -200,7 +215,10 @@ impl From<&RQBuildParams> for RabitQuantization {
         Self {
             num_bits: value.num_bits as u32,
             layered: value.layered,
-            row_layout: RowLayout::Columns as i32,
+            row_layout: match value.row_layout {
+                RQRowLayout::Columns => RowLayout::Columns as i32,
+                RQRowLayout::PlaneRows => RowLayout::PlaneRows as i32,
+            },
             rotation_type: match value.rotation_type {
                 RQRotationType::Fast => RotationType::Fast as i32,
                 RQRotationType::Matrix => RotationType::Matrix as i32,
@@ -222,6 +240,7 @@ impl Default for RQBuildParams {
             rotation_type: RQRotationType::default(),
             rotation: None,
             layered: false,
+            row_layout: RQRowLayout::Columns,
         }
     }
 }
@@ -268,6 +287,21 @@ mod tests {
         assert_eq!(RQBuildParams::default().num_bits, 5);
     }
 
+    #[test]
+    fn test_rq_build_params_row_layout_reaches_details() {
+        use crate::pb::vector_index_details::rabit_quantization::RowLayout;
+        assert_eq!(RQBuildParams::default().row_layout, RQRowLayout::Columns);
+        assert_eq!(RQBuildParams::new(7).row_layout, RQRowLayout::Columns);
+        let columns = RabitQuantization::from(&RQBuildParams::new(7));
+        assert_eq!(columns.row_layout(), RowLayout::Columns);
+        let params = RQBuildParams::new(7)
+            .with_layered(true)
+            .with_row_layout(RQRowLayout::PlaneRows);
+        let plane_rows = RabitQuantization::from(&params);
+        assert_eq!(plane_rows.row_layout(), RowLayout::PlaneRows);
+        assert!(plane_rows.layered);
+    }
+
     #[test]
     fn test_rabit_num_bits_validation() {
         validate_rq_num_bits(1).unwrap();

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

@@ -17,8 +17,8 @@ use rand_distr::Distribution;
 use rayon::prelude::*;
 
 use crate::vector::bq::storage::{
-    RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RabitQuantizationMetadata, RabitQuantizationStorage,
-    RabitQueryEstimator, rabit_binary_code_field, rabit_ex_code_field,
+    RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RQRowLayout, RabitQuantizationMetadata,
+    RabitQuantizationStorage, RabitQueryEstimator, rabit_binary_code_field, rabit_ex_code_field,
 };
 use crate::vector::bq::transform::{
     ADD_FACTORS_FIELD, ERROR_FACTORS_FIELD, EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD,
@@ -578,6 +578,7 @@ impl RabitQuantizer {
                     num_bits,
                     packed: false,
                     layered: false,
+                    row_layout: RQRowLayout::Columns,
                     query_estimator: RabitQueryEstimator::RawQuery,
                 }
             }
@@ -590,6 +591,7 @@ impl RabitQuantizer {
                 num_bits,
                 packed: false,
                 layered: false,
+                row_layout: RQRowLayout::Columns,
                 query_estimator: RabitQueryEstimator::RawQuery,
             },
         };
@@ -1061,6 +1063,7 @@ impl Quantization for RabitQuantizer {
         }
         if let Some(mut q) = Self::from_supplied_rotation(params, dim)? {
             q.metadata.layered = params.layered;
+            q.metadata.row_layout = params.row_layout;
             return Ok(q);
         }
 
@@ -1088,6 +1091,7 @@ impl Quantization for RabitQuantizer {
             }
         };
         q.metadata.layered = params.layered;
+        q.metadata.row_layout = params.row_layout;
         Ok(q)
     }
 
@@ -1168,25 +1172,38 @@ impl Quantization for RabitQuantizer {
     }
 
     fn extra_fields(&self) -> Vec<Field> {
-        let mut fields = vec![ADD_FACTORS_FIELD.clone(), SCALE_FACTORS_FIELD.clone()];
-        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")
-        {
-            fields.push(ex_code_field);
-            fields.push(EX_ADD_FACTORS_FIELD.clone());
-            fields.push(EX_SCALE_FACTORS_FIELD.clone());
-        }
+        let mut fields = rabit_storage_fields(&self.metadata)
+            .expect("RabitQ num_bits and layered layout validated at build");
+        // The first is the binary code field, `Self::field`.
+        fields.remove(0);
         fields
     }
 }
 
+/// The fields an IVF_RQ storage file with `metadata` stores after the row id
+/// in the column layout, in their order: the binary codes and their factors,
+/// then the ex codes and factors, as plain or layered planes.
+pub fn rabit_storage_fields(metadata: &RabitQuantizationMetadata) -> Result<Vec<Field>> {
+    let rotated_dim = metadata.rotated_dim();
+    let mut fields = vec![
+        rabit_binary_code_field(rotated_dim),
+        ADD_FACTORS_FIELD.clone(),
+        SCALE_FACTORS_FIELD.clone(),
+    ];
+    if metadata.query_estimator == RabitQueryEstimator::RawQuery {
+        fields.push(ERROR_FACTORS_FIELD.clone());
+    }
+    if metadata.layered {
+        return super::layered::storage_fields(rotated_dim, metadata.num_bits, fields);
+    }
+    if let Some(ex_code_field) = rabit_ex_code_field(rotated_dim, metadata.num_bits)? {
+        fields.push(ex_code_field);
+        fields.push(EX_ADD_FACTORS_FIELD.clone());
+        fields.push(EX_SCALE_FACTORS_FIELD.clone());
+    }
+    Ok(fields)
+}
+
 impl TryFrom<Quantizer> for RabitQuantizer {
     type Error = Error;
 
@@ -1693,6 +1710,27 @@ mod tests {
         assert_eq!(metadata.num_bits, 3);
         assert_eq!(metadata.rotation_type, RQRotationType::Fast);
         assert_eq!(metadata.fast_rotation_signs, supplied_signs);
+        assert_eq!(metadata.row_layout, RQRowLayout::Columns);
+    }
+
+    /// The build params' row layout reaches the quantizer's metadata, which
+    /// the writer reads it from, whether the rotation is new or supplied.
+    #[test]
+    fn test_rabit_quantizer_takes_row_layout_from_params() {
+        let vectors = Float32Array::from(vec![0.0f32; 4 * 32]);
+        let fsl = FixedSizeListArray::try_new_from_values(vectors, 32).unwrap();
+        let params = RQBuildParams::new(7)
+            .with_layered(true)
+            .with_row_layout(RQRowLayout::PlaneRows);
+        let quantizer = RabitQuantizer::build(&fsl, DistanceType::L2, &params).unwrap();
+        assert_eq!(quantizer.metadata_ref().row_layout, RQRowLayout::PlaneRows);
+        assert_eq!(quantizer.metadata(None).row_layout, RQRowLayout::PlaneRows);
+
+        // A rotation supplied from a plane-row index takes the params' layout.
+        let mut params = RQBuildParams::new(7).with_layered(true);
+        params.rotation = Some(quantizer.metadata(None));
+        let quantizer = RabitQuantizer::build(&fsl, DistanceType::L2, &params).unwrap();
+        assert_eq!(quantizer.metadata_ref().row_layout, RQRowLayout::Columns);
     }
 
     #[test]