jordepic commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3926048672
########## 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: One data point for scoping the fix: on the write path the null sort key is actually consistent with Spark. Spark's `SortExec` builds its key through an `UnsafeRowWriter` too, so the JVM sort sees null for the same rows. The divergence you describe is confined to filters and projections, where codegen evaluates `invoke` without materializing the row. It may be worth keeping the fix scoped to those paths so the sort key stays aligned with what Spark's own sort produces. ########## spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala: ########## @@ -0,0 +1,400 @@ +/* + * 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 + +import java.io.File +import java.math.{BigDecimal => JBigDecimal, BigInteger} +import java.nio.file.Files +import java.time.{Instant, LocalDate, LocalDateTime, ZoneOffset} + +import scala.util.Random + +import org.scalactic.source.Position +import org.scalatest.Tag + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, Literal} +import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke +import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSortExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +import org.apache.comet.serde.{CometExpressionSerde, CometIcebergBucket, CometIcebergTruncate, CometStaticInvoke, Compatible, SupportLevel, Unsupported} + +/** + * Native support for Iceberg's system functions (`bucket`, `truncate`, `years`, `months`, `days`, + * `hours`). + * + * Every comparison runs the same query with Comet on and off, so the reference values come from + * Iceberg's own JVM implementations (`BucketFunction`, `TruncateFunction`, ...) evaluated by + * Spark, over seeded random data plus the boundary values of each type. A native result that + * disagreed with Iceberg would only fail loudly on the write path (the clustered writer rejects + * out-of-order partitions); in a filter or projection it would be a silently wrong answer, which + * is why the coverage is per type rather than a few hand-picked rows. + */ +class CometIcebergSystemFunctionSuite + extends CometTestBase + with CometIcebergTestBase + with AdaptiveSparkPlanHelper { + + override protected def sparkConf: SparkConf = { + super.sparkConf + .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true") + .set(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key, "true") + } + + override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit + pos: Position): Unit = { + super.test(testName, testTags: _*) { + assume(icebergAvailable, "Iceberg not available in classpath") + testFun + } + } + + private val catalog = "ice" + private val source = "system_function_source" + private val bucketColumns = + Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin", "dt", "ts", "ts_ntz") + private val truncateColumns = Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin") + + // The source data is written once per suite; every test reads the same parquet directory. + private var sourceDir: File = _ + private def sourcePath: String = new File(sourceDir, "data").getAbsolutePath + + override def beforeAll(): Unit = { + super.beforeAll() + sourceDir = Files.createTempDirectory("comet-iceberg-system-functions").toFile + sourceData().write.parquet(sourcePath) + } + + override def afterAll(): Unit = { + try deleteRecursively(sourceDir) + finally super.afterAll() + } + + test("bucket matches Iceberg for every supported type") { + withSourceTable { + bucketColumns.foreach { column => + val buckets = + Seq(1, 7, 16, Int.MaxValue).map(n => s"$catalog.system.bucket($n, $column)") + checkSparkAnswerAndOperator(s"SELECT $column, ${buckets.mkString(", ")} FROM $source") + } + } + } + + test("truncate matches Iceberg for every supported type") { + withSourceTable { + truncateColumns.foreach { column => + val truncated = + Seq(1, 3, 10, 1000, Int.MaxValue).map(w => s"$catalog.system.truncate($w, $column)") + checkSparkAnswerAndOperator(s"SELECT $column, ${truncated.mkString(", ")} FROM $source") + } + } + } + + test("years, months, days, and hours match Iceberg regardless of session timezone") { + withSourceTable { + // Iceberg evaluates the temporal transforms in UTC; a shifted session timezone must not + // leak into the native result either. + for (timezone <- Seq("UTC", "America/Los_Angeles", "Asia/Kathmandu")) { + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> timezone) { + Seq("dt", "ts", "ts_ntz").foreach { column => + val functions = Seq("years", "months", "days") ++ (if (column == "dt") Nil + else Seq("hours")) + val transformed = functions.map(f => s"$catalog.system.$f($column)") + checkSparkAnswerAndOperator( + s"SELECT $column, ${transformed.mkString(", ")} FROM $source") + } + } + } + } + } + + test("system functions in filters stay native") { + withSourceTable { + checkSparkAnswerAndOperator( + s"SELECT i32 FROM $source WHERE $catalog.system.bucket(8, i32) IN (0, 3)") + checkSparkAnswerAndOperator( + s"SELECT str FROM $source WHERE $catalog.system.truncate(1, str) = 'a'") + checkSparkAnswerAndOperator( + s"SELECT ts FROM $source WHERE $catalog.system.days(ts) >= DATE '2000-01-01'") + checkSparkAnswerAndOperator(s"SELECT dt FROM $source WHERE $catalog.system.months(dt) < 0") + } + } + + test("sorting on system functions stays native") { + withSourceTable { + val df = sql( + s"SELECT i32, str FROM $source " + + s"ORDER BY $catalog.system.bucket(4, i32), $catalog.system.truncate(2, str), i32, str") + checkSparkAnswerAndOperator(df) + val sorts = collect(stripAQEPlan(df.queryExecution.executedPlan)) { case s: CometSortExec => + s + } + assert(sorts.nonEmpty, "expected a native sort") + } + } + + test("hash partitioning on system functions uses the native shuffle") { + withSourceTable { + val df = sql( + s"SELECT i32, str, ts FROM $source " + + s"DISTRIBUTE BY $catalog.system.bucket(8, i32), $catalog.system.truncate(2, str), " + + s"$catalog.system.hours(ts)") + checkSparkAnswerAndOperator(df) + checkCometExchange(df, 1, native = true) + } + } + + test("partitioned Iceberg write with default distribution mode stays native end to end") { + withSourceTable { + val table = s"$catalog.db.hidden_partitioning" + // No `write.distribution-mode`: Iceberg picks hash distribution for a partitioned table, + // which plans a shuffle and a local sort on the partition transforms. The string column is + // deliberately not a partition source: its multi-byte values would land in partition + // directory names, which iceberg-java reads back through the JVM's platform charset. + sql(s""" + CREATE TABLE $table (i32 INT, i64 BIGINT, ts TIMESTAMP, dt DATE) + USING iceberg + PARTITIONED BY (bucket(4, i32), truncate(1000, i64), days(ts), months(dt))""") + // iceberg-rust's own truncate transform, which the writer uses for partition values, does Review Comment: This works around unchecked arithmetic in iceberg-rust's `truncate_i64`, which also affects the native writer's partition values in debug builds, independent of this PR. Could you file that upstream in iceberg-rust and reference the issue here instead of describing the workaround inline? Otherwise the filter reads as a test-data choice and the writer-side bug gets lost. ########## 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 Review Comment: In Comet, `TimestampType` arrays are always tagged `UTC` (`native/core/src/execution/serde.rs`), and the writer casts every batch to iceberg-rust's Arrow schema, which tags `Timestamptz` as `+00:00`. So `date_part` in iceberg-rust's `Year` and `Month` transforms would also evaluate in UTC on both the sort path and the write path. Is there a plan shape where a non-UTC tag reaches these kernels? If not, it would be good for this doc to say the concern is defensive rather than observed, or these two could reuse iceberg-rust's kernels as well. -- 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]
