graphrecords_query/operations/numeric/
clip.rs

1use crate::{
2    Bare, BareValueDomain, Explain, IndexDomain, Indexed, Labeled, Operand, QueryResult,
3    capabilities::ValueClip,
4    element::{Pipeline, Retention},
5    execution::EvaluationCache,
6    operations::{
7        Apply, ArgumentSource, ElementKernel, ElementPipeline, Keyed, Operation, OperationContext,
8        Prepare, Unaligned,
9    },
10    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
11    registry::{describe::ArgumentRetention, operation_manifest},
12    traits::Clip,
13};
14use graphrecords_core::GraphRecord;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Element)]
18#[explain(label = "Clip")]
19#[plan(optimizer_hints(empty = if_all))]
20pub struct ClipOperation<L, U> {
21    #[argument]
22    lower: L,
23    #[argument]
24    upper: U,
25}
26
27impl<L: Prepare, U: Prepare> Prepare for ClipOperation<L, U> {
28    type Prepared<'a>
29        = (L::Prepared<'a>, U::Prepared<'a>)
30    where
31        Self: 'a;
32
33    fn prepare<'a>(
34        &'a self,
35        graphrecord: &'a GraphRecord,
36        cache: &'a EvaluationCache<'a>,
37    ) -> QueryResult<Self::Prepared<'a>> {
38        Ok((
39            self.lower.prepare(graphrecord, cache)?,
40            self.upper.prepare(graphrecord, cache)?,
41        ))
42    }
43}
44
45impl<I, V, L, U> ElementKernel<Indexed<I, V>> for ClipOperation<L, U>
46where
47    I: IndexDomain,
48    V: ValueClip,
49    L: ArgumentSource<Keyed<I>, V>,
50    U: ArgumentSource<Keyed<I>, V>,
51{
52    type Emission = <L::Retention as Retention>::Then<U::Retention>;
53    type OutShape = Indexed<I, V>;
54
55    fn pipeline<'a>(
56        _graphrecord: &'a GraphRecord,
57        prepared: Self::Prepared<'a>,
58    ) -> QueryResult<ElementPipeline<'a, Indexed<I, V>, Self>> {
59        Ok(Pipeline::keyed(move |index, item| {
60            let value = match item {
61                Ok(value) => value,
62                Err(failure) => {
63                    return Self::Emission::keep(Err(failure));
64                }
65            };
66
67            let lower = L::resolve(&prepared.0, &index, Self::LABEL);
68
69            L::Retention::and_then(lower, |lower| {
70                let upper = U::resolve(&prepared.1, &index, Self::LABEL);
71
72                U::Retention::map_step(upper, |upper| {
73                    upper.and_then(|upper| {
74                        V::clip(Self::LABEL, value, lower, upper)
75                            .map_err(|failure| failure.at::<I>(&index))
76                    })
77                })
78            })
79        }))
80    }
81
82    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
83        input.with_unknown_distinct()
84    }
85}
86
87impl<V, L, U> ElementKernel<Bare<V>> for ClipOperation<L, U>
88where
89    V: ValueClip + BareValueDomain,
90    L: ArgumentSource<Unaligned, V>,
91    U: ArgumentSource<Unaligned, V>,
92{
93    type Emission = <L::Retention as Retention>::Then<U::Retention>;
94    type OutShape = Bare<V>;
95
96    fn pipeline<'a>(
97        _graphrecord: &'a GraphRecord,
98        prepared: Self::Prepared<'a>,
99    ) -> QueryResult<ElementPipeline<'a, Bare<V>, Self>> {
100        Ok(Pipeline::new(move |item| {
101            let value = match item {
102                Ok(value) => value,
103                Err(failure) => {
104                    return Self::Emission::keep(Err(failure));
105                }
106            };
107
108            let lower = L::resolve(&prepared.0, &(), Self::LABEL);
109
110            L::Retention::and_then(lower, |lower| {
111                let upper = U::resolve(&prepared.1, &(), Self::LABEL);
112
113                U::Retention::map_step(upper, |upper| {
114                    upper.and_then(|upper| V::clip(Self::LABEL, value, lower, upper))
115                })
116            })
117        }))
118    }
119
120    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
121        input.with_unknown_distinct()
122    }
123}
124
125impl<O, L, U> Clip<L, U> for O
126where
127    ClipOperation<L, U>: Operation,
128    O: Apply<ClipOperation<L, U>>,
129{
130    type ReturnOperand = O::Output;
131
132    fn clip(&self, lower: L, upper: U) -> Self::ReturnOperand {
133        Self::ReturnOperand::new(OperationContext::new(
134            self.clone(),
135            ClipOperation { lower, upper },
136        ))
137    }
138}
139
140operation_manifest! {
141    ClipOperation<L, U> {
142        method: Clip<L, U>::clip;
143        scope: element;
144
145        kernel {
146            parameters: <I: IndexDomain, V: ValueClip>;
147            argument: L: ArgumentSource<Keyed<I>, V>;
148            argument: U: ArgumentSource<Keyed<I>, V>;
149            input: Indexed<I, V>;
150            output: Indexed<I, V>;
151            emission: ArgumentRetention;
152        }
153
154        kernel {
155            parameters: <V: ValueClip + BareValueDomain>;
156            argument: L: ArgumentSource<Unaligned, V>;
157            argument: U: ArgumentSource<Unaligned, V>;
158            input: Bare<V>;
159            output: Bare<V>;
160            emission: ArgumentRetention;
161        }
162    }
163}
graphrecords_query/operations/numeric/clip.rs

graphrecords_query/operations/numeric/
clip.rs
