From 214d82c71cafd3e9ca7690a6dfdffa7dfe581d64 Mon Sep 17 00:00:00 2001 From: Jojii <15957865+JojiiOfficial@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:12:04 +0200 Subject: [PATCH] IO uring for TQDT (#9852) --- lib/segment/Cargo.toml | 4 + lib/segment/benches/turbo_vector_search.rs | 228 ++++++++++ .../quantized/quantized_storage.rs | 78 +++- .../query_scorer/turbo_custom_query_scorer.rs | 19 + .../query_scorer/turbo_query_scorer.rs | 14 + lib/segment/src/vector_storage/turbo/mod.rs | 416 +++++++++++++++++- .../turbo/turbo_encoded_vectors.rs | 88 +++- .../src/vector_storage/vector_storage_base.rs | 37 +- 8 files changed, 859 insertions(+), 25 deletions(-) create mode 100644 lib/segment/benches/turbo_vector_search.rs diff --git a/lib/segment/Cargo.toml b/lib/segment/Cargo.toml index 772ef24ea8..e05fc9ba4c 100644 --- a/lib/segment/Cargo.toml +++ b/lib/segment/Cargo.toml @@ -128,6 +128,10 @@ io-uring = "0.7.12" name = "vector_search" harness = false +[[bench]] +name = "turbo_vector_search" +harness = false + [[bench]] name = "read_vectors" harness = false diff --git a/lib/segment/benches/turbo_vector_search.rs b/lib/segment/benches/turbo_vector_search.rs new file mode 100644 index 0000000000..77a0c20e15 --- /dev/null +++ b/lib/segment/benches/turbo_vector_search.rs @@ -0,0 +1,228 @@ +//! mmap vs io_uring batch scoring for `TurboVectorStorage` (TQDT). +//! +//! Three modes over the same single-file on-disk dataset: +//! - `unbatched-mmap`: per-point `score_point` loop — the pre-batching behavior; +//! - `batched-mmap`: `score_stored_batch` over the mmap backend; +//! - `batched-uring`: `score_stored_batch` over the io_uring backend (Linux). +//! +//! Groups: +//! - `tq-scoring-warm` / `tq-scoring-cold`: the fair pair — both score the same +//! shuffled random-id subset (HNSW-like scatter); the ONLY difference is that +//! the cold group drops the page cache before every iteration. +//! - `tq-scoring-scan-warm`: full ascending scan, page-cache warm — the plain +//! (non-indexed) search shape. Context only: sequential cached reads are +//! mmap's absolute best case, so this group is not comparable to the +//! scattered-read groups above. +//! +//! The dataset lives under `CARGO_TARGET_TMPDIR`, NOT the system tempdir: +//! `/tmp` is commonly tmpfs, where `clear_cache()` cannot evict anything and +//! cold numbers would silently measure RAM. + +use std::hint::black_box; +use std::path::Path; +use std::time::Duration; + +use common::bitvec::BitSlice; +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use criterion::measurement::WallTime; +use criterion::{BatchSize, BenchmarkGroup, Criterion, criterion_group, criterion_main}; +use rand::distr::StandardUniform; +use rand::seq::{IteratorRandom, SliceRandom}; +use rand::{Rng, RngExt}; +use segment::data_types::vectors::{DenseVector, QueryVector}; +use segment::fixtures::payload_context_fixture::create_id_tracker_fixture; +use segment::id_tracker::IdTrackerRead; +use segment::index::hnsw_index::point_scorer::BatchFilteredSearcher; +use segment::types::Distance; +use segment::vector_storage::turbo::{ + open_appendable_turbo_vector_storage, open_turbo_vector_storage_with_uring, +}; +use segment::vector_storage::{ + DEFAULT_STOPPED, DenseTQVectorStorage, VectorStorage, VectorStorageEnum, new_raw_scorer, +}; +use tempfile::TempDir; + +const DIM: usize = 1024; +const DISTANCE: Distance = Distance::Dot; +const VECTORS: usize = 200_000; +/// Random ids scored per subset iteration — models HNSW's scattered reads +/// while keeping one iteration tractable even in the fully serial unbatched +/// cold mode. +const SUBSET: usize = 4096; +const TOP: usize = 10; + +fn random_vector(rng: &mut impl Rng, size: usize) -> DenseVector { + rng.sample_iter(StandardUniform).take(size).collect() +} + +fn random_query() -> QueryVector { + QueryVector::from(random_vector(&mut rand::rng(), DIM)) +} + +/// Ids to score in one subset iteration: a shuffled sample without replacement. +fn subset_ids() -> Vec { + let mut rng = rand::rng(); + let mut ids: Vec = (0..VECTORS as PointOffsetType).sample(&mut rng, SUBSET); + ids.shuffle(&mut rng); + ids +} + +/// Build the single-file TQ dataset once: encode through an in-RAM appendable +/// storage, then bulk-append the encoded bytes into the single-file layout, +/// exactly as the optimizer does. +fn build_dataset(dir: &Path) { + let mut rng = rand::rng(); + let hw_counter = HardwareCounterCell::new(); + + let encoder_dir = TempDir::new().expect("encoder tempdir created"); + let mut encoder = open_appendable_turbo_vector_storage(encoder_dir.path(), DIM, DISTANCE, true) + .expect("encoder storage created"); + for i in 0..VECTORS { + let vector = random_vector(&mut rng, DIM); + encoder + .insert_vector(i as PointOffsetType, vector.as_slice().into(), &hw_counter) + .expect("vector inserted"); + } + + let mut storage = open_turbo_vector_storage_with_uring(dir, DIM, DISTANCE, false, false) + .expect("single-file storage created"); + let mut encoded = + (0..VECTORS as PointOffsetType).map(|key| (encoder.get_quantized_vector(key), false)); + DenseTQVectorStorage::update_from(&mut storage, &mut encoded, &DEFAULT_STOPPED) + .expect("dataset built"); +} + +/// Score `ids` one point at a time — the exact pre-batching read pattern. +fn score_unbatched(storage: &VectorStorageEnum, ids: impl Iterator) { + let scorer = new_raw_scorer(random_query(), storage, HardwareCounterCell::new()) + .expect("scorer created"); + let mut acc = 0.0; + for id in ids { + acc += scorer.score_point(id); + } + black_box(acc); +} + +/// Score the shuffled subset in every mode. Warm and cold both run through +/// here so the two groups differ in nothing but `clear_cache`. +fn bench_subset( + group: &mut BenchmarkGroup<'_, WallTime>, + modes: &[(&str, bool, &VectorStorageEnum)], + point_deleted: &BitSlice, + clear_cache: bool, +) { + for &(label, batched, storage) in modes { + if !clear_cache { + storage.populate().expect("storage populated"); + } + group.bench_function(label, |b| { + b.iter_batched( + || { + if clear_cache { + storage.clear_cache().expect("cache cleared"); + } + subset_ids() + }, + |ids| { + if batched { + let results = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&random_query()), + storage, + point_deleted, + TOP, + ) + .peek_top_iter(ids.iter().copied(), &DEFAULT_STOPPED) + .expect("points scored"); + black_box(results); + } else { + score_unbatched(storage, ids.iter().copied()); + } + }, + BatchSize::PerIteration, + ) + }); + } +} + +fn benchmark(c: &mut Criterion) { + let data_dir = tempfile::Builder::new() + .prefix("turbo-vector-search-bench") + .tempdir_in(env!("CARGO_TARGET_TMPDIR")) + .expect("bench data dir created"); + build_dataset(data_dir.path()); + + let mmap_storage = VectorStorageEnum::DenseTurbo(Box::new( + open_turbo_vector_storage_with_uring(data_dir.path(), DIM, DISTANCE, false, false) + .expect("mmap storage opened"), + )); + + let mut modes: Vec<(&str, bool, &VectorStorageEnum)> = vec![ + ("unbatched-mmap", false, &mmap_storage), + ("batched-mmap", true, &mmap_storage), + ]; + + #[cfg(target_os = "linux")] + let uring_storage = VectorStorageEnum::DenseTurbo(Box::new( + open_turbo_vector_storage_with_uring(data_dir.path(), DIM, DISTANCE, false, true) + .expect("uring storage opened"), + )); + #[cfg(target_os = "linux")] + modes.push(("batched-uring", true, &uring_storage)); + + let id_tracker = create_id_tracker_fixture(VECTORS); + + let mut warm = c.benchmark_group("tq-scoring-warm"); + warm.sample_size(20); + bench_subset( + &mut warm, + &modes, + id_tracker.deleted_point_bitslice(), + false, + ); + warm.finish(); + + let mut cold = c.benchmark_group("tq-scoring-cold"); + cold.sample_size(10); + cold.warm_up_time(Duration::from_secs(1)); + bench_subset(&mut cold, &modes, id_tracker.deleted_point_bitslice(), true); + cold.finish(); + + // Context group: full ascending scan over everything, page-cache warm — + // the plain (non-indexed) search shape and mmap's best case. + let mut scan = c.benchmark_group("tq-scoring-scan-warm"); + scan.sample_size(20); + for &(label, batched, storage) in &modes { + storage.populate().expect("storage populated"); + scan.bench_function(label, |b| { + b.iter_batched( + random_query, + |query| { + if batched { + let results = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + id_tracker.deleted_point_bitslice(), + TOP, + ) + .peek_top_all(&DEFAULT_STOPPED) + .expect("points scored"); + black_box(results); + } else { + score_unbatched(storage, 0..VECTORS as PointOffsetType); + } + }, + BatchSize::SmallInput, + ) + }); + } + scan.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default(); + targets = benchmark, +} + +criterion_main!(benches); diff --git a/lib/segment/src/vector_storage/quantized/quantized_storage.rs b/lib/segment/src/vector_storage/quantized/quantized_storage.rs index f555c7134d..baba9cf8ce 100644 --- a/lib/segment/src/vector_storage/quantized/quantized_storage.rs +++ b/lib/segment/src/vector_storage/quantized/quantized_storage.rs @@ -1,11 +1,13 @@ use std::borrow::Cow; use std::io::BufWriter; use std::marker::PhantomData; +use std::mem::MaybeUninit; use std::num::NonZeroUsize; use std::path::{Path, PathBuf}; use common::counter::hardware_counter::HardwareCounterCell; -use common::generic_consts::Random; +use common::generic_consts::{AccessPattern, Random, Sequential}; +use common::maybe_uninit::maybe_uninit_fill_from; use common::mmap::{AdviceSetting, MmapFlusher, advice}; use common::types::PointOffsetType; use common::universal_io::{ @@ -16,6 +18,8 @@ use fs_err as fs; use memmap2::MmapMut; use crate::common::operation_error::{OperationError, OperationResult}; +use crate::vector_storage::common::VECTOR_READ_BATCH_SIZE; +use crate::vector_storage::query_scorer::is_read_with_prefetch_efficient; #[derive(Debug)] pub struct QuantizedStorage { @@ -43,16 +47,78 @@ impl QuantizedStorage { } } -impl QuantizedStorage { - /// Open the backing file for build-time bulk appends, bypassing the read-only mmap. +impl QuantizedStorage { + /// Open the backing file for build-time bulk appends, bypassing the read-only handle. pub(crate) fn open_appender(&self) -> std::io::Result> { Ok(BufWriter::new(open_append(&self.path)?)) } - /// Re-mmap after the file grew so reads observe appended vectors. Build-time only. - pub(crate) fn reload(&mut self) -> OperationResult<()> { + /// Reopen after the file grew so reads observe appended vectors. Build-time only. + pub(crate) fn reload(&mut self, fs: &S::Fs) -> OperationResult<()> { let path = self.path.clone(); - *self = Self::from_file(&MmapFs, &path, self.quantized_vector_size.get())?; + *self = Self::from_file(fs, &path, self.quantized_vector_size.get())?; + Ok(()) + } + + /// Read one vector with the given access pattern. + fn read_vector(&self, key: PointOffsetType) -> Cow<'_, [u8]> { + let size = self.quantized_vector_size.get() as u64; + self.storage + .read::(ReadRange { + byte_offset: size * u64::from(key), + length: size, + }) + .expect("vector read from quantized storage failed") + } + + /// Run `f` for each vector in the batch, batching the underlying reads. + /// + /// Async-capable backends (io_uring) get the whole batch submitted in one + /// go; mmap-style backends fetch a batch first and then run `f` over it, + /// which is more cache friendly than interleaving fetch and use. + pub fn for_each_in_batch( + &self, + keys: &[PointOffsetType], + mut f: F, + ) -> OperationResult<()> { + if ReadOnly::::kind().can_be_async() { + let size = self.quantized_vector_size.get() as u64; + let ranges = keys.iter().enumerate().map(|(idx, &key)| { + let range = ReadRange { + byte_offset: size * u64::from(key), + length: size, + }; + (idx, range) + }); + + let callback = |idx, bytes: &[u8]| { + f(idx, bytes); + Ok(()) + }; + + // Access pattern does not matter for io_uring. + self.storage.read_batch::(ranges, callback)?; + return Ok(()); + } + + let mut vectors_buffer = [const { MaybeUninit::uninit() }; VECTOR_READ_BATCH_SIZE]; + + for (batch_idx, keys) in keys.chunks(VECTOR_READ_BATCH_SIZE).enumerate() { + let vectors = if is_read_with_prefetch_efficient(keys) { + let iter = keys.iter().map(|&key| self.read_vector::(key)); + maybe_uninit_fill_from(&mut vectors_buffer, iter).0 + } else { + let iter = keys.iter().map(|&key| self.read_vector::(key)); + maybe_uninit_fill_from(&mut vectors_buffer, iter).0 + }; + + let batch_offset = VECTOR_READ_BATCH_SIZE * batch_idx; + + for (vector_idx, vector) in vectors.iter().enumerate() { + f(batch_offset + vector_idx, vector); + } + } + Ok(()) } } diff --git a/lib/segment/src/vector_storage/query_scorer/turbo_custom_query_scorer.rs b/lib/segment/src/vector_storage/query_scorer/turbo_custom_query_scorer.rs index 62f2d1fa2c..6da8ef2d11 100644 --- a/lib/segment/src/vector_storage/query_scorer/turbo_custom_query_scorer.rs +++ b/lib/segment/src/vector_storage/query_scorer/turbo_custom_query_scorer.rs @@ -76,6 +76,25 @@ where }) } + #[inline] + fn score_stored_batch(&self, ids: &[PointOffsetType], scores: &mut [ScoreType]) { + debug_assert_eq!(ids.len(), scores.len()); + + // One vector of IO per point; CPU is counted per sub-query below, + // matching `score_stored`. + self.hardware_counter.vector_io_read().incr_delta(ids.len()); + let cpu_counter = self.hardware_counter.cpu_counter(); + + self.storage + .for_each_in_dense_tq_batch(ids, |idx, bytes| { + scores[idx] = self.query.score_by(|query| { + cpu_counter.incr(); + self.storage.score_query_bytes(query, bytes) + }); + }) + .expect("read TQ vectors"); + } + fn score_internal(&self, _point_a: PointOffsetType, _point_b: PointOffsetType) -> ScoreType { unimplemented!("Custom scorer compares against multiple vectors, not just one"); } diff --git a/lib/segment/src/vector_storage/query_scorer/turbo_query_scorer.rs b/lib/segment/src/vector_storage/query_scorer/turbo_query_scorer.rs index 94eb44a73d..8e2b1e72c3 100644 --- a/lib/segment/src/vector_storage/query_scorer/turbo_query_scorer.rs +++ b/lib/segment/src/vector_storage/query_scorer/turbo_query_scorer.rs @@ -53,6 +53,20 @@ impl QueryScorer for TurboQueryScorer<'_> { self.storage.score_query_bytes(&self.query, &bytes) } + #[inline] + fn score_stored_batch(&self, ids: &[PointOffsetType], scores: &mut [ScoreType]) { + debug_assert_eq!(ids.len(), scores.len()); + + self.hardware_counter.vector_io_read().incr_delta(ids.len()); + self.hardware_counter.cpu_counter().incr_delta(ids.len()); + + self.storage + .for_each_in_dense_tq_batch(ids, |idx, bytes| { + scores[idx] = self.storage.score_query_bytes(&self.query, bytes); + }) + .expect("read TQ vectors"); + } + fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType { self.hardware_counter.cpu_counter().incr(); self.storage.score_internal_encoded(point_a, point_b) diff --git a/lib/segment/src/vector_storage/turbo/mod.rs b/lib/segment/src/vector_storage/turbo/mod.rs index fa11a6b1be..9abbf8ac58 100644 --- a/lib/segment/src/vector_storage/turbo/mod.rs +++ b/lib/segment/src/vector_storage/turbo/mod.rs @@ -22,7 +22,7 @@ use common::bitvec::BitSlice; use common::counter::hardware_counter::HardwareCounterCell; use common::generic_consts::AccessPattern; use common::types::{PointOffsetType, ScoreType}; -use common::universal_io::{MmapFile, MmapFs, Populate}; +use common::universal_io::{MmapFile, MmapFs, Populate, UserData}; use quantization::turboquant::quantization::TurboQuantizer; use quantization::turboquant::{EncodedQueryTQ, TQBits, TQMode, TQRotation}; @@ -232,20 +232,57 @@ impl std::fmt::Debug for TurboVectorStorage { } } -/// Open (create-or-load) a TurboQuant vector storage backed by a single mmap file (non-appendable). -/// Counterpart to `open_dense_vector_storage`. +/// Open (create-or-load) a TurboQuant vector storage backed by a single file (non-appendable). +/// Counterpart to `open_dense_vector_storage`: reads go through io_uring when the +/// async scorer is enabled (Linux), plain mmap otherwise. pub fn open_turbo_vector_storage( path: &Path, dim: usize, distance: Distance, populate: bool, ) -> OperationResult { + #[cfg(target_os = "linux")] + let with_uring = crate::vector_storage::common::get_async_scorer(); + + #[cfg(not(target_os = "linux"))] + let with_uring = false; + + open_turbo_vector_storage_with_uring(path, dim, distance, populate, with_uring) +} + +/// [`open_turbo_vector_storage`] with an explicit backend choice instead of the +/// global async-scorer flag. Falls back to mmap (with an error log) if the +/// io_uring backend cannot be opened. +pub fn open_turbo_vector_storage_with_uring( + path: &Path, + dim: usize, + distance: Distance, + populate: bool, + with_uring: bool, +) -> OperationResult { + // prevent "unused variable" warning + let _ = with_uring; + open_turbo_vector_storage_impl( path, dim, distance, populate, |vectors_path, quantized_vector_size| { + #[cfg(target_os = "linux")] + if with_uring { + match TurboEncodedVectorStorage::open_uring( + vectors_path, + quantized_vector_size, + populate, + ) { + Ok(storage) => return Ok(storage), + Err(err) => { + log::error!("Failed to open io_uring based TurboQuant storage: {err}"); + } + } + } + TurboEncodedVectorStorage::open_mmap(vectors_path, quantized_vector_size, populate) }, ) @@ -365,6 +402,24 @@ impl VectorStorageRead for TurboVectorStorage { self.dequantize_vector(self.storage.get_quantized_vector(key)) } + fn read_vectors( + &self, + keys: impl IntoIterator, + mut callback: impl FnMut(U, PointOffsetType, CowVector<'_>), + ) { + // Split into parallel arrays in one pass (mirrors the dense storages): + // `for_each_in_batch` needs an offsets slice for batched reads, but we + // still want `user_data[idx]` available inside the callback. + let (user_data, point_offsets): (Vec, Vec) = keys.into_iter().unzip(); + + self.storage + .for_each_in_batch(&point_offsets, |idx, bytes| { + let vector = self.dequantize_vector(Cow::Borrowed(bytes)); + callback(user_data[idx], point_offsets[idx], vector); + }) + .expect("read TQ vectors"); + } + fn get_vector_opt(&self, key: PointOffsetType) -> Option> { Some(self.dequantize_vector(self.storage.get_quantized_vector_opt(key)?)) } @@ -443,6 +498,28 @@ impl DenseTQVectorStorage for TurboVectorStorage { self.storage.get_quantized_vector(key) } + fn for_each_in_dense_tq_batch( + &self, + keys: &[PointOffsetType], + f: F, + ) -> OperationResult<()> { + self.storage.for_each_in_batch(keys, f) + } + + fn read_dense_tq_bytes( + &self, + keys: impl IntoIterator, + mut callback: impl FnMut(U, PointOffsetType, Vec), + ) -> OperationResult<()> { + // Same parallel-arrays split as `read_vectors`, minus the dequantization. + let (user_data, point_offsets): (Vec, Vec) = keys.into_iter().unzip(); + + self.storage + .for_each_in_batch(&point_offsets, |idx, bytes| { + callback(user_data[idx], point_offsets[idx], bytes.to_vec()); + }) + } + fn update_from<'a>( &mut self, other_vectors: &mut impl Iterator, bool)>, @@ -1657,4 +1734,337 @@ mod tests { } } } + + // --------------------------------------------------------------------- + // io_uring backend + batched reads. + // --------------------------------------------------------------------- + + /// Build a flushed single-file storage at `dir` from oracle-encoded + /// `inputs`, exactly as the optimizer does, then drop it — leaving the + /// directory ready to be reopened by either single-file backend. + fn build_single_file(dir: &Path, inputs: &[DenseVector], dim: usize, distance: Distance) { + let oracle = Oracle::new(dim, distance); + let stopped = AtomicBool::new(false); + let encoded: Vec> = inputs.iter().map(|v| oracle.encode(v)).collect(); + + let mut storage = + open_turbo_vector_storage_with_uring(dir, dim, distance, false, false).unwrap(); + let mut it = encoded.iter().map(|b| (Cow::from(b.as_slice()), false)); + storage.update_from(&mut it, &stopped).unwrap(); + storage.flusher()().unwrap(); + } + + /// The io_uring backend must read back exactly what the mmap backend and + /// the independent oracle see: byte-identical encoded vectors and f32-exact + /// dequantized reads. + #[cfg(target_os = "linux")] + #[test] + fn uring_backend_matches_mmap_and_oracle() { + const COUNT: usize = 80; + + for (seed, dim) in [(SEEDS[0], 127), (SEEDS[1], 128), (SEEDS[2], 1024)] { + let distance = Distance::Dot; + let dir = Builder::new() + .prefix("turbo_uring_parity") + .tempdir() + .unwrap(); + let inputs = make_vectors(dim, COUNT, seed); + build_single_file(dir.path(), &inputs, dim, distance); + + let mmap = + open_turbo_vector_storage_with_uring(dir.path(), dim, distance, false, false) + .unwrap(); + let uring = + open_turbo_vector_storage_with_uring(dir.path(), dim, distance, false, true) + .unwrap(); + + // The uring open must not have silently fallen back to mmap — a + // fallback would make every check below vacuous. + assert!( + matches!(uring.storage, TurboEncodedVectorStorage::Uring(_)), + "io_uring backend unavailable in this environment (dim {dim})", + ); + + assert_eq!(uring.total_vector_count(), COUNT); + assert_eq!(uring.deleted_vector_count(), 0); + + let oracle = Oracle::new(dim, distance); + for (i, input) in inputs.iter().enumerate() { + let key = i as PointOffsetType; + let expected = oracle.encode(input); + + let uring_bytes = uring.get_quantized_vector(key); + assert_eq!( + uring_bytes.as_ref(), + expected.as_slice(), + "uring encoded bytes diverge from oracle at {i} (dim {dim})", + ); + assert_eq!( + uring_bytes.as_ref(), + mmap.get_quantized_vector(key).as_ref(), + "uring encoded bytes diverge from mmap at {i} (dim {dim})", + ); + + let via_uring = DenseVector::try_from(uring.get_vector::(key)).unwrap(); + let via_mmap = DenseVector::try_from(mmap.get_vector::(key)).unwrap(); + assert_eq!( + via_uring, via_mmap, + "dequantized read diverges between backends at {i} (dim {dim})", + ); + } + } + } + + /// `score_stored_batch` must agree exactly with per-point `score_stored` + /// on every backend and both scorers. Keys are shuffled, contain + /// duplicates, and exceed `VECTOR_READ_BATCH_SIZE`, so chunking and the + /// idx→key mapping are actually exercised. + #[test] + fn score_stored_batch_matches_score_stored() { + use rand::seq::SliceRandom; + + use crate::vector_storage::query::{RecoBestScoreQuery, RecoQuery}; + use crate::vector_storage::query_scorer::QueryScorer; + use crate::vector_storage::query_scorer::turbo_custom_query_scorer::TurboCustomQueryScorer; + use crate::vector_storage::query_scorer::turbo_query_scorer::TurboQueryScorer; + + const DIM: usize = 128; + const COUNT: usize = 100; + + let distance = Distance::Dot; + let seed = SEEDS[0]; + let inputs = make_vectors(DIM, COUNT, seed); + let hw_counter = HardwareCounterCell::new(); + + let mut rng = StdRng::seed_from_u64(seed); + let mut ids: Vec = (0..COUNT as PointOffsetType) + .chain(0..(COUNT / 2) as PointOffsetType) + .collect(); + ids.shuffle(&mut rng); + + // Backends under test: appendable chunked, single-file mmap, and (on + // Linux) single-file uring. + let chunked_dir = Builder::new() + .prefix("turbo_batch_chunked") + .tempdir() + .unwrap(); + let mut chunked = + open_appendable_turbo_vector_storage(chunked_dir.path(), DIM, distance, true).unwrap(); + insert_all(&mut chunked, &inputs, &hw_counter); + + let single_dir = Builder::new() + .prefix("turbo_batch_single") + .tempdir() + .unwrap(); + build_single_file(single_dir.path(), &inputs, DIM, distance); + let mmap = + open_turbo_vector_storage_with_uring(single_dir.path(), DIM, distance, false, false) + .unwrap(); + + let mut storages = vec![("chunked", chunked), ("mmap", mmap)]; + + #[cfg(target_os = "linux")] + { + let uring = + open_turbo_vector_storage_with_uring(single_dir.path(), DIM, distance, false, true) + .unwrap(); + assert!( + matches!(uring.storage, TurboEncodedVectorStorage::Uring(_)), + "io_uring backend unavailable in this environment", + ); + storages.push(("uring", uring)); + } + + for (backend, storage) in &storages { + let nearest = + TurboQueryScorer::new(inputs[0].clone(), storage, HardwareCounterCell::new()); + let reco = TurboCustomQueryScorer::new( + RecoBestScoreQuery::from(RecoQuery::new( + vec![inputs[1].clone()], + vec![inputs[2].clone()], + )), + storage, + HardwareCounterCell::new(), + ); + + let mut nearest_scores = vec![0.0; ids.len()]; + nearest.score_stored_batch(&ids, &mut nearest_scores); + let mut reco_scores = vec![0.0; ids.len()]; + reco.score_stored_batch(&ids, &mut reco_scores); + + for (idx, &id) in ids.iter().enumerate() { + assert_eq!( + nearest_scores[idx], + nearest.score_stored(id), + "nearest batch score diverges at idx {idx} (key {id}, {backend})", + ); + assert_eq!( + reco_scores[idx], + reco.score_stored(id), + "reco batch score diverges at idx {idx} (key {id}, {backend})", + ); + } + } + } + + /// The batched retrieval readers (`read_vectors` and `read_dense_tq_bytes`) + /// must agree exactly with their per-point counterparts on every backend, + /// thread user data to the right offset, and visit each key exactly once. + /// Keys are shuffled, contain duplicates, and exceed + /// `VECTOR_READ_BATCH_SIZE`, so chunking and the idx→key mapping are + /// actually exercised. + #[test] + fn batched_retrieval_matches_per_point_reads() { + use rand::seq::SliceRandom; + + const DIM: usize = 128; + const COUNT: usize = 100; + + let distance = Distance::Dot; + let seed = SEEDS[1]; + let inputs = make_vectors(DIM, COUNT, seed); + let hw_counter = HardwareCounterCell::new(); + + let mut rng = StdRng::seed_from_u64(seed); + let mut ids: Vec = (0..COUNT as PointOffsetType) + .chain(0..(COUNT / 2) as PointOffsetType) + .collect(); + ids.shuffle(&mut rng); + + let chunked_dir = Builder::new() + .prefix("turbo_retr_chunked") + .tempdir() + .unwrap(); + let mut chunked = + open_appendable_turbo_vector_storage(chunked_dir.path(), DIM, distance, true).unwrap(); + insert_all(&mut chunked, &inputs, &hw_counter); + + let single_dir = Builder::new() + .prefix("turbo_retr_single") + .tempdir() + .unwrap(); + build_single_file(single_dir.path(), &inputs, DIM, distance); + let mmap = + open_turbo_vector_storage_with_uring(single_dir.path(), DIM, distance, false, false) + .unwrap(); + + let mut storages = vec![("chunked", chunked), ("mmap", mmap)]; + + #[cfg(target_os = "linux")] + { + let uring = + open_turbo_vector_storage_with_uring(single_dir.path(), DIM, distance, false, true) + .unwrap(); + assert!( + matches!(uring.storage, TurboEncodedVectorStorage::Uring(_)), + "io_uring backend unavailable in this environment", + ); + storages.push(("uring", uring)); + } + + // Tag each key with its input position so the threading is checkable. + let keys: Vec<(usize, PointOffsetType)> = ids.iter().copied().enumerate().collect(); + + for (backend, storage) in &storages { + // Decoded path: `read_vectors` ≡ `get_vector`, tags ride along. + let mut seen = vec![false; keys.len()]; + storage.read_vectors::(keys.iter().copied(), |tag, offset, vector| { + assert_eq!( + ids[tag], offset, + "user data not threaded to its offset ({backend})", + ); + assert!( + !seen[tag], + "key at position {tag} visited twice ({backend})" + ); + seen[tag] = true; + let direct = storage.get_vector::(offset); + assert_eq!( + DenseVector::try_from(vector).unwrap(), + DenseVector::try_from(direct).unwrap(), + "read_vectors disagrees with get_vector at {offset} ({backend})", + ); + }); + assert!( + seen.iter().all(|&s| s), + "not every key was visited ({backend})", + ); + + // Raw-bytes path: `read_dense_tq_bytes` ≡ `get_dense_tq`. + let mut seen = vec![false; keys.len()]; + storage + .read_dense_tq_bytes::(keys.iter().copied(), |tag, offset, bytes| { + assert_eq!( + ids[tag], offset, + "user data not threaded to its offset ({backend})", + ); + assert!( + !seen[tag], + "key at position {tag} visited twice ({backend})" + ); + seen[tag] = true; + assert_eq!( + bytes.as_slice(), + storage.get_dense_tq::(offset).as_ref(), + "read_dense_tq_bytes disagrees with get_dense_tq at {offset} ({backend})", + ); + }) + .unwrap(); + assert!( + seen.iter().all(|&s| s), + "not every key was visited ({backend})", + ); + } + } + + /// Batch scoring must accrue exactly the same hardware-counter totals as + /// scoring the same keys one by one. + #[test] + fn batch_scoring_accumulates_same_hw_counters() { + use common::counter::hardware_accumulator::HwMeasurementAcc; + + use crate::vector_storage::query_scorer::QueryScorer; + use crate::vector_storage::query_scorer::turbo_query_scorer::TurboQueryScorer; + + const DIM: usize = 128; + const COUNT: usize = 100; + + let distance = Distance::Dot; + let inputs = make_vectors(DIM, COUNT, SEEDS[0]); + + // On-disk single-file storage, so the vector-io-read multiplier is active + // and IO accounting is part of the comparison. + let dir = Builder::new().prefix("turbo_batch_hw").tempdir().unwrap(); + build_single_file(dir.path(), &inputs, DIM, distance); + let storage = + open_turbo_vector_storage_with_uring(dir.path(), DIM, distance, false, false).unwrap(); + + let ids: Vec = (0..COUNT as PointOffsetType).collect(); + + let per_point_acc = HwMeasurementAcc::new(); + { + let scorer = TurboQueryScorer::new( + inputs[0].clone(), + &storage, + per_point_acc.get_counter_cell(), + ); + for &id in &ids { + scorer.score_stored(id); + } + } + + let batch_acc = HwMeasurementAcc::new(); + { + let scorer = + TurboQueryScorer::new(inputs[0].clone(), &storage, batch_acc.get_counter_cell()); + let mut scores = vec![0.0; ids.len()]; + scorer.score_stored_batch(&ids, &mut scores); + } + + assert_eq!(batch_acc.get_cpu(), per_point_acc.get_cpu()); + assert_eq!( + batch_acc.get_vector_io_read(), + per_point_acc.get_vector_io_read(), + ); + } } diff --git a/lib/segment/src/vector_storage/turbo/turbo_encoded_vectors.rs b/lib/segment/src/vector_storage/turbo/turbo_encoded_vectors.rs index e8c3ddd4c7..8e7b0eaa00 100644 --- a/lib/segment/src/vector_storage/turbo/turbo_encoded_vectors.rs +++ b/lib/segment/src/vector_storage/turbo/turbo_encoded_vectors.rs @@ -7,7 +7,9 @@ use std::sync::atomic::AtomicBool; use common::counter::hardware_counter::HardwareCounterCell; use common::mmap::MmapFlusher; use common::types::PointOffsetType; -use common::universal_io::{MmapFile, MmapFs}; +#[cfg(target_os = "linux")] +use common::universal_io::{IoUringFile, IoUringFs}; +use common::universal_io::{MmapFile, MmapFs, UniversalRead}; use quantization::EncodedStorage; use crate::common::operation_error::{OperationResult, check_process_stopped}; @@ -19,6 +21,10 @@ pub(super) enum TurboEncodedVectorStorage { /// Single mem-mapped file of encoded vectors. Mmap(QuantizedStorage), + /// Single file of encoded vectors, read through io_uring (non-appendable). + #[cfg(target_os = "linux")] + Uring(QuantizedStorage), + /// Chunked mem-mapped encoded vectors (appendable). ChunkedMmap(QuantizedChunkedStorage), } @@ -38,6 +44,21 @@ impl TurboEncodedVectorStorage { )?)) } + /// Open (create-or-load) the single-file io_uring backend (non-appendable). + #[cfg(target_os = "linux")] + pub(super) fn open_uring( + path: &Path, + quantized_vector_size: usize, + populate: bool, + ) -> OperationResult { + Ok(Self::Uring(QuantizedStorage::::open( + &IoUringFs, + path, + quantized_vector_size, + populate, + )?)) + } + /// Open (create-or-load) the appendable chunked mem-mapped backend. pub(super) fn open_chunked_mmap( path: &Path, @@ -56,6 +77,8 @@ impl TurboEncodedVectorStorage { pub(super) fn get_quantized_vector(&self, key: PointOffsetType) -> Cow<'_, [u8]> { match self { Self::Mmap(s) => s.get_vector_data(key), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.get_vector_data(key), Self::ChunkedMmap(s) => s.get_vector_data(key), } } @@ -64,14 +87,39 @@ impl TurboEncodedVectorStorage { pub(super) fn get_quantized_vector_opt(&self, key: PointOffsetType) -> Option> { match self { Self::Mmap(s) => s.get_vector_data_opt(key), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.get_vector_data_opt(key), Self::ChunkedMmap(s) => s.get_vector_data_opt(key), } } + /// Run `f` for each vector in the batch, batching the underlying reads + /// (io_uring submission batching / mmap prefetch on the single-file + /// backends; the appendable chunked backend reads one record at a time). + pub(super) fn for_each_in_batch( + &self, + keys: &[PointOffsetType], + mut f: F, + ) -> OperationResult<()> { + match self { + Self::Mmap(s) => s.for_each_in_batch(keys, f), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.for_each_in_batch(keys, f), + Self::ChunkedMmap(s) => { + for (idx, &key) in keys.iter().enumerate() { + f(idx, &s.get_vector_data(key)); + } + Ok(()) + } + } + } + /// Number of encoded vectors (including soft-deleted). pub(super) fn vectors_count(&self) -> usize { match self { Self::Mmap(s) => s.vectors_count(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.vectors_count(), Self::ChunkedMmap(s) => s.vectors_count(), } } @@ -79,6 +127,8 @@ impl TurboEncodedVectorStorage { pub(super) fn is_on_disk(&self) -> bool { match self { Self::Mmap(s) => s.is_on_disk(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.is_on_disk(), Self::ChunkedMmap(s) => s.is_on_disk(), } } @@ -87,6 +137,8 @@ impl TurboEncodedVectorStorage { pub(super) fn files(&self) -> Vec { match self { Self::Mmap(s) => s.files(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.files(), Self::ChunkedMmap(s) => s.files(), } } @@ -94,6 +146,8 @@ impl TurboEncodedVectorStorage { pub(super) fn immutable_files(&self) -> Vec { match self { Self::Mmap(s) => s.immutable_files(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.immutable_files(), Self::ChunkedMmap(s) => s.immutable_files(), } } @@ -101,6 +155,8 @@ impl TurboEncodedVectorStorage { pub(super) fn flusher(&self) -> MmapFlusher { match self { Self::Mmap(s) => s.flusher(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.flusher(), Self::ChunkedMmap(s) => s.flusher(), } } @@ -109,6 +165,8 @@ impl TurboEncodedVectorStorage { pub(super) fn populate(&self) -> OperationResult<()> { match self { Self::Mmap(s) => s.populate(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.populate(), Self::ChunkedMmap(s) => s.populate()?, } Ok(()) @@ -118,6 +176,8 @@ impl TurboEncodedVectorStorage { pub(super) fn clear_cache(&self) -> OperationResult<()> { match self { Self::Mmap(s) => s.clear_cache(), + #[cfg(target_os = "linux")] + Self::Uring(s) => s.clear_cache(), Self::ChunkedMmap(s) => s.clear_cache()?, } Ok(()) @@ -135,6 +195,10 @@ impl TurboEncodedVectorStorage { // Therefore, we don't assume it's read-only here and pretend to write. storage.upsert_vector(id, vector, hw_counter) } + #[cfg(target_os = "linux")] + TurboEncodedVectorStorage::Uring(storage) => { + storage.upsert_vector(id, vector, hw_counter) + } TurboEncodedVectorStorage::ChunkedMmap(storage) => { storage.upsert_vector(id, vector, hw_counter) } @@ -148,15 +212,23 @@ impl TurboEncodedVectorStorage { stopped: &AtomicBool, ) -> OperationResult> { match self { - Self::Mmap(storage) => Self::update_from_mmap(storage, vectors, stopped), + Self::Mmap(storage) => { + Self::update_from_single_file(storage, &MmapFs, vectors, stopped) + } + #[cfg(target_os = "linux")] + Self::Uring(storage) => { + Self::update_from_single_file(storage, &IoUringFs, vectors, stopped) + } Self::ChunkedMmap(storage) => Self::update_from_chunked(storage, vectors, stopped), } } - /// Single-file backend: bulk-append encoded bytes to the file, then re-mmap once - /// (mirrors `DenseVectorStorageImpl::update_from`). - fn update_from_mmap<'a>( - storage: &mut QuantizedStorage, + /// Single-file backends: bulk-append encoded bytes to the file, then reopen once + /// through `fs` so reads observe the appended vectors (mirrors + /// `DenseVectorStorageImpl::update_from`). + fn update_from_single_file<'a, S: UniversalRead>( + storage: &mut QuantizedStorage, + fs: &S::Fs, vectors: impl Iterator>, stopped: &AtomicBool, ) -> OperationResult> { @@ -170,13 +242,13 @@ impl TurboEncodedVectorStorage { end_index += 1; } - // Persist + re-mmap so reads observe the appended vectors. + // Persist + reopen so reads observe the appended vectors. writer.flush()?; let file = writer .into_inner() .map_err(io::IntoInnerError::into_error)?; file.sync_data()?; - storage.reload()?; + storage.reload(fs)?; Ok(start_index..end_index) } diff --git a/lib/segment/src/vector_storage/vector_storage_base.rs b/lib/segment/src/vector_storage/vector_storage_base.rs index 6812546b04..60ef85901d 100644 --- a/lib/segment/src/vector_storage/vector_storage_base.rs +++ b/lib/segment/src/vector_storage/vector_storage_base.rs @@ -446,17 +446,38 @@ pub trait DenseTQVectorStorage: VectorStorageRead { .map_err(|_| OperationError::service_error("Layout is too big")) } - /// Run given function for each vector in the dense batch. - /// - /// Implementation can assume that the keys are consecutive + /// Run given function for each vector in the batch, batching the + /// underlying reads where the storage supports it (io_uring submission + /// batching / mmap prefetch for on-disk storages). fn for_each_in_dense_tq_batch( &self, keys: &[PointOffsetType], mut f: F, - ) { + ) -> OperationResult<()> { for (idx, &key) in keys.iter().enumerate() { f(idx, &self.get_dense_tq::(key)); } + Ok(()) + } + + /// Batched byte counterpart of [`VectorStorageRead::read_vectors`]: calls + /// `callback` with the raw encoded bytes of each vector. TQ counterpart of + /// [`DenseVectorStorageRead::read_dense_bytes`], with the same valid-keys + /// precondition. + /// + /// The default reads one vector at a time; storages with batched readers + /// override this so bulk byte reads keep the same read pipelining as + /// `read_vectors` (io_uring submission batching for on-disk storages). + fn read_dense_tq_bytes( + &self, + keys: impl IntoIterator, + mut callback: impl FnMut(U, PointOffsetType, Vec), + ) -> OperationResult<()> { + for (user_data, key) in keys { + let bytes = self.get_dense_tq::

(key); + callback(user_data, key, bytes.to_vec()); + } + Ok(()) } } @@ -1083,10 +1104,10 @@ impl VectorStorageRead for VectorStorageEnum { VectorStorageEnum::DenseAppendableMemmapHalf(v) => { v.read_dense_bytes::(keys, callback) } - // No batched byte readers here (yet): turbo, sparse and - // multi-dense read one vector at a time. - VectorStorageEnum::DenseTurbo(_) - | VectorStorageEnum::SparseVolatile(_) + VectorStorageEnum::DenseTurbo(v) => v.read_dense_tq_bytes::(keys, callback), + // No batched byte readers here (yet): sparse and multi-dense read + // one vector at a time. + VectorStorageEnum::SparseVolatile(_) | VectorStorageEnum::SparseMmap(_) | VectorStorageEnum::MultiDenseVolatile(_) | VectorStorageEnum::MultiDenseAppendableMemmap(_)