diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 94372ef27cc..5f0818f6d46 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -465,6 +465,7 @@ pub fn get_vector_index_params( Ok(SQBuildParams { num_bits, sample_rate, + bounds: None, }) }, )?; diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 19aa3fa177d..3b155cb3772 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/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> { 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, diff --git a/rust/lance-index/src/vector/sq.rs b/rust/lance-index/src/vector/sq.rs index bdc3204021e..b45d4d1228f 100644 --- a/rust/lance-index/src/vector/sq.rs +++ b/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` built from per-element values, for the SQ + /// build tests. + fn f32_fsl(values: impl IntoIterator, 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, + ¶ms, + ) + .unwrap(); + let sq_b = ScalarQuantizer::build( + &f32_fsl((0..dim).map(|v| (v as f32) * 100.0 - 500.0), dim), + DistanceType::L2, + ¶ms, + ) + .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, ¶ms).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, ¶ms).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, ¶ms).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, + ¶ms, + ) + .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]; diff --git a/rust/lance-index/src/vector/sq/builder.rs b/rust/lance-index/src/vector/sq/builder.rs index 359765040dd..9e8b3f2eb10 100644 --- a/rust/lance-index/src/vector/sq/builder.rs +++ b/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>, } 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) -> Self { + Self { + num_bits, + bounds: Some(bounds), + ..Default::default() } } } diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index bfc4a4f3474..15223660059 100644 --- a/rust/lance/src/index/vector.rs +++ b/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, } } diff --git a/rust/lance/src/index/vector/details.rs b/rust/lance/src/index/vector/details.rs index 63f9375792e..02daccf01e7 100644 --- a/rust/lance/src/index/vector/details.rs +++ b/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 { diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 732e6befd90..61c97e7e775 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -3545,7 +3545,9 @@ mod tests { VectorIndexParams::with_ivf_sq_params( DistanceType::L2, ivf_params, - SQBuildParams::default(), + // Shards must share globally-trained bounds to be + // mergeable; the data is generated in [0, 1). + SQBuildParams::with_bounds(8, 0.0..1.0), ) } other => panic!("unsupported test index type: {}", other), @@ -3807,7 +3809,9 @@ mod tests { DistanceType::L2, prepare_global_ivf(&dataset, "vector").await, HnswBuildParams::default(), - SQBuildParams::default(), + // Shards must share globally-trained bounds to be mergeable; + // the data is generated in [0, 1). + SQBuildParams::with_bounds(8, 0.0..1.0), ), other => panic!("unexpected HNSW index type {other}"), }; @@ -3915,6 +3919,224 @@ mod tests { ); } + fn make_sq_index_params( + index_kind: &str, + ivf_params: IvfBuildParams, + sq_params: SQBuildParams, + ) -> VectorIndexParams { + match index_kind { + "IVF_SQ" => { + VectorIndexParams::with_ivf_sq_params(DistanceType::L2, ivf_params, sq_params) + } + "IVF_HNSW_SQ" => VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + ivf_params, + HnswBuildParams::default(), + sq_params, + ), + other => panic!("unexpected SQ index kind {other}"), + } + } + + #[rstest] + #[case::ivf_sq("IVF_SQ")] + #[case::ivf_hnsw_sq("IVF_HNSW_SQ")] + #[tokio::test] + async fn test_merge_sq_segments_with_injected_bounds(#[case] index_kind: &str) { + const INDEX_NAME: &str = "vector_idx"; + const K: usize = 10; + const NUM_QUERIES: usize = 10; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_two_fragment_batches(); + + let ds_single_uri = format!("{}/single", base_uri); + let ds_split_uri = format!("{}/split", base_uri); + let mut ds_single = + write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; + let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; + + // Globally-trained bounds injected into every build; the data is + // generated in [0, 1) so these bounds cover it entirely. + let injected_bounds = 0.0f64..1.0f64; + let ivf_params = prepare_global_ivf(&ds_single, "vector").await; + let params = make_sq_index_params( + index_kind, + ivf_params, + SQBuildParams::with_bounds(8, injected_bounds.clone()), + ); + + // Baseline: single non-sharded index built with the same injected bounds. + ds_single + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Sharded: one uncommitted segment per fragment, then merge + commit. + let fragment_groups: Vec> = ds_split + .get_fragments() + .iter() + .map(|f| vec![f.id() as u32]) + .collect(); + assert!(fragment_groups.len() >= 2); + let segments = + build_segments_for_fragment_groups(&mut ds_split, fragment_groups, ¶ms, INDEX_NAME) + .await; + let merged = ds_split + .merge_existing_index_segments(segments) + .await + .unwrap(); + ds_split + .commit_existing_index_segments(INDEX_NAME, "vector", vec![merged]) + .await + .unwrap(); + + // Both the baseline and the merged index must carry exactly the + // injected bounds in their SQ storage metadata. + let obj_store = Arc::new(ObjectStore::local()); + let scheduler = ScanScheduler::new(obj_store, SchedulerConfig::default_for_testing()); + for ds in [&ds_single, &ds_split] { + let indices = ds.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(indices.len(), 1); + let sq_meta = + get_sq_metadata(ds, scheduler.clone(), &indices[0].uuid.to_string()).await; + assert_eq!(sq_meta.bounds, injected_bounds); + assert_eq!(sq_meta.num_bits, 8); + } + + // Query both datasets with the same query set. + let query_batch = ds_single + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(NUM_QUERIES as i64), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = query_batch["vector"].as_fixed_size_list(); + + // For HNSW (approximate, non-deterministic build) accumulate overlap + // across all queries and check the aggregate recall once -- a single + // query's overlap is too noisy on a small fixture to threshold. + let mut hnsw_overlap = 0usize; + let mut hnsw_total = 0usize; + for i in 0..vectors.len() { + let q = vectors.value(i); + let collect = |ds: &Dataset| { + let q = q.clone(); + let ds = ds.clone(); + async move { + let result = ds + .scan() + .with_row_id() + .project(&["_rowid"] as &[&str]) + .unwrap() + .nearest("vector", q.as_ref(), K) + .unwrap() + .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) + .try_into_batch() + .await + .unwrap(); + result[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>() + } + }; + let ids_single = collect(&ds_single).await; + let ids_split = collect(&ds_split).await; + if index_kind == "IVF_SQ" { + // With identical bounds the codes are identical, so the merged + // index must return exactly the same Top-K as the baseline. + assert_eq!( + ids_single, ids_split, + "single vs merged IVF_SQ returned different Top-K row ids", + ); + } else { + assert_eq!(ids_split.len(), K); + let baseline: std::collections::HashSet = ids_single.iter().copied().collect(); + hnsw_overlap += ids_split.iter().filter(|id| baseline.contains(id)).count(); + hnsw_total += K; + } + } + if index_kind != "IVF_SQ" { + // Aggregate recall >= 0.5 vs the baseline (fixture measures ~0.75 + // with tight variance, so this is a stable floor); a mismatched + // scale would drop overlap toward random. + assert!( + hnsw_overlap * 2 >= hnsw_total, + "merged IVF_HNSW_SQ aggregate overlap too low: {hnsw_overlap}/{hnsw_total}", + ); + } + } + + #[rstest] + #[case::ivf_sq("IVF_SQ")] + #[case::ivf_hnsw_sq("IVF_HNSW_SQ")] + #[tokio::test] + async fn test_merge_sq_segments_rejects_mismatched_bounds(#[case] index_kind: &str) { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_two_fragment_batches(); + let dataset_uri = format!("{}/merge_sq_rejects_mismatched_bounds", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + + let ivf_params = prepare_global_ivf(&dataset, "vector").await; + let first_params = make_sq_index_params( + index_kind, + ivf_params.clone(), + SQBuildParams::with_bounds(8, 0.0..1.0), + ); + let second_params = make_sq_index_params( + index_kind, + ivf_params, + SQBuildParams::with_bounds(8, 0.0..2.0), + ); + + let first_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, &first_params) + .name("vector_idx".to_string()) + .fragments(vec![fragments[0].id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + let second_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, &second_params) + .name("vector_idx".to_string()) + .fragments(vec![fragments[1].id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + + let error = dataset + .merge_existing_index_segments(vec![first_segment, second_segment]) + .await + .unwrap_err(); + assert!( + matches!(error, lance_core::Error::Index { .. }), + "expected Error::Index, got {error:?}" + ); + assert!( + error + .to_string() + .contains("SQ metadata mismatch across shards"), + "{error}" + ); + } + #[tokio::test] async fn test_merge_index_metadata_reports_progress() { const INDEX_NAME: &str = "vector_idx";