jordepic commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3927764330
########## 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 need to correct what I said above, and @sunchao is right to push on it. `SortExec` does not build its key through an `UnsafeRowWriter`. It creates the comparator with `RowOrdering.create`, which generates code that evaluates the sort expression on each input row and compares the raw result, and the sort prefix for a `decimal(18,4)` is `toUnscaledLong` of that same value. So Spark's sort sees the oversized decimal and Comet's sort sees null. The sentence about the sort key in the new `iceberg.md` paragraph and in the `truncate.rs` comment came from my mistake and should be dropped. The practical effect on the write path is still nil, for a different reason: for a given precision and width exactly one truncated value can exceed the precision (the multiples of the width that land below the negative bound form a window narrower than the width), and it is the column minimum, so it sorts first under either key and the clustered writer sees a single run. The difference is confined to ordering position under a non-default null ordering, and to predicates and hashes as you described. -- 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]
