andygrove commented on code in PR #5638: URL: https://github.com/apache/datafusion-comet/pull/5638#discussion_r3928586770
########## 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: _LLM-assisted: this reply and the changes it describes were written with Claude Code._ Your mechanism is right, and it is already fixed upstream — which is the part I had wrong when I said I had no cause. iceberg-java escapes both field name and value through `URLEncoder.encode` in `PartitionSpec.partitionToPath`; iceberg-rust at our pinned rev formats `{name}={human_string}` raw. apache/iceberg-rust#2875 changed that to `form_urlencoded::Serializer::append_pair` on both sides and merged 2026-07-30 — one day after the rev we pin, `3d84c81` (2026-07-29). So the caveat in `iceberg-writes.md` is accurate for our pin and stale against upstream main. Three things support the locale half without my having read the logs. `[scans]` runs in the `amd64/rust` container, which is Debian with no `LANG`, so `sun.jnu.encoding` resolves to ASCII. The failing set was exactly {4.0, 4.1}, which is exactly the set of jobs that reached the test — 3.4 and 3.5 aborted in `beforeAll` on the `LocalDate` encoder, and 4.2 skips the suite because `icebergAvailable` is false for `isSpark42Plus`. And @sunchao reproduced the filesystem half on Linux/JDK 21. I would still call the historical failure inferred rather than confirmed, per his point. #5651 bumps the pin, so I am not doing it here. Once it lands the string column can go back into the write test's partition spec, and that run is what would actually confirm this. -- 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]
