Skip to main content

mz_expr/scalar/
func.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
14use std::borrow::Cow;
15use std::cmp::Ordering;
16use std::convert::{TryFrom, TryInto};
17use std::str::FromStr;
18use std::{iter, str};
19
20use ::encoding::DecoderTrap;
21use ::encoding::label::encoding_from_whatwg_label;
22use chrono::{DateTime, Duration, NaiveDate, NaiveDateTime, TimeZone, Timelike, Utc};
23use chrono_tz::{OffsetComponents, OffsetName, Tz};
24use dec::OrderedDecimal;
25use itertools::Itertools;
26use md5::{Digest, Md5};
27use mz_expr_derive::sqlfunc;
28use mz_ore::cast::{self, CastFrom};
29use mz_ore::fmt::FormatBuffer;
30use mz_ore::lex::LexBuf;
31use mz_ore::option::OptionExt;
32use mz_pgrepr::Type;
33use mz_pgtz::timezone::{Timezone, TimezoneSpec};
34use mz_repr::adt::array::{Array, ArrayDimension};
35use mz_repr::adt::date::Date;
36use mz_repr::adt::interval::{Interval, RoundBehavior};
37use mz_repr::adt::jsonb::JsonbRef;
38use mz_repr::adt::mz_acl_item::{AclMode, MzAclItem};
39use mz_repr::adt::numeric::{self, Numeric};
40use mz_repr::adt::range::Range;
41use mz_repr::adt::regex::Regex;
42use mz_repr::adt::timestamp::{CheckedTimestamp, TimestampLike};
43use mz_repr::{
44    ArrayRustType, Datum, DatumList, DatumMap, ExcludeNull, FromDatum, InputDatumType, Row,
45    RowArena, SqlScalarType, strconv,
46};
47use mz_sql_parser::ast::display::FormatMode;
48use mz_sql_pretty::{PrettyConfig, pretty_str};
49use num::traits::CheckedNeg;
50use sha1::Sha1;
51use sha2::{Sha224, Sha256, Sha384, Sha512};
52use subtle::ConstantTimeEq;
53
54use crate::scalar::func::format::DateTimeFormat;
55use crate::{EvalError, like_pattern};
56
57#[macro_use]
58mod macros;
59mod binary;
60mod encoding;
61pub(crate) mod format;
62pub(crate) mod impls;
63mod unary;
64mod unmaterializable;
65pub mod variadic;
66
67pub use binary::BinaryFunc;
68pub use impls::*;
69pub use unary::{EagerUnaryFunc, LazyUnaryFunc, UnaryFunc};
70pub use unmaterializable::UnmaterializableFunc;
71pub use variadic::VariadicFunc;
72
73/// The maximum size of the result strings of certain string functions, such as `repeat` and `lpad`.
74/// Chosen to be the smallest number to keep our tests passing without changing. 100MiB is probably
75/// higher than what we want, but it's better than no limit.
76///
77/// Note: This number appears in our user-facing documentation in the function reference for every
78/// function where it applies.
79pub const MAX_STRING_FUNC_RESULT_BYTES: usize = 1024 * 1024 * 100;
80
81pub fn jsonb_stringify<'a>(a: Datum<'a>, temp_storage: &'a RowArena) -> Option<&'a str> {
82    match a {
83        Datum::JsonNull => None,
84        Datum::String(s) => Some(s),
85        _ => {
86            let s = cast_jsonb_to_string(JsonbRef::from_datum(a));
87            Some(temp_storage.push_string(s))
88        }
89    }
90}
91
92#[sqlfunc(
93    is_monotone = "(true, true)",
94    is_infix_op = true,
95    sqlname = "+",
96    propagates_nulls = true
97)]
98fn add_int16(a: i16, b: i16) -> Result<i16, EvalError> {
99    a.checked_add(b).ok_or(EvalError::NumericFieldOverflow)
100}
101
102#[sqlfunc(
103    is_monotone = "(true, true)",
104    is_infix_op = true,
105    sqlname = "+",
106    propagates_nulls = true
107)]
108fn add_int32(a: i32, b: i32) -> Result<i32, EvalError> {
109    a.checked_add(b).ok_or(EvalError::NumericFieldOverflow)
110}
111
112#[sqlfunc(
113    is_monotone = "(true, true)",
114    is_infix_op = true,
115    sqlname = "+",
116    propagates_nulls = true
117)]
118fn add_int64(a: i64, b: i64) -> Result<i64, EvalError> {
119    a.checked_add(b).ok_or(EvalError::NumericFieldOverflow)
120}
121
122#[sqlfunc(
123    is_monotone = "(true, true)",
124    is_infix_op = true,
125    sqlname = "+",
126    propagates_nulls = true
127)]
128fn add_uint16(a: u16, b: u16) -> Result<u16, EvalError> {
129    a.checked_add(b)
130        .ok_or_else(|| EvalError::UInt16OutOfRange(format!("{a} + {b}").into()))
131}
132
133#[sqlfunc(
134    is_monotone = "(true, true)",
135    is_infix_op = true,
136    sqlname = "+",
137    propagates_nulls = true
138)]
139fn add_uint32(a: u32, b: u32) -> Result<u32, EvalError> {
140    a.checked_add(b)
141        .ok_or_else(|| EvalError::UInt32OutOfRange(format!("{a} + {b}").into()))
142}
143
144#[sqlfunc(
145    is_monotone = "(true, true)",
146    is_infix_op = true,
147    sqlname = "+",
148    propagates_nulls = true
149)]
150fn add_uint64(a: u64, b: u64) -> Result<u64, EvalError> {
151    a.checked_add(b)
152        .ok_or_else(|| EvalError::UInt64OutOfRange(format!("{a} + {b}").into()))
153}
154
155#[sqlfunc(
156    is_monotone = "(true, true)",
157    is_infix_op = true,
158    sqlname = "+",
159    propagates_nulls = true
160)]
161fn add_float32(a: f32, b: f32) -> Result<f32, EvalError> {
162    let sum = a + b;
163    if sum.is_infinite() && !a.is_infinite() && !b.is_infinite() {
164        Err(EvalError::FloatOverflow)
165    } else {
166        Ok(sum)
167    }
168}
169
170#[sqlfunc(
171    is_monotone = "(true, true)",
172    is_infix_op = true,
173    sqlname = "+",
174    propagates_nulls = true
175)]
176fn add_float64(a: f64, b: f64) -> Result<f64, EvalError> {
177    let sum = a + b;
178    if sum.is_infinite() && !a.is_infinite() && !b.is_infinite() {
179        Err(EvalError::FloatOverflow)
180    } else {
181        Ok(sum)
182    }
183}
184
185#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "+")]
186fn add_timestamp_interval(
187    a: CheckedTimestamp<NaiveDateTime>,
188    b: Interval,
189) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
190    add_timestamplike_interval(a, b)
191}
192
193#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "+")]
194fn add_timestamp_tz_interval(
195    a: CheckedTimestamp<DateTime<Utc>>,
196    b: Interval,
197) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
198    add_timestamplike_interval(a, b)
199}
200
201fn add_timestamplike_interval<T>(
202    a: CheckedTimestamp<T>,
203    b: Interval,
204) -> Result<CheckedTimestamp<T>, EvalError>
205where
206    T: TimestampLike,
207{
208    let dt = a.date_time();
209    let dt = add_timestamp_months(&dt, b.months)?;
210    let dt = dt
211        .checked_add_signed(b.duration_as_chrono())
212        .ok_or(EvalError::TimestampOutOfRange)?;
213    Ok(CheckedTimestamp::from_timestamplike(T::from_date_time(dt))?)
214}
215
216#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "-")]
217fn sub_timestamp_interval(
218    a: CheckedTimestamp<NaiveDateTime>,
219    b: Interval,
220) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
221    sub_timestamplike_interval(a, b)
222}
223
224#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "-")]
225fn sub_timestamp_tz_interval(
226    a: CheckedTimestamp<DateTime<Utc>>,
227    b: Interval,
228) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
229    sub_timestamplike_interval(a, b)
230}
231
232fn sub_timestamplike_interval<T>(
233    a: CheckedTimestamp<T>,
234    b: Interval,
235) -> Result<CheckedTimestamp<T>, EvalError>
236where
237    T: TimestampLike,
238{
239    neg_interval_inner(b).and_then(|i| add_timestamplike_interval(a, i))
240}
241
242#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "+")]
243fn add_date_time(
244    date: Date,
245    time: chrono::NaiveTime,
246) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
247    let dt = NaiveDate::from(date)
248        .and_hms_nano_opt(time.hour(), time.minute(), time.second(), time.nanosecond())
249        .unwrap();
250    Ok(CheckedTimestamp::from_timestamplike(dt)?)
251}
252
253#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "+")]
254fn add_date_interval(
255    date: Date,
256    interval: Interval,
257) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
258    let dt = NaiveDate::from(date).and_hms_opt(0, 0, 0).unwrap();
259    let dt = add_timestamp_months(&dt, interval.months)?;
260    let dt = dt
261        .checked_add_signed(interval.duration_as_chrono())
262        .ok_or(EvalError::TimestampOutOfRange)?;
263    Ok(CheckedTimestamp::from_timestamplike(dt)?)
264}
265
266#[sqlfunc(
267    // <time> + <interval> wraps!
268    is_monotone = "(false, false)",
269    is_infix_op = true,
270    sqlname = "+",
271    propagates_nulls = true
272)]
273fn add_time_interval(time: chrono::NaiveTime, interval: Interval) -> chrono::NaiveTime {
274    let (t, _) = time.overflowing_add_signed(interval.duration_as_chrono());
275    t
276}
277
278#[sqlfunc(
279    is_monotone = "(true, false)",
280    output_type = "Numeric",
281    sqlname = "round",
282    propagates_nulls = true
283)]
284fn round_numeric_binary(a: OrderedDecimal<Numeric>, mut b: i32) -> Result<Numeric, EvalError> {
285    let mut a = a.0;
286    let mut cx = numeric::cx_datum();
287    let a_exp = a.exponent();
288    if a_exp > 0 && b > 0 || a_exp < 0 && -a_exp < b {
289        // This condition indicates:
290        // - a is a value without a decimal point, b is a positive number
291        // - a has a decimal point, but b is larger than its scale
292        // In both of these situations, right-pad the number with zeroes, which // is most easily done with rescale.
293
294        // Ensure rescale doesn't exceed max precision by putting a ceiling on
295        // b equal to the maximum remaining scale the value can support.
296        let max_remaining_scale = u32::from(numeric::NUMERIC_DATUM_MAX_PRECISION)
297            - (numeric::get_precision(&a) - numeric::get_scale(&a));
298        b = match i32::try_from(max_remaining_scale) {
299            Ok(max_remaining_scale) => std::cmp::min(b, max_remaining_scale),
300            Err(_) => b,
301        };
302        cx.rescale(&mut a, &numeric::Numeric::from(-b));
303    } else {
304        // To avoid invalid operations, clamp b to be within 1 more than the
305        // precision limit.
306        const MAX_P_LIMIT: i32 = 1 + cast::u8_to_i32(numeric::NUMERIC_DATUM_MAX_PRECISION);
307        b = std::cmp::min(MAX_P_LIMIT, b);
308        b = std::cmp::max(-MAX_P_LIMIT, b);
309        let mut b = numeric::Numeric::from(b);
310        // Shift by 10^b; this put digit to round to in the one's place.
311        cx.scaleb(&mut a, &b);
312        cx.round(&mut a);
313        // Negate exponent for shift back
314        cx.neg(&mut b);
315        cx.scaleb(&mut a, &b);
316    }
317
318    if cx.status().overflow() {
319        Err(EvalError::FloatOverflow)
320    } else if a.is_zero() {
321        // simpler than handling cases where exponent has gotten set to some
322        // value greater than the max precision, but all significant digits
323        // were rounded away.
324        Ok(numeric::Numeric::zero())
325    } else {
326        numeric::munge_numeric(&mut a).unwrap();
327        Ok(a)
328    }
329}
330
331#[sqlfunc(sqlname = "convert_from", propagates_nulls = true)]
332fn convert_from<'a>(a: &'a [u8], b: &str) -> Result<&'a str, EvalError> {
333    // Convert PostgreSQL-style encoding names[1] to WHATWG-style encoding names[2],
334    // which the encoding library uses[3].
335    // [1]: https://www.postgresql.org/docs/9.5/multibyte.html
336    // [2]: https://encoding.spec.whatwg.org/
337    // [3]: https://github.com/lifthrasiir/rust-encoding/blob/4e79c35ab6a351881a86dbff565c4db0085cc113/src/label.rs
338    let encoding_name = b.to_lowercase().replace('_', "-").into_boxed_str();
339
340    // Supporting other encodings is tracked by database-issues#797.
341    if encoding_from_whatwg_label(&encoding_name).map(|e| e.name()) != Some("utf-8") {
342        return Err(EvalError::InvalidEncodingName(encoding_name));
343    }
344
345    match str::from_utf8(a) {
346        Ok(from) => Ok(from),
347        Err(e) => Err(EvalError::InvalidByteSequence {
348            byte_sequence: e.to_string().into(),
349            encoding_name,
350        }),
351    }
352}
353
354#[sqlfunc]
355fn encode(bytes: &[u8], format: &str) -> Result<String, EvalError> {
356    let format = encoding::lookup_format(format)?;
357    Ok(format.encode(bytes))
358}
359
360#[sqlfunc]
361fn decode(string: &str, format: &str) -> Result<Vec<u8>, EvalError> {
362    let format = encoding::lookup_format(format)?;
363    let out = format.decode(string)?;
364    if out.len() > MAX_STRING_FUNC_RESULT_BYTES {
365        Err(EvalError::LengthTooLarge)
366    } else {
367        Ok(out)
368    }
369}
370
371#[sqlfunc(sqlname = "length", propagates_nulls = true)]
372fn encoded_bytes_char_length(a: &[u8], b: &str) -> Result<i32, EvalError> {
373    // Convert PostgreSQL-style encoding names[1] to WHATWG-style encoding names[2],
374    // which the encoding library uses[3].
375    // [1]: https://www.postgresql.org/docs/9.5/multibyte.html
376    // [2]: https://encoding.spec.whatwg.org/
377    // [3]: https://github.com/lifthrasiir/rust-encoding/blob/4e79c35ab6a351881a86dbff565c4db0085cc113/src/label.rs
378    let encoding_name = b.to_lowercase().replace('_', "-").into_boxed_str();
379
380    let enc = match encoding_from_whatwg_label(&encoding_name) {
381        Some(enc) => enc,
382        None => return Err(EvalError::InvalidEncodingName(encoding_name)),
383    };
384
385    let decoded_string = match enc.decode(a, DecoderTrap::Strict) {
386        Ok(s) => s,
387        Err(e) => {
388            return Err(EvalError::InvalidByteSequence {
389                byte_sequence: e.into(),
390                encoding_name,
391            });
392        }
393    };
394
395    let count = decoded_string.chars().count();
396    i32::try_from(count).map_err(|_| EvalError::Int32OutOfRange(count.to_string().into()))
397}
398
399// TODO(benesch): remove potentially dangerous usage of `as`.
400#[allow(clippy::as_conversions)]
401pub fn add_timestamp_months<T: TimestampLike>(
402    dt: &T,
403    mut months: i32,
404) -> Result<CheckedTimestamp<T>, EvalError> {
405    if months == 0 {
406        return Ok(CheckedTimestamp::from_timestamplike(dt.clone())?);
407    }
408
409    let (mut year, mut month, mut day) = (dt.year(), dt.month0() as i32, dt.day());
410    let years = months / 12;
411    year = year
412        .checked_add(years)
413        .ok_or(EvalError::TimestampOutOfRange)?;
414
415    months %= 12;
416    // positive modulus is easier to reason about
417    if months < 0 {
418        year -= 1;
419        months += 12;
420    }
421    year += (month + months) / 12;
422    month = (month + months) % 12;
423    // account for dt.month0
424    month += 1;
425
426    // handle going from January 31st to February by saturation
427    let mut new_d = chrono::NaiveDate::from_ymd_opt(year, month as u32, day);
428    while new_d.is_none() {
429        // If we have decremented day past 28 and are still receiving `None`,
430        // then we have generally overflowed `NaiveDate`.
431        if day < 28 {
432            return Err(EvalError::TimestampOutOfRange);
433        }
434        day -= 1;
435        new_d = chrono::NaiveDate::from_ymd_opt(year, month as u32, day);
436    }
437    let new_d = new_d.unwrap();
438
439    // Neither postgres nor mysql support leap seconds, so this should be safe.
440    //
441    // Both my testing and https://dba.stackexchange.com/a/105829 support the
442    // idea that we should ignore leap seconds
443    let new_dt = new_d
444        .and_hms_nano_opt(dt.hour(), dt.minute(), dt.second(), dt.nanosecond())
445        .unwrap();
446    let new_dt = T::from_date_time(new_dt);
447    Ok(CheckedTimestamp::from_timestamplike(new_dt)?)
448}
449
450#[sqlfunc(
451    is_monotone = "(true, true)",
452    is_infix_op = true,
453    sqlname = "+",
454    propagates_nulls = true
455)]
456fn add_numeric(
457    a: OrderedDecimal<Numeric>,
458    b: OrderedDecimal<Numeric>,
459) -> Result<Numeric, EvalError> {
460    let mut cx = numeric::cx_datum();
461    let mut a = a.0;
462    cx.add(&mut a, &b.0);
463    if cx.status().overflow() {
464        Err(EvalError::FloatOverflow)
465    } else {
466        Ok(a)
467    }
468}
469
470#[sqlfunc(
471    is_monotone = "(true, true)",
472    is_infix_op = true,
473    sqlname = "+",
474    propagates_nulls = true
475)]
476fn add_interval(a: Interval, b: Interval) -> Result<Interval, EvalError> {
477    a.checked_add(&b)
478        .ok_or_else(|| EvalError::IntervalOutOfRange(format!("{a} + {b}").into()))
479}
480
481#[sqlfunc(is_infix_op = true, sqlname = "&", propagates_nulls = true)]
482fn bit_and_int16(a: i16, b: i16) -> i16 {
483    a & b
484}
485
486#[sqlfunc(is_infix_op = true, sqlname = "&", propagates_nulls = true)]
487fn bit_and_int32(a: i32, b: i32) -> i32 {
488    a & b
489}
490
491#[sqlfunc(is_infix_op = true, sqlname = "&", propagates_nulls = true)]
492fn bit_and_int64(a: i64, b: i64) -> i64 {
493    a & b
494}
495
496#[sqlfunc(is_infix_op = true, sqlname = "&", propagates_nulls = true)]
497fn bit_and_uint16(a: u16, b: u16) -> u16 {
498    a & b
499}
500
501#[sqlfunc(is_infix_op = true, sqlname = "&", propagates_nulls = true)]
502fn bit_and_uint32(a: u32, b: u32) -> u32 {
503    a & b
504}
505
506#[sqlfunc(is_infix_op = true, sqlname = "&", propagates_nulls = true)]
507fn bit_and_uint64(a: u64, b: u64) -> u64 {
508    a & b
509}
510
511#[sqlfunc(is_infix_op = true, sqlname = "|", propagates_nulls = true)]
512fn bit_or_int16(a: i16, b: i16) -> i16 {
513    a | b
514}
515
516#[sqlfunc(is_infix_op = true, sqlname = "|", propagates_nulls = true)]
517fn bit_or_int32(a: i32, b: i32) -> i32 {
518    a | b
519}
520
521#[sqlfunc(is_infix_op = true, sqlname = "|", propagates_nulls = true)]
522fn bit_or_int64(a: i64, b: i64) -> i64 {
523    a | b
524}
525
526#[sqlfunc(is_infix_op = true, sqlname = "|", propagates_nulls = true)]
527fn bit_or_uint16(a: u16, b: u16) -> u16 {
528    a | b
529}
530
531#[sqlfunc(is_infix_op = true, sqlname = "|", propagates_nulls = true)]
532fn bit_or_uint32(a: u32, b: u32) -> u32 {
533    a | b
534}
535
536#[sqlfunc(is_infix_op = true, sqlname = "|", propagates_nulls = true)]
537fn bit_or_uint64(a: u64, b: u64) -> u64 {
538    a | b
539}
540
541#[sqlfunc(is_infix_op = true, sqlname = "#", propagates_nulls = true)]
542fn bit_xor_int16(a: i16, b: i16) -> i16 {
543    a ^ b
544}
545
546#[sqlfunc(is_infix_op = true, sqlname = "#", propagates_nulls = true)]
547fn bit_xor_int32(a: i32, b: i32) -> i32 {
548    a ^ b
549}
550
551#[sqlfunc(is_infix_op = true, sqlname = "#", propagates_nulls = true)]
552fn bit_xor_int64(a: i64, b: i64) -> i64 {
553    a ^ b
554}
555
556#[sqlfunc(is_infix_op = true, sqlname = "#", propagates_nulls = true)]
557fn bit_xor_uint16(a: u16, b: u16) -> u16 {
558    a ^ b
559}
560
561#[sqlfunc(is_infix_op = true, sqlname = "#", propagates_nulls = true)]
562fn bit_xor_uint32(a: u32, b: u32) -> u32 {
563    a ^ b
564}
565
566#[sqlfunc(is_infix_op = true, sqlname = "#", propagates_nulls = true)]
567fn bit_xor_uint64(a: u64, b: u64) -> u64 {
568    a ^ b
569}
570
571#[sqlfunc(is_infix_op = true, sqlname = "<<", propagates_nulls = true)]
572// TODO(benesch): remove potentially dangerous usage of `as`.
573#[allow(clippy::as_conversions)]
574fn bit_shift_left_int16(a: i16, b: i32) -> i16 {
575    // widen to i32 and then cast back to i16 in order emulate the C promotion rules used in by Postgres
576    // when the rhs in the 16-31 range, e.g. (1 << 17 should evaluate to 0)
577    // see https://github.com/postgres/postgres/blob/REL_14_STABLE/src/backend/utils/adt/int.c#L1460-L1476
578    let lhs: i32 = a as i32;
579    let rhs: u32 = b as u32;
580    lhs.wrapping_shl(rhs) as i16
581}
582
583#[sqlfunc(is_infix_op = true, sqlname = "<<", propagates_nulls = true)]
584// TODO(benesch): remove potentially dangerous usage of `as`.
585#[allow(clippy::as_conversions)]
586fn bit_shift_left_int32(lhs: i32, rhs: i32) -> i32 {
587    let rhs = rhs as u32;
588    lhs.wrapping_shl(rhs)
589}
590
591#[sqlfunc(is_infix_op = true, sqlname = "<<", propagates_nulls = true)]
592// TODO(benesch): remove potentially dangerous usage of `as`.
593#[allow(clippy::as_conversions)]
594fn bit_shift_left_int64(lhs: i64, rhs: i32) -> i64 {
595    let rhs = rhs as u32;
596    lhs.wrapping_shl(rhs)
597}
598
599#[sqlfunc(is_infix_op = true, sqlname = "<<", propagates_nulls = true)]
600// TODO(benesch): remove potentially dangerous usage of `as`.
601#[allow(clippy::as_conversions)]
602fn bit_shift_left_uint16(a: u16, b: u32) -> u16 {
603    // widen to u32 and then cast back to u16 in order emulate the C promotion rules used in by Postgres
604    // when the rhs in the 16-31 range, e.g. (1 << 17 should evaluate to 0)
605    // see https://github.com/postgres/postgres/blob/REL_14_STABLE/src/backend/utils/adt/int.c#L1460-L1476
606    let lhs: u32 = a as u32;
607    let rhs: u32 = b;
608    lhs.wrapping_shl(rhs) as u16
609}
610
611#[sqlfunc(is_infix_op = true, sqlname = "<<", propagates_nulls = true)]
612fn bit_shift_left_uint32(a: u32, b: u32) -> u32 {
613    let lhs = a;
614    let rhs = b;
615    lhs.wrapping_shl(rhs)
616}
617
618#[sqlfunc(
619    output_type = "u64",
620    is_infix_op = true,
621    sqlname = "<<",
622    propagates_nulls = true
623)]
624fn bit_shift_left_uint64(lhs: u64, rhs: u32) -> u64 {
625    lhs.wrapping_shl(rhs)
626}
627
628#[sqlfunc(is_infix_op = true, sqlname = ">>", propagates_nulls = true)]
629// TODO(benesch): remove potentially dangerous usage of `as`.
630#[allow(clippy::as_conversions)]
631fn bit_shift_right_int16(lhs: i16, rhs: i32) -> i16 {
632    // widen to i32 and then cast back to i16 in order emulate the C promotion rules used in by Postgres
633    // when the rhs in the 16-31 range, e.g. (-32767 >> 17 should evaluate to -1)
634    // see https://github.com/postgres/postgres/blob/REL_14_STABLE/src/backend/utils/adt/int.c#L1460-L1476
635    let lhs = lhs as i32;
636    let rhs = rhs as u32;
637    lhs.wrapping_shr(rhs) as i16
638}
639
640#[sqlfunc(is_infix_op = true, sqlname = ">>", propagates_nulls = true)]
641// TODO(benesch): remove potentially dangerous usage of `as`.
642#[allow(clippy::as_conversions)]
643fn bit_shift_right_int32(lhs: i32, rhs: i32) -> i32 {
644    lhs.wrapping_shr(rhs as u32)
645}
646
647#[sqlfunc(is_infix_op = true, sqlname = ">>", propagates_nulls = true)]
648// TODO(benesch): remove potentially dangerous usage of `as`.
649#[allow(clippy::as_conversions)]
650fn bit_shift_right_int64(lhs: i64, rhs: i32) -> i64 {
651    lhs.wrapping_shr(rhs as u32)
652}
653
654#[sqlfunc(is_infix_op = true, sqlname = ">>", propagates_nulls = true)]
655// TODO(benesch): remove potentially dangerous usage of `as`.
656#[allow(clippy::as_conversions)]
657fn bit_shift_right_uint16(lhs: u16, rhs: u32) -> u16 {
658    // widen to u32 and then cast back to u16 in order emulate the C promotion rules used in by Postgres
659    // when the rhs in the 16-31 range, e.g. (-32767 >> 17 should evaluate to -1)
660    // see https://github.com/postgres/postgres/blob/REL_14_STABLE/src/backend/utils/adt/int.c#L1460-L1476
661    let lhs = lhs as u32;
662    lhs.wrapping_shr(rhs) as u16
663}
664
665#[sqlfunc(is_infix_op = true, sqlname = ">>", propagates_nulls = true)]
666fn bit_shift_right_uint32(lhs: u32, rhs: u32) -> u32 {
667    lhs.wrapping_shr(rhs)
668}
669
670#[sqlfunc(is_infix_op = true, sqlname = ">>", propagates_nulls = true)]
671fn bit_shift_right_uint64(lhs: u64, rhs: u32) -> u64 {
672    lhs.wrapping_shr(rhs)
673}
674
675#[sqlfunc(
676    is_monotone = "(true, true)",
677    is_infix_op = true,
678    sqlname = "-",
679    propagates_nulls = true
680)]
681fn sub_int16(a: i16, b: i16) -> Result<i16, EvalError> {
682    a.checked_sub(b).ok_or(EvalError::NumericFieldOverflow)
683}
684
685#[sqlfunc(
686    is_monotone = "(true, true)",
687    is_infix_op = true,
688    sqlname = "-",
689    propagates_nulls = true
690)]
691fn sub_int32(a: i32, b: i32) -> Result<i32, EvalError> {
692    a.checked_sub(b).ok_or(EvalError::NumericFieldOverflow)
693}
694
695#[sqlfunc(
696    is_monotone = "(true, true)",
697    is_infix_op = true,
698    sqlname = "-",
699    propagates_nulls = true
700)]
701fn sub_int64(a: i64, b: i64) -> Result<i64, EvalError> {
702    a.checked_sub(b).ok_or(EvalError::NumericFieldOverflow)
703}
704
705#[sqlfunc(
706    is_monotone = "(true, true)",
707    is_infix_op = true,
708    sqlname = "-",
709    propagates_nulls = true
710)]
711fn sub_uint16(a: u16, b: u16) -> Result<u16, EvalError> {
712    a.checked_sub(b)
713        .ok_or_else(|| EvalError::UInt16OutOfRange(format!("{a} - {b}").into()))
714}
715
716#[sqlfunc(
717    is_monotone = "(true, true)",
718    is_infix_op = true,
719    sqlname = "-",
720    propagates_nulls = true
721)]
722fn sub_uint32(a: u32, b: u32) -> Result<u32, EvalError> {
723    a.checked_sub(b)
724        .ok_or_else(|| EvalError::UInt32OutOfRange(format!("{a} - {b}").into()))
725}
726
727#[sqlfunc(
728    is_monotone = "(true, true)",
729    is_infix_op = true,
730    sqlname = "-",
731    propagates_nulls = true
732)]
733fn sub_uint64(a: u64, b: u64) -> Result<u64, EvalError> {
734    a.checked_sub(b)
735        .ok_or_else(|| EvalError::UInt64OutOfRange(format!("{a} - {b}").into()))
736}
737
738#[sqlfunc(
739    is_monotone = "(true, true)",
740    is_infix_op = true,
741    sqlname = "-",
742    propagates_nulls = true
743)]
744fn sub_float32(a: f32, b: f32) -> Result<f32, EvalError> {
745    let difference = a - b;
746    if difference.is_infinite() && !a.is_infinite() && !b.is_infinite() {
747        Err(EvalError::FloatOverflow)
748    } else {
749        Ok(difference)
750    }
751}
752
753#[sqlfunc(
754    is_monotone = "(true, true)",
755    is_infix_op = true,
756    sqlname = "-",
757    propagates_nulls = true
758)]
759fn sub_float64(a: f64, b: f64) -> Result<f64, EvalError> {
760    let difference = a - b;
761    if difference.is_infinite() && !a.is_infinite() && !b.is_infinite() {
762        Err(EvalError::FloatOverflow)
763    } else {
764        Ok(difference)
765    }
766}
767
768#[sqlfunc(
769    is_monotone = "(true, true)",
770    is_infix_op = true,
771    sqlname = "-",
772    propagates_nulls = true
773)]
774fn sub_numeric(
775    a: OrderedDecimal<Numeric>,
776    b: OrderedDecimal<Numeric>,
777) -> Result<Numeric, EvalError> {
778    let mut cx = numeric::cx_datum();
779    let mut a = a.0;
780    cx.sub(&mut a, &b.0);
781    if cx.status().overflow() {
782        Err(EvalError::FloatOverflow)
783    } else {
784        Ok(a)
785    }
786}
787
788#[sqlfunc(
789    is_monotone = "(true, true)",
790    output_type = "Interval",
791    sqlname = "age",
792    propagates_nulls = true
793)]
794fn age_timestamp(
795    a: CheckedTimestamp<chrono::NaiveDateTime>,
796    b: CheckedTimestamp<chrono::NaiveDateTime>,
797) -> Result<Interval, EvalError> {
798    Ok(a.age(&b)?)
799}
800
801#[sqlfunc(is_monotone = "(true, true)", sqlname = "age", propagates_nulls = true)]
802fn age_timestamp_tz(
803    a: CheckedTimestamp<chrono::DateTime<Utc>>,
804    b: CheckedTimestamp<chrono::DateTime<Utc>>,
805) -> Result<Interval, EvalError> {
806    Ok(a.age(&b)?)
807}
808
809#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "-")]
810fn sub_timestamp(
811    a: CheckedTimestamp<NaiveDateTime>,
812    b: CheckedTimestamp<NaiveDateTime>,
813) -> Result<Interval, EvalError> {
814    Interval::from_chrono_duration(a - b)
815        .map_err(|e| EvalError::IntervalOutOfRange(e.to_string().into()))
816}
817
818#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "-")]
819fn sub_timestamp_tz(
820    a: CheckedTimestamp<chrono::DateTime<Utc>>,
821    b: CheckedTimestamp<chrono::DateTime<Utc>>,
822) -> Result<Interval, EvalError> {
823    Interval::from_chrono_duration(a - b)
824        .map_err(|e| EvalError::IntervalOutOfRange(e.to_string().into()))
825}
826
827#[sqlfunc(
828    is_monotone = "(true, true)",
829    is_infix_op = true,
830    sqlname = "-",
831    propagates_nulls = true
832)]
833fn sub_date(a: Date, b: Date) -> i32 {
834    a - b
835}
836
837#[sqlfunc(is_monotone = "(true, true)", is_infix_op = true, sqlname = "-")]
838fn sub_time(a: chrono::NaiveTime, b: chrono::NaiveTime) -> Result<Interval, EvalError> {
839    Interval::from_chrono_duration(a - b)
840        .map_err(|e| EvalError::IntervalOutOfRange(e.to_string().into()))
841}
842
843#[sqlfunc(
844    is_monotone = "(true, true)",
845    output_type = "Interval",
846    is_infix_op = true,
847    sqlname = "-",
848    propagates_nulls = true
849)]
850fn sub_interval(a: Interval, b: Interval) -> Result<Interval, EvalError> {
851    b.checked_neg()
852        .and_then(|b| b.checked_add(&a))
853        .ok_or_else(|| EvalError::IntervalOutOfRange(format!("{a} - {b}").into()))
854}
855
856#[sqlfunc(
857    is_monotone = "(true, true)",
858    is_infix_op = true,
859    sqlname = "-",
860    propagates_nulls = true
861)]
862fn sub_date_interval(
863    date: Date,
864    interval: Interval,
865) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
866    let dt = NaiveDate::from(date).and_hms_opt(0, 0, 0).unwrap();
867    let dt = interval
868        .months
869        .checked_neg()
870        .ok_or_else(|| EvalError::IntervalOutOfRange(interval.months.to_string().into()))
871        .and_then(|months| add_timestamp_months(&dt, months))?;
872    let dt = dt
873        .checked_sub_signed(interval.duration_as_chrono())
874        .ok_or(EvalError::TimestampOutOfRange)?;
875    Ok(dt.try_into()?)
876}
877
878#[sqlfunc(
879    is_monotone = "(false, false)",
880    is_infix_op = true,
881    sqlname = "-",
882    propagates_nulls = true
883)]
884fn sub_time_interval(time: chrono::NaiveTime, interval: Interval) -> chrono::NaiveTime {
885    let (t, _) = time.overflowing_sub_signed(interval.duration_as_chrono());
886    t
887}
888
889#[sqlfunc(
890    is_monotone = "(true, true)",
891    is_infix_op = true,
892    sqlname = "*",
893    propagates_nulls = true
894)]
895fn mul_int16(a: i16, b: i16) -> Result<i16, EvalError> {
896    a.checked_mul(b).ok_or(EvalError::NumericFieldOverflow)
897}
898
899#[sqlfunc(
900    is_monotone = "(true, true)",
901    is_infix_op = true,
902    sqlname = "*",
903    propagates_nulls = true
904)]
905fn mul_int32(a: i32, b: i32) -> Result<i32, EvalError> {
906    a.checked_mul(b).ok_or(EvalError::NumericFieldOverflow)
907}
908
909#[sqlfunc(
910    is_monotone = "(true, true)",
911    is_infix_op = true,
912    sqlname = "*",
913    propagates_nulls = true
914)]
915fn mul_int64(a: i64, b: i64) -> Result<i64, EvalError> {
916    a.checked_mul(b).ok_or(EvalError::NumericFieldOverflow)
917}
918
919#[sqlfunc(
920    is_monotone = "(true, true)",
921    is_infix_op = true,
922    sqlname = "*",
923    propagates_nulls = true
924)]
925fn mul_uint16(a: u16, b: u16) -> Result<u16, EvalError> {
926    a.checked_mul(b)
927        .ok_or_else(|| EvalError::UInt16OutOfRange(format!("{a} * {b}").into()))
928}
929
930#[sqlfunc(
931    is_monotone = "(true, true)",
932    is_infix_op = true,
933    sqlname = "*",
934    propagates_nulls = true
935)]
936fn mul_uint32(a: u32, b: u32) -> Result<u32, EvalError> {
937    a.checked_mul(b)
938        .ok_or_else(|| EvalError::UInt32OutOfRange(format!("{a} * {b}").into()))
939}
940
941#[sqlfunc(
942    is_monotone = "(true, true)",
943    is_infix_op = true,
944    sqlname = "*",
945    propagates_nulls = true
946)]
947fn mul_uint64(a: u64, b: u64) -> Result<u64, EvalError> {
948    a.checked_mul(b)
949        .ok_or_else(|| EvalError::UInt64OutOfRange(format!("{a} * {b}").into()))
950}
951
952#[sqlfunc(
953    is_monotone = (true, true),
954    is_infix_op = true,
955    sqlname = "*",
956    propagates_nulls = true
957)]
958fn mul_float32(a: f32, b: f32) -> Result<f32, EvalError> {
959    let product = a * b;
960    if product.is_infinite() && !a.is_infinite() && !b.is_infinite() {
961        Err(EvalError::FloatOverflow)
962    } else if product == 0.0f32 && a != 0.0f32 && b != 0.0f32 {
963        Err(EvalError::FloatUnderflow)
964    } else {
965        Ok(product)
966    }
967}
968
969#[sqlfunc(
970    is_monotone = "(true, true)",
971    is_infix_op = true,
972    sqlname = "*",
973    propagates_nulls = true
974)]
975fn mul_float64(a: f64, b: f64) -> Result<f64, EvalError> {
976    let product = a * b;
977    if product.is_infinite() && !a.is_infinite() && !b.is_infinite() {
978        Err(EvalError::FloatOverflow)
979    } else if product == 0.0f64 && a != 0.0f64 && b != 0.0f64 {
980        Err(EvalError::FloatUnderflow)
981    } else {
982        Ok(product)
983    }
984}
985
986#[sqlfunc(
987    is_monotone = "(true, true)",
988    is_infix_op = true,
989    sqlname = "*",
990    propagates_nulls = true
991)]
992fn mul_numeric(mut a: Numeric, b: Numeric) -> Result<Numeric, EvalError> {
993    let mut cx = numeric::cx_datum();
994    cx.mul(&mut a, &b);
995    let cx_status = cx.status();
996    if cx_status.overflow() {
997        Err(EvalError::FloatOverflow)
998    } else if cx_status.subnormal() {
999        Err(EvalError::FloatUnderflow)
1000    } else {
1001        numeric::munge_numeric(&mut a).unwrap();
1002        Ok(a)
1003    }
1004}
1005
1006#[sqlfunc(
1007    is_monotone = "(false, false)",
1008    is_infix_op = true,
1009    sqlname = "*",
1010    propagates_nulls = true
1011)]
1012fn mul_interval(a: Interval, b: f64) -> Result<Interval, EvalError> {
1013    a.checked_mul(b)
1014        .ok_or_else(|| EvalError::IntervalOutOfRange(format!("{a} * {b}").into()))
1015}
1016
1017#[sqlfunc(
1018    is_monotone = "(true, false)",
1019    is_infix_op = true,
1020    sqlname = "/",
1021    propagates_nulls = true
1022)]
1023fn div_int16(a: i16, b: i16) -> Result<i16, EvalError> {
1024    if b == 0 {
1025        Err(EvalError::DivisionByZero)
1026    } else {
1027        a.checked_div(b)
1028            .ok_or_else(|| EvalError::Int16OutOfRange(format!("{a} / {b}").into()))
1029    }
1030}
1031
1032#[sqlfunc(
1033    is_monotone = "(true, false)",
1034    is_infix_op = true,
1035    sqlname = "/",
1036    propagates_nulls = true
1037)]
1038fn div_int32(a: i32, b: i32) -> Result<i32, EvalError> {
1039    if b == 0 {
1040        Err(EvalError::DivisionByZero)
1041    } else {
1042        a.checked_div(b)
1043            .ok_or_else(|| EvalError::Int32OutOfRange(format!("{a} / {b}").into()))
1044    }
1045}
1046
1047#[sqlfunc(
1048    is_monotone = "(true, false)",
1049    is_infix_op = true,
1050    sqlname = "/",
1051    propagates_nulls = true
1052)]
1053fn div_int64(a: i64, b: i64) -> Result<i64, EvalError> {
1054    if b == 0 {
1055        Err(EvalError::DivisionByZero)
1056    } else {
1057        a.checked_div(b)
1058            .ok_or_else(|| EvalError::Int64OutOfRange(format!("{a} / {b}").into()))
1059    }
1060}
1061
1062#[sqlfunc(
1063    is_monotone = "(true, false)",
1064    is_infix_op = true,
1065    sqlname = "/",
1066    propagates_nulls = true
1067)]
1068fn div_uint16(a: u16, b: u16) -> Result<u16, EvalError> {
1069    if b == 0 {
1070        Err(EvalError::DivisionByZero)
1071    } else {
1072        Ok(a / b)
1073    }
1074}
1075
1076#[sqlfunc(
1077    is_monotone = "(true, false)",
1078    is_infix_op = true,
1079    sqlname = "/",
1080    propagates_nulls = true
1081)]
1082fn div_uint32(a: u32, b: u32) -> Result<u32, EvalError> {
1083    if b == 0 {
1084        Err(EvalError::DivisionByZero)
1085    } else {
1086        Ok(a / b)
1087    }
1088}
1089
1090#[sqlfunc(
1091    is_monotone = "(true, false)",
1092    is_infix_op = true,
1093    sqlname = "/",
1094    propagates_nulls = true
1095)]
1096fn div_uint64(a: u64, b: u64) -> Result<u64, EvalError> {
1097    if b == 0 {
1098        Err(EvalError::DivisionByZero)
1099    } else {
1100        Ok(a / b)
1101    }
1102}
1103
1104#[sqlfunc(
1105    is_monotone = "(true, false)",
1106    is_infix_op = true,
1107    sqlname = "/",
1108    propagates_nulls = true
1109)]
1110fn div_float32(a: f32, b: f32) -> Result<f32, EvalError> {
1111    if b == 0.0f32 && !a.is_nan() {
1112        Err(EvalError::DivisionByZero)
1113    } else {
1114        let quotient = a / b;
1115        if quotient.is_infinite() && !a.is_infinite() {
1116            Err(EvalError::FloatOverflow)
1117        } else if quotient == 0.0f32 && a != 0.0f32 && !b.is_infinite() {
1118            Err(EvalError::FloatUnderflow)
1119        } else {
1120            Ok(quotient)
1121        }
1122    }
1123}
1124
1125#[sqlfunc(
1126    is_monotone = "(true, false)",
1127    is_infix_op = true,
1128    sqlname = "/",
1129    propagates_nulls = true
1130)]
1131fn div_float64(a: f64, b: f64) -> Result<f64, EvalError> {
1132    if b == 0.0f64 && !a.is_nan() {
1133        Err(EvalError::DivisionByZero)
1134    } else {
1135        let quotient = a / b;
1136        if quotient.is_infinite() && !a.is_infinite() {
1137            Err(EvalError::FloatOverflow)
1138        } else if quotient == 0.0f64 && a != 0.0f64 && !b.is_infinite() {
1139            Err(EvalError::FloatUnderflow)
1140        } else {
1141            Ok(quotient)
1142        }
1143    }
1144}
1145
1146#[sqlfunc(
1147    is_monotone = "(true, false)",
1148    is_infix_op = true,
1149    sqlname = "/",
1150    propagates_nulls = true
1151)]
1152fn div_numeric(mut a: Numeric, b: Numeric) -> Result<Numeric, EvalError> {
1153    let mut cx = numeric::cx_datum();
1154
1155    cx.div(&mut a, &b);
1156    let cx_status = cx.status();
1157
1158    // checking the status for division by zero errors is insufficient because
1159    // the underlying library treats 0/0 as undefined and not division by zero.
1160    if b.is_zero() {
1161        Err(EvalError::DivisionByZero)
1162    } else if cx_status.overflow() {
1163        Err(EvalError::FloatOverflow)
1164    } else if cx_status.subnormal() {
1165        Err(EvalError::FloatUnderflow)
1166    } else {
1167        numeric::munge_numeric(&mut a).unwrap();
1168        Ok(a)
1169    }
1170}
1171
1172#[sqlfunc(
1173    is_monotone = "(false, false)",
1174    is_infix_op = true,
1175    sqlname = "/",
1176    propagates_nulls = true
1177)]
1178fn div_interval(a: Interval, b: f64) -> Result<Interval, EvalError> {
1179    if b == 0.0 {
1180        Err(EvalError::DivisionByZero)
1181    } else {
1182        a.checked_div(b)
1183            .ok_or_else(|| EvalError::IntervalOutOfRange(format!("{a} / {b}").into()))
1184    }
1185}
1186
1187#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1188fn mod_int16(a: i16, b: i16) -> Result<i16, EvalError> {
1189    if b == 0 {
1190        Err(EvalError::DivisionByZero)
1191    } else {
1192        Ok(a.checked_rem(b).unwrap_or(0))
1193    }
1194}
1195
1196#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1197fn mod_int32(a: i32, b: i32) -> Result<i32, EvalError> {
1198    if b == 0 {
1199        Err(EvalError::DivisionByZero)
1200    } else {
1201        Ok(a.checked_rem(b).unwrap_or(0))
1202    }
1203}
1204
1205#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1206fn mod_int64(a: i64, b: i64) -> Result<i64, EvalError> {
1207    if b == 0 {
1208        Err(EvalError::DivisionByZero)
1209    } else {
1210        Ok(a.checked_rem(b).unwrap_or(0))
1211    }
1212}
1213
1214#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1215fn mod_uint16(a: u16, b: u16) -> Result<u16, EvalError> {
1216    if b == 0 {
1217        Err(EvalError::DivisionByZero)
1218    } else {
1219        Ok(a % b)
1220    }
1221}
1222
1223#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1224fn mod_uint32(a: u32, b: u32) -> Result<u32, EvalError> {
1225    if b == 0 {
1226        Err(EvalError::DivisionByZero)
1227    } else {
1228        Ok(a % b)
1229    }
1230}
1231
1232#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1233fn mod_uint64(a: u64, b: u64) -> Result<u64, EvalError> {
1234    if b == 0 {
1235        Err(EvalError::DivisionByZero)
1236    } else {
1237        Ok(a % b)
1238    }
1239}
1240
1241#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1242fn mod_float32(a: f32, b: f32) -> Result<f32, EvalError> {
1243    if b == 0.0 {
1244        Err(EvalError::DivisionByZero)
1245    } else {
1246        Ok(a % b)
1247    }
1248}
1249
1250#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1251fn mod_float64(a: f64, b: f64) -> Result<f64, EvalError> {
1252    if b == 0.0 {
1253        Err(EvalError::DivisionByZero)
1254    } else {
1255        Ok(a % b)
1256    }
1257}
1258
1259#[sqlfunc(is_infix_op = true, sqlname = "%", propagates_nulls = true)]
1260fn mod_numeric(mut a: Numeric, b: Numeric) -> Result<Numeric, EvalError> {
1261    if b.is_zero() {
1262        return Err(EvalError::DivisionByZero);
1263    }
1264    let mut cx = numeric::cx_datum();
1265    // Postgres does _not_ use IEEE 754-style remainder
1266    cx.rem(&mut a, &b);
1267    numeric::munge_numeric(&mut a).unwrap();
1268    Ok(a)
1269}
1270
1271fn neg_interval_inner(a: Interval) -> Result<Interval, EvalError> {
1272    a.checked_neg()
1273        .ok_or_else(|| EvalError::IntervalOutOfRange(a.to_string().into()))
1274}
1275
1276fn log_guard_numeric(val: &Numeric, function_name: &str) -> Result<(), EvalError> {
1277    if val.is_negative() {
1278        return Err(EvalError::NegativeOutOfDomain(function_name.into()));
1279    }
1280    if val.is_zero() {
1281        return Err(EvalError::ZeroOutOfDomain(function_name.into()));
1282    }
1283    Ok(())
1284}
1285
1286#[sqlfunc(sqlname = "log", propagates_nulls = true)]
1287fn log_base_numeric(mut a: Numeric, mut b: Numeric) -> Result<Numeric, EvalError> {
1288    log_guard_numeric(&a, "log")?;
1289    log_guard_numeric(&b, "log")?;
1290    let mut cx = numeric::cx_datum();
1291    cx.ln(&mut a);
1292    cx.ln(&mut b);
1293    cx.div(&mut b, &a);
1294    if a.is_zero() {
1295        Err(EvalError::DivisionByZero)
1296    } else {
1297        // This division can result in slightly wrong answers due to the
1298        // limitation of dividing irrational numbers. To correct that, see if
1299        // rounding off the value from its `numeric::NUMERIC_DATUM_MAX_PRECISION
1300        // - 1`th position results in an integral value.
1301        cx.set_precision(usize::from(numeric::NUMERIC_DATUM_MAX_PRECISION - 1))
1302            .expect("reducing precision below max always succeeds");
1303        let mut integral_check = b.clone();
1304
1305        // `reduce` rounds to the context's final digit when the number of
1306        // digits in its argument exceeds its precision. We've contrived that to
1307        // happen by shrinking the context's precision by 1.
1308        cx.reduce(&mut integral_check);
1309
1310        // Reduced integral values always have a non-negative exponent.
1311        let mut b = if integral_check.exponent() >= 0 {
1312            // We believe our result should have been an integral
1313            integral_check
1314        } else {
1315            b
1316        };
1317
1318        numeric::munge_numeric(&mut b).unwrap();
1319        Ok(b)
1320    }
1321}
1322
1323#[sqlfunc(propagates_nulls = true)]
1324fn power(a: f64, b: f64) -> Result<f64, EvalError> {
1325    if a == 0.0 && b.is_sign_negative() {
1326        return Err(EvalError::Undefined(
1327            "zero raised to a negative power".into(),
1328        ));
1329    }
1330    if a.is_sign_negative() && b.fract() != 0.0 {
1331        // Equivalent to PG error:
1332        // > a negative number raised to a non-integer power yields a complex result
1333        return Err(EvalError::ComplexOutOfRange("pow".into()));
1334    }
1335    let res = a.powf(b);
1336    if res.is_infinite() {
1337        return Err(EvalError::FloatOverflow);
1338    }
1339    if res == 0.0 && a != 0.0 {
1340        return Err(EvalError::FloatUnderflow);
1341    }
1342    Ok(res)
1343}
1344
1345#[sqlfunc(propagates_nulls = true)]
1346fn uuid_generate_v5(a: uuid::Uuid, b: &str) -> uuid::Uuid {
1347    uuid::Uuid::new_v5(&a, b.as_bytes())
1348}
1349
1350#[sqlfunc(output_type = "Numeric", propagates_nulls = true)]
1351fn power_numeric(mut a: Numeric, b: Numeric) -> Result<Numeric, EvalError> {
1352    if a.is_zero() {
1353        if b.is_zero() {
1354            return Ok(Numeric::from(1));
1355        }
1356        if b.is_negative() {
1357            return Err(EvalError::Undefined(
1358                "zero raised to a negative power".into(),
1359            ));
1360        }
1361    }
1362    if a.is_negative() && b.exponent() < 0 {
1363        // Equivalent to PG error:
1364        // > a negative number raised to a non-integer power yields a complex result
1365        return Err(EvalError::ComplexOutOfRange("pow".into()));
1366    }
1367    let mut cx = numeric::cx_datum();
1368    cx.pow(&mut a, &b);
1369    let cx_status = cx.status();
1370    if cx_status.overflow() || (cx_status.invalid_operation() && !b.is_negative()) {
1371        Err(EvalError::FloatOverflow)
1372    } else if cx_status.subnormal() || cx_status.invalid_operation() {
1373        Err(EvalError::FloatUnderflow)
1374    } else {
1375        numeric::munge_numeric(&mut a).unwrap();
1376        Ok(a)
1377    }
1378}
1379
1380#[sqlfunc(propagates_nulls = true)]
1381fn get_bit(bytes: &[u8], index: i32) -> Result<i32, EvalError> {
1382    let err = EvalError::IndexOutOfRange {
1383        provided: index,
1384        valid_end: i32::try_from(bytes.len().saturating_mul(8)).unwrap() - 1,
1385    };
1386
1387    let index = usize::try_from(index).map_err(|_| err.clone())?;
1388
1389    let byte_index = index / 8;
1390    let bit_index = index % 8;
1391
1392    let i = bytes
1393        .get(byte_index)
1394        .map(|b| (*b >> bit_index) & 1)
1395        .ok_or(err)?;
1396    assert!(i == 0 || i == 1);
1397    Ok(i32::from(i))
1398}
1399
1400#[sqlfunc(propagates_nulls = true)]
1401fn get_byte(bytes: &[u8], index: i32) -> Result<i32, EvalError> {
1402    let err = EvalError::IndexOutOfRange {
1403        provided: index,
1404        valid_end: i32::try_from(bytes.len()).unwrap() - 1,
1405    };
1406    let i: &u8 = bytes
1407        .get(usize::try_from(index).map_err(|_| err.clone())?)
1408        .ok_or(err)?;
1409    Ok(i32::from(*i))
1410}
1411
1412#[sqlfunc(sqlname = "constant_time_compare_bytes", propagates_nulls = true)]
1413pub fn constant_time_eq_bytes(a: &[u8], b: &[u8]) -> bool {
1414    bool::from(a.ct_eq(b))
1415}
1416
1417#[sqlfunc(sqlname = "constant_time_compare_strings", propagates_nulls = true)]
1418pub fn constant_time_eq_string(a: &str, b: &str) -> bool {
1419    bool::from(a.as_bytes().ct_eq(b.as_bytes()))
1420}
1421
1422#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1423fn range_contains_i32<'a>(a: Range<Datum<'a>>, b: i32) -> bool {
1424    a.contains_elem(&b)
1425}
1426
1427#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1428fn range_contains_i64<'a>(a: Range<Datum<'a>>, elem: i64) -> bool {
1429    a.contains_elem(&elem)
1430}
1431
1432#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1433fn range_contains_date<'a>(a: Range<Datum<'a>>, elem: Date) -> bool {
1434    a.contains_elem(&elem)
1435}
1436
1437#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1438fn range_contains_numeric<'a>(a: Range<Datum<'a>>, elem: OrderedDecimal<Numeric>) -> bool {
1439    a.contains_elem(&elem)
1440}
1441
1442#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1443fn range_contains_timestamp<'a>(
1444    a: Range<Datum<'a>>,
1445    elem: CheckedTimestamp<NaiveDateTime>,
1446) -> bool {
1447    a.contains_elem(&elem)
1448}
1449
1450#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1451fn range_contains_timestamp_tz<'a>(
1452    a: Range<Datum<'a>>,
1453    elem: CheckedTimestamp<DateTime<Utc>>,
1454) -> bool {
1455    a.contains_elem(&elem)
1456}
1457
1458#[sqlfunc(is_infix_op = true, sqlname = "<@", propagates_nulls = true)]
1459fn range_contains_i32_rev<'a>(a: Range<Datum<'a>>, b: i32) -> bool {
1460    a.contains_elem(&b)
1461}
1462
1463#[sqlfunc(is_infix_op = true, sqlname = "<@", propagates_nulls = true)]
1464fn range_contains_i64_rev<'a>(a: Range<Datum<'a>>, elem: i64) -> bool {
1465    a.contains_elem(&elem)
1466}
1467
1468#[sqlfunc(is_infix_op = true, sqlname = "<@", propagates_nulls = true)]
1469fn range_contains_date_rev<'a>(a: Range<Datum<'a>>, elem: Date) -> bool {
1470    a.contains_elem(&elem)
1471}
1472
1473#[sqlfunc(is_infix_op = true, sqlname = "<@", propagates_nulls = true)]
1474fn range_contains_numeric_rev<'a>(a: Range<Datum<'a>>, elem: OrderedDecimal<Numeric>) -> bool {
1475    a.contains_elem(&elem)
1476}
1477
1478#[sqlfunc(is_infix_op = true, sqlname = "<@", propagates_nulls = true)]
1479fn range_contains_timestamp_rev<'a>(
1480    a: Range<Datum<'a>>,
1481    elem: CheckedTimestamp<NaiveDateTime>,
1482) -> bool {
1483    a.contains_elem(&elem)
1484}
1485
1486#[sqlfunc(is_infix_op = true, sqlname = "<@", propagates_nulls = true)]
1487fn range_contains_timestamp_tz_rev<'a>(
1488    a: Range<Datum<'a>>,
1489    elem: CheckedTimestamp<DateTime<Utc>>,
1490) -> bool {
1491    a.contains_elem(&elem)
1492}
1493
1494/// Macro to define binary function for various range operations.
1495/// Parameters:
1496/// 1. Unique binary function symbol.
1497/// 2. Range function symbol.
1498/// 3. SQL name for the function.
1499macro_rules! range_fn {
1500    ($fn:expr, $range_fn:expr, $sqlname:expr) => {
1501        paste::paste! {
1502
1503            #[sqlfunc(
1504                output_type = "bool",
1505                is_infix_op = true,
1506                sqlname = $sqlname,
1507                propagates_nulls = true
1508            )]
1509            fn [< range_ $fn >]<'a>(a: Datum<'a>, b: Datum<'a>) -> Datum<'a>
1510            {
1511                if a.is_null() || b.is_null() { return Datum::Null }
1512                let l = a.unwrap_range();
1513                let r = b.unwrap_range();
1514                Datum::from(Range::<Datum<'a>>::$range_fn(&l, &r))
1515            }
1516        }
1517    };
1518}
1519
1520// RangeContainsRange is either @> or <@ depending on the order of the arguments.
1521// It doesn't influence the result, but it does influence the display string.
1522range_fn!(contains_range, contains_range, "@>");
1523range_fn!(contains_range_rev, contains_range, "<@");
1524range_fn!(overlaps, overlaps, "&&");
1525range_fn!(after, after, ">>");
1526range_fn!(before, before, "<<");
1527range_fn!(overleft, overleft, "&<");
1528range_fn!(overright, overright, "&>");
1529range_fn!(adjacent, adjacent, "-|-");
1530
1531#[sqlfunc(is_infix_op = true, sqlname = "+")]
1532fn range_union<T: Copy + Ord>(l: Range<T>, r: Range<T>) -> Result<Range<T>, EvalError> {
1533    Ok(l.union(&r)?)
1534}
1535
1536#[sqlfunc(is_infix_op = true, sqlname = "*")]
1537fn range_intersection<T: Copy + Ord>(l: Range<T>, r: Range<T>) -> Range<T> {
1538    l.intersection(&r)
1539}
1540
1541#[sqlfunc(
1542    output_type_expr = "input_types[0].scalar_type.without_modifiers().nullable(true)",
1543    is_infix_op = true,
1544    sqlname = "-",
1545    propagates_nulls = true,
1546    introduces_nulls = false
1547)]
1548fn range_difference<'a>(
1549    l: Range<Datum<'a>>,
1550    r: Range<Datum<'a>>,
1551) -> Result<Range<Datum<'a>>, EvalError> {
1552    Ok(l.difference(&r)?)
1553}
1554
1555#[sqlfunc(is_infix_op = true, sqlname = "=", negate = "Some(NotEq.into())")]
1556fn eq<'a>(a: ExcludeNull<Datum<'a>>, b: ExcludeNull<Datum<'a>>) -> bool {
1557    // SQL equality demands that if either input is null, then the result should be null. However,
1558    // we don't need to handle this case here; it is handled when `BinaryFunc::eval` checks
1559    // `propagates_nulls`.
1560    a == b
1561}
1562
1563#[sqlfunc(is_infix_op = true, sqlname = "!=", negate = "Some(Eq.into())")]
1564fn not_eq<'a>(a: ExcludeNull<Datum<'a>>, b: ExcludeNull<Datum<'a>>) -> bool {
1565    a != b
1566}
1567
1568#[sqlfunc(
1569    is_monotone = "(true, true)",
1570    is_infix_op = true,
1571    sqlname = "<",
1572    negate = "Some(Gte.into())"
1573)]
1574fn lt<'a>(a: ExcludeNull<Datum<'a>>, b: ExcludeNull<Datum<'a>>) -> bool {
1575    a < b
1576}
1577
1578#[sqlfunc(
1579    is_monotone = "(true, true)",
1580    is_infix_op = true,
1581    sqlname = "<=",
1582    negate = "Some(Gt.into())"
1583)]
1584fn lte<'a>(a: ExcludeNull<Datum<'a>>, b: ExcludeNull<Datum<'a>>) -> bool {
1585    a <= b
1586}
1587
1588#[sqlfunc(
1589    is_monotone = "(true, true)",
1590    is_infix_op = true,
1591    sqlname = ">",
1592    negate = "Some(Lte.into())"
1593)]
1594fn gt<'a>(a: ExcludeNull<Datum<'a>>, b: ExcludeNull<Datum<'a>>) -> bool {
1595    a > b
1596}
1597
1598#[sqlfunc(
1599    is_monotone = "(true, true)",
1600    is_infix_op = true,
1601    sqlname = ">=",
1602    negate = "Some(Lt.into())"
1603)]
1604fn gte<'a>(a: ExcludeNull<Datum<'a>>, b: ExcludeNull<Datum<'a>>) -> bool {
1605    a >= b
1606}
1607
1608#[sqlfunc(sqlname = "tocharts", propagates_nulls = true)]
1609fn to_char_timestamp_format(ts: CheckedTimestamp<chrono::NaiveDateTime>, format: &str) -> String {
1610    let fmt = DateTimeFormat::compile(format);
1611    fmt.render(&*ts)
1612}
1613
1614#[sqlfunc(sqlname = "tochartstz", propagates_nulls = true)]
1615fn to_char_timestamp_tz_format(
1616    ts: CheckedTimestamp<chrono::DateTime<Utc>>,
1617    format: &str,
1618) -> String {
1619    let fmt = DateTimeFormat::compile(format);
1620    fmt.render(&*ts)
1621}
1622
1623#[sqlfunc(sqlname = "->", is_infix_op = true)]
1624fn jsonb_get_int64<'a>(a: JsonbRef<'a>, i: i64) -> Option<JsonbRef<'a>> {
1625    match a.into_datum() {
1626        Datum::List(list) => {
1627            let i = if i >= 0 {
1628                usize::cast_from(i.unsigned_abs())
1629            } else {
1630                // index backwards from the end
1631                let i = usize::cast_from(i.unsigned_abs());
1632                (list.iter().count()).wrapping_sub(i)
1633            };
1634            let v = list.iter().nth(i)?;
1635            // `v` should be valid jsonb because it came from a jsonb list, but we don't
1636            // panic on mismatch to avoid bringing down the whole system on corrupt data.
1637            // Instead, we'll return None.
1638            JsonbRef::try_from_result(Ok::<_, ()>(v)).ok()
1639        }
1640        Datum::Map(_) => None,
1641        _ => {
1642            // I have no idea why postgres does this, but we're stuck with it
1643            (i == 0 || i == -1).then_some(a)
1644        }
1645    }
1646}
1647
1648#[sqlfunc(sqlname = "->>", is_infix_op = true)]
1649fn jsonb_get_int64_stringify<'a>(
1650    a: JsonbRef<'a>,
1651    i: i64,
1652    temp_storage: &'a RowArena,
1653) -> Option<&'a str> {
1654    let json = jsonb_get_int64(a, i)?;
1655    jsonb_stringify(json.into_datum(), temp_storage)
1656}
1657
1658#[sqlfunc(sqlname = "->", is_infix_op = true)]
1659fn jsonb_get_string<'a>(a: JsonbRef<'a>, k: &str) -> Option<JsonbRef<'a>> {
1660    let dict = DatumMap::try_from_result(Ok::<_, ()>(a.into_datum())).ok()?;
1661    let v = dict.iter().find(|(k2, _v)| k == *k2).map(|(_k, v)| v)?;
1662    JsonbRef::try_from_result(Ok::<_, ()>(v)).ok()
1663}
1664
1665#[sqlfunc(sqlname = "->>", is_infix_op = true)]
1666fn jsonb_get_string_stringify<'a>(
1667    a: JsonbRef<'a>,
1668    k: &str,
1669    temp_storage: &'a RowArena,
1670) -> Option<&'a str> {
1671    let v = jsonb_get_string(a, k)?;
1672    jsonb_stringify(v.into_datum(), temp_storage)
1673}
1674
1675#[sqlfunc(sqlname = "#>", is_infix_op = true)]
1676fn jsonb_get_path<'a>(mut json: JsonbRef<'a>, b: Array<'a>) -> Option<JsonbRef<'a>> {
1677    let path = b.elements();
1678    for key in path.iter() {
1679        let key = match key {
1680            Datum::String(s) => s,
1681            Datum::Null => return None,
1682            _ => unreachable!("keys in jsonb_get_path known to be strings"),
1683        };
1684        let v = match json.into_datum() {
1685            Datum::Map(map) => map.iter().find(|(k, _)| key == *k).map(|(_k, v)| v),
1686            Datum::List(list) => {
1687                let i = strconv::parse_int64(key).ok()?;
1688                let i = if i >= 0 {
1689                    usize::cast_from(i.unsigned_abs())
1690                } else {
1691                    // index backwards from the end
1692                    let i = usize::cast_from(i.unsigned_abs());
1693                    (list.iter().count()).wrapping_sub(i)
1694                };
1695                list.iter().nth(i)
1696            }
1697            _ => return None,
1698        }?;
1699        json = JsonbRef::try_from_result(Ok::<_, ()>(v)).ok()?;
1700    }
1701    Some(json)
1702}
1703
1704#[sqlfunc(sqlname = "#>>", is_infix_op = true)]
1705fn jsonb_get_path_stringify<'a>(
1706    a: JsonbRef<'a>,
1707    b: Array<'a>,
1708    temp_storage: &'a RowArena,
1709) -> Option<&'a str> {
1710    let json = jsonb_get_path(a, b)?;
1711    jsonb_stringify(json.into_datum(), temp_storage)
1712}
1713
1714#[sqlfunc(is_infix_op = true, sqlname = "?")]
1715fn jsonb_contains_string<'a>(a: JsonbRef<'a>, k: &str) -> bool {
1716    // https://www.postgresql.org/docs/current/datatype-json.html#JSON-CONTAINMENT
1717    // When the left operand is SQL NULL (NULL::jsonb), JsonbRef::try_from_result rejects it,
1718    // so the binary evaluator never calls this function and returns NULL (see binary.rs).
1719    // So, this function only runs for non-null jsonb; a.into_datum() never sees Datum::Null.
1720    match a.into_datum() {
1721        Datum::List(list) => list.iter().any(|k2| Datum::from(k) == k2),
1722        Datum::Map(dict) => dict.iter().any(|(k2, _v)| k == k2),
1723        Datum::String(string) => string == k,
1724        _ => false,
1725    }
1726}
1727
1728#[sqlfunc(is_infix_op = true, sqlname = "?", propagates_nulls = true)]
1729// Map keys are always text.
1730fn map_contains_key<'a>(map: DatumMap<'a>, k: &str) -> bool {
1731    map.iter().any(|(k2, _v)| k == k2)
1732}
1733
1734#[sqlfunc(is_infix_op = true, sqlname = "?&")]
1735fn map_contains_all_keys<'a>(map: DatumMap<'a>, keys: Array<'a>) -> bool {
1736    keys.elements()
1737        .iter()
1738        .all(|key| !key.is_null() && map.iter().any(|(k, _v)| k == key.unwrap_str()))
1739}
1740
1741#[sqlfunc(is_infix_op = true, sqlname = "?|", propagates_nulls = true)]
1742fn map_contains_any_keys<'a>(map: DatumMap<'a>, keys: Array<'a>) -> bool {
1743    keys.elements()
1744        .iter()
1745        .any(|key| !key.is_null() && map.iter().any(|(k, _v)| k == key.unwrap_str()))
1746}
1747
1748#[sqlfunc(is_infix_op = true, sqlname = "@>", propagates_nulls = true)]
1749fn map_contains_map<'a>(map_a: DatumMap<'a>, b: DatumMap<'a>) -> bool {
1750    b.iter().all(|(b_key, b_val)| {
1751        map_a
1752            .iter()
1753            .any(|(a_key, a_val)| (a_key == b_key) && (a_val == b_val))
1754    })
1755}
1756
1757#[sqlfunc(is_infix_op = true, sqlname = "->", propagates_nulls = true)]
1758fn map_get_value<'a, T: FromDatum<'a>>(a: DatumMap<'a, T>, target_key: &str) -> Option<T> {
1759    a.typed_iter()
1760        .find(|(key, _v)| target_key == *key)
1761        .map(|(_k, v)| v)
1762}
1763
1764#[sqlfunc(is_infix_op = true, sqlname = "@>")]
1765fn list_contains_list<'a>(a: ExcludeNull<DatumList<'a>>, b: ExcludeNull<DatumList<'a>>) -> bool {
1766    // NULL is never equal to NULL. If NULL is an element of b, b cannot be contained in a, even if a contains NULL.
1767    if b.iter().contains(&Datum::Null) {
1768        false
1769    } else {
1770        b.iter()
1771            .all(|item_b| a.iter().any(|item_a| item_a == item_b))
1772    }
1773}
1774
1775#[sqlfunc(is_infix_op = true, sqlname = "<@")]
1776fn list_contains_list_rev<'a>(
1777    a: ExcludeNull<DatumList<'a>>,
1778    b: ExcludeNull<DatumList<'a>>,
1779) -> bool {
1780    list_contains_list(b, a)
1781}
1782
1783// TODO(jamii) nested loops are possibly not the fastest way to do this
1784#[sqlfunc(is_infix_op = true, sqlname = "@>")]
1785fn jsonb_contains_jsonb<'a>(a: JsonbRef<'a>, b: JsonbRef<'a>) -> bool {
1786    // https://www.postgresql.org/docs/current/datatype-json.html#JSON-CONTAINMENT
1787    fn contains(a: Datum, b: Datum, at_top_level: bool) -> bool {
1788        match (a, b) {
1789            (Datum::JsonNull, Datum::JsonNull) => true,
1790            (Datum::False, Datum::False) => true,
1791            (Datum::True, Datum::True) => true,
1792            (Datum::Numeric(a), Datum::Numeric(b)) => a == b,
1793            (Datum::String(a), Datum::String(b)) => a == b,
1794            (Datum::List(a), Datum::List(b)) => b
1795                .iter()
1796                .all(|b_elem| a.iter().any(|a_elem| contains(a_elem, b_elem, false))),
1797            (Datum::Map(a), Datum::Map(b)) => b.iter().all(|(b_key, b_val)| {
1798                a.iter()
1799                    .any(|(a_key, a_val)| (a_key == b_key) && contains(a_val, b_val, false))
1800            }),
1801
1802            // fun special case
1803            (Datum::List(a), b) => {
1804                at_top_level && a.iter().any(|a_elem| contains(a_elem, b, false))
1805            }
1806
1807            _ => false,
1808        }
1809    }
1810    contains(a.into_datum(), b.into_datum(), true)
1811}
1812
1813#[sqlfunc(is_infix_op = true, sqlname = "||")]
1814fn jsonb_concat<'a>(
1815    a: JsonbRef<'a>,
1816    b: JsonbRef<'a>,
1817    temp_storage: &'a RowArena,
1818) -> Option<JsonbRef<'a>> {
1819    let res = match (a.into_datum(), b.into_datum()) {
1820        (Datum::Map(dict_a), Datum::Map(dict_b)) => {
1821            let mut pairs = dict_b.iter().chain(dict_a.iter()).collect::<Vec<_>>();
1822            // stable sort, so if keys collide dedup prefers dict_b
1823            pairs.sort_by(|(k1, _v1), (k2, _v2)| k1.cmp(k2));
1824            pairs.dedup_by(|(k1, _v1), (k2, _v2)| k1 == k2);
1825            temp_storage.make_datum(|packer| packer.push_dict(pairs))
1826        }
1827        (Datum::List(list_a), Datum::List(list_b)) => {
1828            let elems = list_a.iter().chain(list_b.iter());
1829            temp_storage.make_datum(|packer| packer.push_list(elems))
1830        }
1831        (Datum::List(list_a), b) => {
1832            let elems = list_a.iter().chain(Some(b));
1833            temp_storage.make_datum(|packer| packer.push_list(elems))
1834        }
1835        (a, Datum::List(list_b)) => {
1836            let elems = Some(a).into_iter().chain(list_b.iter());
1837            temp_storage.make_datum(|packer| packer.push_list(elems))
1838        }
1839        _ => return None,
1840    };
1841    Some(JsonbRef::from_datum(res))
1842}
1843
1844#[sqlfunc(
1845    output_type_expr = "SqlScalarType::Jsonb.nullable(true)",
1846    is_infix_op = true,
1847    sqlname = "-",
1848    propagates_nulls = true,
1849    introduces_nulls = true
1850)]
1851fn jsonb_delete_int64<'a>(a: Datum<'a>, i: i64, temp_storage: &'a RowArena) -> Datum<'a> {
1852    match a {
1853        Datum::List(list) => {
1854            let i = if i >= 0 {
1855                usize::cast_from(i.unsigned_abs())
1856            } else {
1857                // index backwards from the end
1858                let i = usize::cast_from(i.unsigned_abs());
1859                (list.iter().count()).wrapping_sub(i)
1860            };
1861            let elems = list
1862                .iter()
1863                .enumerate()
1864                .filter(|(i2, _e)| i != *i2)
1865                .map(|(_, e)| e);
1866            temp_storage.make_datum(|packer| packer.push_list(elems))
1867        }
1868        _ => Datum::Null,
1869    }
1870}
1871
1872#[sqlfunc(
1873    output_type_expr = "SqlScalarType::Jsonb.nullable(true)",
1874    is_infix_op = true,
1875    sqlname = "-",
1876    propagates_nulls = true,
1877    introduces_nulls = true
1878)]
1879fn jsonb_delete_string<'a>(a: Datum<'a>, k: &str, temp_storage: &'a RowArena) -> Datum<'a> {
1880    match a {
1881        Datum::List(list) => {
1882            let elems = list.iter().filter(|e| Datum::from(k) != *e);
1883            temp_storage.make_datum(|packer| packer.push_list(elems))
1884        }
1885        Datum::Map(dict) => {
1886            let pairs = dict.iter().filter(|(k2, _v)| k != *k2);
1887            temp_storage.make_datum(|packer| packer.push_dict(pairs))
1888        }
1889        _ => Datum::Null,
1890    }
1891}
1892
1893#[sqlfunc(
1894    sqlname = "extractiv",
1895    propagates_nulls = true,
1896    introduces_nulls = false
1897)]
1898fn date_part_interval_numeric(units: &str, b: Interval) -> Result<Numeric, EvalError> {
1899    match units.parse() {
1900        Ok(units) => Ok(date_part_interval_inner::<Numeric>(units, b)?),
1901        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1902    }
1903}
1904
1905#[sqlfunc(
1906    sqlname = "date_partiv",
1907    propagates_nulls = true,
1908    introduces_nulls = false
1909)]
1910fn date_part_interval_f64(units: &str, b: Interval) -> Result<f64, EvalError> {
1911    match units.parse() {
1912        Ok(units) => Ok(date_part_interval_inner::<f64>(units, b)?),
1913        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1914    }
1915}
1916
1917#[sqlfunc(
1918    sqlname = "extractt",
1919    propagates_nulls = true,
1920    introduces_nulls = false
1921)]
1922fn date_part_time_numeric(units: &str, b: chrono::NaiveTime) -> Result<Numeric, EvalError> {
1923    match units.parse() {
1924        Ok(units) => Ok(date_part_time_inner::<Numeric>(units, b)?),
1925        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1926    }
1927}
1928
1929#[sqlfunc(
1930    sqlname = "date_partt",
1931    propagates_nulls = true,
1932    introduces_nulls = false
1933)]
1934fn date_part_time_f64(units: &str, b: chrono::NaiveTime) -> Result<f64, EvalError> {
1935    match units.parse() {
1936        Ok(units) => Ok(date_part_time_inner::<f64>(units, b)?),
1937        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1938    }
1939}
1940
1941#[sqlfunc(sqlname = "extractts", propagates_nulls = true)]
1942fn date_part_timestamp_timestamp_numeric(
1943    units: &str,
1944    ts: CheckedTimestamp<NaiveDateTime>,
1945) -> Result<Numeric, EvalError> {
1946    match units.parse() {
1947        Ok(units) => Ok(date_part_timestamp_inner::<_, Numeric>(units, &*ts)?),
1948        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1949    }
1950}
1951
1952#[sqlfunc(sqlname = "extracttstz", propagates_nulls = true)]
1953fn date_part_timestamp_timestamp_tz_numeric(
1954    units: &str,
1955    ts: CheckedTimestamp<DateTime<Utc>>,
1956) -> Result<Numeric, EvalError> {
1957    match units.parse() {
1958        Ok(units) => Ok(date_part_timestamp_inner::<_, Numeric>(units, &*ts)?),
1959        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1960    }
1961}
1962
1963#[sqlfunc(sqlname = "date_partts", propagates_nulls = true)]
1964fn date_part_timestamp_timestamp_f64(
1965    units: &str,
1966    ts: CheckedTimestamp<NaiveDateTime>,
1967) -> Result<f64, EvalError> {
1968    match units.parse() {
1969        Ok(units) => date_part_timestamp_inner(units, &*ts),
1970        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1971    }
1972}
1973
1974#[sqlfunc(sqlname = "date_parttstz", propagates_nulls = true)]
1975fn date_part_timestamp_timestamp_tz_f64(
1976    units: &str,
1977    ts: CheckedTimestamp<DateTime<Utc>>,
1978) -> Result<f64, EvalError> {
1979    match units.parse() {
1980        Ok(units) => date_part_timestamp_inner(units, &*ts),
1981        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1982    }
1983}
1984
1985#[sqlfunc(sqlname = "extractd", propagates_nulls = true)]
1986fn extract_date_units(units: &str, b: Date) -> Result<Numeric, EvalError> {
1987    match units.parse() {
1988        Ok(units) => Ok(extract_date_inner(units, b.into())?),
1989        Err(_) => Err(EvalError::UnknownUnits(units.into())),
1990    }
1991}
1992
1993pub fn date_bin<T>(
1994    stride: Interval,
1995    source: CheckedTimestamp<T>,
1996    origin: CheckedTimestamp<T>,
1997) -> Result<CheckedTimestamp<T>, EvalError>
1998where
1999    T: TimestampLike,
2000{
2001    if stride.months != 0 {
2002        return Err(EvalError::DateBinOutOfRange(
2003            "timestamps cannot be binned into intervals containing months or years".into(),
2004        ));
2005    }
2006
2007    let stride_ns = match stride.duration_as_chrono().num_nanoseconds() {
2008        Some(ns) if ns <= 0 => Err(EvalError::DateBinOutOfRange(
2009            "stride must be greater than zero".into(),
2010        )),
2011        Some(ns) => Ok(ns),
2012        None => Err(EvalError::DateBinOutOfRange(
2013            format!("stride cannot exceed {}/{} nanoseconds", i64::MAX, i64::MIN,).into(),
2014        )),
2015    }?;
2016
2017    // Make sure the returned timestamp is at the start of the bin, even if the
2018    // origin is in the future. We do this here because `T` is not `Copy` and
2019    // gets moved by its subtraction operation.
2020    let sub_stride = origin > source;
2021
2022    let tm_diff = (source - origin.clone()).num_nanoseconds().ok_or_else(|| {
2023        EvalError::DateBinOutOfRange(
2024            "source and origin must not differ more than 2^63 nanoseconds".into(),
2025        )
2026    })?;
2027
2028    let mut tm_delta = tm_diff - tm_diff % stride_ns;
2029
2030    if sub_stride {
2031        tm_delta -= stride_ns;
2032    }
2033
2034    let res = origin
2035        .checked_add_signed(Duration::nanoseconds(tm_delta))
2036        .ok_or(EvalError::TimestampOutOfRange)?;
2037    Ok(CheckedTimestamp::from_timestamplike(res)?)
2038}
2039
2040#[sqlfunc(is_monotone = "(true, true)", sqlname = "bin_unix_epoch_timestamp")]
2041fn date_bin_timestamp(
2042    stride: Interval,
2043    source: CheckedTimestamp<NaiveDateTime>,
2044) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
2045    let origin =
2046        CheckedTimestamp::from_timestamplike(DateTime::from_timestamp(0, 0).unwrap().naive_utc())
2047            .expect("must fit");
2048    date_bin(stride, source, origin)
2049}
2050
2051#[sqlfunc(is_monotone = "(true, true)", sqlname = "bin_unix_epoch_timestamptz")]
2052fn date_bin_timestamp_tz(
2053    stride: Interval,
2054    source: CheckedTimestamp<DateTime<Utc>>,
2055) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
2056    let origin = CheckedTimestamp::from_timestamplike(DateTime::from_timestamp(0, 0).unwrap())
2057        .expect("must fit");
2058    date_bin(stride, source, origin)
2059}
2060
2061#[sqlfunc(sqlname = "date_truncts", propagates_nulls = true)]
2062fn date_trunc_units_timestamp(
2063    units: &str,
2064    ts: CheckedTimestamp<NaiveDateTime>,
2065) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
2066    match units.parse() {
2067        Ok(units) => Ok(date_trunc_inner(units, &*ts)?.try_into()?),
2068        Err(_) => Err(EvalError::UnknownUnits(units.into())),
2069    }
2070}
2071
2072#[sqlfunc(sqlname = "date_trunctstz", propagates_nulls = true)]
2073fn date_trunc_units_timestamp_tz(
2074    units: &str,
2075    ts: CheckedTimestamp<DateTime<Utc>>,
2076) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
2077    match units.parse() {
2078        Ok(units) => Ok(date_trunc_inner(units, &*ts)?.try_into()?),
2079        Err(_) => Err(EvalError::UnknownUnits(units.into())),
2080    }
2081}
2082
2083#[sqlfunc(sqlname = "date_trunciv", propagates_nulls = true)]
2084fn date_trunc_interval(units: &str, mut interval: Interval) -> Result<Interval, EvalError> {
2085    let dtf = units
2086        .parse()
2087        .map_err(|_| EvalError::UnknownUnits(units.into()))?;
2088
2089    interval
2090        .truncate_low_fields(dtf, Some(0), RoundBehavior::Truncate)
2091        .expect(
2092            "truncate_low_fields should not fail with max_precision 0 and RoundBehavior::Truncate",
2093        );
2094    Ok(interval)
2095}
2096
2097/// Parses a named timezone like `EST` or `America/New_York`, or a fixed-offset timezone like `-05:00`.
2098///
2099/// The interpretation of fixed offsets depend on whether the POSIX or ISO 8601 standard is being
2100/// used.
2101pub(crate) fn parse_timezone(tz: &str, spec: TimezoneSpec) -> Result<Timezone, EvalError> {
2102    Timezone::parse(tz, spec).map_err(|_| EvalError::InvalidTimezone(tz.into()))
2103}
2104
2105/// Converts the time datum `b`, which is assumed to be in UTC, to the timezone that the interval datum `a` is assumed
2106/// to represent. The interval is not allowed to hold months, but there are no limits on the amount of seconds.
2107/// The interval acts like a `chrono::FixedOffset`, without the `-86,400 < x < 86,400` limitation.
2108#[sqlfunc(sqlname = "timezoneit")]
2109fn timezone_interval_time_binary(
2110    interval: Interval,
2111    time: chrono::NaiveTime,
2112) -> Result<chrono::NaiveTime, EvalError> {
2113    if interval.months != 0 {
2114        Err(EvalError::InvalidTimezoneInterval)
2115    } else {
2116        Ok(time.overflowing_add_signed(interval.duration_as_chrono()).0)
2117    }
2118}
2119
2120/// Converts the timestamp datum `b`, which is assumed to be in the time of the timezone datum `a` to a timestamptz
2121/// in UTC. The interval is not allowed to hold months, but there are no limits on the amount of seconds.
2122/// The interval acts like a `chrono::FixedOffset`, without the `-86,400 < x < 86,400` limitation.
2123#[sqlfunc(sqlname = "timezoneits")]
2124fn timezone_interval_timestamp_binary(
2125    interval: Interval,
2126    ts: CheckedTimestamp<NaiveDateTime>,
2127) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
2128    if interval.months != 0 {
2129        Err(EvalError::InvalidTimezoneInterval)
2130    } else {
2131        match ts.checked_sub_signed(interval.duration_as_chrono()) {
2132            Some(sub) => Ok(DateTime::from_naive_utc_and_offset(sub, Utc).try_into()?),
2133            None => Err(EvalError::TimestampOutOfRange),
2134        }
2135    }
2136}
2137
2138/// Converts the UTC timestamptz datum `b`, to the local timestamp of the timezone datum `a`.
2139/// The interval is not allowed to hold months, but there are no limits on the amount of seconds.
2140/// The interval acts like a `chrono::FixedOffset`, without the `-86,400 < x < 86,400` limitation.
2141#[sqlfunc(sqlname = "timezoneitstz")]
2142fn timezone_interval_timestamp_tz_binary(
2143    interval: Interval,
2144    tstz: CheckedTimestamp<DateTime<Utc>>,
2145) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
2146    if interval.months != 0 {
2147        return Err(EvalError::InvalidTimezoneInterval);
2148    }
2149    match tstz
2150        .naive_utc()
2151        .checked_add_signed(interval.duration_as_chrono())
2152    {
2153        Some(dt) => Ok(dt.try_into()?),
2154        None => Err(EvalError::TimestampOutOfRange),
2155    }
2156}
2157
2158#[sqlfunc(
2159    output_type_expr = r#"SqlScalarType::Record {
2160                fields: [
2161                    ("abbrev".into(), SqlScalarType::String.nullable(false)),
2162                    ("base_utc_offset".into(), SqlScalarType::Interval.nullable(false)),
2163                    ("dst_offset".into(), SqlScalarType::Interval.nullable(false)),
2164                ].into(),
2165                custom_id: None,
2166            }.nullable(true)"#,
2167    propagates_nulls = true,
2168    introduces_nulls = false
2169)]
2170fn timezone_offset<'a>(
2171    tz_str: &str,
2172    b: CheckedTimestamp<chrono::DateTime<Utc>>,
2173    temp_storage: &'a RowArena,
2174) -> Result<Datum<'a>, EvalError> {
2175    let tz = match Tz::from_str_insensitive(tz_str) {
2176        Ok(tz) => tz,
2177        Err(_) => return Err(EvalError::InvalidIanaTimezoneId(tz_str.into())),
2178    };
2179    let offset = tz.offset_from_utc_datetime(&b.naive_utc());
2180    Ok(temp_storage.make_datum(|packer| {
2181        packer.push_list_with(|packer| {
2182            packer.push(Datum::from(offset.abbreviation()));
2183            packer.push(Datum::from(offset.base_utc_offset()));
2184            packer.push(Datum::from(offset.dst_offset()));
2185        });
2186    }))
2187}
2188
2189/// Determines if an mz_aclitem contains one of the specified privileges. This will return true if
2190/// any of the listed privileges are contained in the mz_aclitem.
2191#[sqlfunc(
2192    sqlname = "mz_aclitem_contains_privilege",
2193    output_type = "bool",
2194    propagates_nulls = true
2195)]
2196fn mz_acl_item_contains_privilege(
2197    mz_acl_item: MzAclItem,
2198    privileges: &str,
2199) -> Result<bool, EvalError> {
2200    let acl_mode = AclMode::parse_multiple_privileges(privileges)
2201        .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
2202    let contains = !mz_acl_item.acl_mode.intersection(acl_mode).is_empty();
2203    Ok(contains)
2204}
2205
2206#[sqlfunc]
2207// transliterated from postgres/src/backend/utils/adt/misc.c
2208fn parse_ident<'a>(ident: &'a str, strict: bool) -> Result<ArrayRustType<Cow<'a, str>>, EvalError> {
2209    fn is_ident_start(c: char) -> bool {
2210        matches!(c, 'A'..='Z' | 'a'..='z' | '_' | '\u{80}'..=char::MAX)
2211    }
2212
2213    fn is_ident_cont(c: char) -> bool {
2214        matches!(c, '0'..='9' | '$') || is_ident_start(c)
2215    }
2216
2217    let mut elems = vec![];
2218    let buf = &mut LexBuf::new(ident);
2219
2220    let mut after_dot = false;
2221
2222    buf.take_while(|ch| ch.is_ascii_whitespace());
2223
2224    loop {
2225        let mut missing_ident = true;
2226
2227        let c = buf.next();
2228
2229        if c == Some('"') {
2230            let s = buf.take_while(|ch| !matches!(ch, '"'));
2231
2232            if buf.next() != Some('"') {
2233                return Err(EvalError::InvalidIdentifier {
2234                    ident: ident.into(),
2235                    detail: Some("String has unclosed double quotes.".into()),
2236                });
2237            }
2238            elems.push(Cow::Borrowed(s));
2239            missing_ident = false;
2240        } else if c.map(is_ident_start).unwrap_or(false) {
2241            buf.prev();
2242            let s = buf.take_while(is_ident_cont);
2243            elems.push(Cow::Owned(s.to_ascii_lowercase()));
2244            missing_ident = false;
2245        }
2246
2247        if missing_ident {
2248            if c == Some('.') {
2249                return Err(EvalError::InvalidIdentifier {
2250                    ident: ident.into(),
2251                    detail: Some("No valid identifier before \".\".".into()),
2252                });
2253            } else if after_dot {
2254                return Err(EvalError::InvalidIdentifier {
2255                    ident: ident.into(),
2256                    detail: Some("No valid identifier after \".\".".into()),
2257                });
2258            } else {
2259                return Err(EvalError::InvalidIdentifier {
2260                    ident: ident.into(),
2261                    detail: None,
2262                });
2263            }
2264        }
2265
2266        buf.take_while(|ch| ch.is_ascii_whitespace());
2267
2268        match buf.next() {
2269            Some('.') => {
2270                after_dot = true;
2271
2272                buf.take_while(|ch| ch.is_ascii_whitespace());
2273            }
2274            Some(_) if strict => {
2275                return Err(EvalError::InvalidIdentifier {
2276                    ident: ident.into(),
2277                    detail: None,
2278                });
2279            }
2280            _ => break,
2281        }
2282    }
2283
2284    Ok(elems.into())
2285}
2286
2287fn regexp_split_to_array_re<'a>(
2288    text: &str,
2289    regexp: &Regex,
2290    temp_storage: &'a RowArena,
2291) -> Result<Datum<'a>, EvalError> {
2292    let found = mz_regexp::regexp_split_to_array(text, regexp);
2293    let mut row = Row::default();
2294    let mut packer = row.packer();
2295    packer.try_push_array(
2296        &[ArrayDimension {
2297            lower_bound: 1,
2298            length: found.len(),
2299        }],
2300        found.into_iter().map(Datum::String),
2301    )?;
2302    Ok(temp_storage.push_unary_row(row))
2303}
2304
2305#[sqlfunc(propagates_nulls = true)]
2306fn pretty_sql<'a>(sql: &str, width: i32, temp_storage: &'a RowArena) -> Result<&'a str, EvalError> {
2307    let width =
2308        usize::try_from(width).map_err(|_| EvalError::PrettyError("invalid width".into()))?;
2309    let pretty = pretty_str(
2310        sql,
2311        PrettyConfig {
2312            width,
2313            format_mode: FormatMode::Simple,
2314        },
2315    )
2316    .map_err(|e| EvalError::PrettyError(e.to_string().into()))?;
2317    let pretty = temp_storage.push_string(pretty);
2318    Ok(pretty)
2319}
2320
2321#[sqlfunc(propagates_nulls = true)]
2322fn starts_with(a: &str, b: &str) -> bool {
2323    a.starts_with(b)
2324}
2325
2326#[sqlfunc(
2327    sqlname = "||",
2328    is_infix_op = true,
2329    propagates_nulls = true,
2330    // Text concatenation is monotonic in its second argument, because if I change the
2331    // second argument but don't change the first argument, then we won't find a difference
2332    // in that part of the concatenation result that came from the first argument, so we'll
2333    // find the difference that comes from changing the second argument.
2334    // (It's not monotonic in its first argument, because e.g.,
2335    // 'A' < 'AA' but 'AZ' > 'AAZ'.)
2336    is_monotone = (false, true),
2337)]
2338fn text_concat_binary(a: &str, b: &str) -> Result<String, EvalError> {
2339    if a.len() + b.len() > MAX_STRING_FUNC_RESULT_BYTES {
2340        return Err(EvalError::LengthTooLarge);
2341    }
2342    let mut buf = String::with_capacity(a.len() + b.len());
2343    buf.push_str(a);
2344    buf.push_str(b);
2345    Ok(buf)
2346}
2347
2348#[sqlfunc(propagates_nulls = true, introduces_nulls = false)]
2349fn like_escape<'a>(
2350    pattern: &str,
2351    b: &str,
2352    temp_storage: &'a RowArena,
2353) -> Result<&'a str, EvalError> {
2354    let escape = like_pattern::EscapeBehavior::from_str(b)?;
2355    let normalized = like_pattern::normalize_pattern(pattern, escape)?;
2356    Ok(temp_storage.push_string(normalized))
2357}
2358
2359#[sqlfunc(is_infix_op = true, sqlname = "like")]
2360fn is_like_match_case_sensitive(haystack: &str, pattern: &str) -> Result<bool, EvalError> {
2361    like_pattern::compile(pattern, false).map(|needle| needle.is_match(haystack))
2362}
2363
2364#[sqlfunc(is_infix_op = true, sqlname = "ilike")]
2365fn is_like_match_case_insensitive(haystack: &str, pattern: &str) -> Result<bool, EvalError> {
2366    like_pattern::compile(pattern, true).map(|needle| needle.is_match(haystack))
2367}
2368
2369#[sqlfunc(is_infix_op = true, sqlname = "~")]
2370fn is_regexp_match_case_sensitive(haystack: &str, needle: &str) -> Result<bool, EvalError> {
2371    let regex = build_regex(needle, "")?;
2372    Ok(regex.is_match(haystack))
2373}
2374
2375#[sqlfunc(is_infix_op = true, sqlname = "~*")]
2376fn is_regexp_match_case_insensitive(haystack: &str, needle: &str) -> Result<bool, EvalError> {
2377    let regex = build_regex(needle, "i")?;
2378    Ok(regex.is_match(haystack))
2379}
2380
2381fn regexp_match_static<'a>(
2382    haystack: Datum<'a>,
2383    temp_storage: &'a RowArena,
2384    needle: &regex::Regex,
2385) -> Result<Datum<'a>, EvalError> {
2386    let mut row = Row::default();
2387    let mut packer = row.packer();
2388    if needle.captures_len() > 1 {
2389        // The regex contains capture groups, so return an array containing the
2390        // matched text in each capture group, unless the entire match fails.
2391        // Individual capture groups may also be null if that group did not
2392        // participate in the match.
2393        match needle.captures(haystack.unwrap_str()) {
2394            None => packer.push(Datum::Null),
2395            Some(captures) => packer.try_push_array(
2396                &[ArrayDimension {
2397                    lower_bound: 1,
2398                    length: captures.len() - 1,
2399                }],
2400                // Skip the 0th capture group, which is the whole match.
2401                captures.iter().skip(1).map(|mtch| match mtch {
2402                    None => Datum::Null,
2403                    Some(mtch) => Datum::String(mtch.as_str()),
2404                }),
2405            )?,
2406        }
2407    } else {
2408        // The regex contains no capture groups, so return a one-element array
2409        // containing the match, or null if there is no match.
2410        match needle.find(haystack.unwrap_str()) {
2411            None => packer.push(Datum::Null),
2412            Some(mtch) => packer.try_push_array(
2413                &[ArrayDimension {
2414                    lower_bound: 1,
2415                    length: 1,
2416                }],
2417                iter::once(Datum::String(mtch.as_str())),
2418            )?,
2419        };
2420    };
2421    Ok(temp_storage.push_unary_row(row))
2422}
2423
2424/// Sets `limit` based on the presence of 'g' in `flags` for use in `Regex::replacen`,
2425/// and removes 'g' from `flags` if present.
2426pub(crate) fn regexp_replace_parse_flags(flags: &str) -> (usize, Cow<'_, str>) {
2427    // 'g' means to replace all instead of the first. Use a Cow to avoid allocating in the fast
2428    // path. We could switch build_regex to take an iter which would also achieve that.
2429    let (limit, flags) = if flags.contains('g') {
2430        let flags = flags.replace('g', "");
2431        (0, Cow::Owned(flags))
2432    } else {
2433        (1, Cow::Borrowed(flags))
2434    };
2435    (limit, flags)
2436}
2437
2438pub fn build_regex(needle: &str, flags: &str) -> Result<Regex, EvalError> {
2439    let mut case_insensitive = false;
2440    // Note: Postgres accepts it when both flags are present, taking the last one. We do the same.
2441    for f in flags.chars() {
2442        match f {
2443            'i' => {
2444                case_insensitive = true;
2445            }
2446            'c' => {
2447                case_insensitive = false;
2448            }
2449            _ => return Err(EvalError::InvalidRegexFlag(f)),
2450        }
2451    }
2452    Ok(Regex::new(needle, case_insensitive)?)
2453}
2454
2455#[sqlfunc(sqlname = "repeat")]
2456fn repeat_string(string: &str, count: i32) -> Result<String, EvalError> {
2457    let len = usize::try_from(count).unwrap_or(0);
2458    if (len * string.len()) > MAX_STRING_FUNC_RESULT_BYTES {
2459        return Err(EvalError::LengthTooLarge);
2460    }
2461    Ok(string.repeat(len))
2462}
2463
2464/// Constructs a new zero or one dimensional array out of an arbitrary number of
2465/// scalars.
2466///
2467/// If `datums` is empty, constructs a zero-dimensional array. Otherwise,
2468/// constructs a one dimensional array whose lower bound is one and whose length
2469/// is equal to `datums.len()`.
2470fn array_create_scalar<'a>(
2471    datums: &[Datum<'a>],
2472    temp_storage: &'a RowArena,
2473) -> Result<Datum<'a>, EvalError> {
2474    let mut dims = &[ArrayDimension {
2475        lower_bound: 1,
2476        length: datums.len(),
2477    }][..];
2478    if datums.is_empty() {
2479        // Per PostgreSQL, empty arrays are represented with zero dimensions,
2480        // not one dimension of zero length. We write this condition a little
2481        // strangely to satisfy the borrow checker while avoiding an allocation.
2482        dims = &[];
2483    }
2484    let datum = temp_storage.try_make_datum(|packer| packer.try_push_array(dims, datums))?;
2485    Ok(datum)
2486}
2487
2488fn stringify_datum<'a, B>(
2489    buf: &mut B,
2490    d: Datum<'a>,
2491    ty: &SqlScalarType,
2492) -> Result<strconv::Nestable, EvalError>
2493where
2494    B: FormatBuffer,
2495{
2496    use SqlScalarType::*;
2497    match &ty {
2498        AclItem => Ok(strconv::format_acl_item(buf, d.unwrap_acl_item())),
2499        Bool => Ok(strconv::format_bool(buf, d.unwrap_bool())),
2500        Int16 => Ok(strconv::format_int16(buf, d.unwrap_int16())),
2501        Int32 => Ok(strconv::format_int32(buf, d.unwrap_int32())),
2502        Int64 => Ok(strconv::format_int64(buf, d.unwrap_int64())),
2503        UInt16 => Ok(strconv::format_uint16(buf, d.unwrap_uint16())),
2504        UInt32 | Oid | RegClass | RegProc | RegType => {
2505            Ok(strconv::format_uint32(buf, d.unwrap_uint32()))
2506        }
2507        UInt64 => Ok(strconv::format_uint64(buf, d.unwrap_uint64())),
2508        Float32 => Ok(strconv::format_float32(buf, d.unwrap_float32())),
2509        Float64 => Ok(strconv::format_float64(buf, d.unwrap_float64())),
2510        Numeric { .. } => Ok(strconv::format_numeric(buf, &d.unwrap_numeric())),
2511        Date => Ok(strconv::format_date(buf, d.unwrap_date())),
2512        Time => Ok(strconv::format_time(buf, d.unwrap_time())),
2513        Timestamp { .. } => Ok(strconv::format_timestamp(buf, &d.unwrap_timestamp())),
2514        TimestampTz { .. } => Ok(strconv::format_timestamptz(buf, &d.unwrap_timestamptz())),
2515        Interval => Ok(strconv::format_interval(buf, d.unwrap_interval())),
2516        Bytes => Ok(strconv::format_bytes(buf, d.unwrap_bytes())),
2517        String | VarChar { .. } | PgLegacyName => Ok(strconv::format_string(buf, d.unwrap_str())),
2518        Char { length } => Ok(strconv::format_string(
2519            buf,
2520            &mz_repr::adt::char::format_str_pad(d.unwrap_str(), *length),
2521        )),
2522        PgLegacyChar => {
2523            format_pg_legacy_char(buf, d.unwrap_uint8())?;
2524            Ok(strconv::Nestable::MayNeedEscaping)
2525        }
2526        Jsonb => Ok(strconv::format_jsonb(buf, JsonbRef::from_datum(d))),
2527        Uuid => Ok(strconv::format_uuid(buf, d.unwrap_uuid())),
2528        Record { fields, .. } => {
2529            let mut fields = fields.iter();
2530            strconv::format_record(buf, d.unwrap_list(), |buf, d| {
2531                let (_name, ty) = fields.next().unwrap();
2532                if d.is_null() {
2533                    Ok(buf.write_null())
2534                } else {
2535                    stringify_datum(buf.nonnull_buffer(), d, &ty.scalar_type)
2536                }
2537            })
2538        }
2539        Array(elem_type) => strconv::format_array(
2540            buf,
2541            &d.unwrap_array().dims().into_iter().collect::<Vec<_>>(),
2542            d.unwrap_array().elements(),
2543            |buf, d| {
2544                if d.is_null() {
2545                    Ok(buf.write_null())
2546                } else {
2547                    stringify_datum(buf.nonnull_buffer(), d, elem_type)
2548                }
2549            },
2550        ),
2551        List { element_type, .. } => strconv::format_list(buf, d.unwrap_list(), |buf, d| {
2552            if d.is_null() {
2553                Ok(buf.write_null())
2554            } else {
2555                stringify_datum(buf.nonnull_buffer(), d, element_type)
2556            }
2557        }),
2558        Map { value_type, .. } => strconv::format_map(buf, &d.unwrap_map(), |buf, d| {
2559            if d.is_null() {
2560                Ok(buf.write_null())
2561            } else {
2562                stringify_datum(buf.nonnull_buffer(), d, value_type)
2563            }
2564        }),
2565        Int2Vector => strconv::format_legacy_vector(buf, d.unwrap_array().elements(), |buf, d| {
2566            stringify_datum(buf.nonnull_buffer(), d, &SqlScalarType::Int16)
2567        }),
2568        MzTimestamp { .. } => Ok(strconv::format_mz_timestamp(buf, d.unwrap_mz_timestamp())),
2569        Range { element_type } => strconv::format_range(buf, &d.unwrap_range(), |buf, d| match d {
2570            Some(d) => stringify_datum(buf.nonnull_buffer(), *d, element_type),
2571            None => Ok::<_, EvalError>(buf.write_null()),
2572        }),
2573        MzAclItem => Ok(strconv::format_mz_acl_item(buf, d.unwrap_mz_acl_item())),
2574    }
2575}
2576
2577#[sqlfunc]
2578fn position(substring: &str, string: &str) -> Result<i32, EvalError> {
2579    let char_index = string.find(substring);
2580
2581    if let Some(char_index) = char_index {
2582        // find the index in char space
2583        let string_prefix = &string[0..char_index];
2584
2585        let num_prefix_chars = string_prefix.chars().count();
2586        let num_prefix_chars = i32::try_from(num_prefix_chars)
2587            .map_err(|_| EvalError::Int32OutOfRange(num_prefix_chars.to_string().into()))?;
2588
2589        Ok(num_prefix_chars + 1)
2590    } else {
2591        Ok(0)
2592    }
2593}
2594
2595#[sqlfunc]
2596fn strpos(string: &str, substring: &str) -> Result<i32, EvalError> {
2597    position(substring, string)
2598}
2599
2600#[sqlfunc(
2601    propagates_nulls = true,
2602    // `left` is unfortunately not monotonic (at least for negative second arguments),
2603    // because 'aa' < 'z', but `left(_, -1)` makes 'a' > ''.
2604    is_monotone = (false, false)
2605)]
2606fn left<'a>(string: &'a str, b: i32) -> Result<&'a str, EvalError> {
2607    let n = i64::from(b);
2608
2609    let mut byte_indices = string.char_indices().map(|(i, _)| i);
2610
2611    let end_in_bytes = match n.cmp(&0) {
2612        Ordering::Equal => 0,
2613        Ordering::Greater => {
2614            let n = usize::try_from(n).map_err(|_| {
2615                EvalError::InvalidParameterValue(format!("invalid parameter n: {:?}", n).into())
2616            })?;
2617            // nth from the back
2618            byte_indices.nth(n).unwrap_or(string.len())
2619        }
2620        Ordering::Less => {
2621            let n = usize::try_from(n.abs() - 1).map_err(|_| {
2622                EvalError::InvalidParameterValue(format!("invalid parameter n: {:?}", n).into())
2623            })?;
2624            byte_indices.rev().nth(n).unwrap_or(0)
2625        }
2626    };
2627
2628    Ok(&string[..end_in_bytes])
2629}
2630
2631#[sqlfunc(propagates_nulls = true)]
2632fn right<'a>(string: &'a str, n: i32) -> Result<&'a str, EvalError> {
2633    let mut byte_indices = string.char_indices().map(|(i, _)| i);
2634
2635    let start_in_bytes = if n == 0 {
2636        string.len()
2637    } else if n > 0 {
2638        let n = usize::try_from(n - 1).map_err(|_| {
2639            EvalError::InvalidParameterValue(format!("invalid parameter n: {:?}", n).into())
2640        })?;
2641        // nth from the back
2642        byte_indices.rev().nth(n).unwrap_or(0)
2643    } else if n == i32::MIN {
2644        // this seems strange but Postgres behaves like this
2645        0
2646    } else {
2647        let n = n.abs();
2648        let n = usize::try_from(n).map_err(|_| {
2649            EvalError::InvalidParameterValue(format!("invalid parameter n: {:?}", n).into())
2650        })?;
2651        byte_indices.nth(n).unwrap_or(string.len())
2652    };
2653
2654    Ok(&string[start_in_bytes..])
2655}
2656
2657#[sqlfunc(sqlname = "btrim", propagates_nulls = true)]
2658fn trim<'a>(a: &'a str, trim_chars: &str) -> &'a str {
2659    a.trim_matches(|c| trim_chars.contains(c))
2660}
2661
2662#[sqlfunc(sqlname = "ltrim", propagates_nulls = true)]
2663fn trim_leading<'a>(a: &'a str, trim_chars: &str) -> &'a str {
2664    a.trim_start_matches(|c| trim_chars.contains(c))
2665}
2666
2667#[sqlfunc(sqlname = "rtrim", propagates_nulls = true)]
2668fn trim_trailing<'a>(a: &'a str, trim_chars: &str) -> &'a str {
2669    a.trim_end_matches(|c| trim_chars.contains(c))
2670}
2671
2672#[sqlfunc(
2673    sqlname = "array_length",
2674    propagates_nulls = true,
2675    introduces_nulls = true
2676)]
2677fn array_length<'a>(a: Array<'a>, b: i64) -> Result<Option<i32>, EvalError> {
2678    let i = match usize::try_from(b) {
2679        Ok(0) | Err(_) => return Ok(None),
2680        Ok(n) => n - 1,
2681    };
2682    Ok(match a.dims().into_iter().nth(i) {
2683        None => None,
2684        Some(dim) => Some(
2685            dim.length
2686                .try_into()
2687                .map_err(|_| EvalError::Int32OutOfRange(dim.length.to_string().into()))?,
2688        ),
2689    })
2690}
2691
2692#[sqlfunc(
2693    output_type = "Option<i32>",
2694    is_infix_op = true,
2695    sqlname = "array_lower",
2696    propagates_nulls = true,
2697    introduces_nulls = true
2698)]
2699// TODO(benesch): remove potentially dangerous usage of `as`.
2700#[allow(clippy::as_conversions)]
2701fn array_lower<'a>(a: Array<'a>, i: i64) -> Option<i32> {
2702    if i < 1 {
2703        return None;
2704    }
2705    match a.dims().into_iter().nth(i as usize - 1) {
2706        Some(_) => Some(1),
2707        None => None,
2708    }
2709}
2710
2711#[sqlfunc(
2712    output_type_expr = "input_types[0].scalar_type.without_modifiers().nullable(true)",
2713    sqlname = "array_remove",
2714    propagates_nulls = false,
2715    introduces_nulls = false
2716)]
2717fn array_remove<'a>(
2718    arr: Array<'a>,
2719    b: Datum<'a>,
2720    temp_storage: &'a RowArena,
2721) -> Result<Datum<'a>, EvalError> {
2722    // Zero-dimensional arrays are empty by definition
2723    if arr.dims().len() == 0 {
2724        return Ok(Datum::Array(arr));
2725    }
2726
2727    // array_remove only supports one-dimensional arrays
2728    if arr.dims().len() > 1 {
2729        return Err(EvalError::MultidimensionalArrayRemovalNotSupported);
2730    }
2731
2732    let elems: Vec<_> = arr.elements().iter().filter(|v| v != &b).collect();
2733    let mut dims = arr.dims().into_iter().collect::<Vec<_>>();
2734    // This access is safe because `dims` is guaranteed to be non-empty
2735    dims[0] = ArrayDimension {
2736        lower_bound: 1,
2737        length: elems.len(),
2738    };
2739
2740    Ok(temp_storage.try_make_datum(|packer| packer.try_push_array(&dims, elems))?)
2741}
2742
2743#[sqlfunc(
2744    output_type = "Option<i32>",
2745    is_infix_op = true,
2746    sqlname = "array_upper",
2747    propagates_nulls = true,
2748    introduces_nulls = true
2749)]
2750// TODO(benesch): remove potentially dangerous usage of `as`.
2751#[allow(clippy::as_conversions)]
2752fn array_upper<'a>(a: Array<'a>, i: i64) -> Result<Option<i32>, EvalError> {
2753    if i < 1 {
2754        return Ok(None);
2755    }
2756    a.dims()
2757        .into_iter()
2758        .nth(i as usize - 1)
2759        .map(|dim| {
2760            dim.length
2761                .try_into()
2762                .map_err(|_| EvalError::Int32OutOfRange(dim.length.to_string().into()))
2763        })
2764        .transpose()
2765}
2766
2767#[sqlfunc(
2768    is_infix_op = true,
2769    sqlname = "array_contains",
2770    propagates_nulls = true,
2771    introduces_nulls = false
2772)]
2773fn array_contains<'a>(a: Datum<'a>, array: Array<'a>) -> bool {
2774    array.elements().iter().any(|e| e == a)
2775}
2776
2777#[sqlfunc(is_infix_op = true, sqlname = "@>")]
2778fn array_contains_array<'a>(a: Array<'a>, b: Array<'a>) -> bool {
2779    let a = a.elements();
2780    let b = b.elements();
2781
2782    // NULL is never equal to NULL. If NULL is an element of b, b cannot be contained in a, even if a contains NULL.
2783    if b.iter().contains(&Datum::Null) {
2784        false
2785    } else {
2786        b.iter()
2787            .all(|item_b| a.iter().any(|item_a| item_a == item_b))
2788    }
2789}
2790
2791#[sqlfunc(is_infix_op = true, sqlname = "<@")]
2792fn array_contains_array_rev<'a>(a: Array<'a>, b: Array<'a>) -> bool {
2793    array_contains_array(b, a)
2794}
2795
2796#[sqlfunc(
2797    output_type_expr = "input_types[0].scalar_type.without_modifiers().nullable(true)",
2798    is_infix_op = true,
2799    sqlname = "||",
2800    propagates_nulls = false,
2801    introduces_nulls = false
2802)]
2803fn array_array_concat<'a>(
2804    a: Option<Array<'a>>,
2805    b: Option<Array<'a>>,
2806    temp_storage: &'a RowArena,
2807) -> Result<Option<Array<'a>>, EvalError> {
2808    let Some(a_array) = a else {
2809        return Ok(b);
2810    };
2811    let Some(b_array) = b else {
2812        return Ok(a);
2813    };
2814
2815    let a_dims: Vec<ArrayDimension> = a_array.dims().into_iter().collect();
2816    let b_dims: Vec<ArrayDimension> = b_array.dims().into_iter().collect();
2817
2818    let a_ndims = a_dims.len();
2819    let b_ndims = b_dims.len();
2820
2821    // Per PostgreSQL, if either of the input arrays is zero dimensional,
2822    // the output is the other array, no matter their dimensions.
2823    if a_ndims == 0 {
2824        return Ok(b);
2825    } else if b_ndims == 0 {
2826        return Ok(a);
2827    }
2828
2829    // Postgres supports concatenating arrays of different dimensions,
2830    // as long as one of the arrays has the same type as an element of
2831    // the other array, i.e. `int[2][4] || int[4]` (or `int[4] || int[2][4]`)
2832    // works, because each element of `int[2][4]` is an `int[4]`.
2833    // This check is separate from the one below because Postgres gives a
2834    // specific error message if the number of dimensions differs by more
2835    // than one.
2836    // This cast is safe since MAX_ARRAY_DIMENSIONS is 6
2837    // Can be replaced by .abs_diff once it is stabilized
2838    // TODO(benesch): remove potentially dangerous usage of `as`.
2839    #[allow(clippy::as_conversions)]
2840    if (a_ndims as isize - b_ndims as isize).abs() > 1 {
2841        return Err(EvalError::IncompatibleArrayDimensions {
2842            dims: Some((a_ndims, b_ndims)),
2843        });
2844    }
2845
2846    let mut dims;
2847
2848    // After the checks above, we are certain that:
2849    // - neither array is zero dimensional nor empty
2850    // - both arrays have the same number of dimensions, or differ
2851    //   at most by one.
2852    match a_ndims.cmp(&b_ndims) {
2853        // If both arrays have the same number of dimensions, validate
2854        // that their inner dimensions are the same and concatenate the
2855        // arrays.
2856        Ordering::Equal => {
2857            if &a_dims[1..] != &b_dims[1..] {
2858                return Err(EvalError::IncompatibleArrayDimensions { dims: None });
2859            }
2860            dims = vec![ArrayDimension {
2861                lower_bound: a_dims[0].lower_bound,
2862                length: a_dims[0].length + b_dims[0].length,
2863            }];
2864            dims.extend(&a_dims[1..]);
2865        }
2866        // If `a` has less dimensions than `b`, this is an element-array
2867        // concatenation, which requires that `a` has the same dimensions
2868        // as an element of `b`.
2869        Ordering::Less => {
2870            if &a_dims[..] != &b_dims[1..] {
2871                return Err(EvalError::IncompatibleArrayDimensions { dims: None });
2872            }
2873            dims = vec![ArrayDimension {
2874                lower_bound: b_dims[0].lower_bound,
2875                // Since `a` is treated as an element of `b`, the length of
2876                // the first dimension of `b` is incremented by one, as `a` is
2877                // non-empty.
2878                length: b_dims[0].length + 1,
2879            }];
2880            dims.extend(a_dims);
2881        }
2882        // If `a` has more dimensions than `b`, this is an array-element
2883        // concatenation, which requires that `b` has the same dimensions
2884        // as an element of `a`.
2885        Ordering::Greater => {
2886            if &a_dims[1..] != &b_dims[..] {
2887                return Err(EvalError::IncompatibleArrayDimensions { dims: None });
2888            }
2889            dims = vec![ArrayDimension {
2890                lower_bound: a_dims[0].lower_bound,
2891                // Since `b` is treated as an element of `a`, the length of
2892                // the first dimension of `a` is incremented by one, as `b`
2893                // is non-empty.
2894                length: a_dims[0].length + 1,
2895            }];
2896            dims.extend(b_dims);
2897        }
2898    }
2899
2900    let elems = a_array.elements().iter().chain(b_array.elements().iter());
2901
2902    let datum = temp_storage.try_make_datum(|packer| packer.try_push_array(&dims, elems))?;
2903    Ok(Some(datum.unwrap_array()))
2904}
2905
2906#[sqlfunc(
2907    is_infix_op = true,
2908    sqlname = "||",
2909    propagates_nulls = false,
2910    introduces_nulls = false
2911)]
2912fn list_list_concat<'a, T: FromDatum<'a>>(
2913    a: Option<DatumList<'a, T>>,
2914    b: Option<DatumList<'a, T>>,
2915    temp_storage: &'a RowArena,
2916) -> Option<DatumList<'a, T>> {
2917    let Some(a) = a else {
2918        return b;
2919    };
2920    let Some(b) = b else {
2921        return Some(a);
2922    };
2923
2924    Some(temp_storage.make_datum_list(a.typed_iter().chain(b.typed_iter())))
2925}
2926
2927#[sqlfunc(is_infix_op = true, sqlname = "||", propagates_nulls = false)]
2928fn list_element_concat<'a, T: FromDatum<'a>>(
2929    a: Option<DatumList<'a, T>>,
2930    b: T,
2931    temp_storage: &'a RowArena,
2932) -> DatumList<'a, T> {
2933    let a_elems = a.into_iter().flat_map(|a| a.typed_iter());
2934    temp_storage.make_datum_list(a_elems.chain(std::iter::once(b)))
2935}
2936
2937// Note that the output type corresponds to the _second_ parameter's input type.
2938#[sqlfunc(is_infix_op = true, sqlname = "||", propagates_nulls = false)]
2939fn element_list_concat<'a, T: FromDatum<'a>>(
2940    a: T,
2941    b: Option<DatumList<'a, T>>,
2942    temp_storage: &'a RowArena,
2943) -> DatumList<'a, T> {
2944    let b_elems = b.into_iter().flat_map(|b| b.typed_iter());
2945    temp_storage.make_datum_list(std::iter::once(a).chain(b_elems))
2946}
2947
2948#[sqlfunc(sqlname = "list_remove")]
2949fn list_remove<'a, T: FromDatum<'a>>(
2950    a: DatumList<'a, T>,
2951    b: T,
2952    temp_storage: &'a RowArena,
2953) -> DatumList<'a, T> {
2954    temp_storage.make_datum_list(a.typed_iter().filter(|elem| *elem != b))
2955}
2956
2957#[sqlfunc(sqlname = "digest")]
2958fn digest_string(to_digest: &str, digest_fn: &str) -> Result<Vec<u8>, EvalError> {
2959    digest_inner(to_digest.as_bytes(), digest_fn)
2960}
2961
2962#[sqlfunc(sqlname = "digest")]
2963fn digest_bytes(to_digest: &[u8], digest_fn: &str) -> Result<Vec<u8>, EvalError> {
2964    digest_inner(to_digest, digest_fn)
2965}
2966
2967fn digest_inner(bytes: &[u8], digest_fn: &str) -> Result<Vec<u8>, EvalError> {
2968    match digest_fn {
2969        "md5" => Ok(Md5::digest(bytes).to_vec()),
2970        "sha1" => Ok(Sha1::digest(bytes).to_vec()),
2971        "sha224" => Ok(Sha224::digest(bytes).to_vec()),
2972        "sha256" => Ok(Sha256::digest(bytes).to_vec()),
2973        "sha384" => Ok(Sha384::digest(bytes).to_vec()),
2974        "sha512" => Ok(Sha512::digest(bytes).to_vec()),
2975        other => Err(EvalError::InvalidHashAlgorithm(other.into())),
2976    }
2977}
2978
2979#[sqlfunc]
2980fn mz_render_typmod(oid: u32, typmod: i32) -> String {
2981    match Type::from_oid_and_typmod(oid, typmod) {
2982        Ok(typ) => typ.constraint().display_or("").to_string(),
2983        // Match dubious PostgreSQL behavior of outputting the unmodified
2984        // `typmod` when positive if the type OID/typmod is invalid.
2985        Err(_) if typmod >= 0 => format!("({typmod})"),
2986        Err(_) => "".into(),
2987    }
2988}
2989
2990#[cfg(test)]
2991mod test {
2992    use chrono::prelude::*;
2993    use mz_repr::PropDatum;
2994    use proptest::prelude::*;
2995
2996    use super::*;
2997    use crate::MirScalarExpr;
2998
2999    #[mz_ore::test]
3000    fn add_interval_months() {
3001        let dt = ym(2000, 1);
3002
3003        assert_eq!(add_timestamp_months(&*dt, 0).unwrap(), dt);
3004        assert_eq!(add_timestamp_months(&*dt, 1).unwrap(), ym(2000, 2));
3005        assert_eq!(add_timestamp_months(&*dt, 12).unwrap(), ym(2001, 1));
3006        assert_eq!(add_timestamp_months(&*dt, 13).unwrap(), ym(2001, 2));
3007        assert_eq!(add_timestamp_months(&*dt, 24).unwrap(), ym(2002, 1));
3008        assert_eq!(add_timestamp_months(&*dt, 30).unwrap(), ym(2002, 7));
3009
3010        // and negatives
3011        assert_eq!(add_timestamp_months(&*dt, -1).unwrap(), ym(1999, 12));
3012        assert_eq!(add_timestamp_months(&*dt, -12).unwrap(), ym(1999, 1));
3013        assert_eq!(add_timestamp_months(&*dt, -13).unwrap(), ym(1998, 12));
3014        assert_eq!(add_timestamp_months(&*dt, -24).unwrap(), ym(1998, 1));
3015        assert_eq!(add_timestamp_months(&*dt, -30).unwrap(), ym(1997, 7));
3016
3017        // and going over a year boundary by less than a year
3018        let dt = ym(1999, 12);
3019        assert_eq!(add_timestamp_months(&*dt, 1).unwrap(), ym(2000, 1));
3020        let end_of_month_dt = NaiveDate::from_ymd_opt(1999, 12, 31)
3021            .unwrap()
3022            .and_hms_opt(9, 9, 9)
3023            .unwrap();
3024        assert_eq!(
3025            // leap year
3026            add_timestamp_months(&end_of_month_dt, 2).unwrap(),
3027            NaiveDate::from_ymd_opt(2000, 2, 29)
3028                .unwrap()
3029                .and_hms_opt(9, 9, 9)
3030                .unwrap()
3031                .try_into()
3032                .unwrap(),
3033        );
3034        assert_eq!(
3035            // not leap year
3036            add_timestamp_months(&end_of_month_dt, 14).unwrap(),
3037            NaiveDate::from_ymd_opt(2001, 2, 28)
3038                .unwrap()
3039                .and_hms_opt(9, 9, 9)
3040                .unwrap()
3041                .try_into()
3042                .unwrap(),
3043        );
3044    }
3045
3046    fn ym(year: i32, month: u32) -> CheckedTimestamp<NaiveDateTime> {
3047        NaiveDate::from_ymd_opt(year, month, 1)
3048            .unwrap()
3049            .and_hms_opt(9, 9, 9)
3050            .unwrap()
3051            .try_into()
3052            .unwrap()
3053    }
3054
3055    #[mz_ore::test]
3056    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decNumberFromInt32` on OS `linux`
3057    fn test_is_monotone() {
3058        use proptest::prelude::*;
3059
3060        /// Asserts that the function is either monotonically increasing or decreasing over
3061        /// the given sets of arguments.
3062        fn assert_monotone<'a, const N: usize>(
3063            expr: &MirScalarExpr,
3064            arena: &'a RowArena,
3065            datums: &[[Datum<'a>; N]],
3066        ) {
3067            // TODO: assertions for nulls, errors
3068            let Ok(results) = datums
3069                .iter()
3070                .map(|args| expr.eval(args.as_slice(), arena))
3071                .collect::<Result<Vec<_>, _>>()
3072            else {
3073                return;
3074            };
3075
3076            let forward = results.iter().tuple_windows().all(|(a, b)| a <= b);
3077            let reverse = results.iter().tuple_windows().all(|(a, b)| a >= b);
3078            assert!(
3079                forward || reverse,
3080                "expected {expr} to be monotone, but passing {datums:?} returned {results:?}"
3081            );
3082        }
3083
3084        fn proptest_binary<'a>(
3085            func: BinaryFunc,
3086            arena: &'a RowArena,
3087            left: impl Strategy<Value = PropDatum>,
3088            right: impl Strategy<Value = PropDatum>,
3089        ) {
3090            let (left_monotone, right_monotone) = func.is_monotone();
3091            let expr = MirScalarExpr::CallBinary {
3092                func,
3093                expr1: Box::new(MirScalarExpr::column(0)),
3094                expr2: Box::new(MirScalarExpr::column(1)),
3095            };
3096            proptest!(|(
3097                mut left in proptest::array::uniform3(left),
3098                mut right in proptest::array::uniform3(right),
3099            )| {
3100                left.sort();
3101                right.sort();
3102                if left_monotone {
3103                    for r in &right {
3104                        let args: Vec<[_; 2]> = left
3105                            .iter()
3106                            .map(|l| [Datum::from(l), Datum::from(r)])
3107                            .collect();
3108                        assert_monotone(&expr, arena, &args);
3109                    }
3110                }
3111                if right_monotone {
3112                    for l in &left {
3113                        let args: Vec<[_; 2]> = right
3114                            .iter()
3115                            .map(|r| [Datum::from(l), Datum::from(r)])
3116                            .collect();
3117                        assert_monotone(&expr, arena, &args);
3118                    }
3119                }
3120            });
3121        }
3122
3123        let interesting_strs: Vec<_> = SqlScalarType::String.interesting_datums().collect();
3124        let str_datums = proptest::strategy::Union::new([
3125            proptest::string::string_regex("[A-Z]{0,10}")
3126                .expect("valid regex")
3127                .prop_map(|s| PropDatum::String(s.to_string()))
3128                .boxed(),
3129            (0..interesting_strs.len())
3130                .prop_map(move |i| {
3131                    let Datum::String(val) = interesting_strs[i] else {
3132                        unreachable!("interesting strings has non-strings")
3133                    };
3134                    PropDatum::String(val.to_string())
3135                })
3136                .boxed(),
3137        ]);
3138
3139        let interesting_i32s: Vec<Datum<'static>> =
3140            SqlScalarType::Int32.interesting_datums().collect();
3141        let i32_datums = proptest::strategy::Union::new([
3142            any::<i32>().prop_map(PropDatum::Int32).boxed(),
3143            (0..interesting_i32s.len())
3144                .prop_map(move |i| {
3145                    let Datum::Int32(val) = interesting_i32s[i] else {
3146                        unreachable!("interesting int32 has non-i32s")
3147                    };
3148                    PropDatum::Int32(val)
3149                })
3150                .boxed(),
3151            (-10i32..10).prop_map(PropDatum::Int32).boxed(),
3152        ]);
3153
3154        let arena = RowArena::new();
3155
3156        // It would be interesting to test all funcs here, but we currently need to hardcode
3157        // the generators for the argument types, which makes this tedious. Choose an interesting
3158        // subset for now.
3159        proptest_binary(
3160            BinaryFunc::AddInt32(AddInt32),
3161            &arena,
3162            &i32_datums,
3163            &i32_datums,
3164        );
3165        proptest_binary(SubInt32.into(), &arena, &i32_datums, &i32_datums);
3166        proptest_binary(MulInt32.into(), &arena, &i32_datums, &i32_datums);
3167        proptest_binary(DivInt32.into(), &arena, &i32_datums, &i32_datums);
3168        proptest_binary(TextConcatBinary.into(), &arena, &str_datums, &str_datums);
3169        proptest_binary(Left.into(), &arena, &str_datums, &i32_datums);
3170    }
3171}