feat(python): add shared RaBitQ rotation for distributed IVF_RQ builds

python/python/lance/dataset.py

@@ -3304,6 +3304,7 @@ def _create_index_impl(
         *,
         target_partition_size: Optional[int] = None,
         skip_transpose: bool = False,
+        rq_rotation: Optional[str] = None,
         require_commit: bool = True,
         **kwargs,
     ) -> Index:
@@ -3619,6 +3620,9 @@ def _create_index_impl(
         if skip_transpose:
             kwargs["skip_transpose"] = True
 
+        if rq_rotation is not None:
+            kwargs["rq_rotation"] = rq_rotation
+
         # Add fragment_ids and index_uuid to kwargs if provided for
         # distributed indexing
         if fragment_ids is not None:
@@ -3919,6 +3923,7 @@ def create_index_uncommitted(
         *,
         target_partition_size: Optional[int] = None,
         skip_transpose: bool = False,
+        rq_rotation: Optional[str] = None,
         **kwargs,
     ) -> Index:
         """
@@ -3945,6 +3950,12 @@ def create_index_uncommitted(
         requirement:
 
         - ``fragment_ids`` must be provided
+        - ``rq_rotation`` (``IVF_RQ`` only): a JSON string produced by
+          ``lance.lance.indices.build_rq_rotation``. It must be identical across all
+          workers for their segments to be mergeable, since it pins the RaBitQ
+          rotation so every segment rotates vectors the same way. If omitted, each
+          call generates its own random rotation, which is only safe for a single,
+          non-merged segment.
 
         Returns
         -------
@@ -3974,6 +3985,7 @@ def create_index_uncommitted(
             index_uuid=index_uuid,
             target_partition_size=target_partition_size,
             skip_transpose=skip_transpose,
+            rq_rotation=rq_rotation,
             require_commit=False,
             **kwargs,
         )

python/python/lance/lance/indices/__init__.pyi

@@ -67,6 +67,12 @@ def transform_vectors(
     pq_codebook: pa.Array,
     dst_uri: str,
 ): ...
+def build_rq_rotation(
+    dimension: int,
+    num_bits: int = 1,
+    rotation_type: str = "fast",
+    dtype: str = "float32",
+) -> str: ...
 
 class IndexSegmentDescription:
     uuid: str

python/python/tests/test_vector_index.py

@@ -3023,6 +3023,51 @@ def test_commit_existing_index_segments_accepts_index_metadata(tmp_path):
     assert 0 < len(results) <= 5
 
 
+def test_distributed_ivf_rq_shared_rotation(tmp_path):
+    """Two IVF_RQ segments built on separate fragments with one shared RaBitQ rotation
+    merge into a single committed, queryable index. The shared ``rq_rotation`` (from
+    ``lance.lance.indices.build_rq_rotation``) is what makes the independently built
+    segments mergeable."""
+    from lance.lance import indices
+
+    dim = 32
+    ds = _make_sample_dataset_base(
+        tmp_path, "dist_rq_merge", n_rows=512, dim=dim, max_rows_per_file=256
+    )
+    frags = ds.get_fragments()
+    assert len(frags) == 2
+
+    ivf_model = IndicesBuilder(ds, "vector").train_ivf(
+        num_partitions=2,
+        distance_type="l2",
+        sample_rate=8,
+    )
+    rq_rotation = indices.build_rq_rotation(dimension=dim, num_bits=1)
+    base_kwargs = {
+        "column": "vector",
+        "index_type": "IVF_RQ",
+        "num_partitions": 2,
+        "num_bits": 1,
+        "ivf_centroids": ivf_model.centroids,
+        "rq_rotation": rq_rotation,
+    }
+    first = ds.create_index_uncommitted(
+        **base_kwargs,
+        fragment_ids=[frags[0].fragment_id],
+    )
+    second = ds.create_index_uncommitted(
+        **base_kwargs,
+        fragment_ids=[frags[1].fragment_id],
+    )
+
+    merged = ds.merge_existing_index_segments([first, second])
+    ds = ds.commit_existing_index_segments("vector_idx", "vector", [merged])
+
+    q = np.random.rand(dim).astype(np.float32)
+    results = ds.to_table(nearest={"column": "vector", "q": q, "k": 5})
+    assert 0 < len(results) <= 5
+
+
 def test_index_segment_builder_builds_vector_segments(tmp_path):
     ds = _make_sample_dataset_base(tmp_path, "segment_builder_ds", 2000, 128)
     frags = ds.get_fragments()

python/src/dataset.rs

@@ -20,6 +20,7 @@ use blob::LanceBlobFile;
 use chrono::{Duration, TimeDelta, Utc};
 use futures::{StreamExt, TryFutureExt};
 use lance_index::vector::bq::RQBuildParams;
+use lance_index::vector::bq::storage::RabitQuantizationMetadata;
 use log::error;
 use object_store::path::Path;
 use pyo3::exceptions::{PyStopIteration, PyTypeError};
@@ -4429,6 +4430,13 @@ fn prepare_vector_index_params(
             pq_params.codebook = Some(codebook.values().clone())
         };
 
+        if let Some(r) = kwargs.get_item("rq_rotation")? {
+            let json: String = r.extract()?;
+            let meta: RabitQuantizationMetadata = serde_json::from_str(&json)
+                .map_err(|e| PyValueError::new_err(format!("Invalid rq_rotation JSON: {e}")))?;
+            rq_params.rotation = Some(meta);
+        };
+
         if let Some(version) = kwargs.get_item("index_file_version")? {
             let version: String = version.extract()?;
             index_file_version = IndexFileVersion::try_from(&version)

python/src/indices.rs

@@ -358,6 +358,79 @@ fn train_pq_model<'py>(
     codebook.to_pyarrow(py)
 }
 
+/// Mint one RaBitQ rotation and return it as a JSON string.
+///
+/// Distributed IVF_RQ builds must pin a single rotation across all workers so that
+/// independently built per-fragment segments rotate vectors identically and their
+/// binary codes remain comparable when merged. A driver calls this once and broadcasts
+/// the resulting string to every `create_index_uncommitted(..., rq_rotation=...)` call.
+///
+/// Only the "fast" rotation is supported since its sign vector is JSON-serializable, whereas
+/// the "matrix" rotation stores a dense matrix in a binary buffer that is dropped by the
+/// JSON wire format. `dtype` is accepted for API symmetry but does not affect the fast
+/// rotation.
+///
+/// # Example (Python)
+///
+/// ```python
+/// from lance.lance import indices
+///
+/// # Mint one rotation and broadcast `rot` to every worker.
+/// rot = indices.build_rq_rotation(dimension=128, num_bits=1)
+/// seg = ds.create_index_uncommitted(
+///     column="vector",
+///     index_type="IVF_RQ",
+///     num_partitions=256,
+///     ivf_centroids=centroids,
+///     rq_rotation=rot,
+///     fragment_ids=my_fragments,
+/// )
+/// ```
+#[pyfunction]
+#[pyo3(signature = (dimension, num_bits=1, rotation_type="fast", dtype="float32"))]
+pub fn build_rq_rotation(
+    dimension: usize,
+    num_bits: u8,
+    rotation_type: &str,
+    dtype: &str,
+) -> PyResult<String> {
+    use arrow::datatypes::{Float16Type, Float32Type, Float64Type};
+    use lance_index::vector::bq::builder::RabitQuantizer;
+    use lance_index::vector::bq::RQRotationType;
+    use lance_index::vector::quantizer::Quantization;
+
+    if !dimension.is_multiple_of(u8::BITS as usize) {
+        return Err(PyValueError::new_err(
+            "dimension must be divisible by 8 for IVF_RQ",
+        ));
+    }
+    let rotation = match rotation_type.to_lowercase().as_str() {
+        "fast" => RQRotationType::Fast,
+        "matrix" => {
+            return Err(PyValueError::new_err(
+                "matrix rotation cannot be serialized to JSON for distributed builds; \
+                 use rotation_type='fast'",
+            ));
+        }
+        other => {
+            return Err(PyValueError::new_err(format!(
+                "unknown rotation_type: {other}; expected 'fast'"
+            )));
+        }
+    };
+    let dim = dimension as i32;
+    let quantizer = match dtype.to_lowercase().as_str() {
+        "float16" => RabitQuantizer::new_with_rotation::<Float16Type>(num_bits, dim, rotation),
+        "float32" => RabitQuantizer::new_with_rotation::<Float32Type>(num_bits, dim, rotation),
+        "float64" => RabitQuantizer::new_with_rotation::<Float64Type>(num_bits, dim, rotation),
+        other => {
+            return Err(PyValueError::new_err(format!("unsupported dtype: {other}")));
+        }
+    };
+    serde_json::to_string(&quantizer.metadata(None))
+        .map_err(|e| PyValueError::new_err(format!("failed to serialize RQ rotation: {e}")))
+}
+
 #[allow(clippy::too_many_arguments)]
 async fn do_transform_vectors(
     dataset: &Dataset,
@@ -752,6 +825,7 @@ pub fn register_indices(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let indices = PyModule::new(py, "indices")?;
     indices.add_wrapped(wrap_pyfunction!(train_ivf_model))?;
     indices.add_wrapped(wrap_pyfunction!(train_pq_model))?;
+    indices.add_wrapped(wrap_pyfunction!(build_rq_rotation))?;
     indices.add_wrapped(wrap_pyfunction!(transform_vectors))?;
     indices.add_wrapped(wrap_pyfunction!(shuffle_transformed_vectors))?;
     indices.add_wrapped(wrap_pyfunction!(load_shuffled_vectors))?;

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

@@ -14,6 +14,7 @@ use lance_core::{Error, Result};
 use num_traits::Float;
 use serde::{Deserialize, Serialize};
 
+use crate::vector::bq::storage::RabitQuantizationMetadata;
 use crate::vector::quantizer::QuantizerBuildParams;
 
 pub mod builder;
@@ -100,24 +101,32 @@ impl FromStr for RQRotationType {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug)]
 pub struct RQBuildParams {
     pub num_bits: u8,
     pub rotation_type: RQRotationType,
+    /// Optional pre-built rotation to reuse instead of generating a fresh random one.
+    ///
+    /// Distributed `IVF_RQ` builds mint one rotation and broadcast it so every segment
+    /// rotates vectors identically. This is transient build-time state and is never
+    /// persisted to the `RabitQuantization` params proto.
+    pub rotation: Option<RabitQuantizationMetadata>,
 }
 
 impl RQBuildParams {
     pub fn new(num_bits: u8) -> Self {
         Self {
             num_bits,
             rotation_type: RQRotationType::default(),
+            rotation: None,
         }
     }
 
     pub fn with_rotation_type(num_bits: u8, rotation_type: RQRotationType) -> Self {
         Self {
             num_bits,
             rotation_type,
+            rotation: None,
         }
     }
 }
@@ -146,6 +155,7 @@ impl Default for RQBuildParams {
         Self {
             num_bits: 1,
             rotation_type: RQRotationType::default(),
+            rotation: None,
         }
     }
 }