harsh-ande commented on code in PR #26129:
URL: https://github.com/apache/datafusion/pull/26129#discussion_r4366458209
##########
datafusion/spark/src/function/datetime/unix.rs:
##########
@@ -154,12 +154,38 @@ impl ScalarUDFImpl for SparkUnixTimestamp {
info: &SimplifyContext,
) -> Result<ExprSimplifyResult> {
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,
))
}
Review Comment:
Replaced the duplicated-expression rewrite with runtime floor division on
the already-evaluated argument. Same-unit and upscaling conversions retain the
existing cast simplification. Added deterministic volatile-UDF regressions
covering all six downscaling pairs with scalar and array inputs, checking both
results and exactly one evaluation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
