Skip to main content

graphrecords_query/operations/structure/
in_group.rs

1use crate::{
2    EntityReference, Explain, IndexDomain, Indexed, Labeled, Mask, Operand, QueryResult, Unit,
3    element::{Pipeline, Preserving},
4    execution::EvaluationCache,
5    index::IndicesInGroup,
6    operations::{Apply, ElementKernel, ElementPipeline, Operation, OperationContext, Prepare},
7    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
8    registry::operation_manifest,
9    traits::InGroup,
10};
11use graphrecords_core::{GraphRecord, graphrecord::Group};
12
13#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
14#[operation(scope = Element)]
15#[explain(label = "InGroup")]
16#[plan(optimizer_hints(commutes_with_filter, allows_limit_pushdown, empty = if_any))]
17pub struct InGroupOperation {
18    #[explain(label)]
19    group: Group,
20}
21
22impl Prepare for InGroupOperation {
23    type Prepared<'a> = &'a Group;
24
25    fn prepare<'a>(
26        &'a self,
27        _graphrecord: &'a GraphRecord,
28        _cache: &'a EvaluationCache<'a>,
29    ) -> QueryResult<Self::Prepared<'a>> {
30        Ok(&self.group)
31    }
32}
33
34impl<I: IndicesInGroup> ElementKernel<Indexed<I, Unit>> for InGroupOperation {
35    type Emission = Preserving;
36    type OutShape = Indexed<I, Mask>;
37
38    fn pipeline<'a>(
39        graphrecord: &'a GraphRecord,
40        prepared: Self::Prepared<'a>,
41    ) -> QueryResult<ElementPipeline<'a, Indexed<I, Unit>, Self>> {
42        let members = I::indices_in_group(Self::LABEL, graphrecord, prepared)?;
43
44        Ok(Pipeline::keyed(move |index, membership: QueryResult<_>| {
45            membership.map(|()| members.contains(&index))
46        }))
47    }
48
49    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
50        let size = I::group_size(stats, &self.group);
51        let selectivity = input
52            .elements
53            .map(|elements| size.min(elements) as f64 / elements.max(1) as f64);
54
55        Estimate {
56            selectivity,
57            ..input
58        }
59    }
60}
61
62impl<E: IndicesInGroup, I: IndexDomain> ElementKernel<Indexed<I, EntityReference<E>>>
63    for InGroupOperation
64{
65    type Emission = Preserving;
66    type OutShape = Indexed<I, Mask>;
67
68    fn pipeline<'a>(
69        graphrecord: &'a GraphRecord,
70        prepared: Self::Prepared<'a>,
71    ) -> QueryResult<ElementPipeline<'a, Indexed<I, EntityReference<E>>, Self>> {
72        let members = E::indices_in_group(Self::LABEL, graphrecord, prepared)?;
73
74        Ok(Pipeline::unkeyed(move |reference: QueryResult<_>| {
75            reference.map(|entity| members.contains(&entity))
76        }))
77    }
78
79    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
80        Estimate {
81            selectivity: None,
82            ..input
83        }
84    }
85}
86
87impl<O: Apply<InGroupOperation>> InGroup for O {
88    type ReturnOperand = O::Output;
89
90    fn in_group(&self, group: Group) -> Self::ReturnOperand {
91        Self::ReturnOperand::new(OperationContext::new(
92            self.clone(),
93            InGroupOperation { group },
94        ))
95    }
96}
97
98operation_manifest! {
99    InGroupOperation {
100        method: InGroup::in_group;
101        scope: element;
102
103        kernel {
104            parameters: <I: IndicesInGroup>;
105            field: group: Group;
106            input: Indexed<I, Unit>;
107            output: Indexed<I, Mask>;
108            emission: Preserving;
109        }
110
111        kernel {
112            parameters: <E: IndicesInGroup, I: IndexDomain>;
113            field: group: Group;
114            input: Indexed<I, EntityReference<E>>;
115            output: Indexed<I, Mask>;
116            emission: Preserving;
117        }
118    }
119}