graphrecords_query/operations/ordering/
shuffle.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, IndexDomain, Indexed, Multiple, Operand,
3    OrderState, Ordered, 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::Shuffle,
12};
13use graphrecords_core::GraphRecord;
14use rand::seq::SliceRandom;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Lane)]
18#[explain(label = "Shuffle")]
19#[plan(optimizer_hints(volatile, empty = if_any))]
20pub struct ShuffleOperation;
21
22impl Prepare for ShuffleOperation {
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: IndexDomain, V: ValueDomain, O: OrderState> LaneKernel<Indexed<I, V>, Multiple<O>>
35    for ShuffleOperation
36{
37    type Output = OperandHandle<Indexed<I, V>, Multiple<Ordered>>;
38
39    fn execute<'a>(
40        _graphrecord: &'a GraphRecord,
41        values: KeyedStream<'a, I, V, Multiple<O>>,
42        _prepared: Self::Prepared<'a>,
43    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
44        let mut values: Vec<_> = values.collect();
45        values.shuffle(&mut rand::rng());
46
47        Ok(Box::new(values.into_iter()))
48    }
49
50    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
51        input
52    }
53}
54
55impl<V: BareValueDomain, O: OrderState> LaneKernel<Bare<V>, Multiple<O>> for ShuffleOperation {
56    type Output = OperandHandle<Bare<V>, Multiple<Ordered>>;
57
58    fn execute<'a>(
59        _graphrecord: &'a GraphRecord,
60        values: BareStream<'a, V, Multiple<O>>,
61        _prepared: Self::Prepared<'a>,
62    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
63        let mut values: Vec<_> = values.collect();
64        values.shuffle(&mut rand::rng());
65
66        Ok(Box::new(values.into_iter()))
67    }
68
69    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
70        input
71    }
72}
73
74impl<O: Apply<ShuffleOperation>> Shuffle for O {
75    type ReturnOperand = O::Output;
76
77    fn shuffle(&self) -> Self::ReturnOperand {
78        Self::ReturnOperand::new(OperationContext::new(self.clone(), ShuffleOperation))
79    }
80}
81
82operation_manifest! {
83    ShuffleOperation {
84        method: Shuffle::shuffle;
85        scope: lane;
86
87        kernel {
88            parameters: <I: IndexDomain, V: ValueDomain, O: OrderState>;
89            input: (Indexed<I, V>, Multiple<O>);
90            output: OperandHandle<Indexed<I, V>, Multiple<Ordered>>;
91        }
92        kernel {
93            parameters: <V: BareValueDomain, O: OrderState>;
94            input: (Bare<V>, Multiple<O>);
95            output: OperandHandle<Bare<V>, Multiple<Ordered>>;
96        }
97    }
98}
graphrecords_query/operations/ordering/shuffle.rs

graphrecords_query/operations/ordering/
shuffle.rs
