mz_expr/scalar/func/impls/
case_literal.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A lookup-based evaluation of `CASE expr WHEN lit1 THEN res1 ... ELSE els END`.
11//!
12//! [`CaseLiteral`] replaces chains of `If(Eq(expr, literal), result, If(...))`
13//! with a sorted `Vec` + binary-search lookup, turning O(n) evaluation into O(log n).
14//!
15//! Represented as a `CallVariadic { func: CaseLiteral { lookup, return_type }, exprs }`
16//! where:
17//! * `exprs[0]` = input expression (the `x` in `CASE x WHEN ...`)
18//! * `exprs[1..n]` = case result expressions
19//! * `exprs[last]` = `els` (fallback)
20//! * `lookup: Vec<CaseLiteralEntry>` maps literal values to indices in `exprs` (sorted by `Row`)
21
22use std::fmt;
23
24use mz_repr::{Datum, Row, RowArena, SqlColumnType, StableRow};
25use serde::{Deserialize, Serialize};
26
27use crate::scalar::func::variadic::LazyVariadicFunc;
28use crate::{Eval, EvalError};
29
30/// A single entry in a [`CaseLiteral`] lookup table: a literal `Row` value
31/// paired with the index of the corresponding result expression in `exprs`.
32#[derive(
33    Ord,
34    PartialOrd,
35    Clone,
36    Debug,
37    Eq,
38    PartialEq,
39    Serialize,
40    Deserialize,
41    Hash
42)]
43pub struct CaseLiteralEntry {
44    /// The literal value (as a single-datum `Row`).
45    ///
46    /// A [`StableRow`] because `CaseLiteral` is part of the stable LIR
47    /// serialization surface, where raw `Row` bytes must not appear.
48    pub literal: StableRow,
49    /// Index into the `exprs` vector of the corresponding result expression.
50    pub expr_index: usize,
51}
52
53/// Evaluates a CASE expression by looking up the input datum in a sorted `Vec`.
54///
55/// The input expression (`exprs[0]`) is evaluated once, packed into a temporary
56/// `Row`, and looked up in `lookup` via binary search. If found, the corresponding
57/// result expression (`exprs[idx]`) is evaluated; otherwise the fallback
58/// (`exprs.last()`) is evaluated.
59/// NULL inputs go straight to the fallback (since SQL `NULL = x` is always NULL/falsy).
60#[derive(
61    Ord,
62    PartialOrd,
63    Clone,
64    Debug,
65    Eq,
66    PartialEq,
67    Serialize,
68    Deserialize,
69    Hash
70)]
71pub struct CaseLiteral {
72    /// Sorted vec of literal-to-index entries for binary-search lookup.
73    pub lookup: Vec<CaseLiteralEntry>,
74    /// The output type of this CASE expression.
75    pub return_type: SqlColumnType,
76}
77
78impl LazyVariadicFunc for CaseLiteral {
79    fn eval<'a>(
80        &'a self,
81        datums: &[Datum<'a>],
82        temp_storage: &'a RowArena,
83        exprs: &'a [impl Eval],
84    ) -> Result<Datum<'a>, EvalError> {
85        let input = exprs[0].eval(datums, temp_storage)?;
86        // SQL NULL = x is always NULL/falsy, so go straight to the fallback.
87        if input.is_null() {
88            return exprs.last().unwrap().eval(datums, temp_storage);
89        }
90        let key = Row::pack_slice(&[input]);
91        if let Ok(pos) = self
92            .lookup
93            .binary_search_by(|entry| entry.literal.0.cmp(&key))
94        {
95            exprs[self.lookup[pos].expr_index].eval(datums, temp_storage)
96        } else {
97            exprs.last().unwrap().eval(datums, temp_storage)
98        }
99    }
100
101    fn output_type(&self, _input_types: &[SqlColumnType]) -> SqlColumnType {
102        self.return_type.clone()
103    }
104
105    fn propagates_nulls(&self) -> bool {
106        // NULL input goes to the fallback, not automatically to NULL output.
107        false
108    }
109
110    fn introduces_nulls(&self) -> bool {
111        // Branch results or the fallback may be NULL.
112        true
113    }
114
115    fn could_error(&self) -> bool {
116        // The function itself does not error; errors in sub-expressions are
117        // checked separately by MirScalarExpr::could_error.
118        false
119    }
120
121    fn is_monotone(&self) -> bool {
122        false
123    }
124
125    fn is_associative(&self) -> bool {
126        false
127    }
128}
129
130// Note: this Display impl is unused at runtime because CaseLiteral has
131// custom printing in src/expr/src/explain/text.rs.
132impl CaseLiteral {
133    /// Look up a key in the sorted lookup vec. Returns the expr index if found.
134    pub fn get(&self, key: &Row) -> Option<usize> {
135        self.lookup
136            .binary_search_by(|entry| entry.literal.0.cmp(key))
137            .ok()
138            .map(|pos| self.lookup[pos].expr_index)
139    }
140
141    /// Insert an entry, maintaining sorted order.
142    /// If the literal already exists, overwrites the index and returns the old one.
143    pub fn insert(&mut self, literal: Row, expr_index: usize) -> Option<usize> {
144        match self
145            .lookup
146            .binary_search_by(|entry| entry.literal.0.cmp(&literal))
147        {
148            Ok(pos) => {
149                let old = self.lookup[pos].expr_index;
150                self.lookup[pos].expr_index = expr_index;
151                Some(old)
152            }
153            Err(pos) => {
154                self.lookup.insert(
155                    pos,
156                    CaseLiteralEntry {
157                        literal: StableRow(literal),
158                        expr_index,
159                    },
160                );
161                None
162            }
163        }
164    }
165}
166
167impl fmt::Display for CaseLiteral {
168    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
169        write!(f, "case_literal[{} cases]", self.lookup.len())
170    }
171}
mz_expr/scalar/func/impls/case_literal.rs

mz_expr/scalar/func/impls/
case_literal.rs
