Skip to main content

graphrecords_query/operations/indexing/
select.rs

1use 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}