From a95292f3aaa2e522ab650dfd9223bc7108da08bd Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Wed, 20 May 2026 14:12:06 +0200 Subject: [PATCH] are you happy fmt --- lib/segment/src/data_types/primitive.rs | 28 ++++-- lib/segment/src/data_types/turbo_quant.rs | 97 +++++++------------ .../dense/appendable_dense_vector_storage.rs | 27 ++++-- .../dense/dense_vector_storage.rs | 36 +++++-- .../dense/empty_dense_vector_storage.rs | 4 + .../dense/volatile_dense_vector_storage.rs | 21 ++-- .../quantized/quantized_vectors.rs | 22 ++++- .../src/vector_storage/vector_storage_base.rs | 16 ++- 8 files changed, 146 insertions(+), 105 deletions(-) diff --git a/lib/segment/src/data_types/primitive.rs b/lib/segment/src/data_types/primitive.rs index 97373a6478..454b9573ec 100644 --- a/lib/segment/src/data_types/primitive.rs +++ b/lib/segment/src/data_types/primitive.rs @@ -97,14 +97,28 @@ where fn storage_len_in_elements(api_dim: usize, distance: Distance) -> usize { Self::storage_layout(api_dim, distance).size() / size_of::() } +} - /// Recover the api-level dimension from a slot length in `Self`-elements. - /// - /// Inverse of `storage_len_in_elements`. Default is identity (flat - /// layout); metric-aware types subtract the payload elements. - fn api_dim_from_storage_len(storage_len: usize, distance: Distance) -> usize { - let _ = distance; - storage_len +/// Truncate a decoded api-level vector to the originally requested `api_dim`. +/// +/// Some `T` (TurboQuant) round the slot length up to a codebook-aligned +/// `padded_dim` and thus `slice_to_float_cow` produces a `padded_dim`-long +/// `Cow`. Storage / quantization layers know the canonical `api_dim` and use +/// this helper to drop the trailing padding before exposing the floats. +/// For flat `T` the input is already at `api_dim` and this is a no-op. +pub fn truncate_to_api_dim( + v: Cow<'_, [VectorElementType]>, + api_dim: usize, +) -> Cow<'_, [VectorElementType]> { + if v.len() <= api_dim { + return v; + } + match v { + Cow::Borrowed(s) => Cow::Borrowed(&s[..api_dim]), + Cow::Owned(mut v) => { + v.truncate(api_dim); + Cow::Owned(v) + } } } diff --git a/lib/segment/src/data_types/turbo_quant.rs b/lib/segment/src/data_types/turbo_quant.rs index 9546d832be..9f3adfc808 100644 --- a/lib/segment/src/data_types/turbo_quant.rs +++ b/lib/segment/src/data_types/turbo_quant.rs @@ -1,7 +1,14 @@ //! Storage element wrapping a TurboQuant-encoded byte. Layout queries delegate //! to [`turboquant::quantization::TurboQuantizer`] so the on-storage size and -//! padding stays consistent with the encoder. Encoding / decoding / scoring are -//! still stubbed. +//! padding stays consistent with the encoder. +//! +//! Decoded outputs (`slice_to_float_cow`, `decode_for_quantization`) return the +//! full `padded_dim` of floats in the original (un-rotated) basis. The caller +//! — typically the surrounding [`crate::vector_storage::DenseVectorStorage`] +//! impl which carries the original `api_dim` — is responsible for truncating +//! back to the requested length. Keeping the slice truncation at the storage +//! layer (rather than threading `api_dim` through every trait method) keeps +//! the `PrimitiveVectorElement` API metric-aware but otherwise unchanged. use std::alloc::Layout; use std::borrow::Cow; @@ -70,26 +77,9 @@ impl PrimitiveVectorElement for TurboQuantElement { } fn slice_to_float_cow(vector: Cow<[Self]>, distance: Distance) -> Cow<[VectorElementType]> { - // Public/API path: always returns vectors in the original basis. - decode(&vector, distance, /* apply_inverse_rotation */ true) - } - - fn decode_for_quantization( - vector: Cow<[Self]>, - distance: Distance, - ) -> Cow<[VectorElementType]> { - // For L2/Cosine/Dot the score is rotation-invariant, so we keep the - // data in TurboQuant's rotated basis — the downstream quantizer is - // expected to skip its own rotation step. L1 is not invariant under - // rotation; revert before passing to the next stage. - let revert = matches!(distance, Distance::Manhattan); - decode(&vector, distance, revert) - } - - fn is_prerotated_for_quantization(distance: Distance) -> bool { - // Symmetric to `decode_for_quantization`: rotated basis exposed for - // every rotation-invariant metric; not for L1. - !matches!(distance, Distance::Manhattan) + // Output length is `padded_dim` — caller (storage layer with access to + // the original `api_dim`) is expected to trim before exposing to API. + decode(&vector, distance) } fn quantization_preprocess<'a>( @@ -97,6 +87,9 @@ impl PrimitiveVectorElement for TurboQuantElement { distance: Distance, vector: Cow<'a, [Self]>, ) -> Cow<'a, [f32]> { + // Same contract as `slice_to_float_cow`: padded_dim, original basis. + // The quantization pipeline truncates to `api_dim` before handing the + // floats to the downstream encoder. Self::decode_for_quantization(vector, distance) } @@ -124,45 +117,28 @@ impl PrimitiveVectorElement for TurboQuantElement { TurboQuantizer::quantized_size_for(api_dim, BITS, to_tq_distance(distance), MODE); Layout::from_size_align(bytes, align_of::()).expect("valid layout") } - - /// Recover api-level dimension from the slot length. TQ rounds the - /// requested `dim` up to a `bit_size`-aligned `padded_dim`; once rounded - /// up we cannot tell the original dim apart from any other dim mapping to - /// the same padded_dim. By convention this returns `padded_dim` — the - /// effective api_dim after TQ's rounding. - fn api_dim_from_storage_len(storage_len: usize, distance: Distance) -> usize { - // `quantized_size_for(0, …)` returns only the extras trailer (since - // padded_dim(0) == 0 yields zero packed bytes). Subtract it off to - // recover the packed-data byte count. - let extras_size = - TurboQuantizer::quantized_size_for(0, BITS, to_tq_distance(distance), MODE); - let packed_bytes = storage_len - .checked_sub(extras_size) - .expect("storage_len shorter than TurboQuant extras trailer"); - // `padded_bytes_to_dim` is implemented locally rather than via a - // `TQBits` getter because `TQBits::bit_size` is crate-private. Two - // public `quantized_size_for` calls let us infer the relationship. - padded_bytes_to_dim(packed_bytes) - } } -/// Shared decode helper for [`TurboQuantElement::slice_to_float_cow`] and -/// [`TurboQuantElement::decode_for_quantization`]. Builds a stateless -/// quantizer, calls `dequantize`, optionally applies the inverse rotation, -/// and downcasts to `f32`. The output length is `padded_dim` (see -/// `api_dim_from_storage_len` for the dim convention). -fn decode( - vector: &[TurboQuantElement], - distance: Distance, - apply_inverse_rotation: bool, -) -> Cow<'static, [VectorElementType]> { - let api_dim = TurboQuantElement::api_dim_from_storage_len(vector.len(), distance); - let quantizer = TurboQuantizer::new(api_dim, BITS, MODE, to_tq_distance(distance), None); +/// Build a quantizer from a slot length. The slot encodes `padded_dim` +/// elements plus the metric-specific extras trailer; we strip the trailer, +/// recover `padded_dim` from the remaining packed bytes, and hand that as +/// `dim` to `TurboQuantizer::new` (which is idempotent under further padding). +fn quantizer_for_slot(slot_len: usize, distance: Distance) -> TurboQuantizer { + let tq_distance = to_tq_distance(distance); + let extras_size = TurboQuantizer::quantized_size_for(0, BITS, tq_distance, MODE); + let packed_bytes = slot_len + .checked_sub(extras_size) + .expect("slot shorter than TurboQuant extras trailer"); + let padded_dim = padded_bytes_to_dim(packed_bytes); + TurboQuantizer::new(padded_dim, BITS, MODE, tq_distance, None) +} + +/// Decode a slot into `padded_dim` `f32`s in the original (un-rotated) basis. +fn decode(vector: &[TurboQuantElement], distance: Distance) -> Cow<'static, [VectorElementType]> { + let quantizer = quantizer_for_slot(vector.len(), distance); let bytes: &[u8] = bytemuck::cast_slice(vector); let mut deq = quantizer.dequantize(bytes); - if apply_inverse_rotation { - quantizer.rotation.apply_inverse(&mut deq); - } + quantizer.rotation.apply_inverse(&mut deq); Cow::Owned(deq.into_iter().map(|x| x as f32).collect()) } @@ -183,8 +159,8 @@ fn padded_bytes_to_dim(packed_bytes: usize) -> usize { } /// Symmetric score helper shared by every `Metric` impl. -/// Recovers `api_dim` from the slot length, builds a stateless quantizer, -/// reinterprets both slices as bytes, and delegates to `score_symmetric`. +/// Builds a stateless quantizer from the slot length, reinterprets both +/// slices as bytes, and delegates to `score_symmetric`. fn turbo_score_symmetric( v1: &[TurboQuantElement], v2: &[TurboQuantElement], @@ -195,8 +171,7 @@ fn turbo_score_symmetric( v2.len(), "TurboQuant symmetric score requires matching slot lengths" ); - let api_dim = TurboQuantElement::api_dim_from_storage_len(v1.len(), distance); - let quantizer = TurboQuantizer::new(api_dim, BITS, MODE, to_tq_distance(distance), None); + let quantizer = quantizer_for_slot(v1.len(), distance); let v1_bytes: &[u8] = bytemuck::cast_slice(v1); let v2_bytes: &[u8] = bytemuck::cast_slice(v2); quantizer.score_symmetric(v1_bytes, v2_bytes) diff --git a/lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs b/lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs index c3efcdfe21..0870eaa8f3 100644 --- a/lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs @@ -17,7 +17,7 @@ use crate::common::flags::bitvec_flags::BitvecFlags; use crate::common::flags::dynamic_stored_flags::DynamicStoredFlags; use crate::common::operation_error::{OperationResult, check_process_stopped}; use crate::data_types::named_vectors::CowVector; -use crate::data_types::primitive::PrimitiveVectorElement; +use crate::data_types::primitive::{PrimitiveVectorElement, truncate_to_api_dim}; use crate::data_types::vectors::{VectorElementType, VectorRef}; use crate::types::{Distance, VectorStorageDatatype}; use crate::vector_storage::chunked_vectors::ChunkedVectors; @@ -30,6 +30,11 @@ const DELETED_DIR_PATH: &str = "deleted"; #[derive(Debug)] pub struct AppendableMmapDenseVectorStorage { + /// Api-level vector dimension as requested by the collection config. The + /// underlying `ChunkedVectors` slot may be longer (e.g. TurboQuant pads up + /// to the codebook alignment); `api_dim` is the canonical "vector length" + /// reported to callers and used to truncate decoded outputs. + api_dim: usize, vectors: ChunkedVectors, /// Flags marking deleted vectors /// @@ -79,9 +84,12 @@ impl AppendableMmapDenseVectorStorage { } impl DenseVectorStorage for AppendableMmapDenseVectorStorage { + fn vector_dim(&self) -> usize { + self.api_dim + } + fn vector_layout(&self) -> Layout { - let api_dim = T::api_dim_from_storage_len(self.vectors.dim(), self.distance); - T::storage_layout(api_dim, self.distance) + T::storage_layout(self.api_dim, self.distance) } fn get_dense(&self, key: PointOffsetType) -> Cow<'_, [T]> { @@ -115,14 +123,18 @@ impl VectorStorageRead for AppendableMmapDenseVectorS fn get_vector(&self, key: PointOffsetType) -> CowVector<'_> { self.vectors .get::

(key as VectorOffsetType) - .map(|slice| CowVector::from(T::slice_to_float_cow(slice, self.distance))) + .map(|slice| { + let decoded = T::slice_to_float_cow(slice, self.distance); + CowVector::from(truncate_to_api_dim(decoded, self.api_dim)) + }) .expect("Vector not found") } fn get_vector_opt(&self, key: PointOffsetType) -> Option> { - self.vectors - .get::

(key as VectorOffsetType) - .map(|slice| CowVector::from(T::slice_to_float_cow(slice, self.distance))) + self.vectors.get::

(key as VectorOffsetType).map(|slice| { + let decoded = T::slice_to_float_cow(slice, self.distance); + CowVector::from(truncate_to_api_dim(decoded, self.api_dim)) + }) } fn is_deleted_vector(&self, key: PointOffsetType) -> bool { @@ -277,6 +289,7 @@ pub fn open_appendable_memmap_vector_storage_impl( let deleted_count = deleted.count_trues(); Ok(AppendableMmapDenseVectorStorage { + api_dim: dim, vectors, deleted, distance, diff --git a/lib/segment/src/vector_storage/dense/dense_vector_storage.rs b/lib/segment/src/vector_storage/dense/dense_vector_storage.rs index d3c9e5cf69..e750cea1ac 100644 --- a/lib/segment/src/vector_storage/dense/dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/dense/dense_vector_storage.rs @@ -17,7 +17,7 @@ use fs_err::{File, OpenOptions}; use crate::common::Flusher; use crate::common::operation_error::{OperationError, OperationResult, check_process_stopped}; use crate::data_types::named_vectors::CowVector; -use crate::data_types::primitive::PrimitiveVectorElement; +use crate::data_types::primitive::{PrimitiveVectorElement, truncate_to_api_dim}; use crate::data_types::vectors::VectorRef; use crate::types::{Distance, VectorStorageDatatype}; #[cfg(target_os = "linux")] @@ -44,6 +44,11 @@ where { vectors_path: PathBuf, deleted_path: PathBuf, + /// Api-level vector dimension. The on-disk slot may be longer when `T` + /// rounds the requested dim up (TurboQuant). `api_dim` is the canonical + /// "vector length" callers see and the upper bound used to truncate + /// decoded slots back to the originally requested size. + api_dim: usize, vectors: Option>, distance: Distance, populated: bool, @@ -67,6 +72,7 @@ where let Self { vectors_path: _, deleted_path: _, + api_dim: _, vectors, distance: _, populated: _, @@ -205,6 +211,7 @@ where let storage = DenseVectorStorageImpl { vectors_path, deleted_path, + api_dim: dim, vectors: Some(vectors), distance, populated: populate, @@ -218,10 +225,12 @@ where T: PrimitiveVectorElement, S: UniversalRead, { + fn vector_dim(&self) -> usize { + self.api_dim + } + fn vector_layout(&self) -> Layout { - let slot_len = self.vectors.as_ref().unwrap().dim; - let api_dim = T::api_dim_from_storage_len(slot_len, self.distance); - T::storage_layout(api_dim, self.distance) + T::storage_layout(self.api_dim, self.distance) } fn get_dense(&self, key: PointOffsetType) -> Cow<'_, [T]> { @@ -261,11 +270,15 @@ where fn get_vector(&self, key: PointOffsetType) -> CowVector<'_> { let distance = self.distance; + let api_dim = self.api_dim; self.vectors .as_ref() .unwrap() .get_vector_opt::

(key) - .map(|vector| T::slice_to_float_cow(vector, distance).into()) + .map(|vector| { + let decoded = T::slice_to_float_cow(vector, distance); + CowVector::from(truncate_to_api_dim(decoded, api_dim)) + }) .expect("Vector not found") } @@ -279,24 +292,29 @@ where // `user_data[idx]` available inside the callback. let (user_data, point_offsets): (Vec, Vec) = keys.into_iter().unzip(); let distance = self.distance; + let api_dim = self.api_dim; self.vectors .as_ref() .unwrap() .for_each_in_batch(&point_offsets, |idx, vector| { - let vector = - CowVector::from(T::slice_to_float_cow(Cow::Borrowed(vector), distance)); + let decoded = T::slice_to_float_cow(Cow::Borrowed(vector), distance); + let vector = CowVector::from(truncate_to_api_dim(decoded, api_dim)); callback(user_data[idx], point_offsets[idx], vector); }); } fn get_vector_opt(&self, key: PointOffsetType) -> Option> { let distance = self.distance; + let api_dim = self.api_dim; self.vectors .as_ref() .unwrap() .get_vector_opt::

(key) - .map(|vector| T::slice_to_float_cow(vector, distance).into()) + .map(|vector| { + let decoded = T::slice_to_float_cow(vector, distance); + CowVector::from(truncate_to_api_dim(decoded, api_dim)) + }) } fn is_deleted_vector(&self, key: PointOffsetType) -> bool { @@ -331,7 +349,7 @@ where other_vectors: &'a mut impl Iterator, bool)>, stopped: &AtomicBool, ) -> OperationResult> { - let slot_len = self.vectors.as_ref().unwrap().dim; + let slot_len = T::storage_len_in_elements(self.api_dim, self.distance); let start_index = self.vectors.as_ref().unwrap().num_vectors as PointOffsetType; let mut end_index = start_index; diff --git a/lib/segment/src/vector_storage/dense/empty_dense_vector_storage.rs b/lib/segment/src/vector_storage/dense/empty_dense_vector_storage.rs index 2ccf5c6ef3..95811cc42b 100644 --- a/lib/segment/src/vector_storage/dense/empty_dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/dense/empty_dense_vector_storage.rs @@ -87,6 +87,10 @@ pub fn new_empty_dense_vector_storage( } impl DenseVectorStorage for EmptyDenseVectorStorage { + fn vector_dim(&self) -> usize { + self.dim + } + fn vector_layout(&self) -> Layout { VectorElementType::storage_layout(self.dim, self.distance) } diff --git a/lib/segment/src/vector_storage/dense/volatile_dense_vector_storage.rs b/lib/segment/src/vector_storage/dense/volatile_dense_vector_storage.rs index a5a74cfba7..3c69485a70 100644 --- a/lib/segment/src/vector_storage/dense/volatile_dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/dense/volatile_dense_vector_storage.rs @@ -11,7 +11,7 @@ use common::types::PointOffsetType; use crate::common::Flusher; use crate::common::operation_error::{OperationResult, check_process_stopped}; use crate::data_types::named_vectors::CowVector; -use crate::data_types::primitive::PrimitiveVectorElement; +use crate::data_types::primitive::{PrimitiveVectorElement, truncate_to_api_dim}; use crate::data_types::vectors::{VectorElementType, VectorRef}; use crate::types::{Distance, VectorStorageDatatype}; use crate::vector_storage::volatile_chunked_vectors::VolatileChunkedVectors; @@ -24,6 +24,9 @@ use crate::vector_storage::{ /// This storage is not persisted and intended for temporary use in tests. #[derive(Debug)] pub struct VolatileDenseVectorStorage { + /// Api-level vector dimension. The chunked slot may be longer when `T` + /// pads (TurboQuant); `api_dim` is the canonical "vector length". + api_dim: usize, distance: Distance, vectors: VolatileChunkedVectors, /// BitVec for deleted flags. Grows dynamically upto last set flag. @@ -58,6 +61,7 @@ impl VolatileDenseVectorStorage { pub fn new(dim: usize, distance: Distance) -> Self { let slot_len = T::storage_len_in_elements(dim, distance); Self { + api_dim: dim, distance, vectors: VolatileChunkedVectors::new(slot_len), deleted: BitVec::new(), @@ -84,9 +88,12 @@ impl VolatileDenseVectorStorage { } impl DenseVectorStorage for VolatileDenseVectorStorage { + fn vector_dim(&self) -> usize { + self.api_dim + } + fn vector_layout(&self) -> Layout { - let api_dim = T::api_dim_from_storage_len(self.vectors.dim(), self.distance); - T::storage_layout(api_dim, self.distance) + T::storage_layout(self.api_dim, self.distance) } fn get_dense(&self, key: PointOffsetType) -> Cow<'_, [T]> { @@ -119,9 +126,11 @@ impl VectorStorageRead for VolatileDenseVectorStorage fn get_vector_opt(&self, key: PointOffsetType) -> Option> { // In memory so no optimization to be done for access pattern let distance = self.distance; - self.vectors - .get_opt(key as VectorOffsetType) - .map(|slice| CowVector::from(T::slice_to_float_cow(slice.into(), distance))) + let api_dim = self.api_dim; + self.vectors.get_opt(key as VectorOffsetType).map(|slice| { + let decoded = T::slice_to_float_cow(slice.into(), distance); + CowVector::from(truncate_to_api_dim(decoded, api_dim)) + }) } fn is_deleted_vector(&self, key: PointOffsetType) -> bool { diff --git a/lib/segment/src/vector_storage/quantized/quantized_vectors.rs b/lib/segment/src/vector_storage/quantized/quantized_vectors.rs index 0f0a5350e5..e70507b2b4 100644 --- a/lib/segment/src/vector_storage/quantized/quantized_vectors.rs +++ b/lib/segment/src/vector_storage/quantized/quantized_vectors.rs @@ -24,7 +24,7 @@ use super::quantized_multivector_storage::{ use super::quantized_scorer_builder::QuantizedScorerBuilder; use crate::common::Flusher; use crate::common::operation_error::{OperationError, OperationResult}; -use crate::data_types::primitive::PrimitiveVectorElement; +use crate::data_types::primitive::{PrimitiveVectorElement, truncate_to_api_dim}; use crate::data_types::vectors::{QueryVector, VectorElementType, VectorRef}; use crate::types::{ BinaryQuantization, BinaryQuantizationConfig, BinaryQuantizationEncoding, @@ -808,7 +808,14 @@ impl QuantizedVectors { let datatype = vector_storage.datatype(); let vectors = (0..count as PointOffsetType).map(|i| { let vector = vector_storage.get_dense::(i); - PrimitiveVectorElement::quantization_preprocess(quantization_config, distance, vector) + let preprocessed = PrimitiveVectorElement::quantization_preprocess( + quantization_config, + distance, + vector, + ); + // Storage may decode to a `padded_dim`-long Cow (TurboQuant); the + // downstream quantizer expects api-level vectors. Trim to `dim`. + truncate_to_api_dim(preprocessed, dim) }); let on_disk_vector_storage = vector_storage.is_on_disk(); @@ -870,7 +877,7 @@ impl QuantizedVectors { on_disk_vector_storage, max_threads, stopped, - TElement::is_prerotated_for_quantization(distance), + false, )?, }; @@ -908,7 +915,12 @@ impl QuantizedVectors { let datatype = vector_storage.datatype(); let multi_vector_config = *vector_storage.multi_vector_config(); let vectors = vector_storage.iterate_inner_vectors().map(|vector| { - PrimitiveVectorElement::quantization_preprocess(quantization_config, distance, vector) + let preprocessed = PrimitiveVectorElement::quantization_preprocess( + quantization_config, + distance, + vector, + ); + truncate_to_api_dim(preprocessed, dim) }); let inner_vectors_count = vectors.clone().count(); let vectors_count = vector_storage.total_vector_count(); @@ -1000,7 +1012,7 @@ impl QuantizedVectors { on_disk_vector_storage, max_threads, stopped, - TElement::is_prerotated_for_quantization(distance), + false, )?, }; diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index a67e75f221..6c4ae372b2 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -180,20 +180,16 @@ pub trait VectorStorage: VectorStorageRead { } pub trait DenseVectorStorage: VectorStorageRead { + /// Api-level dimension of vectors stored here — the "vector length" + /// callers passed in via the collection config. May be strictly smaller + /// than the slot's `T`-element count when `T` rounds up internally + /// (TurboQuant pads to the codebook alignment). + fn vector_dim(&self) -> usize; + /// Memory layout of a single on-storage vector slot. Source of truth for /// the slot's byte size and `T`-element count. fn vector_layout(&self) -> Layout; - /// Api-level dimension of vectors stored here. - /// - /// Derived from `vector_layout()` via `T::api_dim_from_storage_len`. - fn vector_dim(&self) -> usize { - T::api_dim_from_storage_len( - self.vector_layout().size() / std::mem::size_of::(), - self.distance(), - ) - } - fn get_dense(&self, key: PointOffsetType) -> Cow<'_, [T]>; /// Call `f` with the raw bytes of the vector if it exists.