1use super::reject_key_failures;
2use crate::{
3 Bare, BareValueDomain, Definite, EvaluateOperand, Explain, Failure, IndexDomain, Indexed,
4 Labeled, Multiple, Operand, QueryResult, Single, Unordered, ValueDomain,
5 error::grouping::MissingGroupAggregate,
6 execution::EvaluationCache,
7 index::GroupKey,
8 operands::{OperandHandle, Partition},
9 operations::{Apply, GroupKernel, Operation, OperationContext, Prepare},
10 optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
11 registry::operation_manifest,
12 traits::UngroupKeyed,
13};
14use graphrecords_core::GraphRecord;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Group)]
18#[explain(label = "UngroupKeyed")]
19#[plan(optimizer_hints(empty = if_any))]
20pub struct UngroupKeyedOperation;
21
22fn bucket_element_estimate(input: &Estimate) -> Estimate {
23 Estimate {
24 elements: input.elements,
25 distinct: input.elements,
26 selectivity: input
27 .per_group
28 .as_deref()
29 .and_then(|inner| inner.selectivity),
30 per_group: None,
31 }
32}
33
34impl Prepare for UngroupKeyedOperation {
35 type Prepared<'a> = ();
36
37 fn prepare<'a>(
38 &'a self,
39 _graphrecord: &'a GraphRecord,
40 _cache: &'a EvaluationCache<'a>,
41 ) -> QueryResult<Self::Prepared<'a>> {
42 Ok(())
43 }
44}
45
46impl<M: IndexDomain, K: GroupKey, I: IndexDomain, V: ValueDomain>
47 GroupKernel<M, K, OperandHandle<Indexed<I, V>, Single>> for UngroupKeyedOperation
48{
49 type Output = OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
50
51 fn execute<'a>(
52 graphrecord: &'a GraphRecord,
53 partition: Partition<'a, M, K, OperandHandle<Indexed<I, V>, Single>>,
54 _prepared: Self::Prepared<'a>,
55 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
56 let (buckets, key_failures) = partition.into_parts();
57
58 reject_key_failures::<M>(key_failures, Self::LABEL)?;
59
60 let elements: Vec<_> = buckets
61 .into_iter()
62 .map(|(key, _, payload)| {
63 let index = K::resolve_key(Self::LABEL, graphrecord, &key)?;
64 let outcome = match payload {
65 Ok(Some((_, outcome))) => outcome,
66 Ok(None) => Err(Failure::new_at::<K, _>(
67 Self::LABEL,
68 MissingGroupAggregate,
69 &index,
70 )),
71 Err(failure) => Err(failure),
72 };
73
74 Ok((index, outcome))
75 })
76 .collect::<QueryResult<_>>()?;
77
78 Ok(Box::new(elements.into_iter()))
79 }
80
81 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
82 bucket_element_estimate(&input)
83 }
84}
85
86impl<M: IndexDomain, K: GroupKey, V: BareValueDomain>
87 GroupKernel<M, K, OperandHandle<Bare<V>, Single>> for UngroupKeyedOperation
88{
89 type Output = OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
90
91 fn execute<'a>(
92 graphrecord: &'a GraphRecord,
93 partition: Partition<'a, M, K, OperandHandle<Bare<V>, Single>>,
94 _prepared: Self::Prepared<'a>,
95 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
96 let (buckets, key_failures) = partition.into_parts();
97
98 reject_key_failures::<M>(key_failures, Self::LABEL)?;
99
100 let elements: Vec<_> = buckets
101 .into_iter()
102 .map(|(key, _, payload)| {
103 let index = K::resolve_key(Self::LABEL, graphrecord, &key)?;
104 let outcome = match payload {
105 Ok(Some(outcome)) => outcome,
106 Ok(None) => Err(Failure::new_at::<K, _>(
107 Self::LABEL,
108 MissingGroupAggregate,
109 &index,
110 )),
111 Err(failure) => Err(failure),
112 };
113
114 Ok((index, outcome))
115 })
116 .collect::<QueryResult<_>>()?;
117
118 Ok(Box::new(elements.into_iter()))
119 }
120
121 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
122 bucket_element_estimate(&input)
123 }
124}
125
126impl<M: IndexDomain, K: GroupKey, I: IndexDomain, V: ValueDomain>
127 GroupKernel<M, K, OperandHandle<Indexed<I, V>, Definite>> for UngroupKeyedOperation
128{
129 type Output = OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
130
131 fn execute<'a>(
132 graphrecord: &'a GraphRecord,
133 partition: Partition<'a, M, K, OperandHandle<Indexed<I, V>, Definite>>,
134 _prepared: Self::Prepared<'a>,
135 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
136 let (buckets, key_failures) = partition.into_parts();
137
138 reject_key_failures::<M>(key_failures, Self::LABEL)?;
139
140 let elements: Vec<_> = buckets
141 .into_iter()
142 .map(|(key, _, payload)| {
143 let index = K::resolve_key(Self::LABEL, graphrecord, &key)?;
144 let outcome = match payload {
145 Ok((_, outcome)) => outcome,
146 Err(failure) => Err(failure),
147 };
148
149 Ok((index, outcome))
150 })
151 .collect::<QueryResult<_>>()?;
152
153 Ok(Box::new(elements.into_iter()))
154 }
155
156 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
157 bucket_element_estimate(&input)
158 }
159}
160
161impl<M: IndexDomain, K: GroupKey, V: BareValueDomain>
162 GroupKernel<M, K, OperandHandle<Bare<V>, Definite>> for UngroupKeyedOperation
163{
164 type Output = OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
165
166 fn execute<'a>(
167 graphrecord: &'a GraphRecord,
168 partition: Partition<'a, M, K, OperandHandle<Bare<V>, Definite>>,
169 _prepared: Self::Prepared<'a>,
170 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
171 let (buckets, key_failures) = partition.into_parts();
172
173 reject_key_failures::<M>(key_failures, Self::LABEL)?;
174
175 let elements: Vec<_> = buckets
176 .into_iter()
177 .map(|(key, _, payload)| {
178 let index = K::resolve_key(Self::LABEL, graphrecord, &key)?;
179 let outcome = match payload {
180 Ok(outcome) => outcome,
181 Err(failure) => Err(failure),
182 };
183
184 Ok((index, outcome))
185 })
186 .collect::<QueryResult<_>>()?;
187
188 Ok(Box::new(elements.into_iter()))
189 }
190
191 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
192 bucket_element_estimate(&input)
193 }
194}
195
196impl<O: Apply<UngroupKeyedOperation>> UngroupKeyed for O {
197 type ReturnOperand = O::Output;
198
199 fn ungroup_keyed(&self) -> Self::ReturnOperand {
200 Self::ReturnOperand::new(OperationContext::new(self.clone(), UngroupKeyedOperation))
201 }
202}
203
204operation_manifest! {
205 UngroupKeyedOperation {
206 method: UngroupKeyed::ungroup_keyed;
207 scope: group;
208
209 kernel {
210 group: <M: IndexDomain, K: GroupKey>;
211 parameters: <I: IndexDomain, V: ValueDomain>;
212 input: OperandHandle<Indexed<I, V>, Single>;
213 output: OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
214 }
215 kernel {
216 group: <M: IndexDomain, K: GroupKey>;
217 parameters: <V: BareValueDomain>;
218 input: OperandHandle<Bare<V>, Single>;
219 output: OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
220 }
221 kernel {
222 group: <M: IndexDomain, K: GroupKey>;
223 parameters: <I: IndexDomain, V: ValueDomain>;
224 input: OperandHandle<Indexed<I, V>, Definite>;
225 output: OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
226 }
227 kernel {
228 group: <M: IndexDomain, K: GroupKey>;
229 parameters: <V: BareValueDomain>;
230 input: OperandHandle<Bare<V>, Definite>;
231 output: OperandHandle<Indexed<K, V>, Multiple<Unordered>>;
232 }
233 }
234}