From da6eb28dd795dc4f8dfffff703acefe792194eec Mon Sep 17 00:00:00 2001 From: Arnaud Gourlay Date: Fri, 10 Jul 2026 19:59:21 +0200 Subject: [PATCH] Add coverage for TurboQuant datatype in model tester (#9449) * Add coverage for TurboQuant datatype in model tester Co-authored-by: Cursor * Address review feedback: tolerance, doc comment, exhaustive match - Tighten dense_matches Turbo4 tolerance to 16 ulps relative and drop the absolute floor, so near-zero sign flips and small systematic quantization drift fail instead of passing - Fix ALL_CANDIDATES doc comment to match INITIAL_ACTIVE (six names start active, "c" and "u" via CreateVectorName) - Make model_vector match exhaustive so new VectorKind variants force a compile error - Use explicit DistanceType::from(distance) in turbo_storage_roundtrip Co-Authored-By: Claude Fable 5 --------- Co-authored-by: Cursor Co-authored-by: Claude Fable 5 --- .../src/model_testing/apply/reads.rs | 11 ++- .../src/model_testing/apply/writes.rs | 8 +- lib/collection/src/model_testing/fixture.rs | 10 +- lib/collection/src/model_testing/mod.rs | 15 ++- .../src/model_testing/op/generators.rs | 4 +- lib/collection/src/model_testing/op/mod.rs | 99 ++++++++++++++++++- lib/collection/src/model_testing/verify.rs | 28 +++++- .../src/operations/vector_params_builder.rs | 7 +- lib/segment/src/vector_storage/turbo/mod.rs | 25 +++++ 9 files changed, 189 insertions(+), 18 deletions(-) diff --git a/lib/collection/src/model_testing/apply/reads.rs b/lib/collection/src/model_testing/apply/reads.rs index 4033244c6c..9907cf6786 100644 --- a/lib/collection/src/model_testing/apply/reads.rs +++ b/lib/collection/src/model_testing/apply/reads.rs @@ -17,8 +17,8 @@ use shard::query::{FusionInternal, ScoringQuery, ShardPrefetch, ShardQueryReques use shard::scroll::ScrollRequestInternal; use super::super::op::{ - FusionKind, NamedVectors, Prefetch, ScrollFilter, canonical_sparse, has_num, - match_has_id_filter, match_has_vector_filter, match_num_filter, match_tag_filter, + FusionKind, NamedVectors, Prefetch, ScrollFilter, canonical_sparse, dense_diff, dense_matches, + has_num, match_has_id_filter, match_has_vector_filter, match_num_filter, match_tag_filter, match_url_prefix_filter, num_matches, optional_read_filter, passes_read_filters, tag_matches, url_prefix_matches, }; @@ -69,7 +69,12 @@ fn assert_named_vectors_match( }); match (ret_vec, exp) { (VectorInternal::Dense(a), VectorValue::Dense(b)) => { - assert_eq!(a, b, "{ctx}: dense vector `{name}` mismatch for id {id:?}"); + assert!( + dense_matches(name, a, b), + "{ctx}: dense vector `{name}` value divergence for id {id:?}: \ + engine {a:?}, model {b:?}; {}", + dense_diff(a, b), + ); } (VectorInternal::Sparse(a), VectorValue::Sparse(b)) => { assert_eq!( diff --git a/lib/collection/src/model_testing/apply/writes.rs b/lib/collection/src/model_testing/apply/writes.rs index d730d7dcf8..8ee7e9a761 100644 --- a/lib/collection/src/model_testing/apply/writes.rs +++ b/lib/collection/src/model_testing/apply/writes.rs @@ -6,7 +6,9 @@ use segment::json_path::JsonPath; use segment::types::{Payload, PayloadFieldSchema, PointIdType, VectorNameBuf}; use super::super::Model; -use super::super::op::{NamedVectors, match_num_filter, model_entry_from, num_matches}; +use super::super::op::{ + NamedVectors, match_num_filter, model_entry_from, model_vector, num_matches, +}; use super::{apply_update, to_named_persisted}; use crate::collection::Collection; use crate::operations::CollectionUpdateOperations; @@ -350,7 +352,9 @@ pub(super) async fn apply_update_vectors( let passes = condition_num.is_none_or(|n| num_matches(&entry.payload, n)); if passes { for (name, value) in partial { - entry.vectors.insert(name.clone(), value.clone()); + entry + .vectors + .insert(name.clone(), model_vector(name, value)); } } } diff --git a/lib/collection/src/model_testing/fixture.rs b/lib/collection/src/model_testing/fixture.rs index eccb7cd06a..b82402b009 100644 --- a/lib/collection/src/model_testing/fixture.rs +++ b/lib/collection/src/model_testing/fixture.rs @@ -20,7 +20,7 @@ use crate::collection::{Collection, RequestShardTransfer}; use crate::config::{CollectionConfigInternal, CollectionParams, WalConfig}; use crate::operations::config_diff::HnswConfigDiff; use crate::operations::shared_storage_config::SharedStorageConfig; -use crate::operations::types::{SparseVectorParams, VectorsConfig}; +use crate::operations::types::{Datatype, SparseVectorParams, VectorsConfig}; use crate::operations::vector_params_builder::VectorParamsBuilder; use crate::optimizers_builder::OptimizersConfig; use crate::shards::channel_service::ChannelService; @@ -110,6 +110,14 @@ pub(super) async fn fixture( params.multivector_config = Some(MultiVectorConfig::default()); dense_vectors.insert(name.to_string(), params); } + VectorKind::DenseTurbo(dim) => { + let mut builder = + VectorParamsBuilder::new(dim, Distance::Dot).with_datatype(Datatype::Turbo4); + if on_disk { + builder = builder.with_on_disk(true); + } + dense_vectors.insert(name.to_string(), builder.build()); + } } } diff --git a/lib/collection/src/model_testing/mod.rs b/lib/collection/src/model_testing/mod.rs index 9ebb4368cd..c63fe02469 100644 --- a/lib/collection/src/model_testing/mod.rs +++ b/lib/collection/src/model_testing/mod.rs @@ -25,9 +25,9 @@ use crate::shards::shard::PeerId; const PEER_ID: PeerId = 1; const COLLECTION_NAME: &str = "test"; -/// Static metadata for every vector name the test might ever activate. Four names start -/// active in the fixture; the remaining two are exclusively reachable through -/// `Op::CreateVectorName`. +/// Static metadata for every vector name the test might ever activate. Six names start +/// active in the fixture ("a", "b", "i", "s", "m", "q", see `INITIAL_ACTIVE`); "c" and +/// "u" are reachable through `Op::CreateVectorName`. pub(super) const ALL_CANDIDATES: &[VectorCandidate] = &[ VectorCandidate { name: "a", @@ -60,10 +60,14 @@ pub(super) const ALL_CANDIDATES: &[VectorCandidate] = &[ name: "m", kind: VectorKind::MultiDense(4), }, + VectorCandidate { + name: "q", + kind: VectorKind::DenseTurbo(8), + }, ]; /// Names present in the collection schema at fixture time. -pub(super) const INITIAL_ACTIVE: &[&str] = &["a", "b", "i", "s", "m"]; +pub(super) const INITIAL_ACTIVE: &[&str] = &["a", "b", "i", "s", "m", "q"]; /// Dense vector name configured with HNSW `inline_storage` + scalar quantization in the fixture. pub(super) const INLINE_STORAGE_VECTOR: &str = "i"; @@ -80,6 +84,9 @@ pub(super) enum VectorKind { /// ColBERT-style multi-vector: each point stores a matrix of `dim`-wide rows. Scoring /// uses MaxSim across query rows × stored rows. MultiDense(u64), + /// Dense vector stored with the `Turbo4` (TurboQuant 4-bit) storage datatype, the + /// primary quantized storage, applied to appendable segments too. + DenseTurbo(u64), } pub(super) fn kind_of(name: &str) -> VectorKind { diff --git a/lib/collection/src/model_testing/op/generators.rs b/lib/collection/src/model_testing/op/generators.rs index 5be0339019..c8dd148f19 100644 --- a/lib/collection/src/model_testing/op/generators.rs +++ b/lib/collection/src/model_testing/op/generators.rs @@ -267,7 +267,9 @@ pub(super) fn random_partial_named_vectors( /// Build a random vector matching the kind metadata associated with `name`. fn random_vector_for_name(rng: &mut impl Rng, name: &str) -> VectorValue { match kind_of(name) { - VectorKind::Dense(dim) => VectorValue::Dense(random_dense_vec(rng, dim)), + VectorKind::Dense(dim) | VectorKind::DenseTurbo(dim) => { + VectorValue::Dense(random_dense_vec(rng, dim)) + } VectorKind::Sparse => VectorValue::Sparse(random_sparse_vector(rng)), VectorKind::MultiDense(dim) => VectorValue::MultiDense(random_multi_dense(rng, dim)), } diff --git a/lib/collection/src/model_testing/op/mod.rs b/lib/collection/src/model_testing/op/mod.rs index 3dc8ca6b7f..1562112e06 100644 --- a/lib/collection/src/model_testing/op/mod.rs +++ b/lib/collection/src/model_testing/op/mod.rs @@ -25,11 +25,12 @@ use segment::json_path::JsonPath; use segment::types::{ Condition, Distance, FieldCondition, Filter, HasIdCondition, HasVectorCondition, Match, MultiVectorConfig, Payload, PayloadFieldSchema, PayloadSchemaParams, PayloadSchemaType, - PointIdType, VectorNameBuf, WithPayloadInterface, WithVector, + PointIdType, VectorNameBuf, VectorStorageDatatype, WithPayloadInterface, WithVector, }; +use segment::vector_storage::turbo::turbo_storage_roundtrip; use sparse::common::sparse_vector::SparseVector; -use super::{ALL_CANDIDATES, Model, ModelEntry, VectorKind, VectorValue}; +use super::{ALL_CANDIDATES, Model, ModelEntry, VectorKind, VectorValue, kind_of}; use crate::operations::point_ops::UpdateMode; /// Operations driven against both the live `Collection` and the model. @@ -611,6 +612,12 @@ impl Op { multivector_config: Some(MultiVectorConfig::default()), datatype: None, }), + VectorKind::DenseTurbo(dim) => VectorNameConfig::dense(DenseVectorConfig { + size: dim as usize, + distance: Distance::Dot, + multivector_config: None, + datatype: Some(VectorStorageDatatype::Turbo4), + }), }; Op::CreateVectorName { name: pick.name.to_string(), @@ -861,11 +868,97 @@ pub(super) fn has_num(payload: &Payload) -> bool { /// Build a new `ModelEntry` from a fresh upsert. pub(super) fn model_entry_from(vecs: &NamedVectors, payload: &Payload) -> ModelEntry { ModelEntry { - vectors: vecs.clone(), + vectors: vecs + .iter() + .map(|(name, value)| (name.clone(), model_vector(name, value))) + .collect(), payload: payload.clone(), } } +/// Predicted engine read-back for `value` stored under `name`. Turbo4-backed dense +/// vectors are lossy: the engine stores 4-bit quantized codes and returns the +/// dequantized vector, so the model must record that round-trip instead of the inserted +/// value. The round-trip is deterministic (fixed rotation seeds), shared across +/// segments and reloads. The engine still receives the original vector: the round-trip +/// is not idempotent (re-quantizing a read-back shifts the stored norm), so +/// canonicalizing at generation time would not converge. +pub(super) fn model_vector(name: &str, value: &VectorValue) -> VectorValue { + match (kind_of(name), value) { + (VectorKind::DenseTurbo(_), VectorValue::Dense(v)) => { + // Every fixture vector uses Dot (see `fixture::fixture` and the + // CreateVectorName generator arm above). + VectorValue::Dense(turbo_storage_roundtrip(v, Distance::Dot)) + } + ( + VectorKind::Dense(_) | VectorKind::Sparse | VectorKind::MultiDense(_), + VectorValue::Dense(_) | VectorValue::Sparse(_) | VectorValue::MultiDense(_), + ) => value.clone(), + (VectorKind::DenseTurbo(_), VectorValue::Sparse(_) | VectorValue::MultiDense(_)) => { + panic!("model_vector: non-dense value for Turbo4 name `{name}`: {value:?}") + } + } +} + +/// Compare a returned dense vector against the model's prediction for `name`. +/// Exact for full-precision names. Turbo4 read-backs are compared with a tiny +/// relative tolerance: a copy-on-write point move re-quantizes the dequantized +/// read-back, and although the codes and the centroid norm are reproduced, the +/// re-measured stored norm passes through two f64 rotation round-trips and can +/// land a few ulps off, uniformly rescaling the read-back at ulp scale. The +/// budget of 16 ulps (relative) absorbs that wobble even accumulated across +/// repeated moves, while real divergences (wrong codes, stale vector, a bad +/// norm divisor) are orders of magnitude larger. Purely relative on purpose: +/// an absolute floor would accept sign flips of near-zero components. +pub(super) fn dense_matches(name: &str, actual: &[f32], expected: &[f32]) -> bool { + if actual.len() != expected.len() { + return false; + } + match kind_of(name) { + VectorKind::DenseTurbo(_) => actual.iter().zip(expected).all(|(&a, &e)| { + let tol = 16.0 * f32::EPSILON * f32::max(a.abs(), e.abs()); + (a - e).abs() <= tol + }), + VectorKind::Dense(_) | VectorKind::Sparse | VectorKind::MultiDense(_) => actual == expected, + } +} + +/// Human-readable breakdown of a dense mismatch for panic messages: per-component +/// deltas plus a uniform-scale probe. A uniform engine/model ratio across all +/// components is the signature of a re-quantization rescale (e.g. the Turbo4 +/// copy-on-write degradation, where read-backs come back scaled by `cn/sqrt(d)`), +/// as opposed to per-component noise or a stale/wrong vector. +pub(super) fn dense_diff(actual: &[f32], expected: &[f32]) -> String { + if actual.len() != expected.len() { + return format!( + "length mismatch: engine {} vs model {}", + actual.len(), + expected.len(), + ); + } + let diffs: Vec = actual.iter().zip(expected).map(|(&a, &e)| a - e).collect(); + let max_abs_diff = diffs.iter().fold(0.0f32, |m, d| m.max(d.abs())); + let ratios: Vec = actual + .iter() + .zip(expected) + .map(|(&a, &e)| if e.abs() > 1e-12 { a / e } else { f32::NAN }) + .collect(); + // Judge uniformity on the finite ratios only: a near-zero expected component + // yields a NaN ratio, and a genuinely uniform rescale should still be labeled + // as such when one component sits at zero. + let finite: Vec = ratios.iter().copied().filter(|r| r.is_finite()).collect(); + let uniform = !finite.is_empty() && finite.iter().all(|r| (r - finite[0]).abs() < 1e-5); + let scale_note = if uniform { + format!( + "UNIFORM engine/model scale {:.6} (single rescale)", + finite[0] + ) + } else { + "non-uniform ratios (per-component divergence)".to_string() + }; + format!("max_abs_diff={max_abs_diff:e}; diffs={diffs:?}; ratios={ratios:?}; {scale_note}") +} + /// Sort sparse indices ascending and drop entries with zero value (mirrors the engine's /// canonicalization on read — see `lib/sparse/src/common/sparse_vector.rs`). pub(super) fn canonical_sparse(sv: &SparseVector) -> SparseVector { diff --git a/lib/collection/src/model_testing/verify.rs b/lib/collection/src/model_testing/verify.rs index f8db31d0e7..87f13434a8 100644 --- a/lib/collection/src/model_testing/verify.rs +++ b/lib/collection/src/model_testing/verify.rs @@ -7,7 +7,7 @@ use common::types::{DetailsLevel, TelemetryDetail}; use segment::types::{PointIdType, VectorNameBuf, WithPayloadInterface, WithVector}; use shard::scroll::ScrollRequestInternal; -use super::op::canonical_sparse; +use super::op::{canonical_sparse, dense_diff, dense_matches}; use super::{Model, ModelEntry, VectorValue}; use crate::collection::Collection; use crate::operations::shard_selector_internal::ShardSelectorInternal; @@ -91,9 +91,31 @@ pub(super) fn assert_matches_model(actual: &Model, expected: &Model, ctx: &str) .get(id) .unwrap_or_else(|| panic!("{ctx}: missing id {id:?}")); assert_eq!( - actual_entry.vectors, expected_entry.vectors, - "{ctx}: vectors mismatch for id {id:?}", + actual_entry.vectors.keys().collect::>(), + expected_entry.vectors.keys().collect::>(), + "{ctx}: vector names mismatch for id {id:?}", ); + for (name, expected_value) in &expected_entry.vectors { + let actual_value = &actual_entry.vectors[name]; + // Turbo4 dense values get a few-ulp tolerance (see `dense_matches`); + // everything else stays exact. + let matches = match (actual_value, expected_value) { + (VectorValue::Dense(a), VectorValue::Dense(e)) => dense_matches(name, a, e), + _ => actual_value == expected_value, + }; + if !matches { + // Dense mismatches get a per-component diff so a uniform rescale + // is distinguishable from noise at a glance. + let detail = match (actual_value, expected_value) { + (VectorValue::Dense(a), VectorValue::Dense(e)) => dense_diff(a, e), + _ => String::new(), + }; + panic!( + "{ctx}: vector `{name}` value divergence for id {id:?}: \ + engine {actual_value:?}, model {expected_value:?}; {detail}", + ); + } + } assert_eq!( actual_entry.payload, expected_entry.payload, "{ctx}: payload mismatch for id {id:?}", diff --git a/lib/collection/src/operations/vector_params_builder.rs b/lib/collection/src/operations/vector_params_builder.rs index c9e5aa30c4..356eafcc4a 100644 --- a/lib/collection/src/operations/vector_params_builder.rs +++ b/lib/collection/src/operations/vector_params_builder.rs @@ -7,7 +7,7 @@ use std::num::NonZeroU64; use segment::types::{Distance, QuantizationConfig}; use crate::operations::config_diff::HnswConfigDiff; -use crate::operations::types::VectorParams; +use crate::operations::types::{Datatype, VectorParams}; pub struct VectorParamsBuilder { vector_params: VectorParams, @@ -49,6 +49,11 @@ impl VectorParamsBuilder { self } + pub fn with_datatype(mut self, datatype: Datatype) -> Self { + self.vector_params.datatype = Some(datatype); + self + } + pub fn build(self) -> VectorParams { self.vector_params } diff --git a/lib/segment/src/vector_storage/turbo/mod.rs b/lib/segment/src/vector_storage/turbo/mod.rs index b518c2b08d..fa11a6b1be 100644 --- a/lib/segment/src/vector_storage/turbo/mod.rs +++ b/lib/segment/src/vector_storage/turbo/mod.rs @@ -315,6 +315,31 @@ fn open_turbo_vector_storage_impl( }) } +/// Quantize then dequantize `vector` exactly as a [`TurboVectorStorage`] with this +/// `distance` does across `insert_vector` + `get_vector`. Pure function of its inputs: +/// the quantizer is fully determined by `(dim, distance)` (the rotation derives from +/// fixed seeds), so the result is identical across storage instances, segment rebuilds, +/// and reloads. Lets model-based tests predict the read-back value of a Turbo4-backed +/// vector without opening a storage. +pub fn turbo_storage_roundtrip(vector: &[f32], distance: Distance) -> Vec { + let dim = vector.len(); + let quantizer = TurboQuantizer::new( + dim, + TQDT_BITS, + TQDT_MODE, + quantization::DistanceType::from(distance), + TQDT_ROTATION, + None, + ); + let mut buf = vec![0.0; quantizer.get_padded_dim()]; + let encoded = quantizer.quantize(vector, &mut buf); + // Mirror of `TurboVectorStorage::dequantize_vector`: dequantize, rotate back, drop + // the padding tail, cast to f32. + let mut dequantized = quantizer.dequantize::(&encoded); + quantizer.apply_inverse_rotation(&mut dequantized); + dequantized[..dim].iter().map(|&x| x as f32).collect() +} + impl VectorStorageRead for TurboVectorStorage { fn size_of_available_vectors_in_bytes(&self) -> usize { self.available_vector_count() * self.quantized_vector_size()