risingwave_expr/expr/
mod.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
15//! Expressions in RisingWave.
16//!
17//! All expressions are implemented under the [`SyncExpression`] or [`AsyncExpression`] trait.
18//!
19//! ## Construction
20//!
21//! Expressions can be constructed by [`build_func()`] function, which returns a
22//! [`BoxedExpression`].
23//!
24//! They can also be transformed from the prost [`ExprNode`] using the [`build_from_prost()`]
25//! function.
26//!
27//! ## Evaluation
28//!
29//! Expressions can be evaluated using the [`eval`] function.
30//!
31//! [`ExprNode`]: risingwave_pb::expr::ExprNode
32//! [`eval`]: BoxedExpression::eval
33
34#[macro_export]
35macro_rules! forward {
36    (sync, $expr:expr, $method:ident($($arg:expr),* $(,)?)) => {
37        ($expr).$method($($arg),*)
38    };
39    (async, $expr:expr, $method:ident($($arg:expr),* $(,)?)) => {
40        ($expr).$method($($arg),*).await
41    };
42}
43
44// These modules define concrete expression structures.
45mod and_or;
46mod expr_input_ref;
47mod expr_literal;
48mod expr_some_all;
49pub(crate) mod expr_udf;
50pub(crate) mod wrapper;
51
52mod build;
53pub mod test_utils;
54mod value;
55
56use std::future::Future;
57use std::sync::Arc;
58
59use risingwave_common::array::{ArrayRef, DataChunk};
60use risingwave_common::row::OwnedRow;
61use risingwave_common::types::{DataType, Datum};
62
63pub use self::build::*;
64pub use self::expr_input_ref::InputRefExpression;
65pub use self::expr_literal::LiteralExpression;
66pub use self::value::{ValueImpl, ValueRef};
67pub use self::wrapper::*;
68pub use super::{ExprError, Result};
69
70/// Common metadata of an expression.
71#[auto_impl::auto_impl(&, Box, Arc)]
72pub trait ExpressionInfo: std::fmt::Debug + Sync + Send {
73    /// Get the return data type.
74    fn return_type(&self) -> DataType;
75
76    /// Get the index if the expression is an `InputRef`.
77    fn input_ref_index(&self) -> Option<usize> {
78        None
79    }
80}
81
82/// Interface of a synchronous expression.
83///
84/// There're two functions to evaluate an expression: `eval` and `eval_v2`, exactly one of them
85/// should be implemented. Prefer calling and implementing `eval_v2` instead of `eval` if possible,
86/// to gain the performance benefit of scalar expression.
87#[auto_impl::auto_impl(&, Box, Arc)]
88pub trait SyncExpression: ExpressionInfo {
89    /// Evaluate the expression in vectorized execution. Returns an array.
90    ///
91    /// The default implementation calls `eval_v2` and always converts the result to an array.
92    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
93        let value = self.eval_v2(input)?;
94        Ok(match value {
95            ValueImpl::Array(array) => array,
96            ValueImpl::Scalar { value, capacity } => {
97                let mut builder = self.return_type().create_array_builder(capacity);
98                builder.append_n(capacity, value);
99                builder.finish().into()
100            }
101        })
102    }
103
104    /// Evaluate the expression in vectorized execution. Returns a value that can be either an
105    /// array, or a scalar if all values in the array are the same.
106    ///
107    /// The default implementation calls `eval` and puts the result into the `Array` variant.
108    fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
109        self.eval(input).map(ValueImpl::Array)
110    }
111
112    /// Evaluate the expression in row-based execution. Returns a nullable scalar.
113    fn eval_row(&self, input: &OwnedRow) -> Result<Datum>;
114
115    /// Evaluate if the expression is constant.
116    fn eval_const(&self) -> Result<Datum> {
117        Err(ExprError::NotConstant)
118    }
119}
120
121/// Interface of an asynchronous expression.
122pub trait AsyncExpression: ExpressionInfo {
123    /// Evaluate the expression in vectorized execution. Returns an array.
124    ///
125    /// The default implementation calls `eval_v2` and always converts the result to an array.
126    fn eval<'a>(
127        &'a self,
128        input: &'a DataChunk,
129    ) -> impl Future<Output = Result<ArrayRef>> + Send + 'a {
130        async move {
131            let value = self.eval_v2(input).await?;
132            Ok(match value {
133                ValueImpl::Array(array) => array,
134                ValueImpl::Scalar { value, capacity } => {
135                    let mut builder = self.return_type().create_array_builder(capacity);
136                    builder.append_n(capacity, value);
137                    builder.finish().into()
138                }
139            })
140        }
141    }
142
143    /// Evaluate the expression in vectorized execution. Returns a value that can be either an
144    /// array, or a scalar if all values in the array are the same.
145    ///
146    /// The default implementation calls `eval` and puts the result into the `Array` variant.
147    fn eval_v2<'a>(
148        &'a self,
149        input: &'a DataChunk,
150    ) -> impl Future<Output = Result<ValueImpl>> + Send + 'a {
151        async move { self.eval(input).await.map(ValueImpl::Array) }
152    }
153
154    /// Evaluate the expression in row-based execution. Returns a nullable scalar.
155    fn eval_row<'a>(
156        &'a self,
157        input: &'a OwnedRow,
158    ) -> impl Future<Output = Result<Datum>> + Send + 'a;
159}
160
161/// Object-safe adapter for asynchronous expressions.
162#[async_trait::async_trait]
163pub trait AsyncDynExpression: ExpressionInfo {
164    /// Evaluate the expression in vectorized execution. Returns an array.
165    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef>;
166
167    /// Evaluate the expression in vectorized execution. Returns a value that can be either an
168    /// array, or a scalar if all values in the array are the same.
169    async fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl>;
170
171    /// Evaluate the expression in row-based execution. Returns a nullable scalar.
172    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum>;
173}
174
175#[async_trait::async_trait]
176impl<E> AsyncDynExpression for E
177where
178    E: AsyncExpression,
179{
180    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
181        AsyncExpression::eval(self, input).await
182    }
183
184    async fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
185        AsyncExpression::eval_v2(self, input).await
186    }
187
188    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
189        AsyncExpression::eval_row(self, input).await
190    }
191}
192
193impl AsyncExpression for Arc<dyn AsyncDynExpression> {
194    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
195        AsyncDynExpression::eval(self.as_ref(), input).await
196    }
197
198    async fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
199        AsyncDynExpression::eval_v2(self.as_ref(), input).await
200    }
201
202    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
203        AsyncDynExpression::eval_row(self.as_ref(), input).await
204    }
205}
206
207/// An owned dynamically typed expression.
208#[derive(Clone, Debug)]
209pub enum BoxedExpression {
210    Sync(Arc<dyn SyncExpression>),
211    Async(Arc<dyn AsyncDynExpression>),
212}
213
214/// Try to unwrap boxed expressions into sync expressions.
215///
216/// Returns the original boxed expression list if any expression is async.
217pub fn try_into_sync_exprs(
218    exprs: Vec<BoxedExpression>,
219) -> std::result::Result<Vec<Arc<dyn SyncExpression>>, Vec<BoxedExpression>> {
220    try_convert_all(
221        exprs,
222        |expr| match expr {
223            BoxedExpression::Sync(expr) => Ok(expr),
224            expr @ BoxedExpression::Async(_) => Err(expr),
225        },
226        BoxedExpression::Sync,
227    )
228}
229
230/// Try to convert all items in a vector, or return the original vector if any conversion fails.
231pub fn try_convert_all<T, U>(
232    items: Vec<T>,
233    mut try_convert: impl FnMut(T) -> std::result::Result<U, T>,
234    recover: impl Fn(U) -> T,
235) -> std::result::Result<Vec<U>, Vec<T>> {
236    let mut converted = Vec::with_capacity(items.len());
237    let mut items = items.into_iter();
238    loop {
239        let Some(item) = items.next() else {
240            return Ok(converted);
241        };
242        match try_convert(item) {
243            Ok(item) => converted.push(item),
244            Err(item) => {
245                let items = converted
246                    .into_iter()
247                    .map(recover)
248                    .chain(std::iter::once(item))
249                    .chain(items)
250                    .collect();
251                return Err(items);
252            }
253        }
254    }
255}
256
257impl BoxedExpression {
258    /// Get the return data type.
259    pub fn return_type(&self) -> DataType {
260        match self {
261            Self::Sync(expr) => expr.return_type(),
262            Self::Async(expr) => expr.return_type(),
263        }
264    }
265
266    /// Get the index if the expression is an `InputRef`.
267    pub fn input_ref_index(&self) -> Option<usize> {
268        match self {
269            Self::Sync(expr) => expr.input_ref_index(),
270            Self::Async(expr) => expr.input_ref_index(),
271        }
272    }
273
274    /// Evaluate if the expression is constant.
275    pub fn eval_const(&self) -> Result<Datum> {
276        match self {
277            Self::Sync(expr) => expr.eval_const(),
278            Self::Async(_) => Err(ExprError::NotConstant),
279        }
280    }
281
282    /// Evaluate the expression in vectorized execution. Returns an array.
283    pub async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
284        match self {
285            Self::Sync(expr) => expr.eval(input),
286            Self::Async(expr) => AsyncDynExpression::eval(expr.as_ref(), input).await,
287        }
288    }
289
290    /// Evaluate the expression in vectorized execution. Returns a value that can be either an
291    /// array, or a scalar if all values in the array are the same.
292    pub async fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
293        match self {
294            Self::Sync(expr) => expr.eval_v2(input),
295            Self::Async(expr) => AsyncDynExpression::eval_v2(expr.as_ref(), input).await,
296        }
297    }
298
299    /// Evaluate the expression in row-based execution. Returns a nullable scalar.
300    pub async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
301        match self {
302            Self::Sync(expr) => expr.eval_row(input),
303            Self::Async(expr) => AsyncDynExpression::eval_row(expr.as_ref(), input).await,
304        }
305    }
306}
307
308impl<E> From<E> for BoxedExpression
309where
310    E: SyncExpression + 'static,
311{
312    fn from(expr: E) -> Self {
313        Self::Sync(Arc::new(expr))
314    }
315}
316
317impl ExpressionInfo for BoxedExpression {
318    fn return_type(&self) -> DataType {
319        self.return_type()
320    }
321
322    fn input_ref_index(&self) -> Option<usize> {
323        self.input_ref_index()
324    }
325}
326
327impl AsyncExpression for BoxedExpression {
328    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
329        match self {
330            Self::Sync(expr) => expr.eval(input),
331            Self::Async(expr) => AsyncDynExpression::eval(expr.as_ref(), input).await,
332        }
333    }
334
335    async fn eval_v2(&self, input: &DataChunk) -> Result<ValueImpl> {
336        match self {
337            Self::Sync(expr) => expr.eval_v2(input),
338            Self::Async(expr) => AsyncDynExpression::eval_v2(expr.as_ref(), input).await,
339        }
340    }
341
342    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
343        match self {
344            Self::Sync(expr) => expr.eval_row(input),
345            Self::Async(expr) => expr.as_ref().eval_row(input).await,
346        }
347    }
348}
349
350/// Extension trait for boxing expressions.
351///
352/// This is not directly made into expression traits because...
353/// - an expression does not have to be `'static`,
354/// - and for the ease of `auto_impl`.
355#[easy_ext::ext(SyncExpressionBoxExt)]
356impl<E: SyncExpression + 'static> E {
357    /// Wrap the expression in a [`BoxedExpression::Sync`].
358    pub fn boxed(self) -> BoxedExpression {
359        BoxedExpression::Sync(Arc::new(self))
360    }
361}
362
363/// Extension trait for boxing async expressions.
364#[easy_ext::ext(AsyncExpressionBoxExt)]
365impl<E: AsyncExpression + 'static> E {
366    /// Wrap the expression in a [`BoxedExpression::Async`].
367    pub fn boxed(self) -> BoxedExpression {
368        BoxedExpression::Async(Arc::new(self))
369    }
370}
371
372/// An type-safe wrapper that indicates the inner expression can be evaluated in a non-strict
373/// manner, i.e., developers can directly call `eval_infallible` and `eval_row_infallible` without
374/// checking the result.
375///
376/// This is usually created by non-strict build functions like [`crate::expr::build_non_strict_from_prost`]
377/// and [`crate::expr::build_func_non_strict`]. It can also be created directly by
378/// [`NonStrictExpression::new_topmost`], where only the evaluation of the topmost level expression
379/// node is non-strict and should be treated as a TODO.
380///
381/// Compared to [`crate::expr::wrapper::non_strict::NonStrict`], this is more like an indicator
382/// applied on the root of an expression tree, while the latter is a wrapper that can be applied on
383/// each node of the tree and actually changes the behavior. As a result, [`NonStrictExpression`]
384/// does not implement expression traits and instead deals directly with developers.
385#[derive(Debug)]
386pub enum NonStrictExpression {
387    Sync(Arc<dyn SyncExpression>),
388    Async(Arc<dyn AsyncDynExpression>),
389}
390
391impl NonStrictExpression {
392    /// Create a non-strict expression directly wrapping the given expression.
393    ///
394    /// Should only be used in tests as evaluation may panic.
395    pub fn for_test(inner: impl Into<BoxedExpression>) -> NonStrictExpression {
396        match inner.into() {
397            BoxedExpression::Sync(inner) => Self::Sync(inner),
398            BoxedExpression::Async(inner) => Self::Async(inner),
399        }
400    }
401
402    /// Create a non-strict expression from the given expression, where only the evaluation of the
403    /// topmost level expression node is non-strict (which is subtly different from
404    /// [`crate::expr::build_non_strict_from_prost`] where every node is non-strict).
405    ///
406    /// This should be used as a TODO.
407    pub fn new_topmost(
408        inner: impl Into<BoxedExpression>,
409        error_report: impl EvalErrorReport + 'static,
410    ) -> NonStrictExpression {
411        match inner.into() {
412            BoxedExpression::Sync(inner) => {
413                let inner = wrapper::non_strict::NonStrict::new(inner, error_report);
414                Self::Sync(Arc::new(inner))
415            }
416            BoxedExpression::Async(inner) => {
417                let inner = wrapper::non_strict::NonStrict::new(inner, error_report);
418                Self::Async(Arc::new(inner))
419            }
420        }
421    }
422
423    /// Get the return data type.
424    pub fn return_type(&self) -> DataType {
425        match self {
426            Self::Sync(expr) => expr.return_type(),
427            Self::Async(expr) => expr.return_type(),
428        }
429    }
430
431    /// Evaluate the expression in vectorized execution and assert it succeeds. Returns an array.
432    ///
433    /// Use with expressions built in non-strict mode.
434    pub async fn eval_infallible(&self, input: &DataChunk) -> ArrayRef {
435        match self {
436            Self::Sync(expr) => expr.eval(input),
437            Self::Async(expr) => AsyncDynExpression::eval(expr.as_ref(), input).await,
438        }
439        .expect("evaluation failed")
440    }
441
442    /// Evaluate the expression in row-based execution and assert it succeeds. Returns a nullable
443    /// scalar.
444    ///
445    /// Use with expressions built in non-strict mode.
446    pub async fn eval_row_infallible(&self, input: &OwnedRow) -> Datum {
447        match self {
448            Self::Sync(expr) => expr.eval_row(input),
449            Self::Async(expr) => AsyncDynExpression::eval_row(expr.as_ref(), input).await,
450        }
451        .expect("evaluation failed")
452    }
453
454    /// Unwrap the inner expression.
455    pub fn into_inner(self) -> BoxedExpression {
456        match self {
457            Self::Sync(expr) => BoxedExpression::Sync(expr),
458            Self::Async(expr) => BoxedExpression::Async(expr),
459        }
460    }
461
462    /// Get a reference to the inner expression.
463    pub fn inner(&self) -> &dyn ExpressionInfo {
464        match self {
465            Self::Sync(expr) => expr.as_ref(),
466            Self::Async(expr) => expr.as_ref(),
467        }
468    }
469}
470
471/// An optional context that can be used in a function.
472///
473/// # Example
474/// ```ignore
475/// #[function("foo(int4) -> int8")]
476/// fn foo(a: i32, ctx: &Context) -> i64 {
477///    assert_eq!(ctx.arg_types[0], DataType::Int32);
478///    assert_eq!(ctx.return_type, DataType::Int64);
479///    // ...
480/// }
481/// ```
482#[derive(Debug)]
483pub struct Context {
484    pub arg_types: Vec<DataType>,
485    pub return_type: DataType,
486    /// Whether the function is variadic.
487    pub variadic: bool,
488}
risingwave_expr/expr/mod.rs

risingwave_expr/expr/
mod.rs
