unikdahal commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3919223848
########## 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: This introduces a narrower date domain than the Iceberg JVM implementation. `Date32Type::to_naive_date_opt` ultimately uses chrono's date representation, whose range is only about 262k years in either direction. Spark `DateType`, however, carries an `i32` epoch-day value, and Iceberg's `DateTimeUtil.daysToYears` / `daysToMonths` uses Java `LocalDate`; that supports the complete `i32` epoch-day domain. For example, epoch day `100000000` is valid for the JVM implementation but falls outside chrono's range, so native `years` / `months` will return an execution error where Spark/Iceberg succeeds. The timestamp versions have the same issue near the extremes of the `i64` microsecond domain. Could we compute the proleptic-Gregorian year/month directly from the epoch-day integer (or fall back outside chrono's range) and add a reference test with epoch days beyond chrono's limits? ########## native/spark-expr/src/iceberg_funcs/truncate.rs: ########## @@ -0,0 +1,392 @@ +// 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 `truncate(width, value)` transform: `v - ((v % W) + W) % W` for integers (with +//! Java's wrapping arithmetic), the same on the unscaled value for decimals, the first `W` code +//! points of a string, and the first `W` bytes of a binary value. + +use super::{apply_unary, positive_int_param, unpacked_type, unsupported_type}; +use crate::utils::is_valid_decimal_precision; +use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array, OffsetSizeTrait}; +use arrow::compute::kernels::substring::{substring, substring_by_char}; +use arrow::datatypes::{DataType, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type}; +use datafusion::common::{utils::take_function_args, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// `TruncateUtil.truncateInt`. Java's `int` arithmetic wraps on overflow, which can happen both in +/// `(v % w) + w` (for widths above 2^30) and in the final subtraction (near `Integer.MIN_VALUE`). +/// `TruncateUtil.truncateByte` / `truncateShort` evaluate the same expression in `int` and then +/// narrow, so tinyint and smallint inputs go through this function and are cast afterwards. +#[inline] +fn truncate_i32(v: i32, w: i32) -> i32 { + v.wrapping_sub((v % w).wrapping_add(w) % w) +} + +/// `TruncateUtil.truncateLong`, with the width promoted to `long` as Java does. +#[inline] +fn truncate_i64(v: i64, w: i64) -> i64 { + v.wrapping_sub((v % w).wrapping_add(w) % w) +} + +/// `TruncateUtil.truncateDecimal` on the unscaled value; `BigInteger` never overflows and neither +/// does an `i128` holding a 38-digit unscaled value minus a 31-bit width. +#[inline] +fn truncate_i128(v: i128, w: i128) -> i128 { + v - ((v % w) + w) % w +} + +/// `UTF8String.substring(0, width)` counts code points, not bytes. A width that covers the whole +/// values buffer cannot truncate anything, so the input is returned as is instead of being copied. +fn truncate_string<O: OffsetSizeTrait>(array: &ArrayRef, width: i32) -> Result<ArrayRef> { + let strings = array.as_string::<O>(); + if width as usize >= strings.value_data().len() { + return Ok(Arc::clone(array)); + } + Ok(Arc::new(substring_by_char(strings, 0, Some(width as u64))?)) +} + +/// `BinaryUtil.truncateBinaryUnsafe` keeps the first `width` bytes. The whole-buffer shortcut +/// matters here beyond avoiding a copy: Arrow's byte `substring` adds the length to each value's +/// offset without checking for overflow, which panics for a width near `i32::MAX`. +fn truncate_binary<O: OffsetSizeTrait>(array: &ArrayRef, width: i32) -> Result<ArrayRef> { + if width as usize >= array.as_binary::<O>().value_data().len() { + return Ok(Arc::clone(array)); + } + Ok(substring(array.as_ref(), 0, Some(width as u64))?) +} + +fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result<ArrayRef> { + let result: ArrayRef = match array.data_type() { + DataType::Int8 => Arc::new( + array + .as_primitive::<Int8Type>() + .unary::<_, Int8Type>(|v| truncate_i32(v as i32, width) as i8), + ), + DataType::Int16 => Arc::new( + array + .as_primitive::<Int16Type>() + .unary::<_, Int16Type>(|v| truncate_i32(v as i32, width) as i16), + ), + DataType::Int32 => Arc::new( + array + .as_primitive::<Int32Type>() + .unary::<_, Int32Type>(|v| truncate_i32(v, width)), + ), + DataType::Int64 => Arc::new( + array + .as_primitive::<Int64Type>() + .unary::<_, Int64Type>(|v| truncate_i64(v, width as i64)), + ), + DataType::Decimal128(precision, scale) => { + // Truncating a negative value grows its magnitude by up to `width - 1` units of the + // last digit, so the result can need one more digit than the column allows. Spark's + // `UnsafeRowWriter` writes such a `Decimal` as null (`changePrecision` fails), so + // match that rather than emit a value the column's precision cannot hold. + let truncated: Decimal128Array = + array.as_primitive::<Decimal128Type>().unary_opt(|v| { + let truncated = truncate_i128(v, width as i128); + is_valid_decimal_precision(truncated, *precision).then_some(truncated) Review Comment: I think this precision check is happening too early relative to Spark's semantics. Iceberg's JVM `TruncateDecimal.invoke` returns `Decimal.apply(truncatedValue)` without coercing it back to the declared `DecimalType`. Spark's `StaticInvoke` therefore exposes a non-null `Decimal` to its parent expression; `changePrecision(precision, scale)` only happens later when an `UnsafeRowWriter` materializes the value. Here `unary_opt` turns that intermediate into NULL during expression evaluation itself. That changes nested/predicate semantics. For example, with `DECIMAL(18,4)` value `-99999999999999.9999`, `truncate(10, v)` mathematically produces `-100000000000000.0000` (19 digits). On the JVM the `StaticInvoke` result is still non-null, whereas this branch returns NULL, so `WHERE truncate(10, v) IS NULL` can silently select a row that Spark does not. Could we either preserve the JVM intermediate semantics or fall back when this overflow case occurs? -- 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]
