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, 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, FromDatum, InputDatumType, OptionalArg, OutputDatumType, Row,
38    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    introduces_nulls = false
163)]
164fn array_create<'a>(
165    &self,
166    datums: Variadic<Datum<'a>>,
167    temp_storage: &'a RowArena,
168) -> Result<Datum<'a>, EvalError> {
169    match &self.elem_type {
170        SqlScalarType::Array(_) => array_create_multidim(&datums, temp_storage),
171        _ => array_create_scalar(&datums, temp_storage),
172    }
173}
174fn array_create_multidim<'a>(
175    datums: &[Datum<'a>],
176    temp_storage: &'a RowArena,
177) -> Result<Datum<'a>, EvalError> {
178    let mut dim: Option<ArrayDimensions> = None;
179    for datum in datums {
180        let actual_dims = match datum {
181            Datum::Null => ArrayDimensions::default(),
182            Datum::Array(arr) => arr.dims(),
183            d => panic!("unexpected datum {d}"),
184        };
185        if let Some(expected) = &dim {
186            if actual_dims.ndims() != expected.ndims() {
187                let actual = actual_dims.ndims().into();
188                let expected = expected.ndims().into();
189                // All input arrays must have the same dimensionality.
190                return Err(InvalidArrayError::WrongCardinality { actual, expected }.into());
191            }
192            if let Some((e, a)) = expected
193                .into_iter()
194                .zip_eq(actual_dims)
195                .find(|(e, a)| e != a)
196            {
197                let actual = a.length;
198                let expected = e.length;
199                // All input arrays must have the same dimensionality.
200                return Err(InvalidArrayError::WrongCardinality { actual, expected }.into());
201            }
202        }
203        dim = Some(actual_dims);
204    }
205    // Per PostgreSQL, if all input arrays are zero dimensional, so is the output.
206    if dim.as_ref().map_or(true, ArrayDimensions::is_empty) {
207        return Ok(temp_storage.try_make_datum(|packer| packer.try_push_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)
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    introduces_nulls = false
245)]
246fn array_fill<'a>(
247    &self,
248    fill: Datum<'a>,
249    dims: Option<Array<'a>>,
250    lower_bounds: OptionalArg<Option<Array<'a>>>,
251    temp_storage: &'a RowArena,
252) -> Result<Datum<'a>, EvalError> {
253    const MAX_SIZE: usize = (1 << 28) - 1;
254    const NULL_ARR_ERR: &str = "dimension array or low bound array";
255    const NULL_ELEM_ERR: &str = "dimension values";
256
257    if matches!(fill, Datum::Array(_)) {
258        return Err(EvalError::Unsupported {
259            feature: "array_fill with arrays".into(),
260            discussion_no: None,
261        });
262    }
263
264    let Some(arr) = dims else {
265        return Err(EvalError::MustNotBeNull(NULL_ARR_ERR.into()));
266    };
267
268    // The dimensions array must be one-dimensional.
269    if arr.dims().ndims() > 1 {
270        return Err(EvalError::ArrayFillWrongArraySubscripts);
271    }
272
273    let dimensions = arr
274        .elements()
275        .iter()
276        .map(|d| match d {
277            Datum::Null => Err(EvalError::MustNotBeNull(NULL_ELEM_ERR.into())),
278            d => Ok(usize::cast_from(u32::reinterpret_cast(d.unwrap_int32()))),
279        })
280        .collect::<Result<Vec<_>, _>>()?;
281
282    let lower_bounds = match *lower_bounds {
283        Some(d) => {
284            let Some(arr) = d else {
285                return Err(EvalError::MustNotBeNull(NULL_ARR_ERR.into()));
286            };
287
288            // The lower bounds array must be one-dimensional.
289            if arr.dims().ndims() > 1 {
290                return Err(EvalError::ArrayFillWrongArraySubscripts);
291            }
292
293            arr.elements()
294                .iter()
295                .map(|l| match l {
296                    Datum::Null => Err(EvalError::MustNotBeNull(NULL_ELEM_ERR.into())),
297                    l => Ok(isize::cast_from(l.unwrap_int32())),
298                })
299                .collect::<Result<Vec<_>, _>>()?
300        }
301        None => {
302            vec![1isize; dimensions.len()]
303        }
304    };
305
306    if lower_bounds.len() != dimensions.len() {
307        return Err(EvalError::ArrayFillWrongArraySubscripts);
308    }
309
310    let fill_count: usize = dimensions
311        .iter()
312        .cloned()
313        .map(Some)
314        .reduce(|a, b| match (a, b) {
315            (Some(a), Some(b)) => a.checked_mul(b),
316            _ => None,
317        })
318        .flatten()
319        .ok_or(EvalError::MaxArraySizeExceeded(MAX_SIZE))?;
320
321    if matches!(
322        mz_repr::datum_size(&fill).checked_mul(fill_count),
323        None | Some(MAX_SIZE..)
324    ) {
325        return Err(EvalError::MaxArraySizeExceeded(MAX_SIZE));
326    }
327
328    // The packed array lands in `temp_storage`, and building it first allocates a transient
329    // `Vec<Datum>` of `fill_count` elements. Unlike the string amplifiers, `array_fill` sizes its
330    // result from a parameter rather than its input, so without this check a single call could
331    // spike far past a budgeted arena before the evaluator's post-call check ever runs. Refuse both
332    // allocations up front instead. Without a budget `budget_remaining` is `usize::MAX` and both
333    // comparisons fold away, so this is free in a dataflow.
334    let budget_remaining = temp_storage.budget_remaining();
335    let packed_size = mz_repr::datum_size(&fill).saturating_mul(fill_count);
336    let build_size = fill_count.saturating_mul(std::mem::size_of::<Datum<'a>>());
337    if packed_size > budget_remaining || build_size > budget_remaining {
338        return Err(EvalError::TempStorageBudgetExceeded);
339    }
340
341    let array_dimensions = if fill_count == 0 {
342        vec![ArrayDimension {
343            lower_bound: 1,
344            length: 0,
345        }]
346    } else {
347        dimensions
348            .into_iter()
349            .zip_eq(lower_bounds)
350            .map(|(length, lower_bound)| ArrayDimension {
351                lower_bound,
352                length,
353            })
354            .collect()
355    };
356
357    Ok(temp_storage.try_make_datum(|packer| {
358        packer.try_push_array(&array_dimensions, vec![fill; fill_count])
359    })?)
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", introduces_nulls = true)]
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
601impl EagerVariadicFunc for RangeCreate {
602    type Input<'a> = (Datum<'a>, Datum<'a>, Datum<'a>);
603    type Output<'a> = Result<Datum<'a>, EvalError>;
604
605    fn call<'a>(
606        &self,
607        (lower, upper, flags_datum): Self::Input<'a>,
608        temp_storage: &'a RowArena,
609    ) -> Self::Output<'a> {
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(temp_storage.make_datum(|row| {
629            row.push_range(range).expect("errors already handled");
630        }))
631    }
632
633    fn output_type(&self, _input_types: &[SqlColumnType]) -> SqlColumnType {
634        SqlScalarType::Range {
635            element_type: Box::new(self.elem_type.clone()),
636        }
637        .nullable(false)
638    }
639
640    fn introduces_nulls(&self) -> bool {
641        false
642    }
643}
644
645#[sqlfunc(sqlname = "datediff")]
646fn date_diff_date(unit_str: &str, a: Date, b: Date) -> Result<i64, EvalError> {
647    let unit = unit_str
648        .parse()
649        .map_err(|_| EvalError::InvalidDatePart(unit_str.into()))?;
650
651    // Convert the Date into a timestamp so we can calculate age.
652    let a_ts = CheckedTimestamp::try_from(NaiveDate::from(a).and_hms_opt(0, 0, 0).unwrap())?;
653    let b_ts = CheckedTimestamp::try_from(NaiveDate::from(b).and_hms_opt(0, 0, 0).unwrap())?;
654    let diff = b_ts.diff_as(&a_ts, unit)?;
655    Ok(diff)
656}
657
658#[sqlfunc(sqlname = "datediff")]
659fn date_diff_time(unit_str: &str, a: NaiveTime, b: NaiveTime) -> Result<i64, EvalError> {
660    let unit = unit_str
661        .parse()
662        .map_err(|_| EvalError::InvalidDatePart(unit_str.into()))?;
663
664    // Convert the Time into a timestamp so we can calculate age.
665    let a_ts =
666        CheckedTimestamp::try_from(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap().and_time(a))?;
667    let b_ts =
668        CheckedTimestamp::try_from(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap().and_time(b))?;
669    let diff = b_ts.diff_as(&a_ts, unit)?;
670    Ok(diff)
671}
672
673#[sqlfunc(sqlname = "datediff")]
674fn date_diff_timestamp(
675    unit: &str,
676    a: CheckedTimestamp<NaiveDateTime>,
677    b: CheckedTimestamp<NaiveDateTime>,
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#[sqlfunc(sqlname = "datediff")]
688fn date_diff_timestamp_tz(
689    unit: &str,
690    a: CheckedTimestamp<DateTime<Utc>>,
691    b: CheckedTimestamp<DateTime<Utc>>,
692) -> Result<i64, EvalError> {
693    let unit = unit
694        .parse()
695        .map_err(|_| EvalError::InvalidDatePart(unit.into()))?;
696
697    let diff = b.diff_as(&a, unit)?;
698    Ok(diff)
699}
700
701#[derive(
702    Ord,
703    PartialOrd,
704    Clone,
705    Debug,
706    Eq,
707    PartialEq,
708    Serialize,
709    Deserialize,
710    Hash
711)]
712pub struct ErrorIfNull;
713
714impl fmt::Display for ErrorIfNull {
715    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
716        f.write_str("error_if_null")
717    }
718}
719
720impl LazyVariadicFunc for ErrorIfNull {
721    fn eval<'a>(
722        &'a self,
723        datums: &[Datum<'a>],
724        temp_storage: &'a RowArena,
725        exprs: &'a [impl Eval],
726    ) -> Result<Datum<'a>, EvalError> {
727        let first = exprs[0].eval(datums, temp_storage)?;
728        match first {
729            Datum::Null => {
730                let err_msg = match exprs[1].eval(datums, temp_storage)? {
731                    Datum::Null => {
732                        return Err(EvalError::Internal(
733                            "unexpected NULL in error side of error_if_null".into(),
734                        ));
735                    }
736                    o => o.unwrap_str(),
737                };
738                Err(EvalError::IfNullError(err_msg.into()))
739            }
740            _ => Ok(first),
741        }
742    }
743
744    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
745        input_types[0].scalar_type.clone().nullable(false)
746    }
747
748    fn propagates_nulls(&self) -> bool {
749        false
750    }
751
752    fn introduces_nulls(&self) -> bool {
753        false
754    }
755}
756
757#[derive(
758    Ord,
759    PartialOrd,
760    Clone,
761    Debug,
762    Eq,
763    PartialEq,
764    Serialize,
765    Deserialize,
766    Hash
767)]
768pub struct Greatest;
769
770impl fmt::Display for Greatest {
771    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
772        f.write_str("greatest")
773    }
774}
775
776impl LazyVariadicFunc for Greatest {
777    fn eval<'a>(
778        &'a self,
779        datums: &[Datum<'a>],
780        temp_storage: &'a RowArena,
781        exprs: &'a [impl Eval],
782    ) -> Result<Datum<'a>, EvalError> {
783        let datums = fallible_iterator::convert(exprs.iter().map(|e| e.eval(datums, temp_storage)));
784        Ok(datums
785            .filter(|d| Ok(!d.is_null()))
786            .max()?
787            .unwrap_or(Datum::Null))
788    }
789
790    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
791        SqlColumnType::union_many(input_types)
792    }
793
794    fn propagates_nulls(&self) -> bool {
795        false
796    }
797
798    fn introduces_nulls(&self) -> bool {
799        false
800    }
801
802    fn could_error(&self) -> bool {
803        false
804    }
805
806    fn is_monotone(&self) -> bool {
807        true
808    }
809
810    fn is_associative(&self) -> bool {
811        true
812    }
813}
814
815#[sqlfunc(sqlname = "hmac")]
816fn hmac_string(to_digest: &str, key: &str, typ: &str) -> Result<Vec<u8>, EvalError> {
817    let to_digest = to_digest.as_bytes();
818    let key = key.as_bytes();
819    hmac_inner(to_digest, key, typ)
820}
821
822#[sqlfunc(sqlname = "hmac")]
823fn hmac_bytes(to_digest: &[u8], key: &[u8], typ: &str) -> Result<Vec<u8>, EvalError> {
824    hmac_inner(to_digest, key, typ)
825}
826
827pub fn hmac_inner(to_digest: &[u8], key: &[u8], typ: &str) -> Result<Vec<u8>, EvalError> {
828    match typ {
829        "md5" => {
830            let mut mac = Hmac::<Md5>::new_from_slice(key).expect("HMAC accepts any key size");
831            mac.update(to_digest);
832            Ok(mac.finalize().into_bytes().to_vec())
833        }
834        "sha1" => {
835            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA1_FOR_LEGACY_USE_ONLY, key);
836            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
837        }
838        "sha224" => {
839            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA224, key);
840            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
841        }
842        "sha256" => {
843            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA256, key);
844            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
845        }
846        "sha384" => {
847            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA384, key);
848            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
849        }
850        "sha512" => {
851            let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA512, key);
852            Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
853        }
854        other => Err(EvalError::InvalidHashAlgorithm(other.into())),
855    }
856}
857
858#[sqlfunc]
859fn jsonb_build_array<'a>(
860    datums: Variadic<Datum<'a>>,
861    temp_storage: &'a RowArena,
862) -> Result<JsonbRef<'a>, EvalError> {
863    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
864    let datum = temp_storage.make_datum(|packer| {
865        packer.push_list(datums.into_iter().map(|d| match d {
866            Datum::Null => Datum::JsonNull,
867            d => d,
868        }))
869    });
870    Ok(JsonbRef::from_datum(datum))
871}
872
873#[sqlfunc]
874fn jsonb_build_object<'a>(
875    mut kvs: Variadic<(Datum<'a>, Datum<'a>)>,
876    temp_storage: &'a RowArena,
877) -> Result<JsonbRef<'a>, EvalError> {
878    kvs.0.sort_by(|kv1, kv2| kv1.0.cmp(&kv2.0));
879    kvs.0.dedup_by(|kv1, kv2| kv1.0 == kv2.0);
880    // Checked after the dedup, so duplicate keys are counted once, as they are packed.
881    check_datums_fit_budget(kvs.0.iter().flat_map(|(k, v)| [*k, *v]), temp_storage)?;
882    let datum = temp_storage.try_make_datum(|packer| {
883        packer.push_dict_with(|packer| {
884            for (k, v) in kvs {
885                if k.is_null() {
886                    return Err(EvalError::KeyCannotBeNull);
887                }
888                let v = match v {
889                    Datum::Null => Datum::JsonNull,
890                    d => d,
891                };
892                packer.push(k);
893                packer.push(v);
894            }
895            Ok(())
896        })
897    })?;
898    Ok(JsonbRef::from_datum(datum))
899}
900
901#[derive(
902    Ord,
903    PartialOrd,
904    Clone,
905    Debug,
906    Eq,
907    PartialEq,
908    Serialize,
909    Deserialize,
910    Hash
911)]
912pub struct Least;
913
914impl fmt::Display for Least {
915    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
916        f.write_str("least")
917    }
918}
919
920impl LazyVariadicFunc for Least {
921    fn eval<'a>(
922        &'a self,
923        datums: &[Datum<'a>],
924        temp_storage: &'a RowArena,
925        exprs: &'a [impl Eval],
926    ) -> Result<Datum<'a>, EvalError> {
927        let datums = fallible_iterator::convert(exprs.iter().map(|e| e.eval(datums, temp_storage)));
928        Ok(datums
929            .filter(|d| Ok(!d.is_null()))
930            .min()?
931            .unwrap_or(Datum::Null))
932    }
933
934    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
935        SqlColumnType::union_many(input_types)
936    }
937
938    fn propagates_nulls(&self) -> bool {
939        false
940    }
941
942    fn introduces_nulls(&self) -> bool {
943        false
944    }
945
946    fn could_error(&self) -> bool {
947        false
948    }
949
950    fn is_monotone(&self) -> bool {
951        true
952    }
953
954    fn is_associative(&self) -> bool {
955        true
956    }
957}
958
959#[derive(
960    Ord,
961    PartialOrd,
962    Clone,
963    Debug,
964    Eq,
965    PartialEq,
966    Serialize,
967    Deserialize,
968    Hash
969)]
970pub struct ListCreate {
971    pub elem_type: SqlScalarType,
972}
973
974#[sqlfunc(
975    output_type_expr = "SqlScalarType::List { element_type: Box::new(self.elem_type.clone()), custom_id: None }.nullable(false)",
976    introduces_nulls = false
977)]
978fn list_create<'a>(
979    &self,
980    datums: Variadic<Datum<'a>>,
981    temp_storage: &'a RowArena,
982) -> Result<Datum<'a>, EvalError> {
983    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
984    Ok(temp_storage.make_datum(|packer| packer.push_list(datums)))
985}
986
987#[derive(
988    Ord,
989    PartialOrd,
990    Clone,
991    Debug,
992    Eq,
993    PartialEq,
994    Serialize,
995    Deserialize,
996    Hash
997)]
998pub struct RecordCreate {
999    pub field_names: Vec<ColumnName>,
1000}
1001
1002#[sqlfunc(
1003    output_type_expr = "SqlScalarType::Record { fields: self.field_names.clone().into_iter().zip_eq(input_types.iter().cloned()).collect(), custom_id: None }.nullable(false)",
1004    introduces_nulls = false
1005)]
1006fn record_create<'a>(
1007    &self,
1008    datums: Variadic<Datum<'a>>,
1009    temp_storage: &'a RowArena,
1010) -> Result<Datum<'a>, EvalError> {
1011    check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
1012    Ok(temp_storage.make_datum(|packer| packer.push_list(datums.iter().copied())))
1013}
1014
1015#[sqlfunc(
1016    output_type_expr = "input_types[0].scalar_type.unwrap_list_nth_layer_type(input_types.len() - 1).clone().nullable(true)",
1017    introduces_nulls = true
1018)]
1019// TODO(benesch): remove potentially dangerous usage of `as`.
1020#[allow(clippy::as_conversions)]
1021fn list_index<'a>(buf: DatumList<'a>, indices: Variadic<i64>) -> Datum<'a> {
1022    let mut buf = Datum::List(buf);
1023    for i in indices {
1024        if buf.is_null() {
1025            break;
1026        }
1027        if i < 1 {
1028            return Datum::Null;
1029        }
1030
1031        buf = match buf.unwrap_list().iter().nth(i as usize - 1) {
1032            Some(datum) => datum,
1033            None => return Datum::Null,
1034        }
1035    }
1036    buf
1037}
1038
1039#[sqlfunc(sqlname = "makeaclitem")]
1040fn make_acl_item(
1041    grantee_oid: u32,
1042    grantor_oid: u32,
1043    privileges: &str,
1044    is_grantable: bool,
1045) -> Result<AclItem, EvalError> {
1046    let grantee = Oid(grantee_oid);
1047    let grantor = Oid(grantor_oid);
1048    let acl_mode = AclMode::parse_multiple_privileges(privileges)
1049        .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
1050    if is_grantable {
1051        return Err(EvalError::Unsupported {
1052            feature: "GRANT OPTION".into(),
1053            discussion_no: None,
1054        });
1055    }
1056
1057    Ok(AclItem {
1058        grantee,
1059        grantor,
1060        acl_mode,
1061    })
1062}
1063
1064#[sqlfunc(sqlname = "make_mz_aclitem")]
1065fn make_mz_acl_item(
1066    grantee_str: &str,
1067    grantor_str: &str,
1068    privileges: &str,
1069) -> Result<MzAclItem, EvalError> {
1070    let grantee: RoleId = grantee_str
1071        .parse()
1072        .map_err(|e: anyhow::Error| EvalError::InvalidRoleId(e.to_string().into()))?;
1073    let grantor: RoleId = grantor_str
1074        .parse()
1075        .map_err(|e: anyhow::Error| EvalError::InvalidRoleId(e.to_string().into()))?;
1076    if grantor == RoleId::Public {
1077        return Err(EvalError::InvalidRoleId(
1078            "mz_aclitem grantor cannot be PUBLIC role".into(),
1079        ));
1080    }
1081    let acl_mode = AclMode::parse_multiple_privileges(privileges)
1082        .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
1083
1084    Ok(MzAclItem {
1085        grantee,
1086        grantor,
1087        acl_mode,
1088    })
1089}
1090
1091#[sqlfunc(sqlname = "makets")]
1092#[allow(clippy::as_conversions)]
1093fn make_timestamp(
1094    year: i64,
1095    month: i64,
1096    day: i64,
1097    hour: i64,
1098    minute: i64,
1099    second_float: f64,
1100) -> Result<Option<CheckedTimestamp<NaiveDateTime>>, EvalError> {
1101    // NOTE: This function never returns `Ok(None)`. The `Option` stays so
1102    // that the derived `introduces_nulls` (and with it the nullability of
1103    // existing objects planned with this function) does not change.
1104
1105    // Negative years are BC and there is no year 0. Chrono's astronomical
1106    // numbering has 1 BC as year 0, so BC years shift by one.
1107    if year == 0 {
1108        return Err(EvalError::DateOutOfRange);
1109    }
1110    let year = if year < 0 { year + 1 } else { year };
1111    let year: i32 = year.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1112    let month: u32 = month.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1113    let day: u32 = day.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1114    let date = NaiveDate::from_ymd_opt(year, month, day).ok_or(EvalError::DateOutOfRange)?;
1115    // Allow exactly 24:00:00 (rolls over to midnight of the next day) and
1116    // 60 seconds (rolls over into the next minute).
1117    if !(0..=24).contains(&hour)
1118        || !(0..=59).contains(&minute)
1119        || !(0.0..=60.0).contains(&second_float)
1120        || (hour == 24 && (minute > 0 || second_float > 0.0))
1121    {
1122        return Err(EvalError::TimestampOutOfRange);
1123    }
1124    // The cast is exact: second_float is in [0, 60].
1125    let micros = (second_float * 1_000_000.0).round() as i64;
1126    let time = chrono::Duration::microseconds((hour * 60 + minute) * 60 * 1_000_000 + micros);
1127    let timestamp = date
1128        .and_hms_opt(0, 0, 0)
1129        .unwrap()
1130        .checked_add_signed(time)
1131        .ok_or(EvalError::TimestampOutOfRange)?;
1132    Ok(Some(timestamp.try_into()?))
1133}
1134
1135#[derive(
1136    Ord,
1137    PartialOrd,
1138    Clone,
1139    Debug,
1140    Eq,
1141    PartialEq,
1142    Serialize,
1143    Deserialize,
1144    Hash
1145)]
1146pub struct MapBuild {
1147    pub value_type: SqlScalarType,
1148}
1149
1150#[sqlfunc(
1151    output_type_expr = "SqlScalarType::Map { value_type: Box::new(self.value_type.clone()), custom_id: None }.nullable(false)",
1152    introduces_nulls = false
1153)]
1154fn map_build<'a>(
1155    &self,
1156    datums: Variadic<(Option<&str>, Datum<'a>)>,
1157    temp_storage: &'a RowArena,
1158) -> Result<Datum<'a>, EvalError> {
1159    // Collect into a `BTreeMap` to provide the same semantics as it.
1160    let map: std::collections::BTreeMap<&str, _> = datums
1161        .into_iter()
1162        .filter_map(|(k, v)| k.map(|k| (k, v)))
1163        .collect();
1164
1165    // Checked after the collect, so duplicate keys are counted once, as they are packed. The
1166    // `BTreeMap` transient needs no pre-check: it holds one entry per argument, a count the query
1167    // text fixes, not the input.
1168    check_datums_fit_budget(
1169        map.iter().flat_map(|(k, v)| [Datum::String(k), *v]),
1170        temp_storage,
1171    )?;
1172    Ok(temp_storage.make_datum(|packer| packer.push_dict(map)))
1173}
1174
1175#[derive(
1176    Ord,
1177    PartialOrd,
1178    Clone,
1179    Debug,
1180    Eq,
1181    PartialEq,
1182    Serialize,
1183    Deserialize,
1184    Hash
1185)]
1186pub struct Or;
1187
1188impl fmt::Display for Or {
1189    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1190        f.write_str("OR")
1191    }
1192}
1193
1194impl LazyVariadicFunc for Or {
1195    fn eval<'a>(
1196        &'a self,
1197        datums: &[Datum<'a>],
1198        temp_storage: &'a RowArena,
1199        exprs: &'a [impl Eval],
1200    ) -> Result<Datum<'a>, EvalError> {
1201        // If any is true, then return true. Else, if any is null, then return null. Else, return false.
1202        let mut null = false;
1203        let mut err = None;
1204        for expr in exprs {
1205            match expr.eval(datums, temp_storage) {
1206                Ok(Datum::False) => {}
1207                Ok(Datum::True) => return Ok(Datum::True), // short-circuit
1208                // No return in these two cases, because we might still see a true
1209                Ok(Datum::Null) => null = true,
1210                Err(this_err) => err = std::cmp::max(err.take(), Some(this_err)),
1211                _ => unreachable!(),
1212            }
1213        }
1214        match (err, null) {
1215            (Some(err), _) => Err(err),
1216            (None, true) => Ok(Datum::Null),
1217            (None, false) => Ok(Datum::False),
1218        }
1219    }
1220
1221    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
1222        let in_nullable = input_types.iter().any(|t| t.nullable);
1223        SqlScalarType::Bool.nullable(in_nullable)
1224    }
1225
1226    fn propagates_nulls(&self) -> bool {
1227        false
1228    }
1229
1230    fn introduces_nulls(&self) -> bool {
1231        false
1232    }
1233
1234    fn could_error(&self) -> bool {
1235        false
1236    }
1237
1238    fn is_monotone(&self) -> bool {
1239        true
1240    }
1241
1242    fn is_associative(&self) -> bool {
1243        true
1244    }
1245
1246    fn is_infix_op(&self) -> bool {
1247        true
1248    }
1249}
1250
1251#[sqlfunc(sqlname = "lpad")]
1252fn pad_leading(
1253    string: &str,
1254    raw_len: i32,
1255    pad: OptionalArg<&str>,
1256    temp_storage: &RowArena,
1257) -> Result<String, EvalError> {
1258    let len = match usize::try_from(raw_len) {
1259        Ok(len) => len,
1260        Err(_) => {
1261            return Err(EvalError::InvalidParameterValue(
1262                "length must be nonnegative".into(),
1263            ));
1264        }
1265    };
1266    let pad_string = pad.unwrap_or(" ");
1267
1268    // `len` is a character count but the ceiling is in bytes, so a multi-byte pad character writes
1269    // several bytes per character the check allowed. Bound the bytes instead. The result is the kept
1270    // prefix of `string` plus `len - string_chars` padding characters, each at most as wide as the
1271    // widest character in `pad_string`. Truncation (`len <= string_chars`) writes at most
1272    // `string.len()`, which the same expression covers.
1273    let widest_pad = pad_string.chars().map(char::len_utf8).max().unwrap_or(0);
1274    let pad_chars = len.saturating_sub(string.chars().count());
1275    let max_result_bytes = string
1276        .len()
1277        .saturating_add(pad_chars.saturating_mul(widest_pad));
1278    if max_result_bytes > max_string_func_result_bytes(temp_storage) {
1279        return Err(EvalError::LengthTooLarge);
1280    }
1281
1282    let (end_char, end_char_byte_offset) = string
1283        .chars()
1284        .take(len)
1285        .fold((0, 0), |acc, char| (acc.0 + 1, acc.1 + char.len_utf8()));
1286
1287    // NOTE: reserve from the bytes actually written, not from `len`, which is a character count.
1288    // Reserving `len` is too little for a multi-byte pad, and unbounded for an empty one, which
1289    // writes nothing and so has no result size for the ceiling above to refuse.
1290    let truncating = len == end_char;
1291    let mut buf = String::with_capacity(if truncating {
1292        end_char_byte_offset
1293    } else {
1294        max_result_bytes
1295    });
1296    if truncating {
1297        buf.push_str(&string[0..end_char_byte_offset]);
1298    } else {
1299        buf.extend(pad_string.chars().cycle().take(len - end_char));
1300        buf.push_str(string);
1301    }
1302
1303    Ok(buf)
1304}
1305
1306#[sqlfunc(
1307    output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(true)",
1308    introduces_nulls = true
1309)]
1310fn regexp_match<'a>(
1311    haystack: &'a str,
1312    needle: &str,
1313    flags: OptionalArg<&str>,
1314    temp_storage: &'a RowArena,
1315) -> Result<Datum<'a>, EvalError> {
1316    let flags = flags.unwrap_or("");
1317    let needle = build_regex(needle, flags)?;
1318    regexp_match_static(Datum::String(haystack), temp_storage, &needle)
1319}
1320
1321#[sqlfunc(
1322    output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)",
1323    introduces_nulls = false
1324)]
1325fn regexp_split_to_array<'a>(
1326    text: &str,
1327    regexp_str: &str,
1328    flags: OptionalArg<&str>,
1329    temp_storage: &'a RowArena,
1330) -> Result<Datum<'a>, EvalError> {
1331    let flags = flags.unwrap_or("");
1332    let regexp = build_regex(regexp_str, flags)?;
1333    regexp_split_to_array_re(text, &regexp, temp_storage)
1334}
1335
1336#[sqlfunc]
1337fn regexp_replace<'a>(
1338    source: &'a str,
1339    pattern: &str,
1340    replacement: &str,
1341    flags_opt: OptionalArg<&str>,
1342    temp_storage: &RowArena,
1343) -> Result<Cow<'a, str>, EvalError> {
1344    let flags = flags_opt.0.unwrap_or("");
1345    let (limit, flags) = regexp_replace_parse_flags(flags);
1346    let regexp = build_regex(pattern, &flags)?;
1347
1348    // The result is the unreplaced parts of `source` plus one replacement per match. A `$ref` in the
1349    // replacement expands to captured text, which is at most the match, and matches are disjoint, so
1350    // all the expansions together cost at most one `source` per `$`.
1351    let dollars = replacement.matches('$').count();
1352    let max_result_bytes = |matches: usize| {
1353        let expansions = match matches {
1354            0 => 0,
1355            _ => dollars.saturating_mul(source.len()),
1356        };
1357        source
1358            .len()
1359            .saturating_add(matches.saturating_mul(replacement.len()))
1360            .saturating_add(expansions)
1361    };
1362    // A `limit` of 0 means replace every match, and a pattern that matches the empty string matches
1363    // at every position. Assume that first, and only pay for counting the real matches when the
1364    // assumption is over the ceiling. Same compromise as `replace`.
1365    let most_matches = match limit {
1366        0 => source.len().saturating_add(1),
1367        limit => limit,
1368    };
1369    let ceiling = max_string_func_result_bytes(temp_storage);
1370    if max_result_bytes(most_matches) > ceiling {
1371        let matches = regexp.find_iter(source).take(most_matches).count();
1372        if max_result_bytes(matches) > ceiling {
1373            return Err(EvalError::LengthTooLarge);
1374        }
1375    }
1376
1377    Ok(regexp.replacen(source, limit, replacement))
1378}
1379
1380#[sqlfunc]
1381fn replace(text: &str, from: &str, to: &str, temp_storage: &RowArena) -> Result<String, EvalError> {
1382    // As a compromise to avoid always nearly duplicating the work of replace by doing size estimation,
1383    // we first check if it's possible for the fully replaced string to exceed the limit by assuming that
1384    // every possible substring is replaced.
1385    //
1386    // If that estimate exceeds the limit, we then do a more precise (and expensive) estimate by counting
1387    // the actual number of replacements that would occur, and using that to calculate the final size.
1388    let max_result_bytes = max_string_func_result_bytes(temp_storage);
1389    let possible_size = text.len() * to.len();
1390    if possible_size > max_result_bytes {
1391        let replacement_count = text.matches(from).count();
1392        let estimated_size = text.len() + replacement_count * (to.len().saturating_sub(from.len()));
1393        if estimated_size > max_result_bytes {
1394            return Err(EvalError::LengthTooLarge);
1395        }
1396    }
1397
1398    Ok(text.replace(from, to))
1399}
1400
1401#[sqlfunc(
1402    output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)",
1403    introduces_nulls = false,
1404    propagates_nulls = false
1405)]
1406fn string_to_array<'a>(
1407    string: &'a str,
1408    delimiter: Option<&'a str>,
1409    null_string: OptionalArg<Option<&'a str>>,
1410    temp_storage: &'a RowArena,
1411) -> Result<Datum<'a>, EvalError> {
1412    if string.is_empty() {
1413        let mut row = Row::default();
1414        let mut packer = row.packer();
1415        packer.try_push_array(&[], std::iter::empty::<Datum>())?;
1416
1417        return Ok(temp_storage.push_unary_row(row));
1418    }
1419
1420    let Some(delimiter) = delimiter else {
1421        let split_all_chars_delimiter = "";
1422        return string_to_array_impl(
1423            string,
1424            split_all_chars_delimiter,
1425            null_string.flatten(),
1426            temp_storage,
1427        );
1428    };
1429
1430    if delimiter.is_empty() {
1431        let mut row = Row::default();
1432        let mut packer = row.packer();
1433        let dims = &[ArrayDimension {
1434            lower_bound: 1,
1435            length: 1,
1436        }];
1437        match null_string.flatten() {
1438            Some(null_string) if null_string == string => {
1439                packer.try_push_array(dims, std::iter::once(Datum::Null))?;
1440            }
1441            _ => {
1442                packer.try_push_array(dims, vec![string].into_iter().map(Datum::String))?;
1443            }
1444        }
1445        Ok(temp_storage.push_unary_row(row))
1446    } else {
1447        string_to_array_impl(string, delimiter, null_string.flatten(), temp_storage)
1448    }
1449}
1450
1451fn string_to_array_impl<'a>(
1452    string: &str,
1453    delimiter: &str,
1454    null_string: Option<&'a str>,
1455    temp_storage: &'a RowArena,
1456) -> Result<Datum<'a>, EvalError> {
1457    // Bound the transient `Vec<&str>` before it is collected. The count walks the same iterator the
1458    // collect does, so the two agree.
1459    check_build_fits_budget(
1460        || {
1461            if delimiter.is_empty() {
1462                string.split(delimiter).filter(|s| !s.is_empty()).count()
1463            } else {
1464                string.split(delimiter).count()
1465            }
1466        },
1467        std::mem::size_of::<&str>(),
1468        temp_storage,
1469    )?;
1470
1471    let mut row = Row::default();
1472    let mut packer = row.packer();
1473
1474    let result = string.split(delimiter);
1475    let found: Vec<&str> = if delimiter.is_empty() {
1476        result.filter(|s| !s.is_empty()).collect()
1477    } else {
1478        result.collect()
1479    };
1480    let array_dimensions = [ArrayDimension {
1481        lower_bound: 1,
1482        length: found.len(),
1483    }];
1484
1485    // Splitting amplifies: each chunk is packed with its own tag and length, so a delimiter that
1486    // splits per character costs several times the input. `Datum::Null` packs smaller than the
1487    // chunk it replaces, so treating every chunk as a string is an upper bound.
1488    check_datums_fit_budget(found.iter().copied().map(Datum::String), temp_storage)?;
1489
1490    if let Some(null_string) = null_string {
1491        let found_datums = found.into_iter().map(|chunk| {
1492            if chunk.eq(null_string) {
1493                Datum::Null
1494            } else {
1495                Datum::String(chunk)
1496            }
1497        });
1498
1499        packer.try_push_array(&array_dimensions, found_datums)?;
1500    } else {
1501        packer.try_push_array(&array_dimensions, found.into_iter().map(Datum::String))?;
1502    }
1503
1504    Ok(temp_storage.push_unary_row(row))
1505}
1506
1507#[sqlfunc]
1508fn substr<'a>(s: &'a str, start: i32, length: OptionalArg<i32>) -> Result<&'a str, EvalError> {
1509    let raw_start_idx = i64::from(start) - 1;
1510    let start_idx = match usize::try_from(cmp::max(raw_start_idx, 0)) {
1511        Ok(i) => i,
1512        Err(_) => {
1513            return Err(EvalError::InvalidParameterValue(
1514                format!(
1515                    "substring starting index ({}) exceeds min/max position",
1516                    raw_start_idx
1517                )
1518                .into(),
1519            ));
1520        }
1521    };
1522
1523    let mut char_indices = s.char_indices();
1524    let get_str_index = |(index, _char)| index;
1525
1526    let str_len = s.len();
1527    let start_char_idx = char_indices.nth(start_idx).map_or(str_len, get_str_index);
1528
1529    if let OptionalArg(Some(len)) = length {
1530        let end_idx = match i64::from(len) {
1531            e if e < 0 => {
1532                return Err(EvalError::InvalidParameterValue(
1533                    "negative substring length not allowed".into(),
1534                ));
1535            }
1536            e if e == 0 || e + raw_start_idx < 1 => return Ok(""),
1537            e => {
1538                let e = cmp::min(raw_start_idx + e - 1, e - 1);
1539                match usize::try_from(e) {
1540                    Ok(i) => i,
1541                    Err(_) => {
1542                        return Err(EvalError::InvalidParameterValue(
1543                            format!("substring length ({}) exceeds max position", e).into(),
1544                        ));
1545                    }
1546                }
1547            }
1548        };
1549
1550        let end_char_idx = char_indices.nth(end_idx).map_or(str_len, get_str_index);
1551
1552        Ok(&s[start_char_idx..end_char_idx])
1553    } else {
1554        Ok(&s[start_char_idx..])
1555    }
1556}
1557
1558#[sqlfunc(sqlname = "split_string")]
1559fn split_part<'a>(string: &'a str, delimiter: &str, field: i32) -> Result<&'a str, EvalError> {
1560    let index = match usize::try_from(i64::from(field) - 1) {
1561        Ok(index) => index,
1562        Err(_) => {
1563            return Err(EvalError::InvalidParameterValue(
1564                "field position must be greater than zero".into(),
1565            ));
1566        }
1567    };
1568
1569    // If the provided delimiter is the empty string,
1570    // PostgreSQL does not break the string into individual
1571    // characters. Instead, it generates the following parts: [string].
1572    if delimiter.is_empty() {
1573        if index == 0 {
1574            return Ok(string);
1575        } else {
1576            return Ok("");
1577        }
1578    }
1579
1580    // If provided index is greater than the number of split parts,
1581    // return an empty string.
1582    Ok(string.split(delimiter).nth(index).unwrap_or(""))
1583}
1584
1585#[sqlfunc(is_associative = true)]
1586fn concat(strs: Variadic<Option<&str>>, temp_storage: &RowArena) -> Result<String, EvalError> {
1587    let max_result_bytes = max_string_func_result_bytes(temp_storage);
1588    let mut total_size = 0;
1589    for s in &strs {
1590        if let Some(s) = s {
1591            total_size += s.len();
1592            if total_size > max_result_bytes {
1593                return Err(EvalError::LengthTooLarge);
1594            }
1595        }
1596    }
1597    let mut buf = String::with_capacity(total_size);
1598    for s in strs {
1599        if let Some(s) = s {
1600            buf.push_str(s);
1601        }
1602    }
1603    Ok(buf)
1604}
1605
1606#[sqlfunc]
1607fn concat_ws(
1608    ws: &str,
1609    rest: Variadic<Option<&str>>,
1610    temp_storage: &RowArena,
1611) -> Result<String, EvalError> {
1612    let max_result_bytes = max_string_func_result_bytes(temp_storage);
1613    let mut total_size = 0;
1614    for s in &rest {
1615        if let Some(s) = s {
1616            total_size += s.len();
1617            total_size += ws.len();
1618            if total_size > max_result_bytes {
1619                return Err(EvalError::LengthTooLarge);
1620            }
1621        }
1622    }
1623
1624    let buf = Itertools::join(&mut rest.into_iter().filter_map(|s| s), ws);
1625
1626    Ok(buf)
1627}
1628
1629#[sqlfunc]
1630fn translate(
1631    string: &str,
1632    from_str: &str,
1633    to_str: &str,
1634    temp_storage: &RowArena,
1635) -> Result<String, EvalError> {
1636    // Translating to a wider character amplifies: each byte of an ASCII input can become four.
1637    // Every output character is either the input character or a `to` character, so it costs at most
1638    // `max(len_utf8(c), widest)` bytes, and at most `len_utf8(c) + widest - 1` since a character is
1639    // at least one byte. Summed, that's the input's byte length plus one `widest - 1` per character.
1640    // Tight when every character widens, exactly the input length when `to` is ASCII, so we refuse
1641    // only calls that really would exceed the ceiling.
1642    let widest = to_str.chars().map(char::len_utf8).max().unwrap_or(0);
1643    let max_result_bytes = match widest.saturating_sub(1) {
1644        // An ASCII `to` widens nothing, so the result fits the input and counting characters isn't
1645        // worth a second pass.
1646        0 => string.len(),
1647        widening => string
1648            .len()
1649            .saturating_add(string.chars().count().saturating_mul(widening)),
1650    };
1651    if max_result_bytes > max_string_func_result_bytes(temp_storage) {
1652        return Err(EvalError::LengthTooLarge);
1653    }
1654
1655    let from = from_str.chars().collect::<Vec<_>>();
1656    let to = to_str.chars().collect::<Vec<_>>();
1657
1658    Ok(string
1659        .chars()
1660        .filter_map(|c| match from.iter().position(|f| f == &c) {
1661            Some(idx) => to.get(idx).copied(),
1662            None => Some(c),
1663        })
1664        .collect())
1665}
1666
1667#[sqlfunc(
1668    output_type_expr = "input_types[0].scalar_type.clone().nullable(false)",
1669    introduces_nulls = false
1670)]
1671// TODO(benesch): remove potentially dangerous usage of `as`.
1672#[allow(clippy::as_conversions)]
1673fn list_slice_linear<'a>(
1674    list: DatumList<'a>,
1675    first: (i64, i64),
1676    remainder: Variadic<(i64, i64)>,
1677    temp_storage: &'a RowArena,
1678) -> Datum<'a> {
1679    let mut start_idx = 0;
1680    let mut total_length = usize::MAX;
1681
1682    for (start, end) in std::iter::once(first).chain(remainder) {
1683        let start = std::cmp::max(start, 1);
1684
1685        // Result should be empty list.
1686        if start > end {
1687            start_idx = 0;
1688            total_length = 0;
1689            break;
1690        }
1691
1692        let start_inner = start as usize - 1;
1693        // Start index only moves to geq positions.
1694        start_idx += start_inner;
1695
1696        // Length index only moves to leq positions
1697        let length_inner = (end - start) as usize + 1;
1698        total_length = std::cmp::min(length_inner, total_length - start_inner);
1699    }
1700
1701    let iter = list.iter().skip(start_idx).take(total_length);
1702
1703    temp_storage.make_datum(|row| {
1704        row.push_list_with(|row| {
1705            // if iter is empty, will get the appropriate empty list.
1706            for d in iter {
1707                row.push(d);
1708            }
1709        });
1710    })
1711}
1712
1713#[sqlfunc(sqlname = "timestamp_bin")]
1714fn date_bin_timestamp(
1715    stride: Interval,
1716    source: CheckedTimestamp<NaiveDateTime>,
1717    origin: CheckedTimestamp<NaiveDateTime>,
1718) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
1719    date_bin(stride, source, origin)
1720}
1721
1722#[sqlfunc(sqlname = "timestamptz_bin")]
1723fn date_bin_timestamp_tz(
1724    stride: Interval,
1725    source: CheckedTimestamp<DateTime<Utc>>,
1726    origin: CheckedTimestamp<DateTime<Utc>>,
1727) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
1728    date_bin(stride, source, origin)
1729}
1730
1731#[sqlfunc(sqlname = "timezonet")]
1732fn timezone_time_variadic(
1733    tz_str: &str,
1734    time: NaiveTime,
1735    wall_time: CheckedTimestamp<DateTime<Utc>>,
1736) -> Result<NaiveTime, EvalError> {
1737    parse_timezone(tz_str, TimezoneSpec::Posix)
1738        .map(|tz| timezone_time(tz, time, &wall_time.naive_utc()))
1739}
1740pub(crate) trait LazyVariadicFunc: fmt::Display {
1741    fn eval<'a>(
1742        &'a self,
1743        datums: &[Datum<'a>],
1744        temp_storage: &'a RowArena,
1745        exprs: &'a [impl Eval],
1746    ) -> Result<Datum<'a>, EvalError>;
1747
1748    /// The output SqlColumnType of this function.
1749    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType;
1750
1751    /// Whether this function will produce NULL on NULL input.
1752    fn propagates_nulls(&self) -> bool;
1753
1754    /// Whether this function will produce NULL on non-NULL input.
1755    fn introduces_nulls(&self) -> bool;
1756
1757    /// Whether this function might error on non-error input.
1758    fn could_error(&self) -> bool {
1759        true
1760    }
1761
1762    /// Returns true if the function is monotone.
1763    fn is_monotone(&self) -> bool {
1764        false
1765    }
1766
1767    /// Returns true if the function is associative.
1768    fn is_associative(&self) -> bool {
1769        false
1770    }
1771
1772    /// Returns true if the function is an infix operator.
1773    fn is_infix_op(&self) -> bool {
1774        false
1775    }
1776}
1777
1778pub(crate) trait EagerVariadicFunc: fmt::Display {
1779    type Input<'a>: InputDatumType<'a, EvalError>;
1780    type Output<'a>: OutputDatumType<'a, EvalError>;
1781
1782    fn call<'a>(&self, input: Self::Input<'a>, temp_storage: &'a RowArena) -> Self::Output<'a>;
1783
1784    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType;
1785
1786    fn propagates_nulls(&self) -> bool {
1787        !Self::Input::nullable()
1788    }
1789
1790    fn introduces_nulls(&self) -> bool {
1791        Self::Output::nullable()
1792    }
1793
1794    fn could_error(&self) -> bool {
1795        Self::Output::fallible()
1796    }
1797
1798    fn is_monotone(&self) -> bool {
1799        false
1800    }
1801
1802    fn is_associative(&self) -> bool {
1803        false
1804    }
1805
1806    fn is_infix_op(&self) -> bool {
1807        false
1808    }
1809}
1810
1811/// Blanket `LazyVariadicFunc` impl for each eager type, bridging
1812/// expression evaluation and null propagation via `InputDatumType::try_from_iter`.
1813impl<T: EagerVariadicFunc> LazyVariadicFunc for T {
1814    fn eval<'a>(
1815        &'a self,
1816        datums: &[Datum<'a>],
1817        temp_storage: &'a RowArena,
1818        exprs: &'a [impl Eval],
1819    ) -> Result<Datum<'a>, EvalError> {
1820        let mut datums = exprs.iter().map(|e| e.eval(datums, temp_storage));
1821        match T::Input::try_from_iter(&mut datums) {
1822            Ok(input) => self.call(input, temp_storage).into_result(temp_storage),
1823            Err(Ok(None)) => Err(EvalError::Internal("missing parameter".into())),
1824            Err(Ok(Some(datum))) if datum.is_null() => Ok(datum),
1825            Err(Ok(Some(_datum))) => {
1826                // datum is _not_ NULL
1827                Err(EvalError::Internal("invalid input type".into()))
1828            }
1829            Err(Err(res)) => Err(res),
1830        }
1831    }
1832
1833    fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
1834        self.output_type(input_types)
1835    }
1836
1837    fn propagates_nulls(&self) -> bool {
1838        self.propagates_nulls()
1839    }
1840
1841    fn introduces_nulls(&self) -> bool {
1842        self.introduces_nulls()
1843    }
1844
1845    fn could_error(&self) -> bool {
1846        self.could_error()
1847    }
1848
1849    fn is_monotone(&self) -> bool {
1850        self.is_monotone()
1851    }
1852
1853    fn is_associative(&self) -> bool {
1854        self.is_associative()
1855    }
1856
1857    fn is_infix_op(&self) -> bool {
1858        self.is_infix_op()
1859    }
1860}
1861
1862derive_variadic! {
1863    Coalesce(Coalesce),
1864    Greatest(Greatest),
1865    Least(Least),
1866    Concat(Concat),
1867    ConcatWs(ConcatWs),
1868    MakeTimestamp(MakeTimestamp),
1869    PadLeading(PadLeading),
1870    Substr(Substr),
1871    Replace(Replace),
1872    JsonbBuildArray(JsonbBuildArray),
1873    JsonbBuildObject(JsonbBuildObject),
1874    MapBuild(MapBuild),
1875    ArrayCreate(ArrayCreate),
1876    ArrayToString(ArrayToString),
1877    ArrayIndex(ArrayIndex),
1878    ListCreate(ListCreate),
1879    RecordCreate(RecordCreate),
1880    ListIndex(ListIndex),
1881    ListSliceLinear(ListSliceLinear),
1882    SplitPart(SplitPart),
1883    RegexpMatch(RegexpMatch),
1884    HmacString(HmacString),
1885    HmacBytes(HmacBytes),
1886    ErrorIfNull(ErrorIfNull),
1887    DateBinTimestamp(DateBinTimestamp),
1888    DateBinTimestampTz(DateBinTimestampTz),
1889    DateDiffTimestamp(DateDiffTimestamp),
1890    DateDiffTimestampTz(DateDiffTimestampTz),
1891    DateDiffDate(DateDiffDate),
1892    DateDiffTime(DateDiffTime),
1893    And(And),
1894    Or(Or),
1895    RangeCreate(RangeCreate),
1896    MakeAclItem(MakeAclItem),
1897    MakeMzAclItem(MakeMzAclItem),
1898    Translate(Translate),
1899    ArrayPosition(ArrayPosition),
1900    ArrayFill(ArrayFill),
1901    StringToArray(StringToArray),
1902    TimezoneTimeVariadic(TimezoneTimeVariadic),
1903    RegexpSplitToArray(RegexpSplitToArray),
1904    RegexpReplace(RegexpReplace),
1905    CaseLiteral(CaseLiteral),
1906}
1907
1908impl VariadicFunc {
1909    pub fn switch_and_or(&self) -> Self {
1910        match self {
1911            VariadicFunc::And(_) => Or.into(),
1912            VariadicFunc::Or(_) => And.into(),
1913            _ => unreachable!(),
1914        }
1915    }
1916
1917    /// Gives the unit (u) of OR or AND, such that `u AND/OR x == x`.
1918    /// Note that a 0-arg AND/OR evaluates to unit_of_and_or.
1919    pub fn unit_of_and_or(&self) -> MirScalarExpr {
1920        match self {
1921            VariadicFunc::And(_) => MirScalarExpr::literal_true(),
1922            VariadicFunc::Or(_) => MirScalarExpr::literal_false(),
1923            _ => unreachable!(),
1924        }
1925    }
1926
1927    /// Gives the zero (z) of OR or AND, such that `z AND/OR x == z`.
1928    pub fn zero_of_and_or(&self) -> MirScalarExpr {
1929        match self {
1930            VariadicFunc::And(_) => MirScalarExpr::literal_false(),
1931            VariadicFunc::Or(_) => MirScalarExpr::literal_true(),
1932            _ => unreachable!(),
1933        }
1934    }
1935}