fix: normalize session timezone IDs before passing them to native code

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
     }
   }
 

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"
+}

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)) {

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

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)

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 =

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())