Skip to main content

graphrecords_query/operations/aggregation/
random.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, IndexDomain, Indexed, Multiple, Operand,
3    OrderState, QueryResult, Single, 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::Random,
12};
13use graphrecords_core::GraphRecord;
14use rand::seq::IteratorRandom;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Lane)]
18#[explain(label = "Random")]
19#[plan(optimizer_hints(volatile, empty = if_any))]
20pub struct RandomOperation;
21
22impl Prepare for RandomOperation {
23    type Prepared<'a> = ();
24
25    fn prepare<'a>(
26        &'a self,
27        _graphrecord: &'a GraphRecord,
28        _cache: &'a EvaluationCache<'a>,
29    ) -> QueryResult<Self::Prepared<'a>> {
30        Ok(())
31    }
32}
33
34impl<I, V, O> LaneKernel<Indexed<I, V>, Multiple<O>> for RandomOperation
35where
36    I: IndexDomain,
37    V: ValueDomain,
38    O: OrderState,
39{
40    type Output = OperandHandle<Indexed<I, V>, Single>;
41
42    fn execute<'a>(
43        _graphrecord: &'a GraphRecord,
44        values: KeyedStream<'a, I, V, Multiple<O>>,
45        _prepared: Self::Prepared<'a>,
46    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
47        Ok(values.choose(&mut rand::rng()))
48    }
49
50    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
51        input.zero_or_one()
52    }
53}
54
55impl<V, O> LaneKernel<Bare<V>, Multiple<O>> for RandomOperation
56where
57    V: BareValueDomain,
58    O: OrderState,
59{
60    type Output = OperandHandle<Bare<V>, Single>;
61
62    fn execute<'a>(
63        _graphrecord: &'a GraphRecord,
64        values: BareStream<'a, V, Multiple<O>>,
65        _prepared: Self::Prepared<'a>,
66    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
67        Ok(values.choose(&mut rand::rng()))
68    }
69
70    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
71        input.zero_or_one()
72    }
73}
74
75impl<O: Apply<RandomOperation>> Random for O {
76    type ReturnOperand = O::Output;
77
78    fn random(&self) -> Self::ReturnOperand {
79        Self::ReturnOperand::new(OperationContext::new(self.clone(), RandomOperation))
80    }
81}
82
83operation_manifest! {
84    RandomOperation {
85        method: Random::random;
86        scope: lane;
87
88        kernel {
89            parameters: <
90                I: IndexDomain,
91                V: ValueDomain,
92                O: OrderState,
93            >;
94            input: (Indexed<I, V>, Multiple<O>);
95            output: OperandHandle<Indexed<I, V>, Single>;
96        }
97
98        kernel {
99            parameters: <
100                V: BareValueDomain,
101                O: OrderState,
102            >;
103            input: (Bare<V>, Multiple<O>);
104            output: OperandHandle<Bare<V>, Single>;
105        }
106    }
107}