glaredb_core/testutil/
operator_wrapper.rs

1use std::sync::Arc;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::task::{Context, Wake, Waker};
4
5use glaredb_error::Result;
6
7use super::database_context::test_db_context;
8use crate::arrays::batch::Batch;
9use crate::execution::operators::{
10    BaseOperator,
11    ExecuteOperator,
12    PollExecute,
13    PollFinalize,
14    PollPull,
15    PollPush,
16    PullOperator,
17    PushOperator,
18};
19
20/// Waker containing a count that gets incremented by one on every wake.
21#[derive(Debug, Default)]
22pub struct CountingWaker {
23    count: AtomicUsize,
24}
25
26impl CountingWaker {
27    pub fn wake_count(&self) -> usize {
28        self.count.load(Ordering::SeqCst)
29    }
30}
31
32impl Wake for CountingWaker {
33    fn wake(self: Arc<Self>) {
34        self.count.fetch_add(1, Ordering::SeqCst);
35    }
36}
37
38/// Wrapper around an operator that provides extra utilities to reduce boiler
39/// plate.
40///
41/// All `poll_` helper method currently shared the same counting waker.
42pub struct OperatorWrapper<O: BaseOperator> {
43    pub waker: Arc<CountingWaker>,
44    pub operator: O,
45}
46
47impl<O> OperatorWrapper<O>
48where
49    O: BaseOperator,
50{
51    pub fn new(operator: O) -> Self {
52        OperatorWrapper {
53            waker: Arc::new(CountingWaker::default()),
54            operator,
55        }
56    }
57}
58
59impl<O> OperatorWrapper<O>
60where
61    O: ExecuteOperator,
62{
63    #[track_caller]
64    pub fn poll_execute(
65        &self,
66        operator_state: &O::OperatorState,
67        state: &mut O::PartitionExecuteState,
68        input: &mut Batch,
69        output: &mut Batch,
70    ) -> Result<PollExecute> {
71        let waker = Waker::from(self.waker.clone());
72        let mut cx = Context::from_waker(&waker);
73        self.operator
74            .poll_execute(&mut cx, operator_state, state, input, output)
75    }
76
77    #[track_caller]
78    pub fn poll_finalize_execute(
79        &self,
80        operator_state: &O::OperatorState,
81        state: &mut O::PartitionExecuteState,
82    ) -> Result<PollFinalize> {
83        let waker = Waker::from(self.waker.clone());
84        let mut cx = Context::from_waker(&waker);
85        self.operator
86            .poll_finalize_execute(&mut cx, operator_state, state)
87    }
88}
89
90impl<O> OperatorWrapper<O>
91where
92    O: PullOperator,
93{
94    #[track_caller]
95    pub fn poll_pull(
96        &self,
97        operator_state: &O::OperatorState,
98        state: &mut O::PartitionPullState,
99        output: &mut Batch,
100    ) -> Result<PollPull> {
101        let waker = Waker::from(self.waker.clone());
102        let mut cx = Context::from_waker(&waker);
103        self.operator
104            .poll_pull(&mut cx, operator_state, state, output)
105    }
106}
107
108impl<O> OperatorWrapper<O>
109where
110    O: PushOperator,
111{
112    #[track_caller]
113    pub fn poll_push(
114        &self,
115        operator_state: &O::OperatorState,
116        state: &mut O::PartitionPushState,
117        input: &mut Batch,
118    ) -> Result<PollPush> {
119        let waker = Waker::from(self.waker.clone());
120        let mut cx = Context::from_waker(&waker);
121        self.operator
122            .poll_push(&mut cx, operator_state, state, input)
123    }
124
125    #[track_caller]
126    pub fn poll_finalize_push(
127        &self,
128        operator_state: &O::OperatorState,
129        state: &mut O::PartitionPushState,
130    ) -> Result<PollFinalize> {
131        let waker = Waker::from(self.waker.clone());
132        let mut cx = Context::from_waker(&waker);
133        self.operator
134            .poll_finalize_push(&mut cx, operator_state, state)
135    }
136}
glaredb_core/testutil/operator_wrapper.rs

glaredb_core/testutil/
operator_wrapper.rs
