risingwave_frontend/expr/
function_call.rs

1// Copyright 2022 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 itertools::Itertools;
16use risingwave_common::catalog::Schema;
17use risingwave_common::types::{DataType, ScalarImpl};
18use risingwave_common::util::iter_util::ZipEqFast;
19
20use super::type_inference::cast;
21use super::{CastContext, CastError, Expr, ExprImpl, Literal, infer_some_all, infer_type};
22use crate::error::Result as RwResult;
23use crate::expr::{ExprDisplay, ExprType, bail_cast_error, is_impure_func_call};
24
25#[derive(Clone, Eq, PartialEq, Hash)]
26pub struct FunctionCall {
27    pub(super) func_type: ExprType,
28    pub(super) return_type: DataType,
29    pub(super) inputs: Vec<ExprImpl>,
30}
31
32fn debug_binary_op(
33    f: &mut std::fmt::Formatter<'_>,
34    op: &str,
35    inputs: &[ExprImpl],
36) -> std::fmt::Result {
37    use std::fmt::Debug;
38
39    assert_eq!(inputs.len(), 2);
40
41    write!(f, "(")?;
42    inputs[0].fmt(f)?;
43    write!(f, " {} ", op)?;
44    inputs[1].fmt(f)?;
45    write!(f, ")")?;
46
47    Ok(())
48}
49
50impl std::fmt::Debug for FunctionCall {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        if f.alternate() {
53            f.debug_struct("FunctionCall")
54                .field("func_type", &self.func_type)
55                .field("return_type", &self.return_type)
56                .field("inputs", &self.inputs)
57                .finish()
58        } else {
59            match &self.func_type {
60                ExprType::Cast => {
61                    assert_eq!(self.inputs.len(), 1);
62                    self.inputs[0].fmt(f)?;
63                    write!(f, "::{:?}", self.return_type)
64                }
65                ExprType::Add => debug_binary_op(f, "+", &self.inputs),
66                ExprType::Subtract => debug_binary_op(f, "-", &self.inputs),
67                ExprType::Multiply => debug_binary_op(f, "*", &self.inputs),
68                ExprType::Divide => debug_binary_op(f, "/", &self.inputs),
69                ExprType::Modulus => debug_binary_op(f, "%", &self.inputs),
70                ExprType::Equal => debug_binary_op(f, "=", &self.inputs),
71                ExprType::NotEqual => debug_binary_op(f, "<>", &self.inputs),
72                ExprType::LessThan => debug_binary_op(f, "<", &self.inputs),
73                ExprType::LessThanOrEqual => debug_binary_op(f, "<=", &self.inputs),
74                ExprType::GreaterThan => debug_binary_op(f, ">", &self.inputs),
75                ExprType::GreaterThanOrEqual => debug_binary_op(f, ">=", &self.inputs),
76                ExprType::And => debug_binary_op(f, "AND", &self.inputs),
77                ExprType::Or => debug_binary_op(f, "OR", &self.inputs),
78                ExprType::BitwiseShiftLeft => debug_binary_op(f, "<<", &self.inputs),
79                ExprType::BitwiseShiftRight => debug_binary_op(f, ">>", &self.inputs),
80                ExprType::BitwiseAnd => debug_binary_op(f, "&", &self.inputs),
81                ExprType::BitwiseOr => debug_binary_op(f, "|", &self.inputs),
82                ExprType::BitwiseXor => debug_binary_op(f, "#", &self.inputs),
83                ExprType::ArrayContains => debug_binary_op(f, "@>", &self.inputs),
84                ExprType::ArrayContained => debug_binary_op(f, "<@", &self.inputs),
85                _ => {
86                    let func_name = format!("{:?}", self.func_type);
87                    let mut builder = f.debug_tuple(&func_name);
88                    self.inputs.iter().for_each(|child| {
89                        builder.field(child);
90                    });
91                    builder.finish()
92                }
93            }
94        }
95    }
96}
97
98impl FunctionCall {
99    /// Create a `FunctionCall` expr with the return type inferred from `func_type` and types of
100    /// `inputs`.
101    // The functions listed here are all variadic.  Type signatures of functions that take a fixed
102    // number of arguments are checked
103    // [elsewhere](crate::expr::type_inference::build_type_derive_map).
104    pub fn new(func_type: ExprType, mut inputs: Vec<ExprImpl>) -> RwResult<Self> {
105        let return_type = infer_type(func_type.into(), &mut inputs)?;
106        Ok(Self::new_unchecked(func_type, inputs, return_type))
107    }
108
109    /// Create a cast expr over `child` to `target` type in `allows` context.
110    /// The input `child` remains unchanged when this returns an error.
111    pub fn cast_mut(
112        child: &mut ExprImpl,
113        target: &DataType,
114        allows: CastContext,
115    ) -> Result<(), CastError> {
116        if let ExprImpl::Parameter(expr) = child
117            && !expr.has_infer()
118        {
119            // Always Ok below. Safe to mutate `expr` (from `child`).
120            expr.cast_infer_type(target);
121            return Ok(());
122        }
123        if let ExprImpl::FunctionCall(func) = child
124            && func.func_type == ExprType::Row
125        {
126            // Row function will have empty fields in Datatype::Struct at this point. Therefore,
127            // we will need to take some special care to generate the cast types. For normal struct
128            // types, they will be handled in `cast_ok`.
129            return Self::cast_row_expr(func, target, allows);
130        }
131        if child.is_untyped() {
132            // `is_unknown` makes sure `as_literal` and `as_utf8` will never panic.
133            let literal = child.as_literal().unwrap();
134            let datum = literal
135                .get_data()
136                .as_ref()
137                .map(|scalar| ScalarImpl::from_text(scalar.as_utf8(), target))
138                .transpose();
139            if let Ok(datum) = datum {
140                *child = Literal::new(datum, target.clone()).into();
141                return Ok(());
142            }
143            // else when eager parsing fails, just proceed as normal.
144            // Some callers are not ready to handle `'a'::int` error here.
145        }
146
147        let source = child.return_type();
148        if &source == target {
149            return Ok(());
150        }
151
152        if child.is_untyped() {
153            // Casting from unknown is allowed in all context. And PostgreSQL actually does the parsing
154            // in frontend.
155        } else {
156            cast(&source, target, allows)?;
157        }
158
159        // Always Ok below. Safe to mutate `child`.
160        let owned = std::mem::replace(child, ExprImpl::literal_bool(false));
161        *child = Self::new_unchecked(ExprType::Cast, vec![owned], target.clone()).into();
162        Ok(())
163    }
164
165    /// Cast a `ROW` expression to the target type. We intentionally disallow casting arbitrary
166    /// expressions, like `ROW(1)::STRUCT<i INTEGER>` to `STRUCT<VARCHAR>`, although an integer
167    /// is castable to VARCHAR. It's to simply the casting rules.
168    fn cast_row_expr(
169        func: &mut FunctionCall,
170        target_type: &DataType,
171        allows: CastContext,
172    ) -> Result<(), CastError> {
173        // Can only cast to a struct type.
174        let DataType::Struct(t) = &target_type else {
175            bail_cast_error!(
176                "cannot cast type \"{}\" to \"{}\"",
177                func.return_type(), // typically "record"
178                target_type,
179            );
180        };
181
182        let expected_len = t.len();
183        let actual_len = func.inputs.len();
184
185        match expected_len.cmp(&actual_len) {
186            std::cmp::Ordering::Equal => {
187                // FIXME: `func` shall not be in a partially mutated state when one of its fields
188                // fails to cast.
189                func.inputs
190                    .iter_mut()
191                    .zip_eq_fast(t.types())
192                    .try_for_each(|(e, t)| Self::cast_mut(e, t, allows))?;
193                func.return_type = target_type.clone();
194                Ok(())
195            }
196            std::cmp::Ordering::Less => bail_cast_error!(
197                "input has too many columns, expected {expected_len} but got {actual_len}"
198            ),
199            std::cmp::Ordering::Greater => bail_cast_error!(
200                "input has too few columns, expected {expected_len} but got {actual_len}"
201            ),
202        }
203    }
204
205    /// Construct a `FunctionCall` expr directly with the provided `return_type`, bypassing type
206    /// inference. Use with caution.
207    pub fn new_unchecked(
208        func_type: ExprType,
209        inputs: Vec<ExprImpl>,
210        return_type: DataType,
211    ) -> Self {
212        FunctionCall {
213            func_type,
214            return_type,
215            inputs,
216        }
217    }
218
219    pub fn new_binary_op_func(
220        mut func_types: Vec<ExprType>,
221        mut inputs: Vec<ExprImpl>,
222    ) -> RwResult<ExprImpl> {
223        let expr_type = func_types.remove(0);
224        match expr_type {
225            ExprType::Some | ExprType::All => {
226                let return_type = infer_some_all(func_types, &mut inputs)?;
227                Ok(FunctionCall::new_unchecked(expr_type, inputs, return_type).into())
228            }
229            ExprType::Not | ExprType::IsNotNull | ExprType::IsNull | ExprType::Neg => {
230                Ok(FunctionCall::new(
231                    expr_type,
232                    vec![Self::new_binary_op_func(func_types, inputs)?],
233                )?
234                .into())
235            }
236            _ => Ok(FunctionCall::new(expr_type, inputs)?.into()),
237        }
238    }
239
240    pub fn decompose(self) -> (ExprType, Vec<ExprImpl>, DataType) {
241        (self.func_type, self.inputs, self.return_type)
242    }
243
244    pub fn decompose_as_binary(self) -> (ExprType, ExprImpl, ExprImpl) {
245        assert_eq!(self.inputs.len(), 2);
246        let mut iter = self.inputs.into_iter();
247        let left = iter.next().unwrap();
248        let right = iter.next().unwrap();
249        (self.func_type, left, right)
250    }
251
252    pub fn decompose_as_unary(self) -> (ExprType, ExprImpl) {
253        assert_eq!(self.inputs.len(), 1);
254        let mut iter = self.inputs.into_iter();
255        let input = iter.next().unwrap();
256        (self.func_type, input)
257    }
258
259    pub fn func_type(&self) -> ExprType {
260        self.func_type
261    }
262
263    /// Get a reference to the function call's inputs.
264    pub fn inputs(&self) -> &[ExprImpl] {
265        self.inputs.as_ref()
266    }
267
268    pub fn inputs_mut(&mut self) -> &mut [ExprImpl] {
269        self.inputs.as_mut()
270    }
271
272    pub(super) fn from_expr_proto(
273        function_call: &risingwave_pb::expr::FunctionCall,
274        func_type: ExprType,
275        return_type: DataType,
276    ) -> RwResult<Self> {
277        let inputs: Vec<_> = function_call
278            .get_children()
279            .iter()
280            .map(ExprImpl::from_expr_proto)
281            .try_collect()?;
282        Ok(Self {
283            func_type,
284            return_type,
285            inputs,
286        })
287    }
288
289    pub fn is_pure(&self) -> bool {
290        !is_impure_func_call(self)
291    }
292}
293
294impl Expr for FunctionCall {
295    fn return_type(&self) -> DataType {
296        self.return_type.clone()
297    }
298
299    fn try_to_expr_proto(&self) -> Result<risingwave_pb::expr::ExprNode, String> {
300        use risingwave_pb::expr::expr_node::*;
301        use risingwave_pb::expr::*;
302
303        let children = self
304            .inputs()
305            .iter()
306            .map(|input| input.try_to_expr_proto())
307            .try_collect()?;
308
309        Ok(ExprNode {
310            function_type: self.func_type().into(),
311            return_type: Some(self.return_type().to_protobuf()),
312            rex_node: Some(RexNode::FuncCall(FunctionCall { children })),
313        })
314    }
315}
316
317pub struct FunctionCallDisplay<'a> {
318    pub function_call: &'a FunctionCall,
319    pub input_schema: &'a Schema,
320}
321
322impl std::fmt::Debug for FunctionCallDisplay<'_> {
323    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
324        let that = self.function_call;
325        match &that.func_type {
326            ExprType::Cast => {
327                assert_eq!(that.inputs.len(), 1);
328                ExprDisplay {
329                    expr: &that.inputs[0],
330                    input_schema: self.input_schema,
331                }
332                .fmt(f)?;
333                write!(f, "::{:?}", that.return_type)
334            }
335            ExprType::Add => explain_verbose_binary_op(f, "+", &that.inputs, self.input_schema),
336            ExprType::Subtract => {
337                explain_verbose_binary_op(f, "-", &that.inputs, self.input_schema)
338            }
339            ExprType::Multiply => {
340                explain_verbose_binary_op(f, "*", &that.inputs, self.input_schema)
341            }
342            ExprType::Divide => explain_verbose_binary_op(f, "/", &that.inputs, self.input_schema),
343            ExprType::Modulus => explain_verbose_binary_op(f, "%", &that.inputs, self.input_schema),
344            ExprType::Equal => explain_verbose_binary_op(f, "=", &that.inputs, self.input_schema),
345            ExprType::NotEqual => {
346                explain_verbose_binary_op(f, "<>", &that.inputs, self.input_schema)
347            }
348            ExprType::LessThan => {
349                explain_verbose_binary_op(f, "<", &that.inputs, self.input_schema)
350            }
351            ExprType::LessThanOrEqual => {
352                explain_verbose_binary_op(f, "<=", &that.inputs, self.input_schema)
353            }
354            ExprType::GreaterThan => {
355                explain_verbose_binary_op(f, ">", &that.inputs, self.input_schema)
356            }
357            ExprType::GreaterThanOrEqual => {
358                explain_verbose_binary_op(f, ">=", &that.inputs, self.input_schema)
359            }
360            ExprType::And => explain_verbose_binary_op(f, "AND", &that.inputs, self.input_schema),
361            ExprType::Or => explain_verbose_binary_op(f, "OR", &that.inputs, self.input_schema),
362            ExprType::BitwiseShiftLeft => {
363                explain_verbose_binary_op(f, "<<", &that.inputs, self.input_schema)
364            }
365            ExprType::BitwiseShiftRight => {
366                explain_verbose_binary_op(f, ">>", &that.inputs, self.input_schema)
367            }
368            ExprType::BitwiseAnd => {
369                explain_verbose_binary_op(f, "&", &that.inputs, self.input_schema)
370            }
371            ExprType::BitwiseOr => {
372                explain_verbose_binary_op(f, "|", &that.inputs, self.input_schema)
373            }
374            ExprType::BitwiseXor => {
375                explain_verbose_binary_op(f, "#", &that.inputs, self.input_schema)
376            }
377            ExprType::ArrayContains => {
378                explain_verbose_binary_op(f, "@>", &that.inputs, self.input_schema)
379            }
380            ExprType::ArrayContained => {
381                explain_verbose_binary_op(f, "<@", &that.inputs, self.input_schema)
382            }
383            ExprType::Proctime => {
384                write!(f, "{:?}", that.func_type)
385            }
386            _ => {
387                let func_name = format!("{:?}", that.func_type);
388                let mut builder = f.debug_tuple(&func_name);
389                that.inputs.iter().for_each(|child| {
390                    builder.field(&ExprDisplay {
391                        expr: child,
392                        input_schema: self.input_schema,
393                    });
394                });
395                builder.finish()
396            }
397        }
398    }
399}
400
401fn explain_verbose_binary_op(
402    f: &mut std::fmt::Formatter<'_>,
403    op: &str,
404    inputs: &[ExprImpl],
405    input_schema: &Schema,
406) -> std::fmt::Result {
407    use std::fmt::Debug;
408
409    assert_eq!(inputs.len(), 2);
410
411    write!(f, "(")?;
412    ExprDisplay {
413        expr: &inputs[0],
414        input_schema,
415    }
416    .fmt(f)?;
417    write!(f, " {} ", op)?;
418    ExprDisplay {
419        expr: &inputs[1],
420        input_schema,
421    }
422    .fmt(f)?;
423    write!(f, ")")?;
424
425    Ok(())
426}
427
428pub fn is_row_function(expr: &ExprImpl) -> bool {
429    if let ExprImpl::FunctionCall(func) = expr
430        && func.func_type() == ExprType::Row
431    {
432        return true;
433    }
434    false
435}