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