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