andygrove commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3928584282
########## 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._ Thanks for catching it. I'd taken the `UnsafeRowWriter` claim from your earlier note and propagated it into `iceberg.md` and the kernel comment without checking, so that one's on me too. Confirmed: `SortExec` orders the child's rows with an ordering from `RowOrdering.create`, which evaluates the sort expression per comparison, and the `decimal(18,4)` prefix is `toUnscaledLong` of that same value — no `UnsafeRowWriter` in the path. Both places are corrected in 0a57b6a, and it is moot anyway now that decimals fall back. ########## spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala: ########## @@ -0,0 +1,229 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal} +import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke +import org.apache.spark.sql.types._ + +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType} + +/** + * Native support for Iceberg's Spark system functions (`bucket`, `truncate`, `years`, `months`, + * `days`, `hours`). + * + * Iceberg exposes each of these through Spark's static magic method, so + * `V2ExpressionUtils.resolveScalarFunction` binds them as `StaticInvoke(cls, "invoke", args)` + * where `cls` is one of the per-type implementations under `org.apache.iceberg.spark.functions` + * (e.g. `BucketFunction$BucketInt`). The same expressions appear in the hash distribution and + * local sort that Iceberg requests in front of a partitioned write, and in predicates and + * projections that users write against hidden partitioning, so routing them through + * [[CometStaticInvoke]] covers shuffle, sort, filter, and projection at once. + * + * The list of classes is Iceberg's; `IcebergVersionFunction` is a zero-argument constant and is + * deliberately left out. + */ +object CometIcebergSystemFunctions { + + private val FunctionsPackage = "org.apache.iceberg.spark.functions." + + /** Every Iceberg system function exposes its static magic method under this name. */ + private val MagicMethod = "invoke" + + private def implementations( + outer: String, + handler: CometExpressionSerde[StaticInvoke], + inner: String*): Seq[((String, String), CometExpressionSerde[StaticInvoke])] = + inner.map(name => (MagicMethod, s"$FunctionsPackage$outer$$$name") -> handler) + + /** + * Handlers keyed by `(functionName, class name)` of the Iceberg implementation class that + * `StaticInvoke` calls, the shape [[CometStaticInvoke]] dispatches on. Iceberg is not on + * Comet's compile classpath, which is why the key carries the class name rather than the class. + */ + val staticInvokeHandlers: Map[(String, String), CometExpressionSerde[StaticInvoke]] = ( + implementations( + "BucketFunction", + CometIcebergBucket, + "BucketInt", + "BucketLong", + "BucketString", + "BucketBinary", + "BucketDecimal") ++ + implementations( + "TruncateFunction", + CometIcebergTruncate, + "TruncateTinyInt", + "TruncateSmallInt", + "TruncateInt", + "TruncateBigInt", + "TruncateString", + "TruncateBinary", + "TruncateDecimal") ++ + implementations( + "YearsFunction", + CometIcebergYears, + "DateToYearsFunction", + "TimestampToYearsFunction", + "TimestampNtzToYearsFunction") ++ + implementations( + "MonthsFunction", + CometIcebergMonths, + "DateToMonthsFunction", + "TimestampToMonthsFunction", + "TimestampNtzToMonthsFunction") ++ + implementations( + "DaysFunction", + CometIcebergDays, + "DateToDaysFunction", + "TimestampToDaysFunction", + "TimestampNtzToDaysFunction") ++ + implementations( + "HoursFunction", + CometIcebergHours, + "TimestampToHoursFunction", + "TimestampNtzToHoursFunction") + ).toMap + + /** + * The `numBuckets` / `width` argument as a positive int, if it is a literal. Iceberg declares + * the parameter as `IntegerType`, so a tinyint or smallint literal arrives already cast and + * folded; the narrower literal types are matched anyway in case folding did not run. + */ + private[serde] def positiveIntLiteral(expr: Expression): Option[Int] = expr match { + case Literal(v: Int, IntegerType) if v > 0 => Some(v) + case Literal(v: Short, ShortType) if v > 0 => Some(v.toInt) + case Literal(v: Byte, ByteType) if v > 0 => Some(v.toInt) + case _ => None + } +} + +/** + * Shared shape of `bucket(numBuckets, value)` and `truncate(width, value)`: a positive integer + * parameter followed by the value. The parameter has to be a literal because the native kernel + * takes it as a constant, and it has to be positive because Iceberg's Java implementation divides + * by it (zero throws, which the fallback preserves by leaving the expression to Spark). + */ +abstract class CometIcebergParameterizedTransform( + nativeName: String, + parameterName: String, + valueTypeSupported: DataType => Boolean) + extends CometExpressionSerde[StaticInvoke] { + + override def getSupportLevel(expr: StaticInvoke): SupportLevel = expr.arguments match { + case Seq(parameter, value) => + if (CometIcebergSystemFunctions.positiveIntLiteral(parameter).isEmpty) { + Unsupported(Some(s"$parameterName must be a positive integer literal, got $parameter")) + } else if (!valueTypeSupported(value.dataType)) { + Unsupported(Some(s"$nativeName does not support input type ${value.dataType}")) + } else { + Compatible() Review Comment: _LLM-assisted: this reply and the changes it describes were written with Claude Code._ Good call on reaching for the existing mechanism. I went one step past it and marked decimal `truncate` `Unsupported` rather than `Incompatible`, for two reasons. The wiring @sunchao flagged is real: `getExprConfigName` resolves against the outer expression class, so the opt-in would have been `spark.comet.expression.StaticInvoke.allowIncompatible`, which also unlocks every future incompatible static invoke. Fixing that properly means changing what `spark.comet.expression.StaticInvoke.enabled` covers for `readSidePadding`, `aesEncrypt`, `decode` and `base64`, and I would rather not smuggle that into this PR. `Unsupported` needs no key, so the question disappears. The docs half of your point I did take. `CometStaticInvoke.getUnsupportedReasons()` now aggregates the per-function handlers' notes, since `GenerateDocs` only asks the serde registered for the expression class — without it the note reached the Iceberg guide but not the generated compatibility page. There is a test asserting it is reachable that way so it cannot silently drop out. If you would rather have the opt-in after all, `Incompatible` is a one-line change once the config-name question is settled on its own. -- 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]
