risingwave_expr/expr/
build.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 std::iter::Peekable;
16
17use itertools::Itertools;
18use risingwave_common::types::{DataType, ScalarImpl};
19use risingwave_expr::expr::LogReport;
20use risingwave_pb::expr::ExprNode;
21use risingwave_pb::expr::expr_node::{PbType, RexNode};
22
23use super::NonStrictExpression;
24use super::expr_some_all::SomeAllExpression;
25use super::expr_udf::UserDefinedFunction;
26use super::strict::Strict;
27use super::wrapper::EvalErrorReport;
28use super::wrapper::checked::Checked;
29use super::wrapper::non_strict::NonStrict;
30use crate::expr::{
31    AsyncExpressionBoxExt, BoxedExpression, InputRefExpression, LiteralExpression, SyncExpression,
32    SyncExpressionBoxExt,
33};
34use crate::expr_context::strict_mode;
35use crate::sig::FUNCTION_REGISTRY;
36use crate::{Result, bail};
37
38/// Build an expression from protobuf.
39pub fn build_from_prost(prost: &ExprNode) -> Result<BoxedExpression> {
40    let expr = ExprBuilder::new_strict().build(prost)?;
41    Ok(match expr {
42        BoxedExpression::Sync(expr) => Strict::new(expr).boxed(),
43        BoxedExpression::Async(expr) => Strict::new(expr).boxed(),
44    })
45}
46
47/// Build an expression from protobuf in non-strict mode.
48pub fn build_non_strict_from_prost(
49    prost: &ExprNode,
50    error_report: impl EvalErrorReport + 'static,
51) -> Result<NonStrictExpression> {
52    let expr = ExprBuilder::new_non_strict(error_report).build(prost)?;
53    Ok(match expr {
54        BoxedExpression::Sync(expr) => NonStrictExpression::Sync(expr),
55        BoxedExpression::Async(expr) => NonStrictExpression::Async(expr),
56    })
57}
58
59/// Build a strict or non-strict expression according to expr context.
60///
61/// When strict mode is off, the expression will not fail but leave a null value as result.
62///
63/// Unlike [`build_non_strict_from_prost`], the returning value here can be either non-strict or
64/// strict. Thus, the caller is supposed to handle potential errors under strict mode.
65pub fn build_batch_expr_from_prost(prost: &ExprNode) -> Result<BoxedExpression> {
66    if strict_mode()? {
67        build_from_prost(prost)
68    } else {
69        // TODO(eric): report errors to users via psql notice
70        Ok(ExprBuilder::new_non_strict(LogReport).build(prost)?)
71    }
72}
73
74/// Build an expression from protobuf with possibly some wrappers attached to each node.
75struct ExprBuilder<R> {
76    /// The error reporting for non-strict mode.
77    ///
78    /// If set, each expression node will be wrapped with a [`NonStrict`] node that reports
79    /// errors to this error reporting.
80    error_report: Option<R>,
81}
82
83impl ExprBuilder<!> {
84    /// Create a new builder in strict mode.
85    fn new_strict() -> Self {
86        Self { error_report: None }
87    }
88}
89
90impl<R> ExprBuilder<R>
91where
92    R: EvalErrorReport + 'static,
93{
94    /// Create a new builder in non-strict mode with the given error reporting.
95    fn new_non_strict(error_report: R) -> Self {
96        Self {
97            error_report: Some(error_report),
98        }
99    }
100
101    /// Attach wrappers to an expression.
102    fn wrap(&self, expr: BoxedExpression) -> BoxedExpression {
103        match expr {
104            BoxedExpression::Sync(expr) => {
105                let checked = Checked(expr);
106                if let Some(error_report) = &self.error_report {
107                    NonStrict::new(checked, error_report.clone()).boxed()
108                } else {
109                    checked.boxed()
110                }
111            }
112            BoxedExpression::Async(expr) => {
113                let checked = Checked(expr);
114                if let Some(error_report) = &self.error_report {
115                    NonStrict::new(checked, error_report.clone()).boxed()
116                } else {
117                    checked.boxed()
118                }
119            }
120        }
121    }
122
123    /// Build an expression with `build_inner` and attach some wrappers.
124    fn build(&self, prost: &ExprNode) -> Result<BoxedExpression> {
125        let expr = self.build_inner(prost)?;
126        Ok(self.wrap(expr))
127    }
128
129    /// Build an expression from protobuf.
130    fn build_inner(&self, prost: &ExprNode) -> Result<BoxedExpression> {
131        use PbType as E;
132
133        let build_child = |prost: &'_ ExprNode| self.build(prost);
134
135        match prost.get_rex_node()? {
136            RexNode::InputRef(_) => InputRefExpression::build_boxed(prost, build_child),
137            RexNode::Constant(_) => LiteralExpression::build_boxed(prost, build_child),
138            RexNode::Udf(_) => UserDefinedFunction::build_boxed(prost, build_child),
139
140            RexNode::FuncCall(_) => match prost.function_type() {
141                // Dedicated types
142                E::All | E::Some => SomeAllExpression::build_boxed(prost, build_child),
143
144                // General types, lookup in the function signature map
145                _ => FuncCallBuilder::build_boxed(prost, build_child),
146            },
147
148            RexNode::Now(_) => unreachable!("now should not be built at backend"),
149
150            RexNode::SecretRef(sr) => {
151                use risingwave_common::secret::LocalSecretManager;
152                use risingwave_pb::secret::SecretRef as PbSecretRef;
153
154                let pb_ref = PbSecretRef {
155                    secret_id: sr.secret_id.into(),
156                    ref_as: sr.ref_as,
157                };
158                let value = LocalSecretManager::global()
159                    .fill_secret(pb_ref)
160                    .map_err(|e| anyhow::anyhow!(e))?;
161                Ok(
162                    LiteralExpression::new(DataType::Varchar, Some(ScalarImpl::Utf8(value.into())))
163                        .boxed(),
164                )
165            }
166        }
167    }
168}
169
170/// Manually build the expression `Self` from protobuf.
171pub(crate) trait Build: SyncExpression + Sized {
172    /// Build the expression `Self` from protobuf.
173    ///
174    /// To build children, call `build_child` on each child instead of [`build_from_prost`].
175    fn build(
176        prost: &ExprNode,
177        build_child: impl Fn(&ExprNode) -> Result<BoxedExpression>,
178    ) -> Result<Self>;
179
180    /// Build the expression `Self` from protobuf for test, where each child is built with
181    /// [`build_from_prost`].
182    #[cfg(test)]
183    fn build_for_test(prost: &ExprNode) -> Result<Self> {
184        Self::build(prost, build_from_prost)
185    }
186}
187
188/// Manually build a boxed expression from protobuf.
189pub(crate) trait BuildBoxed: 'static {
190    /// Build a boxed expression from protobuf.
191    fn build_boxed(
192        prost: &ExprNode,
193        build_child: impl Fn(&ExprNode) -> Result<BoxedExpression>,
194    ) -> Result<BoxedExpression>;
195}
196
197/// Implement [`BuildBoxed`] for all expressions that implement [`Build`].
198impl<E: Build + 'static> BuildBoxed for E {
199    fn build_boxed(
200        prost: &ExprNode,
201        build_child: impl Fn(&ExprNode) -> Result<BoxedExpression>,
202    ) -> Result<BoxedExpression> {
203        Self::build(prost, build_child).map(SyncExpressionBoxExt::boxed)
204    }
205}
206
207/// Build a function call expression from protobuf with [`build_func`].
208struct FuncCallBuilder;
209
210impl BuildBoxed for FuncCallBuilder {
211    fn build_boxed(
212        prost: &ExprNode,
213        build_child: impl Fn(&ExprNode) -> Result<BoxedExpression>,
214    ) -> Result<BoxedExpression> {
215        let func_type = prost.function_type();
216        let ret_type = DataType::from(prost.get_return_type().unwrap());
217        let func_call = prost
218            .get_rex_node()?
219            .as_func_call()
220            .expect("not a func call");
221
222        let children = func_call
223            .get_children()
224            .iter()
225            .map(build_child)
226            .try_collect()?;
227
228        build_func(func_type, ret_type, children)
229    }
230}
231
232/// Build an expression in `FuncCall` variant.
233pub fn build_func(
234    func: PbType,
235    ret_type: DataType,
236    children: Vec<BoxedExpression>,
237) -> Result<BoxedExpression> {
238    let args = children.iter().map(|c| c.return_type()).collect_vec();
239    let desc = FUNCTION_REGISTRY.get(func, &args, &ret_type)?;
240    desc.build_scalar(ret_type, children)
241}
242
243/// Build an expression in `FuncCall` variant in non-strict mode.
244///
245/// Note: This is a workaround, and only the root node are wrappedin non-strict mode.
246/// Prefer [`build_non_strict_from_prost`] if possible.
247pub fn build_func_non_strict(
248    func: PbType,
249    ret_type: DataType,
250    children: Vec<BoxedExpression>,
251    error_report: impl EvalErrorReport + 'static,
252) -> Result<NonStrictExpression> {
253    let expr = build_func(func, ret_type, children)?;
254    let wrapped = ExprBuilder::new_non_strict(error_report).wrap(expr);
255    Ok(match wrapped {
256        BoxedExpression::Sync(expr) => NonStrictExpression::Sync(expr),
257        BoxedExpression::Async(expr) => NonStrictExpression::Async(expr),
258    })
259}
260
261pub(super) fn get_children_and_return_type(prost: &ExprNode) -> Result<(&[ExprNode], DataType)> {
262    let ret_type = DataType::from(prost.get_return_type().unwrap());
263    if let RexNode::FuncCall(func_call) = prost.get_rex_node().unwrap() {
264        Ok((func_call.get_children(), ret_type))
265    } else {
266        bail!("Expected RexNode::FuncCall");
267    }
268}
269
270/// Build an expression from a string.
271///
272/// # Example
273///
274/// ```ignore
275/// # use risingwave_expr::expr::build_from_pretty;
276/// build_from_pretty("42:int2"); // literal
277/// build_from_pretty("$0:int8"); // inputref
278/// build_from_pretty("(add:int8 42:int2 $1:int8)"); // function
279/// build_from_pretty("(add:int8 42:int2 (add:int8 42:int2 $1:int8))");
280/// ```
281///
282/// # Syntax
283///
284/// ```text
285/// <expr>      ::= <literal> | <input_ref> | <function>
286/// <literal>   ::= <value>:<type>
287/// <input_ref> ::= <index>:<type>
288/// <function>  ::= (<name>:<type> <expr>...)
289/// <name>      ::= [a-zA-Z_][a-zA-Z0-9_]*
290/// <index>     ::= $[0-9]+
291/// ```
292pub fn build_from_pretty(s: impl AsRef<str>) -> BoxedExpression {
293    let tokens = lexer(s.as_ref());
294    Parser::new(tokens.into_iter()).parse_expression()
295}
296
297struct Parser<Iter: Iterator> {
298    tokens: Peekable<Iter>,
299}
300
301impl<Iter: Iterator<Item = Token>> Parser<Iter> {
302    fn new(tokens: Iter) -> Self {
303        Self {
304            tokens: tokens.peekable(),
305        }
306    }
307
308    fn parse_expression(&mut self) -> BoxedExpression {
309        match self.tokens.next().expect("Unexpected end of input") {
310            Token::Index(index) => {
311                assert_eq!(self.tokens.next(), Some(Token::Colon), "Expected a Colon");
312                let ty = self.parse_type();
313                InputRefExpression::new(ty, index).boxed()
314            }
315            Token::LParen => {
316                let func = self.parse_function();
317                assert_eq!(self.tokens.next(), Some(Token::Colon), "Expected a Colon");
318                let ty = self.parse_type();
319
320                let mut children = Vec::new();
321                while self.tokens.peek() != Some(&Token::RParen) {
322                    children.push(self.parse_expression());
323                }
324                self.tokens.next(); // Consume the RParen
325
326                build_func(func, ty, children).expect("Failed to build")
327            }
328            Token::Literal(value) => {
329                assert_eq!(self.tokens.next(), Some(Token::Colon), "Expected a Colon");
330                let ty = self.parse_type();
331                let value = match value.as_str() {
332                    "null" | "NULL" => None,
333                    _ => Some(ScalarImpl::from_text(&value, &ty).expect_str("value", &value)),
334                };
335                LiteralExpression::new(ty, value).boxed()
336            }
337            _ => panic!("Unexpected token"),
338        }
339    }
340
341    fn parse_type(&mut self) -> DataType {
342        match self.tokens.next().expect("Unexpected end of input") {
343            Token::Literal(name) => {
344                let mut processed_name = name.replace('_', " ");
345
346                // Special logic to support Map type in `build_from_pretty`.
347                // Please refer to `src/expr/impl/src/scalar/map_filter.rs`.
348                if processed_name.starts_with("map") {
349                    processed_name = processed_name.replace('<', "(").replace('>', ")");
350                }
351
352                processed_name
353                    .parse::<DataType>()
354                    .expect_str("type", &processed_name)
355            }
356            t => panic!("Expected a Literal, got {t:?}"),
357        }
358    }
359
360    fn parse_function(&mut self) -> PbType {
361        match self.tokens.next().expect("Unexpected end of input") {
362            Token::Literal(name) => {
363                PbType::from_str_name(&name.to_uppercase()).expect_str("function", &name)
364            }
365            t => panic!("Expected a Literal, got {t:?}"),
366        }
367    }
368}
369
370#[derive(Debug, PartialEq, Clone)]
371pub(crate) enum Token {
372    LParen,
373    RParen,
374    Colon,
375    Index(usize),
376    Literal(String),
377}
378
379pub(crate) fn lexer(input: &str) -> Vec<Token> {
380    let mut tokens = Vec::new();
381    let mut chars = input.chars().peekable();
382    while let Some(c) = chars.next() {
383        let token = match c {
384            '(' => Token::LParen,
385            ')' => Token::RParen,
386            ':' => Token::Colon,
387            '$' => {
388                let mut number = String::new();
389                while let Some(c) = chars.peek()
390                    && c.is_ascii_digit()
391                {
392                    number.push(chars.next().unwrap());
393                }
394                let index = number.parse::<usize>().expect("Invalid number");
395                Token::Index(index)
396            }
397            ' ' | '\t' | '\r' | '\n' => continue,
398            _ => {
399                let mut literal = String::new();
400                literal.push(c);
401                while let Some(&c) = chars.peek()
402                    && !matches!(c, '(' | ')' | ':' | ' ' | '\t' | '\r' | '\n')
403                {
404                    literal.push(chars.next().unwrap());
405                }
406                Token::Literal(literal)
407            }
408        };
409        tokens.push(token);
410    }
411    tokens
412}
413
414pub(crate) trait ExpectExt<T> {
415    fn expect_str(self, what: &str, s: &str) -> T;
416}
417
418impl<T> ExpectExt<T> for Option<T> {
419    #[track_caller]
420    fn expect_str(self, what: &str, s: &str) -> T {
421        match self {
422            Some(x) => x,
423            None => panic!("expect {what} in {s:?}"),
424        }
425    }
426}
427
428impl<T, E> ExpectExt<T> for std::result::Result<T, E> {
429    #[track_caller]
430    fn expect_str(self, what: &str, s: &str) -> T {
431        match self {
432            Ok(x) => x,
433            Err(_) => panic!("expect {what} in {s:?}"),
434        }
435    }
436}
risingwave_expr/expr/build.rs

risingwave_expr/expr/
build.rs
