graphrecords_query/operations/grouping/
ungroup.rs

1use super::reject_key_failures;
2use crate::{
3    Bare, BareValueDomain, Definite, EvaluateOperand, Explain, Failure, IndexDomain, Indexed,
4    Labeled, Multiple, Operand, QueryResult, Single, Unordered, ValueDomain,
5    error::grouping::UnresolvedBucketFailures,
6    execution::EvaluationCache,
7    index::GroupKey,
8    operands::{CheckedIndexedLaneBuilder, OperandHandle, Partition, PartitionArity},
9    operations::{Apply, GroupKernel, Operation, OperationContext, Prepare},
10    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
11    registry::operation_manifest,
12    traits::Ungroup,
13};
14use graphrecords_core::GraphRecord;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Group)]
18#[explain(label = "Ungroup")]
19#[plan(optimizer_hints(empty = if_any))]
20pub struct UngroupOperation;
21
22fn ungroup_estimate(input: &Estimate) -> Estimate {
23    let per_group = input.per_group.as_deref();
24    let elements = match (input.elements, per_group.and_then(|inner| inner.elements)) {
25        (Some(groups), Some(elements_per_group)) => groups.checked_mul(elements_per_group),
26        _ => None,
27    };
28    let distinct = match (input.elements, per_group.and_then(|inner| inner.distinct)) {
29        (Some(groups), Some(distinct_per_group)) => groups.checked_mul(distinct_per_group),
30        _ => None,
31    };
32
33    Estimate {
34        elements,
35        distinct: match (distinct, elements) {
36            (Some(distinct), Some(elements)) => Some(distinct.min(elements)),
37            (distinct, _) => distinct,
38        },
39        selectivity: per_group.and_then(|inner| inner.selectivity),
40        per_group: None,
41    }
42}
43
44impl Prepare for UngroupOperation {
45    type Prepared<'a> = ();
46
47    fn prepare<'a>(
48        &'a self,
49        _graphrecord: &'a GraphRecord,
50        _cache: &'a EvaluationCache<'a>,
51    ) -> QueryResult<Self::Prepared<'a>> {
52        Ok(())
53    }
54}
55
56impl<M: IndexDomain, K: GroupKey, I: IndexDomain, V: ValueDomain, C: PartitionArity<Indexed<I, V>>>
57    GroupKernel<M, K, OperandHandle<Indexed<I, V>, C>> for UngroupOperation
58{
59    type Output = OperandHandle<Indexed<I, V>, Multiple<Unordered>>;
60
61    fn execute<'a>(
62        _graphrecord: &'a GraphRecord,
63        partition: Partition<'a, M, K, OperandHandle<Indexed<I, V>, C>>,
64        _prepared: Self::Prepared<'a>,
65    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
66        let (buckets, key_failures) = partition.into_parts();
67
68        reject_key_failures::<M>(key_failures, Self::LABEL)?;
69
70        let mut output = CheckedIndexedLaneBuilder::<I, V>::new();
71        let mut bucket_failures = Vec::new();
72
73        for bucket in buckets {
74            match bucket.2 {
75                Ok(payload) => {
76                    for (index, outcome) in C::into_elements(payload) {
77                        output.push(index, outcome)?;
78                    }
79                }
80                Err(failure) => bucket_failures.push(*failure),
81            }
82        }
83
84        if !bucket_failures.is_empty() {
85            return Err(Failure::new(
86                Self::LABEL,
87                UnresolvedBucketFailures::new(bucket_failures),
88            ));
89        }
90
91        Ok(output.finish())
92    }
93
94    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
95        ungroup_estimate(&input)
96    }
97}
98
99impl<M: IndexDomain, K: GroupKey, V: BareValueDomain, C: PartitionArity<Bare<V>>>
100    GroupKernel<M, K, OperandHandle<Bare<V>, C>> for UngroupOperation
101{
102    type Output = OperandHandle<Bare<V>, Multiple<Unordered>>;
103
104    fn execute<'a>(
105        _graphrecord: &'a GraphRecord,
106        partition: Partition<'a, M, K, OperandHandle<Bare<V>, C>>,
107        _prepared: Self::Prepared<'a>,
108    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
109        let (buckets, key_failures) = partition.into_parts();
110
111        reject_key_failures::<M>(key_failures, Self::LABEL)?;
112
113        let mut payloads = Vec::with_capacity(buckets.len());
114        let mut bucket_failures = Vec::new();
115
116        for bucket in buckets {
117            match bucket.2 {
118                Ok(payload) => payloads.push(payload),
119                Err(failure) => bucket_failures.push(*failure),
120            }
121        }
122
123        if !bucket_failures.is_empty() {
124            return Err(Failure::new(
125                Self::LABEL,
126                UnresolvedBucketFailures::new(bucket_failures),
127            ));
128        }
129
130        Ok(Box::new(payloads.into_iter().flat_map(C::into_elements)))
131    }
132
133    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
134        ungroup_estimate(&input)
135    }
136}
137
138impl<O: Apply<UngroupOperation>> Ungroup for O {
139    type ReturnOperand = O::Output;
140
141    fn ungroup(&self) -> Self::ReturnOperand {
142        Self::ReturnOperand::new(OperationContext::new(self.clone(), UngroupOperation))
143    }
144}
145
146operation_manifest! {
147    UngroupOperation {
148        method: Ungroup::ungroup;
149        scope: group;
150
151        kernel {
152            group: <M: IndexDomain, K: GroupKey>;
153            parameters: <I: IndexDomain, V: ValueDomain, O: OrderState>;
154            input: OperandHandle<Indexed<I, V>, Multiple<O>>;
155            output: OperandHandle<Indexed<I, V>, Multiple<Unordered>>;
156        }
157        kernel {
158            group: <M: IndexDomain, K: GroupKey>;
159            parameters: <I: IndexDomain, V: ValueDomain>;
160            input: OperandHandle<Indexed<I, V>, Single>;
161            output: OperandHandle<Indexed<I, V>, Multiple<Unordered>>;
162        }
163        kernel {
164            group: <M: IndexDomain, K: GroupKey>;
165            parameters: <I: IndexDomain, V: ValueDomain>;
166            input: OperandHandle<Indexed<I, V>, Definite>;
167            output: OperandHandle<Indexed<I, V>, Multiple<Unordered>>;
168        }
169        kernel {
170            group: <M: IndexDomain, K: GroupKey>;
171            parameters: <V: BareValueDomain, O: OrderState>;
172            input: OperandHandle<Bare<V>, Multiple<O>>;
173            output: OperandHandle<Bare<V>, Multiple<Unordered>>;
174        }
175        kernel {
176            group: <M: IndexDomain, K: GroupKey>;
177            parameters: <V: BareValueDomain>;
178            input: OperandHandle<Bare<V>, Single>;
179            output: OperandHandle<Bare<V>, Multiple<Unordered>>;
180        }
181        kernel {
182            group: <M: IndexDomain, K: GroupKey>;
183            parameters: <V: BareValueDomain>;
184            input: OperandHandle<Bare<V>, Definite>;
185            output: OperandHandle<Bare<V>, Multiple<Unordered>>;
186        }
187    }
188}
graphrecords_query/operations/grouping/ungroup.rs

graphrecords_query/operations/grouping/
ungroup.rs
