upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-07-23 11:11:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

are you happy fmt

Browse Source
This commit is contained in:
Ivan Pleshkov
2026-05-20 14:12:06 +02:00
committed by jojii
parent 4aae44760d
commit a95292f3aa
8 changed files with 146 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
lib/segment/src/data_types/primitive.rs
View File

@@ -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::<Self>()
}
}
/// 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)
}
}
}

97
lib/segment/src/data_types/turbo_quant.rs
View File

@@ -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::<f32>()).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<TurboQuantElement>` 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)

27
lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs
View File

@@ -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<T: PrimitiveVectorElement> {
/// 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<T, MmapFile>,
/// Flags marking deleted vectors
///
@@ -79,9 +84,12 @@ impl<T: PrimitiveVectorElement> AppendableMmapDenseVectorStorage<T> {
}
impl<T: PrimitiveVectorElement> DenseVectorStorage<T> for AppendableMmapDenseVectorStorage<T> {
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<P: AccessPattern>(&self, key: PointOffsetType) -> Cow<'_, [T]> {
@@ -115,14 +123,18 @@ impl<T: PrimitiveVectorElement> VectorStorageRead for AppendableMmapDenseVectorS
fn get_vector<P: AccessPattern>(&self, key: PointOffsetType) -> CowVector<'_> {
self.vectors
.get::<P>(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<P: AccessPattern>(&self, key: PointOffsetType) -> Option<CowVector<'_>> {
self.vectors
.get::<P>(key as VectorOffsetType)
.map(|slice| CowVector::from(T::slice_to_float_cow(slice, self.distance)))
self.vectors.get::<P>(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<T: PrimitiveVectorElement>(
let deleted_count = deleted.count_trues();
Ok(AppendableMmapDenseVectorStorage {
api_dim: dim,
vectors,
deleted,
distance,

36
lib/segment/src/vector_storage/dense/dense_vector_storage.rs
View File

@@ -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<ImmutableDenseVectors<T, S>>,
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<P: AccessPattern>(&self, key: PointOffsetType) -> Cow<'_, [T]> {
@@ -261,11 +270,15 @@ where
fn get_vector<P: AccessPattern>(&self, key: PointOffsetType) -> CowVector<'_> {
let distance = self.distance;
let api_dim = self.api_dim;
self.vectors
.as_ref()
.unwrap()
.get_vector_opt::<P>(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<U>, Vec<PointOffsetType>) = 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<P: AccessPattern>(&self, key: PointOffsetType) -> Option<CowVector<'_>> {
let distance = self.distance;
let api_dim = self.api_dim;
self.vectors
.as_ref()
.unwrap()
.get_vector_opt::<P>(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<Item = (CowVector<'a>, bool)>,
stopped: &AtomicBool,
) -> OperationResult<Range<PointOffsetType>> {
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;

4
lib/segment/src/vector_storage/dense/empty_dense_vector_storage.rs
View File

@@ -87,6 +87,10 @@ pub fn new_empty_dense_vector_storage(
}
impl DenseVectorStorage<VectorElementType> for EmptyDenseVectorStorage {
fn vector_dim(&self) -> usize {
self.dim
}
fn vector_layout(&self) -> Layout {
VectorElementType::storage_layout(self.dim, self.distance)
}

21
lib/segment/src/vector_storage/dense/volatile_dense_vector_storage.rs
View File

@@ -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<T: PrimitiveVectorElement> {
/// 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<T>,
/// BitVec for deleted flags. Grows dynamically upto last set flag.
@@ -58,6 +61,7 @@ impl<T: PrimitiveVectorElement> VolatileDenseVectorStorage<T> {
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<T: PrimitiveVectorElement> VolatileDenseVectorStorage<T> {
}
impl<T: PrimitiveVectorElement> DenseVectorStorage<T> for VolatileDenseVectorStorage<T> {
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<P: AccessPattern>(&self, key: PointOffsetType) -> Cow<'_, [T]> {
@@ -119,9 +126,11 @@ impl<T: PrimitiveVectorElement> VectorStorageRead for VolatileDenseVectorStorage
fn get_vector_opt<P: AccessPattern>(&self, key: PointOffsetType) -> Option<CowVector<'_>> {
// 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 {

22
lib/segment/src/vector_storage/quantized/quantized_vectors.rs
View File

@@ -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::<Sequential>(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,
)?,
};

16
lib/segment/src/vector_storage/vector_storage_base.rs
View File

@@ -180,20 +180,16 @@ pub trait VectorStorage: VectorStorageRead {
}
pub trait DenseVectorStorage<T: PrimitiveVectorElement>: 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::<T>(),
self.distance(),
)
}
fn get_dense<P: AccessPattern>(&self, key: PointOffsetType) -> Cow<'_, [T]>;
/// Call `f` with the raw bytes of the vector if it exists.