Skip to main content

graphrecords_query/operations/arithmetic/
power.rs

1use super::{arithmetic_bare, arithmetic_indexed};
2use crate::{
3    Bare, BareValueDomain, Explain, IndexDomain, Indexed, Labeled, Operand, QueryResult,
4    capabilities::ValuePower,
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::Power,
13};
14use graphrecords_core::GraphRecord;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Element)]
18#[explain(label = "Power")]
19#[plan(optimizer_hints(empty = if_all))]
20pub struct PowerOperation<A> {
21    #[argument]
22    argument: A,
23}
24
25impl<A: Prepare> Prepare for PowerOperation<A> {
26    type Prepared<'a>
27        = A::Prepared<'a>
28    where
29        Self: 'a;
30
31    fn prepare<'a>(
32        &'a self,
33        graphrecord: &'a GraphRecord,
34        cache: &'a EvaluationCache<'a>,
35    ) -> QueryResult<Self::Prepared<'a>> {
36        self.argument.prepare(graphrecord, cache)
37    }
38}
39
40impl<I, V, A> ElementKernel<Indexed<I, V>> for PowerOperation<A>
41where
42    I: IndexDomain,
43    V: ValuePower,
44    A: ArgumentSource<Keyed<I>, V>,
45{
46    type Emission = A::Retention;
47    type OutShape = Indexed<I, V>;
48
49    fn pipeline<'a>(
50        _graphrecord: &'a GraphRecord,
51        prepared: Self::Prepared<'a>,
52    ) -> QueryResult<ElementPipeline<'a, Indexed<I, V>, Self>> {
53        Ok(arithmetic_indexed::<_, V, A>(
54            prepared,
55            Self::LABEL,
56            V::power,
57        ))
58    }
59
60    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
61        input.with_unknown_distinct()
62    }
63}
64
65impl<V, A> ElementKernel<Bare<V>> for PowerOperation<A>
66where
67    V: ValuePower + BareValueDomain,
68    A: ArgumentSource<Unaligned, V>,
69{
70    type Emission = A::Retention;
71    type OutShape = Bare<V>;
72
73    fn pipeline<'a>(
74        _graphrecord: &'a GraphRecord,
75        prepared: Self::Prepared<'a>,
76    ) -> QueryResult<ElementPipeline<'a, Bare<V>, Self>> {
77        Ok(arithmetic_bare::<V, A>(prepared, Self::LABEL, V::power))
78    }
79
80    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
81        input.with_unknown_distinct()
82    }
83}
84
85impl<O, A> Power<A> for O
86where
87    PowerOperation<A>: Operation,
88    O: Apply<PowerOperation<A>>,
89{
90    type ReturnOperand = O::Output;
91
92    fn power(&self, argument: A) -> Self::ReturnOperand {
93        Self::ReturnOperand::new(OperationContext::new(
94            self.clone(),
95            PowerOperation { argument },
96        ))
97    }
98}
99
100operation_manifest! {
101    PowerOperation<A> {
102        method: Power<A>::power;
103        scope: element;
104
105        kernel {
106            parameters: <I: IndexDomain, V: ValuePower>;
107            argument: A: ArgumentSource<Keyed<I>, V>;
108            input: Indexed<I, V>;
109            output: Indexed<I, V>;
110            emission: ArgumentRetention;
111        }
112
113        kernel {
114            parameters: <V: ValuePower + BareValueDomain>;
115            argument: A: ArgumentSource<Unaligned, V>;
116            input: Bare<V>;
117            output: Bare<V>;
118            emission: ArgumentRetention;
119        }
120    }
121}