graphrecords_query/operations/membership/
is_in.rs

1use crate::{
2    Bare, BareValueDomain, Explain, IndexDomain, Indexed, Mask, Operand, QueryResult,
3    capabilities::ValueEquality,
4    element::{Pipeline, Preserving},
5    execution::EvaluationCache,
6    operations::{
7        Apply, ElementKernel, ElementPipeline, Operation, OperationContext, Prepare, SetSource,
8    },
9    optimizer::{
10        Estimate, Estimated, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats,
11    },
12    registry::operation_manifest,
13    traits::IsIn,
14};
15use graphrecords_core::GraphRecord;
16use std::hash::Hash;
17
18#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
19#[operation(scope = Element)]
20#[explain(label = "IsIn")]
21#[plan(optimizer_hints(empty = if_all))]
22pub struct IsInOperation<A> {
23    #[argument]
24    argument: A,
25}
26
27impl<A: Prepare> Prepare for IsInOperation<A> {
28    type Prepared<'a>
29        = A::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        self.argument.prepare(graphrecord, cache)
39    }
40}
41
42fn membership_estimate<A: Estimated>(
43    operation: &IsInOperation<A>,
44    input: Estimate,
45    stats: &Stats,
46) -> Estimate {
47    membership_estimate_from(input, &operation.argument.estimate(stats))
48}
49
50fn membership_estimate_from(input: Estimate, set: &Estimate) -> Estimate {
51    let selectivity = input
52        .distinct
53        .zip(set.elements)
54        .map(|(distinct, size)| (size as f64 / distinct.max(1) as f64).min(1.0));
55
56    Estimate {
57        selectivity,
58        ..input.with_unknown_distinct()
59    }
60}
61
62impl<I, V, A> ElementKernel<Indexed<I, V>> for IsInOperation<A>
63where
64    I: IndexDomain,
65    V: ValueEquality,
66    A: SetSource<V>,
67    for<'a> V::Value<'a>: Eq + Hash,
68{
69    type Emission = Preserving;
70    type OutShape = Indexed<I, Mask>;
71
72    fn pipeline<'a>(
73        _graphrecord: &'a GraphRecord,
74        prepared: Self::Prepared<'a>,
75    ) -> QueryResult<ElementPipeline<'a, Indexed<I, V>, Self>> {
76        let set = A::set(prepared)?;
77
78        Ok(Pipeline::unkeyed(move |outcome: QueryResult<_>| {
79            outcome.map(|value| set.contains(&value))
80        }))
81    }
82
83    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
84        membership_estimate(self, input, stats)
85    }
86}
87
88impl<V, A> ElementKernel<Bare<V>> for IsInOperation<A>
89where
90    V: ValueEquality + BareValueDomain,
91    A: SetSource<V>,
92    for<'a> V::Value<'a>: Eq + Hash,
93{
94    type Emission = Preserving;
95    type OutShape = Bare<Mask>;
96
97    fn pipeline<'a>(
98        _graphrecord: &'a GraphRecord,
99        prepared: Self::Prepared<'a>,
100    ) -> QueryResult<ElementPipeline<'a, Bare<V>, Self>> {
101        let set = A::set(prepared)?;
102
103        Ok(Pipeline::new(move |outcome: QueryResult<_>| {
104            outcome.map(|value| set.contains(&value))
105        }))
106    }
107
108    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
109        membership_estimate(self, input, stats)
110    }
111}
112
113impl<O, A> IsIn<A> for O
114where
115    IsInOperation<A>: Operation,
116    O: Apply<IsInOperation<A>>,
117{
118    type ReturnOperand = O::Output;
119
120    fn is_in(&self, argument: A) -> Self::ReturnOperand {
121        Self::ReturnOperand::new(OperationContext::new(
122            self.clone(),
123            IsInOperation { argument },
124        ))
125    }
126}
127
128operation_manifest! {
129    IsInOperation<A> {
130        method: IsIn<A>::is_in;
131        scope: element;
132
133        kernel {
134            parameters: <I: IndexDomain, V: ValueEquality>;
135            argument: A: SetSource<V>;
136            input: Indexed<I, V>;
137            output: Indexed<I, Mask>;
138            emission: Preserving;
139        }
140
141        kernel {
142            parameters: <V: ValueEquality + BareValueDomain>;
143            argument: A: SetSource<V>;
144            input: Bare<V>;
145            output: Bare<Mask>;
146            emission: Preserving;
147        }
148    }
149}
graphrecords_query/operations/membership/is_in.rs

graphrecords_query/operations/membership/
is_in.rs
