graphrecords_query/operations/ordering/
take.rs

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

graphrecords_query/operations/ordering/
take.rs
