diff --git a/Cargo.lock b/Cargo.lock index 662188f659..a7908d09d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1961,6 +1961,8 @@ dependencies = [ "fs-err", "indicatif", "reqwest 0.13.3", + "serde", + "serde_json", ] [[package]] @@ -7564,6 +7566,7 @@ dependencies = [ "schemars 0.8.22", "serde", "serde_json", + "sha2 0.11.0", "sparse", "tempfile", "validator", diff --git a/lib/common/dataset/Cargo.toml b/lib/common/dataset/Cargo.toml index b0968e8141..1822d3bc1e 100644 --- a/lib/common/dataset/Cargo.toml +++ b/lib/common/dataset/Cargo.toml @@ -18,6 +18,8 @@ flate2 = { version = "1.1.9" } fs-err = { workspace = true } indicatif = { workspace = true } reqwest = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } [dev-dependencies] -fs-err = { workspace = true, features = ["debug"] } +fs-err = { workspace = true, features = ["debug", "debug_tokio", "tokio"] } diff --git a/lib/common/dataset/src/lib.rs b/lib/common/dataset/src/lib.rs index 4128f8fe58..bb76db0767 100644 --- a/lib/common/dataset/src/lib.rs +++ b/lib/common/dataset/src/lib.rs @@ -1,10 +1,14 @@ +use std::env::var_os; use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::OnceLock; use anyhow::{Context, Result, anyhow}; use flate2::read::GzDecoder; use fs_err as fs; use fs_err::File; use indicatif::{ProgressBar, ProgressDrawTarget}; +use serde::Deserialize; pub enum Dataset { // https://github.com/qdrant/sparse-vectors-experiments @@ -56,9 +60,7 @@ fn download_cached(url: &str) -> Result { }; // Cache directory, e.g. "target/datasets". - let cache_dir = workspace_dir() - .join(std::env::var_os("CARGO_TARGET_DIR").unwrap_or_else(|| "target".into())) - .join("datasets"); + let cache_dir = cargo_target_dir()?.join("datasets"); // Cache file path, e.g. "target/datasets/base_full.csr". let cache_path = cache_dir.join(basename); @@ -102,14 +104,27 @@ fn download_cached(url: &str) -> Result { Ok(cache_path) } -fn workspace_dir() -> PathBuf { - let output = std::process::Command::new(env!("CARGO")) - .arg("locate-project") - .arg("--workspace") - .arg("--message-format=plain") +fn cargo_target_dir() -> Result<&'static PathBuf> { + static CACHED: OnceLock = OnceLock::new(); + + if let Some(cached) = CACHED.get() { + return Ok(cached); + } + + #[derive(Deserialize)] + struct Metadata { + target_directory: PathBuf, + } + + let result = Command::new(var_os("CARGO").context("$CARGO not set")?) + .args(["metadata", "--no-deps", "--format-version=1"]) .output() - .unwrap() - .stdout; - let cargo_path = Path::new(std::str::from_utf8(&output).unwrap().trim()); - cargo_path.parent().unwrap().to_path_buf() + .context("Failed to run `cargo metadata`")?; + if !result.status.success() { + let stderr = String::from_utf8_lossy(&result.stderr); + anyhow::bail!("`cargo metadata` failed:\n{stderr}"); + } + let metadata: Metadata = serde_json::from_slice(&result.stdout) + .context("Failed to parse `cargo metadata` output")?; + Ok(CACHED.get_or_init(|| metadata.target_directory)) } diff --git a/lib/sparse/Cargo.toml b/lib/sparse/Cargo.toml index ea0dd069b4..3bb6da060c 100644 --- a/lib/sparse/Cargo.toml +++ b/lib/sparse/Cargo.toml @@ -40,7 +40,9 @@ dataset = { path = "../common/dataset" } fs-err = { workspace = true, features = ["debug"] } generic-tests = { workspace = true } indicatif = { workspace = true } +sha2 = { workspace = true } sparse = { path = ".", features = ["testing"] } +zerocopy = { workspace = true } [target.'cfg(not(target_os = "windows"))'.dev-dependencies] pprof = { workspace = true } diff --git a/lib/sparse/benches/search.rs b/lib/sparse/benches/search.rs index 791acfd816..f3227f74cc 100644 --- a/lib/sparse/benches/search.rs +++ b/lib/sparse/benches/search.rs @@ -1,6 +1,5 @@ use std::borrow::Cow; -use std::io; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; use common::counter::hardware_counter::HardwareCounterCell; @@ -8,21 +7,26 @@ use common::types::PointOffsetType; use criterion::measurement::Measurement; use criterion::{Criterion, criterion_group, criterion_main}; use dataset::Dataset; +use fs_err as fs; use indicatif::{ProgressBar, ProgressDrawTarget}; use itertools::Itertools; use rand::SeedableRng as _; use rand::rngs::StdRng; +use sha2::Digest; use sparse::common::scores_memory_pool::ScoresMemoryPool; 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; +use sparse::common::types::{QuantizedU8, Weight}; use sparse::index::inverted_index::InvertedIndex; use sparse::index::inverted_index::inverted_index_compressed_immutable_ram::InvertedIndexCompressedImmutableRam; use sparse::index::inverted_index::inverted_index_compressed_mmap::InvertedIndexCompressedMmap; use sparse::index::inverted_index::inverted_index_ram::InvertedIndexRam; use sparse::index::inverted_index::inverted_index_ram_builder::InvertedIndexBuilder; use sparse::index::loaders::{self, Csr}; +use sparse::index::posting_list::PostingList; +use sparse::index::posting_list_common::PostingElementEx; use sparse::index::search_context::SearchContext; +use zerocopy::IntoBytes; mod prof; const NUM_QUERIES: usize = 2048; @@ -40,12 +44,17 @@ pub fn bench_search(c: &mut Criterion) { .unwrap() .collect::>(); - let index_1m = load_csr_index(Dataset::NeurIps2023_1M.download().unwrap(), 1.0).unwrap(); - run_bench(c, "neurips2023-1M", index_1m, &query_vectors); + let name = "neurips2023-1M"; + let index_1m = cached_ram_index(name, || { + load_csr_index(Dataset::NeurIps2023_1M.download().unwrap(), 1.0) + }); + run_bench(c, name, index_1m, &query_vectors); - let index_full = - load_csr_index(Dataset::NeurIps2023Full.download().unwrap(), 0.25).unwrap(); - run_bench(c, "neurips2023-full-25pct", index_full, &query_vectors); + let name = "neurips2023-full-25pct"; + let index_full = cached_ram_index(name, || { + load_csr_index(Dataset::NeurIps2023Full.download().unwrap(), 0.25) + }); + run_bench(c, name, index_full, &query_vectors); } bench_movies(c); @@ -54,12 +63,13 @@ pub fn bench_search(c: &mut Criterion) { fn bench_uniform_random(c: &mut Criterion, name: &str, num_vectors: usize) { let mut rnd = StdRng::seed_from_u64(42); - let index = InvertedIndexBuilder::build_from_iterator((0..num_vectors).map(|idx| { - ( - idx as PointOffsetType, - random_sparse_vector(&mut rnd, MAX_SPARSE_DIM).into_remapped(), - ) - })); + let index = cached_ram_index(name, || { + let mut rnd = rnd.fork(); + InvertedIndexBuilder::build_from_iterator((0..num_vectors).map(|idx| { + let vec = random_sparse_vector(&mut rnd, MAX_SPARSE_DIM).into_remapped(); + (idx as PointOffsetType, vec) + })) + }); let query_vectors = (0..NUM_QUERIES) .map(|_| random_positive_sparse_vector(&mut rnd, MAX_SPARSE_DIM)) @@ -141,28 +151,18 @@ pub fn run_bench( macro_rules! run_bench2 { ($name:literal, $type:ty) => { + let index_path = cached_compressed_index::<$type>(&index, name); + run_bench2( c.benchmark_group(format!("search/ram_{}/{name}", $name)), - &InvertedIndexCompressedImmutableRam::<$type>::from_ram_index( - Cow::Borrowed(&index), - "nonexistent/path", - ) - .unwrap(), + &InvertedIndexCompressedImmutableRam::<$type>::open(&index_path).unwrap(), query_vectors, &hottest_query_vectors, ); run_bench2( c.benchmark_group(format!("search/mmap_{}/{name}", $name)), - &InvertedIndexCompressedMmap::<$type>::from_ram_index( - Cow::Borrowed(&index), - tempfile::Builder::new() - .prefix("test_index_dir") - .tempdir() - .unwrap() - .path(), - ) - .unwrap(), + &InvertedIndexCompressedMmap::<$type>::open(&index_path).unwrap(), query_vectors, &hottest_query_vectors, ); @@ -216,18 +216,104 @@ fn run_bench2( }); } -fn load_csr_index(path: impl AsRef, ratio: f32) -> io::Result { - let csr = Csr::open(path.as_ref())?; +fn load_csr_index(path: impl AsRef, ratio: f32) -> InvertedIndexRam { + let csr = Csr::open(path.as_ref()).unwrap(); let mut builder = InvertedIndexBuilder::new(); assert!(ratio > 0.0 && ratio <= 1.0); let count = (csr.len() as f32 * ratio) as usize; let bar = ProgressBar::with_draw_target(Some(count as u64), ProgressDrawTarget::stderr_with_hz(12)); - for (row, vec) in bar.wrap_iter(csr.iter()?.take(count).enumerate()) { + for (row, vec) in bar.wrap_iter(csr.iter().unwrap().take(count).enumerate()) { builder.add(row as u32, vec.into_remapped()); } bar.finish_and_clear(); - Ok(builder.build()) + builder.build() +} + +fn cache_dir() -> PathBuf { + Path::new(env!("CARGO_TARGET_TMPDIR")) + .join(env!("CARGO_PKG_NAME")) + .join(env!("CARGO_CRATE_NAME")) +} + +/// Load an [`InvertedIndexRam`] from the cache. +/// If not exists, calls `build()` to create it. +fn cached_ram_index(name: &str, build: impl FnOnce() -> InvertedIndexRam) -> InvertedIndexRam { + let path = cache_dir().join(name); + if !path.exists() { + eprintln!("Building cache: {path:?}"); + let index = build(); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + index.test_save(&path.with_extension("tmp")).unwrap(); + fs::rename(path.with_extension("tmp"), &path).unwrap(); + } + eprintln!("Using cache: {path:?}"); + InvertedIndexRam::test_load(&path).unwrap() +} + +/// Loads [`InvertedIndexCompressedMmap`] from the cache. +/// If not exists, converts it from the given [`InvertedIndexRam`]. +fn cached_compressed_index(index: &InvertedIndexRam, name: &str) -> PathBuf { + let path = cache_dir().join(format!( + "{name}-{}-{hash}", + W::NAME, + hash = inverted_index_partial_hash(index) + )); + + if !path.exists() { + eprintln!("Building cache: {path:?}"); + let tmp_path = path.with_extension("tmp"); + if tmp_path.exists() { + fs::remove_dir_all(&tmp_path).unwrap(); + } + fs::create_dir_all(&tmp_path).unwrap(); + InvertedIndexCompressedMmap::::from_ram_index(Cow::Borrowed(index), &tmp_path).unwrap(); + fs::rename(tmp_path, &path).unwrap(); + } + eprintln!("Using cache: {path:?}"); + path +} + +/// Compute hash of the given [`InvertedIndexRam`]. +/// For performance reasons, use only a subset of the data. +fn inverted_index_partial_hash(index: &InvertedIndexRam) -> String { + let mut hasher = sha2::Sha256::new(); + let InvertedIndexRam { + postings, + vector_count, + total_sparse_size, + } = index; + hasher.update(vector_count.as_bytes()); + hasher.update(total_sparse_size.as_bytes()); + + let mut hash_posting = |posting: &PostingList| { + let PostingList { elements } = posting; + for element in elements.iter() { + let PostingElementEx { + record_id, + weight, + max_next_weight, + } = element; + hasher.update(record_id.as_bytes()); + hasher.update(weight.as_bytes()); + hasher.update(max_next_weight.as_bytes()); + } + }; + + // For performance reasons, hash only first and last 100 postings. + let mut postings = postings.iter(); + for posting in postings.by_ref().take(100) { + hash_posting(posting); + } + for posting in postings.rev().take(100) { + hash_posting(posting); + } + + hasher + .finalize() + .iter() + .map(|b| format!("{b:02x}")) + .collect() } #[cfg(not(target_os = "windows"))] diff --git a/lib/sparse/src/common/types.rs b/lib/sparse/src/common/types.rs index 3c3379808c..eba433cae6 100644 --- a/lib/sparse/src/common/types.rs +++ b/lib/sparse/src/common/types.rs @@ -12,6 +12,9 @@ pub type DimWeight = f32; pub trait Weight: PartialEq + Copy + Debug + FromBytes + Immutable + KnownLayout + 'static { type QuantizationParams: Copy + PartialEq + Debug + FromBytes + Immutable + KnownLayout; + #[cfg(feature = "testing")] + const NAME: &'static str; + fn quantization_params_for( values: impl ExactSizeIterator + Clone, ) -> Self::QuantizationParams; @@ -30,6 +33,9 @@ pub trait Weight: PartialEq + Copy + Debug + FromBytes + Immutable + KnownLayout impl Weight for f32 { type QuantizationParams = (); + #[cfg(feature = "testing")] + const NAME: &'static str = "f32"; + #[inline] fn quantization_params_for(_values: impl ExactSizeIterator + Clone) {} @@ -53,6 +59,9 @@ impl Weight for f32 { impl Weight for half::f16 { type QuantizationParams = (); + #[cfg(feature = "testing")] + const NAME: &'static str = "f16"; + #[inline] fn quantization_params_for(_values: impl ExactSizeIterator + Clone) {} @@ -76,6 +85,8 @@ impl Weight for half::f16 { impl Weight for u8 { type QuantizationParams = (); + const NAME: &'static str = "u8"; + #[inline] fn quantization_params_for(_values: impl ExactSizeIterator + Clone) {} @@ -120,6 +131,9 @@ pub struct QuantizedU8Params { impl Weight for QuantizedU8 { type QuantizationParams = QuantizedU8Params; + #[cfg(feature = "testing")] + const NAME: &'static str = "u8"; + #[inline] fn quantization_params_for( values: impl Iterator, 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 0a6dc972c2..0b5e927e66 100644 --- a/lib/sparse/src/index/inverted_index/inverted_index_ram.rs +++ b/lib/sparse/src/index/inverted_index/inverted_index_ram.rs @@ -1,10 +1,16 @@ use std::borrow::Cow; +#[cfg(feature = "testing")] +use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use common::counter::hardware_counter::HardwareCounterCell; use common::storage_version::StorageVersion; use common::types::PointOffsetType; use common::universal_io::Result; +#[cfg(feature = "testing")] +use fs_err as fs; +#[cfg(feature = "testing")] +use zerocopy::{FromBytes, IntoBytes}; use crate::common::sparse_vector::RemappedSparseVector; use crate::common::types::{DimId, DimOffset}; @@ -197,6 +203,72 @@ impl InvertedIndexRam { } } +#[cfg(feature = "testing")] +/// Load/save methods to speed up benchmark setup. +impl InvertedIndexRam { + const TEST_MAGIC: &'static [u8] = b"InvertedIndexRam for benchmarks."; + + pub fn test_save(&self, path: &Path) -> std::io::Result<()> { + let InvertedIndexRam { + postings, + vector_count, + total_sparse_size, + } = self; + + let mut f = BufWriter::new(fs::File::create(path)?); + f.write_all(Self::TEST_MAGIC)?; + f.write_all(postings.len().as_bytes())?; + for posting in postings { + f.write_all(posting.elements.len().as_bytes())?; + let bytes = posting.elements.as_bytes(); + f.write_all(bytes)?; + // padding + f.write_all(&0usize.to_ne_bytes()[..bytes.len() % size_of::()])?; + } + f.write_all(vector_count.as_bytes())?; + f.write_all(total_sparse_size.as_bytes())?; + + Ok(()) + } + + pub fn test_load(path: &Path) -> std::io::Result { + let f = fs::File::open(path)?; + let mmap = unsafe { memmap2::Mmap::map(&f)? }; + let bytes = &mmap[..]; + + let (magic, bytes) = + <[u8]>::ref_from_prefix_with_elems(bytes, Self::TEST_MAGIC.len()).unwrap(); + if magic != Self::TEST_MAGIC { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Invalid magic", + )); + } + + let (postings_len, bytes) = usize::ref_from_prefix(bytes).unwrap(); + let mut postings = Vec::with_capacity(*postings_len); + let mut bytes_cursor = bytes; + for _ in 0..*postings_len { + let (elements_len, bytes) = usize::ref_from_prefix(bytes_cursor).unwrap(); + let (elements, bytes) = + <[PostingElementEx]>::ref_from_prefix_with_elems(bytes, *elements_len).unwrap(); + let elements = elements.to_vec(); + postings.push(PostingList { elements }); + bytes_cursor = &bytes[bytes.as_ptr().align_offset(size_of::())..]; + } + let bytes = bytes_cursor; + + let (&vector_count, bytes) = usize::ref_from_prefix(bytes).unwrap(); + let (&total_sparse_size, _bytes) = usize::ref_from_prefix(bytes).unwrap(); + + Ok(InvertedIndexRam { + postings, + vector_count, + total_sparse_size, + }) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/lib/sparse/src/index/posting_list_common.rs b/lib/sparse/src/index/posting_list_common.rs index 2841de11cc..c68c0594bc 100644 --- a/lib/sparse/src/index/posting_list_common.rs +++ b/lib/sparse/src/index/posting_list_common.rs @@ -1,5 +1,5 @@ use common::types::PointOffsetType; -use zerocopy::{FromBytes, Immutable, KnownLayout}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; use crate::common::types::DimWeight; @@ -23,7 +23,7 @@ pub struct PostingElement { pub weight: DimWeight, } -#[derive(Debug, Clone, PartialEq, FromBytes, Immutable, KnownLayout)] +#[derive(Debug, Clone, PartialEq, FromBytes, Immutable, IntoBytes, KnownLayout)] #[repr(C)] pub struct PostingElementEx { /// Record ID