From c2aa0b427c20bb6396da025e3396051c52fce050 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 28 Sep 2026 13:07:20 -0600 Subject: [PATCH] fix: normalize session timezone IDs before passing them to native code Spark resolves spark.sql.session.timeZone with ZoneId.of(id, ZoneId.SHORT_IDS), which accepts IDs such as Z, +8, +08:00:00, GMT+8 and PST. Native code parses the timezone with arrow's Tz, which accepts only IANA names and +HH, +HHMM or +HH:MM offsets, so with those IDs casts, hour, year and df.show() failed at execution time. Add CometTimeZone, which normalizes the timezone every serde passes to native code: fixed offsets become +HH:MM, zero offsets become UTC, and short IDs become their region. An offset with seconds cannot be written for native code, so those expressions are reported as unsupported and go through the codegen dispatcher or fall back. Closes #6329. --- .../apache/comet/expressions/CometCast.scala | 15 +++-- .../apache/comet/serde/CometTimeZone.scala | 65 +++++++++++++++++++ .../org/apache/comet/serde/datetime.scala | 56 ++++++++-------- .../serde/operator/CometNativeScan.scala | 6 +- .../org/apache/comet/serde/structs.scala | 20 ++++-- .../org/apache/comet/serde/unixtime.scala | 8 ++- .../comet/serde/CometToPrettyString.scala | 5 +- .../comet/serde/CometToPrettyString.scala | 5 +- .../datetime/session_timezone_ids.sql | 47 ++++++++++++++ .../session_timezone_ids_unsupported.sql | 36 ++++++++++ .../comet/CometTemporalExpressionSuite.scala | 29 ++++++++- 11 files changed, 247 insertions(+), 45 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/serde/CometTimeZone.scala create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids_unsupported.sql diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index d29ef7cd3b7..969dc67f72c 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason} import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, Compatible, ExprOuterClass, Incompatible, SupportLevel, Unsupported} +import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, CometTimeZone, Compatible, ExprOuterClass, Incompatible, SupportLevel, Unsupported} import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProtoInternal, serializeDataType} import org.apache.comet.shims.{CometExprShim, CometTypeShim} @@ -97,6 +97,8 @@ object CometCast return unsupported(cast.child.dataType, cast.dataType) } Compatible() + } else if (CometTimeZone.nativeId(cast.timeZoneId).isEmpty) { + CometTimeZone.supportLevel(cast.timeZoneId) } else { isSupported(cast.child.dataType, cast.dataType, cast.timeZoneId, evalMode(cast)) } @@ -149,8 +151,8 @@ object CometCast dt: DataType, childExpr: Expr, evalMode: CometEvalMode.Value): Option[Expr] = { - serializeDataType(dt) match { - case Some(dataType) => + (serializeDataType(dt), CometTimeZone.nativeId(timeZoneId)) match { + case (Some(dataType), Some(timeZone)) => val castBuilder = ExprOuterClass.Cast.newBuilder() castBuilder.setChild(childExpr) castBuilder.setDatatype(dataType) @@ -159,16 +161,19 @@ object CometCast SQLConf.get .getConfString(CometConf.getExprAllowIncompatConfigKey(classOf[Cast]), "false") .toBoolean) - castBuilder.setTimezone(timeZoneId.getOrElse("UTC")) + castBuilder.setTimezone(timeZone) castBuilder.setIsSpark4Plus(isSpark40Plus) Some( ExprOuterClass.Expr .newBuilder() .setCast(castBuilder) .build()) - case _ => + case (None, _) => withFallbackReason(expr, s"Unsupported datatype in castToProto: $dt") None + case (_, None) => + withFallbackReason(expr, CometTimeZone.unsupportedReason(timeZoneId)) + None } } diff --git a/spark/src/main/scala/org/apache/comet/serde/CometTimeZone.scala b/spark/src/main/scala/org/apache/comet/serde/CometTimeZone.scala new file mode 100644 index 00000000000..1ed6967a116 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/CometTimeZone.scala @@ -0,0 +1,65 @@ +/* + * 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 java.time.ZoneOffset + +import scala.util.Try + +import org.apache.spark.sql.catalyst.util.DateTimeUtils + +/** + * Converts the timezone Spark stamps on an expression into an ID that native code can parse. + * + * Spark resolves timezone IDs with `ZoneId.of(id, ZoneId.SHORT_IDS)`, which accepts forms such as + * `Z`, `+8`, `+08:00:00`, `GMT+8` and `PST`. Native code parses them with arrow's `Tz`, which + * accepts only IANA zone names and offsets written as `+HH`, `+HHMM` or `+HH:MM`. + */ +object CometTimeZone { + + /** + * The ID to pass to native code for `timeZoneId`, or None when native code cannot represent the + * zone, which is the case for an offset with seconds. A fixed offset becomes `+HH:MM`, a zero + * offset becomes `UTC`, and a short ID becomes its region. An expression with no timezone gets + * `UTC`: Spark leaves it unset only on casts that do not use it. + */ + def nativeId(timeZoneId: Option[String]): Option[String] = timeZoneId match { + case None => Some("UTC") + case Some(id) => + Try(DateTimeUtils.getZoneId(id).normalized()).toOption.flatMap { + case offset: ZoneOffset if offset.getTotalSeconds == 0 => Some("UTC") + case offset: ZoneOffset if offset.getTotalSeconds % 60 == 0 => Some(offset.getId) + case _: ZoneOffset => None + case region => Some(region.getId) + } + } + + /** Whether `timeZoneId` has a zero offset at every instant. */ + def isUtc(timeZoneId: Option[String]): Boolean = nativeId(timeZoneId).contains("UTC") + + /** `Unsupported` when native code cannot represent `timeZoneId`, `Compatible` otherwise. */ + def supportLevel(timeZoneId: Option[String]): SupportLevel = nativeId(timeZoneId) match { + case Some(_) => Compatible() + case None => Unsupported(Some(unsupportedReason(timeZoneId))) + } + + def unsupportedReason(timeZoneId: Option[String]): String = + s"Timezone ${timeZoneId.getOrElse("")} cannot be represented in native code" +} diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index 5106e0059ac..21899def8e9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -169,18 +169,20 @@ object CometQuarter extends CometExpressionSerde[Quarter] with CometExprGetDateF object CometHour extends CometExpressionSerde[Hour] { + override def getSupportLevel(expr: Hour): SupportLevel = + CometTimeZone.supportLevel(expr.timeZoneId) + override def convert( expr: Hour, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val childExpr = exprToProtoInternal(expr.child, inputs, binding) + val timeZone = CometTimeZone.nativeId(expr.timeZoneId) - if (childExpr.isDefined) { + if (childExpr.isDefined && timeZone.isDefined) { val builder = ExprOuterClass.Hour.newBuilder() builder.setChild(childExpr.get) - - val timeZone = expr.timeZoneId.getOrElse("UTC") - builder.setTimezone(timeZone) + builder.setTimezone(timeZone.get) Some( ExprOuterClass.Expr @@ -195,18 +197,20 @@ object CometHour extends CometExpressionSerde[Hour] { object CometMinute extends CometExpressionSerde[Minute] { + override def getSupportLevel(expr: Minute): SupportLevel = + CometTimeZone.supportLevel(expr.timeZoneId) + override def convert( expr: Minute, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val childExpr = exprToProtoInternal(expr.child, inputs, binding) + val timeZone = CometTimeZone.nativeId(expr.timeZoneId) - if (childExpr.isDefined) { + if (childExpr.isDefined && timeZone.isDefined) { val builder = ExprOuterClass.Minute.newBuilder() builder.setChild(childExpr.get) - - val timeZone = expr.timeZoneId.getOrElse("UTC") - builder.setTimezone(timeZone) + builder.setTimezone(timeZone.get) Some( ExprOuterClass.Expr @@ -221,18 +225,20 @@ object CometMinute extends CometExpressionSerde[Minute] { object CometSecond extends CometExpressionSerde[Second] { + override def getSupportLevel(expr: Second): SupportLevel = + CometTimeZone.supportLevel(expr.timeZoneId) + override def convert( expr: Second, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val childExpr = exprToProtoInternal(expr.child, inputs, binding) + val timeZone = CometTimeZone.nativeId(expr.timeZoneId) - if (childExpr.isDefined) { + if (childExpr.isDefined && timeZone.isDefined) { val builder = ExprOuterClass.Second.newBuilder() builder.setChild(childExpr.get) - - val timeZone = expr.timeZoneId.getOrElse("UTC") - builder.setTimezone(timeZone) + builder.setTimezone(timeZone.get) Some( ExprOuterClass.Expr @@ -278,7 +284,7 @@ object CometUnixTimestamp val inputType = expr.children.head.dataType Unsupported(Some(s"unix_timestamp does not support input type: $inputType")) } else { - Compatible() + CometTimeZone.supportLevel(expr.timeZoneId) } } @@ -288,13 +294,12 @@ object CometUnixTimestamp binding: Boolean): Option[ExprOuterClass.Expr] = { // getSupportLevel reports an unsupported input type before reaching here, so no re-check. val childExpr = exprToProtoInternal(expr.children.head, inputs, binding) + val timeZone = CometTimeZone.nativeId(expr.timeZoneId) - if (childExpr.isDefined) { + if (childExpr.isDefined && timeZone.isDefined) { val builder = ExprOuterClass.UnixTimestamp.newBuilder() builder.setChild(childExpr.get) - - val timeZone = expr.timeZoneId.getOrElse("UTC") - builder.setTimezone(timeZone) + builder.setTimezone(timeZone.get) Some( ExprOuterClass.Expr @@ -605,9 +610,10 @@ object CometTruncTimestamp override def getSupportLevel(expr: TruncTimestamp): SupportLevel = { if (DatetimeCollation.hasNonDefaultCollation(expr)) { Incompatible(Some(collationReason)) + } else if (CometTimeZone.nativeId(expr.timeZoneId).isEmpty) { + CometTimeZone.supportLevel(expr.timeZoneId) } else { - val timezone = expr.timeZoneId.getOrElse("UTC") - val isUtc = timezone == "UTC" || timezone == "Etc/UTC" + val isUtc = CometTimeZone.isUtc(expr.timeZoneId) expr.format match { case Literal(fmt: UTF8String, _) => if (supportedFormats.contains(fmt.toString.toLowerCase(Locale.ROOT))) { @@ -631,14 +637,13 @@ object CometTruncTimestamp binding: Boolean): Option[ExprOuterClass.Expr] = { val childExpr = exprToProtoInternal(expr.timestamp, inputs, binding) val formatExpr = exprToProtoInternal(expr.format, inputs, binding) + val timeZone = CometTimeZone.nativeId(expr.timeZoneId) - if (childExpr.isDefined && formatExpr.isDefined) { + if (childExpr.isDefined && formatExpr.isDefined && timeZone.isDefined) { val builder = ExprOuterClass.TruncTimestamp.newBuilder() builder.setChild(childExpr.get) builder.setFormat(formatExpr.get) - - val timeZone = expr.timeZoneId.getOrElse("UTC") - builder.setTimezone(timeZone) + builder.setTimezone(timeZone.get) Some( ExprOuterClass.Expr @@ -716,10 +721,7 @@ object CometDateFormat case _ => false } - private def isUtc(expr: DateFormatClass): Boolean = { - val timezone = expr.timeZoneId.getOrElse("UTC") - timezone == "UTC" || timezone == "Etc/UTC" - } + private def isUtc(expr: DateFormatClass): Boolean = CometTimeZone.isUtc(expr.timeZoneId) override def getSupportLevel(expr: DateFormatClass): SupportLevel = { if (DatetimeCollation.hasNonDefaultCollation(expr)) { diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala index 080c86085b3..9f861332555 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala @@ -36,7 +36,7 @@ import org.apache.comet.CometConf.COMET_EXEC_ENABLED import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus, isSpark41Plus, withFallbackReason} import org.apache.comet.objectstore.NativeConfig import org.apache.comet.parquet.CometParquetUtils -import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel} +import org.apache.comet.serde.{CometOperatorSerde, CometTimeZone, Compatible, OperatorOuterClass, SupportLevel} import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} @@ -277,7 +277,9 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with CometTypeS commonBuilder.addAllDataSchema(dataSchema.asJava) commonBuilder.addAllRequiredSchema(requiredSchema.asJava) commonBuilder.addAllPartitionSchema(partitionSchema.asJava) - commonBuilder.setSessionTimezone(scan.conf.getConfString("spark.sql.session.timeZone")) + val sessionTimeZone = scan.conf.getConfString("spark.sql.session.timeZone") + commonBuilder.setSessionTimezone( + CometTimeZone.nativeId(Some(sessionTimeZone)).getOrElse(sessionTimeZone)) commonBuilder.setCaseSensitive(scan.conf.getConf[Boolean](SQLConf.CASE_SENSITIVE)) // SPARK-53535 (Spark 4.1+): when reading a struct whose requested fields are all diff --git a/spark/src/main/scala/org/apache/comet/serde/structs.scala b/spark/src/main/scala/org/apache/comet/serde/structs.scala index a1cfcadd57e..bb72e4a6cff 100644 --- a/spark/src/main/scala/org/apache/comet/serde/structs.scala +++ b/spark/src/main/scala/org/apache/comet/serde/structs.scala @@ -130,7 +130,8 @@ object CometStructsToJson extends CometCodegenDispatch[StructsToJson] with Nativ " (https://github.com/apache/datafusion-comet/issues/3016)") private def nativeSupported(expr: StructsToJson): Boolean = - expr.options.isEmpty && isSupportedType(expr.child.dataType) + expr.options.isEmpty && isSupportedType(expr.child.dataType) && + CometTimeZone.nativeId(expr.timeZoneId).isDefined override def getSupportLevel(expr: StructsToJson): SupportLevel = if (!CometConf.isExprAllowIncompat(getExprConfigName(expr)) && nativeSupported(expr)) { @@ -151,7 +152,7 @@ object CometStructsToJson extends CometCodegenDispatch[StructsToJson] with Nativ val toJson = ExprOuterClass.ToJson .newBuilder() .setChild(p) - .setTimezone(expr.timeZoneId.getOrElse("UTC")) + .setTimezone(CometTimeZone.nativeId(expr.timeZoneId).get) .setIgnoreNullFields(ignoreNullFields) .build() Some( @@ -198,7 +199,8 @@ object CometJsonToStructs extends CometCodegenDispatch[JsonToStructs] with Nativ Seq("Partially implemented and not comprehensively tested") private def nativeSupported(expr: JsonToStructs): Boolean = - expr.schema != null && isSupportedSchema(expr.schema) + expr.schema != null && isSupportedSchema(expr.schema) && + CometTimeZone.nativeId(expr.timeZoneId).isDefined override def getSupportLevel(expr: JsonToStructs): SupportLevel = if (!CometConf.isExprAllowIncompat(getExprConfigName(expr)) && nativeSupported(expr)) { @@ -242,7 +244,7 @@ object CometJsonToStructs extends CometCodegenDispatch[JsonToStructs] with Nativ .newBuilder() .setChild(childProto) .setSchema(schemaProto) - .setTimezone(expr.timeZoneId.getOrElse("UTC")) + .setTimezone(CometTimeZone.nativeId(expr.timeZoneId).get) .build() ExprOuterClass.Expr.newBuilder().setFromJson(fromJson).build() } @@ -289,6 +291,9 @@ object CometStructsToCsv extends CometExpressionSerde[StructsToCsv] with Codegen s"The schema ${expr.inputSchema} is not supported because " + s"it includes a incompatible data types: $incompatibleDataTypes")) } + if (CometTimeZone.nativeId(expr.timeZoneId).isEmpty) { + return CometTimeZone.supportLevel(expr.timeZoneId) + } // https://github.com/apache/datafusion-comet/issues/3232 Incompatible() } @@ -299,8 +304,9 @@ object CometStructsToCsv extends CometExpressionSerde[StructsToCsv] with Codegen binding: Boolean): Option[ExprOuterClass.Expr] = { for { childProto <- exprToProtoInternal(expr.child, inputs, binding) + timeZone <- CometTimeZone.nativeId(expr.timeZoneId) } yield { - val optionsProto = options2Proto(expr.options, expr.timeZoneId) + val optionsProto = options2Proto(expr.options, timeZone) val toCsv = ExprOuterClass.ToCsv .newBuilder() .setChild(childProto) @@ -312,14 +318,14 @@ object CometStructsToCsv extends CometExpressionSerde[StructsToCsv] with Codegen private def options2Proto( options: Map[String, String], - timeZoneId: Option[String]): ExprOuterClass.CsvWriteOptions = { + timeZone: String): ExprOuterClass.CsvWriteOptions = { ExprOuterClass.CsvWriteOptions .newBuilder() .setDelimiter(options.getOrElse("delimiter", ",")) .setQuote(options.getOrElse("quote", "\"")) .setEscape(options.getOrElse("escape", "\\")) .setNullValue(options.getOrElse("nullValue", "")) - .setTimezone(timeZoneId.getOrElse("UTC")) + .setTimezone(timeZone) .setIgnoreLeadingWhiteSpace(options .get("ignoreLeadingWhiteSpace") .flatMap(ignoreLeadingWhiteSpace => Try(ignoreLeadingWhiteSpace.toBoolean).toOption) diff --git a/spark/src/main/scala/org/apache/comet/serde/unixtime.scala b/spark/src/main/scala/org/apache/comet/serde/unixtime.scala index 861799fcb9e..cd725d6466e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/unixtime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/unixtime.scala @@ -55,6 +55,8 @@ object CometFromUnixTime extends CometExpressionSerde[FromUnixTime] with Codegen override def getSupportLevel(expr: FromUnixTime): SupportLevel = { if (expr.format != Literal(TimestampFormatter.defaultPattern())) { Unsupported(Some(formatReason)) + } else if (CometTimeZone.nativeId(expr.timeZoneId).isEmpty) { + CometTimeZone.supportLevel(expr.timeZoneId) } else { Incompatible(Some(timestampRangeReason)) } @@ -72,7 +74,11 @@ object CometFromUnixTime extends CometExpressionSerde[FromUnixTime] with Codegen // https://github.com/apache/datafusion/issues/14536 // After fixing these issues, use provided `format` instead of the manual replacement below val formatExpr = exprToProtoInternal(Literal("%Y-%m-%d %H:%M:%S"), inputs, binding) - val timeZone = exprToProtoInternal(Literal(expr.timeZoneId.orNull), inputs, binding) + val timeZone = + exprToProtoInternal( + Literal(CometTimeZone.nativeId(expr.timeZoneId).orNull), + inputs, + binding) if (secExpr.isDefined && formatExpr.isDefined) { val timestampExpr = diff --git a/spark/src/main/spark-3.5/org/apache/comet/serde/CometToPrettyString.scala b/spark/src/main/spark-3.5/org/apache/comet/serde/CometToPrettyString.scala index 147eeb1e29c..7952a8418f6 100644 --- a/spark/src/main/spark-3.5/org/apache/comet/serde/CometToPrettyString.scala +++ b/spark/src/main/spark-3.5/org/apache/comet/serde/CometToPrettyString.scala @@ -31,6 +31,9 @@ object CometToPrettyString extends CometExpressionSerde[ToPrettyString] { Seq("Falls back to Spark when the input type cannot be cast to string.") override def getSupportLevel(expr: ToPrettyString): SupportLevel = { + if (CometTimeZone.nativeId(expr.timeZoneId).isEmpty) { + return CometTimeZone.supportLevel(expr.timeZoneId) + } CometCast.isSupported( expr.child.dataType, DataTypes.StringType, @@ -51,7 +54,7 @@ object CometToPrettyString extends CometExpressionSerde[ToPrettyString] { val tps = ExprOuterClass.ToPrettyString .newBuilder() .setChild(p) - .setTimezone(expr.timeZoneId.getOrElse("UTC")) + .setTimezone(CometTimeZone.nativeId(expr.timeZoneId).get) .setBinaryOutputStyle(binaryOutputStyle) .build() Some(ExprOuterClass.Expr.newBuilder().setToPrettyString(tps).build()) diff --git a/spark/src/main/spark-4.x/org/apache/comet/serde/CometToPrettyString.scala b/spark/src/main/spark-4.x/org/apache/comet/serde/CometToPrettyString.scala index 147eeb1e29c..7952a8418f6 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/serde/CometToPrettyString.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/serde/CometToPrettyString.scala @@ -31,6 +31,9 @@ object CometToPrettyString extends CometExpressionSerde[ToPrettyString] { Seq("Falls back to Spark when the input type cannot be cast to string.") override def getSupportLevel(expr: ToPrettyString): SupportLevel = { + if (CometTimeZone.nativeId(expr.timeZoneId).isEmpty) { + return CometTimeZone.supportLevel(expr.timeZoneId) + } CometCast.isSupported( expr.child.dataType, DataTypes.StringType, @@ -51,7 +54,7 @@ object CometToPrettyString extends CometExpressionSerde[ToPrettyString] { val tps = ExprOuterClass.ToPrettyString .newBuilder() .setChild(p) - .setTimezone(expr.timeZoneId.getOrElse("UTC")) + .setTimezone(CometTimeZone.nativeId(expr.timeZoneId).get) .setBinaryOutputStyle(binaryOutputStyle) .build() Some(ExprOuterClass.Expr.newBuilder().setToPrettyString(tps).build()) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids.sql b/spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids.sql new file mode 100644 index 00000000000..53f497e9345 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids.sql @@ -0,0 +1,47 @@ +-- 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. + +-- Spark accepts session timezone IDs that native code cannot parse as written: `Z`, offsets +-- such as `+8`, `-08` and `+08:00:00`, prefixed offsets such as `GMT+8`, and short IDs such as +-- `PST`. Comet normalizes them before passing them to native code. +-- https://github.com/apache/datafusion-comet/issues/6329 + +-- ConfigMatrix: spark.sql.session.timeZone=GMT+8,UTC+08:00,+8,-08,+08:00:00,Z,PST,IST + +statement +CREATE TABLE test_session_tz_ids(ts timestamp, ntz timestamp_ntz, d date, s string) USING parquet + +statement +INSERT INTO test_session_tz_ids VALUES + (TIMESTAMP'2024-01-15 18:30:45Z', TIMESTAMP_NTZ'2024-01-15 18:30:45', DATE'2024-01-15', '2024-01-15 18:30:45'), + (TIMESTAMP'2024-06-30 23:30:00Z', TIMESTAMP_NTZ'2024-06-30 23:30:00', DATE'2024-06-30', '2024-06-30 23:30:00'), + (NULL, NULL, NULL, NULL) + +query +SELECT CAST(ts AS STRING), CAST(ts AS DATE), hour(ts), minute(ts) FROM test_session_tz_ids + +query +SELECT year(ts), month(ts), dayofmonth(ts) FROM test_session_tz_ids + +query +SELECT CAST(s AS TIMESTAMP), CAST(d AS TIMESTAMP), unix_timestamp(d) FROM test_session_tz_ids + +query +SELECT CAST(ntz AS TIMESTAMP), CAST(ts AS TIMESTAMP_NTZ) FROM test_session_tz_ids + +query +SELECT date_trunc('HOUR', ts), date_format(ts, 'yyyy-MM-dd HH:mm') FROM test_session_tz_ids diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids_unsupported.sql b/spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids_unsupported.sql new file mode 100644 index 00000000000..98ca546c163 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/session_timezone_ids_unsupported.sql @@ -0,0 +1,36 @@ +-- 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. + +-- Native code cannot express an offset with seconds, so expressions that need the session +-- timezone do not run natively. +-- https://github.com/apache/datafusion-comet/issues/6329 + +-- Config: spark.sql.session.timeZone=+05:45:30 + +statement +CREATE TABLE test_session_tz_seconds(ts timestamp) USING parquet + +statement +INSERT INTO test_session_tz_seconds VALUES (TIMESTAMP'2024-01-15 18:30:45Z'), (NULL) + +-- casts go through the codegen dispatcher +query +SELECT CAST(ts AS STRING), CAST(ts AS DATE) FROM test_session_tz_seconds + +-- hour has no dispatcher path, so it falls back to Spark +query expect_fallback(cannot be represented in native code) +SELECT hour(ts) FROM test_session_tz_seconds diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 4d950b8718f..6f62d1f2eac 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus -import org.apache.comet.serde.{CometDateFormat, CometTruncDate, CometTruncTimestamp} +import org.apache.comet.serde.{CometDateFormat, CometTimeZone, CometTruncDate, CometTruncTimestamp} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -108,6 +108,33 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH checkSparkAnswerAndOperator("SELECT c0, trunc(c0, c1) from tbl order by c0, c1") } + test("session timezone IDs are normalized for native code") { + // Spark accepts these through `ZoneId.of(id, ZoneId.SHORT_IDS)`, but native code only parses + // IANA names and `+HH`, `+HHMM` or `+HH:MM` offsets. + val expected = Seq( + "UTC" -> Some("UTC"), + "Etc/UTC" -> Some("UTC"), + "GMT" -> Some("UTC"), + "Z" -> Some("UTC"), + "+00:00" -> Some("UTC"), + "+8" -> Some("+08:00"), + "-08" -> Some("-08:00"), + "+08:00:00" -> Some("+08:00"), + "GMT+8" -> Some("+08:00"), + "UTC+08:00" -> Some("+08:00"), + "EST" -> Some("-05:00"), + "PST" -> Some("America/Los_Angeles"), + "IST" -> Some("Asia/Kolkata"), + "America/Los_Angeles" -> Some("America/Los_Angeles"), + // native code has no way to express an offset with seconds + "+05:45:30" -> None) + for ((id, nativeId) <- expected) { + assert(CometTimeZone.nativeId(Some(id)) == nativeId, id) + } + // Spark leaves the timezone unset only on casts that do not use it + assert(CometTimeZone.nativeId(None).contains("UTC")) + } + test("date_trunc (TruncTimestamp) - reading from DataFrame") { val supportedFormats = CometTruncTimestamp.supportedFormats val unsupportedFormats = Seq("invalid")