qdrant/lib/segment/benches/vector_search.rs

use std::array;
use std::sync::Arc;

use atomic_refcell::AtomicRefCell;
use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
use rand::RngExt;
use rand::distr::StandardUniform;
use rand::rngs::SmallRng;
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, Memory};
use segment::vector_storage::dense::dense_vector_storage::open_dense_vector_storage;
use segment::vector_storage::{DEFAULT_STOPPED, DenseVectorStorage, VectorStorageEnum};
use tempfile::Builder;

#[cfg(not(target_os = "windows"))]
mod prof;

const DIM: usize = 1024;

fn random_vector(size: usize) -> DenseVector {
    rand::make_rng::<SmallRng>()
        .sample_iter(StandardUniform)
        .take(size)
        .collect()
}

fn random_query_batch<const SIZE: usize>() -> [QueryVector; SIZE] {
    array::from_fn(|_| QueryVector::from(random_vector(DIM)))
}

fn benchmark<const IO_URING: bool, const VECTORS: usize, const BATCH: usize>(c: &mut Criterion) {
    let tmp = Builder::new()
        .prefix("vector-search-bench")
        .tempdir()
        .expect("tempdir created");

    #[cfg(target_os = "linux")]
    segment::vector_storage::common::set_async_scorer(IO_URING);

    #[cfg(not(target_os = "linux"))]
    assert!(!IO_URING, "async scorer is only supported on Linux");

    let mut storage = open_dense_vector_storage(tmp.path(), DIM, Distance::Dot, Memory::Cold)
        .expect("vector storage created");

    let mut vectors = (0..VECTORS).map(|_| {
        let vector = random_vector(DIM);
        (std::borrow::Cow::Owned(vector), false)
    });

    let result = match &mut storage {
        VectorStorageEnum::DenseMemmap(v) => v.update_from(&mut vectors, &DEFAULT_STOPPED),
        #[cfg(target_os = "linux")]
        VectorStorageEnum::DenseUring(v) => v.update_from(&mut vectors, &DEFAULT_STOPPED),
        _ => panic!("unexpected dense vector storage variant"),
    };
    result.expect("vector storage populated");

    let id_tracker = Arc::new(AtomicRefCell::new(create_id_tracker_fixture(VECTORS)));
    let id_tracker = id_tracker.borrow();

    let mut group = c.benchmark_group("vector search");

    let benchmark_id = format!(
        "{} storage/{}k vectors/batch of {BATCH}",
        if IO_URING { "io_uring" } else { "mmap" },
        VECTORS / 1000,
    );

    group.bench_function(benchmark_id, |b| {
        b.iter_batched(
            random_query_batch::<BATCH>,
            |vectors| {
                BatchFilteredSearcher::new_for_test(
                    &vectors,
                    &storage,
                    id_tracker.deleted_point_bitslice(),
                    10,
                )
                .peek_top_all(&DEFAULT_STOPPED)
                .expect("points scored")
            },
            BatchSize::SmallInput,
        )
    });
}

#[cfg(target_os = "linux")]
criterion_group! {
    name = benches;
    config = Criterion::default().with_profiler(prof::FlamegraphProfiler::new(1000));
    targets =
        benchmark::<false, 10_000, 1>,
        benchmark::<false, 10_000, 4>,
        benchmark::<false, 100_000, 1>,
        benchmark::<false, 100_000, 4>,

        benchmark::<true, 10_000, 1>,
        benchmark::<true, 10_000, 4>,
        benchmark::<true, 100_000, 1>,
        benchmark::<true, 100_000, 4>,
}

#[cfg(not(any(target_os = "linux", target_os = "windows")))]
criterion_group! {
    name = benches;
    config = Criterion::default().with_profiler(prof::FlamegraphProfiler::new(1000));
    targets =
        benchmark::<false, 10_000, 1>,
        benchmark::<false, 10_000, 4>,
        benchmark::<false, 100_000, 1>,
        benchmark::<false, 100_000, 4>,
}

#[cfg(target_os = "windows")]
criterion_group! {
    name = benches;
    config = Criterion::default();
    targets =
        benchmark::<false, 10_000, 1>,
        benchmark::<false, 10_000, 4>,
        benchmark::<false, 100_000, 1>,
        benchmark::<false, 100_000, 4>,
}

criterion_main!(benches);