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};
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    pub literal: Row,
46    /// Index into the `exprs` vector of the corresponding result expression.
47    pub expr_index: usize,
48}
49
50/// Evaluates a CASE expression by looking up the input datum in a sorted `Vec`.
51///
52/// The input expression (`exprs[0]`) is evaluated once, packed into a temporary
53/// `Row`, and looked up in `lookup` via binary search. If found, the corresponding
54/// result expression (`exprs[idx]`) is evaluated; otherwise the fallback
55/// (`exprs.last()`) is evaluated.
56/// NULL inputs go straight to the fallback (since SQL `NULL = x` is always NULL/falsy).
57#[derive(
58    Ord,
59    PartialOrd,
60    Clone,
61    Debug,
62    Eq,
63    PartialEq,
64    Serialize,
65    Deserialize,
66    Hash
67)]
68pub struct CaseLiteral {
69    /// Sorted vec of literal-to-index entries for binary-search lookup.
70    pub lookup: Vec<CaseLiteralEntry>,
71    /// The output type of this CASE expression.
72    pub return_type: SqlColumnType,
73}
74
75impl LazyVariadicFunc for CaseLiteral {
76    fn eval<'a>(
77        &'a self,
78        datums: &[Datum<'a>],
79        temp_storage: &'a RowArena,
80        exprs: &'a [impl Eval],
81    ) -> Result<Datum<'a>, EvalError> {
82        let input = exprs[0].eval(datums, temp_storage)?;
83        // SQL NULL = x is always NULL/falsy, so go straight to the fallback.
84        if input.is_null() {
85            return exprs.last().unwrap().eval(datums, temp_storage);
86        }
87        let key = Row::pack_slice(&[input]);
88        if let Ok(pos) = self
89            .lookup
90            .binary_search_by(|entry| entry.literal.cmp(&key))
91        {
92            exprs[self.lookup[pos].expr_index].eval(datums, temp_storage)
93        } else {
94            exprs.last().unwrap().eval(datums, temp_storage)
95        }
96    }
97
98    fn output_type(&self, _input_types: &[SqlColumnType]) -> SqlColumnType {
99        self.return_type.clone()
100    }
101
102    fn propagates_nulls(&self) -> bool {
103        // NULL input goes to the fallback, not automatically to NULL output.
104        false
105    }
106
107    fn introduces_nulls(&self) -> bool {
108        // Branch results or the fallback may be NULL.
109        true
110    }
111
112    fn could_error(&self) -> bool {
113        // The function itself does not error; errors in sub-expressions are
114        // checked separately by MirScalarExpr::could_error.
115        false
116    }
117
118    fn is_monotone(&self) -> bool {
119        false
120    }
121
122    fn is_associative(&self) -> bool {
123        false
124    }
125}
126
127// Note: this Display impl is unused at runtime because CaseLiteral has
128// custom printing in src/expr/src/explain/text.rs.
129impl CaseLiteral {
130    /// Look up a key in the sorted lookup vec. Returns the expr index if found.
131    pub fn get(&self, key: &Row) -> Option<usize> {
132        self.lookup
133            .binary_search_by(|entry| entry.literal.cmp(key))
134            .ok()
135            .map(|pos| self.lookup[pos].expr_index)
136    }
137
138    /// Insert an entry, maintaining sorted order.
139    /// If the literal already exists, overwrites the index and returns the old one.
140    pub fn insert(&mut self, literal: Row, expr_index: usize) -> Option<usize> {
141        match self
142            .lookup
143            .binary_search_by(|entry| entry.literal.cmp(&literal))
144        {
145            Ok(pos) => {
146                let old = self.lookup[pos].expr_index;
147                self.lookup[pos].expr_index = expr_index;
148                Some(old)
149            }
150            Err(pos) => {
151                self.lookup.insert(
152                    pos,
153                    CaseLiteralEntry {
154                        literal,
155                        expr_index,
156                    },
157                );
158                None
159            }
160        }
161    }
162}
163
164impl fmt::Display for CaseLiteral {
165    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
166        write!(f, "case_literal[{} cases]", self.lookup.len())
167    }
168}
mz_expr/scalar/func/impls/case_literal.rs

mz_expr/scalar/func/impls/
case_literal.rs
