Skip to main content

mz_expr/scalar/func/
variadic.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// Portions of this file are derived from the PostgreSQL project. The original
11// source code is subject to the terms of the PostgreSQL license, a copy of
12// which can be found in the LICENSE file at the root of this repository.
13
14//! Variadic functions.
15
16use std::borrow::Cow;
17use std::cmp;
18use std::fmt;
19
20use aws_lc_rs::hmac as aws_hmac;
21use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
22use hmac::{Hmac, KeyInit, Mac};
23use itertools::Itertools;
24use md5::Md5;
25use mz_expr_derive::sqlfunc;
26use mz_ore::cast::{CastFrom, ReinterpretCast};
27use mz_pgtz::timezone::TimezoneSpec;
28use mz_repr::ReprColumnType;
29use mz_repr::adt::array::{Array, ArrayDimension, ArrayDimensions, InvalidArrayError};
30use mz_repr::adt::mz_acl_item::{AclItem, AclMode, MzAclItem};
31use mz_repr::adt::range::{InvalidRangeError, Range, RangeBound, parse_range_bound_flags};
32use mz_repr::adt::system::Oid;
33use mz_repr::adt::timestamp::CheckedTimestamp;
34use mz_repr::role_id::RoleId;
35use mz_repr::{
36    ColumnName, Datum, DatumList, DatumMap, FromDatum, InputDatumType, OptionalArg,
37    OutputDatumType, Row, RowArena, SqlColumnType, SqlScalarType, Variadic,
38};
39use serde::{Deserialize, Serialize};
40
41use crate::func::{
42    CaseLiteral, array_create_scalar, build_regex, check_build_fits_budget,
43    check_datums_fit_budget, date_bin, max_string_func_result_bytes, parse_timezone,
44    regexp_match_static, regexp_replace_parse_flags, regexp_split_to_array_re, stringify_datum,
45    timezone_time,
46};
47use crate::{Eval, EvalError, MirScalarExpr};
48use mz_repr::adt::date::Date;
49use mz_repr::adt::interval::Interval;
50use mz_repr::adt::jsonb::JsonbRef;
51
52#[derive(
53    Ord,
54    PartialOrd,
55    Clone,
56    Debug,
57    Eq,
58    PartialEq,
59    Serialize,
60    Deserialize,
61    Hash
62)]
63pub struct And;
64
65impl fmt::Display for And {
66    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
67        f.write_str("AND")
68    }
69}
70
71impl LazyVariadicFunc for And {
72    fn eval<'a>(
73        &'a self,
74        datums: &[Datum<'a>],
75        temp_storage: &'a RowArena,
76        exprs: &'a [impl Eval],
77    ) -> Result<Datum<'a>, EvalError> {
78        // If any is false, then return false. Else, if any is null, then return null. Else, return true.
79        let mut null = false;
80        let mut err = None;
81        for expr in exprs {
82            match expr.eval(datums, temp_storage) {
83                Ok(Datum::False) => return Ok(Datum::False), // short-circuit
84                Ok(Datum::True) => {}
85                // No return in these two cases, because we might still see a false
86                Ok(Datum::Null) => null = true,
87                Err(this_err) => err = std::cmp::max(err.take(), Some(this_err)),
88                _ => unreachable!(),
89            }
90        }
91        match (err, null) {
92            (Some(err), _) => Err(err),
93            (None, true) => Ok(Datum::Null),
94            (None, false) => Ok(Datum::True),
95        }
96    }
97
98    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
99        let in_nullable = input_types.iter().any(|t| t.nullable);
100        SqlScalarType::Bool.nullable(in_nullable)
101    }
102
103    fn propagates_nulls(&self) -> bool {
104        false
105    }
106
107    fn introduces_nulls(&self) -> bool {
108        false
109    }
110
111    fn could_error(&self) -> bool {
112        false
113    }
114
115    fn is_monotone(&self) -> bool {
116        true
117    }
118
119    fn is_associative(&self) -> bool {
120        true
121    }
122
123    fn is_infix_op(&self) -> bool {
124        true
125    }
126}
127
128#[derive(
129    Ord,
130    PartialOrd,
131    Clone,
132    Debug,
133    Eq,
134    PartialEq,
135    Serialize,
136    Deserialize,
137    Hash
138)]
139pub struct ArrayCreate {
140    pub elem_type: SqlScalarType,
141}
142
143/// Constructs a new multidimensional array out of an arbitrary number of
144/// lower-dimensional arrays.
145///
146/// For example, if given three 1D arrays of length 2, this function will
147/// construct a 2D array with dimensions 3x2.
148///
149/// The input datums in `datums` must all be arrays of the same dimensions.
150/// (The arrays must also be of the same element type, but that is checked by
151/// the SQL type system, rather than checked here at runtime.)
152///
153/// If all input arrays are zero-dimensional arrays, then the output is a zero-
154/// dimensional array. Otherwise, the lower bound of the additional dimension is
155/// one and the length of the new dimension is equal to `datums.len()`.
156///
157/// Null elements are allowed and considered to be zero-dimensional arrays.
158#[sqlfunc(
159    ArrayCreate,
160    output_type_expr = "match &self.elem_type { SqlScalarType::Array(_) => self.elem_type.clone().nullable(false), _ => SqlScalarType::Array(Box::new(self.elem_type.clone())).nullable(false) }"
161)]
162fn array_create<'a>(
163    &self,
164    datums: Variadic<Datum<'a>>,
165    temp_storage: &'a RowArena,
166) -> Result<Array<'a>, EvalError> {
167    match &self.elem_type {
168        SqlScalarType::Array(_) => array_create_multidim(&datums, temp_storage),
169        _ => array_create_scalar(&datums, temp_storage),
170    }
171}
172fn array_create_multidim<'a>(
173    datums: &[Datum<'a>],
174    temp_storage: &'a RowArena,
175) -> Result<Array<'a>, EvalError> {
176    let mut dim: Option<ArrayDimensions> = None;
177    for datum in datums {
178        let actual_dims = match datum {
179            Datum::Null => ArrayDimensions::default(),
180            Datum::Array(arr) => arr.dims(),
181            d => panic!("unexpected datum {d}"),
182        };
183        if let Some(expected) = &dim {
184            if actual_dims.ndims() != expected.ndims() {
185                let actual = actual_dims.ndims().into();
186                let expected = expected.ndims().into();
187                // All input arrays must have the same dimensionality.
188                return Err(InvalidArrayError::WrongCardinality { actual, expected }.into());
189            }
190            if let Some((e, a)) = expected
191                .into_iter()
192                .zip_eq(actual_dims)
193                .find(|(e, a)| e != a)
194            {
195                let actual = a.length;
196                let expected = e.length;
197                // All input arrays must have the same dimensionality.
198                return Err(InvalidArrayError::WrongCardinality { actual, expected }.into());
199            }
200        }
201        dim = Some(actual_dims);
202    }
203    // Per PostgreSQL, if all input arrays are zero dimensional, so is the output.
204    if dim.as_ref().map_or(true, ArrayDimensions::is_empty) {
205        let datum = temp_storage.try_make_datum(|packer| packer.try_push_array(&[], &[]))?;
206        return Ok(datum.unwrap_array());
207    }
208
209    let mut dims = vec![ArrayDimension {
210        lower_bound: 1,
211        length: datums.len(),
212    }];
213    if let Some(d) = datums.first() {
214        dims.extend(d.unwrap_array().dims());
215    };
216    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
217    let elements = datums
218        .iter()
219        .flat_map(|d| d.unwrap_array().elements().iter());
220    let datum =
221        temp_storage.try_make_datum(move |packer| packer.try_push_array(&dims, elements))?;
222    Ok(datum.unwrap_array())
223}
224
225#[derive(
226    Ord,
227    PartialOrd,
228    Clone,
229    Debug,
230    Eq,
231    PartialEq,
232    Serialize,
233    Deserialize,
234    Hash
235)]
236pub struct ArrayFill {
237    pub elem_type: SqlScalarType,
238}
239
240#[sqlfunc(
241    ArrayFill,
242    output_type_expr = "SqlScalarType::Array(Box::new(self.elem_type.clone())).nullable(false)"
243)]
244fn array_fill<'a>(
245    &self,
246    fill: Datum<'a>,
247    dims: Option<Array<'a>>,
248    lower_bounds: OptionalArg<Option<Array<'a>>>,
249    temp_storage: &'a RowArena,
250) -> Result<Array<'a>, EvalError> {
251    const MAX_SIZE: usize = (1 << 28) - 1;
252    const NULL_ARR_ERR: &str = "dimension array or low bound array";
253    const NULL_ELEM_ERR: &str = "dimension values";
254
255    if matches!(fill, Datum::Array(_)) {
256        return Err(EvalError::Unsupported {
257            feature: "array_fill with arrays".into(),
258            discussion_no: None,
259        });
260    }
261
262    let Some(arr) = dims else {
263        return Err(EvalError::MustNotBeNull(NULL_ARR_ERR.into()));
264    };
265
266    // The dimensions array must be one-dimensional.
267    if arr.dims().ndims() > 1 {
268        return Err(EvalError::ArrayFillWrongArraySubscripts);
269    }
270
271    let dimensions = arr
272        .elements()
273        .iter()
274        .map(|d| match d {
275            Datum::Null => Err(EvalError::MustNotBeNull(NULL_ELEM_ERR.into())),
276            d => Ok(usize::cast_from(u32::reinterpret_cast(d.unwrap_int32()))),
277        })
278        .collect::<Result<Vec<_>, _>>()?;
279
280    let lower_bounds = match *lower_bounds {
281        Some(d) => {
282            let Some(arr) = d else {
283                return Err(EvalError::MustNotBeNull(NULL_ARR_ERR.into()));
284            };
285
286            // The lower bounds array must be one-dimensional.
287            if arr.dims().ndims() > 1 {
288                return Err(EvalError::ArrayFillWrongArraySubscripts);
289            }
290
291            arr.elements()
292                .iter()
293                .map(|l| match l {
294                    Datum::Null => Err(EvalError::MustNotBeNull(NULL_ELEM_ERR.into())),
295                    l => Ok(isize::cast_from(l.unwrap_int32())),
296                })
297                .collect::<Result<Vec<_>, _>>()?
298        }
299        None => {
300            vec![1isize; dimensions.len()]
301        }
302    };
303
304    if lower_bounds.len() != dimensions.len() {
305        return Err(EvalError::ArrayFillWrongArraySubscripts);
306    }
307
308    let fill_count: usize = dimensions
309        .iter()
310        .cloned()
311        .map(Some)
312        .reduce(|a, b| match (a, b) {
313            (Some(a), Some(b)) => a.checked_mul(b),
314            _ => None,
315        })
316        .flatten()
317        .ok_or(EvalError::MaxArraySizeExceeded(MAX_SIZE))?;
318
319    if matches!(
320        mz_repr::datum_size(&fill).checked_mul(fill_count),
321        None | Some(MAX_SIZE..)
322    ) {
323        return Err(EvalError::MaxArraySizeExceeded(MAX_SIZE));
324    }
325
326    // The packed array lands in `temp_storage`, and building it first allocates a transient
327    // `Vec<Datum>` of `fill_count` elements. Unlike the string amplifiers, `array_fill` sizes its
328    // result from a parameter rather than its input, so without this check a single call could
329    // spike far past a budgeted arena before the evaluator's post-call check ever runs. Refuse both
330    // allocations up front instead. Without a budget `budget_remaining` is `usize::MAX` and both
331    // comparisons fold away, so this is free in a dataflow.
332    let budget_remaining = temp_storage.budget_remaining();
333    let packed_size = mz_repr::datum_size(&fill).saturating_mul(fill_count);
334    let build_size = fill_count.saturating_mul(std::mem::size_of::<Datum<'a>>());
335    if packed_size > budget_remaining || build_size > budget_remaining {
336        return Err(EvalError::TempStorageBudgetExceeded);
337    }
338
339    let array_dimensions = if fill_count == 0 {
340        vec![ArrayDimension {
341            lower_bound: 1,
342            length: 0,
343        }]
344    } else {
345        dimensions
346            .into_iter()
347            .zip_eq(lower_bounds)
348            .map(|(length, lower_bound)| ArrayDimension {
349                lower_bound,
350                length,
351            })
352            .collect()
353    };
354
355    let datum = temp_storage.try_make_datum(|packer| {
356        packer.try_push_array(&array_dimensions, vec![fill; fill_count])
357    })?;
358    Ok(datum.unwrap_array())
359}
360
361#[derive(
362    Ord,
363    PartialOrd,
364    Clone,
365    Debug,
366    Eq,
367    PartialEq,
368    Serialize,
369    Deserialize,
370    Hash
371)]
372pub struct ArrayIndex {
373    pub offset: i64,
374}
375#[sqlfunc(ArrayIndex, sqlname = "array_index")]
376fn array_index<'a, T: FromDatum<'a>>(
377    &self,
378    array: Array<'a, T>,
379    indices: Variadic<i64>,
380) -> Option<T> {
381    mz_ore::soft_assert_no_log!(
382        self.offset == 0 || self.offset == 1,
383        "offset must be either 0 or 1"
384    );
385
386    let dims = array.dims();
387    if dims.len() != indices.len() {
388        // You missed the datums "layer"
389        return None;
390    }
391
392    let mut final_idx = 0;
393
394    for (d, idx) in dims.into_iter().zip_eq(indices.iter()) {
395        // Lower bound is written in terms of 1-based indexing, which offset accounts for.
396        let idx = isize::cast_from(*idx + self.offset);
397
398        let (lower, upper) = d.dimension_bounds();
399
400        // This index missed all of the data at this layer. The dimension bounds are inclusive,
401        // while range checks are exclusive, so adjust.
402        if !(lower..upper + 1).contains(&idx) {
403            return None;
404        }
405
406        // We discover how many indices our last index represents physically.
407        final_idx *= d.length;
408
409        // Because both index and lower bound are handled in 1-based indexing, taking their
410        // difference moves us back into 0-based indexing. Similarly, if the lower bound is
411        // negative, subtracting a negative value >= to itself ensures its non-negativity.
412        final_idx += usize::try_from(idx - d.lower_bound)
413            .expect("previous bounds check ensures physical index is at least 0");
414    }
415
416    array.elements().typed_iter().nth(final_idx)
417}
418
419#[sqlfunc]
420fn array_position<'a>(
421    array: Array<'a>,
422    search: Datum<'a>,
423    initial_pos: OptionalArg<Option<i32>>,
424) -> Result<Option<i32>, EvalError> {
425    if array.dims().len() > 1 {
426        return Err(EvalError::MultiDimensionalArraySearch);
427    }
428
429    // Comparisons use IS NOT DISTINCT FROM semantics, so a NULL search term
430    // matches a NULL element rather than yielding NULL. `Datum` equality already
431    // treats `Null == Null` as true, so the search loop below handles it.
432    //
433    // NOTE: rejecting a NULL initial position before the search is a deliberate
434    // divergence from PostgreSQL. PostgreSQL validates the initial position only
435    // after a fast path that returns NULL when the search term is NULL and the
436    // array holds no NULL element, so it errors for
437    // `array_position(ARRAY['sun', NULL], NULL, NULL)` but not for
438    // `array_position(ARRAY['sun'], NULL, NULL)`. Making the error depend on the
439    // array contents is not worth reproducing.
440    let skip = match initial_pos.0 {
441        None => 0,
442        Some(None) => return Err(EvalError::MustNotBeNull("initial position".into())),
443        Some(Some(o)) => usize::try_from(o).unwrap_or(0).saturating_sub(1),
444    };
445
446    let Some(r) = array.elements().iter().skip(skip).position(|d| d == search) else {
447        return Ok(None);
448    };
449
450    // Adjust count for the amount we skipped, plus 1 for adjusting to PG indexing scheme.
451    let p = i32::try_from(r + skip + 1).expect("fewer than i32::MAX elements in array");
452    Ok(Some(p))
453}
454
455#[derive(
456    Ord,
457    PartialOrd,
458    Clone,
459    Debug,
460    Eq,
461    PartialEq,
462    Serialize,
463    Deserialize,
464    Hash
465)]
466pub struct ArrayToString {
467    pub elem_type: SqlScalarType,
468}
469
470#[sqlfunc]
471fn array_to_string<'a>(
472    &self,
473    array: Array<'a>,
474    delimiter: &str,
475    null_str_arg: OptionalArg<Option<&str>>,
476    temp_storage: &RowArena,
477) -> Result<String, EvalError> {
478    // `flatten` treats absent arguments (`None`) the same as explicit NULL
479    // (`Some(None)`), both becoming `None`.
480    let null_str = null_str_arg.flatten();
481    // A delimiter is repeated once per element, so the array size doesn't bound the result. We
482    // check as it grows rather than up front, since an element's stringified length isn't known
483    // without stringifying it. That bounds the peak at the ceiling plus one element.
484    let max_result_bytes = max_string_func_result_bytes(temp_storage);
485    let mut out = String::new();
486    for elem in array.elements().iter() {
487        if elem.is_null() {
488            if let Some(null_str) = null_str {
489                out.push_str(null_str);
490                out.push_str(delimiter);
491            }
492        } else {
493            stringify_datum(&mut out, elem, &self.elem_type)?;
494            out.push_str(delimiter);
495        }
496        if out.len() > max_result_bytes {
497            return Err(EvalError::LengthTooLarge);
498        }
499    }
500    if out.len() > 0 {
501        // Lop off last delimiter only if string is not empty
502        out.truncate(out.len() - delimiter.len());
503    }
504    Ok(out)
505}
506
507#[derive(
508    Ord,
509    PartialOrd,
510    Clone,
511    Debug,
512    Eq,
513    PartialEq,
514    Serialize,
515    Deserialize,
516    Hash
517)]
518pub struct Coalesce;
519
520impl fmt::Display for Coalesce {
521    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
522        f.write_str("coalesce")
523    }
524}
525
526impl LazyVariadicFunc for Coalesce {
527    fn eval<'a>(
528        &'a self,
529        datums: &[Datum<'a>],
530        temp_storage: &'a RowArena,
531        exprs: &'a [impl Eval],
532    ) -> Result<Datum<'a>, EvalError> {
533        for e in exprs {
534            let d = e.eval(datums, temp_storage)?;
535            if !d.is_null() {
536                return Ok(d);
537            }
538        }
539        Ok(Datum::Null)
540    }
541
542    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
543        // Note that the parser doesn't allow empty argument lists for variadic functions
544        // that use the standard function call syntax (ArrayCreate and co. are different
545        // because of the special syntax for calling them).
546        let nullable = input_types.iter().all(|typ| typ.nullable);
547        SqlColumnType::union_many(input_types).nullable(nullable)
548    }
549
550    fn propagates_nulls(&self) -> bool {
551        false
552    }
553
554    fn introduces_nulls(&self) -> bool {
555        false
556    }
557
558    fn could_error(&self) -> bool {
559        false
560    }
561
562    fn is_monotone(&self) -> bool {
563        true
564    }
565
566    fn is_associative(&self) -> bool {
567        true
568    }
569}
570
571#[derive(
572    Ord,
573    PartialOrd,
574    Clone,
575    Debug,
576    Eq,
577    PartialEq,
578    Serialize,
579    Deserialize,
580    Hash
581)]
582pub struct RangeCreate {
583    pub elem_type: SqlScalarType,
584}
585
586impl fmt::Display for RangeCreate {
587    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
588        f.write_str(match &self.elem_type {
589            SqlScalarType::Int32 => "int4range",
590            SqlScalarType::Int64 => "int8range",
591            SqlScalarType::Date => "daterange",
592            SqlScalarType::Numeric { .. } => "numrange",
593            SqlScalarType::Timestamp { .. } => "tsrange",
594            SqlScalarType::TimestampTz { .. } => "tstzrange",
595            _ => unreachable!(),
596        })
597    }
598}
599
600#[sqlfunc(
601    RangeCreate,
602    skip_display = true,
603    output_type_expr = "SqlScalarType::Range { element_type: Box::new(self.elem_type.clone()) }.nullable(false)"
604)]
605fn range_create<'a>(
606    &self,
607    lower: Datum<'a>,
608    upper: Datum<'a>,
609    flags_datum: Datum<'a>,
610) -> Result<Range<Datum<'a>>, EvalError> {
611    let flags = match flags_datum {
612        Datum::Null => {
613            return Err(EvalError::InvalidRange(
614                InvalidRangeError::NullRangeBoundFlags,
615            ));
616        }
617        o => o.unwrap_str(),
618    };
619
620    let (lower_inclusive, upper_inclusive) = parse_range_bound_flags(flags)?;
621
622    let mut range = Range::new(Some((
623        RangeBound::new(lower, lower_inclusive),
624        RangeBound::new(upper, upper_inclusive),
625    )));
626
627    range.canonicalize()?;
628
629    Ok(range)
630}
631
632#[sqlfunc(sqlname = "datediff")]
633fn date_diff_date(unit_str: &str, a: Date, b: Date) -> Result<i64, EvalError> {
634    let unit = unit_str
635        .parse()
636        .map_err(|_| EvalError::InvalidDatePart(unit_str.into()))?;
637
638    // Convert the Date into a timestamp so we can calculate age.
639    let a_ts = CheckedTimestamp::try_from(NaiveDate::from(a).and_hms_opt(0, 0, 0).unwrap())?;
640    let b_ts = CheckedTimestamp::try_from(NaiveDate::from(b).and_hms_opt(0, 0, 0).unwrap())?;
641    let diff = b_ts.diff_as(&a_ts, unit)?;
642    Ok(diff)
643}
644
645#[sqlfunc(sqlname = "datediff")]
646fn date_diff_time(unit_str: &str, a: NaiveTime, b: NaiveTime) -> Result<i64, EvalError> {
647    let unit = unit_str
648        .parse()
649        .map_err(|_| EvalError::InvalidDatePart(unit_str.into()))?;
650
651    // Convert the Time into a timestamp so we can calculate age.
652    let a_ts =
653        CheckedTimestamp::try_from(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap().and_time(a))?;
654    let b_ts =
655        CheckedTimestamp::try_from(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap().and_time(b))?;
656    let diff = b_ts.diff_as(&a_ts, unit)?;
657    Ok(diff)
658}
659
660#[sqlfunc(sqlname = "datediff")]
661fn date_diff_timestamp(
662    unit: &str,
663    a: CheckedTimestamp<NaiveDateTime>,
664    b: CheckedTimestamp<NaiveDateTime>,
665) -> Result<i64, EvalError> {
666    let unit = unit
667        .parse()
668        .map_err(|_| EvalError::InvalidDatePart(unit.into()))?;
669
670    let diff = b.diff_as(&a, unit)?;
671    Ok(diff)
672}
673
674#[sqlfunc(sqlname = "datediff")]
675fn date_diff_timestamp_tz(
676    unit: &str,
677    a: CheckedTimestamp<DateTime<Utc>>,
678    b: CheckedTimestamp<DateTime<Utc>>,
679) -> Result<i64, EvalError> {
680    let unit = unit
681        .parse()
682        .map_err(|_| EvalError::InvalidDatePart(unit.into()))?;
683
684    let diff = b.diff_as(&a, unit)?;
685    Ok(diff)
686}
687
688#[derive(
689    Ord,
690    PartialOrd,
691    Clone,
692    Debug,
693    Eq,
694    PartialEq,
695    Serialize,
696    Deserialize,
697    Hash
698)]
699pub struct ErrorIfNull;
700
701impl fmt::Display for ErrorIfNull {
702    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
703        f.write_str("error_if_null")
704    }
705}
706
707impl LazyVariadicFunc for ErrorIfNull {
708    fn eval<'a>(
709        &'a self,
710        datums: &[Datum<'a>],
711        temp_storage: &'a RowArena,
712        exprs: &'a [impl Eval],
713    ) -> Result<Datum<'a>, EvalError> {
714        let first = exprs[0].eval(datums, temp_storage)?;
715        match first {
716            Datum::Null => {
717                let err_msg = match exprs[1].eval(datums, temp_storage)? {
718                    Datum::Null => {
719                        return Err(EvalError::Internal(
720                            "unexpected NULL in error side of error_if_null".into(),
721                        ));
722                    }
723                    o => o.unwrap_str(),
724                };
725                Err(EvalError::IfNullError(err_msg.into()))
726            }
727            _ => Ok(first),
728        }
729    }
730
731    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
732        input_types[0].scalar_type.clone().nullable(false)
733    }
734
735    fn propagates_nulls(&self) -> bool {
736        false
737    }
738
739    fn introduces_nulls(&self) -> bool {
740        false
741    }
742}
743
744#[derive(
745    Ord,
746    PartialOrd,
747    Clone,
748    Debug,
749    Eq,
750    PartialEq,
751    Serialize,
752    Deserialize,
753    Hash
754)]
755pub struct Greatest;
756
757impl fmt::Display for Greatest {
758    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
759        f.write_str("greatest")
760    }
761}
762
763impl LazyVariadicFunc for Greatest {
764    fn eval<'a>(
765        &'a self,
766        datums: &[Datum<'a>],
767        temp_storage: &'a RowArena,
768        exprs: &'a [impl Eval],
769    ) -> Result<Datum<'a>, EvalError> {
770        let mut greatest = None;
771        for e in exprs {
772            let d = e.eval(datums, temp_storage)?;
773            // NOTE: The later argument wins a tie. Datums that compare equal can
774            // still differ, e.g. float8 `-0` and `0`.
775            if !d.is_null() && greatest.is_none_or(|g| d >= g) {
776                greatest = Some(d);
777            }
778        }
779        Ok(greatest.unwrap_or(Datum::Null))
780    }
781
782    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
783        SqlColumnType::union_many(input_types)
784    }
785
786    fn propagates_nulls(&self) -> bool {
787        false
788    }
789
790    fn introduces_nulls(&self) -> bool {
791        false
792    }
793
794    fn could_error(&self) -> bool {
795        false
796    }
797
798    fn is_monotone(&self) -> bool {
799        true
800    }
801
802    fn is_associative(&self) -> bool {
803        true
804    }
805}
806
807#[sqlfunc(sqlname = "hmac")]
808fn hmac_string(to_digest: &str, key: &str, typ: &str) -> Result<Vec<u8>, EvalError> {
809    let to_digest = to_digest.as_bytes();
810    let key = key.as_bytes();
811    hmac_inner(to_digest, key, typ)
812}
813
814#[sqlfunc(sqlname = "hmac")]
815fn hmac_bytes(to_digest: &[u8], key: &[u8], typ: &str) -> Result<Vec<u8>, EvalError> {
816    hmac_inner(to_digest, key, typ)
817}
818
819pub fn hmac_inner(to_digest: &[u8], key: &[u8], typ: &str) -> Result<Vec<u8>, EvalError> {
820    match typ {
821        "md5" => {
822            let mut mac = Hmac::<Md5>::new_from_slice(key).expect("HMAC accepts any key size");
823            mac.update(to_digest);
824            Ok(mac.finalize().into_bytes().to_vec())
825        }
826        "sha1" => {
827            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA1_FOR_LEGACY_USE_ONLY, key);
828            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
829        }
830        "sha224" => {
831            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA224, key);
832            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
833        }
834        "sha256" => {
835            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA256, key);
836            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
837        }
838        "sha384" => {
839            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA384, key);
840            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
841        }
842        "sha512" => {
843            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA512, key);
844            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
845        }
846        other => Err(EvalError::InvalidHashAlgorithm(other.into())),
847    }
848}
849
850#[sqlfunc]
851fn jsonb_build_array<'a>(
852    datums: Variadic<Datum<'a>>,
853    temp_storage: &'a RowArena,
854) -> Result<JsonbRef<'a>, EvalError> {
855    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
856    let datum = temp_storage.make_datum(|packer| {
857        packer.push_list(datums.into_iter().map(|d| match d {
858            Datum::Null => Datum::JsonNull,
859            d => d,
860        }))
861    });
862    Ok(JsonbRef::from_datum(datum))
863}
864
865#[sqlfunc]
866fn jsonb_build_object<'a>(
867    mut kvs: Variadic<(Datum<'a>, Datum<'a>)>,
868    temp_storage: &'a RowArena,
869) -> Result<JsonbRef<'a>, EvalError> {
870    kvs.0.sort_by(|kv1, kv2| kv1.0.cmp(&kv2.0));
871    kvs.0.dedup_by(|kv1, kv2| kv1.0 == kv2.0);
872    // Checked after the dedup, so duplicate keys are counted once, as they are packed.
873    check_datums_fit_budget(kvs.0.iter().flat_map(|(k, v)| [*k, *v]), temp_storage)?;
874    let datum = temp_storage.try_make_datum(|packer| {
875        packer.push_dict_with(|packer| {
876            for (k, v) in kvs {
877                if k.is_null() {
878                    return Err(EvalError::KeyCannotBeNull);
879                }
880                let v = match v {
881                    Datum::Null => Datum::JsonNull,
882                    d => d,
883                };
884                packer.push(k);
885                packer.push(v);
886            }
887            Ok(())
888        })
889    })?;
890    Ok(JsonbRef::from_datum(datum))
891}
892
893#[derive(
894    Ord,
895    PartialOrd,
896    Clone,
897    Debug,
898    Eq,
899    PartialEq,
900    Serialize,
901    Deserialize,
902    Hash
903)]
904pub struct Least;
905
906impl fmt::Display for Least {
907    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
908        f.write_str("least")
909    }
910}
911
912impl LazyVariadicFunc for Least {
913    fn eval<'a>(
914        &'a self,
915        datums: &[Datum<'a>],
916        temp_storage: &'a RowArena,
917        exprs: &'a [impl Eval],
918    ) -> Result<Datum<'a>, EvalError> {
919        let mut least = None;
920        for e in exprs {
921            let d = e.eval(datums, temp_storage)?;
922            // NOTE: The later argument wins a tie, see `Greatest::eval`.
923            if !d.is_null() && least.is_none_or(|l| d <= l) {
924                least = Some(d);
925            }
926        }
927        Ok(least.unwrap_or(Datum::Null))
928    }
929
930    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
931        SqlColumnType::union_many(input_types)
932    }
933
934    fn propagates_nulls(&self) -> bool {
935        false
936    }
937
938    fn introduces_nulls(&self) -> bool {
939        false
940    }
941
942    fn could_error(&self) -> bool {
943        false
944    }
945
946    fn is_monotone(&self) -> bool {
947        true
948    }
949
950    fn is_associative(&self) -> bool {
951        true
952    }
953}
954
955#[derive(
956    Ord,
957    PartialOrd,
958    Clone,
959    Debug,
960    Eq,
961    PartialEq,
962    Serialize,
963    Deserialize,
964    Hash
965)]
966pub struct ListCreate {
967    pub elem_type: SqlScalarType,
968}
969
970#[sqlfunc(
971    output_type_expr = "SqlScalarType::List { element_type: Box::new(self.elem_type.clone()), custom_id: None }.nullable(false)"
972)]
973fn list_create<'a>(
974    &self,
975    datums: Variadic<Datum<'a>>,
976    temp_storage: &'a RowArena,
977) -> Result<DatumList<'a>, EvalError> {
978    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
979    Ok(temp_storage
980        .make_datum(|packer| packer.push_list(datums))
981        .unwrap_list())
982}
983
984#[derive(
985    Ord,
986    PartialOrd,
987    Clone,
988    Debug,
989    Eq,
990    PartialEq,
991    Serialize,
992    Deserialize,
993    Hash
994)]
995pub struct RecordCreate {
996    pub field_names: Vec<ColumnName>,
997}
998
999#[sqlfunc(
1000    output_type_expr = "SqlScalarType::Record { fields: self.field_names.clone().into_iter().zip_eq(input_types.iter().cloned()).collect(), custom_id: None }.nullable(false)"
1001)]
1002fn record_create<'a>(
1003    &self,
1004    datums: Variadic<Datum<'a>>,
1005    temp_storage: &'a RowArena,
1006) -> Result<DatumList<'a>, EvalError> {
1007    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
1008    Ok(temp_storage
1009        .make_datum(|packer| packer.push_list(datums.iter().copied()))
1010        .unwrap_list())
1011}
1012
1013#[sqlfunc(
1014    output_type_expr = "input_types[0].scalar_type.unwrap_list_nth_layer_type(input_types.len() - 1).clone().nullable(true)"
1015)]
1016// TODO(benesch): remove potentially dangerous usage of `as`.
1017#[allow(clippy::as_conversions)]
1018fn list_index<'a>(buf: DatumList<'a>, indices: Variadic<i64>) -> Datum<'a> {
1019    let mut buf = Datum::List(buf);
1020    for i in indices {
1021        if buf.is_null() {
1022            break;
1023        }
1024        if i < 1 {
1025            return Datum::Null;
1026        }
1027
1028        buf = match buf.unwrap_list().iter().nth(i as usize - 1) {
1029            Some(datum) => datum,
1030            None => return Datum::Null,
1031        }
1032    }
1033    buf
1034}
1035
1036#[sqlfunc(sqlname = "makeaclitem")]
1037fn make_acl_item(
1038    grantee_oid: u32,
1039    grantor_oid: u32,
1040    privileges: &str,
1041    is_grantable: bool,
1042) -> Result<AclItem, EvalError> {
1043    let grantee = Oid(grantee_oid);
1044    let grantor = Oid(grantor_oid);
1045    let acl_mode = AclMode::parse_multiple_privileges(privileges)
1046        .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
1047    if is_grantable {
1048        return Err(EvalError::Unsupported {
1049            feature: "GRANT OPTION".into(),
1050            discussion_no: None,
1051        });
1052    }
1053
1054    Ok(AclItem {
1055        grantee,
1056        grantor,
1057        acl_mode,
1058    })
1059}
1060
1061#[sqlfunc(sqlname = "make_mz_aclitem")]
1062fn make_mz_acl_item(
1063    grantee_str: &str,
1064    grantor_str: &str,
1065    privileges: &str,
1066) -> Result<MzAclItem, EvalError> {
1067    let grantee: RoleId = grantee_str
1068        .parse()
1069        .map_err(|e: anyhow::Error| EvalError::InvalidRoleId(e.to_string().into()))?;
1070    let grantor: RoleId = grantor_str
1071        .parse()
1072        .map_err(|e: anyhow::Error| EvalError::InvalidRoleId(e.to_string().into()))?;
1073    if grantor == RoleId::Public {
1074        return Err(EvalError::InvalidRoleId(
1075            "mz_aclitem grantor cannot be PUBLIC role".into(),
1076        ));
1077    }
1078    let acl_mode = AclMode::parse_multiple_privileges(privileges)
1079        .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
1080
1081    Ok(MzAclItem {
1082        grantee,
1083        grantor,
1084        acl_mode,
1085    })
1086}
1087
1088#[sqlfunc(sqlname = "makets")]
1089#[allow(clippy::as_conversions)]
1090fn make_timestamp(
1091    year: i64,
1092    month: i64,
1093    day: i64,
1094    hour: i64,
1095    minute: i64,
1096    second_float: f64,
1097) -> Result<Option<CheckedTimestamp<NaiveDateTime>>, EvalError> {
1098    // NOTE: This function never returns `Ok(None)`. The `Option` stays so
1099    // that the derived `introduces_nulls` (and with it the nullability of
1100    // existing objects planned with this function) does not change.
1101
1102    // Negative years are BC and there is no year 0. Chrono's astronomical
1103    // numbering has 1 BC as year 0, so BC years shift by one.
1104    if year == 0 {
1105        return Err(EvalError::DateOutOfRange);
1106    }
1107    let year = if year < 0 { year + 1 } else { year };
1108    let year: i32 = year.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1109    let month: u32 = month.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1110    let day: u32 = day.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1111    let date = NaiveDate::from_ymd_opt(year, month, day).ok_or(EvalError::DateOutOfRange)?;
1112    // Allow exactly 24:00:00 (rolls over to midnight of the next day) and
1113    // 60 seconds (rolls over into the next minute).
1114    if !(0..=24).contains(&hour)
1115        || !(0..=59).contains(&minute)
1116        || !(0.0..=60.0).contains(&second_float)
1117        || (hour == 24 && (minute > 0 || second_float > 0.0))
1118    {
1119        return Err(EvalError::TimestampOutOfRange);
1120    }
1121    // The cast is exact: second_float is in [0, 60].
1122    let micros = (second_float * 1_000_000.0).round() as i64;
1123    let time = chrono::Duration::microseconds((hour * 60 + minute) * 60 * 1_000_000 + micros);
1124    let timestamp = date
1125        .and_hms_opt(0, 0, 0)
1126        .unwrap()
1127        .checked_add_signed(time)
1128        .ok_or(EvalError::TimestampOutOfRange)?;
1129    Ok(Some(timestamp.try_into()?))
1130}
1131
1132#[derive(
1133    Ord,
1134    PartialOrd,
1135    Clone,
1136    Debug,
1137    Eq,
1138    PartialEq,
1139    Serialize,
1140    Deserialize,
1141    Hash
1142)]
1143pub struct MapBuild {
1144    pub value_type: SqlScalarType,
1145}
1146
1147#[sqlfunc(
1148    output_type_expr = "SqlScalarType::Map { value_type: Box::new(self.value_type.clone()), custom_id: None }.nullable(false)"
1149)]
1150fn map_build<'a>(
1151    &self,
1152    datums: Variadic<(Option<&str>, Datum<'a>)>,
1153    temp_storage: &'a RowArena,
1154) -> Result<DatumMap<'a>, EvalError> {
1155    // Collect into a `BTreeMap` to provide the same semantics as it.
1156    let map: std::collections::BTreeMap<&str, _> = datums
1157        .into_iter()
1158        .filter_map(|(k, v)| k.map(|k| (k, v)))
1159        .collect();
1160
1161    // Checked after the collect, so duplicate keys are counted once, as they are packed. The
1162    // `BTreeMap` transient needs no pre-check: it holds one entry per argument, a count the query
1163    // text fixes, not the input.
1164    check_datums_fit_budget(
1165        map.iter().flat_map(|(k, v)| [Datum::String(k), *v]),
1166        temp_storage,
1167    )?;
1168    Ok(temp_storage
1169        .make_datum(|packer| packer.push_dict(map))
1170        .unwrap_map())
1171}
1172
1173#[derive(
1174    Ord,
1175    PartialOrd,
1176    Clone,
1177    Debug,
1178    Eq,
1179    PartialEq,
1180    Serialize,
1181    Deserialize,
1182    Hash
1183)]
1184pub struct Or;
1185
1186impl fmt::Display for Or {
1187    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1188        f.write_str("OR")
1189    }
1190}
1191
1192impl LazyVariadicFunc for Or {
1193    fn eval<'a>(
1194        &'a self,
1195        datums: &[Datum<'a>],
1196        temp_storage: &'a RowArena,
1197        exprs: &'a [impl Eval],
1198    ) -> Result<Datum<'a>, EvalError> {
1199        // If any is true, then return true. Else, if any is null, then return null. Else, return false.
1200        let mut null = false;
1201        let mut err = None;
1202        for expr in exprs {
1203            match expr.eval(datums, temp_storage) {
1204                Ok(Datum::False) => {}
1205                Ok(Datum::True) => return Ok(Datum::True), // short-circuit
1206                // No return in these two cases, because we might still see a true
1207                Ok(Datum::Null) => null = true,
1208                Err(this_err) => err = std::cmp::max(err.take(), Some(this_err)),
1209                _ => unreachable!(),
1210            }
1211        }
1212        match (err, null) {
1213            (Some(err), _) => Err(err),
1214            (None, true) => Ok(Datum::Null),
1215            (None, false) => Ok(Datum::False),
1216        }
1217    }
1218
1219    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
1220        let in_nullable = input_types.iter().any(|t| t.nullable);
1221        SqlScalarType::Bool.nullable(in_nullable)
1222    }
1223
1224    fn propagates_nulls(&self) -> bool {
1225        false
1226    }
1227
1228    fn introduces_nulls(&self) -> bool {
1229        false
1230    }
1231
1232    fn could_error(&self) -> bool {
1233        false
1234    }
1235
1236    fn is_monotone(&self) -> bool {
1237        true
1238    }
1239
1240    fn is_associative(&self) -> bool {
1241        true
1242    }
1243
1244    fn is_infix_op(&self) -> bool {
1245        true
1246    }
1247}
1248
1249#[sqlfunc(sqlname = "lpad")]
1250fn pad_leading(
1251    string: &str,
1252    raw_len: i32,
1253    pad: OptionalArg<&str>,
1254    temp_storage: &RowArena,
1255) -> Result<String, EvalError> {
1256    let len = match usize::try_from(raw_len) {
1257        Ok(len) => len,
1258        Err(_) => {
1259            return Err(EvalError::InvalidParameterValue(
1260                "length must be nonnegative".into(),
1261            ));
1262        }
1263    };
1264    let pad_string = pad.unwrap_or(" ");
1265
1266    // `len` is a character count but the ceiling is in bytes, so a multi-byte pad character writes
1267    // several bytes per character the check allowed. Bound the bytes instead. The result is the kept
1268    // prefix of `string` plus `len - string_chars` padding characters, each at most as wide as the
1269    // widest character in `pad_string`. Truncation (`len <= string_chars`) writes at most
1270    // `string.len()`, which the same expression covers.
1271    let widest_pad = pad_string.chars().map(char::len_utf8).max().unwrap_or(0);
1272    let pad_chars = len.saturating_sub(string.chars().count());
1273    let max_result_bytes = string
1274        .len()
1275        .saturating_add(pad_chars.saturating_mul(widest_pad));
1276    if max_result_bytes > max_string_func_result_bytes(temp_storage) {
1277        return Err(EvalError::LengthTooLarge);
1278    }
1279
1280    let (end_char, end_char_byte_offset) = string
1281        .chars()
1282        .take(len)
1283        .fold((0, 0), |acc, char| (acc.0 + 1, acc.1 + char.len_utf8()));
1284
1285    // NOTE: reserve from the bytes actually written, not from `len`, which is a character count.
1286    // Reserving `len` is too little for a multi-byte pad, and unbounded for an empty one, which
1287    // writes nothing and so has no result size for the ceiling above to refuse.
1288    let truncating = len == end_char;
1289    let mut buf = String::with_capacity(if truncating {
1290        end_char_byte_offset
1291    } else {
1292        max_result_bytes
1293    });
1294    if truncating {
1295        buf.push_str(&string[0..end_char_byte_offset]);
1296    } else {
1297        buf.extend(pad_string.chars().cycle().take(len - end_char));
1298        buf.push_str(string);
1299    }
1300
1301    Ok(buf)
1302}
1303
1304#[sqlfunc(
1305    output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(true)"
1306)]
1307fn regexp_match<'a>(
1308    haystack: &'a str,
1309    needle: &str,
1310    flags: OptionalArg<&str>,
1311    temp_storage: &'a RowArena,
1312) -> Result<Option<Array<'a>>, EvalError> {
1313    let flags = flags.unwrap_or("");
1314    let needle = build_regex(needle, flags)?;
1315    regexp_match_static(haystack, temp_storage, &needle)
1316}
1317
1318#[sqlfunc(
1319    output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)"
1320)]
1321fn regexp_split_to_array<'a>(
1322    text: &str,
1323    regexp_str: &str,
1324    flags: OptionalArg<&str>,
1325    temp_storage: &'a RowArena,
1326) -> Result<Array<'a>, EvalError> {
1327    let flags = flags.unwrap_or("");
1328    let regexp = build_regex(regexp_str, flags)?;
1329    regexp_split_to_array_re(text, &regexp, temp_storage)
1330}
1331
1332#[sqlfunc]
1333fn regexp_replace<'a>(
1334    source: &'a str,
1335    pattern: &str,
1336    replacement: &str,
1337    flags_opt: OptionalArg<&str>,
1338    temp_storage: &RowArena,
1339) -> Result<Cow<'a, str>, EvalError> {
1340    let flags = flags_opt.0.unwrap_or("");
1341    let (limit, flags) = regexp_replace_parse_flags(flags);
1342    let regexp = build_regex(pattern, &flags)?;
1343
1344    // The result is the unreplaced parts of `source` plus one replacement per match. A `$ref` in the
1345    // replacement expands to captured text, which is at most the match, and matches are disjoint, so
1346    // all the expansions together cost at most one `source` per `$`.
1347    let dollars = replacement.matches('$').count();
1348    let max_result_bytes = |matches: usize| {
1349        let expansions = match matches {
1350            0 => 0,
1351            _ => dollars.saturating_mul(source.len()),
1352        };
1353        source
1354            .len()
1355            .saturating_add(matches.saturating_mul(replacement.len()))
1356            .saturating_add(expansions)
1357    };
1358    // A `limit` of 0 means replace every match, and a pattern that matches the empty string matches
1359    // at every position. Assume that first, and only pay for counting the real matches when the
1360    // assumption is over the ceiling. Same compromise as `replace`.
1361    let most_matches = match limit {
1362        0 => source.len().saturating_add(1),
1363        limit => limit,
1364    };
1365    let ceiling = max_string_func_result_bytes(temp_storage);
1366    if max_result_bytes(most_matches) > ceiling {
1367        let matches = regexp.find_iter(source).take(most_matches).count();
1368        if max_result_bytes(matches) > ceiling {
1369            return Err(EvalError::LengthTooLarge);
1370        }
1371    }
1372
1373    Ok(regexp.replacen(source, limit, replacement))
1374}
1375
1376#[sqlfunc]
1377fn replace(text: &str, from: &str, to: &str, temp_storage: &RowArena) -> Result<String, EvalError> {
1378    // As a compromise to avoid always nearly duplicating the work of replace by doing size estimation,
1379    // we first check if it's possible for the fully replaced string to exceed the limit by assuming that
1380    // every possible substring is replaced.
1381    //
1382    // If that estimate exceeds the limit, we then do a more precise (and expensive) estimate by counting
1383    // the actual number of replacements that would occur, and using that to calculate the final size.
1384    let max_result_bytes = max_string_func_result_bytes(temp_storage);
1385    let possible_size = text.len() * to.len();
1386    if possible_size > max_result_bytes {
1387        let replacement_count = text.matches(from).count();
1388        let estimated_size = text.len() + replacement_count * (to.len().saturating_sub(from.len()));
1389        if estimated_size > max_result_bytes {
1390            return Err(EvalError::LengthTooLarge);
1391        }
1392    }
1393
1394    Ok(text.replace(from, to))
1395}
1396
1397#[sqlfunc(
1398    output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)",
1399    propagates_nulls = false
1400)]
1401fn string_to_array<'a>(
1402    string: &'a str,
1403    delimiter: Option<&'a str>,
1404    null_string: OptionalArg<Option<&'a str>>,
1405    temp_storage: &'a RowArena,
1406) -> Result<Array<'a>, EvalError> {
1407    if string.is_empty() {
1408        let mut row = Row::default();
1409        let mut packer = row.packer();
1410        packer.try_push_array(&[], std::iter::empty::<Datum>())?;
1411
1412        return Ok(temp_storage.push_unary_row(row).unwrap_array());
1413    }
1414
1415    let Some(delimiter) = delimiter else {
1416        let split_all_chars_delimiter = "";
1417        return string_to_array_impl(
1418            string,
1419            split_all_chars_delimiter,
1420            null_string.flatten(),
1421            temp_storage,
1422        );
1423    };
1424
1425    if delimiter.is_empty() {
1426        let mut row = Row::default();
1427        let mut packer = row.packer();
1428        let dims = &[ArrayDimension {
1429            lower_bound: 1,
1430            length: 1,
1431        }];
1432        match null_string.flatten() {
1433            Some(null_string) if null_string == string => {
1434                packer.try_push_array(dims, std::iter::once(Datum::Null))?;
1435            }
1436            _ => {
1437                packer.try_push_array(dims, vec![string].into_iter().map(Datum::String))?;
1438            }
1439        }
1440        Ok(temp_storage.push_unary_row(row).unwrap_array())
1441    } else {
1442        string_to_array_impl(string, delimiter, null_string.flatten(), temp_storage)
1443    }
1444}
1445
1446fn string_to_array_impl<'a>(
1447    string: &str,
1448    delimiter: &str,
1449    null_string: Option<&'a str>,
1450    temp_storage: &'a RowArena,
1451) -> Result<Array<'a>, EvalError> {
1452    // Bound the transient `Vec<&str>` before it is collected. The count walks the same iterator the
1453    // collect does, so the two agree.
1454    check_build_fits_budget(
1455        || {
1456            if delimiter.is_empty() {
1457                string.split(delimiter).filter(|s| !s.is_empty()).count()
1458            } else {
1459                string.split(delimiter).count()
1460            }
1461        },
1462        std::mem::size_of::<&str>(),
1463        temp_storage,
1464    )?;
1465
1466    let mut row = Row::default();
1467    let mut packer = row.packer();
1468
1469    let result = string.split(delimiter);
1470    let found: Vec<&str> = if delimiter.is_empty() {
1471        result.filter(|s| !s.is_empty()).collect()
1472    } else {
1473        result.collect()
1474    };
1475    let array_dimensions = [ArrayDimension {
1476        lower_bound: 1,
1477        length: found.len(),
1478    }];
1479
1480    // Splitting amplifies: each chunk is packed with its own tag and length, so a delimiter that
1481    // splits per character costs several times the input. `Datum::Null` packs smaller than the
1482    // chunk it replaces, so treating every chunk as a string is an upper bound.
1483    check_datums_fit_budget(found.iter().copied().map(Datum::String), temp_storage)?;
1484
1485    if let Some(null_string) = null_string {
1486        let found_datums = found.into_iter().map(|chunk| {
1487            if chunk.eq(null_string) {
1488                Datum::Null
1489            } else {
1490                Datum::String(chunk)
1491            }
1492        });
1493
1494        packer.try_push_array(&array_dimensions, found_datums)?;
1495    } else {
1496        packer.try_push_array(&array_dimensions, found.into_iter().map(Datum::String))?;
1497    }
1498
1499    Ok(temp_storage.push_unary_row(row).unwrap_array())
1500}
1501
1502#[sqlfunc]
1503fn substr<'a>(s: &'a str, start: i32, length: OptionalArg<i32>) -> Result<&'a str, EvalError> {
1504    let raw_start_idx = i64::from(start) - 1;
1505    let start_idx = match usize::try_from(cmp::max(raw_start_idx, 0)) {
1506        Ok(i) => i,
1507        Err(_) => {
1508            return Err(EvalError::InvalidParameterValue(
1509                format!(
1510                    "substring starting index ({}) exceeds min/max position",
1511                    raw_start_idx
1512                )
1513                .into(),
1514            ));
1515        }
1516    };
1517
1518    let mut char_indices = s.char_indices();
1519    let get_str_index = |(index, _char)| index;
1520
1521    let str_len = s.len();
1522    let start_char_idx = char_indices.nth(start_idx).map_or(str_len, get_str_index);
1523
1524    if let OptionalArg(Some(len)) = length {
1525        let end_idx = match i64::from(len) {
1526            e if e < 0 => {
1527                return Err(EvalError::InvalidParameterValue(
1528                    "negative substring length not allowed".into(),
1529                ));
1530            }
1531            e if e == 0 || e + raw_start_idx < 1 => return Ok(""),
1532            e => {
1533                let e = cmp::min(raw_start_idx + e - 1, e - 1);
1534                match usize::try_from(e) {
1535                    Ok(i) => i,
1536                    Err(_) => {
1537                        return Err(EvalError::InvalidParameterValue(
1538                            format!("substring length ({}) exceeds max position", e).into(),
1539                        ));
1540                    }
1541                }
1542            }
1543        };
1544
1545        let end_char_idx = char_indices.nth(end_idx).map_or(str_len, get_str_index);
1546
1547        Ok(&s[start_char_idx..end_char_idx])
1548    } else {
1549        Ok(&s[start_char_idx..])
1550    }
1551}
1552
1553#[sqlfunc(sqlname = "split_string")]
1554fn split_part<'a>(string: &'a str, delimiter: &str, field: i32) -> Result<&'a str, EvalError> {
1555    let index = match usize::try_from(i64::from(field) - 1) {
1556        Ok(index) => index,
1557        Err(_) => {
1558            return Err(EvalError::InvalidParameterValue(
1559                "field position must be greater than zero".into(),
1560            ));
1561        }
1562    };
1563
1564    // If the provided delimiter is the empty string,
1565    // PostgreSQL does not break the string into individual
1566    // characters. Instead, it generates the following parts: [string].
1567    if delimiter.is_empty() {
1568        if index == 0 {
1569            return Ok(string);
1570        } else {
1571            return Ok("");
1572        }
1573    }
1574
1575    // If provided index is greater than the number of split parts,
1576    // return an empty string.
1577    Ok(string.split(delimiter).nth(index).unwrap_or(""))
1578}
1579
1580#[sqlfunc(is_associative = true)]
1581fn concat(strs: Variadic<Option<&str>>, temp_storage: &RowArena) -> Result<String, EvalError> {
1582    let max_result_bytes = max_string_func_result_bytes(temp_storage);
1583    let mut total_size = 0;
1584    for s in &strs {
1585        if let Some(s) = s {
1586            total_size += s.len();
1587            if total_size > max_result_bytes {
1588                return Err(EvalError::LengthTooLarge);
1589            }
1590        }
1591    }
1592    let mut buf = String::with_capacity(total_size);
1593    for s in strs {
1594        if let Some(s) = s {
1595            buf.push_str(s);
1596        }
1597    }
1598    Ok(buf)
1599}
1600
1601#[sqlfunc]
1602fn concat_ws(
1603    ws: &str,
1604    rest: Variadic<Option<&str>>,
1605    temp_storage: &RowArena,
1606) -> Result<String, EvalError> {
1607    let max_result_bytes = max_string_func_result_bytes(temp_storage);
1608    let mut total_size = 0;
1609    for s in &rest {
1610        if let Some(s) = s {
1611            total_size += s.len();
1612            total_size += ws.len();
1613            if total_size > max_result_bytes {
1614                return Err(EvalError::LengthTooLarge);
1615            }
1616        }
1617    }
1618
1619    let buf = Itertools::join(&mut rest.into_iter().filter_map(|s| s), ws);
1620
1621    Ok(buf)
1622}
1623
1624#[sqlfunc]
1625fn translate(
1626    string: &str,
1627    from_str: &str,
1628    to_str: &str,
1629    temp_storage: &RowArena,
1630) -> Result<String, EvalError> {
1631    // Translating to a wider character amplifies: each byte of an ASCII input can become four.
1632    // Every output character is either the input character or a `to` character, so it costs at most
1633    // `max(len_utf8(c), widest)` bytes, and at most `len_utf8(c) + widest - 1` since a character is
1634    // at least one byte. Summed, that's the input's byte length plus one `widest - 1` per character.
1635    // Tight when every character widens, exactly the input length when `to` is ASCII, so we refuse
1636    // only calls that really would exceed the ceiling.
1637    let widest = to_str.chars().map(char::len_utf8).max().unwrap_or(0);
1638    let max_result_bytes = match widest.saturating_sub(1) {
1639        // An ASCII `to` widens nothing, so the result fits the input and counting characters isn't
1640        // worth a second pass.
1641        0 => string.len(),
1642        widening => string
1643            .len()
1644            .saturating_add(string.chars().count().saturating_mul(widening)),
1645    };
1646    if max_result_bytes > max_string_func_result_bytes(temp_storage) {
1647        return Err(EvalError::LengthTooLarge);
1648    }
1649
1650    let from = from_str.chars().collect::<Vec<_>>();
1651    let to = to_str.chars().collect::<Vec<_>>();
1652
1653    Ok(string
1654        .chars()
1655        .filter_map(|c| match from.iter().position(|f| f == &c) {
1656            Some(idx) => to.get(idx).copied(),
1657            None => Some(c),
1658        })
1659        .collect())
1660}
1661
1662#[sqlfunc(output_type_expr = "input_types[0].scalar_type.clone().nullable(false)")]
1663// TODO(benesch): remove potentially dangerous usage of `as`.
1664#[allow(clippy::as_conversions)]
1665fn list_slice_linear<'a>(
1666    list: DatumList<'a>,
1667    first: (i64, i64),
1668    remainder: Variadic<(i64, i64)>,
1669    temp_storage: &'a RowArena,
1670) -> DatumList<'a> {
1671    let mut start_idx = 0;
1672    let mut total_length = usize::MAX;
1673
1674    for (start, end) in std::iter::once(first).chain(remainder) {
1675        let start = std::cmp::max(start, 1);
1676
1677        // Result should be empty list.
1678        if start > end {
1679            start_idx = 0;
1680            total_length = 0;
1681            break;
1682        }
1683
1684        let start_inner = start as usize - 1;
1685        // Start index only moves to geq positions.
1686        start_idx += start_inner;
1687
1688        // Length index only moves to leq positions
1689        let length_inner = (end - start) as usize + 1;
1690        total_length = std::cmp::min(length_inner, total_length - start_inner);
1691    }
1692
1693    let iter = list.iter().skip(start_idx).take(total_length);
1694
1695    temp_storage
1696        .make_datum(|row| {
1697            row.push_list_with(|row| {
1698                // if iter is empty, will get the appropriate empty list.
1699                for d in iter {
1700                    row.push(d);
1701                }
1702            });
1703        })
1704        .unwrap_list()
1705}
1706
1707#[sqlfunc(sqlname = "timestamp_bin")]
1708fn date_bin_timestamp(
1709    stride: Interval,
1710    source: CheckedTimestamp<NaiveDateTime>,
1711    origin: CheckedTimestamp<NaiveDateTime>,
1712) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
1713    date_bin(stride, source, origin)
1714}
1715
1716#[sqlfunc(sqlname = "timestamptz_bin")]
1717fn date_bin_timestamp_tz(
1718    stride: Interval,
1719    source: CheckedTimestamp<DateTime<Utc>>,
1720    origin: CheckedTimestamp<DateTime<Utc>>,
1721) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
1722    date_bin(stride, source, origin)
1723}
1724
1725#[sqlfunc(sqlname = "timezonet")]
1726fn timezone_time_variadic(
1727    tz_str: &str,
1728    time: NaiveTime,
1729    wall_time: CheckedTimestamp<DateTime<Utc>>,
1730) -> Result<NaiveTime, EvalError> {
1731    parse_timezone(tz_str, TimezoneSpec::Posix)
1732        .map(|tz| timezone_time(tz, time, &wall_time.naive_utc()))
1733}
1734pub(crate) trait LazyVariadicFunc: fmt::Display {
1735    fn eval<'a>(
1736        &'a self,
1737        datums: &[Datum<'a>],
1738        temp_storage: &'a RowArena,
1739        exprs: &'a [impl Eval],
1740    ) -> Result<Datum<'a>, EvalError>;
1741
1742    /// The output SqlColumnType of this function.
1743    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType;
1744
1745    /// Whether this function will produce NULL on NULL input.
1746    fn propagates_nulls(&self) -> bool;
1747
1748    /// Whether this function will produce NULL on non-NULL input.
1749    fn introduces_nulls(&self) -> bool;
1750
1751    /// Whether this function might error on non-error input.
1752    fn could_error(&self) -> bool {
1753        true
1754    }
1755
1756    /// Returns true if the function is monotone.
1757    fn is_monotone(&self) -> bool {
1758        false
1759    }
1760
1761    /// Returns true if the function is associative.
1762    fn is_associative(&self) -> bool {
1763        false
1764    }
1765
1766    /// Returns true if the function is an infix operator.
1767    fn is_infix_op(&self) -> bool {
1768        false
1769    }
1770}
1771
1772pub(crate) trait EagerVariadicFunc: fmt::Display {
1773    type Input<'a>: InputDatumType<'a, EvalError>;
1774    type Output<'a>: OutputDatumType<'a, EvalError>;
1775
1776    fn call<'a>(&self, input: Self::Input<'a>, temp_storage: &'a RowArena) -> Self::Output<'a>;
1777
1778    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType;
1779
1780    fn propagates_nulls(&self) -> bool {
1781        !Self::Input::nullable()
1782    }
1783
1784    fn introduces_nulls(&self) -> bool {
1785        Self::Output::nullable()
1786    }
1787
1788    fn could_error(&self) -> bool {
1789        Self::Output::fallible()
1790    }
1791
1792    fn is_monotone(&self) -> bool {
1793        false
1794    }
1795
1796    fn is_associative(&self) -> bool {
1797        false
1798    }
1799
1800    fn is_infix_op(&self) -> bool {
1801        false
1802    }
1803}
1804
1805/// Blanket `LazyVariadicFunc` impl for each eager type, bridging
1806/// expression evaluation and null propagation via `InputDatumType::try_from_iter`.
1807impl<T: EagerVariadicFunc> LazyVariadicFunc for T {
1808    fn eval<'a>(
1809        &'a self,
1810        datums: &[Datum<'a>],
1811        temp_storage: &'a RowArena,
1812        exprs: &'a [impl Eval],
1813    ) -> Result<Datum<'a>, EvalError> {
1814        let mut datums = exprs.iter().map(|e| e.eval(datums, temp_storage));
1815        match T::Input::try_from_iter(&mut datums) {
1816            Ok(input) => self.call(input, temp_storage).into_result(temp_storage),
1817            Err(Ok(None)) => Err(EvalError::Internal("missing parameter".into())),
1818            Err(Ok(Some(datum))) if datum.is_null() => Ok(datum),
1819            Err(Ok(Some(_datum))) => {
1820                // datum is _not_ NULL
1821                Err(EvalError::Internal("invalid input type".into()))
1822            }
1823            Err(Err(res)) => Err(res),
1824        }
1825    }
1826
1827    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
1828        self.output_type(input_types)
1829    }
1830
1831    fn propagates_nulls(&self) -> bool {
1832        self.propagates_nulls()
1833    }
1834
1835    fn introduces_nulls(&self) -> bool {
1836        self.introduces_nulls()
1837    }
1838
1839    fn could_error(&self) -> bool {
1840        self.could_error()
1841    }
1842
1843    fn is_monotone(&self) -> bool {
1844        self.is_monotone()
1845    }
1846
1847    fn is_associative(&self) -> bool {
1848        self.is_associative()
1849    }
1850
1851    fn is_infix_op(&self) -> bool {
1852        self.is_infix_op()
1853    }
1854}
1855
1856derive_variadic! {
1857    Coalesce(Coalesce),
1858    Greatest(Greatest),
1859    Least(Least),
1860    Concat(Concat),
1861    ConcatWs(ConcatWs),
1862    MakeTimestamp(MakeTimestamp),
1863    PadLeading(PadLeading),
1864    Substr(Substr),
1865    Replace(Replace),
1866    JsonbBuildArray(JsonbBuildArray),
1867    JsonbBuildObject(JsonbBuildObject),
1868    MapBuild(MapBuild),
1869    ArrayCreate(ArrayCreate),
1870    ArrayToString(ArrayToString),
1871    ArrayIndex(ArrayIndex),
1872    ListCreate(ListCreate),
1873    RecordCreate(RecordCreate),
1874    ListIndex(ListIndex),
1875    ListSliceLinear(ListSliceLinear),
1876    SplitPart(SplitPart),
1877    RegexpMatch(RegexpMatch),
1878    HmacString(HmacString),
1879    HmacBytes(HmacBytes),
1880    ErrorIfNull(ErrorIfNull),
1881    DateBinTimestamp(DateBinTimestamp),
1882    DateBinTimestampTz(DateBinTimestampTz),
1883    DateDiffTimestamp(DateDiffTimestamp),
1884    DateDiffTimestampTz(DateDiffTimestampTz),
1885    DateDiffDate(DateDiffDate),
1886    DateDiffTime(DateDiffTime),
1887    And(And),
1888    Or(Or),
1889    RangeCreate(RangeCreate),
1890    MakeAclItem(MakeAclItem),
1891    MakeMzAclItem(MakeMzAclItem),
1892    Translate(Translate),
1893    ArrayPosition(ArrayPosition),
1894    ArrayFill(ArrayFill),
1895    StringToArray(StringToArray),
1896    TimezoneTimeVariadic(TimezoneTimeVariadic),
1897    RegexpSplitToArray(RegexpSplitToArray),
1898    RegexpReplace(RegexpReplace),
1899    CaseLiteral(CaseLiteral),
1900}
1901
1902impl VariadicFunc {
1903    pub fn switch_and_or(&self) -> Self {
1904        match self {
1905            VariadicFunc::And(_) => Or.into(),
1906            VariadicFunc::Or(_) => And.into(),
1907            _ => unreachable!(),
1908        }
1909    }
1910
1911    /// Gives the unit (u) of OR or AND, such that `u AND/OR x == x`.
1912    /// Note that a 0-arg AND/OR evaluates to unit_of_and_or.
1913    pub fn unit_of_and_or(&self) -> MirScalarExpr {
1914        match self {
1915            VariadicFunc::And(_) => MirScalarExpr::literal_true(),
1916            VariadicFunc::Or(_) => MirScalarExpr::literal_false(),
1917            _ => unreachable!(),
1918        }
1919    }
1920
1921    /// Gives the zero (z) of OR or AND, such that `z AND/OR x == z`.
1922    pub fn zero_of_and_or(&self) -> MirScalarExpr {
1923        match self {
1924            VariadicFunc::And(_) => MirScalarExpr::literal_false(),
1925            VariadicFunc::Or(_) => MirScalarExpr::literal_true(),
1926            _ => unreachable!(),
1927        }
1928    }
1929}