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