jordepic commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3927765515
########## 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: On the decimal `truncate` question @sunchao and you are going back and forth on: Comet already has the mechanism that gives both of you what you want. If `getSupportLevel` returns `Incompatible(Some(...))` for a `DecimalType` value instead of `Compatible()`, the expression falls back by default and a user who accepts the documented difference can opt in with the per-expression `allowIncompatible` config. That is how every other expression with a known divergence from Spark is handled here, and it keeps the native path available for partitioned writes on tables where the operator has read the caveat. It would also need `getIncompatibleReasons` so the generated compatibility page carries the note, rather than only the Iceberg guide. That seems like a better answer to your "would you rather have that?" than a hard fallback. ########## spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala: ########## @@ -0,0 +1,423 @@ +/* + * 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 + // Three Spark 3.x defaults differ from Spark 4's in ways that block writing this corpus. All + // are set to Spark 4's value so one corpus works on every profile, and none affects how a + // result is compared, since every test reads the data back from parquet. + // + // - `datetimeJava8ApiEnabled`: `sourceData` supplies java.time values for the date and + // timestamp columns. Spark 4 resolves those external types; on Spark 3.x the row encoder + // expects java.sql.Date / java.sql.Timestamp instead. The encoder is built here on the + // driver, so setting the flag around the write is enough. + // - `datetimeRebaseModeInWrite`: the timestamp corpus reaches back to 1843, and the corpus + // is deliberately pre-epoch in places, since the temporal transforms go negative before + // 1970. Spark 3.x throws on writing a timestamp before 1900; Spark 4 defaults to + // CORRECTED, which writes the value as-is. + // - `outputTimestampType`: Spark 3.x defaults to INT96, which has its own separate ancient + // timestamp check. Spark 4 defaults to TIMESTAMP_MICROS, which is also what Iceberg + // itself writes. + withSQLConf( + SQLConf.DATETIME_JAVA8API_ENABLED.key -> "true", + SQLConf.PARQUET_REBASE_MODE_IN_WRITE.key -> "CORRECTED", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "TIMESTAMP_MICROS") { + 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: with it in the spec this test failed on the Linux + // CI runners with a missing data file, which does not reproduce locally, so the multi-byte Review Comment: I think the missing-file failure does have an explanation, and it is the partition-path caveat already documented in `iceberg-writes.md`. iceberg-java percent-encodes partition directory names and values through `URLEncoder` in `PartitionSpec.partitionToPath`, and iceberg-rust writes them raw. The Linux CI runners run the JVM under a POSIX locale, so `sun.jnu.encoding` is ASCII and opening a local path containing `日本` or `😀` fails, which reads back as a missing data file. On macOS that encoding is fixed to UTF-8 by the launcher and a `-Dsun.jnu.encoding` override is ignored, which would be why it did not reproduce locally even when forced. Running the original version of this test under `LANG=C` on Linux should confirm it. If it does, this is a native-writer divergence with an operational effect rather than a test artifact, and worth filing against iceberg-rust's `partition_to_path` the same way you filed the other two. -- 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]
