graphrecords_query/operations/indexing/
select.rs1use crate::{
2 Bare, Definite, EntityDomain, EntityReference, EvaluateOperand, Explain, IndexDomain, Indexed,
3 Multiple, Operand, OrderState, QueryResult, Single, Unordered,
4 execution::EvaluationCache,
5 operands::{DefiniteElementOperand, ElementOperand, ElementsOperand},
6 operations::{
7 Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
8 },
9 optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
10 registry::operation_manifest,
11 traits::Select,
12};
13use graphrecords_core::GraphRecord;
14use graphrecords_utils::aliases::GrHashSet;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Lane)]
18#[explain(label = "Select")]
19#[plan(optimizer_hints(empty = if_any))]
20pub struct SelectOperation;
21
22const fn multiple_estimate(input: &Estimate) -> Estimate {
23 Estimate {
24 elements: input.distinct,
25 distinct: input.distinct,
26 selectivity: None,
27 per_group: None,
28 }
29}
30
31const fn single_estimate(input: &Estimate) -> Estimate {
32 Estimate {
33 elements: input.elements,
34 distinct: input.elements,
35 selectivity: None,
36 per_group: None,
37 }
38}
39
40impl Prepare for SelectOperation {
41 type Prepared<'a> = ();
42
43 fn prepare<'a>(
44 &'a self,
45 _graphrecord: &'a GraphRecord,
46 _cache: &'a EvaluationCache<'a>,
47 ) -> QueryResult<Self::Prepared<'a>> {
48 Ok(())
49 }
50}
51
52impl<E: EntityDomain, I: IndexDomain, O: OrderState>
53 LaneKernel<Indexed<I, EntityReference<E>>, Multiple<O>> for SelectOperation
54{
55 type Output = ElementsOperand<E, Unordered>;
56
57 fn execute<'a>(
58 _graphrecord: &'a GraphRecord,
59 values: KeyedStream<'a, I, EntityReference<E>, Multiple<O>>,
60 _prepared: Self::Prepared<'a>,
61 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
62 let targets: GrHashSet<_> = values
63 .map(|(_, reference)| reference)
64 .collect::<QueryResult<_>>()?;
65
66 Ok(Box::new(targets.into_iter().map(|target| (target, Ok(())))))
67 }
68
69 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
70 multiple_estimate(&input)
71 }
72}
73
74impl<E: EntityDomain, I: IndexDomain> LaneKernel<Indexed<I, EntityReference<E>>, Single>
75 for SelectOperation
76{
77 type Output = ElementOperand<E>;
78
79 fn execute<'a>(
80 _graphrecord: &'a GraphRecord,
81 value: KeyedStream<'a, I, EntityReference<E>, Single>,
82 _prepared: Self::Prepared<'a>,
83 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
84 let Some((_, reference)) = value else {
85 return Ok(None);
86 };
87 let target = reference?;
88
89 Ok(Some((target, Ok(()))))
90 }
91
92 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
93 single_estimate(&input)
94 }
95}
96
97impl<E: EntityDomain, I: IndexDomain> LaneKernel<Indexed<I, EntityReference<E>>, Definite>
98 for SelectOperation
99{
100 type Output = DefiniteElementOperand<E>;
101
102 fn execute<'a>(
103 _graphrecord: &'a GraphRecord,
104 value: KeyedStream<'a, I, EntityReference<E>, Definite>,
105 _prepared: Self::Prepared<'a>,
106 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
107 let target = value.1?;
108
109 Ok((target, Ok(())))
110 }
111
112 fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
113 Estimate::singleton()
114 }
115}
116
117impl<E: EntityDomain, O: OrderState> LaneKernel<Bare<EntityReference<E>>, Multiple<O>>
118 for SelectOperation
119{
120 type Output = ElementsOperand<E, Unordered>;
121
122 fn execute<'a>(
123 _graphrecord: &'a GraphRecord,
124 values: BareStream<'a, EntityReference<E>, Multiple<O>>,
125 _prepared: Self::Prepared<'a>,
126 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
127 let targets: GrHashSet<_> = values.collect::<QueryResult<_>>()?;
128
129 Ok(Box::new(targets.into_iter().map(|target| (target, Ok(())))))
130 }
131
132 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
133 multiple_estimate(&input)
134 }
135}
136
137impl<E: EntityDomain> LaneKernel<Bare<EntityReference<E>>, Single> for SelectOperation {
138 type Output = ElementOperand<E>;
139
140 fn execute<'a>(
141 _graphrecord: &'a GraphRecord,
142 value: BareStream<'a, EntityReference<E>, Single>,
143 _prepared: Self::Prepared<'a>,
144 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
145 Ok(value.transpose()?.map(|target| (target, Ok(()))))
146 }
147
148 fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
149 single_estimate(&input)
150 }
151}
152
153impl<E: EntityDomain> LaneKernel<Bare<EntityReference<E>>, Definite> for SelectOperation {
154 type Output = DefiniteElementOperand<E>;
155
156 fn execute<'a>(
157 _graphrecord: &'a GraphRecord,
158 value: BareStream<'a, EntityReference<E>, Definite>,
159 _prepared: Self::Prepared<'a>,
160 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
161 Ok((value?, Ok(())))
162 }
163
164 fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
165 Estimate::singleton()
166 }
167}
168
169impl<O: Apply<SelectOperation>> Select for O {
170 type ReturnOperand = O::Output;
171
172 fn select(&self) -> Self::ReturnOperand {
173 Self::ReturnOperand::new(OperationContext::new(self.clone(), SelectOperation))
174 }
175}
176
177operation_manifest! {
178 SelectOperation {
179 method: Select::select;
180 scope: lane;
181
182 kernel {
183 parameters: <E: EntityDomain, I: IndexDomain, O: OrderState>;
184 input: (Indexed<I, EntityReference<E>>, Multiple<O>);
185 output: ElementsOperand<E, Unordered>;
186 }
187 kernel {
188 parameters: <E: EntityDomain, I: IndexDomain>;
189 input: (Indexed<I, EntityReference<E>>, Single);
190 output: ElementOperand<E>;
191 }
192 kernel {
193 parameters: <E: EntityDomain, I: IndexDomain>;
194 input: (Indexed<I, EntityReference<E>>, Definite);
195 output: DefiniteElementOperand<E>;
196 }
197 kernel {
198 parameters: <E: EntityDomain, O: OrderState>;
199 input: (Bare<EntityReference<E>>, Multiple<O>);
200 output: ElementsOperand<E, Unordered>;
201 }
202 kernel {
203 parameters: <E: EntityDomain>;
204 input: (Bare<EntityReference<E>>, Single);
205 output: ElementOperand<E>;
206 }
207 kernel {
208 parameters: <E: EntityDomain>;
209 input: (Bare<EntityReference<E>>, Definite);
210 output: DefiniteElementOperand<E>;
211 }
212 }
213}