diff --git a/datafusion-pg-functions/Cargo.toml b/datafusion-pg-functions/Cargo.toml index b351c77..dfb1676 100644 --- a/datafusion-pg-functions/Cargo.toml +++ b/datafusion-pg-functions/Cargo.toml @@ -107,4 +107,4 @@ tokio = { workspace = true, features = ["macros", "rt", "fs"] } [[test]] name = "sqllogictest" path = "tests/sqllogictest.rs" -required-features = ["math", "string"] +required-features = ["math", "string", "datetime"] diff --git a/datafusion-pg-functions/functions.md b/datafusion-pg-functions/functions.md index ada2e4b..f0219d4 100644 --- a/datafusion-pg-functions/functions.md +++ b/datafusion-pg-functions/functions.md @@ -65,7 +65,7 @@ implementation strategy in DataFusion. | [Bit String](#functions-bitstring) | 9 | 6 | 0 | 0 | 0 | 0 | 0+3 | 3 | | [Pattern Matching](#functions-matching) | 15 | 7 | 2 | 0 | 2 | 0 | 3+2 | 0 | | [Data Type Formatting](#functions-formatting) | 10 | 3 | 0 | 0 | 7 | 0 | 0+0 | 0 | -| [Date/Time](#functions-datetime) | 33 | 11 | 0 | 0 | 21 | 0 | 0+1 | 0 | +| [Date/Time](#functions-datetime) | 33 | 11 | 3 | 0 | 18 | 0 | 0+1 | 0 | | [Enum Support](#functions-enum) | 3 | 0 | 0 | 0 | 0 | 0 | 0+0 | 3 | | [Geometric](#functions-geometry) | 23 | 1 | 0 | 0 | 0 | 0 | 0+0 | 22 | | [Network Address](#functions-net) | 14 | 1 | 0 | 0 | 0 | 0 | 0+0 | 13 | @@ -85,7 +85,7 @@ implementation strategy in DataFusion. | [System Administration](#functions-admin) | 105 | 0 | 2 | 0 | 5 | 0 | 0+0 | 98 | | [Trigger](#functions-trigger) | 3 | 0 | 0 | 0 | 0 | 0 | 0+0 | 3 | | [Event Trigger](#functions-event-triggers) | 5 | 0 | 0 | 0 | 0 | 0 | 0+0 | 5 | -| **TOTAL** | **768** | **198** | **54** | **4** | **71** | **65** | **21** | **371** | +| **TOTAL** | **768** | **198** | **57** | **4** | **68** | **65** | **21** | **371** | --- @@ -386,9 +386,9 @@ implementation strategy in DataFusion. | `date_trunc` | fn | ✅ | — | 4 | truncate timestamp with time zone to specified units in specified time zone | Native DataFusion | | `extract` | both | 🚧 | P2 | 6 | extract field from timestamp with time zone | | | `isfinite` | fn | 🚧 | P2 | 4 | finite timestamp? | | -| `justify_days` | fn | 🚧 | P2 | 1 | promote groups of 30 days to numbers of months | | -| `justify_hours` | fn | 🚧 | P2 | 1 | promote groups of 24 hours to numbers of days | | -| `justify_interval` | fn | 🚧 | P2 | 1 | promote groups of 24 hours to numbers of days and promote groups of 30 days to numbers of months | | +| `justify_days` | fn | 🔧 | — | 1 | promote groups of 30 days to numbers of months | Accepts any Arrow interval unit; returns `Interval(MonthDayNano)`. Time field passed through unchanged. | +| `justify_hours` | fn | 🔧 | — | 1 | promote groups of 24 hours to numbers of days | Accepts any Arrow interval unit; returns `Interval(MonthDayNano)`. Keeps sub-microsecond nanoseconds (PG stores microseconds). | +| `justify_interval` | fn | 🔧 | — | 1 | promote groups of 24 hours to numbers of days and promote groups of 30 days to numbers of months | Accepts any Arrow interval unit; returns `Interval(MonthDayNano)`. Keeps sub-microsecond nanoseconds (PG stores microseconds). | | `localtime` | fn | 🚧 | P2 | 0 | Current date and time, no time zone | | | `localtimestamp` | fn | 🚧 | P2 | 0 | Current timestamp, no time zone | | | `make_date` | fn | ✅ | — | 1 | construct date | Native DataFusion | diff --git a/datafusion-pg-functions/src/datetime.rs b/datafusion-pg-functions/src/datetime.rs index 052b642..aaf8ef7 100644 --- a/datafusion-pg-functions/src/datetime.rs +++ b/datafusion-pg-functions/src/datetime.rs @@ -1,4 +1,14 @@ -//! datetime functions. +//! PostgreSQL date/time functions. +//! +//! This module hosts the UDFs listed in the "Date/Time Functions and +//! Operators" section of [`functions.md`](../functions.md), one file per +//! function: +//! +//! - [`justify_days`] — promote 30-day periods of an interval to months. +//! - [`justify_hours`] — promote 24-hour periods of an interval to days. +//! - [`justify_interval`] — normalize days and time, with sign adjustment. +//! +//! Shared interval input handling lives in the private `interval` module. //! //! See the corresponding section of `functions.md` for the catalog of //! PostgreSQL built-ins in this category and their implementation status. @@ -6,12 +16,26 @@ use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::ScalarUDF; +mod interval; +pub mod justify_days; +pub mod justify_hours; +pub mod justify_interval; + /// Register every PostgreSQL built-in UDF in the datetime category against /// `registry`. /// /// Returns the number of UDFs that were registered. -pub fn register(_registry: &mut dyn FunctionRegistry) -> usize { - let _udfs: Vec = vec![]; - // registry.register_udf(...); - 0 +pub fn register(registry: &mut dyn FunctionRegistry) -> usize { + let udfs: Vec = vec![ + justify_days::create_justify_days_udf(), + justify_hours::create_justify_hours_udf(), + justify_interval::create_justify_interval_udf(), + ]; + + let mut count = 0; + for udf in udfs { + let _ = registry.register_udf(udf.into()); + count += 1; + } + count } diff --git a/datafusion-pg-functions/src/datetime/interval.rs b/datafusion-pg-functions/src/datetime/interval.rs new file mode 100644 index 0000000..8a5bd48 --- /dev/null +++ b/datafusion-pg-functions/src/datetime/interval.rs @@ -0,0 +1,89 @@ +//! Shared helpers for the PostgreSQL `interval` UDFs in this module. +//! +//! PostgreSQL has a single `interval` type made of a month, a day and a time +//! field. Arrow has three interval units; these helpers accept all of them, +//! convert each value losslessly to `(months, days, nanoseconds)`, and always +//! produce `Interval(MonthDayNano)`, matching PostgreSQL's one result type. + +use std::sync::Arc; + +use datafusion::arrow::array::{Array, ArrayRef, AsArray, PrimitiveBuilder, new_null_array}; +use datafusion::arrow::datatypes::{ + DataType, IntervalDayTimeType, IntervalMonthDayNano, IntervalMonthDayNanoType, IntervalUnit, + IntervalYearMonthType, +}; +use datafusion::common::{Result, exec_err}; +use datafusion::logical_expr::{Signature, TypeSignature, Volatility}; + +/// Nanoseconds in one 24-hour day (PostgreSQL `USECS_PER_DAY`, in nanoseconds). +pub(crate) const NANOS_PER_DAY: i64 = 86_400 * 1_000_000_000; + +/// Days in one month, as assumed by PostgreSQL (`DAYS_PER_MONTH`). +pub(crate) const DAYS_PER_MONTH: i32 = 30; + +/// Nanoseconds in one millisecond, for `Interval(DayTime)` input. +const NANOS_PER_MILLI: i64 = 1_000_000; + +/// The error PostgreSQL raises when an interval field overflows. +pub(crate) const OUT_OF_RANGE: &str = "interval out of range"; + +/// Signature for a function taking one interval of any Arrow unit, or a NULL +/// literal. +pub(crate) fn single_interval_signature() -> Signature { + Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Interval(IntervalUnit::MonthDayNano)]), + TypeSignature::Exact(vec![DataType::Interval(IntervalUnit::DayTime)]), + TypeSignature::Exact(vec![DataType::Interval(IntervalUnit::YearMonth)]), + TypeSignature::Exact(vec![DataType::Null]), + ], + Volatility::Immutable, + ) +} + +/// Apply `f(months, days, nanoseconds)` to every non-null element of an +/// interval array, producing an `Interval(MonthDayNano)` array. NULLs stay +/// NULL; an error from `f` fails the whole call. +pub(crate) fn map_interval_array(name: &str, input: &ArrayRef, f: F) -> Result +where + F: Fn(i32, i32, i64) -> Result, +{ + let len = input.len(); + let mut out = PrimitiveBuilder::::with_capacity(len); + match input.data_type() { + DataType::Null => { + return Ok(new_null_array( + &DataType::Interval(IntervalUnit::MonthDayNano), + len, + )); + } + DataType::Interval(IntervalUnit::MonthDayNano) => { + for v in input.as_primitive::().iter() { + match v { + Some(v) => out.append_value(f(v.months, v.days, v.nanoseconds)?), + None => out.append_null(), + } + } + } + DataType::Interval(IntervalUnit::DayTime) => { + for v in input.as_primitive::().iter() { + match v { + Some(v) => { + out.append_value(f(0, v.days, v.milliseconds as i64 * NANOS_PER_MILLI)?) + } + None => out.append_null(), + } + } + } + DataType::Interval(IntervalUnit::YearMonth) => { + for v in input.as_primitive::().iter() { + match v { + Some(months) => out.append_value(f(months, 0, 0)?), + None => out.append_null(), + } + } + } + other => return exec_err!("{name} expects an interval argument, got {other:?}"), + } + Ok(Arc::new(out.finish())) +} diff --git a/datafusion-pg-functions/src/datetime/justify_days.rs b/datafusion-pg-functions/src/datetime/justify_days.rs new file mode 100644 index 0000000..a628534 --- /dev/null +++ b/datafusion-pg-functions/src/datetime/justify_days.rs @@ -0,0 +1,334 @@ +//! PostgreSQL `justify_days(interval)` — promote groups of 30 days to months. +//! +//! Per the [PostgreSQL manual](https://www.postgresql.org/docs/current/functions-datetime.html): +//! +//! > `justify_days ( interval ) → interval` +//! > +//! > Adjust interval, converting 30-day time periods to months. +//! +//! This follows `interval_justify_days` in PostgreSQL's +//! `src/backend/utils/adt/timestamp.c`: +//! +//! 1. Whole 30-day periods are split off the day field, truncating toward +//! zero, and added to the month field. +//! 2. If the month and day fields then have opposite signs, one month is +//! borrowed so both share a sign (e.g. `1 mon -1 days` becomes `29 days`). +//! +//! The time field is never changed, even when it holds more than 24 hours or +//! has a different sign (`justify_days('1 mon -1 hour')` is unchanged). +//! +//! # Input and output types +//! +//! Accepts any Arrow interval unit (`YearMonth`, `DayTime`, `MonthDayNano`) +//! or a NULL literal, and always returns `Interval(MonthDayNano)`. +//! +//! # Deviations from PostgreSQL +//! +//! - **Precision:** the time field is passed through unchanged, including any +//! sub-microsecond nanoseconds PostgreSQL could not store. +//! - **Overflow:** like PostgreSQL, a month field that overflows `i32` is an +//! error (`interval out of range`), which fails the query. + +use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::datatypes::{DataType, IntervalMonthDayNano, IntervalUnit}; +use datafusion::common::{Result, exec_err}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, +}; + +use super::interval::{ + DAYS_PER_MONTH, OUT_OF_RANGE, map_interval_array, single_interval_signature, +}; + +/// Create the PostgreSQL `justify_days(interval)` UDF. +pub fn create_justify_days_udf() -> ScalarUDF { + ScalarUDF::new_from_impl(JustifyDaysUDF::default()) +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JustifyDaysUDF { + signature: Signature, +} + +impl Default for JustifyDaysUDF { + fn default() -> Self { + Self::new() + } +} + +impl JustifyDaysUDF { + pub fn new() -> Self { + Self { + signature: single_interval_signature(), + } + } +} + +impl ScalarUDFImpl for JustifyDaysUDF { + fn name(&self) -> &str { + "justify_days" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Interval(IntervalUnit::MonthDayNano)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; + if args.len() != 1 { + return exec_err!("justify_days requires 1 argument, got {}", args.len()); + } + Ok(ColumnarValue::Array(justify_days_array(&args[0])?)) + } +} + +/// Apply `justify_days` to every element of an interval array. +fn justify_days_array(input: &ArrayRef) -> Result { + map_interval_array("justify_days", input, justify_days) +} + +/// Core of PostgreSQL's `interval_justify_days`, on Arrow interval fields. +fn justify_days(months: i32, days: i32, nanos: i64) -> Result { + // Rust's `/` truncates toward zero, matching C integer division. + let whole_months = days / DAYS_PER_MONTH; + let mut days = days - whole_months * DAYS_PER_MONTH; + let mut months = match months.checked_add(whole_months) { + Some(m) => m, + None => return exec_err!("{OUT_OF_RANGE}"), + }; + + // Borrow one month so that the month and day fields share a sign. This + // moves `months` toward zero, so it cannot overflow. + if months > 0 && days < 0 { + days += DAYS_PER_MONTH; + months -= 1; + } else if months < 0 && days > 0 { + days -= DAYS_PER_MONTH; + months += 1; + } + + Ok(IntervalMonthDayNano::new(months, days, nanos)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use datafusion::arrow::array::{ + Array, AsArray, IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, + }; + use datafusion::arrow::datatypes::{IntervalDayTime, IntervalMonthDayNanoType}; + use datafusion::prelude::SessionContext; + + const NANOS_PER_MICRO: i64 = 1_000; + const NANOS_PER_HOUR: i64 = 3_600 * 1_000_000_000; + + fn mdn(months: i32, days: i32, nanos: i64) -> IntervalMonthDayNano { + IntervalMonthDayNano::new(months, days, nanos) + } + + fn hours(h: i64) -> i64 { + h * NANOS_PER_HOUR + } + + /// Expected values were recorded from PostgreSQL 16.15. + #[test] + fn matches_postgres_results() { + let cases = [ + // Documented example. + ( + mdn(12, 65, 0), + mdn(14, 5, 0), + "1 year 65 days -> 1 year 2 mons 5 days", + ), + (mdn(0, 30, 0), mdn(1, 0, 0), "30 days -> 1 mon"), + (mdn(0, 29, 0), mdn(0, 29, 0), "29 days -> unchanged"), + (mdn(0, 31, 0), mdn(1, 1, 0), "31 days -> 1 mon 1 day"), + (mdn(0, 60, 0), mdn(2, 0, 0), "60 days -> 2 mons"), + (mdn(0, -30, 0), mdn(-1, 0, 0), "-30 days -> -1 mons"), + ( + mdn(0, -31, 0), + mdn(-1, -1, 0), + "-31 days -> -1 mons -1 days", + ), + // Mixed signs between months and days. + (mdn(1, -35, 0), mdn(0, -5, 0), "1 mon -35 days -> -5 days"), + (mdn(-1, 35, 0), mdn(0, 5, 0), "-1 mons +35 days -> 5 days"), + (mdn(1, -1, 0), mdn(0, 29, 0), "1 mon -1 days -> 29 days"), + (mdn(-1, 1, 0), mdn(0, -29, 0), "-1 mons +1 day -> -29 days"), + ( + mdn(12, -400, 0), + mdn(-1, -10, 0), + "1 year -400 days -> -1 mons -10 days", + ), + ( + mdn(3, -91, hours(1)), + mdn(0, -1, hours(1)), + "3 mons -91 days +01:00:00 -> -1 days +01:00:00", + ), + ( + mdn(2, -59, hours(-48)), + mdn(0, 1, hours(-48)), + "2 mons -59 days -48:00:00 -> 1 day -48:00:00", + ), + ( + mdn(14, -3, hours(100)), + mdn(13, 27, hours(100)), + "1 year 2 mons -3 days 100:00:00 -> 1 year 1 mon 27 days 100:00:00", + ), + // The time field is never touched. + ( + mdn(0, 0, hours(50) + 600 * 1_000_000_000), + mdn(0, 0, hours(50) + 600 * 1_000_000_000), + "50:10:00 -> unchanged", + ), + ( + mdn(1, 0, hours(-1)), + mdn(1, 0, hours(-1)), + "1 mon -01:00:00 -> unchanged", + ), + ( + mdn(0, 35, hours(50)), + mdn(1, 5, hours(50)), + "35 days 50:00:00 -> 1 mon 5 days 50:00:00", + ), + ( + mdn(0, 1, hours(-1)), + mdn(0, 1, hours(-1)), + "1 day -01:00:00 -> unchanged", + ), + ( + mdn(1, 0, -NANOS_PER_MICRO), + mdn(1, 0, -NANOS_PER_MICRO), + "1 mon -00:00:00.000001 -> unchanged", + ), + (mdn(0, 0, 0), mdn(0, 0, 0), "00:00:00 -> 00:00:00"), + ]; + for (input, expected, label) in cases { + let got = justify_days(input.months, input.days, input.nanoseconds).unwrap(); + assert_eq!(got, expected, "{label}"); + } + } + + #[test] + fn extreme_day_values() { + // PostgreSQL 16.15 results. + assert_eq!(justify_days(0, i32::MAX, 0).unwrap(), mdn(71_582_788, 7, 0)); + assert_eq!( + justify_days(0, i32::MIN, 0).unwrap(), + mdn(-71_582_788, -8, 0) + ); + } + + #[test] + fn largest_month_values_that_fit() { + // PostgreSQL 16.15 results. + assert_eq!( + justify_days(i32::MAX - 1, 30, 0).unwrap(), + mdn(i32::MAX, 0, 0) + ); + assert_eq!( + justify_days(i32::MAX, -1, 0).unwrap(), + mdn(i32::MAX - 1, 29, 0) + ); + assert_eq!( + justify_days(i32::MIN + 1, -30, 0).unwrap(), + mdn(i32::MIN, 0, 0) + ); + assert_eq!( + justify_days(i32::MIN, 1, 0).unwrap(), + mdn(i32::MIN + 1, -29, 0) + ); + assert_eq!( + justify_days(i32::MAX, 29, hours(24)).unwrap(), + mdn(i32::MAX, 29, hours(24)) + ); + } + + #[test] + fn overflow_is_an_error() { + // PostgreSQL: justify_days(interval '2147483647 mons 30 days') + // -> ERROR: interval out of range + let err = justify_days(i32::MAX, 30, 0).unwrap_err(); + assert!(err.to_string().contains(OUT_OF_RANGE), "{err}"); + let err = justify_days(i32::MIN, -30, 0).unwrap_err(); + assert!(err.to_string().contains(OUT_OF_RANGE), "{err}"); + } + + #[test] + fn month_day_nano_batch_with_nulls() { + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(mdn(0, 31, 0)), + None, + Some(mdn(1, -1, hours(5))), + ])); + let out = justify_days_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.len(), 3); + assert_eq!(out.value(0), mdn(1, 1, 0)); + assert!(out.is_null(1)); + assert_eq!(out.value(2), mdn(0, 29, hours(5))); + } + + #[test] + fn empty_input() { + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(Vec::< + Option, + >::new())); + let out = justify_days_array(&input).unwrap(); + assert_eq!(out.len(), 0); + assert_eq!( + out.data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + } + + #[test] + fn day_time_input() { + let input: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![ + Some(IntervalDayTime::new(65, 3_600_000)), + None, + ])); + let out = justify_days_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.value(0), mdn(2, 5, hours(1))); + assert!(out.is_null(1)); + } + + #[test] + fn year_month_input_is_unchanged() { + let input: ArrayRef = Arc::new(IntervalYearMonthArray::from(vec![Some(-14), None])); + let out = justify_days_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.value(0), mdn(-14, 0, 0)); + assert!(out.is_null(1)); + } + + #[tokio::test] + async fn sql_null_and_literal() { + let ctx = SessionContext::new(); + ctx.register_udf(create_justify_days_udf()); + let batches = ctx + .sql("SELECT justify_days(NULL), justify_days(INTERVAL '1 year 65 days')") + .await + .unwrap() + .collect() + .await + .unwrap(); + let null_col = batches[0].column(0); + assert_eq!( + null_col.data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + assert!(null_col.is_null(0)); + let col = batches[0] + .column(1) + .as_primitive::(); + assert_eq!(col.value(0), mdn(14, 5, 0)); + } +} diff --git a/datafusion-pg-functions/src/datetime/justify_hours.rs b/datafusion-pg-functions/src/datetime/justify_hours.rs new file mode 100644 index 0000000..11505e5 --- /dev/null +++ b/datafusion-pg-functions/src/datetime/justify_hours.rs @@ -0,0 +1,337 @@ +//! PostgreSQL `justify_hours(interval)` — promote groups of 24 hours to days. +//! +//! Per the [PostgreSQL manual](https://www.postgresql.org/docs/current/functions-datetime.html): +//! +//! > `justify_hours ( interval ) → interval` +//! > +//! > Adjust interval, converting 24-hour time periods to days. +//! +//! This follows `interval_justify_hours` in PostgreSQL's +//! `src/backend/utils/adt/timestamp.c`: +//! +//! 1. Whole days are split off the time field, truncating toward zero, and +//! added to the day field. The month field is never changed. +//! 2. If the day and time fields then have opposite signs, one day is +//! borrowed so both share a sign (e.g. `1 day -01:00:00` becomes +//! `23:00:00`). +//! +//! # Input and output types +//! +//! Accepts any Arrow interval unit (`YearMonth`, `DayTime`, `MonthDayNano`) +//! or a NULL literal, and always returns `Interval(MonthDayNano)`, matching +//! PostgreSQL's single `interval` result type. `DayTime` milliseconds and +//! `YearMonth` months convert to `MonthDayNano` losslessly. +//! +//! # Deviations from PostgreSQL +//! +//! - **Precision:** PostgreSQL stores the time field in microseconds and +//! rounds sub-microsecond input when parsing; Arrow `MonthDayNano` stores +//! nanoseconds, and DataFusion can produce sub-microsecond values. These +//! are kept and justified exactly, so the result always equals the input +//! (with 24-hour days). The difference is only visible for such input: +//! `justify_hours(INTERVAL '1 day -1 nanosecond')` returns +//! `23:59:59.999999999`, which arrow-pg truncates to `23:59:59.999999` on +//! the wire, whereas PostgreSQL would have rounded the input to `1 day` +//! and returned `1 day`. +//! - **Overflow:** like PostgreSQL, a day field that overflows `i32` is an +//! error (`interval out of range`), which fails the query. + +use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::datatypes::{DataType, IntervalMonthDayNano, IntervalUnit}; +use datafusion::common::{Result, exec_err}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, +}; + +use super::interval::{NANOS_PER_DAY, OUT_OF_RANGE, map_interval_array, single_interval_signature}; + +/// Create the PostgreSQL `justify_hours(interval)` UDF. +pub fn create_justify_hours_udf() -> ScalarUDF { + ScalarUDF::new_from_impl(JustifyHoursUDF::default()) +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JustifyHoursUDF { + signature: Signature, +} + +impl Default for JustifyHoursUDF { + fn default() -> Self { + Self::new() + } +} + +impl JustifyHoursUDF { + pub fn new() -> Self { + Self { + signature: single_interval_signature(), + } + } +} + +impl ScalarUDFImpl for JustifyHoursUDF { + fn name(&self) -> &str { + "justify_hours" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Interval(IntervalUnit::MonthDayNano)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; + if args.len() != 1 { + return exec_err!("justify_hours requires 1 argument, got {}", args.len()); + } + Ok(ColumnarValue::Array(justify_hours_array(&args[0])?)) + } +} + +/// Apply `justify_hours` to every element of an interval array. +fn justify_hours_array(input: &ArrayRef) -> Result { + map_interval_array("justify_hours", input, justify_hours) +} + +/// Core of PostgreSQL's `interval_justify_hours`, on Arrow interval fields. +fn justify_hours(months: i32, days: i32, nanos: i64) -> Result { + // Split whole days off the time field. Rust's `/` and `%` truncate + // toward zero, matching PostgreSQL's TMODULO macro. + let whole_days = nanos / NANOS_PER_DAY; + let mut time = nanos % NANOS_PER_DAY; + + // |nanos| / NANOS_PER_DAY is at most ~106,752, so this cast is lossless. + let mut days = match days.checked_add(whole_days as i32) { + Some(d) => d, + None => return exec_err!("{OUT_OF_RANGE}"), + }; + + // Borrow one day so that the day and time fields share a sign. This moves + // `days` toward zero, so it cannot overflow. + if days > 0 && time < 0 { + time += NANOS_PER_DAY; + days -= 1; + } else if days < 0 && time > 0 { + time -= NANOS_PER_DAY; + days += 1; + } + + Ok(IntervalMonthDayNano::new(months, days, time)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use datafusion::arrow::array::{ + Array, AsArray, IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, + }; + use datafusion::arrow::datatypes::{IntervalDayTime, IntervalMonthDayNanoType}; + use datafusion::prelude::SessionContext; + + const NANOS_PER_MICRO: i64 = 1_000; + const NANOS_PER_SEC: i64 = 1_000_000_000; + const NANOS_PER_HOUR: i64 = 3_600 * NANOS_PER_SEC; + + fn mdn(months: i32, days: i32, nanos: i64) -> IntervalMonthDayNano { + IntervalMonthDayNano::new(months, days, nanos) + } + + fn hours(h: i64) -> i64 { + h * NANOS_PER_HOUR + } + + /// Expected values were recorded from PostgreSQL 16.15. + #[test] + fn matches_postgres_results() { + let cases = [ + // (input, expected, PostgreSQL rendering of input -> result) + ( + mdn(0, 0, hours(50) + 10 * 60 * NANOS_PER_SEC), + mdn(0, 2, hours(2) + 10 * 60 * NANOS_PER_SEC), + "50:10:00 -> 2 days 02:10:00", + ), + (mdn(0, 0, hours(24)), mdn(0, 1, 0), "24:00:00 -> 1 day"), + ( + mdn(0, 0, hours(24) - NANOS_PER_MICRO), + mdn(0, 0, hours(24) - NANOS_PER_MICRO), + "23:59:59.999999 -> unchanged", + ), + ( + mdn(0, 0, hours(25)), + mdn(0, 1, hours(1)), + "25:00:00 -> 1 day 01:00:00", + ), + ( + mdn(0, 0, hours(-25)), + mdn(0, -1, hours(-1)), + "-25:00:00 -> -1 days -01:00:00", + ), + ( + mdn(0, 1, hours(-25)), + mdn(0, 0, hours(-1)), + "1 day -25:00:00 -> -01:00:00", + ), + ( + mdn(0, -1, hours(25)), + mdn(0, 0, hours(1)), + "-1 days +25:00:00 -> 01:00:00", + ), + ( + mdn(0, 1, hours(-1)), + mdn(0, 0, hours(23)), + "1 day -01:00:00 -> 23:00:00", + ), + ( + mdn(2, 0, hours(49)), + mdn(2, 2, hours(1)), + "2 mons 49:00:00 -> 2 mons 2 days 01:00:00", + ), + ( + mdn(1, -1, hours(1)), + mdn(1, 0, hours(-23)), + "1 mon -1 days +01:00:00 -> 1 mon -23:00:00", + ), + ( + mdn(0, 0, hours(48) + NANOS_PER_MICRO), + mdn(0, 2, NANOS_PER_MICRO), + "48:00:00.000001 -> 2 days 00:00:00.000001", + ), + ( + mdn(0, 1000, hours(1000)), + mdn(0, 1041, hours(16)), + "1000 days 1000:00:00 -> 1041 days 16:00:00", + ), + (mdn(1, -1, 0), mdn(1, -1, 0), "1 mon -1 days -> unchanged"), + (mdn(0, 0, 0), mdn(0, 0, 0), "00:00:00 -> 00:00:00"), + ]; + for (input, expected, label) in cases { + let got = justify_hours(input.months, input.days, input.nanoseconds).unwrap(); + assert_eq!(got, expected, "{label}"); + } + } + + #[test] + fn overflow_is_an_error() { + // PostgreSQL: justify_hours(interval '2147483647 days 24 hours') + // -> ERROR: interval out of range + let err = justify_hours(0, i32::MAX, hours(24)).unwrap_err(); + assert!(err.to_string().contains("interval out of range"), "{err}"); + + let err = justify_hours(0, i32::MIN, hours(-24)).unwrap_err(); + assert!(err.to_string().contains("interval out of range"), "{err}"); + } + + #[test] + fn largest_day_values_that_fit() { + assert_eq!( + justify_hours(0, i32::MAX - 1, hours(24)).unwrap(), + mdn(0, i32::MAX, 0) + ); + assert_eq!( + justify_hours(0, i32::MAX, hours(23)).unwrap(), + mdn(0, i32::MAX, hours(23)) + ); + } + + #[test] + fn keeps_sub_microsecond_nanos() { + assert_eq!( + justify_hours(0, 0, hours(24) + 500).unwrap(), + mdn(0, 1, 500) + ); + } + + #[test] + fn month_day_nano_batch_with_nulls() { + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(mdn(0, 0, hours(25))), + None, + Some(mdn(0, 1, hours(-1))), + Some(mdn(3, 0, hours(-49))), + ])); + let out = justify_hours_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.len(), 4); + assert_eq!(out.value(0), mdn(0, 1, hours(1))); + assert!(out.is_null(1)); + assert_eq!(out.value(2), mdn(0, 0, hours(23))); + assert_eq!(out.value(3), mdn(3, -2, hours(-1))); + } + + #[test] + fn empty_input() { + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(Vec::< + Option, + >::new())); + let out = justify_hours_array(&input).unwrap(); + assert_eq!(out.len(), 0); + assert_eq!( + out.data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + } + + #[test] + fn day_time_input() { + let input: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![ + Some(IntervalDayTime::new(0, 25 * 3_600_000)), + Some(IntervalDayTime::new(1, -3_600_000)), + None, + ])); + let out = justify_hours_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.value(0), mdn(0, 1, hours(1))); + assert_eq!(out.value(1), mdn(0, 0, hours(23))); + assert!(out.is_null(2)); + } + + #[test] + fn year_month_input_is_unchanged() { + let input: ArrayRef = Arc::new(IntervalYearMonthArray::from(vec![Some(14), None])); + let out = justify_hours_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.value(0), mdn(14, 0, 0)); + assert!(out.is_null(1)); + } + + #[tokio::test] + async fn sql_null_literal() { + let ctx = SessionContext::new(); + ctx.register_udf(create_justify_hours_udf()); + let batches = ctx + .sql("SELECT justify_hours(NULL)") + .await + .unwrap() + .collect() + .await + .unwrap(); + let col = batches[0].column(0); + assert_eq!( + col.data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + assert!(col.is_null(0)); + } + + #[tokio::test] + async fn sql_interval_literal() { + let ctx = SessionContext::new(); + ctx.register_udf(create_justify_hours_udf()); + let batches = ctx + .sql("SELECT justify_hours(INTERVAL '50 hours 10 minutes')") + .await + .unwrap() + .collect() + .await + .unwrap(); + let col = batches[0] + .column(0) + .as_primitive::(); + assert_eq!(col.value(0), mdn(0, 2, hours(2) + 10 * 60 * NANOS_PER_SEC)); + } +} diff --git a/datafusion-pg-functions/src/datetime/justify_interval.rs b/datafusion-pg-functions/src/datetime/justify_interval.rs new file mode 100644 index 0000000..46a1c21 --- /dev/null +++ b/datafusion-pg-functions/src/datetime/justify_interval.rs @@ -0,0 +1,443 @@ +//! PostgreSQL `justify_interval(interval)` — normalize days and time, with +//! sign adjustment. +//! +//! Per the [PostgreSQL manual](https://www.postgresql.org/docs/current/functions-datetime.html): +//! +//! > `justify_interval ( interval ) → interval` +//! > +//! > Adjust interval using `justify_days` and `justify_hours`, with additional +//! > sign adjustments. +//! +//! This is not simply `justify_days(justify_hours(x))`. It follows +//! `interval_justify_interval` in PostgreSQL's +//! `src/backend/utils/adt/timestamp.c` step by step: +//! +//! 1. If days and time share a sign, whole 30-day periods are first moved +//! from days to months (PostgreSQL does this to avoid overflow). +//! 2. Whole days are split off the time field (truncating toward zero) and +//! added to days. +//! 3. Whole 30-day periods are moved from days to months. +//! 4. If months and the remaining days have opposite signs — or days are +//! zero and time has the opposite sign — one month is borrowed. +//! 5. If days and time have opposite signs, one day is borrowed. +//! +//! The result therefore has all three fields sharing one sign, e.g. +//! `1 mon -00:00:00.000001` becomes `29 days 23:59:59.999999`. +//! +//! # Input and output types +//! +//! Accepts any Arrow interval unit (`YearMonth`, `DayTime`, `MonthDayNano`) +//! or a NULL literal, and always returns `Interval(MonthDayNano)`. +//! +//! # Deviations from PostgreSQL +//! +//! - **Precision:** PostgreSQL stores the time field in microseconds and +//! rounds sub-microsecond input when parsing; Arrow `MonthDayNano` stores +//! nanoseconds. Sub-microsecond values are kept and justified exactly +//! rather than rounded first, so e.g. +//! `justify_interval(INTERVAL '1 month -1 nanosecond')` returns +//! `29 days 23:59:59.999999999`, where PostgreSQL would return `1 mon` +//! for the rounded input. See `justify_hours` for details. +//! - **Overflow:** like PostgreSQL, a month field that overflows `i32` is an +//! error (`interval out of range`), which fails the query. + +use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::datatypes::{DataType, IntervalMonthDayNano, IntervalUnit}; +use datafusion::common::{Result, exec_err}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, +}; + +use super::interval::{ + DAYS_PER_MONTH, NANOS_PER_DAY, OUT_OF_RANGE, map_interval_array, single_interval_signature, +}; + +/// Create the PostgreSQL `justify_interval(interval)` UDF. +pub fn create_justify_interval_udf() -> ScalarUDF { + ScalarUDF::new_from_impl(JustifyIntervalUDF::default()) +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JustifyIntervalUDF { + signature: Signature, +} + +impl Default for JustifyIntervalUDF { + fn default() -> Self { + Self::new() + } +} + +impl JustifyIntervalUDF { + pub fn new() -> Self { + Self { + signature: single_interval_signature(), + } + } +} + +impl ScalarUDFImpl for JustifyIntervalUDF { + fn name(&self) -> &str { + "justify_interval" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Interval(IntervalUnit::MonthDayNano)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = ColumnarValue::values_to_arrays(&args.args)?; + if args.len() != 1 { + return exec_err!("justify_interval requires 1 argument, got {}", args.len()); + } + Ok(ColumnarValue::Array(justify_interval_array(&args[0])?)) + } +} + +/// Apply `justify_interval` to every element of an interval array. +fn justify_interval_array(input: &ArrayRef) -> Result { + map_interval_array("justify_interval", input, justify_interval) +} + +/// Move whole 30-day periods from `days` into `months`, as PostgreSQL does. +fn days_to_months(months: i32, days: i32) -> Result<(i32, i32)> { + // Rust's `/` truncates toward zero, matching C integer division. + let whole_months = days / DAYS_PER_MONTH; + match months.checked_add(whole_months) { + Some(m) => Ok((m, days - whole_months * DAYS_PER_MONTH)), + None => exec_err!("{OUT_OF_RANGE}"), + } +} + +/// Core of PostgreSQL's `interval_justify_interval`, on Arrow interval fields. +fn justify_interval(months: i32, days: i32, nanos: i64) -> Result { + let (mut months, mut days, mut time) = (months, days, nanos); + + // 1. Pre-justify days when days and time share a sign. + if (days > 0 && time > 0) || (days < 0 && time < 0) { + (months, days) = days_to_months(months, days)?; + } + + // 2. Split whole days off the time field (truncating toward zero, like + // PostgreSQL's TMODULO). |nanos| / NANOS_PER_DAY is at most ~106,752, so + // the cast is lossless. PostgreSQL shows this addition cannot overflow + // (after step 1, |days| < 30 or days and time differ in sign); it is + // still checked here rather than assumed. + let whole_days = time / NANOS_PER_DAY; + time %= NANOS_PER_DAY; + days = match days.checked_add(whole_days as i32) { + Some(d) => d, + None => return exec_err!("{OUT_OF_RANGE}"), + }; + + // 3. Move whole 30-day periods into months. + (months, days) = days_to_months(months, days)?; + + // 4. Borrow a month so months agree in sign with days, or with time when + // days are zero. This moves `months` toward zero, so it cannot overflow. + if months > 0 && (days < 0 || (days == 0 && time < 0)) { + days += DAYS_PER_MONTH; + months -= 1; + } else if months < 0 && (days > 0 || (days == 0 && time > 0)) { + days -= DAYS_PER_MONTH; + months += 1; + } + + // 5. Borrow a day so days agree in sign with time. + if days > 0 && time < 0 { + time += NANOS_PER_DAY; + days -= 1; + } else if days < 0 && time > 0 { + time -= NANOS_PER_DAY; + days += 1; + } + + Ok(IntervalMonthDayNano::new(months, days, time)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use datafusion::arrow::array::{ + Array, AsArray, IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, + }; + use datafusion::arrow::datatypes::{IntervalDayTime, IntervalMonthDayNanoType}; + use datafusion::prelude::SessionContext; + + const NANOS_PER_MICRO: i64 = 1_000; + const NANOS_PER_HOUR: i64 = 3_600 * 1_000_000_000; + + fn mdn(months: i32, days: i32, nanos: i64) -> IntervalMonthDayNano { + IntervalMonthDayNano::new(months, days, nanos) + } + + fn hours(h: i64) -> i64 { + h * NANOS_PER_HOUR + } + + /// Expected values were recorded from PostgreSQL 16.15. + #[test] + fn matches_postgres_results() { + let cases = [ + // Documented example. + ( + mdn(1, 0, hours(-1)), + mdn(0, 29, hours(23)), + "1 mon -01:00:00 -> 29 days 23:00:00", + ), + ( + mdn(-1, 0, hours(1)), + mdn(0, -29, hours(-23)), + "-1 mons +01:00:00 -> -29 days -23:00:00", + ), + ( + mdn(12, 65, 0), + mdn(14, 5, 0), + "1 year 65 days -> 1 year 2 mons 5 days", + ), + ( + mdn(0, 0, hours(50) + 600 * 1_000_000_000), + mdn(0, 2, hours(2) + 600 * 1_000_000_000), + "50:10:00 -> 2 days 02:10:00", + ), + ( + mdn(0, 35, hours(50)), + mdn(1, 7, hours(2)), + "35 days 50:00:00 -> 1 mon 7 days 02:00:00", + ), + ( + mdn(0, -35, hours(-50)), + mdn(-1, -7, hours(-2)), + "-35 days -50:00:00 -> -1 mons -7 days -02:00:00", + ), + ( + mdn(0, 29, hours(47)), + mdn(1, 0, hours(23)), + "29 days 47:00:00 -> 1 mon 23:00:00", + ), + ( + mdn(0, -29, hours(-47)), + mdn(-1, 0, hours(-23)), + "-29 days -47:00:00 -> -1 mons -23:00:00", + ), + ( + mdn(1, 29, hours(24)), + mdn(2, 0, 0), + "1 mon 29 days 24:00:00 -> 2 mons", + ), + ( + mdn(-1, -29, hours(-24)), + mdn(-2, 0, 0), + "-1 mons -29 days -24:00:00 -> -2 mons", + ), + // Mixed signs. + (mdn(1, -31, 0), mdn(0, -1, 0), "1 mon -31 days -> -1 days"), + (mdn(-1, 31, 0), mdn(0, 1, 0), "-1 mons +31 days -> 1 day"), + ( + mdn(1, 1, hours(-25)), + mdn(0, 29, hours(23)), + "1 mon 1 day -25:00:00 -> 29 days 23:00:00", + ), + ( + mdn(0, 1, hours(-1)), + mdn(0, 0, hours(23)), + "1 day -01:00:00 -> 23:00:00", + ), + ( + mdn(0, -1, hours(1)), + mdn(0, 0, hours(-23)), + "-1 days +01:00:00 -> -23:00:00", + ), + ( + mdn(1, -1, hours(1)), + mdn(0, 29, hours(1)), + "1 mon -1 days +01:00:00 -> 29 days 01:00:00", + ), + ( + mdn(2, -59, hours(-48)), + mdn(0, -1, 0), + "2 mons -59 days -48:00:00 -> -1 days", + ), + ( + mdn(3, -91, hours(1)), + mdn(0, 0, hours(-23)), + "3 mons -91 days +01:00:00 -> -23:00:00", + ), + ( + mdn(14, -3, hours(100)), + mdn(14, 1, hours(4)), + "1 year 2 mons -3 days 100:00:00 -> 1 year 2 mons 1 day 04:00:00", + ), + // Month borrow driven by the time field when days are zero. + ( + mdn(1, 0, -NANOS_PER_MICRO), + mdn(0, 29, hours(24) - NANOS_PER_MICRO), + "1 mon -00:00:00.000001 -> 29 days 23:59:59.999999", + ), + ( + mdn(-1, 0, NANOS_PER_MICRO), + mdn(0, -29, -(hours(24) - NANOS_PER_MICRO)), + "-1 mons +00:00:00.000001 -> -29 days -23:59:59.999999", + ), + (mdn(0, 0, 0), mdn(0, 0, 0), "00:00:00 -> 00:00:00"), + ]; + for (input, expected, label) in cases { + let got = justify_interval(input.months, input.days, input.nanoseconds).unwrap(); + assert_eq!(got, expected, "{label}"); + } + } + + #[test] + fn extreme_values() { + // PostgreSQL 16.15 results. + assert_eq!( + justify_interval(0, i32::MAX, 0).unwrap(), + mdn(71_582_788, 7, 0) + ); + assert_eq!( + justify_interval(0, i32::MIN, 0).unwrap(), + mdn(-71_582_788, -8, 0) + ); + assert_eq!( + justify_interval(0, i32::MAX, hours(1)).unwrap(), + mdn(71_582_788, 7, hours(1)) + ); + assert_eq!( + justify_interval(0, i32::MAX, hours(-1)).unwrap(), + mdn(71_582_788, 6, hours(23)) + ); + assert_eq!( + justify_interval(0, i32::MIN, hours(1)).unwrap(), + mdn(-71_582_788, -7, hours(-23)) + ); + // Largest whole-hour time Arrow can hold: 2562047 hours. + assert_eq!( + justify_interval(0, 0, hours(2_562_047)).unwrap(), + mdn(3_558, 11, hours(23)) + ); + assert_eq!( + justify_interval(0, 0, hours(-2_562_047)).unwrap(), + mdn(-3_558, -11, hours(-23)) + ); + assert_eq!( + justify_interval(i32::MAX - 1, 30, 0).unwrap(), + mdn(i32::MAX, 0, 0) + ); + assert_eq!( + justify_interval(i32::MAX, -1, 0).unwrap(), + mdn(i32::MAX - 1, 29, 0) + ); + assert_eq!( + justify_interval(i32::MIN + 1, -30, 0).unwrap(), + mdn(i32::MIN, 0, 0) + ); + assert_eq!( + justify_interval(i32::MIN, 1, 0).unwrap(), + mdn(i32::MIN + 1, -29, 0) + ); + assert_eq!( + justify_interval(i32::MAX, 0, hours(-1)).unwrap(), + mdn(i32::MAX - 1, 29, hours(23)) + ); + assert_eq!( + justify_interval(i32::MIN, 0, hours(1)).unwrap(), + mdn(i32::MIN + 1, -29, hours(-23)) + ); + assert_eq!( + justify_interval(i32::MAX, 29, 0).unwrap(), + mdn(i32::MAX, 29, 0) + ); + } + + #[test] + fn overflow_is_an_error() { + // PostgreSQL: each of these raises ERROR: interval out of range. + for (m, d, n) in [ + (i32::MAX, 30, 0), + (i32::MIN, -30, 0), + (i32::MAX, 29, hours(24)), + (i32::MIN, -29, hours(-24)), + ] { + let err = justify_interval(m, d, n).unwrap_err(); + assert!(err.to_string().contains(OUT_OF_RANGE), "{m} {d} {n}: {err}"); + } + } + + #[test] + fn month_day_nano_batch_with_nulls() { + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(mdn(1, 0, hours(-1))), + None, + Some(mdn(0, 29, hours(47))), + ])); + let out = justify_interval_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.len(), 3); + assert_eq!(out.value(0), mdn(0, 29, hours(23))); + assert!(out.is_null(1)); + assert_eq!(out.value(2), mdn(1, 0, hours(23))); + } + + #[test] + fn empty_input() { + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(Vec::< + Option, + >::new())); + let out = justify_interval_array(&input).unwrap(); + assert_eq!(out.len(), 0); + assert_eq!( + out.data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + } + + #[test] + fn day_time_input() { + let input: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![ + Some(IntervalDayTime::new(29, 47 * 3_600_000)), + None, + ])); + let out = justify_interval_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.value(0), mdn(1, 0, hours(23))); + assert!(out.is_null(1)); + } + + #[test] + fn year_month_input_is_unchanged() { + let input: ArrayRef = Arc::new(IntervalYearMonthArray::from(vec![Some(14), None])); + let out = justify_interval_array(&input).unwrap(); + let out = out.as_primitive::(); + assert_eq!(out.value(0), mdn(14, 0, 0)); + assert!(out.is_null(1)); + } + + #[tokio::test] + async fn sql_null_and_literal() { + let ctx = SessionContext::new(); + ctx.register_udf(create_justify_interval_udf()); + let batches = ctx + .sql("SELECT justify_interval(NULL), justify_interval(INTERVAL '1 month -1 hour')") + .await + .unwrap() + .collect() + .await + .unwrap(); + let null_col = batches[0].column(0); + assert_eq!( + null_col.data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + assert!(null_col.is_null(0)); + let col = batches[0] + .column(1) + .as_primitive::(); + assert_eq!(col.value(0), mdn(0, 29, hours(23))); + } +} diff --git a/datafusion-pg-functions/tests/sqllogictest/datetime.slt b/datafusion-pg-functions/tests/sqllogictest/datetime.slt new file mode 100644 index 0000000..f41c10b --- /dev/null +++ b/datafusion-pg-functions/tests/sqllogictest/datetime.slt @@ -0,0 +1,341 @@ +# Integration tests for the PostgreSQL date/time functions implemented in +# `datafusion-pg-functions/src/datetime/`. +# +# Each section exercises one UDF via the full SQL -> plan -> execute path +# against a `SessionContext` that has had `register_all` applied. +# +# Conventions: +# * Expected values were checked against PostgreSQL 16.15. Intervals are +# rendered in Arrow's display format (e.g. `2 days 2 hours 10 mins`), +# not PostgreSQL's (`2 days 02:10:00`); the PostgreSQL result is noted +# above each case. +# * NULL inputs must propagate to NULL. +# * Only `Interval(MonthDayNano)` is reachable from SQL here: DataFusion +# cannot CAST to `Interval(DayTime)` or `Interval(YearMonth)`. Those input +# types are covered by each function's unit tests. +# * Arrow displays months without folding them into years, so PostgreSQL's +# `1 year 2 mons` appears here as `14 mons`. + +# ============================================================================ +# justify_hours(interval) -- promote 24-hour periods to days +# ============================================================================ + +# PostgreSQL: 2 days 02:10:00 +query ? +SELECT justify_hours(INTERVAL '50 hours 10 minutes') +---- +2 days 2 hours 10 mins + +# PostgreSQL: 1 day +query ? +SELECT justify_hours(INTERVAL '24 hours') +---- +1 days + +# PostgreSQL: 23:59:59.999999 (unchanged) +query ? +SELECT justify_hours(INTERVAL '23 hours 59 minutes 59.999999 seconds') +---- +23 hours 59 mins 59.999999000 secs + +# PostgreSQL: 1 day 01:00:00 +query ? +SELECT justify_hours(INTERVAL '25 hours') +---- +1 days 1 hours + +# PostgreSQL: -1 days -01:00:00 +query ? +SELECT justify_hours(INTERVAL '-25 hours') +---- +-1 days -1 hours + +# Mixed signs. PostgreSQL: -01:00:00 +query ? +SELECT justify_hours(INTERVAL '1 day -25 hours') +---- +-1 hours + +# Mixed signs. PostgreSQL: 01:00:00 +query ? +SELECT justify_hours(INTERVAL '-1 day 25 hours') +---- +1 hours + +# Sign borrow. PostgreSQL: 23:00:00 +query ? +SELECT justify_hours(INTERVAL '1 day -1 hour') +---- +23 hours + +# Months are never changed. PostgreSQL: 2 mons 2 days 01:00:00 +query ? +SELECT justify_hours(INTERVAL '2 months 49 hours') +---- +2 mons 2 days 1 hours + +# Days and time are reconciled, months are not. PostgreSQL: 1 mon -23:00:00 +query ? +SELECT justify_hours(INTERVAL '1 month -1 day 1 hour') +---- +1 mons -23 hours + +# PostgreSQL: 2 days 00:00:00.000001 +query ? +SELECT justify_hours(INTERVAL '48 hours 0.000001 seconds') +---- +2 days 0.000001000 secs + +# PostgreSQL: 1041 days 16:00:00 +query ? +SELECT justify_hours(INTERVAL '1000 days 1000 hours') +---- +1041 days 16 hours + +# PostgreSQL: 00:00:00 +query ? +SELECT justify_hours(INTERVAL '0 seconds') +---- +0 secs + +# Sub-microsecond time is kept and justified exactly. PostgreSQL rounds the +# input to whole microseconds when parsing, so it returns 1 day here. +query ? +SELECT justify_hours(INTERVAL '1 day -1 nanosecond') +---- +23 hours 59 mins 59.999999999 secs + +# NULL propagates +query ? +SELECT justify_hours(CAST(NULL AS INTERVAL)) +---- +NULL + +query ? +SELECT justify_hours(NULL) +---- +NULL + +# Row-wise vectorized batch +query ? +SELECT justify_hours(v) FROM (VALUES + (INTERVAL '25 hours'), + (CAST(NULL AS INTERVAL)), + (INTERVAL '1 day -1 hour'), + (INTERVAL '3 months -49 hours') +) AS t(v) +---- +1 days 1 hours +NULL +23 hours +3 mons -2 days -1 hours + +# Overflow. PostgreSQL: ERROR: interval out of range +query error interval out of range +SELECT justify_hours(INTERVAL '2147483647 days 24 hours') + +# ============================================================================ +# justify_days(interval) -- promote 30-day periods to months +# ============================================================================ + +# Documented example. PostgreSQL: 1 year 2 mons 5 days +query ? +SELECT justify_days(INTERVAL '1 year 65 days') +---- +14 mons 5 days + +# PostgreSQL: 1 mon +query ? +SELECT justify_days(INTERVAL '30 days') +---- +1 mons + +# PostgreSQL: 29 days (unchanged) +query ? +SELECT justify_days(INTERVAL '29 days') +---- +29 days + +# PostgreSQL: -1 mons -1 days +query ? +SELECT justify_days(INTERVAL '-31 days') +---- +-1 mons -1 days + +# Mixed signs. PostgreSQL: 29 days +query ? +SELECT justify_days(INTERVAL '1 month -1 day') +---- +29 days + +# Mixed signs. PostgreSQL: -29 days +query ? +SELECT justify_days(INTERVAL '-1 month 1 day') +---- +-29 days + +# PostgreSQL: -1 mons -10 days +query ? +SELECT justify_days(INTERVAL '1 year -400 days') +---- +-1 mons -10 days + +# The time field is never changed. PostgreSQL: 1 mon 5 days 50:00:00 +query ? +SELECT justify_days(INTERVAL '35 days 50 hours') +---- +1 mons 5 days 50 hours + +# The time field is never changed. PostgreSQL: 1 mon -01:00:00 +query ? +SELECT justify_days(INTERVAL '1 month -1 hour') +---- +1 mons -1 hours + +# Largest day count. PostgreSQL: 5965232 years 4 mons 7 days +query ? +SELECT justify_days(INTERVAL '2147483647 days') +---- +71582788 mons 7 days + +# PostgreSQL: 178956970 years 7 mons (the largest month count) +query ? +SELECT justify_days(INTERVAL '2147483646 months 30 days') +---- +2147483647 mons + +# NULL propagates +query ? +SELECT justify_days(NULL) +---- +NULL + +# Row-wise vectorized batch +query ? +SELECT justify_days(v) FROM (VALUES + (INTERVAL '31 days'), + (CAST(NULL AS INTERVAL)), + (INTERVAL '1 month -35 days'), + (INTERVAL '3 months -91 days 1 hour') +) AS t(v) +---- +1 mons 1 days +NULL +-5 days +-1 days 1 hours + +# Overflow. PostgreSQL: ERROR: interval out of range +query error interval out of range +SELECT justify_days(INTERVAL '2147483647 months 30 days') + +query error interval out of range +SELECT justify_days(INTERVAL '-2147483648 months -30 days') + +# ============================================================================ +# justify_interval(interval) -- justify days and time, with sign adjustment +# ============================================================================ + +# Documented example. PostgreSQL: 29 days 23:00:00 +query ? +SELECT justify_interval(INTERVAL '1 month -1 hour') +---- +29 days 23 hours + +# PostgreSQL: -29 days -23:00:00 +query ? +SELECT justify_interval(INTERVAL '-1 month 1 hour') +---- +-29 days -23 hours + +# PostgreSQL: 1 mon 7 days 02:00:00 +query ? +SELECT justify_interval(INTERVAL '35 days 50 hours') +---- +1 mons 7 days 2 hours + +# PostgreSQL: -1 mons -7 days -02:00:00 +query ? +SELECT justify_interval(INTERVAL '-35 days -50 hours') +---- +-1 mons -7 days -2 hours + +# PostgreSQL: 1 mon 23:00:00 +query ? +SELECT justify_interval(INTERVAL '29 days 47 hours') +---- +1 mons 23 hours + +# PostgreSQL: 2 mons +query ? +SELECT justify_interval(INTERVAL '1 month 29 days 24 hours') +---- +2 mons + +# Mixed signs. PostgreSQL: 29 days 23:00:00 +query ? +SELECT justify_interval(INTERVAL '1 month 1 day -25 hours') +---- +29 days 23 hours + +# Mixed signs. PostgreSQL: -1 days +query ? +SELECT justify_interval(INTERVAL '2 months -59 days -48 hours') +---- +-1 days + +# Month borrow driven by the time field. PostgreSQL: 29 days 23:59:59.999999 +query ? +SELECT justify_interval(INTERVAL '1 month -0.000001 seconds') +---- +29 days 23 hours 59 mins 59.999999000 secs + +# Sub-microsecond time is kept and justified exactly. PostgreSQL rounds the +# input to whole microseconds when parsing, so it returns 1 mon here. +query ? +SELECT justify_interval(INTERVAL '1 month -1 nanosecond') +---- +29 days 23 hours 59 mins 59.999999999 secs + +# PostgreSQL: 1 year 2 mons 1 day 04:00:00 +query ? +SELECT justify_interval(INTERVAL '1 year 2 months -3 days 100 hours') +---- +14 mons 1 days 4 hours + +# Largest day count with a negative time. PostgreSQL: 5965232 years 4 mons 6 days 23:00:00 +query ? +SELECT justify_interval(INTERVAL '2147483647 days -1 hour') +---- +71582788 mons 6 days 23 hours + +# NULL propagates +query ? +SELECT justify_interval(NULL) +---- +NULL + +# Row-wise vectorized batch +query ? +SELECT justify_interval(v) FROM (VALUES + (INTERVAL '1 month -1 hour'), + (CAST(NULL AS INTERVAL)), + (INTERVAL '1 day -1 hour'), + (INTERVAL '3 months -91 days 1 hour') +) AS t(v) +---- +29 days 23 hours +NULL +23 hours +-23 hours + +# Overflow. PostgreSQL: ERROR: interval out of range +query error interval out of range +SELECT justify_interval(INTERVAL '2147483647 months 30 days') + +# justify_days accepts this, but justify_interval also folds the 24 hours +# into a day and overflows. PostgreSQL: ERROR: interval out of range +query error interval out of range +SELECT justify_interval(INTERVAL '2147483647 months 29 days 24 hours') + +query error interval out of range +SELECT justify_interval(INTERVAL '-2147483648 months -29 days -24 hours')