diff --git a/datafusion/spark/src/function/datetime/unix.rs b/datafusion/spark/src/function/datetime/unix.rs index 6eaf3a08780bc..31a7c4d5fc4db 100644 --- a/datafusion/spark/src/function/datetime/unix.rs +++ b/datafusion/spark/src/function/datetime/unix.rs @@ -24,7 +24,7 @@ use datafusion_common::{Result, internal_err}; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ Coercion, ColumnarValue, Expr, ExprSchemable, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, lit, }; /// Returns the number of days since epoch (1970-01-01) for the given date. @@ -154,12 +154,38 @@ impl ScalarUDFImpl for SparkUnixTimestamp { info: &SimplifyContext, ) -> Result { let [ts] = take_function_args(self.name(), args)?; + // Arrow casts to a coarser timestamp unit truncate toward zero, but Spark + // floors (`Math.floorDiv`). When downscaling, divide the raw ticks and + // subtract one for negative values with a remainder. + let divisor: i64 = match (ts.get_type(info.schema())?, self.time_unit) { + (DataType::Timestamp(TimeUnit::Millisecond, _), TimeUnit::Second) + | (DataType::Timestamp(TimeUnit::Microsecond, _), TimeUnit::Millisecond) + | (DataType::Timestamp(TimeUnit::Nanosecond, _), TimeUnit::Microsecond) => { + 1_000 + } + (DataType::Timestamp(TimeUnit::Microsecond, _), TimeUnit::Second) + | (DataType::Timestamp(TimeUnit::Nanosecond, _), TimeUnit::Millisecond) => { + 1_000_000 + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), TimeUnit::Second) => { + 1_000_000_000 + } + _ => { + return Ok(ExprSimplifyResult::Simplified( + ts.cast_to( + &DataType::Timestamp(self.time_unit, Some("UTC".into())), + info.schema(), + )? + .cast_to(&DataType::Int64, info.schema())?, + )); + } + }; + let ticks = ts.cast_to(&DataType::Int64, info.schema())?; + let correction = (ticks.clone() % lit(divisor)) + .lt(lit(0_i64)) + .cast_to(&DataType::Int64, info.schema())?; Ok(ExprSimplifyResult::Simplified( - ts.cast_to( - &DataType::Timestamp(self.time_unit, Some("UTC".into())), - info.schema(), - )? - .cast_to(&DataType::Int64, info.schema())?, + ticks / lit(divisor) - correction, )) } } diff --git a/datafusion/sqllogictest/test_files/spark/datetime/unix.slt b/datafusion/sqllogictest/test_files/spark/datetime/unix.slt index 3e19b836b4e7a..b753804a3b2fe 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/unix.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/unix.slt @@ -132,3 +132,51 @@ NULL query error Function 'unix_seconds' requires Timestamp, but received String \(DataType: Utf8View\) SELECT unix_seconds('1970-01-01 00:00:01Z'::string); + +# Spark floors (Math.floorDiv) when converting to a coarser unit, so pre-epoch +# timestamps with a sub-unit remainder round toward negative infinity. +statement ok +CREATE TABLE unix_floor_us AS SELECT arrow_cast(column1, 'Timestamp(Microsecond, None)') AS ts FROM (VALUES + (-1), (-1000), (-1001), (-1000000), (-1000001), (1500), (NULL) +); + +query III +SELECT arrow_cast(ts, 'Int64'), unix_millis(ts), unix_seconds(ts) FROM unix_floor_us; +---- +-1 -1 -1 +-1000 -1 -1 +-1001 -2 -1 +-1000000 -1000 -1 +-1000001 -1001 -2 +1500 1 0 +NULL NULL NULL + +statement ok +DROP TABLE unix_floor_us; + +# Same results when the expressions are constant-folded +query II +SELECT + unix_millis(arrow_cast(-1001, 'Timestamp(Microsecond, None)')), + unix_seconds(arrow_cast(-1000001, 'Timestamp(Microsecond, None)')); +---- +-2 -2 + +# Other input units: nanoseconds floor too, milliseconds to seconds floors, +# and conversions to a finer unit are unchanged +query IIII +SELECT + unix_millis(arrow_cast(-1, 'Timestamp(Nanosecond, None)')), + unix_micros(arrow_cast(-1, 'Timestamp(Nanosecond, None)')), + unix_seconds(arrow_cast(-1, 'Timestamp(Millisecond, None)')), + unix_micros(arrow_cast(-1, 'Timestamp(Millisecond, None)')); +---- +-1 -1 -1 -1000 + +# Extreme values do not overflow +query II +SELECT + unix_seconds(arrow_cast(-9223372036854775808, 'Timestamp(Microsecond, None)')), + unix_seconds(arrow_cast(9223372036854775, 'Timestamp(Millisecond, None)')); +---- +-9223372036855 9223372036854