risingwave_expr/expr/wrapper/
checked.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use risingwave_common::array::{ArrayRef, DataChunk};
16use risingwave_common::row::OwnedRow;
17use risingwave_common::types::{DataType, Datum};
18
19use crate::error::Result;
20use crate::expr::{AsyncExpression, ExpressionInfo, SyncExpression, ValueImpl};
21
22/// A wrapper of an expression that does extra checks after evaluation.
23#[derive(Debug)]
24pub(crate) struct Checked<E>(pub E);
25
26impl<E: ExpressionInfo> ExpressionInfo for Checked<E> {
27    fn return_type(&self) -> DataType {
28        self.0.return_type()
29    }
30
31    fn input_ref_index(&self) -> Option<usize> {
32        self.0.input_ref_index()
33    }
34}
35
36macro_rules! checked_eval {
37    ($mode:ident, $this:expr, $input:expr, $method:ident) => {{
38        let res = forward!($mode, $this.0, $method($input))?;
39        assert_eq!(res.len(), $input.capacity());
40        Ok(res)
41    }};
42}
43
44impl<E: SyncExpression> SyncExpression for Checked<E> {
45    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
46        checked_eval!(sync, self, input, eval)
47    }
48
49    fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
50        checked_eval!(sync, self, input, eval_v2)
51    }
52
53    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
54        self.0.eval_row(input)
55    }
56
57    fn eval_const(&self) -> Result<Datum> {
58        self.0.eval_const()
59    }
60}
61
62impl<E: AsyncExpression> AsyncExpression for Checked<E> {
63    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
64        checked_eval!(async, self, input, eval)
65    }
66
67    async fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
68        checked_eval!(async, self, input, eval_v2)
69    }
70
71    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
72        self.0.eval_row(input).await
73    }
74}
risingwave_expr/expr/wrapper/checked.rs

risingwave_expr/expr/wrapper/
checked.rs
