Replace some generic parameters with associated types #6711
Conversation
📝 WalkthroughWalkthroughThis change refactors several core traits and their implementations in the quantization and vector storage modules. The Possibly related PRs
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (13)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (6)
lib/segment/src/vector_storage/quantized/quantized_mmap_storage.rs (2)
61-65:unwrap()s leave the builder fallible path uncheckedBoth
flush()andmake_read_only()may fail (I/O, SIGBUS, …). Panicking here will crash the process at index-build time. Propagate the errors instead:- fn build(self) -> QuantizedMmapStorage { - self.mmap.flush().unwrap(); - let mmap = self.mmap.make_read_only().unwrap(); // TODO: remove unwrap - QuantizedMmapStorage { mmap } + fn build(self) -> QuantizedMmapStorage { + // Flush the data to disk first + self.mmap + .flush() + .expect("failed to flush mmap while building QuantizedMmapStorage"); + + let mmap = self + .mmap + .make_read_only() + .expect("failed to switch mmap to read-only"); + + QuantizedMmapStorage { mmap } }(or return
std::io::Result<QuantizedMmapStorage>and bubble the error upstream).
68-70: Bounds check missing inpush_vector_data
cursor_pos + other.len()can overflow the mmap slice and cause UB.
Consider an assert:fn push_vector_data(&mut self, other: &[u8]) { - self.mmap[self.cursor_pos..self.cursor_pos + other.len()].copy_from_slice(other); + debug_assert!( + self.cursor_pos + other.len() <= self.mmap.len(), + "Encoded vector data exceeds allocated mmap size" + ); + self.mmap[self.cursor_pos..self.cursor_pos + other.len()].copy_from_slice(other); self.cursor_pos += other.len(); }lib/segment/src/vector_storage/query_scorer/multi_metric_query_scorer.rs (1)
48-48:size_ofisn’t in scope – compilation will fail
size_of::<TElement>()is used butstd::mem::size_ofis not imported in this module.
Add the import or fully-qualify the call:-use std::mem::MaybeUninit; +use std::mem::{size_of, MaybeUninit};(or
std::mem::size_of::<TElement>()inline).lib/segment/src/vector_storage/query_scorer/multi_custom_query_scorer.rs (1)
66-70: Missingsize_ofimport mirrors the previous file
size_of::<TElement>()is used to compute multipliers butsize_ofis not imported. Same fix as inmulti_metric_query_scorer.rs.lib/segment/src/vector_storage/quantized/quantized_custom_query_scorer.rs (1)
69-72:size_ofused without import – compilation will fail
size_of::<TElement>()is referenced, butsize_ofis not in scope. Add the import or qualify the call:use std::borrow::Cow; use std::marker::PhantomData; +use std::mem::size_of;or
-hardware_counter.set_cpu_multiplier(size_of::<TElement>()); +hardware_counter.set_cpu_multiplier(std::mem::size_of::<TElement>());lib/segment/src/vector_storage/quantized/quantized_multi_custom_query_scorer.rs (1)
74-78: Same missingsize_ofimport here
size_of::<TElement>()is used innew_multiwithout bringing it into scope.use std::marker::PhantomData; +use std::mem::size_of;(or qualify the call).
🧹 Nitpick comments (13)
lib/segment/src/vector_storage/chunked_vectors.rs (1)
251-253: PreferSelfinstead of repeating the concrete typeUsing
Selfkeeps the implementation resilient to renames and reduces duplication:impl quantization::EncodedStorageBuilder for ChunkedVectors<u8> { - type Storage = ChunkedVectors<u8>; + type Storage = Self;Purely cosmetic but improves maintainability.
lib/segment/src/vector_storage/query_scorer/mod.rs (2)
16-18: Nice move to an associated type – consider documentingTVectorThe trait is now more ergonomic; adding a short rustdoc line describing what
TVectorrepresents will help implementers and future maintainers.
34-35: Inline hint for tiny hot loop
scoreis on the hot-path for every search. An#[inline(always)](or at least#[inline]) could be added to encourage inlining across crate boundaries:- fn score(&self, v2: &Self::TVector) -> ScoreType; + #[inline] // hot path + fn score(&self, v2: &Self::TVector) -> ScoreType;Minor, but can shave a few cycles in tight loops.
lib/quantization/src/encoded_vectors.rs (2)
23-25: Associated type introduction looks good but consider trait-object usabilitySwitching from a generic to an associated type simplifies downstream generics – nice.
If you ever needdyn EncodedVectors, the free associated type makes the trait not object-safe. That limitation existed before, but it becomes more apparent now. Documenting it (or adding#[object_safe]tests) would avoid surprises.
36-44: Add bounds onEncodedQueryfor better compile-time guarantees
encode_query/score_pointimplicitly rely onSelf::EncodedQuerybeingSized(stored in structs) and usuallySend + Sync.
Adding explicit bounds can prevent accidental unsized / non-thread-safe implementations:-pub trait EncodedVectors: Sized { - type EncodedQuery; +pub trait EncodedVectors: Sized { + type EncodedQuery: Sized + Send + Sync + 'static;lib/segment/src/vector_storage/quantized/quantized_query_scorer.rs (2)
13-24: Struct remains generic overTElementalthough it is only a marker
TElementis kept solely for thePhantomDatas – fine for type-level bookkeeping.
If you notice growing compile times, consider removing the parameter and usingphantom: PhantomData<fn() -> TElement>in the few scorers that actually depend on it.
59-67: Implementscorewithunreachable!()instead ofunimplemented!()
unimplemented!()triggers a panic without optimisation hints.
Usingunreachable!()conveys intent and allows LLVM to optimise better:- fn score(&self, _v2: &[TElement]) -> ScoreType { - unimplemented!("This method is not expected to be called for quantized scorer"); - } + fn score(&self, _v2: &[TElement]) -> ScoreType { + unreachable!("`score` should never be called for quantized scorers"); + }lib/segment/src/vector_storage/quantized/quantized_multivector_storage.rs (2)
261-270: Capacity hint can avoid reallocations when encoding query
collect()on an iterator of known length reallocates.
Pre-allocate withVec::with_capacity(inner_count)to save a few mallocs:- multi_vector - .multi_vectors() - .map(|inner_vector| self.quantized_storage.encode_query(inner_vector)) - .collect() + { + let inner_cnt = multi_vector.multi_vectors().count(); + let mut result = Vec::with_capacity(inner_cnt); + for v in multi_vector.multi_vectors() { + result.push(self.quantized_storage.encode_query(v)); + } + result + }
233-241:type EncodedQueryalias is clear – consider documenting multi-vector orderBecause the returned
Vecflattens inner vectors in insertion order, a short doc comment on the associated type or onencode_queryexplaining the ordering assumption would help implementors avoid subtle bugs.lib/segment/src/vector_storage/raw_scorer.rs (1)
443-447: Returnimpl RawScorerto avoid unnecessary boxingYou could make the helper zero-cost by returning
impl RawScorer(stilldyn-compatible if needed by callers):-pub fn raw_scorer_from_query_scorer<'a>( - query_scorer: impl QueryScorer + 'a, -) -> OperationResult<Box<dyn RawScorer + 'a>> { - Ok(Box::new(RawScorerImpl { query_scorer })) +pub fn raw_scorer_from_query_scorer<'a>( + query_scorer: impl QueryScorer + 'a, +) -> OperationResult<RawScorerImpl<impl QueryScorer + 'a>> { + Ok(RawScorerImpl { query_scorer }) }Only convert to
Box<dyn RawScorer>at the outermost API boundary if necessary.lib/segment/src/vector_storage/quantized/quantized_multi_query_scorer.rs (2)
39-51: Pre-allocate flattened query buffer
querygrows withextend_from_sliceinside the loop.
Pre-allocating:let mut query = Vec::with_capacity(raw_query.multi_vectors().count() * self.dim);(or derive capacity from
raw_query) reduces reallocations on large inputs.
66-74: Sameunimplemented!()remark as for single-vector scorerConsider replacing with
unreachable!()for optimisation consistency.lib/segment/src/vector_storage/quantized/quantized_scorer_builder.rs (1)
132-136: Consider naming the “shadowed” struct field for clarity
quantized_storage: _same_as_quantized_storage_in_argsexists solely to satisfy the move out ofself.
Although the leading underscore silences lints, the name is opaque. A shorter_unused(or simply omit by pattern..) improves readability.-let Self { - quantized_storage: _same_as_quantized_storage_in_args, - ... -} = self; +let Self { .. } = self;Also applies to: 203-207
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (23)
lib/quantization/src/encoded_storage.rs(2 hunks)lib/quantization/src/encoded_vectors.rs(2 hunks)lib/quantization/src/encoded_vectors_binary.rs(2 hunks)lib/quantization/src/encoded_vectors_pq.rs(4 hunks)lib/quantization/src/encoded_vectors_u8.rs(2 hunks)lib/segment/src/index/hnsw_index/gpu/gpu_vector_storage/gpu_multivectors.rs(1 hunks)lib/segment/src/vector_storage/async_raw_scorer.rs(3 hunks)lib/segment/src/vector_storage/chunked_vectors.rs(1 hunks)lib/segment/src/vector_storage/quantized/quantized_custom_query_scorer.rs(3 hunks)lib/segment/src/vector_storage/quantized/quantized_mmap_storage.rs(1 hunks)lib/segment/src/vector_storage/quantized/quantized_multi_custom_query_scorer.rs(3 hunks)lib/segment/src/vector_storage/quantized/quantized_multi_query_scorer.rs(2 hunks)lib/segment/src/vector_storage/quantized/quantized_multivector_storage.rs(6 hunks)lib/segment/src/vector_storage/quantized/quantized_query_scorer.rs(2 hunks)lib/segment/src/vector_storage/quantized/quantized_scorer_builder.rs(12 hunks)lib/segment/src/vector_storage/quantized/quantized_vectors.rs(2 hunks)lib/segment/src/vector_storage/query_scorer/custom_query_scorer.rs(1 hunks)lib/segment/src/vector_storage/query_scorer/metric_query_scorer.rs(1 hunks)lib/segment/src/vector_storage/query_scorer/mod.rs(2 hunks)lib/segment/src/vector_storage/query_scorer/multi_custom_query_scorer.rs(1 hunks)lib/segment/src/vector_storage/query_scorer/multi_metric_query_scorer.rs(1 hunks)lib/segment/src/vector_storage/query_scorer/sparse_custom_query_scorer.rs(1 hunks)lib/segment/src/vector_storage/raw_scorer.rs(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
lib/segment/src/vector_storage/quantized/quantized_multivector_storage.rs (3)
lib/quantization/src/encoded_vectors_u8.rs (3)
query(346-346)query(352-355)encode_query(327-369)lib/quantization/src/encoded_vectors_pq.rs (1)
encode_query(484-506)lib/quantization/src/encoded_vectors.rs (1)
encode_query(36-36)
lib/segment/src/vector_storage/quantized/quantized_scorer_builder.rs (4)
lib/segment/src/vector_storage/quantized/quantized_custom_query_scorer.rs (1)
new(36-80)lib/segment/src/vector_storage/quantized/quantized_query_scorer.rs (1)
new(33-57)lib/segment/src/vector_storage/quantized/quantized_multi_custom_query_scorer.rs (1)
new_multi(36-86)lib/segment/src/vector_storage/quantized/quantized_multi_query_scorer.rs (1)
new_multi(33-63)
⏰ Context from checks skipped due to timeout of 90000ms (14)
- GitHub Check: Basic TLS/HTTPS tests
- GitHub Check: test-shard-snapshot-api-s3-minio
- GitHub Check: test-consistency
- GitHub Check: test-snapshot-operations-s3-minio
- GitHub Check: test-consensus-compose
- GitHub Check: test-low-resources
- GitHub Check: rust-tests (windows-latest)
- GitHub Check: rust-tests (ubuntu-latest)
- GitHub Check: rust-tests (macos-latest)
- GitHub Check: integration-tests-consensus
- GitHub Check: rust-tests-no-rocksdb (ubuntu-latest)
- GitHub Check: storage-compat-test
- GitHub Check: integration-tests
- GitHub Check: lint
🔇 Additional comments (23)
lib/segment/src/vector_storage/query_scorer/custom_query_scorer.rs (1)
80-83: Same missingsize_ofimport here
size_of::<TElement>()is used on l. 58, but onlyMaybeUninitis imported fromstd::mem. Bringsize_ofinto scope:-use std::mem::MaybeUninit; +use std::mem::{MaybeUninit, size_of};Without this, the crate will not compile.
⛔ Skipped due to learnings
Learnt from: JojiiOfficial PR: qdrant/qdrant#5954 File: lib/segment/src/index/field_index/bool_index/simple_bool_index.rs:219-228 Timestamp: 2025-03-12T13:40:56.980Z Learning: In Rust 1.80.0 and later, `std::mem::size_of` is included in the prelude and doesn't need to be explicitly imported.Learnt from: JojiiOfficial PR: qdrant/qdrant#5954 File: lib/segment/src/index/field_index/bool_index/simple_bool_index.rs:219-228 Timestamp: 2025-03-12T13:40:56.980Z Learning: In Rust 1.80.0 and later, `std::mem::size_of` is included in the prelude and doesn't need to be explicitly imported.lib/segment/src/vector_storage/quantized/quantized_mmap_storage.rs (1)
58-60: Trait migration looks correct
QuantizedMmapStorageBuildernow implementsEncodedStorageBuildervia an associatedStoragetype, matching the new trait definition. No issues spotted.lib/segment/src/vector_storage/query_scorer/multi_metric_query_scorer.rs (1)
90-93: Associated-type migration LGTMThe scorer now declares
type TVectorand compiles against the newQueryScorersignature. No issues with the change itself.lib/segment/src/vector_storage/query_scorer/multi_custom_query_scorer.rs (1)
113-116: Associated-type declaration correct
type TVector = TypedMultiDenseVector<TElement>correctly adapts to the newQueryScorertrait. Looks good.lib/quantization/src/encoded_storage.rs (2)
23-27: Trait refactor is soundIntroducing the associated type
Storageremoves redundancy and keeps the API intuitive. Method signature updated accordingly – 👍
68-70: Implementation forVec<u8>matches the new trait
type Storage = Vec<u8>and unchanged logic compile cleanly.lib/segment/src/index/hnsw_index/gpu/gpu_vector_storage/gpu_multivectors.rs (1)
61-62: Generic bound updated correctly
QuantizedStorage: EncodedVectorsnow relies on the trait’s associatedEncodedQuerytype. Change aligns with the new API.lib/segment/src/vector_storage/quantized/quantized_multivector_storage.rs (1)
159-190: Possible double-count of vector I/O inscore_point_max_similarityYou pre-increment
vector_io_readonce (👍) but each inner call to
self.quantized_storage.score_point(...)is likely to increment the same counter again, effectively over-counting.
Verify thatscore_pointdoesn’t touchvector_io_readwhen the caller already did. If it does, consider adding a “caller already counted” flag or a specialised internal method to avoid duplication.lib/segment/src/vector_storage/raw_scorer.rs (2)
43-47: Nice generic simplificationDropping the explicit
TVectorparameter cleans up the wrapper without behavioural change – good refactor.
709-734: Batch scoring loop is safe and cache-friendly – LGTMChunking with
VECTOR_READ_BATCH_SIZEavoids oversized mutable borrows; implementation is sound.lib/segment/src/vector_storage/async_raw_scorer.rs (1)
27-35: Trait-bound update looks correctThe new
QueryScorer<TVector = [VectorElementType]>bound cleanly follows the associated-type refactor and keeps the compiler‐level guarantees unchanged. No functional or coherence issues spotted.Also applies to: 46-47, 179-180
lib/segment/src/vector_storage/quantized/quantized_vectors.rs (3)
70-72:Binary*Multinow usesu8instead ofu128– verify performance/size trade-offPrevious single-vector variants keep
u128, while the new multivector aliases switch tou8.
This compiles, but it silently changes:
- bit-packing granularity (8 × more elements processed per scalar op, fewer SIMD opportunities);
- memory layout returned by
get_quantized_vector_size_from_params.If the change is intentional, consider documenting the rationale and benchmarking; otherwise reverting to
u128keeps consistency with the non-multi variants.-type BinaryRamMulti = - QuantizedMultivectorStorage<EncodedVectorsBin<u8, ChunkedVectors<u8>>, Vec<MultivectorOffset>>; +// Was u128 in the single-vector storage – revisit if SIMD throughput is important
54-56: Aliases updated to drop the redundant encoded-query parameter – good
The new type aliases read cleaner and match the trait changes.Also applies to: 62-64
9-11: Import cleanup is fine
Unused imports removed; the module compiles leaner.lib/quantization/src/encoded_vectors_binary.rs (2)
176-178: Correct use of the builder’s associatedStoragetypeSwitching the bound to
impl EncodedStorageBuilder<Storage = TStorage>reflects the new trait shape and keeps the type information explicit.
278-283:EncodedVectorsimpl updated withEncodedQueryassociated type – looks goodThe associated-type definition lines up with the trait and downstream code. No further action needed.
lib/segment/src/vector_storage/quantized/quantized_scorer_builder.rs (1)
94-110: Helper calls simplified – no functional changeReplacing the three-parameter generic with
&impl EncodedVectorskeeps monomorphisation while dropping redundant types. 👍Also applies to: 112-128
lib/quantization/src/encoded_vectors_u8.rs (2)
43-45: Builder bound migrated to associated type – correct
The newEncodedStorageBuilder<Storage = TStorage>signature matches the refactor; implementation untouched.
294-296: AssociatedEncodedQuerytype set – implementation consistent
Interface users now rely on the trait’s associated type; no issues found.lib/segment/src/vector_storage/quantized/quantized_custom_query_scorer.rs (2)
14-20: Good migration to associated type.Switching from the explicit
TEncodedQuerygeneric toTEncodedVectors::EncodedQuerytrims noise and makes the struct signature easier to read.
83-92: Verify trait bound vs. unsized associated type
type TVector = [TElement];introduces an unsized associated type. EnsureQueryScorerdefines
type TVector: ?Sized(or similar) – otherwise this will not compile.lib/segment/src/vector_storage/quantized/quantized_multi_custom_query_scorer.rs (1)
14-20: Clean generic reduction acknowledgedThe move to
TEncodedVectors::EncodedQuerykeeps the public surface smaller and avoids phantom types. Nice.lib/quantization/src/encoded_vectors_pq.rs (1)
58-61: Updated builder bound aligns with newEncodedStorageBuilderdesignUsing
impl EncodedStorageBuilder<Storage = TStorage>reflects the new associated-type pattern and removes a redundant generic. Looks correct.
lib/segment/src/vector_storage/query_scorer/sparse_custom_query_scorer.rs
Show resolved
Hide resolved
* Make EncodedVectors::EncodedQuery associated type * Make EncodedStorageBuilder::Storage associated type * Make QueryScorer::TVector associated type * Remove TVector from RawScorerImpl generic parameters
This PR replaces some trait generic parameters with associated types. Consequentially, these generic parameters could be omitted in places where they are redundant.