Skip to main content

mz_expr/scalar/func/
macros.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// Convenience macro for generating `inverse` values.
11macro_rules! to_unary {
12    ($f:expr) => {
13        Some(crate::UnaryFunc::from($f))
14    };
15}
16
17#[cfg(test)]
18mod test {
19    use mz_expr_derive::sqlfunc;
20    use mz_repr::SqlScalarType;
21
22    use crate::EvalError;
23    use crate::scalar::func::LazyUnaryFunc;
24
25    #[sqlfunc(sqlname = "INFALLIBLE")]
26    fn infallible1(a: f32) -> f32 {
27        a
28    }
29
30    #[sqlfunc]
31    fn infallible2(a: Option<f32>) -> f32 {
32        a.unwrap_or_default()
33    }
34
35    #[sqlfunc]
36    fn infallible3(a: f32) -> Option<f32> {
37        Some(a)
38    }
39
40    #[mz_ore::test]
41    fn elision_rules_infallible() {
42        assert_eq!(format!("{}", Infallible1), "INFALLIBLE");
43        assert!(Infallible1.propagates_nulls());
44        assert!(!Infallible1.introduces_nulls());
45
46        assert!(!Infallible2.propagates_nulls());
47        assert!(!Infallible2.introduces_nulls());
48
49        assert!(Infallible3.propagates_nulls());
50        assert!(Infallible3.introduces_nulls());
51    }
52
53    #[mz_ore::test]
54    fn output_types_infallible() {
55        assert_eq!(
56            Infallible1.output_sql_type(SqlScalarType::Float32.nullable(true)),
57            SqlScalarType::Float32.nullable(true)
58        );
59        assert_eq!(
60            Infallible1.output_sql_type(SqlScalarType::Float32.nullable(false)),
61            SqlScalarType::Float32.nullable(false)
62        );
63
64        assert_eq!(
65            Infallible2.output_sql_type(SqlScalarType::Float32.nullable(true)),
66            SqlScalarType::Float32.nullable(false)
67        );
68        assert_eq!(
69            Infallible2.output_sql_type(SqlScalarType::Float32.nullable(false)),
70            SqlScalarType::Float32.nullable(false)
71        );
72
73        assert_eq!(
74            Infallible3.output_sql_type(SqlScalarType::Float32.nullable(true)),
75            SqlScalarType::Float32.nullable(true)
76        );
77        assert_eq!(
78            Infallible3.output_sql_type(SqlScalarType::Float32.nullable(false)),
79            SqlScalarType::Float32.nullable(true)
80        );
81    }
82
83    #[sqlfunc]
84    fn fallible1(a: f32) -> Result<f32, EvalError> {
85        Ok(a)
86    }
87
88    #[sqlfunc]
89    fn fallible2(a: Option<f32>) -> Result<f32, EvalError> {
90        Ok(a.unwrap_or_default())
91    }
92
93    #[sqlfunc]
94    fn fallible3(a: f32) -> Result<Option<f32>, EvalError> {
95        Ok(Some(a))
96    }
97
98    #[mz_ore::test]
99    fn elision_rules_fallible() {
100        assert!(Fallible1.propagates_nulls());
101        assert!(!Fallible1.introduces_nulls());
102
103        assert!(!Fallible2.propagates_nulls());
104        assert!(!Fallible2.introduces_nulls());
105
106        assert!(Fallible3.propagates_nulls());
107        assert!(Fallible3.introduces_nulls());
108    }
109
110    #[mz_ore::test]
111    fn output_types_fallible() {
112        assert_eq!(
113            Fallible1.output_sql_type(SqlScalarType::Float32.nullable(true)),
114            SqlScalarType::Float32.nullable(true)
115        );
116        assert_eq!(
117            Fallible1.output_sql_type(SqlScalarType::Float32.nullable(false)),
118            SqlScalarType::Float32.nullable(false)
119        );
120
121        assert_eq!(
122            Fallible2.output_sql_type(SqlScalarType::Float32.nullable(true)),
123            SqlScalarType::Float32.nullable(false)
124        );
125        assert_eq!(
126            Fallible2.output_sql_type(SqlScalarType::Float32.nullable(false)),
127            SqlScalarType::Float32.nullable(false)
128        );
129
130        assert_eq!(
131            Fallible3.output_sql_type(SqlScalarType::Float32.nullable(true)),
132            SqlScalarType::Float32.nullable(true)
133        );
134        assert_eq!(
135            Fallible3.output_sql_type(SqlScalarType::Float32.nullable(false)),
136            SqlScalarType::Float32.nullable(true)
137        );
138    }
139}
140
141/// Temporary macro that generates the equivalent of what enum_dispatch will do in the future. We
142/// need this manual macro implementation to delegate to the previous manual implementation for
143/// variants that use the old definitions.
144///
145/// Once everything is handled by this macro we can remove it and replace it with `enum_dispatch`
146///
147/// Variants marked `(E)` have payload structs generic over the stored expression type. The
148/// marker must be the literal ident `E`, matching the generic parameter the macro declares.
149macro_rules! derive_unary {
150    // Internal rules that rebuild one variant payload during expression
151    // conversion. Payloads without a marker are expression free and clone.
152    (@try_map_payload $f:ident) => { $f.clone() };
153    (@try_map_payload $f:ident $marker:ident) => { $f.try_map_expr()? };
154    (@map_payload $f:ident) => { $f.clone() };
155    (@map_payload $f:ident $marker:ident) => { $f.map_expr() };
156    // Internal rules that name a payload type at the MIR instantiation, for
157    // use in the impl block where `E` is not in scope.
158    (@mir_ty $name:ident) => { $name };
159    (@mir_ty $name:ident $marker:ident) => { $name<crate::MirScalarExpr> };
160    ($($name:ident $(($marker:ident))?),* $(,)?) => {
161        #[derive(
162            Ord, PartialOrd, Clone, Debug, Eq, PartialEq,
163            serde::Serialize, serde::Deserialize, Hash,
164                )]
165        pub enum UnaryFunc<E = crate::MirScalarExpr> {
166            $($name($name $(<$marker>)?),)*
167        }
168
169        // The `Eval` bound is required by the expression bearing payloads'
170        // `LazyUnaryFunc` impls, even for methods that never evaluate.
171        impl<E: Eval> UnaryFunc<E> {
172            pub fn eval<'a>(
173                &'a self,
174                datums: &[Datum<'a>],
175                temp_storage: &'a RowArena,
176                a: &'a impl Eval,
177            ) -> Result<Datum<'a>, EvalError> {
178                match self {
179                    $(Self::$name(f) => f.eval(datums, temp_storage, a),)*
180                }
181            }
182
183            pub fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
184                match self {
185                    $(Self::$name(f) => LazyUnaryFunc::output_sql_type(f, input_type),)*
186                }
187            }
188            pub fn output_type(&self, input_type: ReprColumnType) -> ReprColumnType {
189                match self {
190                    $(Self::$name(f) => LazyUnaryFunc::output_type(f, input_type),)*
191                }
192            }
193            pub fn propagates_nulls(&self) -> bool {
194                match self {
195                    $(Self::$name(f) => LazyUnaryFunc::propagates_nulls(f),)*
196                }
197            }
198            pub fn introduces_nulls(&self) -> bool {
199                match self {
200                    $(Self::$name(f) => LazyUnaryFunc::introduces_nulls(f),)*
201                }
202            }
203            pub fn preserves_uniqueness(&self) -> bool {
204                match self {
205                    $(Self::$name(f) => LazyUnaryFunc::preserves_uniqueness(f),)*
206                }
207            }
208            pub fn is_monotone(&self) -> bool {
209                match self {
210                    $(Self::$name(f) => LazyUnaryFunc::is_monotone(f),)*
211                }
212            }
213            pub fn could_error(&self) -> bool {
214                match self {
215                    $(Self::$name(f) => LazyUnaryFunc::could_error(f),)*
216                }
217            }
218            pub fn is_eliminable_cast(&self) -> bool {
219                match self {
220                    $(Self::$name(f) => LazyUnaryFunc::is_eliminable_cast(f),)*
221                }
222            }
223        }
224
225        impl<E> UnaryFunc<E> {
226            /// The canonical name of this variant, as declared by its
227            /// [`FuncName`](crate::func::FuncName) impl.
228            pub fn variant_name(&self) -> &'static str {
229                match self {
230                    $(Self::$name(_) =>
231                        <$name $(<$marker>)? as crate::func::FuncName>::NAME,)*
232                }
233            }
234
235            /// Rebuilds this function with any stored expressions converted to
236            /// `E2`. Fails if any expression conversion fails, reporting the
237            /// first failure.
238            pub fn try_map_expr<'a, E2: TryFrom<&'a E>>(
239                &'a self,
240            ) -> Result<UnaryFunc<E2>, E2::Error> {
241                Ok(match self {
242                    $(Self::$name(f) => UnaryFunc::$name(
243                        derive_unary!(@try_map_payload f $($marker)?),
244                    ),)*
245                })
246            }
247
248            /// Rebuilds this function with any stored expressions converted to
249            /// `E2`.
250            pub fn map_expr<'a, E2: From<&'a E>>(&'a self) -> UnaryFunc<E2> {
251                match self {
252                    $(Self::$name(f) => UnaryFunc::$name(
253                        derive_unary!(@map_payload f $($marker)?),
254                    ),)*
255                }
256            }
257        }
258
259        // Methods that only exist at the MIR instantiation. `inverse` is an
260        // optimizer concept and the name lookup helpers serve MIR test
261        // tooling.
262        impl UnaryFunc {
263            pub fn inverse(&self) -> Option<UnaryFunc> {
264                match self {
265                    $(Self::$name(f) => LazyUnaryFunc::inverse(f),)*
266                }
267            }
268
269            /// Attempts to construct a `UnaryFunc` from the canonical name of
270            /// one of its variants, as declared by the variant's
271            /// [`FuncName`](crate::func::FuncName) impl (the name of the
272            /// underlying Rust function, e.g. `cast_int32_to_numeric`).
273            ///
274            /// Intended for test tooling that needs to name exact function
275            /// variants. Only variants whose inner function deserializes from
276            /// no data (unit functions and functions all of whose parameters
277            /// may default to `None`) are constructible. Returns `None` for
278            /// other variants and for unknown names. Returns the MIR
279            /// instantiation.
280            pub fn from_variant_name(name: &str) -> Option<Self> {
281                use crate::func::FuncName;
282                $(
283                    if name == <derive_unary!(@mir_ty $name $($marker)?) as FuncName>::NAME {
284                        return serde_json::from_value(
285                            serde_json::json!({ stringify!($name): null }),
286                        )
287                        .ok();
288                    }
289                )*
290                None
291            }
292
293            /// The canonical names of all variants, in declaration order.
294            pub fn variant_names() -> impl Iterator<Item = &'static str> {
295                use crate::func::FuncName;
296                [$(
297                    <derive_unary!(@mir_ty $name $($marker)?) as FuncName>::NAME,
298                )*]
299                .into_iter()
300            }
301        }
302
303        impl<E> fmt::Display for UnaryFunc<E> {
304            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
305                match self {
306                    $(Self::$name(func) => func.fmt(f),)*
307                }
308            }
309        }
310
311        $(
312            impl<E> From<$name $(<$marker>)?> for crate::UnaryFunc<E> {
313                fn from(variant: $name $(<$marker>)?) -> Self {
314                    Self::$name(variant)
315                }
316            }
317        )*
318    }
319}
320
321/// Generates the `VariadicFunc` enum, its `impl` block,
322/// `Display` impl, and `From<InnerType>` impls for each variant.
323///
324/// All variants must use explicit `Name(Type)` syntax. When the variant name equals
325/// the inner type name, write e.g. `And(And)`.
326macro_rules! derive_variadic {
327    ($($name:ident ( $variant:ident )),* $(,)?) => {
328        #[derive(
329            Ord, PartialOrd, Clone, Debug, Eq, PartialEq,
330            serde::Serialize, serde::Deserialize, Hash,
331        )]
332        pub enum VariadicFunc {
333            $($name($variant),)*
334        }
335
336        impl VariadicFunc {
337            pub fn eval<'a>(
338                &'a self,
339                datums: &[Datum<'a>],
340                temp_storage: &'a RowArena,
341                exprs: &'a [impl Eval],
342            ) -> Result<Datum<'a>, EvalError> {
343                match self {
344                    $(Self::$name(f) => f.eval(datums, temp_storage, exprs),)*
345                }
346            }
347
348            pub fn output_sql_type(&self, input_types: Vec<SqlColumnType>) -> SqlColumnType {
349                match self {
350                    $(Self::$name(f) => LazyVariadicFunc::output_type(f, &input_types),)*
351                }
352            }
353
354            /// Computes the representation type of this variadic function.
355            ///
356            /// Wrapper around [`Self::output_sql_type`] that converts to representation types.
357            pub fn output_type(&self, input_types: Vec<ReprColumnType>) -> ReprColumnType {
358                let sql_types = input_types.iter().map(SqlColumnType::from_repr).collect();
359                ReprColumnType::from(&self.output_sql_type(sql_types))
360            }
361
362            pub fn propagates_nulls(&self) -> bool {
363                match self {
364                    $(Self::$name(f) => LazyVariadicFunc::propagates_nulls(f),)*
365                }
366            }
367
368            pub fn introduces_nulls(&self) -> bool {
369                match self {
370                    $(Self::$name(f) => LazyVariadicFunc::introduces_nulls(f),)*
371                }
372            }
373
374            pub fn could_error(&self) -> bool {
375                match self {
376                    $(Self::$name(f) => LazyVariadicFunc::could_error(f),)*
377                }
378            }
379
380            pub fn is_monotone(&self) -> bool {
381                match self {
382                    $(Self::$name(f) => LazyVariadicFunc::is_monotone(f),)*
383                }
384            }
385
386            pub fn is_associative(&self) -> bool {
387                match self {
388                    $(Self::$name(f) => LazyVariadicFunc::is_associative(f),)*
389                }
390            }
391
392            pub fn is_infix_op(&self) -> bool {
393                match self {
394                    $(Self::$name(f) => LazyVariadicFunc::is_infix_op(f),)*
395                }
396            }
397
398            /// The canonical name of this variant, as declared by its
399            /// [`FuncName`](crate::func::FuncName) impl.
400            pub fn variant_name(&self) -> &'static str {
401                match self {
402                    $(Self::$name(_) => <$variant as crate::func::FuncName>::NAME,)*
403                }
404            }
405
406            /// Attempts to construct a `VariadicFunc` from the canonical name
407            /// of one of its variants, as declared by the variant's
408            /// [`FuncName`](crate::func::FuncName) impl (the name of the
409            /// underlying Rust function, e.g. `record_create`).
410            ///
411            /// Intended for test tooling that needs to name exact function
412            /// variants. Only variants whose inner function deserializes from
413            /// no data (unit functions and functions all of whose parameters
414            /// may default to `None`) are constructible. Returns `None` for
415            /// other variants and for unknown names.
416            pub fn from_variant_name(name: &str) -> Option<Self> {
417                $(
418                    if name == <$variant as crate::func::FuncName>::NAME {
419                        return serde_json::from_value(
420                            serde_json::json!({ stringify!($name): null }),
421                        )
422                        .ok();
423                    }
424                )*
425                None
426            }
427
428            /// The canonical names of all variants, in declaration order.
429            pub fn variant_names() -> impl Iterator<Item = &'static str> {
430                [$(<$variant as crate::func::FuncName>::NAME,)*].into_iter()
431            }
432        }
433
434        impl fmt::Display for VariadicFunc {
435            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
436                match self {
437                    $(Self::$name(func) => func.fmt(f),)*
438                }
439            }
440        }
441
442        $(
443            impl From<$variant> for crate::VariadicFunc {
444                fn from(variant: $variant) -> Self {
445                    Self::$name(variant)
446                }
447            }
448        )*
449    }
450}
451
452/// Generates the `BinaryFunc` enum, its `impl` block (with 8 delegating methods),
453/// `Display` impl, and `From<InnerType>` impls for each variant.
454///
455/// All variants must use explicit `Name(Type)` syntax. When the variant name equals
456/// the inner type name, write e.g. `AddInt16(AddInt16)`.
457macro_rules! derive_binary {
458    ($($name:ident ( $variant:ident )),* $(,)?) => {
459        #[derive(
460            Ord, PartialOrd, Clone, Debug, Eq, PartialEq,
461            serde::Serialize, serde::Deserialize, Hash,
462                )]
463        pub enum BinaryFunc {
464            $($name($variant),)*
465        }
466
467        impl BinaryFunc {
468            pub fn eval<'a>(
469                &'a self,
470                datums: &[Datum<'a>],
471                temp_storage: &'a RowArena,
472                exprs: &[&'a impl Eval],
473            ) -> Result<Datum<'a>, EvalError> {
474                match self {
475                    $(Self::$name(f) => f.eval(datums, temp_storage, exprs),)*
476                }
477            }
478
479            pub fn output_sql_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
480                match self {
481                    $(Self::$name(f) => {
482                        LazyBinaryFunc::output_sql_type(f, input_types)
483                    },)*
484                }
485            }
486
487            pub fn output_type(&self, input_types: &[ReprColumnType]) -> ReprColumnType {
488                match self {
489                    $(Self::$name(f) => LazyBinaryFunc::output_type(f, input_types),)*
490                }
491            }
492
493            pub fn propagates_nulls(&self) -> bool {
494                match self {
495                    $(Self::$name(f) => LazyBinaryFunc::propagates_nulls(f),)*
496                }
497            }
498
499            pub fn introduces_nulls(&self) -> bool {
500                match self {
501                    $(Self::$name(f) => LazyBinaryFunc::introduces_nulls(f),)*
502                }
503            }
504
505            pub fn is_infix_op(&self) -> bool {
506                match self {
507                    $(Self::$name(f) => LazyBinaryFunc::is_infix_op(f),)*
508                }
509            }
510
511            pub fn negate(&self) -> Option<BinaryFunc> {
512                match self {
513                    $(Self::$name(f) => LazyBinaryFunc::negate(f),)*
514                }
515            }
516
517            pub fn could_error(&self) -> bool {
518                match self {
519                    $(Self::$name(f) => LazyBinaryFunc::could_error(f),)*
520                }
521            }
522
523            pub fn is_monotone(&self) -> (bool, bool) {
524                match self {
525                    $(Self::$name(f) => LazyBinaryFunc::is_monotone(f),)*
526                }
527            }
528
529            pub fn is_infinity_monotone(&self) -> bool {
530                match self {
531                    $(Self::$name(f) => LazyBinaryFunc::is_infinity_monotone(f),)*
532                }
533            }
534
535            /// The canonical name of this variant, as declared by its
536            /// [`FuncName`](crate::func::FuncName) impl.
537            pub fn variant_name(&self) -> &'static str {
538                match self {
539                    $(Self::$name(_) => <$variant as crate::func::FuncName>::NAME,)*
540                }
541            }
542
543            /// Attempts to construct a `BinaryFunc` from the canonical name of
544            /// one of its variants, as declared by the variant's
545            /// [`FuncName`](crate::func::FuncName) impl (the name of the
546            /// underlying Rust function, e.g. `add_int32`).
547            ///
548            /// Intended for test tooling that needs to name exact function
549            /// variants. Only variants whose inner function deserializes from
550            /// no data (unit functions and functions all of whose parameters
551            /// may default to `None`) are constructible. Returns `None` for
552            /// other variants and for unknown names.
553            pub fn from_variant_name(name: &str) -> Option<Self> {
554                $(
555                    if name == <$variant as crate::func::FuncName>::NAME {
556                        return serde_json::from_value(
557                            serde_json::json!({ stringify!($name): null }),
558                        )
559                        .ok();
560                    }
561                )*
562                None
563            }
564
565            /// The canonical names of all variants, in declaration order.
566            pub fn variant_names() -> impl Iterator<Item = &'static str> {
567                [$(<$variant as crate::func::FuncName>::NAME,)*].into_iter()
568            }
569        }
570
571        impl fmt::Display for BinaryFunc {
572            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
573                match self {
574                    $(Self::$name(func) => func.fmt(f),)*
575                }
576            }
577        }
578
579        $(
580            impl From<$variant> for crate::BinaryFunc {
581                fn from(variant: $variant) -> Self {
582                    Self::$name(variant)
583                }
584            }
585        )*
586    }
587}