1use super::reject_key_failures;
2use crate::{
3 Arity, Bare, BareValueDomain, Definite, EvaluateOperand, Explain, Failure, IndexDomain,
4 Indexed, Labeled, Operand, QueryResult, Single, ValueDomain,
5 capabilities::GroupingValue,
6 error::grouping::MissingGroupAggregate,
7 execution::EvaluationCache,
8 index::GroupKey,
9 operands::{OperandHandle, Partition},
10 operations::{
11 Apply, GroupKernel, IndexedElementContainer, IndexedElementSource, Operation,
12 OperationContext, Prepare,
13 },
14 optimizer::{
15 Estimate, Estimated, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats,
16 },
17 registry::operation_manifest,
18 traits::BroadcastVia,
19};
20use graphrecords_core::GraphRecord;
21use graphrecords_utils::aliases::GrHashMap;
22
23#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
24#[operation(scope = Group)]
25#[explain(label = "BroadcastVia")]
26#[plan(optimizer_hints(empty = if_all))]
27pub struct BroadcastViaOperation<A> {
28 #[argument]
29 via: A,
30}
31
32impl<A: Prepare> Prepare for BroadcastViaOperation<A> {
33 type Prepared<'a>
34 = A::Prepared<'a>
35 where
36 Self: 'a;
37
38 fn prepare<'a>(
39 &'a self,
40 graphrecord: &'a GraphRecord,
41 cache: &'a EvaluationCache<'a>,
42 ) -> QueryResult<Self::Prepared<'a>> {
43 self.via.prepare(graphrecord, cache)
44 }
45}
46
47fn broadcast_via<'a, K, V, A>(
48 prepared: A::Prepared<'a>,
49 aggregates: GrHashMap<K::Owned, Option<QueryResult<V::Value<'a>>>>,
50 label: &'static str,
51) -> IndexedElementContainer<'a, A::IndexDomain, V::Value<'a>, A::Arity>
52where
53 K: GroupKey,
54 V: ValueDomain,
55 A: IndexedElementSource + 'a,
56 A::ValueDomain: GroupingValue<Key = K>,
57{
58 let elements = A::elements(prepared);
59
60 A::Arity::map_elements(elements, move |(index, via_outcome)| {
61 let outcome = match via_outcome {
62 Err(failure) => Err(failure),
63 Ok(value) => match aggregates.get(&A::ValueDomain::to_group_key(&value)) {
64 Some(Some(aggregate)) => aggregate.clone(),
65 Some(None) | None => Err(Failure::new_at::<A::IndexDomain, _>(
66 label,
67 MissingGroupAggregate,
68 &index,
69 )),
70 },
71 };
72
73 (index, outcome)
74 })
75}
76
77fn broadcast_via_estimate<A: Estimated>(
78 operation: &BroadcastViaOperation<A>,
79 input: &Estimate,
80 stats: &Stats,
81) -> Estimate {
82 let via = operation.via.estimate(stats);
83
84 Estimate {
85 elements: via.elements,
86 distinct: input.elements,
87 selectivity: via.selectivity,
88 per_group: None,
89 }
90}
91
92impl<M, K, J, V, A> GroupKernel<M, K, OperandHandle<Indexed<J, V>, Single>>
93 for BroadcastViaOperation<A>
94where
95 M: IndexDomain,
96 K: GroupKey,
97 J: IndexDomain,
98 V: ValueDomain,
99 A: IndexedElementSource,
100 A::ValueDomain: GroupingValue<Key = K>,
101{
102 type Output = OperandHandle<Indexed<A::IndexDomain, V>, A::Arity>;
103
104 fn execute<'a>(
105 _graphrecord: &'a GraphRecord,
106 partition: Partition<'a, M, K, OperandHandle<Indexed<J, V>, Single>>,
107 prepared: Self::Prepared<'a>,
108 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
109 let (buckets, key_failures) = partition.into_parts();
110 reject_key_failures::<M>(key_failures, Self::LABEL)?;
111
112 let aggregates = buckets
113 .into_iter()
114 .map(|(key, _, payload)| {
115 let aggregate = match payload {
116 Ok(Some((_, outcome))) => Some(outcome),
117 Ok(None) => None,
118 Err(failure) => Some(Err(failure)),
119 };
120
121 (key, aggregate)
122 })
123 .collect();
124
125 Ok(broadcast_via::<K, V, A>(prepared, aggregates, Self::LABEL))
126 }
127
128 fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
129 broadcast_via_estimate(self, &input, stats)
130 }
131}
132
133impl<M, K, V, A> GroupKernel<M, K, OperandHandle<Bare<V>, Single>> for BroadcastViaOperation<A>
134where
135 M: IndexDomain,
136 K: GroupKey,
137 V: BareValueDomain,
138 A: IndexedElementSource,
139 A::ValueDomain: GroupingValue<Key = K>,
140{
141 type Output = OperandHandle<Indexed<A::IndexDomain, V>, A::Arity>;
142
143 fn execute<'a>(
144 _graphrecord: &'a GraphRecord,
145 partition: Partition<'a, M, K, OperandHandle<Bare<V>, Single>>,
146 prepared: Self::Prepared<'a>,
147 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
148 let (buckets, key_failures) = partition.into_parts();
149 reject_key_failures::<M>(key_failures, Self::LABEL)?;
150
151 let aggregates = buckets
152 .into_iter()
153 .map(|(key, _, payload)| {
154 let aggregate = match payload {
155 Ok(Some(outcome)) => Some(outcome),
156 Ok(None) => None,
157 Err(failure) => Some(Err(failure)),
158 };
159
160 (key, aggregate)
161 })
162 .collect();
163
164 Ok(broadcast_via::<K, V, A>(prepared, aggregates, Self::LABEL))
165 }
166
167 fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
168 broadcast_via_estimate(self, &input, stats)
169 }
170}
171
172impl<M, K, J, V, A> GroupKernel<M, K, OperandHandle<Indexed<J, V>, Definite>>
173 for BroadcastViaOperation<A>
174where
175 M: IndexDomain,
176 K: GroupKey,
177 J: IndexDomain,
178 V: ValueDomain,
179 A: IndexedElementSource,
180 A::ValueDomain: GroupingValue<Key = K>,
181{
182 type Output = OperandHandle<Indexed<A::IndexDomain, V>, A::Arity>;
183
184 fn execute<'a>(
185 _graphrecord: &'a GraphRecord,
186 partition: Partition<'a, M, K, OperandHandle<Indexed<J, V>, Definite>>,
187 prepared: Self::Prepared<'a>,
188 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
189 let (buckets, key_failures) = partition.into_parts();
190 reject_key_failures::<M>(key_failures, Self::LABEL)?;
191
192 let aggregates = buckets
193 .into_iter()
194 .map(|(key, _, payload)| {
195 let aggregate = match payload {
196 Ok((_, outcome)) => outcome,
197 Err(failure) => Err(failure),
198 };
199
200 (key, Some(aggregate))
201 })
202 .collect();
203
204 Ok(broadcast_via::<K, V, A>(prepared, aggregates, Self::LABEL))
205 }
206
207 fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
208 broadcast_via_estimate(self, &input, stats)
209 }
210}
211
212impl<M, K, V, A> GroupKernel<M, K, OperandHandle<Bare<V>, Definite>> for BroadcastViaOperation<A>
213where
214 M: IndexDomain,
215 K: GroupKey,
216 V: BareValueDomain,
217 A: IndexedElementSource,
218 A::ValueDomain: GroupingValue<Key = K>,
219{
220 type Output = OperandHandle<Indexed<A::IndexDomain, V>, A::Arity>;
221
222 fn execute<'a>(
223 _graphrecord: &'a GraphRecord,
224 partition: Partition<'a, M, K, OperandHandle<Bare<V>, Definite>>,
225 prepared: Self::Prepared<'a>,
226 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
227 let (buckets, key_failures) = partition.into_parts();
228 reject_key_failures::<M>(key_failures, Self::LABEL)?;
229
230 let aggregates = buckets
231 .into_iter()
232 .map(|(key, _, payload)| {
233 let aggregate = match payload {
234 Ok(outcome) => outcome,
235 Err(failure) => Err(failure),
236 };
237
238 (key, Some(aggregate))
239 })
240 .collect();
241
242 Ok(broadcast_via::<K, V, A>(prepared, aggregates, Self::LABEL))
243 }
244
245 fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
246 broadcast_via_estimate(self, &input, stats)
247 }
248}
249
250impl<O, A> BroadcastVia<A::IndexDomain, A> for O
251where
252 A: IndexedElementSource,
253 BroadcastViaOperation<A>: Operation,
254 O: Apply<BroadcastViaOperation<A>>,
255{
256 type ReturnOperand = O::Output;
257
258 fn broadcast_via(&self, via: A) -> Self::ReturnOperand {
259 Self::ReturnOperand::new(OperationContext::new(
260 self.clone(),
261 BroadcastViaOperation { via },
262 ))
263 }
264}
265
266operation_manifest! {
267 BroadcastViaOperation<A> {
268 method: BroadcastVia<J, A>::broadcast_via;
269 scope: group;
270
271 kernel {
272 group: <M: IndexDomain, K: GroupKey>;
273 parameters: <
274 P: IndexDomain,
275 V: ValueDomain,
276 J: IndexDomain,
277 X: GroupingValue<K>,
278 C: EnumerableArity,
279 >;
280 argument: A: IndexedElementSource<Indexed<J, X>, C>;
281 input: OperandHandle<Indexed<P, V>, Single>;
282 output: OperandHandle<Indexed<J, V>, C>;
283 }
284
285 kernel {
286 group: <M: IndexDomain, K: GroupKey>;
287 parameters: <
288 P: IndexDomain,
289 V: ValueDomain,
290 J: IndexDomain,
291 X: GroupingValue<K>,
292 C: EnumerableArity,
293 >;
294 argument: A: IndexedElementSource<Indexed<J, X>, C>;
295 input: OperandHandle<Indexed<P, V>, Definite>;
296 output: OperandHandle<Indexed<J, V>, C>;
297 }
298
299 kernel {
300 group: <M: IndexDomain, K: GroupKey>;
301 parameters: <
302 V: BareValueDomain,
303 J: IndexDomain,
304 X: GroupingValue<K>,
305 C: EnumerableArity,
306 >;
307 argument: A: IndexedElementSource<Indexed<J, X>, C>;
308 input: OperandHandle<Bare<V>, Single>;
309 output: OperandHandle<Indexed<J, V>, C>;
310 }
311
312 kernel {
313 group: <M: IndexDomain, K: GroupKey>;
314 parameters: <
315 V: BareValueDomain,
316 J: IndexDomain,
317 X: GroupingValue<K>,
318 C: EnumerableArity,
319 >;
320 argument: A: IndexedElementSource<Indexed<J, X>, C>;
321 input: OperandHandle<Bare<V>, Definite>;
322 output: OperandHandle<Indexed<J, V>, C>;
323 }
324 }
325}