qdrant_edge/shard/grouping/
mod.rs

1//! Building blocks for group-by queries, shared between the full server implementation
2//! (multi-request group filling in `collection`) and the single-request edge implementation,
3//! so both interpret group keys and shape candidate queries identically.
4
5pub mod aggregator;
6mod driver;
7
8use ahash::AHashMap;
9use fnv::FnvBuildHasher;
10use indexmap::IndexSet;
11use crate::segment::data_types::groups::GroupId;
12use crate::segment::json_path::JsonPath;
13use crate::segment::types::{
14    AnyVariants, Condition, FieldCondition, Filter, Match, PointIdType, ScoredPoint,
15    WithPayloadInterface,
16};
17use serde_json::Value;
18
19pub use self::aggregator::GroupsAggregator;
20pub use self::driver::{GroupByDriver, RequestBudget};
21use crate::shard::query::{ShardPrefetch, ShardQueryRequest};
22
23#[derive(PartialEq, Debug)]
24pub enum AggregatorError {
25    BadKeyType,
26    KeyNotFound,
27}
28
29/// A group of points that share the same value of the `group_by` field.
30#[derive(Debug, Clone)]
31pub struct Group {
32    pub hits: Vec<ScoredPoint>,
33    pub key: GroupId,
34}
35
36impl Group {
37    pub fn hydrate_from(&mut self, map: &AHashMap<PointIdType, ScoredPoint>) {
38        self.hits.iter_mut().for_each(|hit| {
39            if let Some(point) = map.get(&hit.id) {
40                hit.payload.clone_from(&point.payload);
41                hit.vector.clone_from(&point.vector);
42            }
43        });
44    }
45}
46
47/// Make `group_by` field selector work with as `with_payload`.
48fn group_by_to_payload_selector(group_by: &JsonPath) -> WithPayloadInterface {
49    WithPayloadInterface::Fields(vec![group_by.strip_wildcard_suffix()])
50}
51
52/// Merge an extra filter into an optional existing one.
53fn merge_filter(target: &mut Option<Filter>, extra: Filter) {
54    *target = Some(match target.take() {
55        Some(filter) => filter.merge_owned(extra),
56        None => extra,
57    });
58}
59
60/// Rewrite a base scoring query into the query used to fetch group candidates: restrict it to
61/// points that carry the `group_by` field, fetch `limit` candidates with only the `group_by`
62/// payload, and scale nested prefetch limits by `group_size` so enough candidates survive
63/// every rescoring stage.
64fn shape_candidates_query(
65    query: &mut ShardQueryRequest,
66    group_by: &JsonPath,
67    limit: usize,
68    group_size: usize,
69) {
70    query.limit = limit;
71    query.offset = 0;
72    query
73        .prefetches
74        .iter_mut()
75        .for_each(|prefetch| increase_limit_for_group(prefetch, group_size));
76
77    let key_not_empty = Filter::new_must_not(Condition::IsEmpty(group_by.clone().into()));
78    merge_filter(&mut query.filter, key_not_empty);
79
80    query.with_payload = group_by_to_payload_selector(group_by);
81}
82
83fn increase_limit_for_group(shard_prefetch: &mut ShardPrefetch, group_size: usize) {
84    shard_prefetch.limit *= group_size;
85    shard_prefetch.prefetches.iter_mut().for_each(|prefetch| {
86        increase_limit_for_group(prefetch, group_size);
87    });
88}
89
90/// Uses the set of values to create Match::Except's, if possible
91fn except_on(path: &JsonPath, values: &[Value]) -> Vec<Condition> {
92    values_to_any_variants(values)
93        .into_iter()
94        .map(|v| {
95            Condition::Field(FieldCondition::new_match(
96                path.clone(),
97                Match::new_except(v),
98            ))
99        })
100        .collect()
101}
102
103/// Uses the set of values to create Match::Any's, if possible
104fn match_on(path: &JsonPath, values: &[Value]) -> Vec<Condition> {
105    values_to_any_variants(values)
106        .into_iter()
107        .map(|any_variants| {
108            Condition::Field(FieldCondition::new_match(
109                path.clone(),
110                Match::new_any(any_variants),
111            ))
112        })
113        .collect()
114}
115
116fn values_to_any_variants(values: &[Value]) -> Vec<AnyVariants> {
117    let mut any_variants = Vec::new();
118
119    // gather int values
120    let ints: IndexSet<_, FnvBuildHasher> = values.iter().filter_map(|v| v.as_i64()).collect();
121
122    if !ints.is_empty() {
123        any_variants.push(AnyVariants::Integers(ints));
124    }
125
126    // gather string values
127    let strs: IndexSet<_, FnvBuildHasher> = values
128        .iter()
129        .filter_map(|v| v.as_str().map(Into::into))
130        .collect();
131
132    if !strs.is_empty() {
133        any_variants.push(AnyVariants::Strings(strs));
134    }
135
136    any_variants
137}
qdrant_edge/shard/grouping/mod.rs

qdrant_edge/shard/grouping/
mod.rs
