risingwave_expr_impl/scalar/
case.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::collections::HashMap;
16use std::sync::Arc;
17
18use risingwave_common::array::{ArrayRef, DataChunk};
19use risingwave_common::bail;
20use risingwave_common::row::{OwnedRow, Row};
21use risingwave_common::types::{DataType, Datum, ScalarImpl};
22use risingwave_expr::expr::{
23    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
24    SyncExpressionBoxExt, try_convert_all,
25};
26use risingwave_expr::{Result, build_function};
27
28#[derive(Debug)]
29struct WhenClause<E> {
30    when: E,
31    then: E,
32}
33
34#[derive(Debug)]
35struct CaseExpression<E> {
36    return_type: DataType,
37    when_clauses: Vec<WhenClause<E>>,
38    else_clause: Option<E>,
39}
40
41impl<E> CaseExpression<E> {
42    fn new(
43        return_type: DataType,
44        when_clauses: Vec<WhenClause<E>>,
45        else_clause: Option<E>,
46    ) -> Self {
47        Self {
48            return_type,
49            when_clauses,
50            else_clause,
51        }
52    }
53}
54
55impl<E: ExpressionInfo> ExpressionInfo for CaseExpression<E> {
56    fn return_type(&self) -> DataType {
57        self.return_type.clone()
58    }
59}
60
61macro_rules! eval_case {
62    ($mode:ident, $this:expr, $input:expr) => {{
63        let mut input = $input.clone();
64        let input_len = input.capacity();
65        let mut selection = vec![None; input_len];
66        let when_len = $this.when_clauses.len();
67        let mut result_array = Vec::with_capacity(when_len + 1);
68        for (when_idx, WhenClause { when, then }) in $this.when_clauses.iter().enumerate() {
69            let input_vis = input.visibility().clone();
70            // note that evaluated result from when clause may contain bits that are not visible,
71            // so we need to mask it with input visibility.
72            let calc_then_vis = risingwave_expr::forward!($mode, when, eval(&input))?
73                .as_bool()
74                .to_bitmap()
75                & &input_vis;
76            input.set_visibility(calc_then_vis.clone());
77            let then_res = risingwave_expr::forward!($mode, then, eval(&input))?;
78            calc_then_vis
79                .iter_ones()
80                .for_each(|pos| selection[pos] = Some(when_idx));
81            input.set_visibility(&input_vis & (!calc_then_vis));
82            result_array.push(then_res);
83        }
84        if let Some(ref else_expr) = $this.else_clause {
85            let else_res = risingwave_expr::forward!($mode, else_expr, eval(&input))?;
86            input
87                .visibility()
88                .iter_ones()
89                .for_each(|pos| selection[pos] = Some(when_len));
90            result_array.push(else_res);
91        }
92        let mut builder = $this.return_type().create_array_builder(input.capacity());
93        for (i, sel) in selection.into_iter().enumerate() {
94            if let Some(when_idx) = sel {
95                builder.append(result_array[when_idx].value_at(i));
96            } else {
97                builder.append_null();
98            }
99        }
100        Ok(Arc::new(builder.finish()))
101    }};
102}
103
104macro_rules! eval_row_case {
105    ($mode:ident, $this:expr, $input:expr) => {{
106        for WhenClause { when, then } in &$this.when_clauses {
107            if risingwave_expr::forward!($mode, when, eval_row($input))?
108                .is_some_and(|w| w.into_bool())
109            {
110                return risingwave_expr::forward!($mode, then, eval_row($input));
111            }
112        }
113        if let Some(ref else_expr) = $this.else_clause {
114            risingwave_expr::forward!($mode, else_expr, eval_row($input))
115        } else {
116            Ok(None)
117        }
118    }};
119}
120
121impl<E: SyncExpression> SyncExpression for CaseExpression<E> {
122    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
123        eval_case!(sync, self, input)
124    }
125
126    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
127        eval_row_case!(sync, self, input)
128    }
129}
130
131impl<E: AsyncExpression> AsyncExpression for CaseExpression<E> {
132    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
133        eval_case!(async, self, input)
134    }
135
136    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
137        eval_row_case!(async, self, input)
138    }
139}
140
141/// With large scale of simple form match arms in case-when expression,
142/// we could optimize the `CaseExpression` to `ConstantLookupExpression`,
143/// which could significantly facilitate the evaluation of case-when.
144#[derive(Debug)]
145struct ConstantLookupExpression<E> {
146    return_type: DataType,
147    arms: HashMap<ScalarImpl, E>,
148    fallback: Option<E>,
149    /// `operand` must exist at present
150    operand: E,
151}
152
153impl<E> ConstantLookupExpression<E> {
154    fn new(
155        return_type: DataType,
156        arms: HashMap<ScalarImpl, E>,
157        fallback: Option<E>,
158        operand: E,
159    ) -> Self {
160        Self {
161            return_type,
162            arms,
163            fallback,
164            operand,
165        }
166    }
167}
168
169impl<E: ExpressionInfo> ExpressionInfo for ConstantLookupExpression<E> {
170    fn return_type(&self) -> DataType {
171        self.return_type.clone()
172    }
173}
174
175macro_rules! eval_fallback_lookup {
176    ($mode:ident, $this:expr, $input:expr) => {{
177        if let Some(ref fallback) = $this.fallback {
178            let Ok(res) = risingwave_expr::forward!($mode, fallback, eval_row($input)) else {
179                bail!("failed to evaluate the input for fallback arm");
180            };
181            Ok::<Datum, risingwave_expr::ExprError>(res)
182        } else {
183            Ok::<Datum, risingwave_expr::ExprError>(None)
184        }
185    }};
186}
187
188macro_rules! lookup {
189    ($mode:ident, $this:expr, $datum:expr, $input:expr) => {{
190        match $datum.as_ref() {
191            Some(datum) => {
192                if let Some(expr) = $this.arms.get(datum) {
193                    let Ok(res) = risingwave_expr::forward!($mode, expr, eval_row($input)) else {
194                        bail!("failed to evaluate the input for normal arm");
195                    };
196                    Ok(res)
197                } else {
198                    eval_fallback_lookup!($mode, $this, $input)
199                }
200            }
201            None => eval_fallback_lookup!($mode, $this, $input),
202        }
203    }};
204}
205
206macro_rules! eval_constant_lookup {
207    ($mode:ident, $this:expr, $input:expr) => {{
208        let input_len = $input.capacity();
209        let mut builder = $this.return_type().create_array_builder(input_len);
210
211        // Evaluate the input DataChunk at first
212        let eval_result = risingwave_expr::forward!($mode, $this.operand, eval($input))?;
213
214        for i in 0..input_len {
215            let datum = eval_result.datum_at(i);
216            let (row, vis) = $input.row_at(i);
217
218            // Check for visibility
219            if !vis {
220                builder.append_null();
221                continue;
222            }
223
224            // Note that the `owned_row` here is extracted from input
225            // rather than from `eval_result`
226            let owned_row = row.into_owned_row();
227
228            // Lookup and evaluate the current input datum
229            if let Ok(datum) = lookup!($mode, $this, datum, &owned_row) {
230                builder.append(datum.as_ref());
231            } else {
232                bail!("failed to lookup and evaluate the expression in `eval`");
233            }
234        }
235
236        Ok(Arc::new(builder.finish()))
237    }};
238}
239
240impl<E: SyncExpression> SyncExpression for ConstantLookupExpression<E> {
241    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
242        eval_constant_lookup!(sync, self, input)
243    }
244
245    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
246        let datum = self.operand.eval_row(input)?;
247        lookup!(sync, self, datum, input)
248    }
249}
250
251impl<E: AsyncExpression> AsyncExpression for ConstantLookupExpression<E> {
252    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
253        eval_constant_lookup!(async, self, input)
254    }
255
256    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
257        let datum = self.operand.eval_row(input).await?;
258        lookup!(async, self, datum, input)
259    }
260}
261
262#[build_function("constant_lookup(...) -> any", type_infer = "unreachable")]
263fn build_constant_lookup_expr(
264    return_type: DataType,
265    children: Vec<BoxedExpression>,
266) -> Result<BoxedExpression> {
267    if children.is_empty() {
268        bail!("children expression must not be empty for constant lookup expression");
269    }
270
271    let mut children = children;
272
273    let operand = children.remove(0);
274
275    let mut arms = HashMap::new();
276
277    // Build the `arms` with iterating over `when` & `then` clauses
278    let mut iter = children.into_iter().array_chunks();
279    for [when, then] in iter.by_ref() {
280        let Ok(Some(s)) = when.eval_const() else {
281            bail!("expect when expression to be const");
282        };
283        arms.insert(s, then);
284    }
285
286    let fallback = if let Some(else_clause) = iter.into_remainder().next() {
287        if else_clause.return_type() != return_type {
288            bail!("Type mismatched between else and case.");
289        }
290        Some(else_clause)
291    } else {
292        None
293    };
294
295    let BoxedExpression::Sync(operand) = operand else {
296        return Ok(ConstantLookupExpression::new(return_type, arms, fallback, operand).boxed());
297    };
298    let arms: Vec<_> = arms.into_iter().collect();
299    let sync_arms: HashMap<ScalarImpl, Arc<dyn SyncExpression>> = match try_convert_all(
300        arms,
301        |(key, expr)| match expr {
302            BoxedExpression::Sync(expr) => Ok((key, expr)),
303            expr @ BoxedExpression::Async(_) => Err((key, expr)),
304        },
305        |(key, expr)| (key, BoxedExpression::Sync(expr)),
306    ) {
307        Ok(arms) => arms.into_iter().collect(),
308        Err(arms) => {
309            return Ok(ConstantLookupExpression::new(
310                return_type,
311                arms.into_iter().collect(),
312                fallback,
313                BoxedExpression::Sync(operand),
314            )
315            .boxed());
316        }
317    };
318    let fallback = match fallback {
319        Some(BoxedExpression::Sync(expr)) => Some(expr),
320        Some(expr @ BoxedExpression::Async(_)) => {
321            let arms = sync_arms
322                .into_iter()
323                .map(|(key, expr)| (key, BoxedExpression::Sync(expr)))
324                .collect();
325            return Ok(ConstantLookupExpression::new(
326                return_type,
327                arms,
328                Some(expr),
329                BoxedExpression::Sync(operand),
330            )
331            .boxed());
332        }
333        None => None,
334    };
335    Ok(ConstantLookupExpression::new(return_type, sync_arms, fallback, operand).boxed())
336}
337
338#[build_function("case(...) -> any", type_infer = "unreachable")]
339fn build_case_expr(
340    return_type: DataType,
341    children: Vec<BoxedExpression>,
342) -> Result<BoxedExpression> {
343    // children: (when, then)+, (else_clause)?
344    let len = children.len();
345    let mut when_clauses = Vec::with_capacity(len / 2);
346    let mut iter = children.into_iter().array_chunks();
347    for [when, then] in iter.by_ref() {
348        if when.return_type() != DataType::Boolean {
349            bail!("Type mismatched between when clause and condition");
350        }
351        if then.return_type() != return_type {
352            bail!("Type mismatched between then clause and case");
353        }
354        when_clauses.push(WhenClause { when, then });
355    }
356    let else_clause = if let Some(else_clause) = iter.into_remainder().next() {
357        if else_clause.return_type() != return_type {
358            bail!("Type mismatched between else and case.");
359        }
360        Some(else_clause)
361    } else {
362        None
363    };
364
365    let sync_when_clauses = match try_convert_all(
366        when_clauses,
367        |WhenClause { when, then }| match (when, then) {
368            (BoxedExpression::Sync(when), BoxedExpression::Sync(then)) => {
369                Ok(WhenClause { when, then })
370            }
371            (when, then) => Err(WhenClause { when, then }),
372        },
373        |WhenClause { when, then }| WhenClause {
374            when: BoxedExpression::Sync(when),
375            then: BoxedExpression::Sync(then),
376        },
377    ) {
378        Ok(when_clauses) => when_clauses,
379        Err(when_clauses) => {
380            return Ok(CaseExpression::new(return_type, when_clauses, else_clause).boxed());
381        }
382    };
383    let else_clause = match else_clause {
384        Some(BoxedExpression::Sync(expr)) => Some(expr),
385        Some(expr @ BoxedExpression::Async(_)) => {
386            let when_clauses = sync_when_clauses
387                .into_iter()
388                .map(|WhenClause { when, then }| WhenClause {
389                    when: BoxedExpression::Sync(when),
390                    then: BoxedExpression::Sync(then),
391                })
392                .collect();
393            return Ok(CaseExpression::new(return_type, when_clauses, Some(expr)).boxed());
394        }
395        None => None,
396    };
397    Ok(CaseExpression::new(return_type, sync_when_clauses, else_clause).boxed())
398}
399
400#[cfg(test)]
401mod tests {
402    use risingwave_common::test_prelude::DataChunkTestExt;
403    use risingwave_common::types::ToOwnedDatum;
404    use risingwave_common::util::iter_util::ZipEqDebug;
405    use risingwave_expr::expr::build_from_pretty;
406
407    use super::*;
408
409    #[tokio::test]
410    async fn test_eval_searched_case() {
411        // when x then 1 else 2
412        let case = build_from_pretty("(case:int4 $0:boolean 1:int4 2:int4)");
413        let (input, expected) = DataChunk::from_pretty(
414            "B i
415             t 1
416             f 2
417             t 1
418             t 1
419             f 2",
420        )
421        .split_column_at(1);
422
423        // test eval
424        let output = case.eval(&input).await.unwrap();
425        assert_eq!(&output, expected.column_at(0));
426
427        // test eval_row
428        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
429            let result = case.eval_row(&row.to_owned_row()).await.unwrap();
430            assert_eq!(result, expected.datum_at(0).to_owned_datum());
431        }
432    }
433
434    #[tokio::test]
435    async fn test_eval_without_else() {
436        // when x then 1 when y then 2
437        let case = build_from_pretty("(case:int4 $0:boolean 1:int4 $1:boolean 2:int4)");
438        let (input, expected) = DataChunk::from_pretty(
439            "B B i
440             f f .
441             f t 2
442             t f 1
443             t t 1",
444        )
445        .split_column_at(2);
446
447        // test eval
448        let output = case.eval(&input).await.unwrap();
449        assert_eq!(&output, expected.column_at(0));
450
451        // test eval_row
452        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
453            let result = case.eval_row(&row.to_owned_row()).await.unwrap();
454            assert_eq!(result, expected.datum_at(0).to_owned_datum());
455        }
456    }
457}
risingwave_expr_impl/scalar/case.rs

risingwave_expr_impl/scalar/
case.rs
