Skip to main content

graphrecords_query/operations/conversion/
enumerate.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, IndexDomain, Indexed, Multiple, Operand,
3    Ordered, Positional, QueryResult, ValueDomain,
4    execution::EvaluationCache,
5    operands::OperandHandle,
6    operations::{
7        Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
8    },
9    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
10    registry::operation_manifest,
11    traits::Enumerate,
12};
13use graphrecords_core::GraphRecord;
14
15#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
16#[operation(scope = Lane)]
17#[explain(label = "Enumerate")]
18#[plan(optimizer_hints(empty = if_any))]
19pub struct EnumerateOperation;
20
21impl Prepare for EnumerateOperation {
22    type Prepared<'a> = ();
23
24    fn prepare<'a>(
25        &'a self,
26        _graphrecord: &'a GraphRecord,
27        _cache: &'a EvaluationCache<'a>,
28    ) -> QueryResult<Self::Prepared<'a>> {
29        Ok(())
30    }
31}
32
33impl<V: BareValueDomain> LaneKernel<Bare<V>, Multiple<Ordered>> for EnumerateOperation {
34    type Output = OperandHandle<Indexed<Positional, V>, Multiple<Ordered>>;
35
36    fn execute<'a>(
37        _graphrecord: &'a GraphRecord,
38        values: BareStream<'a, V, Multiple<Ordered>>,
39        _prepared: Self::Prepared<'a>,
40    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
41        Ok(Box::new(values.enumerate()))
42    }
43
44    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
45        input
46    }
47}
48
49impl<I: IndexDomain, V: ValueDomain> LaneKernel<Indexed<I, V>, Multiple<Ordered>>
50    for EnumerateOperation
51{
52    type Output = OperandHandle<Indexed<Positional, V>, Multiple<Ordered>>;
53
54    fn execute<'a>(
55        _graphrecord: &'a GraphRecord,
56        values: KeyedStream<'a, I, V, Multiple<Ordered>>,
57        _prepared: Self::Prepared<'a>,
58    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
59        Ok(Box::new(
60            values
61                .enumerate()
62                .map(|(position, (_, value))| (position, value)),
63        ))
64    }
65
66    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
67        input
68    }
69}
70
71impl<O: Apply<EnumerateOperation>> Enumerate for O {
72    type ReturnOperand = O::Output;
73
74    fn enumerate(&self) -> Self::ReturnOperand {
75        Self::ReturnOperand::new(OperationContext::new(self.clone(), EnumerateOperation))
76    }
77}
78
79operation_manifest! {
80    EnumerateOperation {
81        method: Enumerate::enumerate;
82        scope: lane;
83
84        kernel {
85            parameters: <V: BareValueDomain>;
86            input: (Bare<V>, Multiple<Ordered>);
87            output: OperandHandle<Indexed<Positional, V>, Multiple<Ordered>>;
88        }
89        kernel {
90            parameters: <
91                I: IndexDomain,
92                V: ValueDomain,
93            >;
94            input: (Indexed<I, V>, Multiple<Ordered>);
95            output: OperandHandle<Indexed<Positional, V>, Multiple<Ordered>>;
96        }
97    }
98}