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