mz_expr/scalar/func/
registry.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//! Registry of the declared properties of every scalar function variant.
11//!
12//! The stable LIR schema (see `lir_schema.rs` in `mz-compute-types`) pins the
13//! *shape* of [`UnaryFunc`], [`BinaryFunc`] and [`VariadicFunc`]: their
14//! variants and payloads. It cannot see the properties the optimizer and the
15//! renderer rely on, such as null propagation, error behavior, monotonicity
16//! and output typing, nor the function bodies themselves. A change to any of
17//! those changes the meaning of a stored LIR plan without changing its
18//! serialized form.
19//!
20//! [`FuncRegistry::build`] records those properties for one representative
21//! instance of every variant, and `tests/func_registry.rs` in
22//! `mz-compute-types` compares the result against checked-in snapshots, one
23//! per half of [`Record`]. The module exists only under the `func-registry` feature,
24//! which that test enables. Production builds carry none of it. For `#[sqlfunc]`
25//! functions the record also carries the declaration text, the types-only
26//! signature and a fingerprint of the function body, see [`SqlFuncSource`],
27//! and the output type is probed at the column types the parameter types
28//! map to, see [`ColumnTypeProbe`].
29//!
30//! Variants whose payload cannot be constructed without data need a
31//! representative [`Sample`] in this module. Building the registry panics
32//! and names the variant otherwise. Samples may also be given for a
33//! constructible variant, to probe its output type at chosen input types or
34//! to record payloads whose properties differ, see [`Sample`].
35
36use std::collections::{BTreeMap, BTreeSet};
37use std::fmt;
38
39use mz_repr::adt::char::CharLength;
40use mz_repr::adt::datetime::DateTimeUnits;
41use mz_repr::adt::numeric::NumericMaxScale;
42use mz_repr::adt::regex::Regex;
43use mz_repr::adt::timestamp::TimestampPrecision;
44use mz_repr::adt::varchar::VarCharMaxLength;
45use mz_repr::{ColumnName, Datum, Row, SqlColumnType, SqlScalarType, StableRow};
46use serde::Serialize;
47
48use crate::func::format::DateTimeFormat;
49use crate::func::variadic::{
50    And, ArrayCreate, ArrayFill, ArrayIndex, ArrayToString, Coalesce, ErrorIfNull, Greatest, Least,
51    ListCreate, MapBuild, Or, RangeCreate, RecordCreate,
52};
53use crate::func::*;
54use crate::{BinaryFunc, MirScalarExpr, UnaryFunc, VariadicFunc, like_pattern};
55
56/// A representative instance of a function variant.
57///
58/// `input_types` are the column types the output type is probed at. Leave
59/// them empty to skip the probe, for functions whose output typing does not
60/// depend on a specific input shape.
61///
62/// A variant's properties can depend on its payload, for example a numeric
63/// cast errors only when it has a scale. One sample sees one point of each
64/// such property function, so a variant may have several samples, told apart
65/// by `label`. The unlabeled sample is the variant's primary record, keyed by
66/// its canonical name. Labeled samples are keyed `name[label]`.
67#[derive(Debug)]
68pub struct Sample<F> {
69    pub func: F,
70    pub input_types: Vec<SqlColumnType>,
71    pub label: &'static str,
72}
73
74impl<F> Sample<F> {
75    /// Marks this as an additional sample of its variant, recorded alongside
76    /// the primary one under `name[label]`.
77    pub fn labeled(self, label: &'static str) -> Self {
78        assert!(!label.is_empty(), "sample labels must be non-empty");
79        Sample { label, ..self }
80    }
81}
82
83/// The registry of every variant of the three scalar function enums, keyed
84/// by canonical variant name (see [`FuncName`]) as described on [`Sample`].
85#[derive(Debug)]
86pub struct FuncRegistry {
87    pub unary: BTreeMap<String, Record<UnaryFuncProperties>>,
88    pub binary: BTreeMap<String, Record<BinaryFuncProperties>>,
89    pub variadic: BTreeMap<String, Record<VariadicFuncProperties>>,
90}
91
92/// One registry entry: the properties that decide what a stored plan using
93/// the function computes, and the source they were read from.
94///
95/// The two halves are snapshotted separately by the test in
96/// `mz-compute-types`. A change to `properties` of a shipped LIR version
97/// requires a version bump. A change to `source` alone is informational.
98#[derive(Debug)]
99pub struct Record<P> {
100    pub properties: P,
101    pub source: FuncSource,
102}
103
104/// Where a record's properties came from.
105///
106/// `display` only feeds EXPLAIN output, since LIR stores variant names.
107/// `sqlfunc_decl` is the declaration text, whose semantic content the
108/// properties already carry as their own fields. `body_fingerprint` tracks
109/// the implementation, whose semantics only a reader can judge.
110#[derive(Debug, Serialize)]
111pub struct FuncSource {
112    pub display: String,
113    pub sqlfunc_decl: Option<&'static str>,
114    pub body_fingerprint: Option<String>,
115}
116
117impl FuncSource {
118    fn of(display: String, source: Option<SqlFuncSource>) -> Self {
119        FuncSource {
120            display,
121            sqlfunc_decl: source.map(|s| s.decl),
122            body_fingerprint: source.map(|s| format!("{:016x}", s.body_fingerprint)),
123        }
124    }
125}
126
127/// Resolves a Rust parameter type to its column type by autoref
128/// specialization: `(&ColumnTypeProbe::<T>(PhantomData)).column_type()` picks
129/// [`ProbeColumnType`] when `T: AsColumnType` and [`ProbeColumnTypeFallback`]
130/// otherwise. `#[sqlfunc]` emits that expression for each parameter, which is
131/// how it can ask for a column type without knowing whether one exists.
132#[derive(Debug)]
133pub struct ColumnTypeProbe<T>(pub std::marker::PhantomData<T>);
134
135/// The specialized arm of [`ColumnTypeProbe`].
136pub trait ProbeColumnType {
137    fn column_type(&self) -> Option<SqlColumnType>;
138}
139
140impl<T: mz_repr::AsColumnType> ProbeColumnType for ColumnTypeProbe<T> {
141    fn column_type(&self) -> Option<SqlColumnType> {
142        Some(T::as_column_type())
143    }
144}
145
146/// The fallback arm of [`ColumnTypeProbe`], reached through one more autoref
147/// than [`ProbeColumnType`] so it only applies when that one does not.
148pub trait ProbeColumnTypeFallback {
149    fn column_type(&self) -> Option<SqlColumnType>;
150}
151
152impl<T> ProbeColumnTypeFallback for &ColumnTypeProbe<T> {
153    fn column_type(&self) -> Option<SqlColumnType> {
154        None
155    }
156}
157
158#[derive(Debug, Serialize)]
159pub struct UnaryFuncProperties {
160    pub variant: String,
161    pub propagates_nulls: bool,
162    pub introduces_nulls: bool,
163    pub could_error: bool,
164    pub preserves_uniqueness: bool,
165    pub is_monotone: bool,
166    pub is_eliminable_cast: bool,
167    /// Canonical name of the inverse function, if one is declared.
168    pub inverse: Option<&'static str>,
169    pub input_types: Vec<String>,
170    pub output_type: Option<String>,
171    /// The `#[sqlfunc]` parameter and return types as written, see
172    /// [`SqlFuncSource::signature`].
173    pub sqlfunc_signature: Option<&'static str>,
174}
175
176#[derive(Debug, Serialize)]
177pub struct BinaryFuncProperties {
178    pub variant: String,
179    pub propagates_nulls: bool,
180    pub introduces_nulls: bool,
181    pub could_error: bool,
182    pub is_monotone: (bool, bool),
183    pub is_infinity_monotone: bool,
184    pub is_infix_op: bool,
185    /// Canonical name of the negated function, if one is declared.
186    pub negate: Option<&'static str>,
187    pub input_types: Vec<String>,
188    pub output_type: Option<String>,
189    /// The `#[sqlfunc]` parameter and return types as written, see
190    /// [`SqlFuncSource::signature`].
191    pub sqlfunc_signature: Option<&'static str>,
192}
193
194#[derive(Debug, Serialize)]
195pub struct VariadicFuncProperties {
196    pub variant: String,
197    pub propagates_nulls: bool,
198    pub introduces_nulls: bool,
199    pub could_error: bool,
200    pub is_monotone: bool,
201    pub is_associative: bool,
202    pub is_infix_op: bool,
203    pub input_types: Vec<String>,
204    pub output_type: Option<String>,
205    /// The `#[sqlfunc]` parameter and return types as written, see
206    /// [`SqlFuncSource::signature`].
207    pub sqlfunc_signature: Option<&'static str>,
208}
209
210impl FuncRegistry {
211    /// Records the properties of every variant of the three function enums.
212    ///
213    /// Panics if a variant can neither be constructed from its canonical name
214    /// nor has a [`Sample`] in this module, or if a sample's input count does
215    /// not fit its function.
216    pub fn build() -> FuncRegistry {
217        FuncRegistry {
218            unary: collect(
219                "UnaryFunc",
220                UnaryFunc::variant_names(),
221                UnaryFunc::from_variant_name,
222                UnaryFunc::sqlfunc_input_types,
223                unary_samples(),
224                UnaryFunc::variant_name,
225                UnaryFuncProperties::of,
226            ),
227            binary: collect(
228                "BinaryFunc",
229                BinaryFunc::variant_names(),
230                BinaryFunc::from_variant_name,
231                BinaryFunc::sqlfunc_input_types,
232                binary_samples(),
233                BinaryFunc::variant_name,
234                BinaryFuncProperties::of,
235            ),
236            variadic: collect(
237                "VariadicFunc",
238                VariadicFunc::variant_names(),
239                VariadicFunc::from_variant_name,
240                VariadicFunc::sqlfunc_input_types,
241                variadic_samples(),
242                VariadicFunc::variant_name,
243                VariadicFuncProperties::of,
244            ),
245        }
246    }
247}
248
249/// Records the properties of every variant of one function enum, keyed as
250/// described on [`Sample`].
251///
252/// A variant's primary record comes from its unlabeled hand-written sample if
253/// there is one, otherwise from a payload-free instance built by `construct`
254/// (see `from_variant_name` on the enums), probed at the column types
255/// `natural_inputs` reports for it, if any. Hand-written primaries take
256/// precedence so a constructible variant can still be probed at chosen input
257/// types.
258///
259/// A variant whose payload has fields must have at least one hand-written
260/// sample, even when the payload is constructible with every field defaulted,
261/// because a defaulted payload sees only one branch of any property that
262/// depends on it.
263///
264/// Panics if a variant has no primary or no required sample, if two samples
265/// share a name and label, or if probing a sample panics, in which case the
266/// panic names the sample.
267fn collect<F: fmt::Debug, P>(
268    enum_name: &str,
269    names: impl Iterator<Item = &'static str>,
270    construct: fn(&str) -> Option<F>,
271    natural_inputs: fn(&F) -> Option<Vec<SqlColumnType>>,
272    samples: Vec<Sample<F>>,
273    variant_name: fn(&F) -> &'static str,
274    properties: fn(&Sample<F>) -> P,
275) -> BTreeMap<String, P> {
276    let samples_fn = format!(
277        "{}_samples() in src/expr/src/scalar/func/registry.rs",
278        enum_name.trim_end_matches("Func").to_lowercase()
279    );
280    let mut by_name: BTreeMap<(&'static str, &'static str), Sample<F>> = BTreeMap::new();
281    let mut hand_written = BTreeSet::new();
282    for sample in samples {
283        let name = variant_name(&sample.func);
284        let label = sample.label;
285        hand_written.insert(name);
286        let duplicate = by_name.insert((name, label), sample).is_some();
287        assert!(
288            !duplicate,
289            "duplicate {enum_name} sample for `{name}` with label `{label}`"
290        );
291    }
292
293    let mut records = BTreeMap::new();
294    for name in names {
295        let primary = by_name.remove(&(name, "")).or_else(|| {
296            construct(name).map(|func| Sample {
297                input_types: natural_inputs(&func).unwrap_or_default(),
298                func,
299                label: "",
300            })
301        });
302        let primary = primary.unwrap_or_else(|| {
303            panic!(
304                "{enum_name} variant `{name}` cannot be constructed from its name because \
305                 its payload needs data. Add a representative Sample for it to {samples_fn}."
306            )
307        });
308        assert!(
309            hand_written.contains(name) || !payload_has_fields(&primary.func),
310            "{enum_name} variant `{name}` has a payload with fields but only its defaulted \
311             instance is recorded. Add a Sample with a non-default payload to {samples_fn}, \
312             labeled if the defaulted instance should stay the primary record."
313        );
314        records.insert(
315            name.to_string(),
316            probe(enum_name, name, &primary, properties),
317        );
318    }
319    for ((name, label), sample) in by_name {
320        assert!(
321            records.contains_key(name),
322            "{enum_name} sample `{name}[{label}]` names an unknown variant"
323        );
324        records.insert(
325            format!("{name}[{label}]"),
326            probe(enum_name, &format!("{name}[{label}]"), &sample, properties),
327        );
328    }
329    records
330}
331
332/// Whether a variant's payload struct has any fields.
333///
334/// Decided from the `Debug` rendering, which every payload derives: a unit
335/// payload renders as `Variant(Payload)`, one with fields opens a second
336/// parenthesis or a brace.
337fn payload_has_fields<F: fmt::Debug>(func: &F) -> bool {
338    let rendered = format!("{func:?}");
339    rendered.matches('(').count() > 1 || rendered.contains('{')
340}
341
342/// Runs `properties` on a sample, attributing any panic (typically an
343/// `output_sql_type` impl rejecting the sample's input types) to the sample.
344fn probe<F, P>(
345    enum_name: &str,
346    key: &str,
347    sample: &Sample<F>,
348    properties: fn(&Sample<F>) -> P,
349) -> P {
350    mz_ore::panic::catch_unwind_str(std::panic::AssertUnwindSafe(|| properties(sample)))
351        .unwrap_or_else(|message| {
352            panic!(
353                "probing {enum_name} sample `{key}` panicked: {message}\n\
354                 Check the sample's input types in src/expr/src/scalar/func/registry.rs."
355            )
356        })
357}
358
359/// The serde variant name, which is what the stable LIR format stores.
360fn variant_ident<F: Serialize>(func: &F) -> String {
361    let value = serde_json::to_value(func).expect("function variants serialize");
362    match value {
363        serde_json::Value::Object(map) if map.len() == 1 => {
364            map.into_iter().next().expect("one entry").0
365        }
366        serde_json::Value::String(name) => name,
367        other => panic!("unexpected function variant encoding: {other}"),
368    }
369}
370
371fn type_strings(types: &[SqlColumnType]) -> Vec<String> {
372    types.iter().map(ToString::to_string).collect()
373}
374
375impl UnaryFuncProperties {
376    fn of(sample: &Sample<UnaryFunc>) -> Record<Self> {
377        let func = &sample.func;
378        let output_type = match sample.input_types.as_slice() {
379            [] => None,
380            [input] => Some(func.output_sql_type(input.clone()).to_string()),
381            _ => panic!("unary sample `{}` needs exactly one input type", func),
382        };
383        let source = func.sqlfunc_source();
384        let properties = UnaryFuncProperties {
385            variant: variant_ident(func),
386            propagates_nulls: func.propagates_nulls(),
387            introduces_nulls: func.introduces_nulls(),
388            could_error: func.could_error(),
389            preserves_uniqueness: func.preserves_uniqueness(),
390            is_monotone: func.is_monotone(),
391            is_eliminable_cast: func.is_eliminable_cast(),
392            inverse: func.inverse().map(|f| f.variant_name()),
393            input_types: type_strings(&sample.input_types),
394            output_type,
395            sqlfunc_signature: source.map(|s| s.signature),
396        };
397        Record {
398            properties,
399            source: FuncSource::of(func.to_string(), source),
400        }
401    }
402}
403
404impl BinaryFuncProperties {
405    fn of(sample: &Sample<BinaryFunc>) -> Record<Self> {
406        let func = &sample.func;
407        let output_type = match sample.input_types.as_slice() {
408            [] => None,
409            [_, _] => Some(func.output_sql_type(&sample.input_types).to_string()),
410            _ => panic!("binary sample `{}` needs exactly two input types", func),
411        };
412        let source = func.sqlfunc_source();
413        let properties = BinaryFuncProperties {
414            variant: variant_ident(func),
415            propagates_nulls: func.propagates_nulls(),
416            introduces_nulls: func.introduces_nulls(),
417            could_error: func.could_error(),
418            is_monotone: func.is_monotone(),
419            is_infinity_monotone: func.is_infinity_monotone(),
420            is_infix_op: func.is_infix_op(),
421            negate: func.negate().map(|f| f.variant_name()),
422            input_types: type_strings(&sample.input_types),
423            output_type,
424            sqlfunc_signature: source.map(|s| s.signature),
425        };
426        Record {
427            properties,
428            source: FuncSource::of(func.to_string(), source),
429        }
430    }
431}
432
433impl VariadicFuncProperties {
434    fn of(sample: &Sample<VariadicFunc>) -> Record<Self> {
435        let func = &sample.func;
436        let output_type = (!sample.input_types.is_empty())
437            .then(|| func.output_sql_type(sample.input_types.clone()).to_string());
438        let source = func.sqlfunc_source();
439        let properties = VariadicFuncProperties {
440            variant: variant_ident(func),
441            propagates_nulls: func.propagates_nulls(),
442            introduces_nulls: func.introduces_nulls(),
443            could_error: func.could_error(),
444            is_monotone: func.is_monotone(),
445            is_associative: func.is_associative(),
446            is_infix_op: func.is_infix_op(),
447            input_types: type_strings(&sample.input_types),
448            output_type,
449            sqlfunc_signature: source.map(|s| s.signature),
450        };
451        Record {
452            properties,
453            source: FuncSource::of(func.to_string(), source),
454        }
455    }
456}
457
458fn column(i: usize) -> Box<MirScalarExpr> {
459    Box::new(MirScalarExpr::column(i))
460}
461
462fn record_type() -> SqlScalarType {
463    SqlScalarType::Record {
464        fields: [
465            ("a".into(), SqlScalarType::Int32.nullable(false)),
466            ("b".into(), SqlScalarType::String.nullable(true)),
467        ]
468        .into(),
469        custom_id: None,
470    }
471}
472
473fn time_record_type() -> SqlScalarType {
474    SqlScalarType::Record {
475        fields: [("a".into(), SqlScalarType::Time.nullable(false))].into(),
476        custom_id: None,
477    }
478}
479
480fn list_type(element: SqlScalarType) -> SqlScalarType {
481    SqlScalarType::List {
482        element_type: Box::new(element),
483        custom_id: None,
484    }
485}
486
487fn map_type(value: SqlScalarType) -> SqlScalarType {
488    SqlScalarType::Map {
489        value_type: Box::new(value),
490        custom_id: None,
491    }
492}
493
494fn range_type(element: SqlScalarType) -> SqlScalarType {
495    SqlScalarType::Range {
496        element_type: Box::new(element),
497    }
498}
499
500fn array_type(element: SqlScalarType) -> SqlScalarType {
501    SqlScalarType::Array(Box::new(element))
502}
503
504fn regex() -> Regex {
505    Regex::new("a+", false).expect("valid regex")
506}
507
508fn unary<F: Into<UnaryFunc>>(func: F, input: SqlScalarType) -> Sample<UnaryFunc> {
509    Sample {
510        func: func.into(),
511        input_types: vec![input.nullable(false)],
512        label: "",
513    }
514}
515
516fn unary_samples() -> Vec<Sample<UnaryFunc>> {
517    let timestamp_precision = Some(TimestampPrecision::try_from(3i64).expect("valid precision"));
518    let numeric_scale = NumericMaxScale::try_from(2i64).expect("valid scale");
519    let char_length = Some(CharLength::try_from(5i64).expect("valid length"));
520    let varchar_length = Some(VarCharMaxLength::try_from(5i64).expect("valid length"));
521    let tz = mz_pgtz::timezone::Timezone::Tz(chrono_tz::Tz::UTC);
522    vec![
523        // Casts between container types, which carry their target type.
524        unary(
525            CastArrayToString {
526                ty: array_type(SqlScalarType::Int32),
527            },
528            array_type(SqlScalarType::Int32),
529        ),
530        unary(
531            CastArrayToJsonb {
532                cast_element: column(0),
533            },
534            array_type(SqlScalarType::Int32),
535        ),
536        unary(
537            CastArrayToArray {
538                return_ty: array_type(SqlScalarType::Int64),
539                cast_expr: column(0),
540            },
541            array_type(SqlScalarType::Int32),
542        ),
543        unary(
544            CastListToString {
545                ty: list_type(SqlScalarType::Int32),
546            },
547            list_type(SqlScalarType::Int32),
548        ),
549        unary(
550            CastListToJsonb {
551                cast_element: column(0),
552            },
553            list_type(SqlScalarType::Int32),
554        ),
555        unary(
556            CastList1ToList2 {
557                return_ty: list_type(SqlScalarType::Int64),
558                cast_expr: column(0),
559            },
560            list_type(SqlScalarType::Int32),
561        ),
562        unary(
563            CastMapToString {
564                ty: map_type(SqlScalarType::Int32),
565            },
566            map_type(SqlScalarType::Int32),
567        ),
568        unary(
569            MapBuildFromRecordList {
570                value_type: SqlScalarType::Int32,
571            },
572            list_type(record_type()),
573        ),
574        unary(
575            CastRangeToString {
576                ty: range_type(SqlScalarType::Int32),
577            },
578            range_type(SqlScalarType::Int32),
579        ),
580        unary(CastRecordToString { ty: record_type() }, record_type()),
581        unary(
582            CastRecord1ToRecord2 {
583                return_ty: record_type(),
584                cast_exprs: vec![*column(0), *column(1)].into(),
585            },
586            record_type(),
587        ),
588        unary(RecordGet(1), record_type()),
589        // The null-fallback wrapper around a cast. Its properties derive from
590        // the wrapped function, so the samples cover the branches that differ:
591        // a fallible parse, a narrowing cast with an inverse (which the wrapper
592        // must not inherit), a nullable output the wrapper only forces nullable,
593        // and a container cast whose element cast lives inside the wrapper.
594        unary(
595            UnaryFunc::try_cast_if_fallible(CastStringToInt32),
596            SqlScalarType::String,
597        ),
598        unary(
599            UnaryFunc::try_cast_if_fallible(CastInt64ToInt32),
600            SqlScalarType::Int64,
601        )
602        .labeled("narrowing"),
603        unary(
604            UnaryFunc::try_cast_if_fallible(CastJsonbToNumeric(None)),
605            SqlScalarType::Jsonb,
606        )
607        .labeled("jsonb"),
608        unary(
609            UnaryFunc::try_cast_if_fallible(CastList1ToList2 {
610                return_ty: list_type(SqlScalarType::Int32),
611                cast_expr: column(0),
612            }),
613            list_type(SqlScalarType::Int64),
614        )
615        .labeled("list"),
616        unary(
617            CastStringToArray {
618                return_ty: array_type(SqlScalarType::Int32),
619                cast_expr: column(0),
620            },
621            SqlScalarType::String,
622        ),
623        unary(
624            CastStringToList {
625                return_ty: list_type(SqlScalarType::Int32),
626                cast_expr: column(0),
627            },
628            SqlScalarType::String,
629        ),
630        unary(
631            CastStringToMap {
632                return_ty: map_type(SqlScalarType::Int32),
633                cast_expr: column(0),
634            },
635            SqlScalarType::String,
636        ),
637        unary(
638            CastStringToRange {
639                return_ty: range_type(SqlScalarType::Int32),
640                cast_expr: column(0),
641            },
642            SqlScalarType::String,
643        ),
644        // Length-parameterized string casts.
645        unary(
646            CastStringToChar {
647                length: char_length,
648                fail_on_len: true,
649            },
650            SqlScalarType::String,
651        ),
652        unary(
653            PadChar {
654                length: char_length,
655            },
656            SqlScalarType::String,
657        ),
658        unary(
659            CastStringToVarChar {
660                length: varchar_length,
661                fail_on_len: true,
662            },
663            SqlScalarType::String,
664        ),
665        // Precision and scale parameterized casts.
666        unary(
667            CastTimestampToTimestampTz {
668                from: None,
669                to: timestamp_precision,
670            },
671            SqlScalarType::Timestamp { precision: None },
672        ),
673        unary(
674            CastTimestampTzToTimestamp {
675                from: None,
676                to: timestamp_precision,
677            },
678            SqlScalarType::TimestampTz { precision: None },
679        ),
680        unary(
681            AdjustTimestampPrecision {
682                from: None,
683                to: timestamp_precision,
684            },
685            SqlScalarType::Timestamp { precision: None },
686        ),
687        unary(
688            AdjustTimestampTzPrecision {
689                from: None,
690                to: timestamp_precision,
691            },
692            SqlScalarType::TimestampTz { precision: None },
693        ),
694        unary(
695            AdjustNumericScale(numeric_scale),
696            SqlScalarType::Numeric { max_scale: None },
697        ),
698        // Pattern matching.
699        unary(
700            IsLikeMatch(like_pattern::compile("%a%", false).expect("valid pattern")),
701            SqlScalarType::String,
702        ),
703        unary(IsRegexpMatch(regex()), SqlScalarType::String),
704        unary(RegexpMatch(regex()), SqlScalarType::String),
705        unary(
706            Jq(JqProgram::new(".a").expect("valid program")),
707            SqlScalarType::Jsonb,
708        ),
709        unary(RegexpSplitToArray(regex()), SqlScalarType::String),
710        // Date and time functions carrying units, time zones and formats.
711        unary(
712            ExtractInterval(DateTimeUnits::Epoch),
713            SqlScalarType::Interval,
714        ),
715        unary(ExtractTime(DateTimeUnits::Epoch), SqlScalarType::Time),
716        unary(
717            ExtractTimestamp(DateTimeUnits::Epoch),
718            SqlScalarType::Timestamp { precision: None },
719        ),
720        unary(
721            ExtractTimestampTz(DateTimeUnits::Epoch),
722            SqlScalarType::TimestampTz { precision: None },
723        ),
724        unary(ExtractDate(DateTimeUnits::Epoch), SqlScalarType::Date),
725        unary(
726            DatePartInterval(DateTimeUnits::Epoch),
727            SqlScalarType::Interval,
728        ),
729        unary(DatePartTime(DateTimeUnits::Epoch), SqlScalarType::Time),
730        unary(
731            DatePartTimestamp(DateTimeUnits::Epoch),
732            SqlScalarType::Timestamp { precision: None },
733        ),
734        unary(
735            DatePartTimestampTz(DateTimeUnits::Epoch),
736            SqlScalarType::TimestampTz { precision: None },
737        ),
738        unary(
739            DateTruncTimestamp(DateTimeUnits::Day),
740            SqlScalarType::Timestamp { precision: None },
741        ),
742        unary(
743            DateTruncTimestampTz(DateTimeUnits::Day),
744            SqlScalarType::TimestampTz { precision: None },
745        ),
746        unary(
747            TimezoneTimestamp(tz.clone()),
748            SqlScalarType::Timestamp { precision: None },
749        ),
750        unary(
751            TimezoneTimestampTz(tz.clone()),
752            SqlScalarType::TimestampTz { precision: None },
753        ),
754        unary(
755            TimezoneTime {
756                tz,
757                wall_time: chrono::NaiveDateTime::default(),
758            },
759            SqlScalarType::Time,
760        ),
761        unary(
762            ToCharTimestamp {
763                format_string: "YYYY".into(),
764                format: DateTimeFormat::compile("YYYY"),
765            },
766            SqlScalarType::Timestamp { precision: None },
767        ),
768        unary(
769            ToCharTimestampTz {
770                format_string: "YYYY".into(),
771                format: DateTimeFormat::compile("YYYY"),
772            },
773            SqlScalarType::TimestampTz { precision: None },
774        ),
775        // Payload-free hand-written casts, probed at their natural input.
776        unary(CastStringToInt2Vector, SqlScalarType::String),
777        unary(CastDateToTimestamp(None), SqlScalarType::Date),
778        unary(CastDateToTimestampTz(None), SqlScalarType::Date),
779        unary(CastStringToTimestamp(None), SqlScalarType::String),
780        unary(CastStringToTimestampTz(None), SqlScalarType::String),
781        unary(
782            CastDateToTimestamp(timestamp_precision),
783            SqlScalarType::Date,
784        )
785        .labeled("precision"),
786        unary(
787            CastDateToTimestampTz(timestamp_precision),
788            SqlScalarType::Date,
789        )
790        .labeled("precision"),
791        unary(
792            CastStringToTimestamp(timestamp_precision),
793            SqlScalarType::String,
794        )
795        .labeled("precision"),
796        unary(
797            CastStringToTimestampTz(timestamp_precision),
798            SqlScalarType::String,
799        )
800        .labeled("precision"),
801        // Without a length to enforce, these casts cannot error.
802        unary(
803            CastStringToChar {
804                length: None,
805                fail_on_len: false,
806            },
807            SqlScalarType::String,
808        )
809        .labeled("unbounded"),
810        unary(PadChar { length: None }, SqlScalarType::String).labeled("unbounded"),
811        unary(
812            CastStringToVarChar {
813                length: None,
814                fail_on_len: false,
815            },
816            SqlScalarType::String,
817        )
818        .labeled("unbounded"),
819        // Widening a precision preserves uniqueness, narrowing does not.
820        unary(
821            CastTimestampToTimestampTz {
822                from: timestamp_precision,
823                to: None,
824            },
825            SqlScalarType::Timestamp {
826                precision: timestamp_precision,
827            },
828        )
829        .labeled("widening"),
830        unary(
831            CastTimestampTzToTimestamp {
832                from: timestamp_precision,
833                to: None,
834            },
835            SqlScalarType::TimestampTz {
836                precision: timestamp_precision,
837            },
838        )
839        .labeled("widening"),
840        unary(
841            AdjustTimestampPrecision {
842                from: timestamp_precision,
843                to: None,
844            },
845            SqlScalarType::Timestamp {
846                precision: timestamp_precision,
847            },
848        )
849        .labeled("widening"),
850        unary(
851            AdjustTimestampTzPrecision {
852                from: timestamp_precision,
853                to: None,
854            },
855            SqlScalarType::TimestampTz {
856                precision: timestamp_precision,
857            },
858        )
859        .labeled("widening"),
860        // Text rounds `time` elements to microseconds, so these casts do not
861        // preserve uniqueness.
862        unary(
863            CastArrayToString {
864                ty: array_type(SqlScalarType::Time),
865            },
866            array_type(SqlScalarType::Time),
867        )
868        .labeled("time"),
869        unary(
870            CastListToString {
871                ty: list_type(SqlScalarType::Time),
872            },
873            list_type(SqlScalarType::Time),
874        )
875        .labeled("time"),
876        unary(
877            CastMapToString {
878                ty: map_type(SqlScalarType::Time),
879            },
880            map_type(SqlScalarType::Time),
881        )
882        .labeled("time"),
883        unary(
884            CastRangeToString {
885                ty: range_type(SqlScalarType::Time),
886            },
887            range_type(SqlScalarType::Time),
888        )
889        .labeled("time"),
890        unary(
891            CastRecordToString {
892                ty: time_record_type(),
893            },
894            time_record_type(),
895        )
896        .labeled("time"),
897        // Units below the most significant ones are not monotone.
898        unary(
899            ExtractInterval(DateTimeUnits::Month),
900            SqlScalarType::Interval,
901        )
902        .labeled("month"),
903        unary(ExtractTime(DateTimeUnits::Minute), SqlScalarType::Time).labeled("minute"),
904        unary(
905            ExtractTimestamp(DateTimeUnits::Month),
906            SqlScalarType::Timestamp { precision: None },
907        )
908        .labeled("month"),
909        unary(
910            ExtractTimestampTz(DateTimeUnits::Month),
911            SqlScalarType::TimestampTz { precision: None },
912        )
913        .labeled("month"),
914        unary(ExtractDate(DateTimeUnits::Month), SqlScalarType::Date).labeled("month"),
915        unary(
916            DatePartInterval(DateTimeUnits::Month),
917            SqlScalarType::Interval,
918        )
919        .labeled("month"),
920        unary(DatePartTime(DateTimeUnits::Minute), SqlScalarType::Time).labeled("minute"),
921        unary(
922            DatePartTimestamp(DateTimeUnits::Month),
923            SqlScalarType::Timestamp { precision: None },
924        )
925        .labeled("month"),
926        unary(
927            DatePartTimestampTz(DateTimeUnits::Month),
928            SqlScalarType::TimestampTz { precision: None },
929        )
930        .labeled("month"),
931    ]
932    .into_iter()
933    .chain(numeric_cast_samples(numeric_scale))
934    .collect()
935}
936
937/// The numeric casts, each without a scale (the primary record) and with
938/// one. Only the scaled cast can error, because it rounds.
939fn numeric_cast_samples(scale: NumericMaxScale) -> Vec<Sample<UnaryFunc>> {
940    let casts: [(fn(Option<NumericMaxScale>) -> UnaryFunc, SqlScalarType); 10] = [
941        (|s| CastInt16ToNumeric(s).into(), SqlScalarType::Int16),
942        (|s| CastInt32ToNumeric(s).into(), SqlScalarType::Int32),
943        (|s| CastInt64ToNumeric(s).into(), SqlScalarType::Int64),
944        (|s| CastUint16ToNumeric(s).into(), SqlScalarType::UInt16),
945        (|s| CastUint32ToNumeric(s).into(), SqlScalarType::UInt32),
946        (|s| CastUint64ToNumeric(s).into(), SqlScalarType::UInt64),
947        (|s| CastFloat32ToNumeric(s).into(), SqlScalarType::Float32),
948        (|s| CastFloat64ToNumeric(s).into(), SqlScalarType::Float64),
949        (|s| CastStringToNumeric(s).into(), SqlScalarType::String),
950        (|s| CastJsonbToNumeric(s).into(), SqlScalarType::Jsonb),
951    ];
952    casts
953        .into_iter()
954        .flat_map(|(cast, input)| {
955            [
956                unary(cast(None), input.clone()),
957                unary(cast(Some(scale)), input).labeled("scale"),
958            ]
959        })
960        .collect()
961}
962
963fn binary<F: Into<BinaryFunc>>(
964    func: F,
965    left: SqlScalarType,
966    right: SqlScalarType,
967) -> Sample<BinaryFunc> {
968    Sample {
969        func: func.into(),
970        input_types: vec![left.nullable(false), right.nullable(false)],
971        label: "",
972    }
973}
974
975fn binary_samples() -> Vec<Sample<BinaryFunc>> {
976    vec![
977        binary(
978            ListLengthMax { max_layer: 1 },
979            list_type(SqlScalarType::Int32),
980            SqlScalarType::Int64,
981        ),
982        binary(
983            RegexpReplace {
984                regex: regex(),
985                limit: 1,
986            },
987            SqlScalarType::String,
988            SqlScalarType::String,
989        ),
990    ]
991}
992
993fn variadic<F: Into<VariadicFunc>>(func: F, inputs: Vec<SqlScalarType>) -> Sample<VariadicFunc> {
994    Sample {
995        func: func.into(),
996        input_types: inputs.into_iter().map(|ty| ty.nullable(false)).collect(),
997        label: "",
998    }
999}
1000
1001fn variadic_samples() -> Vec<Sample<VariadicFunc>> {
1002    vec![
1003        variadic(
1004            ArrayCreate {
1005                elem_type: SqlScalarType::Int32,
1006            },
1007            vec![SqlScalarType::Int32, SqlScalarType::Int32],
1008        ),
1009        variadic(
1010            ArrayFill {
1011                elem_type: SqlScalarType::Int32,
1012            },
1013            vec![SqlScalarType::Int32, array_type(SqlScalarType::Int32)],
1014        ),
1015        variadic(
1016            ArrayIndex { offset: 1 },
1017            vec![array_type(SqlScalarType::Int32), SqlScalarType::Int64],
1018        ),
1019        variadic(
1020            ArrayToString {
1021                elem_type: SqlScalarType::Int32,
1022            },
1023            vec![array_type(SqlScalarType::Int32), SqlScalarType::String],
1024        ),
1025        variadic(
1026            ListCreate {
1027                elem_type: SqlScalarType::Int32,
1028            },
1029            vec![SqlScalarType::Int32, SqlScalarType::Int32],
1030        ),
1031        variadic(
1032            RangeCreate {
1033                elem_type: SqlScalarType::Int32,
1034            },
1035            vec![
1036                SqlScalarType::Int32,
1037                SqlScalarType::Int32,
1038                SqlScalarType::String,
1039            ],
1040        ),
1041        variadic(
1042            RecordCreate {
1043                field_names: vec![ColumnName::from("a".to_string())],
1044            },
1045            vec![SqlScalarType::Int32],
1046        ),
1047        variadic(
1048            MapBuild {
1049                value_type: SqlScalarType::Int32,
1050            },
1051            vec![SqlScalarType::String, SqlScalarType::Int32],
1052        ),
1053        variadic(
1054            CaseLiteral {
1055                lookup: vec![CaseLiteralEntry {
1056                    literal: StableRow::from(Row::pack_slice(&[Datum::Int32(1)])),
1057                    expr_index: 0,
1058                }],
1059                return_type: SqlScalarType::String.nullable(true),
1060            },
1061            vec![SqlScalarType::Int32, SqlScalarType::String],
1062        ),
1063        // Payload-free hand-written functions, probed at their natural inputs.
1064        variadic(And, vec![SqlScalarType::Bool, SqlScalarType::Bool]),
1065        variadic(Or, vec![SqlScalarType::Bool, SqlScalarType::Bool]),
1066        variadic(Coalesce, vec![SqlScalarType::Int32, SqlScalarType::Int32]),
1067        variadic(Greatest, vec![SqlScalarType::Int32, SqlScalarType::Int32]),
1068        variadic(Least, vec![SqlScalarType::Int32, SqlScalarType::Int32]),
1069        variadic(
1070            ErrorIfNull,
1071            vec![SqlScalarType::Int32, SqlScalarType::String],
1072        ),
1073    ]
1074}
1075
1076#[cfg(test)]
1077mod tests {
1078    use super::*;
1079
1080    #[mz_ore::test]
1081    fn payload_field_detection() {
1082        fn unary(func: impl Into<UnaryFunc>) -> UnaryFunc {
1083            func.into()
1084        }
1085        assert!(!payload_has_fields(&unary(Not)));
1086        assert!(payload_has_fields(&unary(CastInt32ToNumeric(None))));
1087        assert!(payload_has_fields(&unary(PadChar { length: None })));
1088        assert!(!payload_has_fields(&VariadicFunc::from(And)));
1089        assert!(payload_has_fields(&VariadicFunc::from(ArrayIndex {
1090            offset: 0
1091        })));
1092    }
1093}
mz_expr/scalar/func/registry.rs

mz_expr/scalar/func/
registry.rs
