qdrant/lib/collection/benches/batch_query_bench.rs

use std::sync::Arc;

use api::rest::SearchRequestInternal;
use collection::common::adaptive_handle::AdaptiveSearchHandle;
use collection::config::{CollectionConfigInternal, CollectionParams, WalConfig};
use collection::operations::CollectionUpdateOperations;
use collection::operations::point_ops::{
    PointInsertOperationsInternal, PointOperations, PointStructPersisted,
};
use collection::operations::universal_query::shard_query::{
    FusionInternal, ScoringQuery, ShardPrefetch, ShardQueryRequest,
};
use collection::operations::vector_params_builder::VectorParamsBuilder;
use collection::optimizers_builder::OptimizersConfig;
use collection::shards::local_shard::LocalShard;
use collection::shards::shard_trait::{ShardOperation, WaitUntil};
use common::budget::ResourceBudget;
use common::counter::hardware_accumulator::HwMeasurementAcc;
use common::save_on_disk::SaveOnDisk;
use criterion::{Criterion, criterion_group, criterion_main};
use ordered_float::OrderedFloat;
use rand::rngs::SmallRng;
use segment::common::reciprocal_rank_fusion::DEFAULT_RRF_K;
use segment::data_types::vectors::{VectorStructInternal, only_default_vector};
use segment::fixtures::payload_fixtures::random_vector;
use segment::types::{
    Condition, Distance, FieldCondition, Filter, Payload, Range, WithPayloadInterface, WithVector,
};
use serde_json::Map;
use shard::search::CoreSearchRequestBatch;
use tempfile::{Builder, TempDir};
use tokio::runtime::Runtime;
use tokio::sync::RwLock;

fn setup() -> (TempDir, LocalShard, Runtime) {
    let storage_dir = Builder::new().prefix("storage").tempdir().unwrap();

    let runtime = Runtime::new().unwrap();
    let handle = runtime.handle().clone();

    let wal_config = WalConfig {
        wal_capacity_mb: 1,
        wal_segments_ahead: 0,
        wal_retain_closed: 1,
    };

    let collection_params = CollectionParams {
        vectors: VectorParamsBuilder::new(100, Distance::Dot).build().into(),
        ..CollectionParams::empty()
    };

    let collection_config = CollectionConfigInternal {
        params: collection_params,
        optimizer_config: OptimizersConfig {
            deleted_threshold: 0.9,
            vacuum_min_vector_number: 1000,
            default_segment_number: 2,
            max_segment_size: Some(100_000),
            #[expect(deprecated)]
            memmap_threshold: Some(100_000),
            indexing_threshold: Some(50_000),
            flush_interval_sec: 30,
            max_optimization_threads: Some(2),
            prevent_unoptimized: None,
        },
        wal_config,
        hnsw_config: Default::default(),
        quantization_config: Default::default(),
        strict_mode_config: Default::default(),
        uuid: None,
        metadata: None,
    };

    let optimizers_config = collection_config.optimizer_config.clone();

    let shared_config = Arc::new(RwLock::new(collection_config));

    let payload_index_schema_dir = Builder::new().prefix("qdrant-test").tempdir().unwrap();
    let payload_index_schema_file = payload_index_schema_dir.path().join("payload-schema.json");
    let payload_index_schema =
        Arc::new(SaveOnDisk::load_or_init_default(payload_index_schema_file).unwrap());

    let shard = handle
        .block_on(LocalShard::build_local(
            0,
            "test_collection".to_string(),
            storage_dir.path(),
            shared_config,
            Default::default(),
            payload_index_schema,
            handle.clone(),
            AdaptiveSearchHandle::new_fixed(handle.clone()),
            ResourceBudget::default(),
            optimizers_config,
        ))
        .unwrap();

    let rnd_batch = create_rnd_batch();

    handle
        .block_on(shard.update(
            rnd_batch.into(),
            WaitUntil::Visible,
            None,
            HwMeasurementAcc::new(),
        ))
        .unwrap();

    (storage_dir, shard, runtime)
}

fn create_rnd_batch() -> CollectionUpdateOperations {
    let mut rng = rand::make_rng::<SmallRng>();
    let num_points = 2000;
    let dim = 100;
    let mut points = Vec::with_capacity(num_points);
    for i in 0..num_points {
        let mut payload_map = Map::new();
        payload_map.insert("a".to_string(), (i % 5).into());
        let vector = random_vector(&mut rng, dim);
        let vectors = only_default_vector(&vector);
        let point = PointStructPersisted {
            id: (i as u64).into(),
            vector: VectorStructInternal::from(vectors).into(),
            payload: Some(Payload(payload_map)),
        };
        points.push(point);
    }
    CollectionUpdateOperations::PointOperation(PointOperations::UpsertPoints(
        PointInsertOperationsInternal::PointsList(points),
    ))
}

fn some_filters() -> Vec<Option<Filter>> {
    vec![
        None,
        Some(Filter::new_must(Condition::Field(
            FieldCondition::new_match("a".parse().unwrap(), 3.into()),
        ))),
        Some(Filter::new_must(Condition::Field(
            FieldCondition::new_range(
                "a".parse().unwrap(),
                Range {
                    lt: None,
                    gt: Some(OrderedFloat(-1.)),
                    gte: None,
                    lte: Some(OrderedFloat(100.0)),
                },
            ),
        ))),
    ]
}

/// Compare nearest neighbors query vs normal search
fn batch_search_bench(c: &mut Criterion) {
    let (_tempdir, shard, search_runtime) = setup();

    let search_runtime_handle = AdaptiveSearchHandle::new_fixed(search_runtime.handle().clone());

    let mut group = c.benchmark_group("batch-search-bench");

    let batch_size = 100;

    for (fid, filter) in some_filters().into_iter().enumerate() {
        group.bench_function(format!("query-batch-{fid}"), |b| {
            b.iter(|| {
                search_runtime.block_on(async {
                    let mut rng = rand::make_rng::<SmallRng>();
                    let mut searches = Vec::with_capacity(batch_size);
                    for _i in 0..batch_size {
                        let query = random_vector(&mut rng, 100);
                        let search_query = ShardQueryRequest {
                            prefetches: vec![],
                            query: Some(ScoringQuery::Vector(query.into())),
                            filter: filter.clone(),
                            params: None,
                            limit: 10,
                            offset: 0,
                            with_payload: WithPayloadInterface::Bool(true),
                            with_vector: WithVector::Bool(false),
                            score_threshold: None,
                        };
                        searches.push(search_query);
                    }

                    let hw_acc = HwMeasurementAcc::new();
                    let result = shard
                        .query_batch(Arc::new(searches), &search_runtime_handle, None, hw_acc)
                        .await
                        .unwrap();
                    assert!(!result.is_empty());
                });
            })
        });

        group.bench_function(format!("search-batch-{fid}"), |b| {
            b.iter(|| {
                search_runtime.block_on(async {
                    let mut rng = rand::make_rng::<SmallRng>();
                    let mut searches = Vec::with_capacity(batch_size);
                    for _i in 0..batch_size {
                        let query = random_vector(&mut rng, 100);
                        let search_query = SearchRequestInternal {
                            vector: query.into(),
                            filter: filter.clone(),
                            params: None,
                            limit: 10,
                            offset: None,
                            with_payload: Some(WithPayloadInterface::Bool(true)),
                            with_vector: None,
                            score_threshold: None,
                        };
                        searches.push(search_query.into());
                    }

                    let hw_acc = HwMeasurementAcc::new();
                    let search_query = CoreSearchRequestBatch { searches };
                    let result = shard
                        .core_search(Arc::new(search_query), &search_runtime_handle, None, hw_acc)
                        .await
                        .unwrap();
                    assert!(!result.is_empty());
                });
            })
        });
    }

    group.finish();

    search_runtime.block_on(async {
        shard.stop_gracefully().await;
    });
}

fn batch_rrf_query_bench(c: &mut Criterion) {
    let (_tempdir, shard, search_runtime) = setup();

    let search_runtime_handle = AdaptiveSearchHandle::new_fixed(search_runtime.handle().clone());

    let mut group = c.benchmark_group("batch-rrf-bench");

    let batch_size = 100;

    for (fid, filter) in some_filters().into_iter().enumerate() {
        group.bench_function(format!("hybrid-query-batch-{fid}"), |b| {
            b.iter(|| {
                search_runtime.block_on(async {
                    let mut rng = rand::make_rng::<SmallRng>();
                    let mut searches = Vec::with_capacity(batch_size);
                    for _i in 0..batch_size {
                        let query1 = random_vector(&mut rng, 100);
                        let query2 = random_vector(&mut rng, 100);
                        let search_query = ShardQueryRequest {
                            prefetches: vec![
                                ShardPrefetch {
                                    prefetches: vec![],
                                    query: Some(ScoringQuery::Vector(query1.into())),
                                    limit: 100,
                                    params: None,
                                    filter: None,
                                    score_threshold: None,
                                },
                                ShardPrefetch {
                                    prefetches: vec![],
                                    query: Some(ScoringQuery::Vector(query2.into())),
                                    limit: 100,
                                    params: None,
                                    filter: None,
                                    score_threshold: None,
                                },
                            ],
                            query: Some(ScoringQuery::Fusion(FusionInternal::Rrf {
                                k: DEFAULT_RRF_K,
                                weights: None,
                            })),
                            filter: filter.clone(),
                            params: None,
                            limit: 10,
                            offset: 0,
                            with_payload: WithPayloadInterface::Bool(true),
                            with_vector: WithVector::Bool(false),
                            score_threshold: None,
                        };
                        searches.push(search_query);
                    }

                    let hw_acc = HwMeasurementAcc::new();
                    let result = shard
                        .query_batch(Arc::new(searches), &search_runtime_handle, None, hw_acc)
                        .await
                        .unwrap();
                    assert!(!result.is_empty());
                });
            })
        });
    }

    group.finish();

    search_runtime.block_on(async {
        shard.stop_gracefully().await;
    });
}

fn batch_rescore_bench(c: &mut Criterion) {
    let (_tempdir, shard, search_runtime) = setup();

    let search_runtime_handle = AdaptiveSearchHandle::new_fixed(search_runtime.handle().clone());

    let mut group = c.benchmark_group("batch-rescore-bench");

    let batch_size = 100;

    for (fid, filter) in some_filters().into_iter().enumerate() {
        group.bench_function(format!("rescore-query-batch-{fid}"), |b| {
            b.iter(|| {
                search_runtime.block_on(async {
                    let mut rng = rand::make_rng::<SmallRng>();
                    let mut searches = Vec::with_capacity(batch_size);
                    for _i in 0..batch_size {
                        let query1 = random_vector(&mut rng, 100);
                        let query2 = random_vector(&mut rng, 100);
                        let search_query = ShardQueryRequest {
                            prefetches: vec![ShardPrefetch {
                                prefetches: vec![],
                                query: Some(ScoringQuery::Vector(query1.into())),
                                limit: 100,
                                params: None,
                                filter: None,
                                score_threshold: None,
                            }],
                            query: Some(ScoringQuery::Vector(query2.into())),
                            filter: filter.clone(),
                            params: None,
                            limit: 10,
                            offset: 0,
                            with_payload: WithPayloadInterface::Bool(true),
                            with_vector: WithVector::Bool(false),
                            score_threshold: None,
                        };
                        searches.push(search_query);
                    }

                    let hw_acc = HwMeasurementAcc::new();
                    let result = shard
                        .query_batch(Arc::new(searches), &search_runtime_handle, None, hw_acc)
                        .await
                        .unwrap();
                    assert!(!result.is_empty());
                });
            })
        });
    }

    group.finish();

    search_runtime.block_on(async {
        shard.stop_gracefully().await;
    });
}

criterion_group! {
    name = benches;
    config = Criterion::default().significance_level(0.01).sample_size(500);
    targets = batch_search_bench, batch_rrf_query_bench, batch_rescore_bench
}

criterion_main!(benches);