Skip to main content

graphrecords_query/operations/grouping/
broadcast.rs

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}