dataflow-expression: Add support for the EXTRACT built-in function

- This commit adds support for the PostgreSQL built-in `EXTRACT` function.
This is a combination of Ethan's and Vassili's commits, squashed for
easier review.

Fixes: REA-4370
Change-Id: I9676fe63ab0416e95f9b698e66b02055adccafee
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/7676
Tested-by: Buildkite CI
Reviewed-by: Jason Brown <[email protected]>

Author: ethowitz

Cargo.lock

@@ -31,3 +31,6 @@ tokio = { workspace = true, features = ["full"] }
 tokio-postgres = { workspace = true, features = ["with-chrono-0_4", "with-eui48-1", "with-uuid-0_8", "with-serde_json-1", "with-bit-vec-0_6"] }
 postgres = { workspace = true, features = ["with-chrono-0_4", "with-eui48-1", "with-uuid-0_8", "with-serde_json-1", "with-bit-vec-0_6"] }
 mysql_async = { workspace = true }
+anyhow = { workspace = true }
+regex = { workspace = true }
+bytes = { workspace = true }

dataflow-expression/Cargo.toml

@@ -5,13 +5,15 @@ use std::ops::{Add, Div, Mul, Sub};
 use std::str::FromStr;
 
 use chrono::{
-    Datelike, LocalResult, Month, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Timelike, Weekday,
+    Datelike, Duration, LocalResult, Month, NaiveDate, NaiveDateTime, NaiveTime, Offset, TimeZone,
+    Timelike, Weekday,
 };
 use chrono_tz::Tz;
 use itertools::Either;
 use mysql_time::MySqlTime;
-use readyset_data::{DfType, DfValue};
-use readyset_errors::{invalid_query_err, unsupported, ReadySetError, ReadySetResult};
+use nom_sql::TimestampField;
+use readyset_data::{DfType, DfValue, TimestampTz};
+use readyset_errors::{internal, invalid_query_err, unsupported, ReadySetError, ReadySetResult};
 use readyset_util::math::integer_rnd;
 use rust_decimal::prelude::{FromPrimitive, ToPrimitive};
 use rust_decimal::Decimal;
@@ -21,8 +23,15 @@ use vec1::Vec1;
 
 use crate::{BuiltinFunction, Expr};
 
+const MICROS_IN_SECOND: u32 = 1_000_000;
+const MILLIS_IN_SECOND: u32 = 1_000;
 const NANOS_IN_MICRO: u32 = 1_000;
+const NANOS_IN_SECOND: u32 = 1_000_000_000;
 const NANOS_IN_MILLI: u32 = 1_000_000;
+const SECONDS_IN_HOUR: u64 = 60 * 60;
+const SECONDS_IN_MINUTE: u64 = 60;
+const SECONDS_IN_DAY: u64 = SECONDS_IN_HOUR * 24;
+const MINUTES_IN_HOUR: u64 = 60;
 
 macro_rules! try_cast_or_none {
     ($df_value:expr, $to_ty:expr, $from_ty:expr) => {{
@@ -272,6 +281,220 @@ fn date_trunc(precision: DateTruncPrecision, dt: NaiveDateTime) -> ReadySetResul
     }
 }
 
+// Source:
+// https://github.com/postgres/postgres/blob/c6cf6d353c2865d82356ac86358622a101fde8ca/src/interfaces/ecpg/pgtypeslib/dt_common.c#L581-L582
+fn ymd_to_julian_date(y: i32, m: i32, d: i32) -> i32 {
+    let month: i32;
+    let year: i32;
+
+    if m > 2 {
+        month = m + 1;
+        year = y + 4800;
+    } else {
+        month = m + 13;
+        year = y + 4799;
+    }
+
+    let century = year / 100;
+    let mut julian_date = year * 365 - 32167;
+    julian_date += year / 4 - century + century / 4;
+    julian_date += 7834 * month / 256 + d;
+
+    julian_date
+}
+
+#[inline(always)]
+fn years_term(year: i32, term: i32) -> i32 {
+    if year > 0 {
+        (year + (term - 1)) / term
+    } else {
+        -(((term - 1) - (year - 1)) / term)
+    }
+}
+
+#[inline(always)]
+fn years_decade(year: i32) -> i32 {
+    if year >= 0 {
+        year / 10
+    } else {
+        -((8 - (year - 1)) / 10)
+    }
+}
+
+#[inline(always)]
+fn seconds_term(secs: u32, nanos: u32, units_per_sec: u32) -> Decimal {
+    Decimal::from(secs * units_per_sec)
+        + (Decimal::from(nanos) / Decimal::from(NANOS_IN_SECOND / units_per_sec))
+}
+
+#[inline(always)]
+fn hms_to_seconds(h: u32, m: u32, s: u32) -> u32 {
+    ((h * MINUTES_IN_HOUR as u32) + m) * SECONDS_IN_MINUTE as u32 + s
+}
+
+#[inline(always)]
+fn hms_with_nanos_to_days(h: u32, m: u32, s: u32, nanos: u32) -> Decimal {
+    seconds_term(hms_to_seconds(h, m, s), nanos, 1) / Decimal::from(SECONDS_IN_DAY).round()
+}
+
+#[inline(always)]
+fn tz_to_seconds(dt: &NaiveDateTime) -> i32 {
+    dt.and_utc().timezone().fix().local_minus_utc()
+}
+
+/* The error message we compose here, is compatible with Postgres 15, but not with 13/14
+ */
+fn invalid_extract_call_error_message(cal_type: &str, field: TimestampField) -> String {
+    let mut msg = format!(
+        "ERROR: unit \"{}\" not supported for type {}",
+        format!("{field}").to_lowercase(),
+        cal_type
+    );
+    if cal_type.eq_ignore_ascii_case("time") || cal_type.eq_ignore_ascii_case("timestamp") {
+        msg.push_str(" without time zone");
+    }
+    msg
+}
+
+fn invalid_extract_call_error(calendar_type: &str, field: TimestampField) -> ReadySetError {
+    invalid_query_err!(
+        "{}",
+        invalid_extract_call_error_message(calendar_type, field)
+    )
+}
+
+fn extract_from_time(field: TimestampField, tm: &MySqlTime) -> ReadySetResult<DfValue> {
+    macro_rules! seconds_term {
+        ($tm:expr, $units_in_second:expr) => {
+            seconds_term(
+                $tm.seconds() as u32,
+                ($tm.microseconds() * NANOS_IN_MICRO),
+                $units_in_second,
+            )
+        };
+    }
+
+    let result: DfValue = match field {
+        TimestampField::Hour => tm.hour().into(),
+        TimestampField::Minute => tm.minutes().into(),
+        TimestampField::Second => seconds_term!(tm, 1).into(),
+        TimestampField::Milliseconds => seconds_term!(tm, MILLIS_IN_SECOND).into(),
+        TimestampField::Microseconds => seconds_term!(tm, MICROS_IN_SECOND).round().into(),
+        TimestampField::Epoch => Duration::from(*tm).num_seconds().into(),
+        _ => return Err(invalid_extract_call_error("time", field)),
+    };
+
+    Ok(result)
+}
+
+fn extract_from_timestamptz(field: TimestampField, tz: &TimestampTz) -> ReadySetResult<DfValue> {
+    macro_rules! has_time_else_error {
+        ($tz:expr, $field:expr) => {
+            if $tz.has_date_only() {
+                return Err(invalid_extract_call_error("date", $field));
+            }
+        };
+    }
+
+    macro_rules! has_timezone_else_error {
+        ($tz:expr, $field:expr) => {
+            if !$tz.has_timezone() {
+                return Err(invalid_extract_call_error(
+                    if $tz.has_date_only() {
+                        "date"
+                    } else {
+                        "timestamp"
+                    },
+                    $field,
+                ));
+            }
+        };
+    }
+
+    macro_rules! seconds_term {
+        ($dt_utc:expr, $units_in_second:expr) => {
+            seconds_term($dt_utc.second(), $dt_utc.nanosecond(), $units_in_second)
+        };
+    }
+
+    let dt_utc = tz.to_chrono().naive_utc();
+    let result: DfValue = match field {
+        TimestampField::Millennium => years_term(dt_utc.year(), 1000).into(),
+        TimestampField::Century => years_term(dt_utc.year(), 100).into(),
+        TimestampField::Decade => years_decade(dt_utc.year()).into(),
+        TimestampField::Year => {
+            let year = dt_utc.year();
+            if year < 0 { year - 1 } else { year }.into()
+        }
+        TimestampField::Isoyear => {
+            let year = dt_utc.iso_week().year();
+            if year <= 0 { year - 1 } else { year }.into()
+        }
+        TimestampField::Quarter => (dt_utc.month0() / 3 + 1).into(),
+        TimestampField::Month => dt_utc.month().into(),
+        TimestampField::Week => dt_utc.iso_week().week().into(),
+        TimestampField::Day => dt_utc.day().into(),
+        TimestampField::Dow => dt_utc.weekday().num_days_from_sunday().into(),
+        TimestampField::Isodow => dt_utc.weekday().number_from_monday().into(),
+        TimestampField::Doy => dt_utc.ordinal().into(),
+        TimestampField::Hour => {
+            has_time_else_error!(tz, field);
+            dt_utc.hour().into()
+        }
+        TimestampField::Minute => {
+            has_time_else_error!(tz, field);
+            dt_utc.minute().into()
+        }
+        TimestampField::Second => {
+            has_time_else_error!(tz, field);
+            seconds_term!(dt_utc, 1).into()
+        }
+        TimestampField::Milliseconds => {
+            has_time_else_error!(tz, field);
+            seconds_term!(dt_utc, MILLIS_IN_SECOND).into()
+        }
+        TimestampField::Microseconds => {
+            has_time_else_error!(tz, field);
+            seconds_term!(dt_utc, MICROS_IN_SECOND).round().into()
+        }
+        TimestampField::Epoch => (Decimal::from(dt_utc.and_utc().timestamp_micros())
+            / Decimal::from(MICROS_IN_SECOND))
+        .round_dp(6)
+        .into(),
+        TimestampField::Julian => {
+            let julian_date: Decimal =
+                ymd_to_julian_date(dt_utc.year(), dt_utc.month() as i32, dt_utc.day() as i32)
+                    .into();
+            if tz.has_date_only() {
+                julian_date
+            } else {
+                julian_date
+                    + hms_with_nanos_to_days(
+                        dt_utc.hour(),
+                        dt_utc.minute(),
+                        dt_utc.second(),
+                        dt_utc.nanosecond(),
+                    )
+            }
+            .into()
+        }
+        TimestampField::Timezone => {
+            has_timezone_else_error!(tz, field);
+            tz_to_seconds(&dt_utc).into()
+        }
+        TimestampField::TimezoneHour => {
+            has_timezone_else_error!(tz, field);
+            (tz_to_seconds(&dt_utc) / SECONDS_IN_HOUR as i32).into()
+        }
+        TimestampField::TimezoneMinute => {
+            has_timezone_else_error!(tz, field);
+            ((tz_to_seconds(&dt_utc) % SECONDS_IN_HOUR as i32) / SECONDS_IN_MINUTE as i32).into()
+        }
+    };
+
+    Ok(result)
+}
+
 /// Format the given time value according to the given `format_string`, using the [MySQL date
 /// formatting rules][mysql-docs]. Since these rules don't match up well with anything available in
 /// the Rust crate ecosystem, this is done manually.
@@ -1033,6 +1256,17 @@ impl BuiltinFunction {
                     date_trunc(precision, datetime.naive_utc()).unwrap(),
                 ))
             }
+            BuiltinFunction::Extract(field, expr) => {
+                let ts = non_null!(expr.eval(record)?);
+                if let DfValue::TimestampTz(tz) = ts {
+                    extract_from_timestamptz(*field, &tz)
+                } else if let DfValue::Time(tm) = ts {
+                    extract_from_time(*field, &tm)
+                } else {
+                    internal!("EXTRACT function input expected to be DfValue::TimestampTz or DfValue::Time. Found {}", ts);
+                }
+                .and_then(|value| value.coerce_to(ty, &DfType::Unknown))
+            }
         }
     }
 }

dataflow-expression/src/eval/builtins.rs

@@ -11,6 +11,7 @@ use std::fmt::{self, Display, Formatter};
 
 pub use eval::builtins::DateTruncPrecision;
 use itertools::Itertools;
+use nom_sql::TimestampField;
 pub use readyset_data::Dialect;
 use readyset_data::{DfType, DfValue};
 use serde::{Deserialize, Serialize};
@@ -112,6 +113,9 @@ pub enum BuiltinFunction {
 
     /// [`date_trunc`](https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC)
     DateTrunc(Expr, Expr),
+
+    /// [`date_trunc`](https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC)
+    Extract(TimestampField, Expr),
 }
 
 impl BuiltinFunction {
@@ -146,6 +150,7 @@ impl BuiltinFunction {
             Least { .. } => "least",
             ArrayToString { .. } => "array_to_string",
             DateTrunc { .. } => "date_trunc",
+            Extract { .. } => "extract",
         }
     }
 }
@@ -244,6 +249,9 @@ impl Display for BuiltinFunction {
             DateTrunc(field, source) => {
                 write!(f, "({}, {})", field, source)
             }
+            Extract(field, expr) => {
+                write!(f, "({} FROM {})", field, expr)
+            }
         }
     }
 }

dataflow-expression/src/lib.rs

@@ -862,6 +862,13 @@ impl Expr {
 
                 Ok(Self::Call { func, ty })
             }
+            AstExpr::Call(FunctionExpr::Extract { field, expr }) => {
+                let expr = Self::lower(*expr, dialect, context)?;
+                let ty = DfType::Numeric { prec: 20, scale: 6 };
+                let func = Box::new(BuiltinFunction::Extract(field, expr));
+
+                Ok(Self::Call { func, ty })
+            }
             AstExpr::Call(call) => internal!(
                 "Unexpected (aggregate?) call node in project expression: {:?}",
                 Sensitive(&call)

dataflow-expression/src/lower.rs

@@ -29,13 +29,8 @@ pub fn parse_lower_eval(
     expr: &str,
     parser_dialect: nom_sql::Dialect,
     expr_dialect: dataflow_expression::Dialect,
-) -> DfValue {
+) -> ReadySetResult<DfValue> {
     let ast = parse_expr(parser_dialect, expr).unwrap();
     let lowered = Expr::lower(ast, expr_dialect, &TestLowerContext).unwrap();
-    match lowered.eval::<DfValue>(&[]) {
-        Ok(res) => res,
-        Err(e) => {
-            panic!("Error evaluating `{expr}`: {e}")
-        }
-    }
+    lowered.eval::<DfValue>(&[])
 }

dataflow-expression/tests/common/mod.rs

@@ -79,7 +79,8 @@ async fn compare_eval(expr: &str, conn: &mut Conn) {
         expr,
         nom_sql::Dialect::MySQL,
         dataflow_expression::Dialect::DEFAULT_MYSQL,
-    );
+    )
+    .unwrap_or_else(|e| panic!("Error evaluating `{expr}`: {e}"));
     assert_eq!(
         our_result, mysql_result,
         "mismatched results for {expr} (left: us, right: mysql)"

dataflow-expression/tests/mysql_oracle.rs

@@ -0,0 +1,18 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 2babad91d07b8955fc644887a196de3f10a2d5906841ecf2266022cab7ada9df # shrinks to field = Century, datetime = 1970-01-01T00:00:00
+cc 74a8620219bfa0ce332d4888bd451e13ebae34411b31cb3c698100aa36b7ed36 # shrinks to field = Second, datetime = 1970-01-01T00:00:01
+cc 7cb714461ce8352e2586db8c9d70dc6473741533cff95ab38f93945034ee9082 # shrinks to field = Isoyear, datetime = 1970-01-01T00:00:00
+cc 0e4282717b1b027255dedda6a563a9e476c859217c68eb85c662c7fd9f222e0a # shrinks to field = Decade, datetime = 1970-01-01T00:00:00
+cc 03b2ab6ce5e6c9a8625203de6bff831de34734dbccce73558e8cc34113459319 # shrinks to field = Epoch, datetime = 1970-01-01T00:00:01
+cc c11828b40b19ecaf688bf5e0da47f3615913ada65a2c103f9d06f06a6c6415d2 # shrinks to field = Dow, datetime = 2035-09-09T00:00:00
+cc d893989c83aaae04070d3d4f3a5268baaee176330458a0a73cb738a64940cc08 # shrinks to field = Julian, datetime = 2038-01-01T00:00:01
+cc f0e758dc44d25cc02f50841e54d5a9ce2808d9d7c3d2bf3d499eba04b893b9f2 # shrinks to field = Julian, datetime = 1970-01-01T00:00:09
+cc ad2e5c4942b266e05f9337bef7b482985b0dd83d986c2be7265a3cb45385931f # shrinks to field = Century, date = -0001-01-01
+cc ae8a16e5a382edc488ecb5d21c31ac5f9c816e47d9b15f1e86bfe1951a218655 # shrinks to field = Isoyear, date = -0001-01-03
+cc 0d521b15cabc8ee12b89ec2532e9b59923158340de74ca00a7580fb78f7c68ba # shrinks to field = Decade, date = -0001-01-02
+cc b9588ba5419366556017de48721258ee0f8568893550e25503c14d637b1744fb # shrinks to field = Century, date = -4290-01-01