From 190677ac892fa9f6cada7ce2a0f5a2801f08bb93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Mon, 6 Jul 2026 11:56:39 +0200 Subject: [PATCH 1/7] fix(date_trunc): always return timestamp; NULL on invalid/null unit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spark's `date_trunc` (TruncTimestamp) always returns TimestampType regardless of input, but Sail returned `timestamp_ntz` for DATE and TIMESTAMP_NTZ inputs (it preserved the coerced input's missing timezone). This made `date_trunc('year', d) = TIMESTAMP '...'` return 0 rows, since the ntz result never matched a tz-aware literal. Cast the value argument to `Timestamp(µs, Some(session_tz))` before truncating, mirroring Spark's coercion to TimestampType. The cast is gated on date-family types (DATE/TIMESTAMP/TIMESTAMP_NTZ/STRING) so other types (int, time) still reject like Spark instead of being silently coerced. For a plain TIMESTAMP column the cast is identity and is elided by SimplifyExpressions, so there is no runtime overhead in the common path. Also return a typed-NULL timestamp when the unit is NULL or unrecognized (Spark returns NULL; DataFusion's date_trunc errors). `is_null_arg` sees through a `CAST(NULL AS STRING)` wrapper that is not folded at plan time. Verified against Spark JVM 4.1.1 --- .../sail-plan/src/function/scalar/datetime.rs | 62 ++++ .../function/features/date_trunc.feature | 306 +++++++++++++++++- 2 files changed, 361 insertions(+), 7 deletions(-) diff --git a/crates/sail-plan/src/function/scalar/datetime.rs b/crates/sail-plan/src/function/scalar/datetime.rs index 06804ac75c..3ae73fbba6 100644 --- a/crates/sail-plan/src/function/scalar/datetime.rs +++ b/crates/sail-plan/src/function/scalar/datetime.rs @@ -85,7 +85,59 @@ fn trunc(date: Expr, part: Expr) -> Expr { } fn date_trunc(input: ScalarFunctionInput) -> PlanResult { + let session_tz = input.function_context.plan_config.session_timezone.clone(); let (part, timestamp) = input.arguments.two()?; + // Spark returns NULL (not an error) when the unit is NULL or unrecognized, whereas + // DataFusion's date_trunc errors. Short-circuit those cases to a typed-NULL timestamp. + let null_timestamp = || { + lit(ScalarValue::TimestampMicrosecond( + None, + Some(session_tz.clone()), + )) + }; + match &part { + e if is_null_arg(e) => return Ok(null_timestamp()), + Expr::Literal(ScalarValue::Utf8(Some(s)), _) + | Expr::Literal(ScalarValue::LargeUtf8(Some(s)), _) + | Expr::Literal(ScalarValue::Utf8View(Some(s)), _) => { + if !matches!( + s.to_lowercase().as_str(), + "year" + | "yyyy" + | "yy" + | "quarter" + | "month" + | "mon" + | "mm" + | "week" + | "day" + | "dd" + | "hour" + | "minute" + | "second" + | "millisecond" + | "microsecond" + ) { + return Ok(null_timestamp()); + } + } + _ => {} + } + // Spark's date_trunc coerces its value argument to TimestampType and always returns + // TimestampType, whether the input is DATE, TIMESTAMP, TIMESTAMP_NTZ, or a string. + // Leave other types untouched so they are rejected like Spark instead of silently coerced. + let timestamp = match timestamp.get_type(input.function_context.schema)? { + DataType::Date32 + | DataType::Date64 + | DataType::Timestamp(_, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View => cast( + timestamp, + DataType::Timestamp(TimeUnit::Microsecond, Some(session_tz.clone())), + ), + _ => timestamp, + }; let truncated = ScalarUDF::from(SparkDateTrunc::new()).call(vec![trunc_part_conversion(part), timestamp]); let truncated = match truncated.get_type(input.function_context.schema)? { @@ -446,6 +498,16 @@ fn is_null_literal(expr: &Expr) -> bool { matches!(expr, Expr::Literal(value, _) if value.is_null()) } +/// Like [`is_null_literal`] but also sees through a `CAST(NULL AS ...)` wrapper, +/// which is not folded to a bare literal at plan time. +fn is_null_arg(expr: &Expr) -> bool { + match expr { + Expr::Cast(cast) => is_null_arg(&cast.expr), + Expr::TryCast(cast) => is_null_arg(&cast.expr), + other => is_null_literal(other), + } +} + fn to_timestamp(input: ScalarFunctionInput, timestamp_ntz: bool) -> PlanResult { timestamp_with_try(input, timestamp_ntz, false) } diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index 11a964d741..9405450885 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -1,8 +1,9 @@ -Feature: DATE_TRUNC preserves timestamp type +@date_trunc +Feature: DATE_TRUNC always returns a timestamp - Rule: date_trunc on timestamp preserves type + Rule: date_trunc returns timestamp for every date-family input type - Scenario: date_trunc on timestamp column preserves timestamp type + Scenario: date_trunc on timestamp column returns timestamp type When query """ WITH t(ts) AS (VALUES (TIMESTAMP '2026-02-02 00:00:00 UTC')) @@ -14,7 +15,7 @@ Feature: DATE_TRUNC preserves timestamp type |-- result: timestamp (nullable = true) """ - Scenario: date_trunc on timestamp_ntz column preserves timestamp_ntz type + Scenario: date_trunc on timestamp_ntz column returns timestamp type When query """ WITH t(ts) AS (VALUES (TIMESTAMP_NTZ '2026-02-02 00:00:00')) @@ -23,10 +24,10 @@ Feature: DATE_TRUNC preserves timestamp type Then query schema """ root - |-- result: timestamp_ntz (nullable = true) + |-- result: timestamp (nullable = true) """ - Scenario: date_trunc on timestamp_ntz literal preserves timestamp_ntz type + Scenario: date_trunc on timestamp_ntz literal returns timestamp type When query """ SELECT date_trunc('YEAR', TIMESTAMP_NTZ '2026-02-02 00:00:00') AS result @@ -34,8 +35,287 @@ Feature: DATE_TRUNC preserves timestamp type Then query schema """ root - |-- result: timestamp_ntz (nullable = true) + |-- result: timestamp (nullable = true) + """ + + Scenario: date_trunc on date literal returns timestamp type + When query + """ + SELECT date_trunc('DAY', DATE '2024-01-15') AS result + """ + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ + Then query result + | result | + | 2024-01-15 00:00:00 | + + Scenario: date_trunc on date column returns timestamp type + When query + """ + WITH t(d) AS (VALUES (DATE '2026-02-02')) + SELECT date_trunc('YEAR', d) AS result FROM t + """ + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ + Then query result + | result | + | 2026-01-01 00:00:00 | + + Scenario: date_trunc on timestamp_ntz literal value + When query + """ + SELECT date_trunc('DAY', TIMESTAMP_NTZ '2024-01-15 12:34:56') AS result + """ + Then query result + | result | + | 2024-01-15 00:00:00 | + + Scenario: date_trunc on string date value coerces to timestamp + When query + """ + SELECT date_trunc('DAY', '2024-01-15') AS result + """ + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ + Then query result + | result | + | 2024-01-15 00:00:00 | + + Scenario: date_trunc on string timestamp value coerces to timestamp + When query + """ + SELECT date_trunc('DAY', '2024-01-15 12:34:56') AS result + """ + Then query result + | result | + | 2024-01-15 00:00:00 | + + Rule: truncation unit selects the granularity + + Scenario: date_trunc YEAR truncates to start of year + When query + """ + SELECT date_trunc('YEAR', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-01-01 00:00:00 | + + Scenario: date_trunc QUARTER truncates to start of quarter + When query + """ + SELECT date_trunc('QUARTER', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-04-01 00:00:00 | + + Scenario: date_trunc MONTH truncates to start of month + When query + """ + SELECT date_trunc('MONTH', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-01 00:00:00 | + + Scenario: date_trunc WEEK truncates to Monday + When query + """ + SELECT date_trunc('WEEK', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-13 00:00:00 | + + Scenario: date_trunc DAY truncates to start of day + When query + """ + SELECT date_trunc('DAY', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-15 00:00:00 | + + Scenario: date_trunc HOUR truncates to start of hour + When query + """ + SELECT date_trunc('HOUR', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-15 13:00:00 | + + Scenario: date_trunc MINUTE truncates to start of minute + When query + """ + SELECT date_trunc('MINUTE', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:00 | + + Scenario: date_trunc SECOND truncates fractional seconds + When query """ + SELECT date_trunc('SECOND', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:30 | + + Scenario: date_trunc MILLISECOND keeps milliseconds + When query + """ + SELECT date_trunc('MILLISECOND', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:30.123 | + + Scenario: date_trunc MICROSECOND keeps microseconds + When query + """ + SELECT date_trunc('MICROSECOND', TIMESTAMP '2024-05-15 13:45:30.123456') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:30.123456 | + + Rule: unit aliases and case are accepted + + Scenario: date_trunc accepts YYYY alias for year + When query + """ + SELECT date_trunc('YYYY', TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | 2024-01-01 00:00:00 | + + Scenario: date_trunc accepts MM alias for month + When query + """ + SELECT date_trunc('MM', TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | 2024-05-01 00:00:00 | + + Scenario: date_trunc accepts DD alias for day + When query + """ + SELECT date_trunc('DD', TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | 2024-05-15 00:00:00 | + + Scenario: date_trunc unit is case insensitive + When query + """ + SELECT date_trunc('yEaR', TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | 2024-01-01 00:00:00 | + + Rule: invalid or null unit returns null + + Scenario: date_trunc with an unrecognized unit returns null + When query + """ + SELECT date_trunc('INVALID', TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | NULL | + + Scenario: date_trunc with an empty unit returns null + When query + """ + SELECT date_trunc('', TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | NULL | + + Scenario: date_trunc with a null unit returns null + When query + """ + SELECT date_trunc(CAST(NULL AS STRING), TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | NULL | + + Rule: null value returns null + + Scenario: date_trunc on a null date returns null + When query + """ + SELECT date_trunc('DAY', CAST(NULL AS DATE)) AS result + """ + Then query result + | result | + | NULL | + + Scenario: date_trunc on a null timestamp returns null + When query + """ + SELECT date_trunc('DAY', CAST(NULL AS TIMESTAMP)) AS result + """ + Then query result + | result | + | NULL | + + Rule: date_trunc rejects non-date input types + + Scenario: date_trunc rejects a time value + When query + """ + SELECT date_trunc('DAY', TIME '12:30:00') AS result + """ + Then query error .* + + Scenario: date_trunc rejects an integer value + When query + """ + SELECT date_trunc('DAY', 1) AS result + """ + Then query error .* + + Rule: date boundaries and leap days + + @sail-bug + Scenario: date_trunc YEAR at the maximum supported year + When query + """ + SELECT date_trunc('YEAR', DATE '9999-12-31') AS result + """ + Then query result + | result | + | 9999-01-01 00:00:00 | + + Scenario: date_trunc DAY on a leap day + When query + """ + SELECT date_trunc('DAY', DATE '2024-02-29') AS result + """ + Then query result + | result | + | 2024-02-29 00:00:00 | + + Rule: timezone-aware truncation on timestamps Scenario: date_trunc YEAR on timestamp values When query @@ -92,3 +372,15 @@ Feature: DATE_TRUNC preserves timestamp type Then query result | result | | 2026-03-15 18:00:00 | + + Rule: truncated date compares equal to a timestamp literal + + Scenario: date_trunc on date equals timestamp literal + When query + """ + SELECT count(*) AS result FROM (SELECT DATE '2019-08-04' AS d) + WHERE date_trunc('year', d) = TIMESTAMP '2019-01-01 00:00:00' + """ + Then query result + | result | + | 1 | From 6a6c0ca47752aec2a7a315f4211a0cddb18bd7a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Mon, 13 Jul 2026 21:36:31 +0200 Subject: [PATCH 2/7] fix column --- .../src/scalar/datetime/spark_date_trunc.rs | 121 +++++++++++++++++- .../sail-plan/src/function/scalar/datetime.rs | 8 +- .../function/features/date_trunc.feature | 47 +++++++ 3 files changed, 169 insertions(+), 7 deletions(-) diff --git a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs index c6be037360..c0b1829ead 100644 --- a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs +++ b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs @@ -1,12 +1,17 @@ use std::sync::Arc; -use datafusion::arrow::datatypes::{DataType, FieldRef}; -use datafusion_common::Result; +use datafusion::arrow::array::{Array, ArrayRef, AsArray, new_null_array}; +use datafusion::arrow::compute::interleave; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, }; use datafusion_functions::datetime::date_trunc::DateTruncFunc; +use sail_common_datafusion::utils::items::ItemTaker; + +use crate::error::unsupported_data_type_exec_err; #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkDateTrunc { @@ -19,12 +24,119 @@ impl Default for SparkDateTrunc { } } +/// The granularities DataFusion's `date_trunc` understands. Spark returns NULL for anything +/// else, so an unrecognized unit never reaches the inner function. +const GRANULARITIES: [&str; 10] = [ + "year", + "quarter", + "month", + "week", + "day", + "hour", + "minute", + "second", + "millisecond", + "microsecond", +]; + impl SparkDateTrunc { pub fn new() -> Self { Self { inner: DateTruncFunc::new(), } } + + /// Spark resolves the unit row by row, so a column of units truncates each row by its own + /// unit and nullifies only the rows whose unit is NULL or unrecognized. DataFusion's + /// `date_trunc` demands a scalar unit, so evaluate it once per DISTINCT unit — at most ten, + /// and one or two in practice — over the whole timestamp array, then assemble the rows. + /// That keeps the vectorized kernel instead of degrading to a per-row call. + fn invoke_with_unit_array(&self, args: ScalarFunctionArgs) -> Result { + let ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field, + config_options, + } = args; + let (unit, timestamp) = args.two()?; + let timestamp_field = arg_fields.two()?.1; + let ColumnarValue::Array(unit) = unit else { + return internal_err!("`date_trunc` expected an array of units"); + }; + let units = unit_strings(&unit)?; + + // Spark matches the unit case-insensitively, and the planner's `CASE` only rewrites the + // aliases (`mm`, `yy`, `dd`), so a column still carries the user's own casing. + let truncated_by_unit = GRANULARITIES + .iter() + .filter(|granularity| { + units + .iter() + .flatten() + .any(|unit| unit.eq_ignore_ascii_case(granularity)) + }) + .map(|granularity| { + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::from(*granularity)), + timestamp.clone(), + ], + arg_fields: vec![ + Arc::new(Field::new("granularity", DataType::Utf8, false)), + Arc::clone(×tamp_field), + ], + number_rows, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }; + let truncated = self.inner.invoke_with_args(args)?.to_array(number_rows)?; + Ok((*granularity, truncated)) + }) + .collect::>>()?; + + // The last source is the NULL row every unrecognized or NULL unit points at. + let nulls = new_null_array(return_field.data_type(), 1); + let sources = truncated_by_unit + .iter() + .map(|(_, array)| array.as_ref()) + .chain(std::iter::once(nulls.as_ref())) + .collect::>(); + let null_source = truncated_by_unit.len(); + + let indices = units + .iter() + .enumerate() + .map(|(row, unit)| { + unit.and_then(|unit| { + truncated_by_unit + .iter() + .position(|(granularity, _)| unit.eq_ignore_ascii_case(granularity)) + }) + .map_or((null_source, 0), |source| (source, row)) + }) + .collect::>(); + + Ok(ColumnarValue::Array(interleave(&sources, &indices)?)) + } +} + +/// The unit column reaches the function as any of the string types, since the planner builds it +/// with a `CASE` expression over the user's own column. +fn unit_strings(unit: &ArrayRef) -> Result>> { + let units: Vec> = match unit.data_type() { + DataType::Utf8 => unit.as_string::().iter().collect(), + DataType::LargeUtf8 => unit.as_string::().iter().collect(), + DataType::Utf8View => unit.as_string_view().iter().collect(), + other => { + return Err(unsupported_data_type_exec_err( + "date_trunc", + "a string unit", + other, + )); + } + }; + Ok(units) } impl ScalarUDFImpl for SparkDateTrunc { @@ -46,7 +158,10 @@ impl ScalarUDFImpl for SparkDateTrunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - self.inner.invoke_with_args(args) + match args.args.first() { + Some(ColumnarValue::Array(_)) => self.invoke_with_unit_array(args), + _ => self.inner.invoke_with_args(args), + } } fn aliases(&self) -> &[String] { diff --git a/crates/sail-plan/src/function/scalar/datetime.rs b/crates/sail-plan/src/function/scalar/datetime.rs index 961a1635c1..4569ca3247 100644 --- a/crates/sail-plan/src/function/scalar/datetime.rs +++ b/crates/sail-plan/src/function/scalar/datetime.rs @@ -99,7 +99,7 @@ fn date_trunc(input: ScalarFunctionInput) -> PlanResult { e if is_null_arg(e) => return Ok(null_timestamp()), Expr::Literal(ScalarValue::Utf8(Some(s)), _) | Expr::Literal(ScalarValue::LargeUtf8(Some(s)), _) - | Expr::Literal(ScalarValue::Utf8View(Some(s)), _) => { + | Expr::Literal(ScalarValue::Utf8View(Some(s)), _) if !matches!( s.to_lowercase().as_str(), "year" @@ -117,9 +117,9 @@ fn date_trunc(input: ScalarFunctionInput) -> PlanResult { | "second" | "millisecond" | "microsecond" - ) { - return Ok(null_timestamp()); - } + ) => + { + return Ok(null_timestamp()); } _ => {} } diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index 9405450885..370724e55a 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -258,6 +258,53 @@ Feature: DATE_TRUNC always returns a timestamp | result | | NULL | + Rule: the unit is resolved per row when it comes from a column + + # The unit is not required to be a literal: Spark resolves it row by row, so each row is + # truncated by its own unit and an unrecognized or NULL unit nullifies only that row. + + Scenario: date_trunc with the unit coming from a column + When query + """ + SELECT date_trunc(u, TIMESTAMP '2026-02-02 10:11:12') AS result FROM VALUES ('YEAR'), ('MONTH') AS t(u) + """ + Then query result ordered + | result | + | 2026-01-01 00:00:00 | + | 2026-02-01 00:00:00 | + + Scenario: an unrecognized unit in a column nullifies only its own row + When query + """ + SELECT date_trunc(u, TIMESTAMP '2026-02-02 10:11:12') AS result FROM VALUES ('YEAR'), ('INVALID') AS t(u) + """ + Then query result ordered + | result | + | 2026-01-01 00:00:00 | + | NULL | + + Scenario: a null unit in a column nullifies only its own row + When query + """ + SELECT date_trunc(u, TIMESTAMP '2026-02-02 10:11:12') AS result FROM VALUES ('YEAR'), (NULL) AS t(u) + """ + Then query result ordered + | result | + | 2026-01-01 00:00:00 | + | NULL | + + Scenario: units in a column keep their aliases and are case insensitive + When query + """ + SELECT date_trunc(u, TIMESTAMP '2026-02-02 10:11:12') AS result FROM VALUES ('yyyy'), ('MM'), ('Day'), ('hour') AS t(u) + """ + Then query result ordered + | result | + | 2026-01-01 00:00:00 | + | 2026-02-01 00:00:00 | + | 2026-02-02 00:00:00 | + | 2026-02-02 10:00:00 | + Rule: null value returns null Scenario: date_trunc on a null date returns null From acb5f622c18a11e4286833aa559f34b326ad8b88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Mon, 13 Jul 2026 23:19:07 +0200 Subject: [PATCH 3/7] fix copilot --- .../src/scalar/datetime/spark_date_trunc.rs | 106 ++++++++++------- .../sail-plan/src/function/scalar/datetime.rs | 77 +++++-------- .../function/features/date_trunc.feature | 108 ++++++++++++++++++ 3 files changed, 198 insertions(+), 93 deletions(-) diff --git a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs index c0b1829ead..4e07d85ef9 100644 --- a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs +++ b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs @@ -64,56 +64,62 @@ impl SparkDateTrunc { let ColumnarValue::Array(unit) = unit else { return internal_err!("`date_trunc` expected an array of units"); }; - let units = unit_strings(&unit)?; - - // Spark matches the unit case-insensitively, and the planner's `CASE` only rewrites the - // aliases (`mm`, `yy`, `dd`), so a column still carries the user's own casing. - let truncated_by_unit = GRANULARITIES - .iter() - .filter(|granularity| { - units - .iter() - .flatten() - .any(|unit| unit.eq_ignore_ascii_case(granularity)) - }) - .map(|granularity| { - let args = ScalarFunctionArgs { - args: vec![ - ColumnarValue::Scalar(ScalarValue::from(*granularity)), - timestamp.clone(), - ], - arg_fields: vec![ - Arc::new(Field::new("granularity", DataType::Utf8, false)), - Arc::clone(×tamp_field), - ], - number_rows, - return_field: Arc::clone(&return_field), - config_options: Arc::clone(&config_options), - }; - let truncated = self.inner.invoke_with_args(args)?.to_array(number_rows)?; - Ok((*granularity, truncated)) + // Resolve every row's granularity in a single pass. Spark matches the unit + // case-insensitively, and the planner's `CASE` only rewrites the aliases (`mm`, `yy`, + // `dd`), so a column still carries the user's own casing. A row whose unit is NULL or + // unrecognized has no granularity, and Spark nullifies it. + let granularity_of_row = unit_strings(&unit)? + .into_iter() + .map(|unit| { + unit.and_then(|unit| { + GRANULARITIES + .iter() + .position(|granularity| unit.eq_ignore_ascii_case(granularity)) + }) }) - .collect::>>()?; + .collect::>(); + + // Truncate once per DISTINCT granularity the column actually uses, over the whole + // timestamp array, and remember which source array each one landed in. + let mut source_of_granularity = [None; GRANULARITIES.len()]; + let mut truncated = Vec::new(); + for granularity in granularity_of_row.iter().flatten() { + if source_of_granularity[*granularity].is_some() { + continue; + } + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::from(GRANULARITIES[*granularity])), + timestamp.clone(), + ], + arg_fields: vec![ + Arc::new(Field::new("granularity", DataType::Utf8, false)), + Arc::clone(×tamp_field), + ], + number_rows, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }; + source_of_granularity[*granularity] = Some(truncated.len()); + truncated.push(self.inner.invoke_with_args(args)?.to_array(number_rows)?); + } // The last source is the NULL row every unrecognized or NULL unit points at. let nulls = new_null_array(return_field.data_type(), 1); - let sources = truncated_by_unit + let null_source = truncated.len(); + let sources = truncated .iter() - .map(|(_, array)| array.as_ref()) + .map(|array| array.as_ref()) .chain(std::iter::once(nulls.as_ref())) .collect::>(); - let null_source = truncated_by_unit.len(); - let indices = units + let indices = granularity_of_row .iter() .enumerate() - .map(|(row, unit)| { - unit.and_then(|unit| { - truncated_by_unit - .iter() - .position(|(granularity, _)| unit.eq_ignore_ascii_case(granularity)) - }) - .map_or((null_source, 0), |source| (source, row)) + .map(|(row, granularity)| { + granularity + .and_then(|granularity| source_of_granularity[granularity]) + .map_or((null_source, 0), |source| (source, row)) }) .collect::>(); @@ -157,11 +163,27 @@ impl ScalarUDFImpl for SparkDateTrunc { Ok(Arc::new(field.as_ref().clone().with_nullable(true))) } + /// Spark yields NULL for a unit it does not recognize, and for a NULL unit, whereas + /// DataFusion errors. Resolve that here rather than in the planner, so the behavior does + /// not depend on whether the unit reached the function as a literal or as a column. fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { match args.args.first() { - Some(ColumnarValue::Array(_)) => self.invoke_with_unit_array(args), - _ => self.inner.invoke_with_args(args), + Some(ColumnarValue::Array(_)) => return self.invoke_with_unit_array(args), + Some(ColumnarValue::Scalar(unit)) => { + let recognized = unit.try_as_str().flatten().is_some_and(|unit| { + GRANULARITIES + .iter() + .any(|granularity| unit.eq_ignore_ascii_case(granularity)) + }); + if !recognized { + return Ok(ColumnarValue::Scalar(ScalarValue::try_from( + args.return_field.data_type(), + )?)); + } + } + None => {} } + self.inner.invoke_with_args(args) } fn aliases(&self) -> &[String] { diff --git a/crates/sail-plan/src/function/scalar/datetime.rs b/crates/sail-plan/src/function/scalar/datetime.rs index 4569ca3247..ab840f5ae0 100644 --- a/crates/sail-plan/src/function/scalar/datetime.rs +++ b/crates/sail-plan/src/function/scalar/datetime.rs @@ -86,57 +86,42 @@ fn trunc(date: Expr, part: Expr) -> Expr { fn date_trunc(input: ScalarFunctionInput) -> PlanResult { let session_tz = input.function_context.plan_config.session_timezone.clone(); + let ansi_mode = input.function_context.plan_config.ansi_mode; let (part, timestamp) = input.arguments.two()?; - // Spark returns NULL (not an error) when the unit is NULL or unrecognized, whereas - // DataFusion's date_trunc errors. Short-circuit those cases to a typed-NULL timestamp. - let null_timestamp = || { - lit(ScalarValue::TimestampMicrosecond( - None, - Some(session_tz.clone()), - )) - }; - match &part { - e if is_null_arg(e) => return Ok(null_timestamp()), - Expr::Literal(ScalarValue::Utf8(Some(s)), _) - | Expr::Literal(ScalarValue::LargeUtf8(Some(s)), _) - | Expr::Literal(ScalarValue::Utf8View(Some(s)), _) - if !matches!( - s.to_lowercase().as_str(), - "year" - | "yyyy" - | "yy" - | "quarter" - | "month" - | "mon" - | "mm" - | "week" - | "day" - | "dd" - | "hour" - | "minute" - | "second" - | "millisecond" - | "microsecond" - ) => - { - return Ok(null_timestamp()); + // A string that is not a timestamp errors under ANSI, and is NULL when ANSI is off. + let to_timestamp = |expr: Expr| { + let timestamp = DataType::Timestamp(TimeUnit::Microsecond, Some(session_tz.clone())); + if ansi_mode { + cast(expr, timestamp) + } else { + try_cast(expr, timestamp) } - _ => {} - } + }; // Spark's date_trunc coerces its value argument to TimestampType and always returns // TimestampType, whether the input is DATE, TIMESTAMP, TIMESTAMP_NTZ, or a string. // Leave other types untouched so they are rejected like Spark instead of silently coerced. let timestamp = match timestamp.get_type(input.function_context.schema)? { - DataType::Date32 + DataType::Null + | DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) | DataType::Utf8 | DataType::LargeUtf8 - | DataType::Utf8View => cast( - timestamp, - DataType::Timestamp(TimeUnit::Microsecond, Some(session_tz.clone())), - ), - _ => timestamp, + | DataType::Utf8View => to_timestamp(timestamp), + // Spark rejects the value for its type when it analyzes the query, so the rejection + // must not depend on the unit: an unrecognized or NULL unit does not turn it into NULL. + other => { + return Err(PlanError::invalid(format!( + "`date_trunc` does not support {other} input" + ))); + } + }; + // Spark coerces a non-string unit to its string form, which then matches no granularity and + // yields NULL, rather than rejecting it. `SparkDateTrunc` resolves the unrecognized and NULL + // units, per row, so nothing here depends on the unit being a literal. + let part = match part.get_type(input.function_context.schema)? { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => part, + _ => cast(part, DataType::Utf8), }; let truncated = ScalarUDF::from(SparkDateTrunc::new()).call(vec![trunc_part_conversion(part), timestamp]); @@ -498,16 +483,6 @@ fn is_null_literal(expr: &Expr) -> bool { matches!(expr, Expr::Literal(value, _) if value.is_null()) } -/// Like [`is_null_literal`] but also sees through a `CAST(NULL AS ...)` wrapper, -/// which is not folded to a bare literal at plan time. -fn is_null_arg(expr: &Expr) -> bool { - match expr { - Expr::Cast(cast) => is_null_arg(&cast.expr), - Expr::TryCast(cast) => is_null_arg(&cast.expr), - other => is_null_literal(other), - } -} - fn to_timestamp(input: ScalarFunctionInput, timestamp_ntz: bool) -> PlanResult { timestamp_with_try(input, timestamp_ntz, false) } diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index 370724e55a..87037de89b 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -325,8 +325,88 @@ Feature: DATE_TRUNC always returns a timestamp | result | | NULL | + Scenario: date_trunc on an untyped null value returns a null timestamp + When query + """ + SELECT date_trunc('DAY', NULL) AS result + """ + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ + Then query result + | result | + | NULL | + + Rule: a non-string unit is measured as its string form + + Scenario: an integer unit matches no granularity and returns null + When query + """ + SELECT date_trunc(1, TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | NULL | + + Rule: an unparseable string value under ANSI on errors + + Scenario: an unparseable string value errors under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT date_trunc('DAY', 'not a date') AS result + """ + Then query error .* + + Scenario: an empty string value errors under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT date_trunc('DAY', '') AS result + """ + Then query error .* + + Rule: an unparseable string value under ANSI off returns NULL + + Scenario: an unparseable string value returns NULL under ANSI off + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT date_trunc('DAY', 'not a date') AS result + """ + Then query result + | result | + | NULL | + + Scenario: an empty string value returns NULL under ANSI off + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT date_trunc('DAY', '') AS result + """ + Then query result + | result | + | NULL | + + Scenario: an unparseable string value nulls only its own row under ANSI off + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT date_trunc('DAY', s) AS result FROM VALUES ('2024-05-15'), ('not a date'), (NULL) AS t(s) + """ + Then query result ordered + | result | + | 2024-05-15 00:00:00 | + | NULL | + | NULL | + Rule: date_trunc rejects non-date input types + # The value is rejected for its TYPE, whatever the unit is: an unrecognized or NULL unit + # does not turn the rejection into a NULL. + Scenario: date_trunc rejects a time value When query """ @@ -341,6 +421,34 @@ Feature: DATE_TRUNC always returns a timestamp """ Then query error .* + Scenario: date_trunc rejects an integer value even when the unit is unrecognized + When query + """ + SELECT date_trunc('INVALID', 1) AS result + """ + Then query error .* + + Scenario: date_trunc rejects an integer value even when the unit is null + When query + """ + SELECT date_trunc(CAST(NULL AS STRING), 1) AS result + """ + Then query error .* + + Scenario: date_trunc rejects a boolean value even when the unit is unrecognized + When query + """ + SELECT date_trunc('INVALID', true) AS result + """ + Then query error .* + + Scenario: date_trunc rejects a time value even when the unit is null + When query + """ + SELECT date_trunc(NULL, TIME '12:30:00') AS result + """ + Then query error .* + Rule: date boundaries and leap days @sail-bug From 6636ffbb32d35b382523fdbd901ad546456028a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Tue, 14 Jul 2026 06:36:37 +0200 Subject: [PATCH 4/7] trunc (TruncDate) and date_trunc (TruncTimestamp) are different Spark expressions and do not accept the same units, but Sail's plan builder feeds both from date_trunc's unit table. Five behaviours diverge: - trunc(date, 'DAY') and trunc(date, 'HOUR') truncate and return the date itself, where Spark returns NULL: a unit finer than a day is not valid for trunc. - an unrecognized unit errors instead of returning NULL. - a NULL unit errors instead of returning NULL. - a unit coming from a column errors, because the plan builder only matches Expr::Literal, so the query falls through to DataFusion's date_trunc, which demands a scalar unit. They live in date_trunc.feature because a single unit table governs both. Ten new scenarios, values captured on Spark JVM 4.1.1: five green, five tagged @sail-bug (strict xfail) and removed by the fix that follows. --- .../__snapshots__/features/date_trunc.yaml | 17 +++ .../function/features/date_trunc.feature | 141 +++++++++++++++++- 2 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml diff --git a/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml b/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml new file mode 100644 index 0000000000..a34e363a45 --- /dev/null +++ b/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml @@ -0,0 +1,17 @@ +# serializer version: 1 +--- +name: "test_the_plan_for_a_literal_unit" +data: + |- + == Physical Plan == + ProjectionExec: expr=[#2@0 as result] + ProjectionExec: expr=[date_trunc(YEAR, column1@0) as #2] + DataSourceExec: partitions=1, partition_sizes=[1] +--- +name: "test_the_plan_for_a_unit_coming_from_a_column" +data: + |- + == Physical Plan == + ProjectionExec: expr=[#4@0 as result] + ProjectionExec: expr=[date_trunc(CASE WHEN column2@1 ILIKE mon OR column2@1 ILIKE mm THEN month WHEN column2@1 ILIKE yy OR column2@1 ILIKE yyyy THEN year WHEN column2@1 ILIKE dd THEN day ELSE column2@1 END, column1@0) as #4] + DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index 87037de89b..0e914af5ec 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -1,5 +1,9 @@ @date_trunc -Feature: DATE_TRUNC always returns a timestamp +Feature: DATE_TRUNC and TRUNC truncate to a unit + # `date_trunc` (TruncTimestamp) and `trunc` (TruncDate) are different Spark expressions and do + # NOT accept the same units: `trunc` only truncates to date-level units, and anything finer -- + # or unrecognized, or NULL -- is NULL, not an error. They are covered together because a single + # unit table in the plan builder governs both. Rule: date_trunc returns timestamp for every date-family input type @@ -539,3 +543,138 @@ Feature: DATE_TRUNC always returns a timestamp Then query result | result | | 1 | + + Rule: trunc truncates to a date-level unit + + Scenario: trunc to the week + When query + """ + SELECT trunc(DATE '2019-08-04', 'week') AS result + """ + Then query result + | result | + | 2019-07-29 | + + Scenario: trunc to the quarter + When query + """ + SELECT trunc(DATE '2019-08-04', 'quarter') AS result + """ + Then query result + | result | + | 2019-07-01 | + + Scenario: trunc matches the unit case-insensitively + When query + """ + SELECT trunc(DATE '2019-08-04', 'MoN') AS result + """ + Then query result + | result | + | 2019-08-01 | + + Scenario: trunc accepts the yy alias + When query + """ + SELECT trunc(DATE '2019-08-04', 'yy') AS result + """ + Then query result + | result | + | 2019-01-01 | + + Scenario: trunc takes the date from a column + When query + """ + SELECT trunc(d, 'month') AS result FROM VALUES (1, DATE '2019-08-04'), (2, DATE '2020-02-29') AS t(i, d) ORDER BY i + """ + Then query result ordered + | result | + | 2019-08-01 | + | 2020-02-01 | + + Rule: a unit finer than a day is NULL, not a truncation + + # Sail reuses date_trunc's unit table, so it truncates to the day and returns the date itself. + @sail-bug + Scenario: trunc to the day is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', 'DAY') AS result + """ + Then query result + | result | + | NULL | + + # Sail returns the date itself. + @sail-bug + Scenario: trunc to the hour is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', 'HOUR') AS result + """ + Then query result + | result | + | NULL | + + Rule: an unrecognized or NULL unit is NULL, not an error + + # Sail errors: Unsupported date_trunc granularity: 'bogus'. + @sail-bug + Scenario: trunc with an unrecognized unit is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', 'bogus') AS result + """ + Then query result + | result | + | NULL | + + # Sail errors: Granularity of `date_trunc` must be non-null scalar Utf8. + @sail-bug + Scenario: trunc with a NULL unit is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', NULL) AS result + """ + Then query result + | result | + | NULL | + + Rule: the unit may come from a column + + # The plan builder only matches `Expr::Literal`, so a unit in a column falls through to + # DataFusion's `date_trunc`, which demands a scalar. Sail errors: Granularity of `date_trunc` + # must be non-null scalar Utf8. + @sail-bug + Scenario: trunc resolves the unit of each row + When query + """ + SELECT trunc(DATE '2019-08-04', c) AS result FROM VALUES (1, 'year'), (2, 'month'), (3, 'DAY'), (4, NULL) AS t(i, c) ORDER BY i + """ + Then query result ordered + | result | + | 2019-01-01 | + | 2019-08-01 | + | NULL | + | NULL | + + Rule: the plan resolves the unit + + # The plan is Sail's, so it cannot be compared against Spark. These snapshots exist to make the + # shape of the plan reviewable: how the unit is matched, and whether a literal unit still + # collapses to a single call once the units are enumerated in the plan. + @sail-only + Scenario: the plan for a literal unit + When query + """ + EXPLAIN SELECT date_trunc('YEAR', ts) AS result FROM VALUES (TIMESTAMP '2024-05-15 13:45:30') AS t(ts) + """ + Then query plan matches snapshot + + @sail-only + Scenario: the plan for a unit coming from a column + When query + """ + EXPLAIN SELECT date_trunc(c, ts) AS result FROM VALUES (TIMESTAMP '2024-05-15 13:45:30', 'YEAR') AS t(ts, c) + """ + Then query plan matches snapshot From 5cd49918827d6f477bd3063003959c85e42040fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Tue, 14 Jul 2026 06:46:37 +0200 Subject: [PATCH 5/7] fix The ClickBench plan snapshot changes only in the internal display name of the expression: the operator that actually runs is still date_trunc(minute, #4@0), a single call, because the CASE folds away for a literal unit. No plan regression. --- .../src/scalar/datetime/spark_date_trunc.rs | 148 ++---------------- .../sail-plan/src/function/scalar/datetime.rs | 113 ++++++++----- .../__snapshots__/test_clickbench.plan.yaml | 10 +- .../__snapshots__/features/date_trunc.yaml | 7 +- .../function/features/date_trunc.feature | 9 -- 5 files changed, 95 insertions(+), 192 deletions(-) diff --git a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs index 4e07d85ef9..3615ada916 100644 --- a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs +++ b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs @@ -1,18 +1,18 @@ use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, AsArray, new_null_array}; -use datafusion::arrow::compute::interleave; -use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; -use datafusion_common::{Result, ScalarValue, internal_err}; +use datafusion::arrow::datatypes::{DataType, FieldRef}; +use datafusion_common::Result; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, }; use datafusion_functions::datetime::date_trunc::DateTruncFunc; -use sail_common_datafusion::utils::items::ItemTaker; - -use crate::error::unsupported_data_type_exec_err; +/// DataFusion's `date_trunc`, with Spark's nullability. +/// +/// The unit is resolved in the plan builder, which enumerates the units and calls this function +/// with a literal one per branch (see `truncate_by_unit`), so an unrecognized, NULL, or columnar +/// unit never reaches the inner function. #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkDateTrunc { inner: DateTruncFunc, @@ -24,125 +24,12 @@ impl Default for SparkDateTrunc { } } -/// The granularities DataFusion's `date_trunc` understands. Spark returns NULL for anything -/// else, so an unrecognized unit never reaches the inner function. -const GRANULARITIES: [&str; 10] = [ - "year", - "quarter", - "month", - "week", - "day", - "hour", - "minute", - "second", - "millisecond", - "microsecond", -]; - impl SparkDateTrunc { pub fn new() -> Self { Self { inner: DateTruncFunc::new(), } } - - /// Spark resolves the unit row by row, so a column of units truncates each row by its own - /// unit and nullifies only the rows whose unit is NULL or unrecognized. DataFusion's - /// `date_trunc` demands a scalar unit, so evaluate it once per DISTINCT unit — at most ten, - /// and one or two in practice — over the whole timestamp array, then assemble the rows. - /// That keeps the vectorized kernel instead of degrading to a per-row call. - fn invoke_with_unit_array(&self, args: ScalarFunctionArgs) -> Result { - let ScalarFunctionArgs { - args, - arg_fields, - number_rows, - return_field, - config_options, - } = args; - let (unit, timestamp) = args.two()?; - let timestamp_field = arg_fields.two()?.1; - let ColumnarValue::Array(unit) = unit else { - return internal_err!("`date_trunc` expected an array of units"); - }; - // Resolve every row's granularity in a single pass. Spark matches the unit - // case-insensitively, and the planner's `CASE` only rewrites the aliases (`mm`, `yy`, - // `dd`), so a column still carries the user's own casing. A row whose unit is NULL or - // unrecognized has no granularity, and Spark nullifies it. - let granularity_of_row = unit_strings(&unit)? - .into_iter() - .map(|unit| { - unit.and_then(|unit| { - GRANULARITIES - .iter() - .position(|granularity| unit.eq_ignore_ascii_case(granularity)) - }) - }) - .collect::>(); - - // Truncate once per DISTINCT granularity the column actually uses, over the whole - // timestamp array, and remember which source array each one landed in. - let mut source_of_granularity = [None; GRANULARITIES.len()]; - let mut truncated = Vec::new(); - for granularity in granularity_of_row.iter().flatten() { - if source_of_granularity[*granularity].is_some() { - continue; - } - let args = ScalarFunctionArgs { - args: vec![ - ColumnarValue::Scalar(ScalarValue::from(GRANULARITIES[*granularity])), - timestamp.clone(), - ], - arg_fields: vec![ - Arc::new(Field::new("granularity", DataType::Utf8, false)), - Arc::clone(×tamp_field), - ], - number_rows, - return_field: Arc::clone(&return_field), - config_options: Arc::clone(&config_options), - }; - source_of_granularity[*granularity] = Some(truncated.len()); - truncated.push(self.inner.invoke_with_args(args)?.to_array(number_rows)?); - } - - // The last source is the NULL row every unrecognized or NULL unit points at. - let nulls = new_null_array(return_field.data_type(), 1); - let null_source = truncated.len(); - let sources = truncated - .iter() - .map(|array| array.as_ref()) - .chain(std::iter::once(nulls.as_ref())) - .collect::>(); - - let indices = granularity_of_row - .iter() - .enumerate() - .map(|(row, granularity)| { - granularity - .and_then(|granularity| source_of_granularity[granularity]) - .map_or((null_source, 0), |source| (source, row)) - }) - .collect::>(); - - Ok(ColumnarValue::Array(interleave(&sources, &indices)?)) - } -} - -/// The unit column reaches the function as any of the string types, since the planner builds it -/// with a `CASE` expression over the user's own column. -fn unit_strings(unit: &ArrayRef) -> Result>> { - let units: Vec> = match unit.data_type() { - DataType::Utf8 => unit.as_string::().iter().collect(), - DataType::LargeUtf8 => unit.as_string::().iter().collect(), - DataType::Utf8View => unit.as_string_view().iter().collect(), - other => { - return Err(unsupported_data_type_exec_err( - "date_trunc", - "a string unit", - other, - )); - } - }; - Ok(units) } impl ScalarUDFImpl for SparkDateTrunc { @@ -158,31 +45,14 @@ impl ScalarUDFImpl for SparkDateTrunc { self.inner.return_type(arg_types) } + /// Spark truncates a NULL timestamp to NULL, so the result is nullable even when the input is + /// not, which is what DataFusion would infer. fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { let field = self.inner.return_field_from_args(args)?; Ok(Arc::new(field.as_ref().clone().with_nullable(true))) } - /// Spark yields NULL for a unit it does not recognize, and for a NULL unit, whereas - /// DataFusion errors. Resolve that here rather than in the planner, so the behavior does - /// not depend on whether the unit reached the function as a literal or as a column. fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - match args.args.first() { - Some(ColumnarValue::Array(_)) => return self.invoke_with_unit_array(args), - Some(ColumnarValue::Scalar(unit)) => { - let recognized = unit.try_as_str().flatten().is_some_and(|unit| { - GRANULARITIES - .iter() - .any(|granularity| unit.eq_ignore_ascii_case(granularity)) - }); - if !recognized { - return Ok(ColumnarValue::Scalar(ScalarValue::try_from( - args.return_field.data_type(), - )?)); - } - } - None => {} - } self.inner.invoke_with_args(args) } diff --git a/crates/sail-plan/src/function/scalar/datetime.rs b/crates/sail-plan/src/function/scalar/datetime.rs index ab840f5ae0..3bcea14bb2 100644 --- a/crates/sail-plan/src/function/scalar/datetime.rs +++ b/crates/sail-plan/src/function/scalar/datetime.rs @@ -48,39 +48,78 @@ fn years(arg: Expr) -> Expr { integer_part(arg, "YEAR") } -fn trunc_part_conversion(part: Expr) -> Expr { +/// The units `date_trunc` truncates to, each with the aliases Spark accepts for it. +const DATE_TRUNC_UNITS: &[(&str, &[&str])] = &[ + ("year", &["year", "yyyy", "yy"]), + ("quarter", &["quarter"]), + ("month", &["month", "mon", "mm"]), + ("week", &["week"]), + ("day", &["day", "dd"]), + ("hour", &["hour"]), + ("minute", &["minute"]), + ("second", &["second"]), + ("millisecond", &["millisecond"]), + ("microsecond", &["microsecond"]), +]; + +/// `trunc` is a different Spark expression (`TruncDate`) and truncates only to a date-level unit. +/// A finer unit is not an error: it matches nothing and yields NULL, so `trunc` must not share +/// `date_trunc`'s table. +const TRUNC_UNITS: &[(&str, &[&str])] = &[ + ("year", &["year", "yyyy", "yy"]), + ("quarter", &["quarter"]), + ("month", &["month", "mon", "mm"]), + ("week", &["week"]), +]; + +/// Spark resolves the truncation unit **per row** and yields NULL for a unit it does not +/// recognize, whereas DataFusion's `date_trunc` demands a literal unit and errors otherwise. +/// +/// So enumerate the units in the plan rather than inspecting the argument: every branch calls the +/// kernel with a literal unit, which makes a unit arriving in a column behave exactly like a +/// literal one, and lets an unrecognized, NULL, or non-string unit fall through to NULL. Matching +/// the unit lives here and nowhere else, so the two comparisons cannot drift apart. +/// +/// This costs nothing in the common case: with a literal unit every predicate is foldable, so the +/// simplifier collapses the whole `CASE` back into a single call. The plan snapshots in +/// `date_trunc.feature` lock both shapes. +fn truncate_by_unit( + unit: Expr, + units: &[(&str, &[&str])], + truncate: impl Fn(&str) -> Expr, + null: Expr, +) -> Expr { + // Spark matches the unit case-insensitively, and coerces a non-string unit to its string form, + // which then matches no unit and yields NULL instead of being rejected. + let unit = expr_fn::lower(cast(unit, DataType::Utf8)); + let when_then_expr = units + .iter() + .filter_map(|(granularity, aliases)| { + let matches = aliases + .iter() + .map(|alias| unit.clone().eq(lit(*alias))) + .reduce(Expr::or)?; + Some((Box::new(matches), Box::new(truncate(granularity)))) + }) + .collect(); Expr::Case(expr::Case { expr: None, - when_then_expr: vec![ - ( - Box::new( - part.clone() - .ilike(lit("mon")) - .or(part.clone().ilike(lit("mm"))), - ), - Box::new(lit("month")), - ), - ( - Box::new( - part.clone() - .ilike(lit("yy")) - .or(part.clone().ilike(lit("yyyy"))), - ), - Box::new(lit("year")), - ), - ( - Box::new(part.clone().ilike(lit("dd"))), - Box::new(lit("day")), - ), - ], - else_expr: Some(Box::new(part)), + when_then_expr, + else_expr: Some(Box::new(null)), }) } fn trunc(date: Expr, part: Expr) -> Expr { - cast( - expr_fn::date_trunc(trunc_part_conversion(part), date), - DataType::Date32, + truncate_by_unit( + part, + TRUNC_UNITS, + |granularity| { + cast( + expr_fn::date_trunc(lit(granularity), date.clone()), + DataType::Date32, + ) + }, + lit(ScalarValue::Date32(None)), ) } @@ -116,15 +155,17 @@ fn date_trunc(input: ScalarFunctionInput) -> PlanResult { ))); } }; - // Spark coerces a non-string unit to its string form, which then matches no granularity and - // yields NULL, rather than rejecting it. `SparkDateTrunc` resolves the unrecognized and NULL - // units, per row, so nothing here depends on the unit being a literal. - let part = match part.get_type(input.function_context.schema)? { - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => part, - _ => cast(part, DataType::Utf8), - }; - let truncated = - ScalarUDF::from(SparkDateTrunc::new()).call(vec![trunc_part_conversion(part), timestamp]); + let truncated = truncate_by_unit( + part, + DATE_TRUNC_UNITS, + |granularity| { + ScalarUDF::from(SparkDateTrunc::new()).call(vec![lit(granularity), timestamp.clone()]) + }, + lit(ScalarValue::TimestampMicrosecond( + None, + Some(session_tz.clone()), + )), + ); let truncated = match truncated.get_type(input.function_context.schema)? { DataType::Timestamp(TimeUnit::Microsecond, _) => truncated, DataType::Timestamp(_, tz) => { diff --git a/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml b/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml index 3223b8fc95..e7d0d1113d 100644 --- a/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml +++ b/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml @@ -593,12 +593,12 @@ data: ProjectionExec: expr=[#105@0 as M, #106@1 as PageViews] ProjectionExec: expr=[#105@0 as #105, #106@1 as #106] GlobalLimitExec: skip=1000, fetch=10 - SortPreservingMergeExec: [date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)@2 ASC], fetch=1010 + SortPreservingMergeExec: [CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@2 ASC], fetch=1010 SortExec: TopK(fetch=1010), expr=[#105@0 ASC], preserve_partitioning=[true] - ProjectionExec: expr=[date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)@0 as #105, count(Int64(1))@1 as #106, date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)@0 as date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)] - AggregateExec: mode=FinalPartitioned, gby=[date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)@0 as date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)], aggr=[count(Int64(1))] - RepartitionExec: partitioning=Hash([date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)@0], 4), input_partitions=4 - AggregateExec: mode=Partial, gby=[date_trunc(minute, #4@0) as date_trunc(CASE WHEN Utf8("minute") ILIKE Utf8("mon") OR Utf8("minute") ILIKE Utf8("mm") THEN Utf8("month") WHEN Utf8("minute") ILIKE Utf8("yy") OR Utf8("minute") ILIKE Utf8("yyyy") THEN Utf8("year") WHEN Utf8("minute") ILIKE Utf8("dd") THEN Utf8("day") ELSE Utf8("minute") END,hits.#4)], aggr=[count(Int64(1))] + ProjectionExec: expr=[CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0 as #105, count(Int64(1))@1 as #106, CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0 as CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END] + AggregateExec: mode=FinalPartitioned, gby=[CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0 as CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END], aggr=[count(Int64(1))] + RepartitionExec: partitioning=Hash([CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0], 4), input_partitions=4 + AggregateExec: mode=Partial, gby=[date_trunc(minute, #4@0) as CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END], aggr=[count(Int64(1))] ProjectionExec: expr=[CAST(EventTime@0 * 1000000 AS Timestamp(µs, "UTC")) as #4] FilterExec: CounterID@2 = 62 AND EventDate@1 >= 15900 AND EventDate@1 <= 15901 AND IsRefresh@3 = 0 AND DontCountHits@4 = 0, projection=[EventTime@0] RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 diff --git a/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml b/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml index a34e363a45..0d621e824a 100644 --- a/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml +++ b/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml @@ -5,7 +5,7 @@ data: |- == Physical Plan == ProjectionExec: expr=[#2@0 as result] - ProjectionExec: expr=[date_trunc(YEAR, column1@0) as #2] + ProjectionExec: expr=[date_trunc(year, column1@0) as #2] DataSourceExec: partitions=1, partition_sizes=[1] --- name: "test_the_plan_for_a_unit_coming_from_a_column" @@ -13,5 +13,6 @@ data: |- == Physical Plan == ProjectionExec: expr=[#4@0 as result] - ProjectionExec: expr=[date_trunc(CASE WHEN column2@1 ILIKE mon OR column2@1 ILIKE mm THEN month WHEN column2@1 ILIKE yy OR column2@1 ILIKE yyyy THEN year WHEN column2@1 ILIKE dd THEN day ELSE column2@1 END, column1@0) as #4] - DataSourceExec: partitions=1, partition_sizes=[1] + ProjectionExec: expr=[CASE WHEN __common_expr_1@0 = year OR __common_expr_1@0 = yyyy OR __common_expr_1@0 = yy THEN date_trunc(year, #2@1) WHEN __common_expr_1@0 = quarter THEN date_trunc(quarter, #2@1) WHEN __common_expr_1@0 = month OR __common_expr_1@0 = mon OR __common_expr_1@0 = mm THEN date_trunc(month, #2@1) WHEN __common_expr_1@0 = week THEN date_trunc(week, #2@1) WHEN __common_expr_1@0 = day OR __common_expr_1@0 = dd THEN date_trunc(day, #2@1) WHEN __common_expr_1@0 = hour THEN date_trunc(hour, #2@1) WHEN __common_expr_1@0 = minute THEN date_trunc(minute, #2@1) WHEN __common_expr_1@0 = second THEN date_trunc(second, #2@1) WHEN __common_expr_1@0 = millisecond THEN date_trunc(millisecond, #2@1) WHEN __common_expr_1@0 = microsecond THEN date_trunc(microsecond, #2@1) END as #4] + ProjectionExec: expr=[lower(column2@1) as __common_expr_1, column1@0 as #2] + DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index 0e914af5ec..03c30c1ea2 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -594,8 +594,6 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit Rule: a unit finer than a day is NULL, not a truncation - # Sail reuses date_trunc's unit table, so it truncates to the day and returns the date itself. - @sail-bug Scenario: trunc to the day is NULL When query """ @@ -605,8 +603,6 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | result | | NULL | - # Sail returns the date itself. - @sail-bug Scenario: trunc to the hour is NULL When query """ @@ -618,8 +614,6 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit Rule: an unrecognized or NULL unit is NULL, not an error - # Sail errors: Unsupported date_trunc granularity: 'bogus'. - @sail-bug Scenario: trunc with an unrecognized unit is NULL When query """ @@ -629,8 +623,6 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | result | | NULL | - # Sail errors: Granularity of `date_trunc` must be non-null scalar Utf8. - @sail-bug Scenario: trunc with a NULL unit is NULL When query """ @@ -645,7 +637,6 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit # The plan builder only matches `Expr::Literal`, so a unit in a column falls through to # DataFusion's `date_trunc`, which demands a scalar. Sail errors: Granularity of `date_trunc` # must be non-null scalar Utf8. - @sail-bug Scenario: trunc resolves the unit of each row When query """ From 36023260c2f117c184e329757783c0dedf0a1894 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Fri, 17 Jul 2026 23:39:30 +0200 Subject: [PATCH 6/7] clean bug --- .../pysail/tests/spark/function/features/date_trunc.feature | 6 ++---- python/pysail/tests/spark/function/features/trunc.feature | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index af8ea562b5..48964fcde1 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -682,8 +682,7 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | result | | 2015-03-01 00:00:00 | - # Sail rejects the column: Sail errors: Granularity of `date_trunc` must be non-null scalar Utf8 - @column_args @sail-bug + @column_args Scenario: date_trunc takes argument 1 from a column holding two different values When query """ @@ -694,8 +693,7 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | 2015-01-01 00:00:00 | | 2015-03-01 00:00:00 | - # Sail rejects the column: Sail errors: Granularity of `date_trunc` must be non-null scalar Utf8 - @column_args @sail-bug + @column_args Scenario: date_trunc takes argument 1 from a column When query """ diff --git a/python/pysail/tests/spark/function/features/trunc.feature b/python/pysail/tests/spark/function/features/trunc.feature index 8e3463098c..759e178ddd 100644 --- a/python/pysail/tests/spark/function/features/trunc.feature +++ b/python/pysail/tests/spark/function/features/trunc.feature @@ -15,8 +15,7 @@ Feature: trunc with an argument coming from a column | result | | 2009-02-01 | - # Sail rejects the column: Sail errors: Granularity of `date_trunc` must be non-null scalar Utf8 - @column_args @sail-bug + @column_args Scenario: trunc takes argument 2 from a column holding two different values When query """ @@ -27,8 +26,7 @@ Feature: trunc with an argument coming from a column | 2009-02-01 | | 2009-02-09 | - # Sail rejects the column: Sail errors: Granularity of `date_trunc` must be non-null scalar Utf8 - @column_args @sail-bug + @column_args Scenario: trunc takes argument 2 from a column When query """ From 10cc0632534c2366bfec50996b3f583f280eec05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Thu, 23 Jul 2026 18:58:37 +0200 Subject: [PATCH 7/7] fix revision --- crates/sail-execution/src/proto/codec.rs | 3 + .../sail-function/src/scalar/datetime/mod.rs | 1 + .../src/scalar/datetime/spark_date_trunc.rs | 249 +++++- .../src/scalar/datetime/spark_timestamp.rs | 51 +- .../src/scalar/datetime/spark_trunc_level.rs | 139 +++ .../sail-plan/src/function/scalar/datetime.rs | 397 +++++++-- .../__snapshots__/test_clickbench.plan.yaml | 10 +- .../__snapshots__/features/date_trunc.yaml | 5 +- .../function/features/date_trunc.feature | 795 ++++++++++++++++-- .../spark/function/features/trunc.feature | 502 ++++++++++- .../function/test_date_trunc_encoding.py | 55 ++ .../function/test_date_trunc_error_class.py | 30 + 12 files changed, 2072 insertions(+), 165 deletions(-) create mode 100644 crates/sail-function/src/scalar/datetime/spark_trunc_level.rs create mode 100644 python/pysail/tests/spark/function/test_date_trunc_encoding.py create mode 100644 python/pysail/tests/spark/function/test_date_trunc_error_class.py diff --git a/crates/sail-execution/src/proto/codec.rs b/crates/sail-execution/src/proto/codec.rs index fe5ab1b8cf..cf233fed85 100644 --- a/crates/sail-execution/src/proto/codec.rs +++ b/crates/sail-execution/src/proto/codec.rs @@ -160,6 +160,7 @@ use sail_function::scalar::datetime::spark_time_diff::SparkTimeDiff; use sail_function::scalar::datetime::spark_time_trunc::SparkTimeTrunc; use sail_function::scalar::datetime::spark_timestamp::SparkTimestamp; use sail_function::scalar::datetime::spark_to_chrono_fmt::SparkToChronoFmt; +use sail_function::scalar::datetime::spark_trunc_level::SparkTruncLevel; use sail_function::scalar::datetime::spark_unix_timestamp::SparkUnixTimestamp; use sail_function::scalar::datetime::spark_window_buckets::SparkWindowBuckets; use sail_function::scalar::datetime::spark_year::SparkYear; @@ -2576,6 +2577,7 @@ impl PhysicalExtensionCodec for RemoteExecutionCodec { Ok(Arc::new(ScalarUDF::from(SparkDatePart::new()))) } "date_trunc" => Ok(Arc::new(ScalarUDF::from(SparkDateTrunc::new()))), + "spark_trunc_level" => Ok(Arc::new(ScalarUDF::from(SparkTruncLevel::new()))), "spark_time_diff" | "time_diff" => Ok(Arc::new(ScalarUDF::from(SparkTimeDiff::new()))), "spark_time_trunc" | "time_trunc" => { Ok(Arc::new(ScalarUDF::from(SparkTimeTrunc::new()))) @@ -2678,6 +2680,7 @@ impl PhysicalExtensionCodec for RemoteExecutionCodec { || node_inner.is::() || node_inner.is::() || node_inner.is::() + || node_inner.is::() || node_inner.is::() || node_inner.is::() || node_inner.is::() diff --git a/crates/sail-function/src/scalar/datetime/mod.rs b/crates/sail-function/src/scalar/datetime/mod.rs index 6537ed3d42..4597066a5d 100644 --- a/crates/sail-function/src/scalar/datetime/mod.rs +++ b/crates/sail-function/src/scalar/datetime/mod.rs @@ -14,6 +14,7 @@ pub mod spark_time_diff; pub mod spark_time_trunc; pub mod spark_timestamp; pub mod spark_to_chrono_fmt; +pub mod spark_trunc_level; pub mod spark_unix_timestamp; pub mod spark_window_buckets; pub mod spark_year; diff --git a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs index 3615ada916..334d0d17a6 100644 --- a/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs +++ b/crates/sail-function/src/scalar/datetime/spark_date_trunc.rs @@ -1,12 +1,21 @@ use std::sync::Arc; -use datafusion::arrow::datatypes::{DataType, FieldRef}; -use datafusion_common::Result; +use chrono::{ + DateTime, Datelike, Duration, NaiveDate, NaiveDateTime, NaiveTime, Offset, TimeZone, Timelike, + Utc, +}; +use datafusion::arrow::array::timezone::Tz; +use datafusion::arrow::array::{AsArray, PrimitiveArray}; +use datafusion::arrow::datatypes::{DataType, FieldRef, TimeUnit, TimestampMicrosecondType}; +use datafusion_common::{Result, ScalarValue, exec_datafusion_err}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, }; use datafusion_functions::datetime::date_trunc::DateTruncFunc; +use sail_common_datafusion::utils::datetime::localize_with_fallback; + +use crate::error::invalid_arg_count_exec_err; /// DataFusion's `date_trunc`, with Spark's nullability. /// @@ -41,8 +50,23 @@ impl ScalarUDFImpl for SparkDateTrunc { self.inner.signature() } + /// `DateTruncFunc::return_type` always fails -- it defers to `return_field_from_args` -- so + /// declaring the type here is what makes this function usable through the plain `return_type` + /// path. The result carries the value's own timestamp type, and a NULL value truncates to a + /// timestamp rather than to NULL. fn return_type(&self, arg_types: &[DataType]) -> Result { - self.inner.return_type(arg_types) + let Some(value) = arg_types.get(1) else { + return Err(invalid_arg_count_exec_err( + "date_trunc", + (2, 2), + arg_types.len(), + )); + }; + if value.is_null() { + Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)) + } else { + Ok(value.clone()) + } } /// Spark truncates a NULL timestamp to NULL, so the result is nullable even when the input is @@ -52,8 +76,66 @@ impl ScalarUDFImpl for SparkDateTrunc { Ok(Arc::new(field.as_ref().clone().with_nullable(true))) } + /// Spark truncates entirely in microseconds ([`DateTimeUtils.truncTimestamp`]), while + /// DataFusion converts the value to nanoseconds first and fails for anything past the year + /// 2262 -- a range Spark's own timestamp type covers. Truncate microsecond timestamps here + /// and leave every other input type to the inner function. + /// + /// [`DateTimeUtils.truncTimestamp`]: https://github.com/apache/spark/blob/v4.1.1/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala#L533-L553 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - self.inner.invoke_with_args(args) + if args.args.len() != 2 { + return Err(invalid_arg_count_exec_err( + "date_trunc", + (2, 2), + args.args.len(), + )); + } + // Decide before touching the arguments, so nothing is still borrowed from `args` if the + // inner function has to take them. + let microsecond_timestamp = match args.args[1].data_type() { + DataType::Timestamp(TimeUnit::Microsecond, timezone) => Some(timezone), + _ => None, + }; + let granularity = scalar_granularity(&args.args[0]) + .as_deref() + .and_then(Granularity::parse); + let Some((granularity, timezone)) = granularity.zip(microsecond_timestamp) else { + return self.inner.invoke_with_args(args); + }; + let data_type = args.args[1].data_type(); + let timezone = parse_timezone(timezone.as_deref())?; + match &args.args[1] { + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(value, tz)) => { + let truncated = value + .map(|value| granularity.truncate(value, &timezone)) + .transpose()?; + Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( + truncated, + tz.clone(), + ))) + } + ColumnarValue::Array(array) => { + // Iterate rather than `try_unary`, which runs the closure over null slots too: + // truncating rejects a value outside chrono's range, and a null slot may hold any + // payload at all, so a garbage byte pattern under a null bit would fail the batch. + let array: PrimitiveArray = array + .as_primitive::() + .iter() + .map(|value| { + value + .map(|value| granularity.truncate(value, &timezone)) + .transpose() + }) + .collect::>()?; + Ok(ColumnarValue::Array(Arc::new( + array.with_data_type(data_type), + ))) + } + other => Err(exec_datafusion_err!( + "`date_trunc` expected a microsecond timestamp, got {:?}", + other.data_type() + )), + } } fn aliases(&self) -> &[String] { @@ -68,3 +150,162 @@ impl ScalarUDFImpl for SparkDateTrunc { self.inner.documentation() } } + +const MICROS_PER_MILLI: i64 = 1_000; +const MICROS_PER_SECOND: i64 = 1_000_000; +const MICROS_PER_MINUTE: i64 = 60 * MICROS_PER_SECOND; +const MICROS_PER_HOUR: i64 = 60 * MICROS_PER_MINUTE; +const MICROS_PER_DAY: i64 = 24 * MICROS_PER_HOUR; + +/// The granularities Spark truncates a timestamp to. A unit and its aliases are resolved by +/// `SparkTruncLevel`, and the plan builder dispatches on the result, so only the canonical names +/// reach this function -- always as a literal, one per branch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Granularity { + Microsecond, + Millisecond, + Second, + Minute, + Hour, + Day, + Week, + Month, + Quarter, + Year, +} + +impl Granularity { + fn parse(granularity: &str) -> Option { + match granularity.to_uppercase().as_str() { + "MICROSECOND" => Some(Self::Microsecond), + "MILLISECOND" => Some(Self::Millisecond), + "SECOND" => Some(Self::Second), + "MINUTE" => Some(Self::Minute), + "HOUR" => Some(Self::Hour), + "DAY" => Some(Self::Day), + "WEEK" => Some(Self::Week), + "MONTH" => Some(Self::Month), + "QUARTER" => Some(Self::Quarter), + "YEAR" => Some(Self::Year), + _ => None, + } + } + + /// Time zone offsets are whole seconds, so microsecond, millisecond and second truncate + /// without consulting the zone at all -- which is why Spark does them first. + fn truncate(self, micros: i64, timezone: &Tz) -> Result { + match self { + Self::Microsecond => Ok(micros), + Self::Millisecond => Ok(micros - micros.rem_euclid(MICROS_PER_MILLI)), + Self::Second => Ok(micros - micros.rem_euclid(MICROS_PER_SECOND)), + Self::Minute => truncate_to_unit(micros, timezone, MICROS_PER_MINUTE), + Self::Hour => truncate_to_unit(micros, timezone, MICROS_PER_HOUR), + Self::Day => truncate_to_unit(micros, timezone, MICROS_PER_DAY), + Self::Week | Self::Month | Self::Quarter | Self::Year => { + let date = local_date(micros, timezone)?; + days_to_micros(self.truncate_date(date)?, timezone) + } + } + } + + fn truncate_date(self, date: NaiveDate) -> Result { + let truncated = match self { + // Spark's week starts on Monday. + Self::Week => date + .checked_sub_signed(Duration::days(date.weekday().num_days_from_monday() as i64)), + Self::Month => date.with_day(1), + Self::Quarter => date + .with_day(1) + .and_then(|date| date.with_month((date.month0() / 3) * 3 + 1)), + Self::Year => date.with_day(1).and_then(|date| date.with_month(1)), + _ => Some(date), + } + .ok_or_else(|| exec_datafusion_err!("cannot truncate date {date} to {self:?}"))?; + Ok(truncated) + } +} + +/// The offset in microseconds that the time zone applies at the given instant. +fn offset_micros(micros: i64, timezone: &Tz) -> Result { + let seconds = micros.div_euclid(MICROS_PER_SECOND); + let instant = DateTime::from_timestamp(seconds, 0) + .ok_or_else(|| exec_datafusion_err!("timestamp out of range: {micros}"))?; + let offset = timezone + .offset_from_utc_datetime(&instant.naive_utc()) + .fix() + .local_minus_utc() as i64; + Ok(offset * MICROS_PER_SECOND) +} + +/// Truncate to a unit that divides the local day, by offset arithmetic rather than by building a +/// zoned datetime per row. Mirrors Spark's `truncToUnitFast`: when truncating moves the value +/// across an offset transition the arithmetic no longer holds, and the local wall clock decides. +fn truncate_to_unit(micros: i64, timezone: &Tz, unit_micros: i64) -> Result { + let offset = offset_micros(micros, timezone)?; + let local = micros + .checked_add(offset) + .ok_or_else(|| exec_datafusion_err!("timestamp out of range: {micros}"))?; + let truncated = local - local.rem_euclid(unit_micros); + let candidate = truncated + .checked_sub(offset) + .ok_or_else(|| exec_datafusion_err!("timestamp out of range: {micros}"))?; + if offset_micros(candidate, timezone)? == offset { + return Ok(candidate); + } + let local = local_datetime(micros, timezone)?; + let time = local.time().num_seconds_from_midnight() as i64 * MICROS_PER_SECOND + + local.time().nanosecond() as i64 / 1_000; + let truncated = local.date().and_time(NaiveTime::MIN) + + Duration::microseconds(time - time.rem_euclid(unit_micros)); + Ok(instant_micros(&localize_with_fallback( + timezone, &truncated, + )?)) +} + +fn local_datetime(micros: i64, timezone: &Tz) -> Result { + let offset = offset_micros(micros, timezone)?; + let local = micros + .checked_add(offset) + .ok_or_else(|| exec_datafusion_err!("timestamp out of range: {micros}"))?; + DateTime::from_timestamp_micros(local) + .map(|datetime| datetime.naive_utc()) + .ok_or_else(|| exec_datafusion_err!("timestamp out of range: {micros}")) +} + +fn local_date(micros: i64, timezone: &Tz) -> Result { + Ok(local_datetime(micros, timezone)?.date()) +} + +/// Midnight of the local date, as microseconds since the epoch. A local midnight that a DST +/// transition skips does not exist, and Spark's `atStartOfDay` moves it forward to the first +/// valid instant, which is what the fallback does. +fn days_to_micros(date: NaiveDate, timezone: &Tz) -> Result { + let midnight = date.and_time(NaiveTime::MIN); + Ok(instant_micros(&localize_with_fallback( + timezone, &midnight, + )?)) +} + +fn instant_micros(instant: &DateTime) -> i64 { + instant.timestamp_micros() +} + +/// The granularity only ever reaches this function as a literal, because the plan builder +/// enumerates the units and calls it once per branch. +fn scalar_granularity(granularity: &ColumnarValue) -> Option { + match granularity { + ColumnarValue::Scalar( + ScalarValue::Utf8(Some(granularity)) + | ScalarValue::LargeUtf8(Some(granularity)) + | ScalarValue::Utf8View(Some(granularity)), + ) => Some(granularity.clone()), + _ => None, + } +} + +fn parse_timezone(timezone: Option<&str>) -> Result { + timezone + .unwrap_or("UTC") + .parse() + .map_err(|_| exec_datafusion_err!("invalid timezone: {timezone:?}")) +} diff --git a/crates/sail-function/src/scalar/datetime/spark_timestamp.rs b/crates/sail-function/src/scalar/datetime/spark_timestamp.rs index bf4dd1621a..3021083e43 100644 --- a/crates/sail-function/src/scalar/datetime/spark_timestamp.rs +++ b/crates/sail-function/src/scalar/datetime/spark_timestamp.rs @@ -1,10 +1,10 @@ use std::fmt::Debug; use std::sync::Arc; -use chrono::{NaiveDate, NaiveDateTime}; +use chrono::{DateTime, NaiveDate, NaiveDateTime}; use datafusion::arrow::array::timezone::Tz; -use datafusion::arrow::array::{Array, ArrayRef, TimestampMicrosecondArray}; -use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::arrow::array::{Array, ArrayRef, AsArray, TimestampMicrosecondArray}; +use datafusion::arrow::datatypes::{DataType, TimeUnit, TimestampMicrosecondType}; use datafusion_common::cast::{as_large_string_array, as_string_array, as_string_view_array}; use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; @@ -148,12 +148,42 @@ impl SparkTimestamp { } } + /// A TIMESTAMP_NTZ value is already a wall clock, so it is localized rather than parsed. + /// Arrow's own cast cannot do this: it resolves the local time with `.single()`, which has no + /// answer inside a DST gap or inside the repeated hour of a fall-back, and so errors or + /// nullifies exactly where Spark moves the value forward, or takes the earlier offset. + fn localize_naive( + parser: &TimestampParser, + safe: bool, + array: &ArrayRef, + ) -> Result { + array + .as_primitive::() + .iter() + .map(|value| match value { + Some(micros) => { + let datetime = DateTime::from_timestamp_micros(micros) + .ok_or_else(|| exec_datafusion_err!("timestamp out of range: {micros}"))? + .naive_utc(); + parser.localize(datetime, "", safe) + } + None => Ok(None), + }) + .collect::>() + } + fn kernel( parser: &TimestampParser, safe: bool, args: &[ArrayRef], ) -> Result { let value_arr = &args[0]; + if matches!( + value_arr.data_type(), + DataType::Timestamp(TimeUnit::Microsecond, None) + ) { + return Self::localize_naive(parser, safe, value_arr); + } match args.get(1) { Some(format_arr) => { if value_arr.len() != format_arr.len() { @@ -211,19 +241,26 @@ impl ScalarUDFImpl for SparkTimestamp { )); } match &arg_types[0] { - // String-only, matching the kernel (which parses strings) and the - // sibling parsers `SparkDate`/`SparkTime`. The planner casts/handles - // DATE/TIMESTAMP inputs directly, so they never reach this UDF. + // Strings, matching the kernel (which parses them) and the sibling parsers + // `SparkDate`/`SparkTime`. A zoned TIMESTAMP the planner casts directly and never + // sends here; a DATE it converts to a naive timestamp first, so that one arrives + // through the arm below. DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View | DataType::Null => {} + // A TIMESTAMP_NTZ has to be localized the way Spark localizes it, which the Arrow + // cast cannot do across a DST transition, so the planner routes it here instead. + DataType::Timestamp(_, None) => {} other => { return Err(unsupported_data_type_exec_err( self.name(), - "STRING or NULL", + "STRING, TIMESTAMP_NTZ or NULL", other, )); } } let mut coerced = arg_types.to_vec(); + if matches!(arg_types[0], DataType::Timestamp(_, None)) { + coerced[0] = DataType::Timestamp(TimeUnit::Microsecond, None); + } if let Some(format) = arg_types.get(1) { match format { DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {} diff --git a/crates/sail-function/src/scalar/datetime/spark_trunc_level.rs b/crates/sail-function/src/scalar/datetime/spark_trunc_level.rs new file mode 100644 index 0000000000..e9204266a2 --- /dev/null +++ b/crates/sail-function/src/scalar/datetime/spark_trunc_level.rs @@ -0,0 +1,139 @@ +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, AsArray, StringArray}; +use datafusion::arrow::datatypes::DataType; +use datafusion_common::{Result, ScalarValue, exec_datafusion_err}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; + +use crate::error::invalid_arg_count_exec_err; + +/// The truncation unit resolved to the granularity it names, in ONE evaluation of the unit. +/// +/// The plan builder dispatches on the result with a `CASE`, and that `CASE` has to tell three +/// outcomes apart: the unit named a granularity, the unit was NULL, or the unit matched nothing. +/// An expression cannot: `CASE x WHEN NULL` never matches, so a NULL selector and an unmatched +/// one both fall to the same `ELSE`, and every way of splitting them in the plan -- `coalesce`, +/// `IS NULL` -- mentions the unit a second time. DataFusion will not share a volatile subtree +/// between those mentions (`common_subexpr_eliminate.rs`, `!node.is_volatile_node()`), so a unit +/// such as `IF(rand() < 0.5, NULL, 'YEAR')` would be drawn once per mention and rows would take a +/// branch neither draw called for. Resolving it here keeps it to a single evaluation. +/// +/// The table is `date_trunc`'s. `trunc` accepts a subset, and the plan builder omits the finer +/// granularities from its dispatch: a literal one it filters out while building the plan, and a +/// columnar one lands in the `ELSE` -- which is where a unit `trunc` does not accept belongs. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkTruncLevel { + signature: Signature, +} + +/// Returned for a NULL unit, to tell it apart from a unit that matched nothing (which yields +/// NULL). Lower case, which `to_uppercase` can never produce, so no real unit collides with it. +pub const TRUNC_LEVEL_NULL: &str = "null"; + +/// Each granularity with the aliases Spark accepts for it, upper-cased because Spark folds the +/// unit with `toUpperCase(Locale.ROOT)` before matching (`DateTimeUtils.parseTruncLevel`). +const LEVELS: &[(&str, &[&str])] = &[ + ("year", &["YEAR", "YYYY", "YY"]), + ("quarter", &["QUARTER"]), + ("month", &["MONTH", "MON", "MM"]), + ("week", &["WEEK"]), + ("day", &["DAY", "DD"]), + ("hour", &["HOUR"]), + ("minute", &["MINUTE"]), + ("second", &["SECOND"]), + ("millisecond", &["MILLISECOND"]), + ("microsecond", &["MICROSECOND"]), +]; + +impl Default for SparkTruncLevel { + fn default() -> Self { + Self::new() + } +} + +impl SparkTruncLevel { + pub fn new() -> Self { + Self { + signature: Signature::string(1, Volatility::Immutable), + } + } + + /// `None` for a unit that names no granularity; Spark yields NULL for it rather than + /// rejecting it. + /// + /// Public because the plan builder resolves a LITERAL unit at plan-build time and has to + /// reach the same answer this function gives at run time. Two alias tables would drift. + pub fn level(unit: &str) -> Option<&'static str> { + let folded = unit.to_uppercase(); + LEVELS + .iter() + .find(|(_, aliases)| aliases.contains(&folded.as_str())) + .map(|(level, _)| *level) + } + + fn resolve(unit: Option<&str>) -> Option<&'static str> { + match unit { + None => Some(TRUNC_LEVEL_NULL), + Some(unit) => Self::level(unit), + } + } +} + +impl ScalarUDFImpl for SparkTruncLevel { + fn name(&self) -> &str { + "spark_trunc_level" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [unit] = args.args.as_slice() else { + return Err(invalid_arg_count_exec_err( + "spark_trunc_level", + (1, 1), + args.args.len(), + )); + }; + match unit { + ColumnarValue::Scalar(unit) => { + let unit = match unit { + ScalarValue::Utf8(unit) + | ScalarValue::LargeUtf8(unit) + | ScalarValue::Utf8View(unit) => unit.as_deref(), + ScalarValue::Null => None, + other => { + return Err(exec_datafusion_err!( + "`spark_trunc_level` expected a string unit, got {other:?}" + )); + } + }; + Ok(ColumnarValue::Scalar(ScalarValue::Utf8( + Self::resolve(unit).map(str::to_string), + ))) + } + ColumnarValue::Array(array) => { + let levels: StringArray = match array.data_type() { + DataType::Utf8 => array.as_string::().iter().map(Self::resolve).collect(), + DataType::LargeUtf8 => { + array.as_string::().iter().map(Self::resolve).collect() + } + DataType::Utf8View => { + array.as_string_view().iter().map(Self::resolve).collect() + } + other => { + return Err(exec_datafusion_err!( + "`spark_trunc_level` expected a string unit, got {other:?}" + )); + } + }; + Ok(ColumnarValue::Array(Arc::new(levels) as ArrayRef)) + } + } + } +} diff --git a/crates/sail-plan/src/function/scalar/datetime.rs b/crates/sail-plan/src/function/scalar/datetime.rs index 3bcea14bb2..ed25bcf90d 100644 --- a/crates/sail-plan/src/function/scalar/datetime.rs +++ b/crates/sail-plan/src/function/scalar/datetime.rs @@ -4,6 +4,7 @@ use datafusion::arrow::datatypes::{ DataType, IntervalDayTimeType, IntervalUnit, IntervalYearMonthType, TimeUnit, }; use datafusion::functions::expr_fn; +use datafusion_common::types::NativeType; use datafusion_common::{DFSchemaRef, ScalarValue}; use datafusion_expr::expr::{self, Expr}; use datafusion_expr::{BinaryExpr, ExprSchemable, Operator, ScalarUDF, cast, lit, try_cast, when}; @@ -13,6 +14,7 @@ use datafusion_spark::function::datetime::make_interval::SparkMakeInterval; use sail_common::datetime::time_unit_to_multiplier; use sail_common_datafusion::utils::items::ItemTaker; use sail_function::scalar::datetime::convert_tz::ConvertTz; +use sail_function::scalar::datetime::spark_date::SparkDate; use sail_function::scalar::datetime::spark_date_part::SparkDatePart; use sail_function::scalar::datetime::spark_date_trunc::SparkDateTrunc; use sail_function::scalar::datetime::spark_last_day::SparkLastDay; @@ -25,6 +27,7 @@ use sail_function::scalar::datetime::spark_time_diff::SparkTimeDiff; use sail_function::scalar::datetime::spark_time_trunc::SparkTimeTrunc; use sail_function::scalar::datetime::spark_timestamp::SparkTimestamp; use sail_function::scalar::datetime::spark_to_chrono_fmt::SparkToChronoFmt; +use sail_function::scalar::datetime::spark_trunc_level::{SparkTruncLevel, TRUNC_LEVEL_NULL}; use sail_function::scalar::datetime::spark_unix_timestamp::SparkUnixTimestamp; use sail_function::scalar::datetime::spark_window_buckets::SparkWindowBuckets; use sail_function::scalar::datetime::spark_year::SparkYear; @@ -48,29 +51,26 @@ fn years(arg: Expr) -> Expr { integer_part(arg, "YEAR") } -/// The units `date_trunc` truncates to, each with the aliases Spark accepts for it. -const DATE_TRUNC_UNITS: &[(&str, &[&str])] = &[ - ("year", &["year", "yyyy", "yy"]), - ("quarter", &["quarter"]), - ("month", &["month", "mon", "mm"]), - ("week", &["week"]), - ("day", &["day", "dd"]), - ("hour", &["hour"]), - ("minute", &["minute"]), - ("second", &["second"]), - ("millisecond", &["millisecond"]), - ("microsecond", &["microsecond"]), +/// The granularities `date_trunc` truncates to. The ALIASES each one accepts live in +/// `SparkTruncLevel`, which is the single place a unit is matched -- whether it arrives as a +/// literal or in a column -- so the two paths cannot drift apart. +const DATE_TRUNC_LEVELS: &[&str] = &[ + "year", + "quarter", + "month", + "week", + "day", + "hour", + "minute", + "second", + "millisecond", + "microsecond", ]; /// `trunc` is a different Spark expression (`TruncDate`) and truncates only to a date-level unit. -/// A finer unit is not an error: it matches nothing and yields NULL, so `trunc` must not share -/// `date_trunc`'s table. -const TRUNC_UNITS: &[(&str, &[&str])] = &[ - ("year", &["year", "yyyy", "yy"]), - ("quarter", &["quarter"]), - ("month", &["month", "mon", "mm"]), - ("week", &["week"]), -]; +/// A finer unit is not an error: it names a granularity this list omits, so it falls through to +/// NULL exactly like a unit that matches nothing. +const TRUNC_LEVELS: &[&str] = &["year", "quarter", "month", "week"]; /// Spark resolves the truncation unit **per row** and yields NULL for a unit it does not /// recognize, whereas DataFusion's `date_trunc` demands a literal unit and errors otherwise. @@ -80,53 +80,284 @@ const TRUNC_UNITS: &[(&str, &[&str])] = &[ /// literal one, and lets an unrecognized, NULL, or non-string unit fall through to NULL. Matching /// the unit lives here and nowhere else, so the two comparisons cannot drift apart. /// -/// This costs nothing in the common case: with a literal unit every predicate is foldable, so the -/// simplifier collapses the whole `CASE` back into a single call. The plan snapshots in -/// `date_trunc.feature` lock both shapes. +/// A literal STRING unit is resolved here and emits the single call it names, so the common case +/// costs nothing: no branch, and nothing per-unit shipped to the workers. A literal of any other +/// type is measured as its string form, which this stage cannot compute, so it takes the dispatch +/// below like a column would. +/// +/// A unit that only exists as a column is dispatched by a `CASE` over the resolved level. That +/// `CASE` repeats the VALUE subtree once per branch -- ten for `date_trunc`, four for `trunc`, +/// plus the unmatched one -- and the whole thing travels to every worker in the physical plan. +/// Runtime cost is bounded (`CaseExpr` evaluates each branch on a disjoint slice of rows, so the +/// value is still computed at most once per row), but plan SIZE is not. Collapsing it would mean +/// passing the level as a column to a single call, which would also make the value an ordinary +/// argument -- always evaluated -- and that is exactly what Spark's short circuit forbids. +/// +/// Whether a unit that matches nothing converts the value is not the same in the two cases, and +/// Spark's codegen is where the difference lives (`TruncInstant.codeGenHelper`): +/// +/// * a **literal** unit that names no granularity short-circuits to NULL and never evaluates the +/// value, so `null` is a bare NULL literal; +/// * a unit **in a column** is evaluated together with the value before the level is checked, so +/// `unmatched` has to convert the value and only then yield NULL -- a bare literal there would +/// swallow the error Spark raises under ANSI for a value it cannot convert. fn truncate_by_unit( unit: Expr, - units: &[(&str, &[&str])], + levels: &[&str], truncate: impl Fn(&str) -> Expr, null: Expr, + unmatched: Expr, + null_unit_converts_value: bool, ) -> Expr { - // Spark matches the unit case-insensitively, and coerces a non-string unit to its string form, - // which then matches no unit and yields NULL instead of being rejected. - let unit = expr_fn::lower(cast(unit, DataType::Utf8)); - let when_then_expr = units - .iter() - .filter_map(|(granularity, aliases)| { - let matches = aliases - .iter() - .map(|alias| unit.clone().eq(lit(*alias))) - .reduce(Expr::or)?; - Some((Box::new(matches), Box::new(truncate(granularity)))) - }) - .collect(); + // A NULL unit matches nothing whatever its type, and Spark folds the cast that types it -- + // `date_trunc(CAST(NULL AS STRING), ...)` short-circuits just like a bare NULL. + if is_folded_null(&unit) { + return null; + } + // Spark's short circuit keys off `foldable`, not off being a literal: it resolves anything + // without an input at compile time, so `concat('BOG','US')` skips the value too. Sail cannot + // name the granularity for those at plan-build time -- the optimizer that would evaluate them + // runs later -- but it can still decline to convert the value, which is the part that shows. + // + // KNOWN GAP -- a scalar subquery is treated as foldable here, and that is not always right. + // `ScalarSubquery.foldable` is false in Spark; what makes the short circuit happen is Spark's + // own optimizer folding the subquery into a literal before codegen. Whether it does depends + // on the subquery, and this check cannot tell: it sees the UNOPTIMIZED plan. + // + // Measured on JVM 4.1.1. `(SELECT 'BOGUS')` and a one-row temporary view are both folded, and + // Spark returns NULL for them -- Sail agrees. A subquery over a relation (`... FROM range(1)`, + // a table) is NOT folded, and Spark raises `CAST_INVALID_INPUT` where Sail still returns NULL. + // That case is reproducible and covered by a `@sail-bug` scenario in `date_trunc.feature`. + // + // The boundary IS detectable from here: Spark folds exactly when every leaf of the subquery + // is a one-row relation. `(SELECT 'BOGUS')`, a view over a constant and `FROM (SELECT 1)` all + // fold; `FROM (VALUES ('x'))`, `FROM range(1)` and a table do not. Left alone anyway, because + // that boundary is a quirk of which optimizer rule fires rather than a rule of the language -- + // nothing explains why `VALUES` of one row is not folded while `SELECT 1` is -- so encoding it + // buys one exotic case and pins Sail to an internal detail of a specific Spark version. + let foldable = unit.column_refs().is_empty() && !unit.is_volatile() && !unit.contains_outer(); + let unmatched = if foldable { null.clone() } else { unmatched }; + // A NULL unit and a unit that matches nothing are NOT the same case, and which one skips the + // value follows the ARGUMENT ORDER of the Spark expression. `nullSafeCodeGen` evaluates the + // left child, and only evaluates the right one inside the null check: `TruncTimestamp` is + // `(format, timestamp)`, so a NULL format never reaches the timestamp; `TruncDate` is + // `(date, format)`, so the date is always evaluated. Verified on JVM 4.1.1. + let null_unit = if null_unit_converts_value { + unmatched.clone() + } else { + null.clone() + }; + // The plan builder runs before types are resolved, so only a literal can be folded here. + // Only a STRING literal resolves here. Any other literal is measured as its string form, and + // that form is not known at this stage: `X'59454152'` is the string `YEAR` once Spark coerces + // it, so treating a non-string literal as "matches nothing" would answer NULL for a unit that + // names a granularity. Let it fall through to the dispatch instead, where the coercion + // happens -- a literal is foldable, so the unmatched branch there is the bare NULL anyway and + // the value is still not converted. + if let Some( + ScalarValue::Utf8(Some(literal)) + | ScalarValue::LargeUtf8(Some(literal)) + | ScalarValue::Utf8View(Some(literal)), + ) = folded_literal(&unit) + { + // The same matcher the columnar path runs, so a literal and a column resolve a unit + // identically. A granularity this function does not accept is filtered out here for the + // same reason the dispatch below omits it. + let granularity = SparkTruncLevel::level(literal).filter(|level| levels.contains(level)); + return granularity.map_or(null, truncate); + } + // The unit resolves to its granularity in ONE evaluation. Doing it with expressions cannot: + // the dispatch has to tell three outcomes apart -- named a granularity, was NULL, matched + // nothing -- and `CASE x WHEN NULL` never matches, so a NULL selector shares the `ELSE` with + // an unmatched one. Splitting them in the plan (`coalesce`, `IS NULL`) mentions the unit + // twice, and DataFusion does not share a volatile subtree between the mentions, so a unit + // like `IF(rand() < 0.5, NULL, 'YEAR')` gets drawn once per mention. + let level = ScalarUDF::from(SparkTruncLevel::new()).call(vec![unit]); + Expr::Case(expr::Case { + expr: Some(Box::new(level)), + when_then_expr: levels + .iter() + .map(|level| (Box::new(lit(*level)), Box::new(truncate(level)))) + .chain(std::iter::once(( + Box::new(lit(TRUNC_LEVEL_NULL)), + Box::new(null_unit), + ))) + .collect(), + // A granularity this function does not accept -- `trunc` takes only the date-level ones -- + // resolves to a level that no arm above lists, and lands here with the unmatched units. + else_expr: Some(Box::new(unmatched)), + }) +} + +/// The branch a unit that matches nothing takes when the unit is NOT a literal: Spark has already +/// evaluated the value by then, so the conversion has to happen and only then yield NULL. +/// +/// `IS NULL` evaluates the value **once** and throws the answer away; both branches are the same +/// NULL, so the row is NULL whatever the value held. `nullif(value, value)` would evaluate it +/// TWICE — DataFusion does not hoist a volatile subtree into a common expression +/// (`common_subexpr_eliminate.rs`, `!node.is_volatile_node()`), so a value built from `rand()` +/// draws a different number per copy, the two do not compare equal, and the row keeps a value +/// where Spark yields NULL. +/// +/// This leans on the value being nullable: against a non-nullable one the simplifier rewrites +/// `IsNull(e)` to `false` and can then drop `e` with the branch, taking the conversion — and the +/// ANSI error that rides on it — with it. Every value reaching here is the result of a UDF call +/// or a cast, both nullable, so it holds today; a future path that converts to a non-nullable +/// value would need a different guard. +fn null_after_converting(value: Expr, null: Expr) -> Expr { Expr::Case(expr::Case { expr: None, - when_then_expr, + when_then_expr: vec![(Box::new(value.is_null()), Box::new(null.clone()))], else_expr: Some(Box::new(null)), }) } -fn trunc(date: Expr, part: Expr) -> Expr { - truncate_by_unit( +/// The literal an expression folds to, looking through the casts that only give it a type. Spark +/// folds those too, so `CAST(NULL AS STRING)` and an integer unit resolve here rather than +/// becoming a per-row match. +fn folded_literal(expr: &Expr) -> Option<&ScalarValue> { + match expr { + Expr::Literal(value, _) => Some(value), + Expr::Cast(expr::Cast { expr, .. }) | Expr::TryCast(expr::TryCast { expr, .. }) => { + folded_literal(expr) + } + _ => None, + } +} + +/// Whether an expression is a NULL that Spark would fold. `is_null_literal` sees the bare literal +/// alone, and a unit written as `CAST(NULL AS STRING)` has to short-circuit just like a bare NULL. +fn is_folded_null(expr: &Expr) -> bool { + folded_literal(expr).is_some_and(ScalarValue::is_null) +} + +/// Spark reports an argument it cannot use for its TYPE while it analyzes the query, as +/// `DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE`. A Spark client matches on that error class, so the +/// rejection has to carry it rather than a plain invalid-argument error. +fn unexpected_input_type_err( + function: &str, + parameter: &str, + required: &str, + actual: &DataType, +) -> PlanError { + // `AnalysisError`, not `invalid`: Spark raises this while it analyzes the query, and a Spark + // client sees the throwable class. `PlanError::InvalidArgument` surfaces as + // `IllegalArgumentException`, where Spark raises `AnalysisException`. + PlanError::AnalysisError(format!( + "[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve `{function}` due to data type \ + mismatch: The {parameter} parameter requires the \"{required}\" type, however the input \ + has the type \"{actual}\"." + )) +} + +/// Spark requires a STRING unit and rejects a complex one while it analyzes the query. Any other +/// atomic type is measured as its string form and simply matches no granularity. +/// +/// A unit that is already a string is handed back untouched. Casting it would rewrite a +/// `Utf8View` -- what Parquet produces by default -- into `Utf8`, copying every value and giving +/// up the view fast path in `upper`, all to reach a type it already had. +fn coerce_unit( + function: &str, + parameter: &str, + unit: Expr, + schema: &DFSchemaRef, +) -> PlanResult { + let data_type = unit.get_type(schema)?; + match (&data_type).into() { + // Only the plain string types go through untouched. An encoding that merely WRAPS a + // string still has to be unwrapped: the matcher's signature brings a dictionary down to a + // string but not a run-end-encoded one, and rejecting a unit for its layout is not a + // thing Spark does. Casting a `Utf8View` here would copy every value to reach a type it + // already is, which is why the plain ones are excluded. + // + // The run-end-encoded case is traced, not exercised: the tests reach a dictionary through + // a pandas categorical, and nothing in the Python client produces a run-end-encoded + // column. It mirrors the value path, which IS covered. + NativeType::String + if matches!( + data_type, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) => + { + Ok(unit) + } + NativeType::String => Ok(cast(unit, DataType::Utf8)), + NativeType::List(_) + | NativeType::FixedSizeList(_, _) + | NativeType::Struct(_) + | NativeType::Map(_) + | NativeType::Union(_) => Err(unexpected_input_type_err( + function, parameter, "STRING", &data_type, + )), + // `try_cast`, not `cast`: Spark's coercion to string never fails -- `Cast(binary, string)` + // is `UTF8String.fromBytes`, which does not validate UTF-8 -- so an argument Arrow refuses + // to convert has to become a unit that matches nothing, not an error. Spark returns NULL + // for `date_trunc(X'FF', ts)`. + _ => Ok(try_cast(unit, DataType::Utf8)), + } +} + +fn trunc(input: ScalarFunctionInput) -> PlanResult { + let ansi_mode = input.function_context.plan_config.ansi_mode; + let schema = input.function_context.schema; + let (date, part) = input.arguments.two()?; + let part = coerce_unit("trunc", "second", part, schema)?; + // Spark's `trunc` coerces its value argument to DateType, and does so before it validates the + // unit: the conversion errors under ANSI and is NULL when ANSI is off, whether or not the + // unit names a granularity. The type is gated first, and not only for the error class: Arrow + // REINTERPRETS an integer as a day count, so `trunc(1, 'YEAR')` would quietly answer + // 1970-01-01 where Spark rejects the query. + let data_type = date.get_type(schema)?; + let date = match (&data_type).into() { + // Strings parse with Spark's own grammar, which takes `2009-2-12` and a bare `2024` that + // Arrow's fixed-position parser rejects. This is the parser Sail's `CAST(… AS DATE)` + // already uses, so the two surfaces agree. + NativeType::String => ScalarUDF::from(SparkDate::new(!ansi_mode)).call(vec![date]), + NativeType::Null | NativeType::Date | NativeType::Timestamp(_, _) => { + if ansi_mode { + cast(date, DataType::Date32) + } else { + try_cast(date, DataType::Date32) + } + } + _ => { + return Err(unexpected_input_type_err( + "trunc", "first", "DATE", &data_type, + )); + } + }; + // Truncate through a naive microsecond timestamp rather than handing the DATE straight to the + // kernel. `SparkDateTrunc` only takes over for a microsecond timestamp; a DATE falls through + // to DataFusion, which coerces it to NANOsecond and multiplies days out with an unchecked + // `unary`, so `trunc(DATE '2300-01-01', 'YEAR')` wraps around into the past instead of + // truncating. A DATE carries no zone, so converting it to a naive timestamp is exact. + let naive = cast(date, DataType::Timestamp(TimeUnit::Microsecond, None)); + Ok(truncate_by_unit( part, - TRUNC_UNITS, + TRUNC_LEVELS, + // Truncating goes through the same kernel `date_trunc` uses, which carries Spark's + // nullability: `trunc` is nullable whatever its arguments are, because a unit that names + // no granularity turns the row NULL. |granularity| { cast( - expr_fn::date_trunc(lit(granularity), date.clone()), + ScalarUDF::from(SparkDateTrunc::new()).call(vec![lit(granularity), naive.clone()]), DataType::Date32, ) }, lit(ScalarValue::Date32(None)), - ) + null_after_converting(naive.clone(), lit(ScalarValue::Date32(None))), + // `TruncDate` is `(date, format)`: the date is the left child, so it is evaluated even + // when the format is NULL. + true, + )) } fn date_trunc(input: ScalarFunctionInput) -> PlanResult { let session_tz = input.function_context.plan_config.session_timezone.clone(); let ansi_mode = input.function_context.plan_config.ansi_mode; + let schema = input.function_context.schema; let (part, timestamp) = input.arguments.two()?; + let part = coerce_unit("date_trunc", "first", part, schema)?; // A string that is not a timestamp errors under ANSI, and is NULL when ANSI is off. let to_timestamp = |expr: Expr| { let timestamp = DataType::Timestamp(TimeUnit::Microsecond, Some(session_tz.clone())); @@ -139,25 +370,67 @@ fn date_trunc(input: ScalarFunctionInput) -> PlanResult { // Spark's date_trunc coerces its value argument to TimestampType and always returns // TimestampType, whether the input is DATE, TIMESTAMP, TIMESTAMP_NTZ, or a string. // Leave other types untouched so they are rejected like Spark instead of silently coerced. - let timestamp = match timestamp.get_type(input.function_context.schema)? { - DataType::Null - | DataType::Date32 - | DataType::Date64 - | DataType::Timestamp(_, _) - | DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Utf8View => to_timestamp(timestamp), + // The match is on the LOGICAL type, so an encoding that only wraps one of these -- a + // dictionary or a run-end-encoded string -- is accepted rather than rejected for its layout. + let data_type = timestamp.get_type(schema)?; + let timestamp = match (&data_type).into() { + // A DATE and a TIMESTAMP_NTZ are both wall clocks, and turning a wall clock into an + // instant is the conversion Arrow's cast gets wrong: it resolves the local time with + // `.single()`, which has no answer inside a DST gap or inside the repeated hour of a + // fall-back. Spark moves the first forward and takes the earlier offset for the second, + // which is what `SparkTimestamp` does. A DATE reaches it as a naive timestamp, a + // conversion that is pure arithmetic and involves no zone: midnight of that date. Zones + // whose transition happens AT midnight (`America/Havana`) make the DATE path fail + // exactly like the NTZ one. + NativeType::Date | NativeType::Timestamp(_, None) => { + let naive = cast(timestamp, DataType::Timestamp(TimeUnit::Microsecond, None)); + ScalarUDF::from(SparkTimestamp::try_new( + Some(session_tz.clone()), + ansi_mode, + false, + )?) + .call(vec![naive]) + } + // Strings parse with Spark's own grammar, which takes `2009-2-12` and a bare `2024`. + // Arrow's parser demands a fixed-position `YYYY-MM-DD` and at least 10 bytes, so it + // rejects both -- and Sail's `CAST(… AS TIMESTAMP)` already routes here, so leaving this + // on the Arrow parser made the two surfaces disagree with each other. + // + // `SparkTimestamp` declares a user-defined signature, which does no coercion of its own, + // so an encoding that merely WRAPS a string -- a dictionary, a run-end-encoded column -- + // has to be unwrapped here or the call is rejected for its layout. Only those are cast: + // casting a `Utf8View` would copy every value to reach a type it already is. + NativeType::String => { + let value = if matches!( + data_type, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + timestamp + } else { + cast(timestamp, DataType::Utf8) + }; + ScalarUDF::from(SparkTimestamp::try_new( + Some(session_tz.clone()), + ansi_mode, + false, + )?) + .call(vec![value]) + } + NativeType::Null | NativeType::Timestamp(_, _) => to_timestamp(timestamp), // Spark rejects the value for its type when it analyzes the query, so the rejection // must not depend on the unit: an unrecognized or NULL unit does not turn it into NULL. - other => { - return Err(PlanError::invalid(format!( - "`date_trunc` does not support {other} input" - ))); + _ => { + return Err(unexpected_input_type_err( + "date_trunc", + "second", + "TIMESTAMP", + &data_type, + )); } }; let truncated = truncate_by_unit( part, - DATE_TRUNC_UNITS, + DATE_TRUNC_LEVELS, |granularity| { ScalarUDF::from(SparkDateTrunc::new()).call(vec![lit(granularity), timestamp.clone()]) }, @@ -165,6 +438,16 @@ fn date_trunc(input: ScalarFunctionInput) -> PlanResult { None, Some(session_tz.clone()), )), + null_after_converting( + timestamp.clone(), + lit(ScalarValue::TimestampMicrosecond( + None, + Some(session_tz.clone()), + )), + ), + // `TruncTimestamp` is `(format, timestamp)`: the format is the left child, so a NULL + // format short-circuits and the timestamp is never evaluated. + false, ); let truncated = match truncated.get_type(input.function_context.schema)? { DataType::Timestamp(TimeUnit::Microsecond, _) => truncated, @@ -1246,7 +1529,7 @@ pub(super) fn list_built_in_datetime_functions() -> Vec<(&'static str, ScalarFun ), ("to_unix_timestamp", F::custom(to_unix_timestamp)), ("to_utc_timestamp", F::custom(to_utc_timestamp)), - ("trunc", F::binary(trunc)), + ("trunc", F::custom(trunc)), ("try_make_interval", F::unknown("try_make_interval")), ( "try_make_timestamp", diff --git a/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml b/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml index 2439297289..31505b9ea1 100644 --- a/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml +++ b/python/pysail/tests/spark/__snapshots__/test_clickbench.plan.yaml @@ -593,12 +593,12 @@ data: ProjectionExec: expr=[#105@0 as M, #106@1 as PageViews] ProjectionExec: expr=[#105@0 as #105, #106@1 as #106] GlobalLimitExec: skip=1000, fetch=10 - SortPreservingMergeExec: [CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@2 ASC], fetch=1010 + SortPreservingMergeExec: [date_trunc(Utf8("minute"),hits.#4)@2 ASC], fetch=1010 SortExec: TopK(fetch=1010), expr=[#105@0 ASC], preserve_partitioning=[true] - ProjectionExec: expr=[CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0 as #105, count(Int64(1))@1 as #106, CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0 as CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END] - AggregateExec: mode=FinalPartitioned, gby=[CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0 as CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END], aggr=[count(Int64(1))] - RepartitionExec: partitioning=Hash([CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END@0], 4), input_partitions=4 - AggregateExec: mode=Partial, gby=[date_trunc(minute, #4@0) as CASE WHEN lower(Utf8("minute")) = Utf8("year") OR lower(Utf8("minute")) = Utf8("yyyy") OR lower(Utf8("minute")) = Utf8("yy") THEN date_trunc(Utf8("year"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("quarter") THEN date_trunc(Utf8("quarter"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("month") OR lower(Utf8("minute")) = Utf8("mon") OR lower(Utf8("minute")) = Utf8("mm") THEN date_trunc(Utf8("month"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("week") THEN date_trunc(Utf8("week"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("day") OR lower(Utf8("minute")) = Utf8("dd") THEN date_trunc(Utf8("day"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("hour") THEN date_trunc(Utf8("hour"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("minute") THEN date_trunc(Utf8("minute"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("second") THEN date_trunc(Utf8("second"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("millisecond") THEN date_trunc(Utf8("millisecond"),hits.#4) WHEN lower(Utf8("minute")) = Utf8("microsecond") THEN date_trunc(Utf8("microsecond"),hits.#4) ELSE TimestampMicrosecond(NULL, Some("UTC")) END], aggr=[count(Int64(1))] + ProjectionExec: expr=[date_trunc(Utf8("minute"),hits.#4)@0 as #105, count(Int64(1))@1 as #106, date_trunc(Utf8("minute"),hits.#4)@0 as date_trunc(Utf8("minute"),hits.#4)] + AggregateExec: mode=FinalPartitioned, gby=[date_trunc(Utf8("minute"),hits.#4)@0 as date_trunc(Utf8("minute"),hits.#4)], aggr=[count(Int64(1))] + RepartitionExec: partitioning=Hash([date_trunc(Utf8("minute"),hits.#4)@0], 4), input_partitions=4 + AggregateExec: mode=Partial, gby=[date_trunc(minute, #4@0) as date_trunc(Utf8("minute"),hits.#4)], aggr=[count(Int64(1))] ProjectionExec: expr=[CAST(EventTime@0 * 1000000 AS Timestamp(µs, "UTC")) as #4] FilterExec: CounterID@2 = 62 AND EventDate@1 >= 15900 AND EventDate@1 <= 15901 AND IsRefresh@3 = 0 AND DontCountHits@4 = 0, projection=[EventTime@0] RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 diff --git a/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml b/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml index 0d621e824a..50286b5749 100644 --- a/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml +++ b/python/pysail/tests/spark/function/__snapshots__/features/date_trunc.yaml @@ -13,6 +13,5 @@ data: |- == Physical Plan == ProjectionExec: expr=[#4@0 as result] - ProjectionExec: expr=[CASE WHEN __common_expr_1@0 = year OR __common_expr_1@0 = yyyy OR __common_expr_1@0 = yy THEN date_trunc(year, #2@1) WHEN __common_expr_1@0 = quarter THEN date_trunc(quarter, #2@1) WHEN __common_expr_1@0 = month OR __common_expr_1@0 = mon OR __common_expr_1@0 = mm THEN date_trunc(month, #2@1) WHEN __common_expr_1@0 = week THEN date_trunc(week, #2@1) WHEN __common_expr_1@0 = day OR __common_expr_1@0 = dd THEN date_trunc(day, #2@1) WHEN __common_expr_1@0 = hour THEN date_trunc(hour, #2@1) WHEN __common_expr_1@0 = minute THEN date_trunc(minute, #2@1) WHEN __common_expr_1@0 = second THEN date_trunc(second, #2@1) WHEN __common_expr_1@0 = millisecond THEN date_trunc(millisecond, #2@1) WHEN __common_expr_1@0 = microsecond THEN date_trunc(microsecond, #2@1) END as #4] - ProjectionExec: expr=[lower(column2@1) as __common_expr_1, column1@0 as #2] - DataSourceExec: partitions=1, partition_sizes=[1] + ProjectionExec: expr=[CASE spark_trunc_level(column2@1) WHEN year THEN date_trunc(year, column1@0) WHEN quarter THEN date_trunc(quarter, column1@0) WHEN month THEN date_trunc(month, column1@0) WHEN week THEN date_trunc(week, column1@0) WHEN day THEN date_trunc(day, column1@0) WHEN hour THEN date_trunc(hour, column1@0) WHEN minute THEN date_trunc(minute, column1@0) WHEN second THEN date_trunc(second, column1@0) WHEN millisecond THEN date_trunc(millisecond, column1@0) WHEN microsecond THEN date_trunc(microsecond, column1@0) WHEN null THEN NULL ELSE CASE WHEN column1@0 IS NULL THEN NULL END END as #4] + DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/python/pysail/tests/spark/function/features/date_trunc.feature b/python/pysail/tests/spark/function/features/date_trunc.feature index 48964fcde1..39f88df15b 100644 --- a/python/pysail/tests/spark/function/features/date_trunc.feature +++ b/python/pysail/tests/spark/function/features/date_trunc.feature @@ -1,9 +1,11 @@ @date_trunc -Feature: DATE_TRUNC and TRUNC truncate to a unit +Feature: DATE_TRUNC truncates a timestamp to a unit # `date_trunc` (TruncTimestamp) and `trunc` (TruncDate) are different Spark expressions and do # NOT accept the same units: `trunc` only truncates to date-level units, and anything finer -- - # or unrecognized, or NULL -- is NULL, not an error. They are covered together because a single - # unit table in the plan builder governs both. + # or unrecognized, or NULL -- is NULL, not an error. Each has its own unit table in the plan + # builder, and they share the matching machinery, so `trunc.feature` covers the same ground for + # the date expression and the two must be kept in step. + # All expected values were captured on Spark JVM 4.1.1. Rule: date_trunc returns timestamp for every date-family input type @@ -272,6 +274,13 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit """ SELECT date_trunc(u, TIMESTAMP '2026-02-02 10:11:12') AS result FROM VALUES ('YEAR'), ('MONTH') AS t(u) """ + # The columnar path builds its own result type rather than reusing the one a literal unit + # produces, and the optimizer may drop the branch that carries it, so pin it down here too. + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ Then query result ordered | result | | 2026-01-01 00:00:00 | @@ -311,11 +320,19 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit Rule: null value returns null + # A NULL prints the same whatever its type, so every scenario here also pins the type down: + # a result typed `void` or `timestamp_ntz` would otherwise pass unnoticed. + Scenario: date_trunc on a null date returns null When query """ SELECT date_trunc('DAY', CAST(NULL AS DATE)) AS result """ + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ Then query result | result | | NULL | @@ -325,6 +342,11 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit """ SELECT date_trunc('DAY', CAST(NULL AS TIMESTAMP)) AS result """ + Then query schema + """ + root + |-- result: timestamp (nullable = true) + """ Then query result | result | | NULL | @@ -356,21 +378,29 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit Rule: an unparseable string value under ANSI on errors + # Sail raises here, but carries DataFusion's parser message ("Error parsing timestamp from + # ...") instead of Spark's CAST_INVALID_INPUT class. That is the ANSI string-to-timestamp + # cast surface, not this expression: every function that casts a string under ANSI reports + # it the same way. Asserting on Spark's class keeps the gap visible instead of letting a + # blank `.*` pass anything. + + @sail-bug Scenario: an unparseable string value errors under ANSI on Given config spark.sql.ansi.enabled = true When query """ SELECT date_trunc('DAY', 'not a date') AS result """ - Then query error .* + Then query error CAST_INVALID_INPUT + @sail-bug Scenario: an empty string value errors under ANSI on Given config spark.sql.ansi.enabled = true When query """ SELECT date_trunc('DAY', '') AS result """ - Then query error .* + Then query error CAST_INVALID_INPUT Rule: an unparseable string value under ANSI off returns NULL @@ -409,53 +439,59 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit Rule: date_trunc rejects non-date input types # The value is rejected for its TYPE, whatever the unit is: an unrecognized or NULL unit - # does not turn the rejection into a NULL. + # does not turn the rejection into a NULL. Spark raises this while it analyzes the query, so + # the scenarios match on Spark's own error class rather than on any error at all: a Sail + # error of a different class is a different observable behaviour for a Spark client. Scenario: date_trunc rejects a time value When query """ SELECT date_trunc('DAY', TIME '12:30:00') AS result """ - Then query error .* + Then query error DATATYPE_MISMATCH Scenario: date_trunc rejects an integer value When query """ SELECT date_trunc('DAY', 1) AS result """ - Then query error .* + Then query error DATATYPE_MISMATCH Scenario: date_trunc rejects an integer value even when the unit is unrecognized When query """ SELECT date_trunc('INVALID', 1) AS result """ - Then query error .* + Then query error DATATYPE_MISMATCH Scenario: date_trunc rejects an integer value even when the unit is null When query """ SELECT date_trunc(CAST(NULL AS STRING), 1) AS result """ - Then query error .* + Then query error DATATYPE_MISMATCH Scenario: date_trunc rejects a boolean value even when the unit is unrecognized When query """ SELECT date_trunc('INVALID', true) AS result """ - Then query error .* + Then query error DATATYPE_MISMATCH Scenario: date_trunc rejects a time value even when the unit is null When query """ SELECT date_trunc(NULL, TIME '12:30:00') AS result """ - Then query error .* + Then query error DATATYPE_MISMATCH Rule: date boundaries and leap days - @sail-bug + # Truncating a timestamp whose microseconds do not fit in a 64-bit nanosecond count + # overflows. The bound is not the maximum year but the granularity: a granularity that has + # to be resolved against a time zone goes through the nanosecond conversion, so the year at + # which it breaks depends on which units take that path, not on the type's range. + Scenario: date_trunc YEAR at the maximum supported year When query """ @@ -465,6 +501,42 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | result | | 9999-01-01 00:00:00 | + Scenario: date_trunc DAY beyond the nanosecond range + When query + """ + SELECT date_trunc('DAY', TIMESTAMP '2300-01-01 05:30:00 UTC') AS result + """ + Then query result + | result | + | 2300-01-01 00:00:00 | + + Scenario: date_trunc HOUR beyond the nanosecond range + When query + """ + SELECT date_trunc('HOUR', TIMESTAMP '2300-01-01 05:30:00 UTC') AS result + """ + Then query result + | result | + | 2300-01-01 05:00:00 | + + Scenario: date_trunc SECOND beyond the nanosecond range + When query + """ + SELECT date_trunc('SECOND', TIMESTAMP '2300-01-01 05:30:00 UTC') AS result + """ + Then query result + | result | + | 2300-01-01 05:30:00 | + + Scenario: date_trunc DAY at the maximum supported year + When query + """ + SELECT date_trunc('DAY', TIMESTAMP '9999-12-31 23:59:59 UTC') AS result + """ + Then query result + | result | + | 9999-12-31 00:00:00 | + Scenario: date_trunc DAY on a leap day When query """ @@ -544,110 +616,608 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | result | | 1 | - Rule: trunc truncates to a date-level unit + Rule: the unit is resolved once per row, even when it is not foldable - Scenario: trunc to the week + # A unit built with `rand` is volatile: matching it once per candidate unit draws a new value + # for every comparison, so a row can match nothing and turn NULL. Spark resolves the unit once + # per row and never yields NULL here. A unit from `VALUES` is folded by Spark and hides this. + + Scenario: a volatile unit never nullifies a row When query - """ - SELECT trunc(DATE '2019-08-04', 'week') AS result - """ + """ + SELECT count(*) AS nulls FROM ( + SELECT date_trunc(IF(rand(7) < 0.5, 'YEAR', 'MONTH'), TIMESTAMP '2024-05-15 13:45:30') AS r + FROM range(10000) + ) WHERE r IS NULL + """ Then query result - | result | - | 2019-07-29 | + | nulls | + | 0 | + + # The scenario above cannot see a unit that is evaluated more than once: both draws name a + # granularity, so any number of evaluations still truncates. Draw a NULL against a valid unit + # instead and the count itself becomes the evidence -- one evaluation gives ~5000, and every + # extra one pushes it up (two give ~7500, since the row turns NULL if EITHER draw does). - Scenario: trunc to the quarter + # The count is asserted as a RANGE, not a value: `rand` is seeded, but nothing says the two + # engines draw the same sequence from that seed. The range is far tighter than the gap it has + # to separate -- one evaluation lands near 5000, two near 7500. + + Scenario: a volatile unit is drawn exactly once per row When query - """ - SELECT trunc(DATE '2019-08-04', 'quarter') AS result - """ + """ + SELECT count(*) BETWEEN 4700 AND 5300 AS drawn_once FROM ( + SELECT date_trunc(IF(rand(7) < 0.5, CAST(NULL AS STRING), 'YEAR'), + TIMESTAMP '2024-05-15 13:45:30') AS r + FROM range(10000) + ) WHERE r IS NULL + """ Then query result - | result | - | 2019-07-01 | + | drawn_once | + | true | + + Rule: the value is converted even when the unit matches nothing + + # Spark evaluates both arguments before it validates the unit, so an unconvertible value + # raises under ANSI even on a row whose unit matches no granularity. Resolving the unit first + # and skipping the conversion for unmatched rows would silently return NULL instead. Both + # arguments are non-foldable so that Spark cannot fold the rows away. + + # Two scenarios, because one assertion cannot carry both claims. The tagged one states the + # target -- Spark's error CLASS -- and xfails until Sail's string cast reports it. But an + # xfail is satisfied by ANY failure, including Sail simply returning NULL, so on its own it + # would not notice this rule breaking. The untagged one below matches what each engine + # actually prints, so it holds on BOTH and fails the moment either stops raising. + @sail-bug + Scenario: an unconvertible value reports Spark's error class when its unit matches nothing + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc(IF(id = 0, 'INVALID', 'DAY'), IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query error CAST_INVALID_INPUT - Scenario: trunc matches the unit case-insensitively + Scenario: an unconvertible value errors under ANSI on even when its unit matches nothing + Given config spark.sql.ansi.enabled = true When query - """ - SELECT trunc(DATE '2019-08-04', 'MoN') AS result - """ + """ + SELECT id, date_trunc(IF(id = 0, 'INVALID', 'DAY'), IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + Scenario: an unconvertible value is NULL under ANSI off even when its unit matches nothing + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT id, date_trunc(IF(id = 0, 'INVALID', 'DAY'), IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | 2024-01-02 00:00:00 | + + # A LITERAL unit that names no granularity is the other half of the rule: Spark's codegen + # short-circuits to NULL without evaluating the value at all, so the same unconvertible value + # does NOT raise here. Converting the value unconditionally would turn these into errors. + + Scenario: a literal unit that matches nothing skips the value entirely under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc('INVALID', IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | NULL | + + Scenario: a literal NULL unit skips the value entirely under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc(CAST(NULL AS STRING), IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | NULL | + + @sail-bug + Scenario: a literal unit that names a granularity reports Spark's error class + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc('DAY', IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query error CAST_INVALID_INPUT + + Scenario: a literal unit that names a granularity still converts the value under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc('DAY', IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + # A NULL unit is NOT the same case as a unit that matches nothing, and which one skips the + # value follows the argument ORDER of the Spark expression: `TruncTimestamp` is + # `(format, timestamp)`, so a NULL format never reaches the timestamp and no ANSI error can + # come out of it. `trunc` is `(date, format)` and does the opposite -- see trunc.feature. + + Scenario: a NULL unit in a column never converts the value + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc(IF(id = 0, CAST(NULL AS STRING), 'DAY'), + IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | 2024-01-02 00:00:00 | + + Rule: the string "null" is an ordinary unit, not the absence of one + + # Telling "the unit was NULL" apart from "the unit matched nothing" needs a value that is not + # NULL, since NULL is already what an unmatched unit yields. A unit is upper-cased before it + # is matched, so a user unit can never come out lower case, and the marker is unreachable from + # data -- but that is a property to demonstrate, not to assume: the string "null" has to + # behave like any other unit that names no granularity, INCLUDING converting the value. + + Scenario: the string null is just a unit that matches nothing + When query + """ + SELECT id, date_trunc(IF(id >= 0, 'null', 'x'), TIMESTAMP '2024-05-15 13:45:30 UTC') AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | NULL | + + Scenario: the string null converts the value, unlike an actual NULL unit + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc(IF(id >= 0, 'null', 'x'), IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + Scenario: an actual NULL unit skips the value where the string null does not + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, date_trunc(IF(id >= 0, CAST(NULL AS STRING), 'x'), + IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | NULL | + + Rule: a unit of a type that cannot become a string matches nothing + + # Spark's coercion to string never fails -- `Cast(binary, string)` is `UTF8String.fromBytes`, + # which does not validate UTF-8 -- so a unit Arrow refuses to convert has to become a unit + # that matches nothing, not an error. + + Scenario: a binary unit is NULL, not an error + When query + """ + SELECT date_trunc(X'FF', TIMESTAMP '2024-01-01 00:00:00 UTC') AS result + """ Then query result - | result | - | 2019-08-01 | + | result | + | NULL | - Scenario: trunc accepts the yy alias + # But a unit is measured as its STRING form, and a binary one has a perfectly good string + # form when its bytes spell a unit. Treating any non-string literal as "matches nothing" + # would answer NULL here, so a literal only short-circuits when it is already a string. + + Scenario: a binary unit whose bytes spell a granularity truncates When query - """ - SELECT trunc(DATE '2019-08-04', 'yy') AS result - """ + """ + SELECT date_trunc(X'59454152', TIMESTAMP '2024-05-15 13:45:30 UTC') AS result + """ Then query result - | result | - | 2019-01-01 | + | result | + | 2024-01-01 00:00:00 | - Scenario: trunc takes the date from a column + # Sail turns a binary unit it cannot decode into a NULL unit, which takes the short circuit, + # where Spark turns it into a string that matches nothing and therefore converts the value. + # Only observable with a NON-foldable unit, an unconvertible value and ANSI on. + @sail-bug + Scenario: an undecodable binary unit converts the value under ANSI on + Given config spark.sql.ansi.enabled = true When query - """ - SELECT trunc(d, 'month') AS result FROM VALUES (1, DATE '2019-08-04'), (2, DATE '2020-02-29') AS t(i, d) ORDER BY i - """ - Then query result ordered - | result | - | 2019-08-01 | - | 2020-02-01 | + """ + SELECT id, date_trunc(IF(id >= 0, X'FF', X'00'), IF(id = 0, 'not a date', '2024-01-02')) AS result + FROM range(2) ORDER BY id + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + Rule: a unit from a scalar subquery is resolved before the query runs - Rule: a unit finer than a day is NULL, not a truncation + # `ScalarSubquery.foldable` is false in Spark, but a subquery over no relation is resolved + # into a literal before codegen, so it takes the same short circuit a literal does: no ANSI + # error comes out of the value even though the unit matches nothing. - Scenario: trunc to the day is NULL + Scenario: a subquery unit that matches nothing skips the value + Given config spark.sql.ansi.enabled = true When query - """ - SELECT trunc(DATE '2019-08-04', 'DAY') AS result - """ + """ + SELECT date_trunc((SELECT 'BOGUS'), 'not a date') AS result + """ Then query result - | result | - | NULL | + | result | + | NULL | - Scenario: trunc to the hour is NULL + Scenario: a subquery unit still names its granularity When query - """ - SELECT trunc(DATE '2019-08-04', 'HOUR') AS result - """ + """ + SELECT date_trunc((SELECT 'DAY'), TIMESTAMP '2024-05-15 13:45:30 UTC') AS result + """ Then query result - | result | - | NULL | + | result | + | 2024-05-15 00:00:00 | + + Rule: a zoned timestamp truncates across a DST transition - Rule: an unrecognized or NULL unit is NULL, not an error + # Every other DST scenario feeds a TIMESTAMP_NTZ or a DATE, which reach the localization the + # plan builder inserts. A zoned timestamp skips that path entirely and is truncated by the + # kernel alone, so it needs its own case -- and the instant, not just the wall clock, because + # the day it belongs to is the one that lasts 25 hours. - Scenario: trunc with an unrecognized unit is NULL + Scenario: truncating a zoned timestamp to the day inside the fall-back day + Given config spark.sql.session.timeZone = America/Los_Angeles When query - """ - SELECT trunc(DATE '2019-08-04', 'bogus') AS result - """ + """ + SELECT date_trunc('DAY', TIMESTAMP '2024-11-03 10:30:00 UTC') AS result + """ Then query result - | result | - | NULL | + | result | + | 2024-11-03 00:00:00 | - Scenario: trunc with a NULL unit is NULL + Scenario: the truncated zoned timestamp lands on the right instant + Given config spark.sql.session.timeZone = America/Los_Angeles When query - """ - SELECT trunc(DATE '2019-08-04', NULL) AS result - """ + """ + SELECT unix_timestamp(date_trunc('DAY', TIMESTAMP '2024-11-03 10:30:00 UTC')) AS result + """ Then query result - | result | - | NULL | + | result | + | 1730617200 | + + Rule: a unit Spark folds but the plan builder cannot - Rule: the unit may come from a column + # Spark's `foldable` covers any expression with no input, so `concat('BOG','US')` is resolved + # at compile time and takes the short circuit. Sail cannot name the granularity for one of + # these at plan-build time -- the optimizer that would evaluate it runs later -- so the unit + # is still matched per row. What it does share with a literal is the part that shows: a unit + # with no column reference and no volatility does not convert the value, so no ANSI error + # comes out of a row it fails to match. - # The plan builder only matches `Expr::Literal`, so a unit in a column falls through to - # DataFusion's `date_trunc`, which demands a scalar. Sail errors: Granularity of `date_trunc` - # must be non-null scalar Utf8. - Scenario: trunc resolves the unit of each row + Scenario: a unit folded from an expression short-circuits like a literal + Given config spark.sql.ansi.enabled = true When query - """ - SELECT trunc(DATE '2019-08-04', c) AS result FROM VALUES (1, 'year'), (2, 'month'), (3, 'DAY'), (4, NULL) AS t(i, c) ORDER BY i - """ + """ + SELECT date_trunc(concat('BOG','US'), 'not a date') AS result + """ + Then query result + | result | + | NULL | + + Scenario: a unit folded from an expression still names its granularity + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT date_trunc(concat('DA','Y'), TIMESTAMP '2024-05-15 13:45:30') AS result + """ + Then query result + | result | + | 2024-05-15 00:00:00 | + + Rule: truncation is resolved in the session time zone + + # Every other scenario runs under the harness default of UTC, where a session time zone that + # was ignored, or hardcoded to UTC, would pass unnoticed. These pin the plumbing down by + # truncating across a DST transition, where the local time zone changes the answer. + + Scenario: a timestamp_ntz in the spring-forward gap moves to the next valid time + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('SECOND', TIMESTAMP_NTZ '2024-03-10 02:30:00') AS result + """ + Then query result + | result | + | 2024-03-10 03:30:00 | + + Scenario: a timestamp_ntz in the fall-back overlap takes the earlier time + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('SECOND', TIMESTAMP_NTZ '2024-11-03 01:30:00') AS result + """ + Then query result + | result | + | 2024-11-03 01:30:00 | + + # The rendered wall clock CANNOT tell the two resolutions apart: 01:30 during the overlap + # prints as 01:30 whether it is resolved to the earlier offset or the later one, so the + # scenario above would pass just as well if the earlier one were dropped. The instant does + # tell them apart -- an hour lies between them. + + Scenario: the fall-back overlap resolves to the earlier of the two instants + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT unix_timestamp(date_trunc('SECOND', TIMESTAMP_NTZ '2024-11-03 01:30:00')) AS result + """ + Then query result + | result | + | 1730622600 | + + Scenario: truncating to the hour inside the spring-forward gap + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('HOUR', TIMESTAMP_NTZ '2024-03-10 02:30:00') AS result + """ + Then query result + | result | + | 2024-03-10 03:00:00 | + + Scenario: truncating to the hour inside the fall-back overlap + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('HOUR', TIMESTAMP_NTZ '2024-11-03 01:30:00') AS result + """ + Then query result + | result | + | 2024-11-03 01:00:00 | + + Scenario: truncating to the hour inside the overlap keeps the earlier instant + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT unix_timestamp(date_trunc('HOUR', TIMESTAMP_NTZ '2024-11-03 01:30:00')) AS result + """ + Then query result + | result | + | 1730620800 | + + Scenario: truncating a timestamp_ntz to the day across a DST transition + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('DAY', TIMESTAMP_NTZ '2024-03-10 02:30:00') AS result + """ + Then query result + | result | + | 2024-03-10 00:00:00 | + + Scenario: truncating a date under a non-UTC session time zone + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('DAY', DATE '2024-03-10') AS result + """ + Then query result + | result | + | 2024-03-10 00:00:00 | + + Scenario: truncating a timestamp_ntz from a column under a non-UTC session time zone + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('SECOND', ts) AS result FROM VALUES (TIMESTAMP_NTZ '2024-03-10 02:30:00') AS t(ts) + """ + Then query result + | result | + | 2024-03-10 03:30:00 | + + Scenario: a literal unit under a non-UTC session time zone + Given config spark.sql.session.timeZone = America/Los_Angeles + When query + """ + SELECT date_trunc('DAY', TIMESTAMP '2024-03-15 08:30:00 UTC') AS result + """ + Then query result + | result | + | 2024-03-15 00:00:00 | + + Rule: the unit is matched with Unicode case folding + + # Spark upper-cases the unit with `toUpperCase(Locale.ROOT)`, which folds non-ASCII letters: + # the dotless i folds to I and the long s folds to S. Matching only ASCII case would return + # NULL for these instead of truncating. + + Scenario: a unit spelled with the dotless i matches MINUTE + When query + """ + SELECT date_trunc('mınute', TIMESTAMP '2024-05-15 13:45:30.123456 UTC') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:00 | + + Scenario: a unit spelled with the long s matches SECOND + When query + """ + SELECT date_trunc('ſecond', TIMESTAMP '2024-05-15 13:45:30.123456 UTC') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:30 | + + Scenario: a unit spelled with the dotless i matches MICROSECOND + When query + """ + SELECT date_trunc('mıcrosecond', TIMESTAMP '2024-05-15 13:45:30.123456 UTC') AS result + """ + Then query result + | result | + | 2024-05-15 13:45:30.123456 | + + # The scenarios above resolve the unit while the plan is built, because it is a literal. A + # unit in a column is resolved per row instead. Both go through the same matcher, and this + # case is what keeps it that way. + + Scenario: a non-ASCII unit in a column folds the same way + When query + """ + SELECT date_trunc(u, TIMESTAMP '2024-05-15 13:45:30.123456 UTC') AS result + FROM VALUES ('mınute'), ('ſecond') AS t(u) + """ Then query result ordered - | result | - | 2019-01-01 | - | 2019-08-01 | - | NULL | - | NULL | + | result | + | 2024-05-15 13:45:00 | + | 2024-05-15 13:45:30 | + + Rule: a complex argument is rejected when the query is analyzed + + # Spark rejects an argument for its TYPE while it analyzes the query. Coercing a complex type + # to its string form instead would match no granularity and quietly return NULL. + + Scenario: date_trunc rejects an array unit + When query + """ + SELECT date_trunc(array('YEAR'), TIMESTAMP '2024-01-01 00:00:00 UTC') AS result + """ + Then query error DATATYPE_MISMATCH + + Scenario: date_trunc rejects a map unit + When query + """ + SELECT date_trunc(map('a', 'YEAR'), TIMESTAMP '2024-01-01 00:00:00 UTC') AS result + """ + Then query error DATATYPE_MISMATCH + + Scenario: date_trunc rejects a struct unit + When query + """ + SELECT date_trunc(named_struct('a', 'YEAR'), TIMESTAMP '2024-01-01 00:00:00 UTC') AS result + """ + Then query error DATATYPE_MISMATCH + + Scenario: date_trunc rejects an array value + When query + """ + SELECT date_trunc('YEAR', array(1)) AS result + """ + Then query error DATATYPE_MISMATCH + + Rule: a string value parses with Spark's grammar + + # Spark's own parser accepts a single-digit month or day and a date with no month or day at + # all. Arrow's parser demands a fixed-position `YYYY-MM-DD` of at least 10 bytes and rejects + # all three -- and Sail's `CAST(… AS TIMESTAMP)` already uses Spark's parser, so leaving this + # on Arrow's made date_trunc disagree with CAST on the very same string. + + Scenario: a value with a single-digit month + When query + """ + SELECT date_trunc('MM', '2009-2-12') AS result + """ + Then query result + | result | + | 2009-02-01 00:00:00 | + + Scenario: a value with no month or day + When query + """ + SELECT date_trunc('YEAR', '2024') AS result + """ + Then query result + | result | + | 2024-01-01 00:00:00 | + + Scenario: a value with no day + When query + """ + SELECT date_trunc('MM', '2024-3') AS result + """ + Then query result + | result | + | 2024-03-01 00:00:00 | + + Rule: a date localizes across a transition that happens at midnight + + # America/Los_Angeles moves its clock at 02:00, so local midnight always exists there and the + # scenarios above cannot see this. In these zones the clock moves AT midnight, so the day the + # transition falls on has no 00:00 at all and Spark reports the first instant that does exist. + + Scenario: a date whose local midnight does not exist + Given config spark.sql.session.timeZone = Africa/Cairo + When query + """ + SELECT date_trunc('DAY', DATE '2024-04-26') AS result + """ + Then query result + | result | + | 2024-04-26 01:00:00 | + + Scenario: a date from a column whose local midnight does not exist + Given config spark.sql.session.timeZone = Africa/Cairo + When query + """ + SELECT date_trunc('DAY', d) AS result FROM VALUES (DATE '2024-04-26') AS t(d) + """ + Then query result + | result | + | 2024-04-26 01:00:00 | + + Scenario: a date whose local midnight does not exist in another zone + Given config spark.sql.session.timeZone = America/Havana + When query + """ + SELECT date_trunc('DAY', DATE '2024-03-10') AS result + """ + Then query result + | result | + | 2024-03-10 01:00:00 | + + Scenario: a timestamp_ntz inside a midnight gap + Given config spark.sql.session.timeZone = Africa/Cairo + When query + """ + SELECT date_trunc('DAY', TIMESTAMP_NTZ '2024-04-26 00:30:00') AS result + """ + Then query result + | result | + | 2024-04-26 01:00:00 | + + Rule: the value is converted once, even when it is volatile + + # The branch an unmatched unit takes has to convert the value -- and convert it ONCE. Reaching + # for `nullif(value, value)` would evaluate it twice, and DataFusion does not hoist a volatile + # subtree into a shared expression, so the two draws would differ and the row would keep a + # value where Spark yields NULL. This is the mirror of the volatile-unit scenario. + + Scenario: a volatile value with a unit that matches nothing is always NULL + When query + """ + SELECT count(*) AS non_nulls FROM ( + SELECT date_trunc(u, IF(rand(7) < 0.5, TIMESTAMP '2020-01-01 00:00:00', + TIMESTAMP '2021-01-01 00:00:00')) AS r + FROM (SELECT 'BOGUS' AS u FROM range(10000)) + ) WHERE r IS NOT NULL + """ + Then query result + | non_nulls | + | 0 | Rule: the plan resolves the unit @@ -703,3 +1273,58 @@ Feature: DATE_TRUNC and TRUNC truncate to a unit | result | | 2015-01-01 00:00:00 | | 2015-01-01 00:00:00 | + + Rule: a unit from a subquery Spark does not fold + + # Whether the short circuit applies depends on Spark's OPTIMIZER, not on the expression: + # `ScalarSubquery.foldable` is always false, and what decides is whether ConstantFolding + # replaced the subquery with a literal before codegen. Measured on JVM 4.1.1, it folds a + # subquery whose plan ends in a one-row relation -- `(SELECT 'BOGUS')`, a view over a + # constant, `FROM (SELECT 1)` -- and does NOT fold one over `VALUES`, `range`, or a table. + # + # Sail decides while the plan is built, before any of that, so it treats every subquery as + # foldable and skips the value. Detectable, since the boundary is structural, but it encodes + # an optimizer quirk rather than a semantic rule -- `FROM (VALUES ('x'))` not folding while + # `FROM (SELECT 1)` does answers to no principle -- so it is left for its own change. + # + # Only ANSI on diverges, and only with a value that cannot be converted: the pair below is + # what shows that, and everything else about a non-folded subquery already agrees. + + @sail-bug + Scenario: a non-folded subquery unit converts the value under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT date_trunc((SELECT 'BOGUS' FROM range(1) LIMIT 1), 'not a date') AS result + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + Scenario: a non-folded subquery unit is NULL under ANSI off + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT date_trunc((SELECT 'BOGUS' FROM range(1) LIMIT 1), 'not a date') AS result + """ + Then query result + | result | + | NULL | + + Scenario: a non-folded subquery unit that names a granularity truncates under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT date_trunc((SELECT 'DAY' FROM range(1) LIMIT 1), TIMESTAMP '2024-05-15 13:45:30 UTC') AS result + """ + Then query result + | result | + | 2024-05-15 00:00:00 | + + Scenario: a non-folded subquery unit that names a granularity truncates under ANSI off + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT date_trunc((SELECT 'DAY' FROM range(1) LIMIT 1), TIMESTAMP '2024-05-15 13:45:30 UTC') AS result + """ + Then query result + | result | + | 2024-05-15 00:00:00 | diff --git a/python/pysail/tests/spark/function/features/trunc.feature b/python/pysail/tests/spark/function/features/trunc.feature index 759e178ddd..813e641706 100644 --- a/python/pysail/tests/spark/function/features/trunc.feature +++ b/python/pysail/tests/spark/function/features/trunc.feature @@ -1,7 +1,469 @@ -Feature: trunc with an argument coming from a column - # A behaviour-governing argument given as a literal is constant-folded, so the literal - # scenarios never exercise the columnar kernel. These scenarios pass the same argument - # through a column. All expected values were captured on Spark JVM 4.1.1. +@trunc +Feature: TRUNC truncates a date to a date-level unit + # `trunc` (TruncDate) and `date_trunc` (TruncTimestamp) are different Spark expressions and do + # NOT accept the same units: `trunc` only truncates to date-level units, and anything finer -- + # or unrecognized, or NULL -- is NULL, not an error. Each has its own unit table in the plan + # builder, and they share the matching machinery, so `date_trunc.feature` covers the same ground + # for the timestamp expression. All expected values were captured on Spark JVM 4.1.1. + + Rule: trunc truncates to a date-level unit + + Scenario: trunc to the week + When query + """ + SELECT trunc(DATE '2019-08-04', 'week') AS result + """ + Then query result + | result | + | 2019-07-29 | + + Scenario: trunc to the quarter + When query + """ + SELECT trunc(DATE '2019-08-04', 'quarter') AS result + """ + Then query result + | result | + | 2019-07-01 | + + Scenario: trunc matches the unit case-insensitively + When query + """ + SELECT trunc(DATE '2019-08-04', 'MoN') AS result + """ + Then query result + | result | + | 2019-08-01 | + + Scenario: trunc accepts the yy alias + When query + """ + SELECT trunc(DATE '2019-08-04', 'yy') AS result + """ + Then query result + | result | + | 2019-01-01 | + + Scenario: trunc takes the date from a column + When query + """ + SELECT trunc(d, 'month') AS result FROM VALUES (1, DATE '2019-08-04'), (2, DATE '2020-02-29') AS t(i, d) ORDER BY i + """ + Then query result ordered + | result | + | 2019-08-01 | + | 2020-02-01 | + + Rule: trunc returns a date + + # `trunc` returns DATE, not the TIMESTAMP that `date_trunc` returns. Nothing else pins the + # return type down, so a change to the shared unit table cannot silently widen it. + + Scenario: trunc on a date literal returns date type + When query + """ + SELECT trunc(DATE '2024-05-15', 'MM') AS result + """ + Then query schema + """ + root + |-- result: date (nullable = true) + """ + Then query result + | result | + | 2024-05-01 | + + Scenario: trunc with the unit in a column returns date type + When query + """ + SELECT trunc(DATE '2024-05-15', c) AS result FROM VALUES ('MM') AS t(c) + """ + Then query schema + """ + root + |-- result: date (nullable = true) + """ + Then query result + | result | + | 2024-05-01 | + + Rule: a unit finer than a day is NULL, not a truncation + + Scenario: trunc to the day is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', 'DAY') AS result + """ + Then query result + | result | + | NULL | + + Scenario: trunc to the hour is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', 'HOUR') AS result + """ + Then query result + | result | + | NULL | + + Rule: an unrecognized or NULL unit is NULL, not an error + + Scenario: trunc with an unrecognized unit is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', 'bogus') AS result + """ + Then query result + | result | + | NULL | + + Scenario: trunc with a NULL unit is NULL + When query + """ + SELECT trunc(DATE '2019-08-04', NULL) AS result + """ + Then query result + | result | + | NULL | + + Rule: the unit may come from a column + + Scenario: trunc resolves the unit of each row + When query + """ + SELECT trunc(DATE '2019-08-04', c) AS result FROM VALUES (1, 'year'), (2, 'month'), (3, 'DAY'), (4, NULL) AS t(i, c) ORDER BY i + """ + Then query result ordered + | result | + | 2019-01-01 | + | 2019-08-01 | + | NULL | + | NULL | + + Rule: the unit is resolved once per row, even when it is not foldable + + # A unit built with `rand` is volatile: matching it once per candidate unit draws a new value + # for every comparison, so a row can match nothing and turn NULL. Spark resolves the unit once + # per row and never yields NULL here. A unit from `VALUES` is folded by Spark and hides this. + + Scenario: a volatile unit never nullifies a row + When query + """ + SELECT count(*) AS nulls FROM ( + SELECT trunc(DATE '2024-05-15', IF(rand(7) < 0.5, 'YEAR', 'MONTH')) AS r FROM range(10000) + ) WHERE r IS NULL + """ + Then query result + | nulls | + | 0 | + + # The scenario above cannot see a unit evaluated more than once: both draws name a + # granularity. Draw a NULL against a valid unit and the count becomes the evidence -- one + # evaluation lands near 5000, two near 7500. Asserted as a range because nothing says the two + # engines draw the same sequence from the same seed. + + Scenario: a volatile unit is drawn exactly once per row + When query + """ + SELECT count(*) BETWEEN 4700 AND 5300 AS drawn_once FROM ( + SELECT trunc(DATE '2024-05-15', IF(rand(7) < 0.5, CAST(NULL AS STRING), 'YEAR')) AS r + FROM range(10000) + ) WHERE r IS NULL + """ + Then query result + | drawn_once | + | true | + + Rule: the value is converted even when the unit matches nothing + + # Spark evaluates both arguments before it validates the unit, so an unconvertible value + # raises under ANSI even on a row whose unit matches no granularity. Resolving the unit first + # and skipping the conversion for unmatched rows would silently return NULL instead. + + # Two scenarios, because one assertion cannot carry both claims. The tagged one states the + # target -- Spark's error CLASS -- and xfails until Sail's string cast reports it. But an + # xfail is satisfied by ANY failure, including Sail simply returning NULL, so on its own it + # would not notice this rule breaking. The untagged one below matches what each engine + # actually prints, so it holds on BOTH and fails the moment either stops raising. + @sail-bug + Scenario: an unconvertible value reports Spark's error class when its unit matches nothing + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, trunc(IF(id = 0, 'not a date', '2024-01-02'), IF(id = 0, 'INVALID', 'YEAR')) AS result + FROM range(2) ORDER BY id + """ + Then query error CAST_INVALID_INPUT + + Scenario: an unconvertible value errors under ANSI on even when its unit matches nothing + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, trunc(IF(id = 0, 'not a date', '2024-01-02'), IF(id = 0, 'INVALID', 'YEAR')) AS result + FROM range(2) ORDER BY id + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + # `trunc` is `(date, format)`, so the date is the LEFT child and Spark evaluates it before it + # looks at the format -- a NULL format does NOT save an unconvertible value here. `date_trunc` + # has the arguments the other way round and behaves the opposite way; both are pinned down. + + Scenario: a NULL unit in a column still converts the value + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, trunc(IF(id = 0, 'not a date', '2024-01-02'), IF(id = 0, CAST(NULL AS STRING), 'YEAR')) AS result + FROM range(2) ORDER BY id + """ + Then query error (CAST_INVALID_INPUT|error in SQL parser: found .*expected '-') + + Scenario: an unconvertible value is NULL under ANSI off even when its unit matches nothing + Given config spark.sql.ansi.enabled = false + When query + """ + SELECT id, trunc(IF(id = 0, 'not a date', '2024-01-02'), IF(id = 0, 'INVALID', 'YEAR')) AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | 2024-01-01 | + + # A LITERAL unit that names no granularity is the other half of the rule: Spark's codegen + # short-circuits to NULL without evaluating the value at all, so the same unconvertible value + # does NOT raise here. For `trunc` a unit finer than a day names no level either, so it takes + # the same short circuit. + + Scenario: a literal unit that matches nothing skips the value entirely under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, trunc(IF(id = 0, 'not a date', '2024-01-02'), 'INVALID') AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | NULL | + + Scenario: a literal unit finer than a day skips the value entirely under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT id, trunc(IF(id = 0, 'not a date', '2024-01-02'), 'DAY') AS result + FROM range(2) ORDER BY id + """ + Then query result ordered + | id | result | + | 0 | NULL | + | 1 | NULL | + + Rule: a date beyond the nanosecond range truncates without wrapping + + # Truncating through a nanosecond timestamp overflows a 64-bit count past ~2262 and, for a + # DATE, does so with an unchecked multiply: the result WRAPS AROUND into the past instead of + # failing. A wrong value with no error is the worst outcome, and only a far-future date + # exposes it. + + Scenario: trunc YEAR beyond the nanosecond range + When query + """ + SELECT trunc(DATE '2300-01-01', 'YEAR') AS result + """ + Then query result + | result | + | 2300-01-01 | + + Scenario: trunc MONTH beyond the nanosecond range + When query + """ + SELECT trunc(DATE '2300-06-15', 'MONTH') AS result + """ + Then query result + | result | + | 2300-06-01 | + + Scenario: trunc YEAR at the maximum supported year + When query + """ + SELECT trunc(DATE '9999-12-31', 'YEAR') AS result + """ + Then query result + | result | + | 9999-01-01 | + + Rule: a string value parses with Spark's grammar + + # Spark's parser takes a single-digit month and a date with no month or day; Arrow's demands + # a fixed-position `YYYY-MM-DD`. Sail's `CAST(… AS DATE)` already uses Spark's. + + Scenario: a value with a single-digit month + When query + """ + SELECT trunc('2009-2-12', 'MM') AS result + """ + Then query result + | result | + | 2009-02-01 | + + Scenario: a value with no month or day + When query + """ + SELECT trunc('2024', 'YEAR') AS result + """ + Then query result + | result | + | 2024-01-01 | + + Rule: the value is converted once, even when it is volatile + + Scenario: a volatile value with a unit that matches nothing is always NULL + When query + """ + SELECT count(*) AS non_nulls FROM ( + SELECT trunc(IF(rand(7) < 0.5, DATE '2020-01-01', DATE '2021-01-01'), u) AS r + FROM (SELECT 'BOGUS' AS u FROM range(10000)) + ) WHERE r IS NOT NULL + """ + Then query result + | non_nulls | + | 0 | + + Rule: a timestamp value truncates on its local date + + # `trunc` takes the DATE of the value, and for a zoned timestamp that is the date in the + # SESSION zone, not in UTC. Under a negative-offset zone an instant late in the day still + # belongs to the previous local day, and under a positive-offset one it is already the next. + # Nothing else in this file uses a timestamp value or a non-UTC session. + + Scenario: a zoned timestamp uses the local date behind UTC + Given config spark.sql.session.timeZone = America/New_York + When query + """ + SELECT trunc(TIMESTAMP '2024-06-30 20:00:00', 'MONTH') AS result + """ + Then query result + | result | + | 2024-06-01 | + + Scenario: a zoned timestamp uses the local date ahead of UTC + Given config spark.sql.session.timeZone = Asia/Tokyo + When query + """ + SELECT trunc(TIMESTAMP '2024-07-01 02:00:00', 'MONTH') AS result + """ + Then query result + | result | + | 2024-07-01 | + + Scenario: a zoned timestamp from a column uses the local date + Given config spark.sql.session.timeZone = America/New_York + When query + """ + SELECT trunc(t, 'MONTH') AS result FROM VALUES (TIMESTAMP '2024-06-30 20:00:00') AS v(t) + """ + Then query result + | result | + | 2024-06-01 | + + Scenario: a timestamp_ntz uses its own wall clock + Given config spark.sql.session.timeZone = America/New_York + When query + """ + SELECT trunc(TIMESTAMP_NTZ '2024-06-30 20:00:00', 'MONTH') AS result + """ + Then query result + | result | + | 2024-06-01 | + + Rule: a unit of a type that cannot become a string matches nothing + + # Spark's coercion to string never fails -- `Cast(binary, string)` does not validate UTF-8 -- + # so a unit Arrow refuses to convert has to become a unit that matches nothing, not an error. + + Scenario: a binary unit is NULL, not an error + When query + """ + SELECT trunc(DATE '2024-01-01', X'FF') AS result + """ + Then query result + | result | + | NULL | + + Scenario: a binary unit whose bytes spell a granularity truncates + When query + """ + SELECT trunc(DATE '2024-05-15', X'4D4D') AS result + """ + Then query result + | result | + | 2024-05-01 | + + Rule: a literal unit short-circuits even where a column would not + + # `trunc` converts its value whatever the unit is -- the date is the LEFT child -- but that + # is the COLUMN rule. A literal unit Spark can fold short-circuits before either child is + # generated, so these do NOT raise, unlike their columnar twins above. + + Scenario: a literal NULL unit skips an unconvertible value under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT trunc('not a date', CAST(NULL AS STRING)) AS result + """ + Then query result + | result | + | NULL | + + Scenario: a literal unit finer than a day skips an unconvertible value under ANSI on + Given config spark.sql.ansi.enabled = true + When query + """ + SELECT trunc('not a date', 'DAY') AS result + """ + Then query result + | result | + | NULL | + + Rule: a value of the wrong type is rejected when the query is analyzed + + # Spark rejects the value for its TYPE. Arrow instead REINTERPRETS an integer as a day count, + # so without the gate `trunc(1, 'YEAR')` answers 1970-01-01 -- an invented value, with no + # error, which is worse than the wrong error class. + + Scenario: trunc rejects an integer value + When query + """ + SELECT trunc(1, 'YEAR') AS result + """ + Then query error DATATYPE_MISMATCH + + Scenario: trunc rejects a boolean value + When query + """ + SELECT trunc(true, 'YEAR') AS result + """ + Then query error DATATYPE_MISMATCH + + Scenario: trunc rejects an array value + When query + """ + SELECT trunc(array(1), 'YEAR') AS result + """ + Then query error DATATYPE_MISMATCH + + Rule: a complex unit is rejected when the query is analyzed + + # Spark rejects the unit for its TYPE while it analyzes the query. Coercing it to its string + # form instead would match no granularity and quietly return NULL. + + Scenario: trunc rejects an array unit + When query + """ + SELECT trunc(DATE '2024-01-01', array('YEAR')) AS result + """ + Then query error DATATYPE_MISMATCH Rule: trunc — the argument may come from a column @@ -36,3 +498,35 @@ Feature: trunc with an argument coming from a column | result | | 2019-07-29 | | 2019-07-29 | + + Rule: a date beyond Sail's supported range + + # Spark's DATE reaches year 5881580 and `trunc` truncates one correctly; Sail's range stops + # much earlier and rejects the value before `trunc` ever sees it. The gap is not in `trunc` + # itself -- every way of building such a date is rejected (`DATE` literal, `make_date`, + # `date_add`, `CAST(str AS DATE)`, `to_date`, date + int) -- so it belongs to the DATE surface. + # + # It is pinned here because it guards something specific: truncating past year ~294,247 + # overflows the microsecond conversion, and today that is unreachable ONLY because the range + # stops first. Widen the range and `trunc` would wrap silently on a date Spark handles, which + # is the same failure this PR fixed at year 2262. + + @sail-bug + Scenario: trunc on a date past the supported range + When query + """ + SELECT trunc(CAST('294248-01-01' AS DATE), 'YEAR') AS result + """ + Then query result + | result | + | +294248-01-01 | + + @sail-bug + Scenario: trunc on a date past the supported range built with to_date + When query + """ + SELECT trunc(to_date('294248-01-01'), 'YEAR') AS result + """ + Then query result + | result | + | +294248-01-01 | diff --git a/python/pysail/tests/spark/function/test_date_trunc_encoding.py b/python/pysail/tests/spark/function/test_date_trunc_encoding.py new file mode 100644 index 0000000000..b5fd4670f0 --- /dev/null +++ b/python/pysail/tests/spark/function/test_date_trunc_encoding.py @@ -0,0 +1,55 @@ +"""`date_trunc` / `trunc` on a string column whose ENCODING merely wraps the string. + +A pandas `category` column arrives as a dictionary-encoded Arrow array. The plan builder matches +the value on its LOGICAL type, so such a column is accepted -- but the parser it routes to +declares a user-defined signature and does no coercion of its own, so the encoding has to be +unwrapped before the call or the query is rejected for its layout alone. + +This cannot be written as a `.feature` scenario: SQL has no syntax that produces a +dictionary-encoded column. +""" + +import pandas as pd +import pytest +from pyspark.sql.types import Row + + +def _categorical(spark, values): + return spark.createDataFrame(pd.DataFrame({"v": pd.Series(values, dtype="category")})) + + +@pytest.mark.date_trunc +def test_date_trunc_on_dictionary_encoded_string(spark): + df = _categorical(spark, ["2024-05-15 13:45:30", "2024-05-15 13:45:30"]) + df.createOrReplaceTempView("dict_values") + # Compare the rendered string: collecting a timestamp converts it to the CLIENT's local zone, + # which would make the expected value depend on where the test runs. + actual = spark.sql("SELECT CAST(date_trunc('MM', v) AS STRING) AS result FROM dict_values").collect() + assert actual == [Row(result="2024-05-01 00:00:00")] * 2 + + +@pytest.mark.trunc +def test_trunc_on_dictionary_encoded_string(spark): + df = _categorical(spark, ["2024-05-15", "2024-05-15"]) + df.createOrReplaceTempView("dict_dates") + actual = spark.sql("SELECT CAST(trunc(v, 'MM') AS STRING) AS result FROM dict_dates").collect() + assert actual == [Row(result="2024-05-01")] * 2 + + +@pytest.mark.date_trunc +def test_date_trunc_with_dictionary_encoded_unit(spark): + """The UNIT, not the value: it reaches a different function, with its own signature.""" + df = _categorical(spark, ["MM", "YEAR"]) + df.createOrReplaceTempView("dict_units") + actual = spark.sql( + "SELECT CAST(date_trunc(v, TIMESTAMP '2024-05-15 13:45:30') AS STRING) AS result FROM dict_units" + ).collect() + assert actual == [Row(result="2024-05-01 00:00:00"), Row(result="2024-01-01 00:00:00")] + + +@pytest.mark.trunc +def test_trunc_with_dictionary_encoded_unit(spark): + df = _categorical(spark, ["MM", "YEAR"]) + df.createOrReplaceTempView("dict_trunc_units") + actual = spark.sql("SELECT CAST(trunc(DATE '2024-05-15', v) AS STRING) AS result FROM dict_trunc_units").collect() + assert actual == [Row(result="2024-05-01"), Row(result="2024-01-01")] diff --git a/python/pysail/tests/spark/function/test_date_trunc_error_class.py b/python/pysail/tests/spark/function/test_date_trunc_error_class.py new file mode 100644 index 0000000000..1bc0dd7853 --- /dev/null +++ b/python/pysail/tests/spark/function/test_date_trunc_error_class.py @@ -0,0 +1,30 @@ +"""`date_trunc` / `trunc` reject a bad argument type with Spark's THROWABLE CLASS. + +Spark raises `AnalysisException` while it analyzes the query. A `.feature` scenario cannot pin +this down: `Then query error` is `pytest.raises(Exception, match=...)`, which accepts any +exception whose message matches, so an `IllegalArgumentException` carrying the right text passes +it identically. Only asserting on the type catches a regression in the class. +""" + +import pytest +from pyspark.errors import AnalysisException + +# The base class both the classic and the Spark Connect clients derive their own from, so this +# holds whether the tests run against Spark JVM or against Sail. +DATATYPE_MISMATCH = [ + "SELECT date_trunc('DAY', 1) AS result", + "SELECT date_trunc('INVALID', true) AS result", + "SELECT date_trunc(array('YEAR'), TIMESTAMP '2024-01-01 00:00:00 UTC') AS result", + "SELECT date_trunc('YEAR', array(1)) AS result", + "SELECT trunc(1, 'YEAR') AS result", + "SELECT trunc(array(1), 'YEAR') AS result", +] + + +@pytest.mark.date_trunc +@pytest.mark.trunc +@pytest.mark.parametrize("query", DATATYPE_MISMATCH) +def test_bad_argument_type_raises_analysis_exception(spark, query): + with pytest.raises(AnalysisException) as error: + spark.sql(query).collect() + assert "DATATYPE_MISMATCH" in str(error.value)