qdrant-edge 0.7.0

A lightweight, in-process vector search engine designed for embedded devices, autonomous systems, and mobile agents.
Documentation
use itertools::Itertools;
use rand::{Rng, RngExt};

use crate::segment::data_types::vectors::{QueryVector, VectorInternal};
use crate::segment::fixtures::payload_fixtures::random_multi_vector;
use crate::segment::vector_storage::query::{ContextPair, ContextQuery, DiscoverQuery, RecoQuery};

const MAX_EXAMPLE_PAIRS: usize = 4;

pub enum QueryVariant {
    Nearest,
    RecoBestScore,
    RecoSumScores,
    Discover,
    Context,
}

pub fn random_query<R: Rng + ?Sized>(
    variant: &QueryVariant,
    rng: &mut R,
    rand_vec: impl Fn(&mut R) -> VectorInternal,
) -> QueryVector {
    match variant {
        QueryVariant::Nearest => rand_vec(rng).into(),
        QueryVariant::Discover => random_discover_query(rng, rand_vec),
        QueryVariant::Context => random_context_query(rng, rand_vec),
        QueryVariant::RecoBestScore => {
            QueryVector::RecommendBestScore(random_reco_query(rng, rand_vec))
        }
        QueryVariant::RecoSumScores => {
            QueryVector::RecommendSumScores(random_reco_query(rng, rand_vec))
        }
    }
}

pub fn random_multi_vec_query<R: Rng + ?Sized>(
    variant: &QueryVariant,
    rng: &mut R,
    dim: usize,
    num_vector_per_points: usize,
) -> QueryVector {
    let rand_vec = move |rng: &mut R| -> VectorInternal {
        random_multi_vector(rng, dim, num_vector_per_points).into()
    };

    random_query(variant, rng, rand_vec)
}

fn random_discover_query<R: Rng + ?Sized>(
    rng: &mut R,
    rand_vec: impl Fn(&mut R) -> VectorInternal,
) -> QueryVector {
    let num_pairs: usize = rng.random_range(1..MAX_EXAMPLE_PAIRS);

    let target = rand_vec(rng);

    let pairs = (0..num_pairs)
        .map(|_| {
            let positive = rand_vec(rng);
            let negative = rand_vec(rng);
            ContextPair { positive, negative }
        })
        .collect_vec();

    DiscoverQuery::new(target, pairs).into()
}

fn random_context_query<R: Rng + ?Sized>(
    rng: &mut R,
    rand_vec: impl Fn(&mut R) -> VectorInternal,
) -> QueryVector {
    let num_pairs: usize = rng.random_range(0..MAX_EXAMPLE_PAIRS);

    let pairs = (0..num_pairs)
        .map(|_| {
            let positive = rand_vec(rng);
            let negative = rand_vec(rng);
            ContextPair { positive, negative }
        })
        .collect_vec();

    QueryVector::Context(ContextQuery::new(pairs))
}

fn random_reco_query<R: Rng + ?Sized>(
    rng: &mut R,
    rand_vec: impl Fn(&mut R) -> VectorInternal,
) -> RecoQuery<VectorInternal> {
    let num_examples: usize = rng.random_range(1..MAX_EXAMPLE_PAIRS);

    let positive = (0..num_examples).map(|_| rand_vec(rng)).collect_vec();
    let negative = (0..num_examples).map(|_| rand_vec(rng)).collect_vec();

    RecoQuery::new(positive, negative)
}