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