From 779a66f78382bbc0b90ad2c8a6c23cb51eac0377 Mon Sep 17 00:00:00 2001 From: xzfc <5121426+xzfc@users.noreply.github.com> Date: Fri, 29 May 2026 17:49:49 +0000 Subject: [PATCH] Migrate InvertedIndexCompressedMmap to UIO (#9144) * Migrate InvertedIndexCompressedMmap to UniversalRead * Rehaul search_scratch.rs (was scores_memory_pool.rs) - Use `typed_arena::Arena` instead of `bumpalo::Bump`. Reason: bump don't drop. Drawback: `Arena::new()` is not free, and `Arena::clear()` doesn't exist; but this commit has workarounds. - Use names that make more sense. * Misc fixes * InvertedIndexCompressedMmap: explicit S type parameter --- Cargo.lock | 7 + lib/common/common/src/ext/aligned_vec.rs | 3 +- lib/segment/benches/sparse_index_search.rs | 3 +- lib/segment/src/fixtures/sparse_fixtures.rs | 4 +- .../index/sparse_index/sparse_vector_index.rs | 12 +- .../sparse_vector_index/search.rs | 8 +- lib/segment/src/index/vector_index_base.rs | 7 +- .../sparse_vector_index_search_tests.rs | 23 +- lib/sparse/Cargo.toml | 1 + lib/sparse/benches/search.rs | 16 +- lib/sparse/src/common/mod.rs | 1 - lib/sparse/src/common/scores_memory_pool.rs | 56 ---- .../src/index/compressed_posting_list.rs | 2 +- ...inverted_index_compressed_immutable_ram.rs | 51 ++-- .../inverted_index_compressed_mmap.rs | 286 ++++++++++-------- .../inverted_index/inverted_index_ram.rs | 8 +- lib/sparse/src/index/inverted_index/mod.rs | 9 +- lib/sparse/src/index/search_context.rs | 27 +- lib/sparse/src/index/tests/common.rs | 15 +- lib/sparse/src/index/tests/hw_counter_test.rs | 9 +- .../src/index/tests/indexed_vs_plain_test.rs | 9 +- .../src/index/tests/search_context_tests.rs | 78 +++-- lib/sparse/src/lib.rs | 3 + lib/sparse/src/search_scratch.rs | 102 +++++++ 24 files changed, 444 insertions(+), 296 deletions(-) delete mode 100644 lib/sparse/src/common/scores_memory_pool.rs create mode 100644 lib/sparse/src/search_scratch.rs diff --git a/Cargo.lock b/Cargo.lock index fe1e3370b1..a445c84f57 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7530,6 +7530,7 @@ dependencies = [ "sha2 0.11.0", "sparse", "tempfile", + "typed-arena", "validator", "zerocopy", ] @@ -8382,6 +8383,12 @@ dependencies = [ "rand 0.9.2", ] +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typeid" version = "1.0.3" diff --git a/lib/common/common/src/ext/aligned_vec.rs b/lib/common/common/src/ext/aligned_vec.rs index 64338b43cd..153b96d215 100644 --- a/lib/common/common/src/ext/aligned_vec.rs +++ b/lib/common/common/src/ext/aligned_vec.rs @@ -1,7 +1,8 @@ use std::borrow::Cow; use std::ops::Deref; -use aligned_vec::{AVec, Alignment, RuntimeAlign}; +use aligned_vec::Alignment; +pub use aligned_vec::{AVec, RuntimeAlign}; use bytemuck::PodCastError; #[derive(Debug)] diff --git a/lib/segment/benches/sparse_index_search.rs b/lib/segment/benches/sparse_index_search.rs index cfebd9e3ea..6ae06a53e1 100644 --- a/lib/segment/benches/sparse_index_search.rs +++ b/lib/segment/benches/sparse_index_search.rs @@ -2,6 +2,7 @@ use std::sync::atomic::AtomicBool; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; +use common::universal_io::MmapFile; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; use dataset::Dataset; use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; @@ -104,7 +105,7 @@ fn sparse_vector_index_search_benchmark_impl( let sparse_index_config = SparseIndexConfig::new(Some(FULL_SCAN_THRESHOLD), SparseIndexType::Mmap, None); let pb = progress("Indexing (2/2)", vectors_len); - let sparse_vector_index_mmap: SparseVectorIndex> = + let sparse_vector_index_mmap: SparseVectorIndex> = SparseVectorIndex::open(SparseVectorIndexOpenArgs { config: sparse_index_config, id_tracker: sparse_vector_index.id_tracker().clone(), diff --git a/lib/segment/src/fixtures/sparse_fixtures.rs b/lib/segment/src/fixtures/sparse_fixtures.rs index 871fdb44d4..5dbd8407a8 100644 --- a/lib/segment/src/fixtures/sparse_fixtures.rs +++ b/lib/segment/src/fixtures/sparse_fixtures.rs @@ -118,9 +118,9 @@ macro_rules! fixture_for_all_indices { ::sparse::index::inverted_index::inverted_index_compressed_immutable_ram::InvertedIndexCompressedImmutableRam >($($args)*); - eprintln!("InvertedIndexCompressedMmap"); + eprintln!("InvertedIndexCompressedMmap"); $test::< - ::sparse::index::inverted_index::inverted_index_compressed_mmap::InvertedIndexCompressedMmap + ::sparse::index::inverted_index::inverted_index_compressed_mmap::InvertedIndexCompressedMmap >($($args)*); }; } diff --git a/lib/segment/src/index/sparse_index/sparse_vector_index.rs b/lib/segment/src/index/sparse_index/sparse_vector_index.rs index 7036afd656..59df9e14e8 100644 --- a/lib/segment/src/index/sparse_index/sparse_vector_index.rs +++ b/lib/segment/src/index/sparse_index/sparse_vector_index.rs @@ -9,7 +9,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::generic_consts::Random; use common::storage_version::StorageVersion as _; use fs_err as fs; -use sparse::common::scores_memory_pool::ScoresMemoryPool; +use sparse::SearchScratchPool; use sparse::common::sparse_vector::SparseVector; use sparse::index::inverted_index::InvertedIndex; use sparse::index::inverted_index::inverted_index_ram_builder::InvertedIndexBuilder; @@ -35,7 +35,7 @@ pub struct SparseVectorIndex { inverted_index: TInvertedIndex, searches_telemetry: SparseSearchesTelemetry, indices_tracker: IndicesTracker, - scores_memory_pool: ScoresMemoryPool, + search_scratch_pool: SearchScratchPool, } /// Getters for internals, used for testing. @@ -130,7 +130,7 @@ impl SparseVectorIndex { let searches_telemetry = SparseSearchesTelemetry::new(); let path = path.to_path_buf(); - let scores_memory_pool = ScoresMemoryPool::new(); + let search_scratch_pool = SearchScratchPool::new(); Ok(Self { config, id_tracker, @@ -140,7 +140,7 @@ impl SparseVectorIndex { inverted_index, searches_telemetry, indices_tracker, - scores_memory_pool, + search_scratch_pool, }) } @@ -227,12 +227,14 @@ impl SparseVectorIndex { let hw_counter = HardwareCounterCell::disposable(); let mut unique_record_ids = std::collections::HashSet::new(); + let mut arena = sparse::SearchScratchArena::new_slow(); for dim_id in &query_vector.indices { if let Some(dim_id) = self.indices_tracker.remap_index(*dim_id) { - let posting_list_iter = self.inverted_index.get(dim_id, &hw_counter)?; + let posting_list_iter = self.inverted_index.get(dim_id, &arena, &hw_counter)?; for element in posting_list_iter.into_std_iter() { unique_record_ids.insert(element.record_id); } + arena.gc(); } } Ok(unique_record_ids.len()) diff --git a/lib/segment/src/index/sparse_index/sparse_vector_index/search.rs b/lib/segment/src/index/sparse_index/sparse_vector_index/search.rs index f1cc0cda27..28a0662a88 100644 --- a/lib/segment/src/index/sparse_index/sparse_vector_index/search.rs +++ b/lib/segment/src/index/sparse_index/sparse_vector_index/search.rs @@ -132,7 +132,7 @@ impl SparseVectorIndex { .collect_vec(); let sparse_vector = self.indices_tracker.remap_vector(sparse_vector.clone()); - let memory_handle = self.scores_memory_pool.get(); + let mut scratch = self.search_scratch_pool.get(); let mut hw_counter = vector_query_context.hardware_counter(); let is_index_on_disk = self.config.index_type.is_on_disk(); if is_index_on_disk { @@ -145,7 +145,7 @@ impl SparseVectorIndex { sparse_vector, top, &self.inverted_index, - memory_handle, + &mut scratch, &is_stopped, &hw_counter, )?; @@ -175,7 +175,7 @@ impl SparseVectorIndex { let is_stopped = vector_query_context.is_stopped(); let sparse_vector = self.indices_tracker.remap_vector(sparse_vector.clone()); - let memory_handle = self.scores_memory_pool.get(); + let mut scratch = self.search_scratch_pool.get(); let mut hw_counter = vector_query_context.hardware_counter(); let is_index_on_disk = self.config.index_type.is_on_disk(); if is_index_on_disk { @@ -188,7 +188,7 @@ impl SparseVectorIndex { sparse_vector, top, &self.inverted_index, - memory_handle, + &mut scratch, &is_stopped, &hw_counter, )?; diff --git a/lib/segment/src/index/vector_index_base.rs b/lib/segment/src/index/vector_index_base.rs index 8ab8796e5f..21d84e4756 100644 --- a/lib/segment/src/index/vector_index_base.rs +++ b/lib/segment/src/index/vector_index_base.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use common::counter::hardware_counter::HardwareCounterCell; use common::types::{PointOffsetType, ScoredPointOffset, TelemetryDetail}; +use common::universal_io::MmapFile; use half::f16; use sparse::common::types::{DimId, QuantizedU8}; use sparse::index::inverted_index::InvertedIndex; @@ -95,9 +96,9 @@ pub enum VectorIndexEnum { SparseCompressedImmutableRamU8( SparseVectorIndex>, ), - SparseCompressedMmapF32(SparseVectorIndex>), - SparseCompressedMmapF16(SparseVectorIndex>), - SparseCompressedMmapU8(SparseVectorIndex>), + SparseCompressedMmapF32(SparseVectorIndex>), + SparseCompressedMmapF16(SparseVectorIndex>), + SparseCompressedMmapU8(SparseVectorIndex>), } impl VectorIndexEnum { diff --git a/lib/segment/tests/integration/sparse_vector_index_search_tests.rs b/lib/segment/tests/integration/sparse_vector_index_search_tests.rs index 8a6212ded3..8cbddca7d1 100644 --- a/lib/segment/tests/integration/sparse_vector_index_search_tests.rs +++ b/lib/segment/tests/integration/sparse_vector_index_search_tests.rs @@ -6,6 +6,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::generic_consts::Random; use common::storage_version::VERSION_FILE; use common::types::{PointOffsetType, TelemetryDetail}; +use common::universal_io::MmapFile; use fs_err as fs; use itertools::Itertools; use rand::SeedableRng; @@ -170,9 +171,10 @@ fn check_index_storage_consistency(sparse_vector_index: &Spars .iter() .zip(remapped_vector.values.iter()) { + let arena = sparse::SearchScratchArena::new_slow(); let posting_list = sparse_vector_index .inverted_index() - .get(*dim_id, &hw_counter) + .get(*dim_id, &arena, &hw_counter) .unwrap(); // assert posting list sorted by record id assert!( @@ -221,7 +223,7 @@ fn sparse_vector_index_consistent_with_storage() { // create mmap sparse vector index let mut sparse_index_config = sparse_vector_ram_index.config(); sparse_index_config.index_type = SparseIndexType::Mmap; - let sparse_vector_mmap_index: SparseVectorIndex> = + let sparse_vector_mmap_index: SparseVectorIndex> = SparseVectorIndex::open(SparseVectorIndexOpenArgs { config: sparse_index_config, id_tracker: sparse_vector_ram_index.id_tracker().clone(), @@ -247,7 +249,7 @@ fn sparse_vector_index_consistent_with_storage() { // load index from memmap file let mut sparse_index_config = sparse_vector_ram_index.config(); sparse_index_config.index_type = SparseIndexType::Mmap; - let sparse_vector_mmap_index: SparseVectorIndex> = + let sparse_vector_mmap_index: SparseVectorIndex> = SparseVectorIndex::open(SparseVectorIndexOpenArgs { config: sparse_index_config, id_tracker: sparse_vector_ram_index.id_tracker().clone(), @@ -271,13 +273,14 @@ fn sparse_vector_index_consistent_with_storage() { #[test] fn sparse_vector_index_load_missing_mmap() { let data_dir = Builder::new().prefix("data_dir").tempdir().unwrap(); - let sparse_vector_index: OperationResult>> = - fixture_sparse_index_from_iter( - data_dir.path(), - [].iter().cloned(), - 10_000, - SparseIndexType::Mmap, - ); + let sparse_vector_index: OperationResult< + SparseVectorIndex>, + > = fixture_sparse_index_from_iter( + data_dir.path(), + [].iter().cloned(), + 10_000, + SparseIndexType::Mmap, + ); // absent configuration file for mmap are ignored // a new index is created assert!(sparse_vector_index.is_ok()) diff --git a/lib/sparse/Cargo.toml b/lib/sparse/Cargo.toml index 3bb6da060c..b4e830df2e 100644 --- a/lib/sparse/Cargo.toml +++ b/lib/sparse/Cargo.toml @@ -33,6 +33,7 @@ itertools = { workspace = true } parking_lot = { workspace = true } log = { workspace = true } zerocopy = { workspace = true } +typed-arena = "2.0.2" [dev-dependencies] criterion = { workspace = true } diff --git a/lib/sparse/benches/search.rs b/lib/sparse/benches/search.rs index f3227f74cc..a8ab0267ce 100644 --- a/lib/sparse/benches/search.rs +++ b/lib/sparse/benches/search.rs @@ -4,6 +4,7 @@ use std::sync::atomic::AtomicBool; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; +use common::universal_io::MmapFile; use criterion::measurement::Measurement; use criterion::{Criterion, criterion_group, criterion_main}; use dataset::Dataset; @@ -13,7 +14,7 @@ use itertools::Itertools; use rand::SeedableRng as _; use rand::rngs::StdRng; use sha2::Digest; -use sparse::common::scores_memory_pool::ScoresMemoryPool; +use sparse::SearchScratchPool; use sparse::common::sparse_vector::{RemappedSparseVector, SparseVector}; use sparse::common::sparse_vector_fixture::{random_positive_sparse_vector, random_sparse_vector}; use sparse::common::types::{QuantizedU8, Weight}; @@ -162,7 +163,7 @@ pub fn run_bench( run_bench2( c.benchmark_group(format!("search/mmap_{}/{name}", $name)), - &InvertedIndexCompressedMmap::<$type>::open(&index_path).unwrap(), + &InvertedIndexCompressedMmap::<$type, MmapFile>::open(&index_path).unwrap(), query_vectors, &hottest_query_vectors, ); @@ -181,7 +182,7 @@ fn run_bench2( query_vectors: &[SparseVector], hottest_query_vectors: &[RemappedSparseVector], ) { - let pool = ScoresMemoryPool::new(); + let pool = SearchScratchPool::new(); let stopped = AtomicBool::new(false); let mut it = query_vectors.iter().cycle(); @@ -192,7 +193,8 @@ fn run_bench2( b.iter_batched( || it.next().unwrap().clone().into_remapped(), |vec| { - SearchContext::new(vec, TOP, index, pool.get(), &stopped, &hardware_counter) + let mut scratch = pool.get(); + SearchContext::new(vec, TOP, index, &mut scratch, &stopped, &hardware_counter) .unwrap() .search(&|_| true) }, @@ -207,7 +209,8 @@ fn run_bench2( b.iter_batched( || it.next().unwrap().clone(), |vec| { - SearchContext::new(vec, TOP, index, pool.get(), &stopped, &hardware_counter) + let mut scratch = pool.get(); + SearchContext::new(vec, TOP, index, &mut scratch, &stopped, &hardware_counter) .unwrap() .search(&|_| true) }, @@ -267,7 +270,8 @@ fn cached_compressed_index(index: &InvertedIndexRam, name: &str) -> P fs::remove_dir_all(&tmp_path).unwrap(); } fs::create_dir_all(&tmp_path).unwrap(); - InvertedIndexCompressedMmap::::from_ram_index(Cow::Borrowed(index), &tmp_path).unwrap(); + InvertedIndexCompressedMmap::::from_ram_index(Cow::Borrowed(index), &tmp_path) + .unwrap(); fs::rename(tmp_path, &path).unwrap(); } eprintln!("Using cache: {path:?}"); diff --git a/lib/sparse/src/common/mod.rs b/lib/sparse/src/common/mod.rs index f3e63978a3..85a3937225 100644 --- a/lib/sparse/src/common/mod.rs +++ b/lib/sparse/src/common/mod.rs @@ -1,4 +1,3 @@ -pub mod scores_memory_pool; pub mod sparse_vector; #[cfg(feature = "testing")] pub mod sparse_vector_fixture; diff --git a/lib/sparse/src/common/scores_memory_pool.rs b/lib/sparse/src/common/scores_memory_pool.rs deleted file mode 100644 index d15813eefe..0000000000 --- a/lib/sparse/src/common/scores_memory_pool.rs +++ /dev/null @@ -1,56 +0,0 @@ -use common::defaults::POOL_KEEP_LIMIT; -use common::types::ScoreType; -use parking_lot::Mutex; - -type PooledScores = Vec; - -#[derive(Debug)] -pub struct PooledScoresHandle<'a> { - pool: &'a ScoresMemoryPool, - pub scores: PooledScores, -} - -impl<'a> PooledScoresHandle<'a> { - fn new(pool: &'a ScoresMemoryPool, scores: PooledScores) -> Self { - PooledScoresHandle { pool, scores } - } -} - -impl Drop for PooledScoresHandle<'_> { - fn drop(&mut self) { - self.pool.return_back(std::mem::take(&mut self.scores)); - } -} - -#[derive(Debug)] -pub struct ScoresMemoryPool { - pool: Mutex>, -} - -impl ScoresMemoryPool { - pub fn new() -> Self { - ScoresMemoryPool { - pool: Mutex::new(Vec::with_capacity(*POOL_KEEP_LIMIT)), - } - } - - pub fn get(&self) -> PooledScoresHandle<'_> { - match self.pool.lock().pop() { - None => PooledScoresHandle::new(self, vec![]), - Some(data) => PooledScoresHandle::new(self, data), - } - } - - fn return_back(&self, data: PooledScores) { - let mut pool = self.pool.lock(); - if pool.len() < *POOL_KEEP_LIMIT { - pool.push(data); - } - } -} - -impl Default for ScoresMemoryPool { - fn default() -> Self { - Self::new() - } -} diff --git a/lib/sparse/src/index/compressed_posting_list.rs b/lib/sparse/src/index/compressed_posting_list.rs index 432e6d37a5..904a88e6d8 100644 --- a/lib/sparse/src/index/compressed_posting_list.rs +++ b/lib/sparse/src/index/compressed_posting_list.rs @@ -17,7 +17,7 @@ use crate::common::types::{DimWeight, Weight}; type BitPackerImpl = bitpacking::BitPacker4x; /// How many elements are packed in a single chunk. -const CHUNK_SIZE: usize = BitPackerImpl::BLOCK_LEN; +pub(super) const CHUNK_SIZE: usize = BitPackerImpl::BLOCK_LEN; #[derive(Default, Debug, Clone, PartialEq)] pub struct CompressedPostingList { diff --git a/lib/sparse/src/index/inverted_index/inverted_index_compressed_immutable_ram.rs b/lib/sparse/src/index/inverted_index/inverted_index_compressed_immutable_ram.rs index a09aee992e..304fdbb9b2 100644 --- a/lib/sparse/src/index/inverted_index/inverted_index_compressed_immutable_ram.rs +++ b/lib/sparse/src/index/inverted_index/inverted_index_compressed_immutable_ram.rs @@ -3,17 +3,18 @@ use std::path::Path; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; -use common::universal_io::Result; +use common::universal_io::{MmapFs, Result}; use super::inverted_index_compressed_mmap::InvertedIndexCompressedMmap; use super::inverted_index_ram::InvertedIndexRam; use super::{InvertedIndex, out_of_bounds}; +use crate::SearchScratchArena; use crate::common::sparse_vector::RemappedSparseVector; use crate::common::types::{DimId, DimOffset, Weight}; use crate::index::compressed_posting_list::{ CompressedPostingBuilder, CompressedPostingList, CompressedPostingListIterator, + CompressedPostingListView, }; -use crate::index::posting_list_common::PostingListIter as _; #[derive(Debug, Clone, PartialEq)] pub struct InvertedIndexCompressedImmutableRam { @@ -22,17 +23,19 @@ pub struct InvertedIndexCompressedImmutableRam { pub(super) total_sparse_size: usize, } +type Storage = common::universal_io::MmapFile; + impl InvertedIndex for InvertedIndexCompressedImmutableRam { type Iter<'a> = CompressedPostingListIterator<'a, W>; - type Version = as InvertedIndex>::Version; + type Version = as InvertedIndex>::Version; fn is_on_disk(&self) -> bool { false } fn open(path: &Path) -> Result { - let mmap_inverted_index = InvertedIndexCompressedMmap::load(path)?; + let mmap_inverted_index = InvertedIndexCompressedMmap::::load(&MmapFs, path)?; let mut inverted_index = InvertedIndexCompressedImmutableRam { postings: Vec::with_capacity(mmap_inverted_index.file_header.posting_count), vector_count: mmap_inverted_index.file_header.vector_count, @@ -41,9 +44,11 @@ impl InvertedIndex for InvertedIndexCompressedImmutableRam { let hw_counter = HardwareCounterCell::disposable(); + let mut arena = SearchScratchArena::new_slow(); for i in 0..mmap_inverted_index.file_header.posting_count as DimId { - let posting_list = mmap_inverted_index.get(i, &hw_counter)?; + let posting_list = mmap_inverted_index.get(i, &arena, &hw_counter)?; inverted_index.postings.push(posting_list.to_owned()); + arena.gc(); } mmap_inverted_index.clear_cache()?; @@ -51,21 +56,18 @@ impl InvertedIndex for InvertedIndexCompressedImmutableRam { Ok(inverted_index) } - fn save(&self, path: &Path) -> std::io::Result<()> { - InvertedIndexCompressedMmap::convert_and_save(self, path)?; + fn save(&self, path: &Path) -> Result<()> { + InvertedIndexCompressedMmap::::convert_and_save(&MmapFs, self, path)?; Ok(()) } fn get<'a>( &'a self, id: DimOffset, + _arena: &'a SearchScratchArena, hw_counter: &'a HardwareCounterCell, // Ignored for in-ram index ) -> Result> { - let posting = self - .postings - .get(id as usize) - .ok_or(out_of_bounds(id, self.len()))?; - Ok(posting.iter(hw_counter)) + Ok(self.get(id, hw_counter)?.iter()) } fn len(&self) -> usize { @@ -73,16 +75,16 @@ impl InvertedIndex for InvertedIndexCompressedImmutableRam { } fn posting_list_len(&self, id: DimOffset, hw_counter: &HardwareCounterCell) -> Result { - Ok(self.get(id, hw_counter)?.len_to_end()) + Ok(self.get(id, hw_counter)?.len()) } fn files(path: &Path) -> Vec { - InvertedIndexCompressedMmap::::files(path) + InvertedIndexCompressedMmap::::files(path) } fn immutable_files(path: &Path) -> Vec { // `InvertedIndexCompressedImmutableRam` is always immutable - InvertedIndexCompressedMmap::::immutable_files(path) + InvertedIndexCompressedMmap::::immutable_files(path) } fn remove(&mut self, _id: PointOffsetType, _old_vector: RemappedSparseVector) { @@ -98,10 +100,7 @@ impl InvertedIndex for InvertedIndexCompressedImmutableRam { panic!("Cannot upsert into a read-only RAM inverted index") } - fn from_ram_index>( - ram_index: Cow, - _path: P, - ) -> std::io::Result { + fn from_ram_index>(ram_index: Cow, _path: P) -> Result { let mut postings = Vec::with_capacity(ram_index.postings.len()); for old_posting_list in &ram_index.postings { let mut new_posting_list = CompressedPostingBuilder::new(); @@ -141,6 +140,20 @@ impl InvertedIndex for InvertedIndexCompressedImmutableRam { } } +impl InvertedIndexCompressedImmutableRam { + #[inline] + fn get<'a>( + &'a self, + id: DimOffset, + hw_counter: &'a HardwareCounterCell, + ) -> Result> { + let Some(posting) = self.postings.get(id as usize) else { + return Err(out_of_bounds(id, self.len())); + }; + Ok(posting.view(hw_counter)) + } +} + #[cfg(test)] mod tests { use tempfile::Builder; diff --git a/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs b/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs index 9886ce4756..0d0d31ee88 100644 --- a/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs +++ b/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs @@ -1,28 +1,30 @@ use std::borrow::Cow; +use std::fmt::Debug; use std::io::{BufWriter, Write as _}; use std::marker::PhantomData; use std::mem::size_of; use std::path::{Path, PathBuf}; -use std::sync::Arc; use common::counter::hardware_counter::HardwareCounterCell; +use common::ext::aligned_vec::ACow; use common::fs::{atomic_save_json, read_json}; -use common::mmap::{Advice, AdviceSetting, Madviseable, create_and_ensure_length, open_read_mmap}; +use common::generic_consts::Random; +use common::mmap::{Advice, AdviceSetting, create_and_ensure_length}; #[expect(deprecated, reason = "legacy code")] use common::mmap::{transmute_to_u8, transmute_to_u8_slice}; use common::storage_version::StorageVersion; use common::types::PointOffsetType; -use common::universal_io::Result; -use memmap2::Mmap; +use common::universal_io::{OpenOptions, Populate, Result, UniversalRead, UniversalReadFs}; use serde::{Deserialize, Serialize}; use zerocopy::{FromBytes, Immutable, KnownLayout}; use super::inverted_index_compressed_immutable_ram::InvertedIndexCompressedImmutableRam; use super::{INDEX_FILE_NAME, corrupted_index, out_of_bounds}; +use crate::SearchScratchArena; use crate::common::sparse_vector::RemappedSparseVector; use crate::common::types::{DimId, DimOffset, Weight}; use crate::index::compressed_posting_list::{ - CompressedPostingChunk, CompressedPostingListIterator, CompressedPostingListView, + CHUNK_SIZE, CompressedPostingChunk, CompressedPostingListIterator, CompressedPostingListView, }; use crate::index::inverted_index::InvertedIndex; use crate::index::inverted_index::inverted_index_ram::InvertedIndexRam; @@ -53,14 +55,14 @@ pub struct InvertedIndexFileHeader { /// Inverted flatten index from dimension id to posting list #[derive(Debug)] -pub struct InvertedIndexCompressedMmap { +pub struct InvertedIndexCompressedMmap { path: PathBuf, - mmap: Arc, + storage: S, pub file_header: InvertedIndexFileHeader, _phantom: PhantomData, } -#[derive(Debug, Default, Clone, FromBytes, Immutable, KnownLayout)] +#[derive(Debug, Default, Copy, Clone, FromBytes, Immutable, KnownLayout)] #[repr(C)] struct PostingListFileHeader { pub ids_start: u64, @@ -73,7 +75,24 @@ struct PostingListFileHeader { pub quantization_params: W::QuantizationParams, } -impl InvertedIndex for InvertedIndexCompressedMmap { +impl PostingListFileHeader { + fn postings_count(&self, remainders_end: u64) -> Option { + let chunks_bytes = self.chunks_count as usize * size_of::>(); + let data_len = remainders_end.checked_sub(self.ids_start)? as usize; + let remainders_bytes = data_len.checked_sub(self.ids_len as usize + chunks_bytes)?; + let elem_size = size_of::>(); + if !remainders_bytes.is_multiple_of(elem_size) { + return None; + } + let remainders_count = remainders_bytes / elem_size; + Some(self.chunks_count as usize * CHUNK_SIZE + remainders_count) + } +} + +impl InvertedIndex for InvertedIndexCompressedMmap +where + S::Fs: Default, +{ type Iter<'a> = CompressedPostingListIterator<'a, W>; type Version = Version; @@ -83,10 +102,10 @@ impl InvertedIndex for InvertedIndexCompressedMmap { } fn open(path: &Path) -> Result { - Self::load(path) + Self::load(&S::Fs::default(), path) } - fn save(&self, path: &Path) -> std::io::Result<()> { + fn save(&self, path: &Path) -> Result<()> { debug_assert_eq!(path, self.path); // If Self instance exists, it's either constructed by using `open()` (which reads index @@ -102,9 +121,10 @@ impl InvertedIndex for InvertedIndexCompressedMmap { fn get<'a>( &'a self, id: DimOffset, + arena: &'a SearchScratchArena, hw_counter: &'a HardwareCounterCell, ) -> Result> { - Ok(self.get(id, hw_counter)?.iter()) + Ok(self.get(id, arena, hw_counter)?.iter()) } fn len(&self) -> usize { @@ -112,7 +132,10 @@ impl InvertedIndex for InvertedIndexCompressedMmap { } fn posting_list_len(&self, id: DimOffset, hw_counter: &HardwareCounterCell) -> Result { - Ok(self.get(id, hw_counter)?.len()) + let (header, remainders_end) = self.read_posting_header(id, hw_counter)?; + header + .postings_count(remainders_end) + .ok_or_else(corrupted_index) } fn files(path: &Path) -> Vec { @@ -140,12 +163,9 @@ impl InvertedIndex for InvertedIndexCompressedMmap { panic!("Cannot upsert into a read-only Mmap inverted index") } - fn from_ram_index>( - ram_index: Cow, - path: P, - ) -> std::io::Result { + fn from_ram_index>(ram_index: Cow, path: P) -> Result { let index = InvertedIndexCompressedImmutableRam::from_ram_index(ram_index, &path)?; - Self::convert_and_save(&index, path) + Self::convert_and_save(&S::Fs::default(), &index, path) } fn vector_count(&self) -> usize { @@ -168,7 +188,7 @@ impl InvertedIndex for InvertedIndexCompressedMmap { } } -impl InvertedIndexCompressedMmap { +impl InvertedIndexCompressedMmap { const HEADER_SIZE: usize = size_of::>(); pub fn index_file_path(path: &Path) -> PathBuf { @@ -182,42 +202,31 @@ impl InvertedIndexCompressedMmap { pub fn get<'a>( &'a self, id: DimId, + arena: &'a SearchScratchArena, hw_counter: &'a HardwareCounterCell, ) -> Result> { - // check that the id is not out of bounds (posting_count includes the empty zeroth entry) - if id >= self.file_header.posting_count as DimId { - return Err(out_of_bounds(id, self.file_header.posting_count)); - } + let (header, remainders_end) = self.read_posting_header(id, hw_counter)?; - let header = self.read_header(id)?.clone(); - - hw_counter.vector_io_read().incr_delta(Self::HEADER_SIZE); - - let remainders_start = header.ids_start - + u64::from(header.ids_len) - + u64::from(header.chunks_count) * size_of::>() as u64; - - let remainders_end = if id + 1 < self.file_header.posting_count as DimId { - self.read_header(id + 1)?.ids_start - } else { - self.mmap.len() as u64 + let data_bytes = self.storage.read_bytes::( + header.ids_start..remainders_end, + align_of::<(CompressedPostingChunk, GenericPostingElement)>(), + )?; + let data = match data_bytes { + ACow::Borrowed(b) => b, + ACow::Owned(avec) => arena.alloc(avec), }; - if remainders_end - .checked_sub(remainders_start) - .is_some_and(|len| len % size_of::>() as u64 != 0) - { - return Err(corrupted_index()); - } - - let id_data = self.byte_slice(header.ids_start, u64::from(header.ids_len))?; - let chunks = <[CompressedPostingChunk]>::ref_from_bytes(self.byte_slice( - header.ids_start + u64::from(header.ids_len), - u64::from(header.chunks_count) * size_of::>() as u64, - )?) + let ids_len = header.ids_len as usize; + let chunks_bytes = header.chunks_count as usize * size_of::>(); + let id_data = data.get(..ids_len).ok_or_else(corrupted_index)?; + let chunks = <[CompressedPostingChunk]>::ref_from_bytes( + data.get(ids_len..ids_len + chunks_bytes) + .ok_or_else(corrupted_index)?, + ) .map_err(|_| corrupted_index())?; let remainders = <[GenericPostingElement]>::ref_from_bytes( - &self.mmap[remainders_start as usize..remainders_end as usize], + data.get(ids_len + chunks_bytes..) + .ok_or_else(corrupted_index)?, ) .map_err(|_| corrupted_index())?; @@ -231,25 +240,42 @@ impl InvertedIndexCompressedMmap { )) } - fn read_header(&self, id: DimId) -> Result<&PostingListFileHeader> { - let start = id as usize * Self::HEADER_SIZE; - let end = start + Self::HEADER_SIZE; - let bytes = self.mmap.get(start..end).ok_or_else(corrupted_index)?; - PostingListFileHeader::::ref_from_bytes(bytes).map_err(|_| corrupted_index()) - } + /// Read the header for `id`, plus the next header's `ids_start` + /// (which doubles as this entry's remainders end). + fn read_posting_header( + &self, + id: DimId, + hw_counter: &HardwareCounterCell, + ) -> Result<(PostingListFileHeader, u64)> { + if id >= self.file_header.posting_count as DimId { + return Err(out_of_bounds(id, self.file_header.posting_count)); + } - fn byte_slice(&self, start: u64, len: u64) -> Result<&[u8]> { - let start = start as usize; - let end = start - .checked_add(len as usize) - .ok_or_else(corrupted_index)?; - self.mmap.get(start..end).ok_or_else(corrupted_index) + let header_start = u64::from(id) * Self::HEADER_SIZE as u64; + let has_next = id + 1 < self.file_header.posting_count as DimId; + let read_size = Self::HEADER_SIZE + if has_next { size_of::() } else { 0 }; + let header_bytes = self.storage.read_bytes::( + header_start..header_start + read_size as u64, + align_of::>(), + )?; + let (&header, rest) = PostingListFileHeader::::ref_from_prefix(&header_bytes) + .map_err(|_| corrupted_index())?; + let remainders_end = if has_next { + *u64::ref_from_bytes(rest).map_err(|_| corrupted_index())? + } else { + self.storage.len::()? + }; + + hw_counter.vector_io_read().incr_delta(read_size); + + Ok((header, remainders_end)) } pub fn convert_and_save>( + fs: &S::Fs, index: &InvertedIndexCompressedImmutableRam, path: P, - ) -> std::io::Result { + ) -> Result { let total_posting_headers_size = index.postings.as_slice().len() * size_of::>(); @@ -264,6 +290,7 @@ impl InvertedIndexCompressedMmap { .map(|p| p.view(&hw_counter).store_size().total) .sum::(); let file_path = Self::index_file_path(path.as_ref()); + // TODO(uio): use UIO for writing as well. let file = create_and_ensure_length(file_path.as_ref(), file_length)?; let mut buf = BufWriter::new(file); @@ -312,34 +339,46 @@ impl InvertedIndexCompressedMmap { atomic_save_json(&Self::index_config_file_path(path.as_ref()), &file_header)?; + let storage = fs.open( + &file_path, + OpenOptions { + writeable: false, + need_sequential: false, + populate: Populate::No, + advice: AdviceSetting::Global, + }, + Default::default(), + )?; + Ok(Self { path: path.as_ref().to_owned(), - mmap: Arc::new(open_read_mmap( - file_path.as_ref(), - AdviceSetting::Global, - false, - )?), + storage, file_header, _phantom: PhantomData, }) } - pub fn load>(path: P) -> Result { + pub fn load>(fs: &S::Fs, path: P) -> Result { // read index config file let config_file_path = Self::index_config_file_path(path.as_ref()); // if the file header does not exist, the index is malformed let file_header: InvertedIndexFileHeader = read_json(&config_file_path)?; - // read index data into mmap + // open index data via universal IO let file_path = Self::index_file_path(path.as_ref()); - let mmap = open_read_mmap( - file_path.as_ref(), - AdviceSetting::from(Advice::Normal), - false, + let storage = fs.open( + &file_path, + OpenOptions { + writeable: false, + need_sequential: false, + populate: Populate::No, + advice: AdviceSetting::Advice(Advice::Normal), + }, + Default::default(), )?; let mut index = Self { path: path.as_ref().to_owned(), - mmap: Arc::new(mmap), + storage, file_header, _phantom: PhantomData, }; @@ -348,63 +387,56 @@ impl InvertedIndexCompressedMmap { if index.file_header.total_sparse_size.is_none() { index.file_header.total_sparse_size = - Some(index.calculate_total_sparse_size(&hw_counter)); + Some(index.calculate_total_sparse_size(&hw_counter)?); atomic_save_json(&config_file_path, &index.file_header)?; } Ok(index) } - fn calculate_total_sparse_size(&self, hw_counter: &HardwareCounterCell) -> usize { - (0..self.file_header.posting_count as DimId) - .filter_map(|id| { - self.get(id, hw_counter) - .map(|posting| posting.store_size().total) - .ok() - }) - .sum() + fn calculate_total_sparse_size(&self, hw_counter: &HardwareCounterCell) -> Result { + let mut total = 0; + let mut arena = SearchScratchArena::new_slow(); + for id in 0..self.file_header.posting_count as DimId { + total += self.get(id, &arena, hw_counter)?.store_size().total; + arena.gc(); + } + Ok(total) } - /// Populate all pages in the mmap. - /// Block until all pages are populated. - pub fn populate(&self) -> std::io::Result<()> { - self.mmap.populate(); - Ok(()) + /// Populate the underlying storage in RAM cache. Block until completed. + pub fn populate(&self) -> Result<()> { + self.storage.populate() } /// Drop disk cache. - pub fn clear_cache(&self) -> std::io::Result<()> { - let Self { - path: _, - mmap, - file_header: _, - _phantom, - } = self; - mmap.clear_cache(); - Ok(()) + pub fn clear_cache(&self) -> Result<()> { + self.storage.clear_ram_cache() } } #[cfg(test)] mod tests { + use common::universal_io::MmapFile; use tempfile::Builder; use super::*; use crate::common::types::QuantizedU8; use crate::index::inverted_index::inverted_index_ram_builder::InvertedIndexBuilder; - fn compare_indexes( + fn compare_indexes( inverted_index_ram: &InvertedIndexCompressedImmutableRam, - inverted_index_mmap: &InvertedIndexCompressedMmap, + inverted_index_mmap: &InvertedIndexCompressedMmap, ) { let hw_counter = HardwareCounterCell::new(); + let arena = SearchScratchArena::new_slow(); for id in 0..inverted_index_ram.postings.len() as DimId { let posting_list_ram = inverted_index_ram .postings .get(id as usize) .unwrap() .view(&hw_counter); - let posting_list_mmap = inverted_index_mmap.get(id, &hw_counter).unwrap(); + let posting_list_mmap = inverted_index_mmap.get(id, &arena, &hw_counter).unwrap(); let mmap_parts = posting_list_mmap.parts(); let ram_parts = posting_list_ram.parts(); @@ -415,13 +447,25 @@ mod tests { #[test] fn test_inverted_index_mmap() { - check_inverted_index_mmap::(); - check_inverted_index_mmap::(); - check_inverted_index_mmap::(); - check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + + #[cfg(target_os = "linux")] + { + use common::universal_io::IoUringFile; + check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + check_inverted_index_mmap::(); + } } - fn check_inverted_index_mmap() { + fn check_inverted_index_mmap() + where + S::Fs: Default, + { let hw_counter = HardwareCounterCell::new(); // skip 4th dimension @@ -446,7 +490,8 @@ mod tests { let tmp_dir_path = Builder::new().prefix("test_index_dir2").tempdir().unwrap(); { - let inverted_index_mmap = InvertedIndexCompressedMmap::::convert_and_save( + let inverted_index_mmap = InvertedIndexCompressedMmap::::convert_and_save( + &S::Fs::default(), &inverted_index_ram, &tmp_dir_path, ) @@ -454,22 +499,25 @@ mod tests { compare_indexes(&inverted_index_ram, &inverted_index_mmap); } - let inverted_index_mmap = InvertedIndexCompressedMmap::::load(&tmp_dir_path).unwrap(); + let index = + InvertedIndexCompressedMmap::::load(&Default::default(), &tmp_dir_path) + .unwrap(); // posting_count: 0th entry is always empty + 1st + 2nd + 3rd + 4th empty + 5th - assert_eq!(inverted_index_mmap.file_header.posting_count, 6); - assert_eq!(inverted_index_mmap.file_header.vector_count, 9); + assert_eq!(index.file_header.posting_count, 6); + assert_eq!(index.file_header.vector_count, 9); - compare_indexes(&inverted_index_ram, &inverted_index_mmap); + compare_indexes(&inverted_index_ram, &index); - assert!(inverted_index_mmap.get(0, &hw_counter).unwrap().is_empty()); // the first entry is always empty as dimension ids start at 1 - assert_eq!(inverted_index_mmap.get(1, &hw_counter).unwrap().len(), 9); - assert_eq!(inverted_index_mmap.get(2, &hw_counter).unwrap().len(), 4); - assert_eq!(inverted_index_mmap.get(3, &hw_counter).unwrap().len(), 3); - assert!(inverted_index_mmap.get(4, &hw_counter).unwrap().is_empty()); // return empty posting list info for intermediary empty ids - assert_eq!(inverted_index_mmap.get(5, &hw_counter).unwrap().len(), 2); - // index after the last values are None - assert!(inverted_index_mmap.get(6, &hw_counter).is_err()); - assert!(inverted_index_mmap.get(7, &hw_counter).is_err()); - assert!(inverted_index_mmap.get(100, &hw_counter).is_err()); + let arena = SearchScratchArena::new_slow(); + assert!(index.get(0, &arena, &hw_counter).unwrap().is_empty()); // the first entry is always empty as dimension ids start at 1 + assert_eq!(index.get(1, &arena, &hw_counter).unwrap().len(), 9); + assert_eq!(index.get(2, &arena, &hw_counter).unwrap().len(), 4); + assert_eq!(index.get(3, &arena, &hw_counter).unwrap().len(), 3); + assert!(index.get(4, &arena, &hw_counter).unwrap().is_empty()); // return empty posting list info for intermediary empty ids + assert_eq!(index.get(5, &arena, &hw_counter).unwrap().len(), 2); + // index after the last values are errors (out of bounds) + assert!(index.get(6, &arena, &hw_counter).is_err()); + assert!(index.get(7, &arena, &hw_counter).is_err()); + assert!(index.get(100, &arena, &hw_counter).is_err()); } } diff --git a/lib/sparse/src/index/inverted_index/inverted_index_ram.rs b/lib/sparse/src/index/inverted_index/inverted_index_ram.rs index 0b5e927e66..838b0d0851 100644 --- a/lib/sparse/src/index/inverted_index/inverted_index_ram.rs +++ b/lib/sparse/src/index/inverted_index/inverted_index_ram.rs @@ -52,13 +52,14 @@ impl InvertedIndex for InvertedIndexRam { panic!("InvertedIndexRam is not supposed to be loaded"); } - fn save(&self, _path: &Path) -> std::io::Result<()> { + fn save(&self, _path: &Path) -> Result<()> { panic!("InvertedIndexRam is not supposed to be saved"); } fn get<'a>( &'a self, id: DimOffset, + _arena: &'a crate::SearchScratchArena, _hw_counter: &'a HardwareCounterCell, ) -> Result> { Ok(self.get(id)?.iter()) @@ -104,10 +105,7 @@ impl InvertedIndex for InvertedIndexRam { self.upsert(id, vector, old_vector); } - fn from_ram_index>( - ram_index: Cow, - _path: P, - ) -> std::io::Result { + fn from_ram_index>(ram_index: Cow, _path: P) -> Result { Ok(ram_index.into_owned()) } diff --git a/lib/sparse/src/index/inverted_index/mod.rs b/lib/sparse/src/index/inverted_index/mod.rs index befb81f66c..465d8f291f 100644 --- a/lib/sparse/src/index/inverted_index/mod.rs +++ b/lib/sparse/src/index/inverted_index/mod.rs @@ -8,6 +8,7 @@ use common::types::PointOffsetType; use common::universal_io::{Result, UniversalIoError}; use super::posting_list_common::PostingListIter; +use crate::SearchScratchArena; use crate::common::sparse_vector::RemappedSparseVector; use crate::common::types::DimOffset; use crate::index::inverted_index::inverted_index_ram::InvertedIndexRam; @@ -36,12 +37,13 @@ pub trait InvertedIndex: Sized + Debug + 'static { fn open(path: &Path) -> Result; /// Save index - fn save(&self, path: &Path) -> std::io::Result<()>; + fn save(&self, path: &Path) -> Result<()>; /// Get posting list for dimension id fn get<'a>( &'a self, id: DimOffset, + arena: &'a SearchScratchArena, hw_counter: &'a HardwareCounterCell, ) -> Result>; @@ -72,10 +74,7 @@ pub trait InvertedIndex: Sized + Debug + 'static { ); /// Create inverted index from ram index - fn from_ram_index>( - ram_index: Cow, - path: P, - ) -> std::io::Result; + fn from_ram_index>(ram_index: Cow, path: P) -> Result; /// Number of indexed vectors fn vector_count(&self) -> usize; diff --git a/lib/sparse/src/index/search_context.rs b/lib/sparse/src/index/search_context.rs index ea6fbd5c56..2f81caa2ab 100644 --- a/lib/sparse/src/index/search_context.rs +++ b/lib/sparse/src/index/search_context.rs @@ -4,11 +4,11 @@ use std::sync::atomic::Ordering::Relaxed; use common::counter::hardware_counter::HardwareCounterCell; use common::top_k::TopK; -use common::types::{PointOffsetType, ScoredPointOffset}; +use common::types::{PointOffsetType, ScoreType, ScoredPointOffset}; use common::universal_io::Result; use super::posting_list_common::PostingListIter; -use crate::common::scores_memory_pool::PooledScoresHandle; +use crate::SearchScratch; use crate::common::sparse_vector::{RemappedSparseVector, score_vectors}; use crate::common::types::{DimId, DimWeight}; use crate::index::inverted_index::InvertedIndex; @@ -24,7 +24,7 @@ pub struct IndexedPostingListIterator { /// Making this larger makes the search faster but uses more (pooled) memory const ADVANCE_BATCH_SIZE: usize = 10_000; -pub struct SearchContext<'a, 'b, T: PostingListIter = PostingListIterator<'a>> { +pub struct SearchContext<'a, T: PostingListIter = PostingListIterator<'a>> { postings_iterators: Vec>, query: RemappedSparseVector, top: usize, @@ -32,27 +32,28 @@ pub struct SearchContext<'a, 'b, T: PostingListIter = PostingListIterator<'a>> { top_results: TopK, min_record_id: Option, // min_record_id ids across all posting lists max_record_id: PointOffsetType, // max_record_id ids across all posting lists - pooled: PooledScoresHandle<'b>, // handle to pooled scores + /// Scores buffer from [`SearchScratch`]. + scores: &'a mut Vec, use_pruning: bool, hardware_counter: &'a HardwareCounterCell, } -impl<'a, 'b, T: PostingListIter> SearchContext<'a, 'b, T> { +impl<'a, T: PostingListIter> SearchContext<'a, T> { pub fn new( query: RemappedSparseVector, top: usize, inverted_index: &'a impl InvertedIndex = T>, - pooled: PooledScoresHandle<'b>, + scratch: &'a mut SearchScratch<'_>, is_stopped: &'a AtomicBool, hardware_counter: &'a HardwareCounterCell, - ) -> Result> { + ) -> Result> { let mut postings_iterators = Vec::new(); // track min and max record ids across all posting lists let mut max_record_id = 0; let mut min_record_id = u32::MAX; // iterate over query indices for (query_weight_offset, id) in query.indices.iter().enumerate() { - let mut it = inverted_index.get(*id, hardware_counter)?; + let mut it = inverted_index.get(*id, &scratch.arena, hardware_counter)?; if let (Some(first), Some(last_id)) = (it.peek(), it.last_id()) { // check if new min let min_record_id_posting = first.record_id; @@ -87,7 +88,7 @@ impl<'a, 'b, T: PostingListIter> SearchContext<'a, 'b, T> { top_results, min_record_id, max_record_id, - pooled, + scores: &mut scratch.scores, use_pruning, hardware_counter, }) @@ -158,13 +159,13 @@ impl<'a, 'b, T: PostingListIter> SearchContext<'a, 'b, T> { ) { // init batch scores let batch_len = batch_last_id - batch_start_id + 1; - self.pooled.scores.clear(); // keep underlying allocated memory - self.pooled.scores.resize(batch_len as usize, 0.0); + self.scores.clear(); // keep underlying allocated memory + self.scores.resize(batch_len as usize, 0.0); for posting in self.postings_iterators.iter_mut() { posting.posting_list_iterator.for_each_till_id( batch_last_id, - self.pooled.scores.as_mut_slice(), + self.scores.as_mut_slice(), #[inline(always)] |scores, id, weight| { let element_score = weight * posting.query_weight; @@ -176,7 +177,7 @@ impl<'a, 'b, T: PostingListIter> SearchContext<'a, 'b, T> { ); } - for (local_index, &score) in self.pooled.scores.iter().enumerate() { + for (local_index, &score) in self.scores.iter().enumerate() { // publish only the non-zero scores above the current min to beat if score != 0.0 && score > self.top_results.threshold() { let real_id = batch_start_id + local_index as PointOffsetType; diff --git a/lib/sparse/src/index/tests/common.rs b/lib/sparse/src/index/tests/common.rs index 5f9f3ddf9f..fb8f549b7a 100644 --- a/lib/sparse/src/index/tests/common.rs +++ b/lib/sparse/src/index/tests/common.rs @@ -1,11 +1,10 @@ use std::borrow::Cow; -use std::sync::OnceLock; use common::types::PointOffsetType; +use common::universal_io::MmapFile; use rand::{RngExt, SeedableRng}; use tempfile::TempDir; -use crate::common::scores_memory_pool::{PooledScoresHandle, ScoresMemoryPool}; use crate::common::sparse_vector::RemappedSparseVector; use crate::common::types::{DimOffset, Weight}; use crate::index::inverted_index::InvertedIndex; @@ -13,14 +12,6 @@ use crate::index::inverted_index::inverted_index_compressed_mmap::InvertedIndexC use crate::index::inverted_index::inverted_index_ram::InvertedIndexRam; use crate::index::inverted_index::inverted_index_ram_builder::InvertedIndexBuilder; -static TEST_SCORES_POOL: OnceLock = OnceLock::new(); - -pub fn get_pooled_scores() -> PooledScoresHandle<'static> { - TEST_SCORES_POOL - .get_or_init(ScoresMemoryPool::default) - .get() -} - pub struct TestIndex { pub index: I, _temp_dir: TempDir, @@ -78,7 +69,7 @@ pub fn generate_sparse_index( density: usize, vocab1: usize, vocab2: usize, -) -> TestIndex> +) -> TestIndex> where W: Weight + 'static, R: rand::Rng, @@ -98,7 +89,7 @@ pub fn build_index( density: usize, vocab1: usize, vocab2: usize, -) -> TestIndex> +) -> TestIndex> where W: Weight + 'static, { diff --git a/lib/sparse/src/index/tests/hw_counter_test.rs b/lib/sparse/src/index/tests/hw_counter_test.rs index 51dfc28b10..ccafb030bb 100644 --- a/lib/sparse/src/index/tests/hw_counter_test.rs +++ b/lib/sparse/src/index/tests/hw_counter_test.rs @@ -4,21 +4,23 @@ use common::counter::hardware_accumulator::HwMeasurementAcc; use common::types::PointOffsetType; use itertools::Itertools; +use crate::SearchScratch; use crate::common::sparse_vector::RemappedSparseVector; use crate::index::inverted_index::InvertedIndex; use crate::index::search_context::SearchContext; -use crate::index::tests::common::{build_index, get_pooled_scores, match_all}; +use crate::index::tests::common::{build_index, match_all}; fn do_search(index: &I, query: RemappedSparseVector) -> HwMeasurementAcc { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); let top = 10; + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( query, top, index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -41,11 +43,12 @@ fn do_plain_search( let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); let top = 10; + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( query, top, index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) diff --git a/lib/sparse/src/index/tests/indexed_vs_plain_test.rs b/lib/sparse/src/index/tests/indexed_vs_plain_test.rs index c459f18952..049cf309e7 100644 --- a/lib/sparse/src/index/tests/indexed_vs_plain_test.rs +++ b/lib/sparse/src/index/tests/indexed_vs_plain_test.rs @@ -2,21 +2,23 @@ use std::sync::atomic::AtomicBool; use common::counter::hardware_accumulator::HwMeasurementAcc; +use crate::SearchScratch; use crate::common::sparse_vector::RemappedSparseVector; use crate::index::inverted_index::InvertedIndex; use crate::index::search_context::SearchContext; -use crate::index::tests::common::{build_index, get_pooled_scores, match_all}; +use crate::index::tests::common::{build_index, match_all}; fn query(index: &I, query: RemappedSparseVector) { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); let top = 10; + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( query.clone(), top, index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -24,12 +26,13 @@ fn query(index: &I, query: RemappedSparseVector) { let result = search_context.search(&match_all); let docs: Vec<_> = result.iter().map(|x| x.idx).collect(); + drop(search_context); let mut search_context = SearchContext::new( query, top, index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) diff --git a/lib/sparse/src/index/tests/search_context_tests.rs b/lib/sparse/src/index/tests/search_context_tests.rs index 7ac34165b0..8c53896130 100644 --- a/lib/sparse/src/index/tests/search_context_tests.rs +++ b/lib/sparse/src/index/tests/search_context_tests.rs @@ -3,15 +3,17 @@ mod tests { use std::any::TypeId; use std::borrow::Cow; - use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use common::counter::hardware_accumulator::HwMeasurementAcc; use common::counter::hardware_counter::HardwareCounterCell; use common::types::{PointOffsetType, ScoredPointOffset}; + #[cfg(target_os = "linux")] + use common::universal_io::IoUringFile; + use common::universal_io::MmapFile; use tempfile::TempDir; - use crate::common::scores_memory_pool::{PooledScoresHandle, ScoresMemoryPool}; + use crate::SearchScratch; use crate::common::sparse_vector::RemappedSparseVector; use crate::common::types::QuantizedU8; use crate::index::inverted_index::InvertedIndex; @@ -38,28 +40,36 @@ mod tests { #[instantiate_tests(>)] mod iram_q8 {} - #[instantiate_tests(>)] + #[instantiate_tests(>)] mod mmap_f32 {} - #[instantiate_tests(>)] + #[instantiate_tests(>)] mod mmap_f16 {} - #[instantiate_tests(>)] + #[instantiate_tests(>)] mod mmap_u8 {} - #[instantiate_tests(>)] + #[instantiate_tests(>)] mod mmap_q8 {} + #[cfg(target_os = "linux")] + #[instantiate_tests(>)] + mod uring_f32 {} + + #[cfg(target_os = "linux")] + #[instantiate_tests(>)] + mod uring_f16 {} + + #[cfg(target_os = "linux")] + #[instantiate_tests(>)] + mod uring_u8 {} + + #[cfg(target_os = "linux")] + #[instantiate_tests(>)] + mod uring_q8 {} + // --- End of test instantiations --- - static TEST_SCORES_POOL: OnceLock = OnceLock::new(); - - fn get_pooled_scores() -> PooledScoresHandle<'static> { - TEST_SCORES_POOL - .get_or_init(ScoresMemoryPool::default) - .get() - } - /// Match all filter condition for testing fn match_all(_p: PointOffsetType) -> bool { true @@ -88,7 +98,9 @@ mod tests { fn round_scores(mut scores: Vec) -> Vec { let errors_allowed_for = [ TypeId::of::>(), - TypeId::of::>(), + TypeId::of::>(), + #[cfg(target_os = "linux")] + TypeId::of::>(), ]; if errors_allowed_for.contains(&TypeId::of::()) { let precision = 0.25; @@ -108,11 +120,12 @@ mod tests { let hw_counter = HardwareCounterCell::disposable(); let is_stopped = AtomicBool::new(false); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector::default(), // empty query vector 10, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hw_counter, ) @@ -133,6 +146,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -140,7 +154,7 @@ mod tests { }, 10, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -192,6 +206,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -199,7 +214,7 @@ mod tests { }, 10, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -235,6 +250,7 @@ mod tests { None, ); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -242,7 +258,7 @@ mod tests { }, 10, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -290,6 +306,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -297,7 +314,7 @@ mod tests { }, 3, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -332,6 +349,7 @@ mod tests { let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -339,7 +357,7 @@ mod tests { }, 4, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -385,6 +403,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1], @@ -392,7 +411,7 @@ mod tests { }, 1, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -420,6 +439,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -427,7 +447,7 @@ mod tests { }, 1, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -460,6 +480,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -467,7 +488,7 @@ mod tests { }, 1, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -498,6 +519,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -505,7 +527,7 @@ mod tests { }, 3, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -531,6 +553,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 2, 3], @@ -538,7 +561,7 @@ mod tests { }, 3, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) @@ -586,6 +609,7 @@ mod tests { let is_stopped = AtomicBool::new(false); let accumulator = HwMeasurementAcc::new(); let hardware_counter = accumulator.get_counter_cell(); + let mut scratch = SearchScratch::new_for_test(); let mut search_context = SearchContext::new( RemappedSparseVector { indices: vec![1, 3], @@ -593,7 +617,7 @@ mod tests { }, 3, &index.index, - get_pooled_scores(), + &mut scratch, &is_stopped, &hardware_counter, ) diff --git a/lib/sparse/src/lib.rs b/lib/sparse/src/lib.rs index dba2624271..2e6db5dda2 100644 --- a/lib/sparse/src/lib.rs +++ b/lib/sparse/src/lib.rs @@ -1,2 +1,5 @@ pub mod common; pub mod index; +pub(crate) mod search_scratch; + +pub use search_scratch::{SearchScratch, SearchScratchArena, SearchScratchPool}; diff --git a/lib/sparse/src/search_scratch.rs b/lib/sparse/src/search_scratch.rs new file mode 100644 index 0000000000..7105fc36f6 --- /dev/null +++ b/lib/sparse/src/search_scratch.rs @@ -0,0 +1,102 @@ +use common::defaults::POOL_KEEP_LIMIT; +use common::ext::aligned_vec::{AVec, RuntimeAlign}; +use common::types::ScoreType; +use parking_lot::Mutex; +use typed_arena::Arena; + +#[derive(Debug)] +pub struct SearchScratchPool { + pool: Mutex>, +} + +impl SearchScratchPool { + #[expect(clippy::new_without_default)] + pub fn new() -> SearchScratchPool { + SearchScratchPool { + pool: Mutex::new(Vec::with_capacity(*POOL_KEEP_LIMIT)), + } + } + + /// Take a single [`SearchScratch`] from the pool. + pub fn get(&self) -> SearchScratch<'_> { + let scores = self.pool.lock().pop().unwrap_or_default(); + SearchScratch { + pool: self, + scores, + arena: SearchScratchArena::new_slow(), + } + } +} + +/// Scratch space for use by [`crate::index::search_context::SearchContext`]. +pub struct SearchScratch<'a> { + /// The pool to return this scratch on drop. + pool: &'a SearchScratchPool, + /// Used for batched scoring. + pub(crate) scores: SearchScratchScores, + /// Used to own/store posting list bytes while reading them from the file. + pub(crate) arena: SearchScratchArena, +} + +impl SearchScratch<'_> { + #[cfg(test)] + pub fn new_for_test() -> SearchScratch<'static> { + use std::sync::OnceLock; + static POOL: OnceLock = OnceLock::new(); + POOL.get_or_init(SearchScratchPool::new).get() + } +} + +type SearchScratchScores = Vec; + +impl Drop for SearchScratch<'_> { + fn drop(&mut self) { + let SearchScratch { + pool: SearchScratchPool { pool }, + scores, + arena: _, + } = self; + let mut pool = pool.lock(); + if pool.len() < *POOL_KEEP_LIMIT { + pool.push(std::mem::take(scores)); + } + } +} + +/// Hacky wrapper around [`Arena`] with [`Self::gc`] method. +/// +/// TODO: get rid of this struct once [`Arena`] gets `clear()` method, and use +/// type alias instead. +/// https://github.com/thomcc/rust-typed-arena/issues/64 +// pub type SearchScratchArena = Arena>; +pub struct SearchScratchArena(Arena>); + +impl SearchScratchArena { + /// Same as [`Arena::new`]. + /// Renamed to `new_slow` to point out that [`Arena::new`] allocates. + pub fn new_slow() -> SearchScratchArena { + SearchScratchArena(Arena::new()) + } + + /// Free memory if the arena got too big. + /// Poor's man version of `Arena::clear()`. + pub fn gc(&mut self) { + // Too small => frequent reallocations + // Too big => memory bloat + const ARBITRARY_LIMIT: usize = 256; + if self.0.len() > ARBITRARY_LIMIT { + self.0 = Arena::new(); + } + + for vec in self.0.iter_mut() { + *vec = AVec::new(1); + } + } +} + +impl std::ops::Deref for SearchScratchArena { + type Target = Arena>; + fn deref(&self) -> &Arena> { + &self.0 + } +}