andygrove commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3927024140
########## 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: _LLM-assisted: this reply and the changes it describes were written with Claude Code._ You're right about the intermediate: `TruncateDecimal.invoke` returns `Decimal.apply(...)` with no `changePrecision`, so the `StaticInvoke` result is non-null and only `UnsafeRowWriter` nulls it. I don't think the kernel can reproduce that, though. It has to return a `Decimal128(precision, scale)` array, and that array has no encoding for "exceeds the declared precision but isn't null yet" — whatever it holds is what both the plan output and any parent expression see. Emitting the oversized value would make the common case wrong: `SELECT truncate(10, dec18)` is null in Spark, and the corpus in the new suite already contains `-99999999999999.9999`, so `checkSparkAnswerAndOperator` fails on it. Nulling eagerly is also what Spark does for decimal overflow elsewhere, via `CheckOverflow(..., nullOnOverflow = true)`; Iceberg's `invoke` simply isn't wrapped in one. So it's a choice of which side to be right on, and I've taken the materialized side. As @jordepic notes, `SortExec` builds its key through an `UnsafeRowWriter` too, so the sort agrees, and so does the projection output. What differs is a truncated decimal feeding another expression directly: `WHERE truncate(10, v) IS NULL`, and the `Murmur3Hash` behind `DISTRIBUTE BY` (that one only moves rows between partitions, it doesn't change an answer). The Iceberg partition value isn't affected either way — the writer derives it from the untruncated column. I've written the difference down rather than leaving it implicit: 3dc1a2d in the Iceberg user guide, plus the kernel comment. The alternative is to fall back for decimal `truncate` entirely, which costs native decimal partitioning for a case that needs a value within `width - 1` ULP of the negative precision boundary. Would you rather have that? -- 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]
