andygrove commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3927026630
########## native/spark-expr/src/iceberg_funcs/temporal.rs: ########## @@ -0,0 +1,339 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Iceberg's `years`, `months`, `days`, and `hours` transforms. +//! +//! Iceberg's `DateTimeUtil` evaluates all four in UTC regardless of the Spark session timezone +//! (`TimestampType` and `TimestampNTZType` are handled identically), and all four floor: a value +//! before the epoch maps to a negative period. `years` and `months` are calendar-aware, `days` and +//! `hours` are plain floor division of the epoch value. `days` returns a date (Iceberg's +//! `DaysFunction.resultType()` is `DateType`), the other three return an int. +//! +//! The kernels work on the raw epoch values rather than going through Arrow's timezone-aware +//! `date_part`, which would otherwise shift a `TimestampType` column by the session offset. + +use super::{apply_unary, unsupported_type}; +use arrow::array::{ArrayRef, AsArray, Int32Array}; +use arrow::datatypes::{DataType, Date32Type, Int32Type, TimeUnit, TimestampMicrosecondType}; +use chrono::Datelike; +use datafusion::common::{utils::take_function_args, DataFusionError, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use num::integer::div_floor; +use std::sync::Arc; + +const MICROS_PER_HOUR: i64 = 3_600_000_000; +const MICROS_PER_DAY: i64 = 86_400_000_000; +const UNIX_EPOCH_YEAR: i32 = 1970; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum TemporalUnit { + Years, + Months, + Days, + Hours, +} + +impl TemporalUnit { + fn fn_name(self) -> &'static str { + match self { + TemporalUnit::Years => "iceberg_years", + TemporalUnit::Months => "iceberg_months", + TemporalUnit::Days => "iceberg_days", + TemporalUnit::Hours => "iceberg_hours", + } + } + + fn return_type(self) -> DataType { + match self { + TemporalUnit::Days => DataType::Date32, + _ => DataType::Int32, + } + } +} + +/// `DateTimeUtil.microsToDays`: floor division, so `-1` micros is day `-1`. +#[inline] +fn micros_to_days(micros: i64) -> i32 { + div_floor(micros, MICROS_PER_DAY) as i32 +} + +/// `DateTimeUtil.microsToHours`. +#[inline] +fn micros_to_hours(micros: i64) -> i32 { + div_floor(micros, MICROS_PER_HOUR) as i32 +} + +/// `DateTimeUtil.daysToYears`: whole calendar years between the epoch and the day, floored. +fn days_to_years(days: i32) -> Result<i32> { + Ok(civil_date(days)?.year() - UNIX_EPOCH_YEAR) +} + +/// `DateTimeUtil.daysToMonths`: whole calendar months between the epoch and the day, floored. +fn days_to_months(days: i32) -> Result<i32> { + let date = civil_date(days)?; + Ok((date.year() - UNIX_EPOCH_YEAR) * 12 + date.month0() as i32) +} + +fn civil_date(days: i32) -> Result<chrono::NaiveDate> { + Date32Type::to_naive_date_opt(days).ok_or_else(|| { Review Comment: _LLM-assisted: this reply and the changes it describes were written with Claude Code._ Good catch, fixed in 71e03be. `years` and `months` no longer touch `chrono`: `civil_from_days` does the proleptic Gregorian split in `i64` integer arithmetic, exact over the whole `i32` epoch-day domain — `i32::MAX` days is +5881580-07-11 and `i32::MIN` is -5877641-06-23, both well inside `LocalDate`. The kernels are infallible now as a result. The expected values come from running Iceberg's own `DateTimeUtil.convertDays` / `convertMicros` on a JDK 17 JVM, and the tables pin epoch days ±1e8, ±1e9 and both `i32` extremes, plus `i64::MIN` / `i64::MAX` micros. There is also a test checking the integer split against `chrono` for every day in a 400-year window around the epoch and on a stride over the rest, so it cannot drift in between. On the timestamp side you're right that the same concern applies, though it lands differently per function. `days` was already safe: `i64::MAX` micros is 1.07e8 days, comfortably inside an `i32`. `years` and `months` went through `chrono` after that conversion and had exactly the problem you describe. `hours` is the one that genuinely overflows — `i64::MAX` micros is 2.56e9 hours — and Iceberg narrows it with a plain `(int)` cast, so `as i32` wraps to the same value; that is now pinned rather than incidental. -- 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]
