[SPARK-32272][SQL] Add SQL standard command SET TIME ZONE

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -240,6 +240,10 @@ statement
     | MSCK REPAIR TABLE multipartIdentifier                            #repairTable
     | op=(ADD | LIST) identifier (STRING | .*?)                        #manageResource
     | SET ROLE .*?                                                     #failNativeCommand
+    | SET TIME ZONE ALL                                                #listTimeZones
+    | SET TIME ZONE interval                                           #setTimeZone
+    | SET TIME ZONE timezone=(STRING | LOCAL)                          #setTimeZone
+    | SET TIME ZONE .*?                                                #setTimeZone
     | SET .*?                                                          #setConfiguration
     | RESET                                                            #resetConfiguration
     | unsupportedHiveNativeCommands .*?                                #failNativeCommand
@@ -1190,6 +1194,7 @@ ansiNonReserved
     | VIEW
     | VIEWS
     | WINDOW
+    | ZONE
 //--ANSI-NON-RESERVED-END
     ;
 
@@ -1431,6 +1436,7 @@ nonReserved
     | TEMPORARY
     | TERMINATED
     | THEN
+    | TIME
     | TO
     | TOUCH
     | TRAILING
@@ -1459,6 +1465,7 @@ nonReserved
     | WINDOW
     | WITH
     | YEAR
+    | ZONE
     ;
 
 // NOTE: If you add a new token in the list below, you should update the list of keywords
@@ -1690,6 +1697,7 @@ TBLPROPERTIES: 'TBLPROPERTIES';
 TEMPORARY: 'TEMPORARY' | 'TEMP';
 TERMINATED: 'TERMINATED';
 THEN: 'THEN';
+TIME: 'TIME';
 TO: 'TO';
 TOUCH: 'TOUCH';
 TRAILING: 'TRAILING';
@@ -1720,6 +1728,7 @@ WHERE: 'WHERE';
 WINDOW: 'WINDOW';
 WITH: 'WITH';
 YEAR: 'YEAR';
+ZONE: 'ZONE';
 //--SPARK-KEYWORD-LIST-END
 //============================
 // End of the keywords list

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -2089,14 +2089,21 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
    * - from-to unit, for instance: interval '1-2' year to month.
    */
   override def visitInterval(ctx: IntervalContext): Literal = withOrigin(ctx) {
+    Literal(parseIntervalLiteral(ctx), CalendarIntervalType)
+  }
+
+  /**
+   * Create a [[CalendarInterval]] object
+   */
+  protected def parseIntervalLiteral(ctx: IntervalContext): CalendarInterval = withOrigin(ctx) {
     if (ctx.errorCapturingMultiUnitsInterval != null) {
       val innerCtx = ctx.errorCapturingMultiUnitsInterval
       if (innerCtx.unitToUnitInterval != null) {
         throw new ParseException(
           "Can only have a single from-to unit in the interval literal syntax",
           innerCtx.unitToUnitInterval)
       }
-      Literal(visitMultiUnitsInterval(innerCtx.multiUnitsInterval), CalendarIntervalType)
+      visitMultiUnitsInterval(innerCtx.multiUnitsInterval)
     } else if (ctx.errorCapturingUnitToUnitInterval != null) {
       val innerCtx = ctx.errorCapturingUnitToUnitInterval
       if (innerCtx.error1 != null || innerCtx.error2 != null) {
@@ -2105,7 +2112,7 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
           "Can only have a single from-to unit in the interval literal syntax",
           errorCtx)
       }
-      Literal(visitUnitToUnitInterval(innerCtx.body), CalendarIntervalType)
+      visitUnitToUnitInterval(innerCtx.body)
     } else {
       throw new ParseException("at least one time unit should be given for interval literal", ctx)
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala

@@ -17,7 +17,8 @@
 
 package org.apache.spark.sql.execution
 
-import java.util.Locale
+import java.time.ZoneOffset
+import java.util.{Locale, TimeZone}
 import javax.ws.rs.core.UriBuilder
 
 import scala.collection.JavaConverters._
@@ -32,10 +33,12 @@ import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.catalyst.parser._
 import org.apache.spark.sql.catalyst.parser.SqlBaseParser._
 import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.util.DateTimeConstants
 import org.apache.spark.sql.execution.command._
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.internal.{HiveSerDe, SQLConf, VariableSubstitution}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.unsafe.types.CalendarInterval
 
 /**
  * Concrete parser for Spark SQL statements.
@@ -90,6 +93,51 @@ class SparkSqlAstBuilder(conf: SQLConf) extends AstBuilder(conf) {
     ResetCommand
   }
 
+  /**
+   * Create a [[SetCommand]] logical plan to set [[SQLConf.SESSION_LOCAL_TIMEZONE]]
+   * Example SQL :
+   * {{{
+   *   SET TIME ZONE LOCAL;
+   *   SET TIME ZONE 'Asia/Shanghai';
+   *   SET TIME ZONE INTERVAL 10 HOURS;
+   * }}}
+   */
+  override def visitSetTimeZone(ctx: SetTimeZoneContext): LogicalPlan = withOrigin(ctx) {
+    val key = SQLConf.SESSION_LOCAL_TIMEZONE.key
+    if (ctx.interval != null) {
+      val interval = parseIntervalLiteral(ctx.interval)
+      if (interval.months != 0 || interval.days != 0 ||
+        math.abs(interval.microseconds) > 18 * DateTimeConstants.MICROS_PER_HOUR) {
+        throw new ParseException("The interval value must be in the range of [-18, +18] hours",
+          ctx.interval())
+      } else {
+        val seconds = (interval.microseconds / DateTimeConstants.MICROS_PER_SECOND).toInt
+        SetCommand(Some(key -> Some(ZoneOffset.ofTotalSeconds(seconds).toString)))
+      }
+    } else if (ctx.timezone != null) {
+      ctx.timezone.getType match {
+        case SqlBaseParser.LOCAL =>
+          SetCommand(Some(key -> Some(TimeZone.getDefault.getID)))
+        case _ =>
+          SetCommand(Some(key -> Some(string(ctx.STRING))))
+      }
+    } else {
+      throw new ParseException("Invalid time zone displacement value", ctx)
+    }
+  }
+
+  /**
+   * Create a [[ListTimeZonesCommand]] to retrieve all the supported region-based Zone IDs
+   * supported.
+   * Example SQL :
+   * {{{
+   *   SET TIME ZONE ALL;
+   * }}}
+   */
+  override def visitListTimeZones(ctx: ListTimeZonesContext): LogicalPlan = withOrigin(ctx) {
+    ListTimeZonesCommand
+  }
+
   /**
    * Create a [[RefreshResource]] logical plan.
    */

sql/core/src/main/scala/org/apache/spark/sql/execution/command/ListTimeZonesCommand.scala

@@ -0,0 +1,49 @@
+/*
+ * 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.spark.sql.execution.command
+
+import java.util.TimeZone
+
+import org.apache.spark.sql.{Row, SparkSession}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference}
+import org.apache.spark.sql.catalyst.util.{DateTimeUtils, TimestampFormatter}
+import org.apache.spark.sql.types.{BooleanType, IntegerType, StringType}
+
+/**
+ * This command is for retrieving all the supported region-based Zone IDs supported.
+ * Command that runs
+ * {{{
+ *   SET TIME ZONE ALL;
+ * }}}
+ */
+case object ListTimeZonesCommand extends RunnableCommand {
+  override val output: Seq[Attribute] =
+    Seq(AttributeReference("id", StringType, nullable = false)(),
+      AttributeReference("name", StringType, nullable = false)(),
+      AttributeReference("offset", IntegerType, nullable = false)(),
+      AttributeReference("useDaylight", BooleanType, nullable = false)())
+  override def run(sparkSession: SparkSession): Seq[Row] = {
+    TimeZone.getAvailableIDs().map { zid =>
+      val zone = DateTimeUtils.getTimeZone(zid)
+      val zoneName = zone.getDisplayName(TimestampFormatter.defaultLocale)
+      val rawOffset = zone.getRawOffset
+      val useDaylight = zone.useDaylightTime()
+      Row(zid, zoneName, rawOffset, useDaylight)
+    }
+  }
+}

sql/core/src/test/resources/sql-tests/inputs/timezone.sql

@@ -0,0 +1,14 @@
+-- valid time zones
+SET TIME ZONE 'Asia/Hong_Kong';
+SET TIME ZONE 'GMT+1';
+SET TIME ZONE INTERVAL 10 HOURS;
+SET TIME ZONE INTERVAL '15:40:32' HOUR TO SECOND;
+SET TIME ZONE LOCAL;
+
+-- invalid time zone
+SET TIME ZONE;
+SET TIME ZONE 'invalid/zone';
+SET TIME ZONE INTERVAL 3 DAYS;
+SET TIME ZONE INTERVAL 24 HOURS;
+SET TIME ZONE INTERVAL '19:40:32' HOUR TO SECOND;
+SET TIME ZONE INTERVAL 10 HOURS 'GMT+1';

sql/core/src/test/resources/sql-tests/results/timezone.sql.out

@@ -0,0 +1,121 @@
+-- Automatically generated by SQLQueryTestSuite
+-- Number of queries: 11
+
+
+-- !query
+SET TIME ZONE 'Asia/Hong_Kong'
+-- !query schema
+struct<key:string,value:string>
+-- !query output
+spark.sql.session.timeZone	Asia/Hong_Kong
+
+
+-- !query
+SET TIME ZONE 'GMT+1'
+-- !query schema
+struct<key:string,value:string>
+-- !query output
+spark.sql.session.timeZone	GMT+1
+
+
+-- !query
+SET TIME ZONE INTERVAL 10 HOURS
+-- !query schema
+struct<key:string,value:string>
+-- !query output
+spark.sql.session.timeZone	+10:00
+
+
+-- !query
+SET TIME ZONE INTERVAL '15:40:32' HOUR TO SECOND
+-- !query schema
+struct<key:string,value:string>
+-- !query output
+spark.sql.session.timeZone	+15:40:32
+
+
+-- !query
+SET TIME ZONE LOCAL
+-- !query schema
+struct<key:string,value:string>
+-- !query output
+spark.sql.session.timeZone	America/Los_Angeles
+
+
+-- !query
+SET TIME ZONE
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.catalyst.parser.ParseException
+
+Invalid time zone displacement value(line 1, pos 0)
+
+== SQL ==
+SET TIME ZONE
+^^^
+
+
+-- !query
+SET TIME ZONE 'invalid/zone'
+-- !query schema
+struct<>
+-- !query output
+java.lang.IllegalArgumentException
+Cannot resolve the given timezone with ZoneId.of(_, ZoneId.SHORT_IDS)
+
+
+-- !query
+SET TIME ZONE INTERVAL 3 DAYS
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.catalyst.parser.ParseException
+
+The interval value must be in the range of [-18, +18] hours(line 1, pos 14)
+
+== SQL ==
+SET TIME ZONE INTERVAL 3 DAYS
+--------------^^^
+
+
+-- !query
+SET TIME ZONE INTERVAL 24 HOURS
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.catalyst.parser.ParseException
+
+The interval value must be in the range of [-18, +18] hours(line 1, pos 14)
+
+== SQL ==
+SET TIME ZONE INTERVAL 24 HOURS
+--------------^^^
+
+
+-- !query
+SET TIME ZONE INTERVAL '19:40:32' HOUR TO SECOND
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.catalyst.parser.ParseException
+
+The interval value must be in the range of [-18, +18] hours(line 1, pos 14)
+
+== SQL ==
+SET TIME ZONE INTERVAL '19:40:32' HOUR TO SECOND
+--------------^^^
+
+
+-- !query
+SET TIME ZONE INTERVAL 10 HOURS 'GMT+1'
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.catalyst.parser.ParseException
+
+Invalid time zone displacement value(line 1, pos 0)
+
+== SQL ==
+SET TIME ZONE INTERVAL 10 HOURS 'GMT+1'
+^^^

sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.internal
 
+import java.util.TimeZone
+
 import scala.language.reflectiveCalls
 
 import org.apache.hadoop.fs.Path
@@ -383,4 +385,20 @@ class SQLConfSuite extends QueryTest with SharedSparkSession {
     }
     assert(e.getMessage === "Cannot resolve the given timezone with ZoneId.of(_, ZoneId.SHORT_IDS)")
   }
+
+  test("set time zone") {
+    TimeZone.getAvailableIDs().foreach { zid =>
+      sql(s"set time zone '$zid'")
+      assert(spark.conf.get(SQLConf.SESSION_LOCAL_TIMEZONE) === zid)
+    }
+    sql("set time zone local")
+    assert(spark.conf.get(SQLConf.SESSION_LOCAL_TIMEZONE) === TimeZone.getDefault.getID)
+
+    val df = sql("set time zone all").where("id='Asia/Chongqing'")
+    checkAnswer(df, Row("Asia/Chongqing", "China Standard Time", 28800000, false))
+
+    val e = intercept[IllegalArgumentException](sql("set time zone 'invalid'"))
+    assert(e.getMessage === "Cannot resolve the given timezone with" +
+      " ZoneId.of(_, ZoneId.SHORT_IDS)")
+  }
 }