graphrecords_query/operations/structure/
in_group.rs1use 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}