Skip to main content

graphrecords_query/operations/logic/
or.rs

1use super::{combine_masks_bare, combine_masks_indexed};
2use crate::{
3    Arity, Bare, ElementShape, Explain, IndexDomain, Indexed, Labeled, Mask, Operand, Or,
4    QueryResult,
5    execution::EvaluationCache,
6    operands::OperandHandle,
7    operations::{
8        Apply, ArgumentSource, ElementKernel, ElementPipeline, Keyed, Operation, OperationContext,
9        Prepare, Unaligned,
10    },
11    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
12    registry::{describe::ArgumentRetention, operation_manifest},
13};
14use graphrecords_core::GraphRecord;
15use std::ops::BitOr;
16
17#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
18#[operation(scope = Element)]
19#[explain(label = "Or")]
20#[plan(optimizer_hints(empty = if_all))]
21pub struct OrOperation<M> {
22    #[argument]
23    other: M,
24}
25
26impl<M: Prepare> Prepare for OrOperation<M> {
27    type Prepared<'a>
28        = M::Prepared<'a>
29    where
30        Self: 'a;
31
32    fn prepare<'a>(
33        &'a self,
34        graphrecord: &'a GraphRecord,
35        cache: &'a EvaluationCache<'a>,
36    ) -> QueryResult<Self::Prepared<'a>> {
37        self.other.prepare(graphrecord, cache)
38    }
39}
40
41impl<I, M> ElementKernel<Indexed<I, Mask>> for OrOperation<M>
42where
43    I: IndexDomain,
44    M: ArgumentSource<Keyed<I>, Mask>,
45{
46    type Emission = M::Retention;
47    type OutShape = Indexed<I, Mask>;
48
49    fn pipeline<'a>(
50        _graphrecord: &'a GraphRecord,
51        prepared: Self::Prepared<'a>,
52    ) -> QueryResult<ElementPipeline<'a, Indexed<I, Mask>, Self>> {
53        Ok(combine_masks_indexed::<_, M>(
54            prepared,
55            Self::LABEL,
56            |left, right| left || right,
57        ))
58    }
59
60    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
61        let selectivity = input
62            .selectivity
63            .zip(self.other.estimate(stats).selectivity)
64            .map(|(left, right)| left.mul_add(-right, left + right));
65
66        Estimate {
67            selectivity,
68            ..input
69        }
70    }
71}
72
73impl<M> ElementKernel<Bare<Mask>> for OrOperation<M>
74where
75    M: ArgumentSource<Unaligned, Mask>,
76{
77    type Emission = M::Retention;
78    type OutShape = Bare<Mask>;
79
80    fn pipeline<'a>(
81        _graphrecord: &'a GraphRecord,
82        prepared: Self::Prepared<'a>,
83    ) -> QueryResult<ElementPipeline<'a, Bare<Mask>, Self>> {
84        Ok(combine_masks_bare::<M>(
85            prepared,
86            Self::LABEL,
87            |left, right| left || right,
88        ))
89    }
90
91    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
92        let selectivity = input
93            .selectivity
94            .zip(self.other.estimate(stats).selectivity)
95            .map(|(left, right)| left.mul_add(-right, left + right));
96
97        Estimate {
98            selectivity,
99            ..input
100        }
101    }
102}
103
104impl<O, M> Or<M> for O
105where
106    OrOperation<M>: Operation,
107    O: Apply<OrOperation<M>>,
108{
109    type ReturnOperand = O::Output;
110
111    fn or(&self, other: M) -> Self::ReturnOperand {
112        Self::ReturnOperand::new(OperationContext::new(self.clone(), OrOperation { other }))
113    }
114}
115
116impl<S, C, M> BitOr<M> for OperandHandle<S, C>
117where
118    S: ElementShape,
119    C: Arity,
120    Self: Or<M>,
121{
122    type Output = <Self as Or<M>>::ReturnOperand;
123
124    fn bitor(self, rhs: M) -> Self::Output {
125        self.or(rhs)
126    }
127}
128
129operation_manifest! {
130    OrOperation<M> {
131        method: Or<M>::or;
132        scope: element;
133
134        kernel {
135            parameters: <I: IndexDomain>;
136            argument: M: ArgumentSource<Keyed<I>, Mask>;
137            input: Indexed<I, Mask>;
138            output: Indexed<I, Mask>;
139            emission: ArgumentRetention;
140        }
141
142        kernel {
143            parameters: <>;
144            argument: M: ArgumentSource<Unaligned, Mask>;
145            input: Bare<Mask>;
146            output: Bare<Mask>;
147            emission: ArgumentRetention;
148        }
149    }
150}