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

Cache sparse benchmark setup (#9141)

Browse Source 
* dataset: add features to fs-err

Reason: fix `cargo check --all-targets -p dataset`. Feature unification
pulls `fs-err/debug` + `fs-err/tokio` without `fs-err/debug_tokio`.

* dataset: run `cargo metadata` to get target dir

Reason: avoid re-downloading in separate worktrees. The old
`cargo locate-project` ignores custom target dirs.

* sparse/benches/search: cache vectors/indices
This commit is contained in:
xzfc
2026-05-23 16:45:42 +00:00
committed by timvisee
parent a57df81e1f
commit 3729c62117
8 changed files with 241 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
Cargo.lock generated
View File

@@ -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",

4
lib/common/dataset/Cargo.toml
View File

@@ -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"] }

39
lib/common/dataset/src/lib.rs
View File

@@ -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<PathBuf> {
};
// 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<PathBuf> {
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<PathBuf> = 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))
}

2
lib/sparse/Cargo.toml
View File

@@ -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 }

150
lib/sparse/benches/search.rs
View File

@@ -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::<Vec<_>>();
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<Path>, ratio: f32) -> io::Result<InvertedIndexRam> {
let csr = Csr::open(path.as_ref())?;
fn load_csr_index(path: impl AsRef<Path>, 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<W: Weight>(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::<W>::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"))]

14
lib/sparse/src/common/types.rs
View File

@@ -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<Item = DimWeight> + 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<Item = DimWeight> + 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<Item = DimWeight> + 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<Item = DimWeight> + 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<Item = DimWeight>,

72
lib/sparse/src/index/inverted_index/inverted_index_ram.rs
View File

@@ -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::<usize>()])?;
}
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<Self> {
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::<usize>())..];
}
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::*;

4
lib/sparse/src/index/posting_list_common.rs
View File

@@ -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