1use crate::{
2 Arity, BoxedIterator, Definite, ElementShape, EvaluateOperand, Failure, IndexDomain, Indexed,
3 Multiple, Operand, OperandContext, OrderState, QueryResult, Single, ValueDomain,
4 error::{grouping::InvalidPartitionBucketArity, index::DuplicateIndex},
5 execution::{CacheableOperand, EvaluationCache},
6 index::GroupKey,
7 operands::OperandHandle,
8 optimizer::{Estimate, Estimated, PlanNode, Stats},
9};
10use graphrecords_core::GraphRecord;
11use graphrecords_utils::aliases::{GrHashMap, GrHashSet};
12use std::{marker::PhantomData, sync::Arc};
13
14pub struct GroupOperand<M: IndexDomain, K: GroupKey, O: Operand> {
15 context: Arc<dyn OperandContext<Self>>,
16}
17
18impl<M: IndexDomain, K: GroupKey, O: Operand> Clone for GroupOperand<M, K, O> {
19 fn clone(&self) -> Self {
20 Self {
21 context: Arc::clone(&self.context),
22 }
23 }
24}
25
26impl<M: IndexDomain, K: GroupKey, O: Operand> EvaluateOperand for GroupOperand<M, K, O> {
27 type ReturnValue<'a>
28 = Partition<'a, M, K, O>
29 where
30 Self: 'a;
31
32 fn evaluate<'a>(
33 &'a self,
34 graphrecord: &'a GraphRecord,
35 cache: &'a EvaluationCache<'a>,
36 ) -> QueryResult<Self::ReturnValue<'a>> {
37 self.context.evaluate(graphrecord, cache)
38 }
39}
40
41impl<M: IndexDomain, K: GroupKey, O: Operand> Estimated for GroupOperand<M, K, O> {
42 fn estimate(&self, stats: &Stats) -> Estimate {
43 self.context().estimate(stats)
44 }
45}
46
47impl<M: IndexDomain, K: GroupKey, O: Operand> Operand for GroupOperand<M, K, O> {
48 fn context(&self) -> &dyn OperandContext<Self> {
49 self.context.as_ref()
50 }
51
52 fn as_plan_node(&self) -> &dyn PlanNode {
53 self.context.as_ref()
54 }
55
56 fn from_context(context: Arc<dyn OperandContext<Self>>) -> Self {
57 Self { context }
58 }
59}
60
61pub struct Partition<'a, M: IndexDomain, K: GroupKey, O: Operand> {
62 buckets: Vec<Bucket<'a, M, K, O>>,
63 key_failures: Vec<KeyFailure<'a, M>>,
64}
65
66pub struct Bucket<'a, M: IndexDomain, K: GroupKey, O: Operand> {
67 key: K::Owned,
68 members: Vec<M::Index<'a>>,
69 payload: QueryResult<O::ReturnValue<'a>>,
70}
71
72impl<'a, M: IndexDomain, K: GroupKey, O: Operand> Bucket<'a, M, K, O> {
73 #[must_use]
74 pub const fn key(&self) -> &K::Owned {
75 &self.key
76 }
77
78 #[must_use]
79 pub fn members(&self) -> &[M::Index<'a>] {
80 &self.members
81 }
82
83 pub const fn payload(&self) -> &QueryResult<O::ReturnValue<'a>> {
84 &self.payload
85 }
86}
87
88pub struct KeyFailure<'a, M: IndexDomain> {
89 member: M::Index<'a>,
90 failure: Box<Failure>,
91}
92
93impl<'a, M: IndexDomain> KeyFailure<'a, M> {
94 #[must_use]
95 pub const fn member(&self) -> &M::Index<'a> {
96 &self.member
97 }
98
99 #[must_use]
100 pub fn failure(&self) -> &Failure {
101 &self.failure
102 }
103}
104
105pub type PartitionBucketParts<'a, M, K, O> = (
106 <K as IndexDomain>::Owned,
107 Vec<<M as IndexDomain>::Index<'a>>,
108 QueryResult<<O as EvaluateOperand>::ReturnValue<'a>>,
109);
110pub type PartitionKeyFailureParts<'a, M> = (<M as IndexDomain>::Index<'a>, Box<Failure>);
111pub type PartitionParts<'a, M, K, O> = (
112 Vec<PartitionBucketParts<'a, M, K, O>>,
113 Vec<PartitionKeyFailureParts<'a, M>>,
114);
115
116pub enum BucketChange<'a, O: Operand> {
117 Drop,
118 ReplacePayload(QueryResult<O::ReturnValue<'a>>),
119}
120
121pub enum KeyFailureChange {
122 Drop,
123 Raise,
124}
125
126impl<'a, M: IndexDomain, K: GroupKey, O: Operand> Partition<'a, M, K, O> {
127 #[must_use]
128 pub fn buckets(&self) -> impl ExactSizeIterator<Item = &Bucket<'a, M, K, O>> + '_ {
129 self.buckets.iter()
130 }
131
132 #[must_use]
133 pub fn key_failures(&self) -> impl ExactSizeIterator<Item = &KeyFailure<'a, M>> + '_ {
134 self.key_failures.iter()
135 }
136
137 pub fn map_payloads<N: Operand>(
138 self,
139 mut function: impl FnMut(
140 &K::Owned,
141 &[M::Index<'a>],
142 QueryResult<O::ReturnValue<'a>>,
143 ) -> QueryResult<N::ReturnValue<'a>>,
144 ) -> Partition<'a, M, K, N> {
145 Partition {
146 buckets: self
147 .buckets
148 .into_iter()
149 .map(|bucket| Bucket {
150 payload: function(&bucket.key, &bucket.members, bucket.payload),
151 key: bucket.key,
152 members: bucket.members,
153 })
154 .collect(),
155 key_failures: self.key_failures,
156 }
157 }
158
159 #[cfg(feature = "dynamic")]
160 pub(crate) fn map_domains<N, L>(
161 self,
162 mut map_member: impl FnMut(M::Index<'a>) -> N::Index<'a>,
163 mut map_key: impl FnMut(K::Owned) -> L::Owned,
164 ) -> Partition<'a, N, L, O>
165 where
166 N: IndexDomain,
167 L: GroupKey,
168 {
169 Partition {
170 buckets: self
171 .buckets
172 .into_iter()
173 .map(|bucket| Bucket {
174 key: map_key(bucket.key),
175 members: bucket.members.into_iter().map(&mut map_member).collect(),
176 payload: bucket.payload,
177 })
178 .collect(),
179 key_failures: self
180 .key_failures
181 .into_iter()
182 .map(|key_failure| KeyFailure {
183 member: map_member(key_failure.member),
184 failure: key_failure.failure,
185 })
186 .collect(),
187 }
188 }
189
190 #[must_use]
191 pub fn change_buckets(
192 self,
193 mut function: impl FnMut(&Bucket<'a, M, K, O>) -> Option<BucketChange<'a, O>>,
194 ) -> Self {
195 let mut buckets = Vec::with_capacity(self.buckets.len());
196
197 for mut bucket in self.buckets {
198 let change = function(&bucket);
199
200 match change {
201 None => buckets.push(bucket),
202 Some(BucketChange::Drop) => {}
203 Some(BucketChange::ReplacePayload(payload)) => {
204 bucket.payload = payload;
205 buckets.push(bucket);
206 }
207 }
208 }
209
210 Self {
211 buckets,
212 key_failures: self.key_failures,
213 }
214 }
215
216 pub fn change_key_failures(
217 self,
218 mut function: impl FnMut(&KeyFailure<'a, M>) -> Option<KeyFailureChange>,
219 ) -> QueryResult<Self> {
220 let mut key_failures = Vec::with_capacity(self.key_failures.len());
221
222 for key_failure in self.key_failures {
223 let change = function(&key_failure);
224
225 match change {
226 None => key_failures.push(key_failure),
227 Some(KeyFailureChange::Drop) => {}
228 Some(KeyFailureChange::Raise) => return Err(key_failure.failure),
229 }
230 }
231
232 Ok(Self {
233 buckets: self.buckets,
234 key_failures,
235 })
236 }
237
238 #[must_use]
239 pub fn into_parts(self) -> PartitionParts<'a, M, K, O> {
240 (
241 self.buckets
242 .into_iter()
243 .map(|bucket| (bucket.key, bucket.members, bucket.payload))
244 .collect(),
245 self.key_failures
246 .into_iter()
247 .map(|key_failure| (key_failure.member, key_failure.failure))
248 .collect(),
249 )
250 }
251
252 pub fn into_return_partition<T>(
253 self,
254 mut convert_payload: impl FnMut(QueryResult<O::ReturnValue<'a>>) -> QueryResult<T>,
255 ) -> ReturnPartition<'a, M, K, T> {
256 ReturnPartition {
257 buckets: self
258 .buckets
259 .into_iter()
260 .map(|bucket| ReturnBucket {
261 key: bucket.key,
262 members: bucket.members,
263 payload: convert_payload(bucket.payload),
264 })
265 .collect(),
266 key_failures: self
267 .key_failures
268 .into_iter()
269 .map(|key_failure| ReturnKeyFailure {
270 member: key_failure.member,
271 failure: key_failure.failure,
272 })
273 .collect(),
274 }
275 }
276
277 pub fn into_owned<T>(
278 self,
279 mut convert_payload: impl FnMut(QueryResult<O::ReturnValue<'a>>) -> QueryResult<T>,
280 ) -> PartitionOwned<M, K, T> {
281 PartitionOwned {
282 buckets: self
283 .buckets
284 .into_iter()
285 .map(|bucket| BucketOwned {
286 key: bucket.key,
287 members: bucket
288 .members
289 .into_iter()
290 .map(|member| M::to_owned(&member))
291 .collect(),
292 payload: convert_payload(bucket.payload),
293 })
294 .collect(),
295 key_failures: self
296 .key_failures
297 .into_iter()
298 .map(|key_failure| KeyFailureOwned {
299 member: M::to_owned(&key_failure.member),
300 failure: key_failure.failure,
301 })
302 .collect(),
303 }
304 }
305}
306
307pub struct ReturnPartition<'a, M: IndexDomain, K: GroupKey, T> {
308 buckets: Vec<ReturnBucket<'a, M, K, T>>,
309 key_failures: Vec<ReturnKeyFailure<'a, M>>,
310}
311
312pub struct ReturnBucket<'a, M: IndexDomain, K: GroupKey, T> {
313 key: K::Owned,
314 members: Vec<M::Index<'a>>,
315 payload: QueryResult<T>,
316}
317
318impl<'a, M: IndexDomain, K: GroupKey, T> ReturnBucket<'a, M, K, T> {
319 #[must_use]
320 pub const fn key(&self) -> &K::Owned {
321 &self.key
322 }
323
324 #[must_use]
325 pub fn members(&self) -> &[M::Index<'a>] {
326 &self.members
327 }
328
329 pub const fn payload(&self) -> &QueryResult<T> {
330 &self.payload
331 }
332
333 pub fn into_parts(self) -> (K::Owned, Vec<M::Index<'a>>, QueryResult<T>) {
334 (self.key, self.members, self.payload)
335 }
336}
337
338pub struct ReturnKeyFailure<'a, M: IndexDomain> {
339 member: M::Index<'a>,
340 failure: Box<Failure>,
341}
342
343impl<'a, M: IndexDomain> ReturnKeyFailure<'a, M> {
344 #[must_use]
345 pub const fn member(&self) -> &M::Index<'a> {
346 &self.member
347 }
348
349 #[must_use]
350 pub fn failure(&self) -> &Failure {
351 &self.failure
352 }
353
354 #[must_use]
355 pub fn into_parts(self) -> (M::Index<'a>, Box<Failure>) {
356 (self.member, self.failure)
357 }
358}
359
360pub type ReturnPartitionParts<'a, M, K, T> =
361 (Vec<ReturnBucket<'a, M, K, T>>, Vec<ReturnKeyFailure<'a, M>>);
362
363impl<'a, M: IndexDomain, K: GroupKey, T> ReturnPartition<'a, M, K, T> {
364 #[must_use]
365 pub fn buckets(&self) -> &[ReturnBucket<'a, M, K, T>] {
366 &self.buckets
367 }
368
369 #[must_use]
370 pub fn key_failures(&self) -> &[ReturnKeyFailure<'a, M>] {
371 &self.key_failures
372 }
373
374 #[must_use]
375 pub fn into_parts(self) -> ReturnPartitionParts<'a, M, K, T> {
376 (self.buckets, self.key_failures)
377 }
378}
379
380pub struct PartitionOwned<M: IndexDomain, K: GroupKey, T> {
381 buckets: Vec<BucketOwned<M, K, T>>,
382 key_failures: Vec<KeyFailureOwned<M>>,
383}
384
385pub struct BucketOwned<M: IndexDomain, K: GroupKey, T> {
386 key: K::Owned,
387 members: Vec<M::Owned>,
388 payload: QueryResult<T>,
389}
390
391impl<M: IndexDomain, K: GroupKey, T> BucketOwned<M, K, T> {
392 #[must_use]
393 pub const fn key(&self) -> &K::Owned {
394 &self.key
395 }
396
397 #[must_use]
398 pub fn members(&self) -> &[M::Owned] {
399 &self.members
400 }
401
402 pub const fn payload(&self) -> &QueryResult<T> {
403 &self.payload
404 }
405
406 pub fn into_parts(self) -> (K::Owned, Vec<M::Owned>, QueryResult<T>) {
407 (self.key, self.members, self.payload)
408 }
409}
410
411pub struct KeyFailureOwned<M: IndexDomain> {
412 member: M::Owned,
413 failure: Box<Failure>,
414}
415
416impl<M: IndexDomain> KeyFailureOwned<M> {
417 #[must_use]
418 pub const fn member(&self) -> &M::Owned {
419 &self.member
420 }
421
422 #[must_use]
423 pub fn failure(&self) -> &Failure {
424 &self.failure
425 }
426
427 #[must_use]
428 pub fn into_parts(self) -> (M::Owned, Box<Failure>) {
429 (self.member, self.failure)
430 }
431}
432
433pub type PartitionOwnedParts<M, K, T> = (Vec<BucketOwned<M, K, T>>, Vec<KeyFailureOwned<M>>);
434
435impl<M: IndexDomain, K: GroupKey, T> PartitionOwned<M, K, T> {
436 #[must_use]
437 pub fn buckets(&self) -> &[BucketOwned<M, K, T>] {
438 &self.buckets
439 }
440
441 #[must_use]
442 pub fn key_failures(&self) -> &[KeyFailureOwned<M>] {
443 &self.key_failures
444 }
445
446 #[must_use]
447 pub fn into_parts(self) -> PartitionOwnedParts<M, K, T> {
448 (self.buckets, self.key_failures)
449 }
450}
451
452pub trait PartitionShape<M: IndexDomain>: ElementShape {
453 fn member<'a>(element: &Self::Element<'a>) -> M::Index<'a>;
454}
455
456impl<M: IndexDomain, V: ValueDomain> PartitionShape<M> for Indexed<M, V> {
457 fn member<'a>(element: &Self::Element<'a>) -> M::Index<'a> {
458 element.0.clone()
459 }
460}
461
462pub trait PartitionArity<S: ElementShape>: Arity {
463 fn into_elements<'a>(
464 container: Self::Container<'a, S::Element<'a>>,
465 ) -> BoxedIterator<'a, S::Element<'a>>
466 where
467 S: 'a;
468
469 fn from_bucket<'a>(
470 elements: Vec<S::Element<'a>>,
471 ) -> QueryResult<Self::Container<'a, S::Element<'a>>>
472 where
473 S: 'a;
474}
475
476impl<S: ElementShape> PartitionArity<S> for Definite {
477 fn into_elements<'a>(
478 container: Self::Container<'a, S::Element<'a>>,
479 ) -> BoxedIterator<'a, S::Element<'a>>
480 where
481 S: 'a,
482 {
483 Box::new(std::iter::once(container))
484 }
485
486 fn from_bucket<'a>(
487 elements: Vec<S::Element<'a>>,
488 ) -> QueryResult<Self::Container<'a, S::Element<'a>>>
489 where
490 S: 'a,
491 {
492 match <[S::Element<'a>; 1]>::try_from(elements) {
493 Ok([element]) => Ok(element),
494 Err(elements) => Err(Failure::new(
495 "partition construction",
496 InvalidPartitionBucketArity::new("exactly one", elements.len()),
497 )),
498 }
499 }
500}
501
502impl<S: ElementShape> PartitionArity<S> for Single {
503 fn into_elements<'a>(
504 container: Self::Container<'a, S::Element<'a>>,
505 ) -> BoxedIterator<'a, S::Element<'a>>
506 where
507 S: 'a,
508 {
509 Box::new(container.into_iter())
510 }
511
512 fn from_bucket<'a>(
513 elements: Vec<S::Element<'a>>,
514 ) -> QueryResult<Self::Container<'a, S::Element<'a>>>
515 where
516 S: 'a,
517 {
518 if elements.len() > 1 {
519 return Err(Failure::new(
520 "partition construction",
521 InvalidPartitionBucketArity::new("at most one", elements.len()),
522 ));
523 }
524
525 Ok(elements.into_iter().next())
526 }
527}
528
529impl<S: ElementShape, O: OrderState> PartitionArity<S> for Multiple<O> {
530 fn into_elements<'a>(
531 container: Self::Container<'a, S::Element<'a>>,
532 ) -> BoxedIterator<'a, S::Element<'a>>
533 where
534 S: 'a,
535 {
536 container
537 }
538
539 fn from_bucket<'a>(
540 elements: Vec<S::Element<'a>>,
541 ) -> QueryResult<Self::Container<'a, S::Element<'a>>>
542 where
543 S: 'a,
544 {
545 Ok(Box::new(elements.into_iter()))
546 }
547}
548
549pub enum PartitionClassification<K: GroupKey> {
550 Key(K::Owned),
551 KeyFailure(Box<Failure>),
552 Omit,
553}
554
555pub struct PartitionBuilder<'a, M, K, S, C>
556where
557 M: IndexDomain,
558 K: GroupKey,
559 S: PartitionShape<M>,
560 C: PartitionArity<S>,
561{
562 source: C::Container<'a, S::Element<'a>>,
563 marker: PhantomData<fn() -> (M, K)>,
564}
565
566impl<'a, M, K, S, C> PartitionBuilder<'a, M, K, S, C>
567where
568 M: IndexDomain,
569 K: GroupKey,
570 S: PartitionShape<M>,
571 C: PartitionArity<S>,
572{
573 #[must_use]
574 pub fn new(source: C::Container<'a, S::Element<'a>>) -> Self {
575 Self {
576 source,
577 marker: PhantomData,
578 }
579 }
580
581 pub fn build(
582 self,
583 mut classify: impl FnMut(&S::Element<'a>) -> PartitionClassification<K>,
584 ) -> QueryResult<Partition<'a, M, K, OperandHandle<S, C>>> {
585 let mut seen_members = GrHashSet::default();
586 let mut key_positions: GrHashMap<_, _> = GrHashMap::default();
587 let mut buckets = Vec::new();
588 let mut key_failures = Vec::new();
589
590 for element in C::into_elements(self.source) {
591 let member = S::member(&element);
592
593 if !seen_members.insert(M::to_owned(&member)) {
594 return Err(Failure::new_at::<M, _>(
595 "partition construction",
596 DuplicateIndex::<M>::new(M::to_owned(&member)),
597 &member,
598 ));
599 }
600
601 match classify(&element) {
602 PartitionClassification::Key(key) => {
603 let position = if let Some(position) = key_positions.get(&key) {
604 *position
605 } else {
606 let position = buckets.len();
607 key_positions.insert(key.clone(), position);
608 buckets.push((key, Vec::new(), Vec::new()));
609 position
610 };
611
612 buckets[position].1.push(member);
613 buckets[position].2.push(element);
614 }
615 PartitionClassification::KeyFailure(failure) => {
616 key_failures.push(KeyFailure { member, failure });
617 }
618 PartitionClassification::Omit => {}
619 }
620 }
621
622 let buckets = buckets
623 .into_iter()
624 .map(|(key, members, elements)| {
625 C::from_bucket(elements).map(|payload| Bucket {
626 key,
627 members,
628 payload: Ok(payload),
629 })
630 })
631 .collect::<QueryResult<_>>()?;
632
633 Ok(Partition {
634 buckets,
635 key_failures,
636 })
637 }
638}
639
640impl<M: IndexDomain, K: GroupKey, O: CacheableOperand> CacheableOperand for GroupOperand<M, K, O> {
641 type Cached = PartitionOwned<M, K, O::Cached>;
642
643 fn into_cached(values: Self::ReturnValue<'_>) -> Self::Cached {
644 values.into_owned(|payload| payload.map(O::into_cached))
645 }
646
647 fn from_cached(cached: &Self::Cached) -> Self::ReturnValue<'_> {
648 Partition {
649 buckets: cached
650 .buckets
651 .iter()
652 .map(|bucket| Bucket {
653 key: bucket.key.clone(),
654 members: bucket.members.iter().map(M::from_owned).collect(),
655 payload: match &bucket.payload {
656 Ok(payload) => Ok(O::from_cached(payload)),
657 Err(failure) => Err(failure.clone()),
658 },
659 })
660 .collect(),
661 key_failures: cached
662 .key_failures
663 .iter()
664 .map(|key_failure| KeyFailure {
665 member: M::from_owned(&key_failure.member),
666 failure: key_failure.failure.clone(),
667 })
668 .collect(),
669 }
670 }
671}