Skip to main content

graphrecords_query/operations/uniqueness/
is_duplicated.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, IndexDomain, Indexed, Mask, Multiple, Operand,
3    OrderState, QueryResult,
4    capabilities::ValueEquivalence,
5    execution::EvaluationCache,
6    operands::OperandHandle,
7    operations::{
8        Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
9    },
10    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
11    registry::operation_manifest,
12    traits::IsDuplicated,
13};
14use graphrecords_core::GraphRecord;
15use graphrecords_utils::aliases::GrHashMap;
16
17#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
18#[operation(scope = Lane)]
19#[explain(label = "IsDuplicated")]
20#[plan(optimizer_hints(empty = if_any))]
21pub struct IsDuplicatedOperation;
22
23impl Prepare for IsDuplicatedOperation {
24    type Prepared<'a> = ();
25
26    fn prepare<'a>(
27        &'a self,
28        _graphrecord: &'a GraphRecord,
29        _cache: &'a EvaluationCache<'a>,
30    ) -> QueryResult<Self::Prepared<'a>> {
31        Ok(())
32    }
33}
34
35impl<I: IndexDomain, V: ValueEquivalence, O: OrderState> LaneKernel<Indexed<I, V>, Multiple<O>>
36    for IsDuplicatedOperation
37{
38    type Output = OperandHandle<Indexed<I, Mask>, Multiple<O>>;
39
40    fn execute<'a>(
41        _graphrecord: &'a GraphRecord,
42        values: KeyedStream<'a, I, V, Multiple<O>>,
43        _prepared: Self::Prepared<'a>,
44    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
45        let values: Vec<_> = values.collect();
46        let mut counts: GrHashMap<_, usize> = GrHashMap::default();
47
48        for (_, outcome) in &values {
49            if let Ok(value) = outcome {
50                *counts.entry(V::equivalence_key(value)).or_insert(0) += 1;
51            }
52        }
53
54        Ok(Box::new(values.into_iter().map(move |(index, outcome)| {
55            let outcome = outcome.map(|value| {
56                counts
57                    .get(&V::equivalence_key(&value))
58                    .copied()
59                    .expect("every successful value was counted")
60                    > 1
61            });
62
63            (index, outcome)
64        })))
65    }
66
67    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
68        Estimate {
69            elements: input.elements,
70            distinct: None,
71            selectivity: None,
72            per_group: None,
73        }
74    }
75}
76
77impl<V: ValueEquivalence + BareValueDomain, O: OrderState> LaneKernel<Bare<V>, Multiple<O>>
78    for IsDuplicatedOperation
79{
80    type Output = OperandHandle<Bare<Mask>, Multiple<O>>;
81
82    fn execute<'a>(
83        _graphrecord: &'a GraphRecord,
84        values: BareStream<'a, V, Multiple<O>>,
85        _prepared: Self::Prepared<'a>,
86    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
87        let values: Vec<_> = values.collect();
88        let mut counts: GrHashMap<_, usize> = GrHashMap::default();
89
90        for value in values.iter().flatten() {
91            *counts.entry(V::equivalence_key(value)).or_insert(0) += 1;
92        }
93
94        Ok(Box::new(values.into_iter().map(move |outcome| {
95            outcome.map(|value| {
96                counts
97                    .get(&V::equivalence_key(&value))
98                    .copied()
99                    .expect("every successful value was counted")
100                    > 1
101            })
102        })))
103    }
104
105    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
106        Estimate {
107            elements: input.elements,
108            distinct: None,
109            selectivity: None,
110            per_group: None,
111        }
112    }
113}
114
115impl<O: Apply<IsDuplicatedOperation>> IsDuplicated for O {
116    type ReturnOperand = O::Output;
117
118    fn is_duplicated(&self) -> Self::ReturnOperand {
119        Self::ReturnOperand::new(OperationContext::new(self.clone(), IsDuplicatedOperation))
120    }
121}
122
123operation_manifest! {
124    IsDuplicatedOperation {
125        method: IsDuplicated::is_duplicated;
126        scope: lane;
127
128        kernel {
129            parameters: <
130                I: IndexDomain,
131                V: ValueEquivalence,
132                O: OrderState,
133            >;
134            input: (Indexed<I, V>, Multiple<O>);
135            output: OperandHandle<Indexed<I, Mask>, Multiple<O>>;
136        }
137        kernel {
138            parameters: <
139                V: ValueEquivalence + BareValueDomain,
140                O: OrderState,
141            >;
142            input: (Bare<V>, Multiple<O>);
143            output: OperandHandle<Bare<Mask>, Multiple<O>>;
144        }
145    }
146}