feat(index): inject globally-trained SQ bounds through the build params

java/lance-jni/src/utils.rs

@@ -465,6 +465,7 @@ pub fn get_vector_index_params(
                     Ok(SQBuildParams {
                         num_bits,
                         sample_rate,
+                        bounds: None,
                     })
                 },
             )?;

rust/lance-index/src/vector/distributed/index_merger.rs

@@ -178,6 +178,26 @@ fn ensure_fixed_size_list_compatible(
     Ok(())
 }
 
+/// SQ shards must share every metadata field that governs how a code is
+/// interpreted: `num_bits` (the scale denominator) and `bounds` (the range).
+/// Unlike PQ codebooks / IVF centroids (compared with a tolerance above, since
+/// training can drift), these are injected verbatim into every shard's build,
+/// so any difference is not benign drift -- it means a shard used a different
+/// scale and its codes are incomparable. Exact inequality is the check
+/// (`num_bits` matters especially once SQ4 lands).
+fn ensure_sq_metadata_compatible(
+    reference: &ScalarQuantizationMetadata,
+    candidate: &ScalarQuantizationMetadata,
+) -> Result<()> {
+    if reference.num_bits != candidate.num_bits || reference.bounds != candidate.bounds {
+        return Err(Error::index(format!(
+            "SQ metadata mismatch across shards: num_bits {} vs {}, bounds {:?} vs {:?}",
+            reference.num_bits, candidate.num_bits, reference.bounds, candidate.bounds
+        )));
+    }
+    Ok(())
+}
+
 async fn try_read_ivf_proto(reader: &V2Reader) -> Result<Option<pb::Ivf>> {
     let Some(ivf_idx) = reader.metadata().file_schema.metadata.get(IVF_METADATA_KEY) else {
         return Ok(None);
@@ -1007,7 +1027,9 @@ pub async fn merge_partial_vector_auxiliary_files(
                     return Err(Error::index("Dimension mismatch across shards".to_string()));
                 }
 
-                if sq_meta.is_none() {
+                if let Some(prev) = &sq_meta {
+                    ensure_sq_metadata_compatible(prev, &sq_meta_parsed)?;
+                } else {
                     sq_meta = Some(sq_meta_parsed.clone());
                 }
                 if v2w_opt.is_none() {
@@ -1425,7 +1447,9 @@ pub async fn merge_partial_vector_auxiliary_files(
                 {
                     return Err(Error::index("Dimension mismatch across shards".to_string()));
                 }
-                if sq_meta.is_none() {
+                if let Some(prev) = &sq_meta {
+                    ensure_sq_metadata_compatible(prev, &sq_meta_parsed)?;
+                } else {
                     sq_meta = Some(sq_meta_parsed.clone());
                 }
                 if v2w_opt.is_none() {
@@ -1652,6 +1676,20 @@ mod tests {
         lance_core::Result<()>
     );
 
+    #[test]
+    fn ensure_sq_metadata_compatible_checks_bounds_and_num_bits() {
+        let meta = |num_bits, bounds| ScalarQuantizationMetadata {
+            dim: 8,
+            num_bits,
+            bounds,
+        };
+        // Identical metadata is compatible.
+        assert!(ensure_sq_metadata_compatible(&meta(8, 0.0..1.0), &meta(8, 0.0..1.0)).is_ok());
+        // Different bounds OR different num_bits (same bounds) are both rejected.
+        assert!(ensure_sq_metadata_compatible(&meta(8, 0.0..1.0), &meta(8, 0.0..2.0)).is_err());
+        assert!(ensure_sq_metadata_compatible(&meta(8, 0.0..1.0), &meta(4, 0.0..1.0)).is_err());
+    }
+
     async fn write_flat_partial_aux(
         store: &ObjectStore,
         aux_path: &Path,

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

@@ -142,6 +142,34 @@ impl Quantization for ScalarQuantizer {
             data.data_type()
         )))?;
 
+        // Injected bounds take precedence over local training, so codes stay
+        // comparable across independently built shards.
+        if let Some(bounds) = params.bounds.clone() {
+            // Validate what the skipped training scan validates implicitly:
+            // supported float type, and finite non-decreasing bounds (NaN/inf
+            // collapse codes; start == end is valid for a constant column).
+            if !matches!(
+                fsl.value_type(),
+                DataType::Float16 | DataType::Float32 | DataType::Float64
+            ) {
+                return Err(Error::invalid_input(format!(
+                    "SQ builder: unsupported data type: {}",
+                    fsl.value_type()
+                )));
+            }
+            if !bounds.start.is_finite() || !bounds.end.is_finite() || bounds.start > bounds.end {
+                return Err(Error::invalid_input(format!(
+                    "SQ builder: injected bounds must be finite and non-decreasing, got {}..{}",
+                    bounds.start, bounds.end
+                )));
+            }
+            return Ok(Self::with_bounds(
+                params.num_bits,
+                fsl.value_length() as usize,
+                bounds,
+            ));
+        }
+
         let mut quantizer = Self::new(params.num_bits, fsl.value_length() as usize);
 
         match fsl.value_type() {
@@ -349,6 +377,107 @@ mod tests {
         });
     }
 
+    /// A `FixedSizeList<Float32>` built from per-element values, for the SQ
+    /// build tests.
+    fn f32_fsl(values: impl IntoIterator<Item = f32>, dim: i32) -> FixedSizeListArray {
+        FixedSizeListArray::try_new_from_values(Float32Array::from_iter_values(values), dim)
+            .unwrap()
+    }
+
+    #[test]
+    fn test_build_with_injected_bounds() {
+        let dim = 16;
+        let injected = -3.0..42.0;
+        let params = SQBuildParams::with_bounds(8, injected.clone());
+
+        // Two samples with very different value ranges; both must adopt the
+        // injected bounds, not their own.
+        let sq_a = ScalarQuantizer::build(
+            &f32_fsl((0..dim).map(|v| v as f32), dim),
+            DistanceType::L2,
+            &params,
+        )
+        .unwrap();
+        let sq_b = ScalarQuantizer::build(
+            &f32_fsl((0..dim).map(|v| (v as f32) * 100.0 - 500.0), dim),
+            DistanceType::L2,
+            &params,
+        )
+        .unwrap();
+        assert_eq!(sq_a.bounds(), injected);
+        assert_eq!(sq_b.bounds(), injected);
+
+        // Identical bounds ⟹ identical codes for the same input.
+        let query = f32_fsl((0..dim).map(|v| v as f32 * 2.0), dim);
+        assert_eq!(
+            sq_a.quantize(&query).unwrap().as_ref(),
+            sq_b.quantize(&query).unwrap().as_ref(),
+        );
+    }
+
+    #[test]
+    fn test_build_rejects_invalid_injected_bounds() {
+        let sample = f32_fsl((0..8).map(|v| v as f32), 8);
+        for bad in [
+            f64::NAN..1.0,
+            0.0..f64::NAN,
+            f64::NEG_INFINITY..1.0,
+            0.0..f64::INFINITY,
+            2.0..1.0, // reversed
+        ] {
+            let params = SQBuildParams::with_bounds(8, bad.clone());
+            let err = ScalarQuantizer::build(&sample, DistanceType::L2, &params).unwrap_err();
+            assert!(
+                matches!(err, Error::InvalidInput { .. }),
+                "invalid injected bounds {bad:?} must be InvalidInput, got {err:?}",
+            );
+            assert!(err.to_string().contains("bounds must be finite"), "{err}");
+        }
+        // Equal endpoints are valid (a globally constant column has min == max).
+        let params = SQBuildParams::with_bounds(8, 3.0..3.0);
+        assert!(ScalarQuantizer::build(&sample, DistanceType::L2, &params).is_ok());
+    }
+
+    #[test]
+    fn test_build_with_injected_bounds_rejects_unsupported_type() {
+        // Int32 element type: the injected-bounds fast path must still reject
+        // it, not defer the failure to quantization time.
+        let fsl = FixedSizeListArray::try_new_from_values(
+            arrow_array::Int32Array::from_iter_values(0..4),
+            4,
+        )
+        .unwrap();
+        let params = SQBuildParams::with_bounds(8, 0.0..10.0);
+        let err = ScalarQuantizer::build(&fsl, DistanceType::L2, &params).unwrap_err();
+        assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}");
+        assert!(err.to_string().contains("unsupported data type"), "{err}");
+    }
+
+    #[test]
+    fn test_build_without_bounds_self_trains() {
+        let dim = 16;
+        let params = SQBuildParams::default();
+        assert!(params.bounds.is_none());
+
+        let sq = ScalarQuantizer::build(
+            &f32_fsl((0..dim).map(|v| v as f32), dim),
+            DistanceType::L2,
+            &params,
+        )
+        .unwrap();
+        assert_eq!(sq.bounds(), 0.0..(dim - 1) as f64);
+    }
+
+    #[test]
+    fn test_sq_build_params_with_bounds() {
+        // 8 bits: SQ currently only encodes u8 (SQ4 is a TODO), so don't
+        // present 4 as a supported width in the test.
+        let params = SQBuildParams::with_bounds(8, 1.0..2.0);
+        assert_eq!(params.num_bits, 8);
+        assert_eq!(params.bounds, Some(1.0..2.0));
+        assert_eq!(params.sample_rate, SQBuildParams::default().sample_rate);
+    }
+
     #[tokio::test]
     async fn test_scale_to_u8_with_nan() {
         let values = vec![0.0, 1.0, 2.0, 3.0, f64::NAN];

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

@@ -10,6 +10,18 @@ pub struct SQBuildParams {
 
     /// Sample rate for training.
     pub sample_rate: usize,
+
+    /// User-provided, globally-trained quantization bounds. When set, per-build
+    /// training is skipped and these exact bounds are used, so codes are
+    /// comparable across independently built shards (mirrors
+    /// `PQBuildParams::codebook`).
+    ///
+    /// Must be the global min/max in the quantizer's **post-transform** space:
+    /// cosine indexes L2-normalize first, so cosine bounds cover the normalized
+    /// components (`-1.0..=1.0`), not the raw values; L2/dot use the raw
+    /// min/max. Bounds that don't cover that space clip every code to 0 and
+    /// silently collapse the index.
+    pub bounds: Option<std::ops::Range<f64>>,
 }
 
 impl From<&SQBuildParams> for crate::pb::vector_index_details::ScalarQuantization {
@@ -25,6 +37,30 @@ impl Default for SQBuildParams {
         Self {
             num_bits: 8,
             sample_rate: 256,
+            bounds: None,
+        }
+    }
+}
+
+impl SQBuildParams {
+    /// Create build params carrying pre-trained bounds, skipping training.
+    ///
+    /// Use this so independently built shards of a distributed index quantize
+    /// on the same scale and their segments stay mergeable (see [`bounds`]).
+    ///
+    /// ```
+    /// use lance_index::vector::sq::builder::SQBuildParams;
+    /// // Every shard builds against the globally-trained min/max:
+    /// let params = SQBuildParams::with_bounds(8, 0.0..1.0);
+    /// assert_eq!(params.bounds, Some(0.0..1.0));
+    /// ```
+    ///
+    /// [`bounds`]: SQBuildParams::bounds
+    pub fn with_bounds(num_bits: u16, bounds: std::ops::Range<f64>) -> Self {
+        Self {
+            num_bits,
+            bounds: Some(bounds),
+            ..Default::default()
         }
     }
 }

rust/lance/src/index/vector.rs

@@ -2013,6 +2013,7 @@ fn derive_sq_params(sq_quantizer: &ScalarQuantizer) -> SQBuildParams {
     SQBuildParams {
         num_bits: sq_quantizer.num_bits(),
         sample_rate: 256, // Default
+        bounds: None,
     }
 }
 

rust/lance/src/index/vector/details.rs

@@ -1163,6 +1163,7 @@ mod tests {
             SQBuildParams {
                 num_bits: 8,
                 sample_rate: 128,
+                bounds: None,
             },
         );
         params.skip_transpose = true;
@@ -1313,6 +1314,7 @@ mod tests {
         let sq = SQBuildParams {
             num_bits: 8,
             sample_rate: 128,
+            bounds: None,
         };
 
         match combo {