graphrecords_query/operations/grouping/
group_by.rs

1use crate::{
2    Definite, EvaluateOperand, Explain, IndexDomain, Indexed, Labeled, Multiple, Operand,
3    QueryResult, Single, ValueDomain,
4    capabilities::GroupingValue,
5    element::Retention,
6    execution::EvaluationCache,
7    operands::{
8        GroupOperand, OperandHandle, PartitionArity, PartitionBuilder, PartitionClassification,
9    },
10    operations::{
11        Apply, ArgumentSource, Keyed, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
12    },
13    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
14    registry::operation_manifest,
15    traits::GroupBy,
16};
17use graphrecords_core::GraphRecord;
18
19#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
20#[operation(scope = Lane)]
21#[explain(label = "GroupBy")]
22pub struct GroupByOperation<K> {
23    #[argument]
24    key: K,
25}
26
27impl<K: Prepare> Prepare for GroupByOperation<K> {
28    type Prepared<'a>
29        = K::Prepared<'a>
30    where
31        Self: 'a;
32
33    fn prepare<'a>(
34        &'a self,
35        graphrecord: &'a GraphRecord,
36        cache: &'a EvaluationCache<'a>,
37    ) -> QueryResult<Self::Prepared<'a>> {
38        self.key.prepare(graphrecord, cache)
39    }
40}
41
42impl<I, V, K, C> LaneKernel<Indexed<I, V>, C> for GroupByOperation<K>
43where
44    I: IndexDomain,
45    V: ValueDomain,
46    K: ArgumentSource<Keyed<I>>,
47    K::ValueDomain: GroupingValue,
48    C: PartitionArity<Indexed<I, V>>,
49{
50    type Output =
51        GroupOperand<I, <K::ValueDomain as GroupingValue>::Key, OperandHandle<Indexed<I, V>, C>>;
52
53    fn execute<'a>(
54        _graphrecord: &'a GraphRecord,
55        values: KeyedStream<'a, I, V, C>,
56        prepared: Self::Prepared<'a>,
57    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
58        let label = Self::LABEL;
59
60        PartitionBuilder::<_, _, _, C>::new(values).build(|element| {
61            let member = &element.0;
62            let step = K::resolve(&prepared, member, label);
63
64            match K::Retention::collapse(step) {
65                None => PartitionClassification::Omit,
66                Some(Err(failure)) => PartitionClassification::KeyFailure(failure),
67                Some(Ok(value)) => {
68                    PartitionClassification::Key(K::ValueDomain::to_group_key(&value))
69                }
70            }
71        })
72    }
73
74    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
75        let (Some(elements), Some(keys)) = (input.elements, self.key.estimate(stats).distinct)
76        else {
77            return Estimate {
78                per_group: Some(Box::new(Estimate::UNKNOWN)),
79                ..Estimate::UNKNOWN
80            };
81        };
82
83        let groups = keys.min(elements);
84        let per_group_elements = if groups == 0 {
85            0
86        } else {
87            elements.div_ceil(groups)
88        };
89
90        Estimate {
91            elements: Some(groups),
92            distinct: Some(groups),
93            selectivity: None,
94            per_group: Some(Box::new(Estimate {
95                elements: Some(per_group_elements),
96                distinct: input
97                    .distinct
98                    .map(|distinct| distinct.min(per_group_elements)),
99                selectivity: input.selectivity,
100                per_group: None,
101            })),
102        }
103    }
104}
105
106impl<O, K> GroupBy<K> for O
107where
108    O: Apply<GroupByOperation<K>>,
109    GroupByOperation<K>: Operation,
110{
111    type ReturnOperand = O::Output;
112
113    fn group_by(&self, key: K) -> Self::ReturnOperand {
114        Self::ReturnOperand::new(OperationContext::new(
115            self.clone(),
116            GroupByOperation { key },
117        ))
118    }
119}
120
121operation_manifest! {
122    GroupByOperation<K> {
123        method: GroupBy<K>::group_by;
124        scope: lane;
125
126        kernel {
127            parameters: <I: IndexDomain, V: ValueDomain, X: GroupingValue, O: OrderState>;
128            argument: K: ArgumentSource<Keyed<I>, X>;
129            input: (Indexed<I, V>, Multiple<O>);
130            output: GroupOperand<
131                I,
132                <X as GroupingValue>::Key,
133                OperandHandle<Indexed<I, V>, Multiple<O>>,
134            >;
135        }
136
137        kernel {
138            parameters: <I: IndexDomain, V: ValueDomain, X: GroupingValue>;
139            argument: K: ArgumentSource<Keyed<I>, X>;
140            input: (Indexed<I, V>, Single);
141            output: GroupOperand<
142                I,
143                <X as GroupingValue>::Key,
144                OperandHandle<Indexed<I, V>, Single>,
145            >;
146        }
147
148        kernel {
149            parameters: <I: IndexDomain, V: ValueDomain, X: GroupingValue>;
150            argument: K: ArgumentSource<Keyed<I>, X>;
151            input: (Indexed<I, V>, Definite);
152            output: GroupOperand<
153                I,
154                <X as GroupingValue>::Key,
155                OperandHandle<Indexed<I, V>, Definite>,
156            >;
157        }
158    }
159}
graphrecords_query/operations/grouping/group_by.rs

graphrecords_query/operations/grouping/
group_by.rs
