[SPARK-28420][SQL] Support the INTERVAL type in date_part()

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.catalyst.expressions
 
 import java.sql.Timestamp
-import java.time.{DateTimeException, Instant, LocalDate, LocalDateTime, ZoneId}
+import java.time.{DateTimeException, LocalDate, LocalDateTime, ZoneId}
 import java.time.temporal.IsoFields
 import java.util.{Locale, TimeZone}
 
@@ -2032,10 +2032,11 @@ object DatePart {
 }
 
 @ExpressionDescription(
-  usage = "_FUNC_(field, source) - Extracts a part of the date/timestamp.",
+  usage = "_FUNC_(field, source) - Extracts a part of the date/timestamp or interval source.",
   arguments = """
     Arguments:
-      * field - selects which part of the source should be extracted. Supported string values are:
+      * field - selects which part of the source should be extracted.
+               Supported string values of `field` for dates and timestamps are:
                 ["MILLENNIUM", ("MILLENNIA", "MIL", "MILS"),
                  "CENTURY", ("CENTURIES", "C", "CENT"),
                  "DECADE", ("DECADES", "DEC", "DECS"),
@@ -2055,7 +2056,21 @@ object DatePart {
                  "MILLISECONDS", ("MSEC", "MSECS", "MILLISECON", "MSECONDS", "MS"),
                  "MICROSECONDS", ("USEC", "USECS", "USECONDS", "MICROSECON", "US"),
                  "EPOCH"]
-      * source - a date (or timestamp) column from where `field` should be extracted
+                Supported string values of `field` for intervals are:
+                 ["MILLENNIUM", ("MILLENNIA", "MIL", "MILS"),
+                   "CENTURY", ("CENTURIES", "C", "CENT"),
+                   "DECADE", ("DECADES", "DEC", "DECS"),
+                   "YEAR", ("Y", "YEARS", "YR", "YRS"),
+                   "QUARTER", ("QTR"),
+                   "MONTH", ("MON", "MONS", "MONTHS"),
+                   "DAY", ("D", "DAYS"),
+                   "HOUR", ("H", "HOURS", "HR", "HRS"),
+                   "MINUTE", ("M", "MIN", "MINS", "MINUTES"),
+                   "SECOND", ("S", "SEC", "SECONDS", "SECS"),
+                   "MILLISECONDS", ("MSEC", "MSECS", "MILLISECON", "MSECONDS", "MS"),
+                   "MICROSECONDS", ("USEC", "USECS", "USECONDS", "MICROSECON", "US"),
+                   "EPOCH"]
+      * source - a date/timestamp or interval column from where `field` should be extracted
   """,
   examples = """
     Examples:
@@ -2067,6 +2082,10 @@ object DatePart {
        224
       > SELECT _FUNC_('SECONDS', timestamp'2019-10-01 00:00:01.000001');
        1.000001
+      > SELECT _FUNC_('days', interval 1 year 10 months 5 days);
+       5
+      > SELECT _FUNC_('seconds', interval 5 hours 30 seconds 1 milliseconds 1 microseconds);
+       30.001001
   """,
   since = "3.0.0")
 case class DatePart(field: Expression, source: Expression, child: Expression)
@@ -2082,9 +2101,16 @@ case class DatePart(field: Expression, source: Expression, child: Expression)
         Literal(null, DoubleType)
       } else {
         val fieldStr = fieldEval.asInstanceOf[UTF8String].toString
-        DatePart.parseExtractField(fieldStr, source, {
-          throw new AnalysisException(s"Literals of type '$fieldStr' are currently not supported.")
-        })
+        val errMsg = s"Literals of type '$fieldStr' are currently not supported " +
+          s"for the ${source.dataType.catalogString} type."
+        if (source.dataType == CalendarIntervalType) {
+          ExtractIntervalPart.parseExtractField(
+            fieldStr,
+            source,
+            throw new AnalysisException(errMsg))
+        } else {
+          DatePart.parseExtractField(fieldStr, source, throw new AnalysisException(errMsg))
+        }
       }
     })
   }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/intervalExpressions.scala

@@ -0,0 +1,111 @@
+/*
+ * 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.catalyst.expressions
+
+import java.util.Locale
+
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
+import org.apache.spark.sql.catalyst.util.IntervalUtils
+import org.apache.spark.sql.catalyst.util.IntervalUtils._
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.CalendarInterval
+
+abstract class ExtractIntervalPart(
+    child: Expression,
+    val dataType: DataType,
+    func: CalendarInterval => Any,
+    funcName: String)
+  extends UnaryExpression with ExpectsInputTypes with Serializable {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(CalendarIntervalType)
+
+  override protected def nullSafeEval(interval: Any): Any = {
+    func(interval.asInstanceOf[CalendarInterval])
+  }
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val iu = IntervalUtils.getClass.getName.stripSuffix("$")
+    defineCodeGen(ctx, ev, c => s"$iu.$funcName($c)")
+  }
+}
+
+case class ExtractIntervalMillenniums(child: Expression)
+    extends ExtractIntervalPart(child, IntegerType, getMillenniums, "getMillenniums")
+
+case class ExtractIntervalCenturies(child: Expression)
+    extends ExtractIntervalPart(child, IntegerType, getCenturies, "getCenturies")
+
+case class ExtractIntervalDecades(child: Expression)
+    extends ExtractIntervalPart(child, IntegerType, getDecades, "getDecades")
+
+case class ExtractIntervalYears(child: Expression)
+    extends ExtractIntervalPart(child, IntegerType, getYears, "getYears")
+
+case class ExtractIntervalQuarters(child: Expression)
+    extends ExtractIntervalPart(child, ByteType, getQuarters, "getQuarters")
+
+case class ExtractIntervalMonths(child: Expression)
+    extends ExtractIntervalPart(child, ByteType, getMonths, "getMonths")
+
+case class ExtractIntervalDays(child: Expression)
+    extends ExtractIntervalPart(child, LongType, getDays, "getDays")
+
+case class ExtractIntervalHours(child: Expression)
+    extends ExtractIntervalPart(child, ByteType, getHours, "getHours")
+
+case class ExtractIntervalMinutes(child: Expression)
+    extends ExtractIntervalPart(child, ByteType, getMinutes, "getMinutes")
+
+case class ExtractIntervalSeconds(child: Expression)
+    extends ExtractIntervalPart(child, DecimalType(8, 6), getSeconds, "getSeconds")
+
+case class ExtractIntervalMilliseconds(child: Expression)
+    extends ExtractIntervalPart(child, DecimalType(8, 3), getMilliseconds, "getMilliseconds")
+
+case class ExtractIntervalMicroseconds(child: Expression)
+    extends ExtractIntervalPart(child, LongType, getMicroseconds, "getMicroseconds")
+
+// Number of seconds in 10000 years is 315576000001 (30 days per one month)
+// which is 12 digits + 6 digits for the fractional part of seconds.
+case class ExtractIntervalEpoch(child: Expression)
+    extends ExtractIntervalPart(child, DecimalType(18, 6), getEpoch, "getEpoch")
+
+object ExtractIntervalPart {
+
+  def parseExtractField(
+      extractField: String,
+      source: Expression,
+      errorHandleFunc: => Nothing): Expression = extractField.toUpperCase(Locale.ROOT) match {
+    case "MILLENNIUM" | "MILLENNIA" | "MIL" | "MILS" => ExtractIntervalMillenniums(source)
+    case "CENTURY" | "CENTURIES" | "C" | "CENT" => ExtractIntervalCenturies(source)
+    case "DECADE" | "DECADES" | "DEC" | "DECS" => ExtractIntervalDecades(source)
+    case "YEAR" | "Y" | "YEARS" | "YR" | "YRS" => ExtractIntervalYears(source)
+    case "QUARTER" | "QTR" => ExtractIntervalQuarters(source)
+    case "MONTH" | "MON" | "MONS" | "MONTHS" => ExtractIntervalMonths(source)
+    case "DAY" | "D" | "DAYS" => ExtractIntervalDays(source)
+    case "HOUR" | "H" | "HOURS" | "HR" | "HRS" => ExtractIntervalHours(source)
+    case "MINUTE" | "M" | "MIN" | "MINS" | "MINUTES" => ExtractIntervalMinutes(source)
+    case "SECOND" | "S" | "SEC" | "SECONDS" | "SECS" => ExtractIntervalSeconds(source)
+    case "MILLISECONDS" | "MSEC" | "MSECS" | "MILLISECON" | "MSECONDS" | "MS" =>
+      ExtractIntervalMilliseconds(source)
+    case "MICROSECONDS" | "USEC" | "USECS" | "USECONDS" | "MICROSECON" | "US" =>
+      ExtractIntervalMicroseconds(source)
+    case "EPOCH" => ExtractIntervalEpoch(source)
+    case _ => errorHandleFunc
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/IntervalUtils.scala

@@ -0,0 +1,91 @@
+/*
+ * 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.catalyst.util
+
+import org.apache.spark.sql.types.Decimal
+import org.apache.spark.unsafe.types.CalendarInterval
+
+object IntervalUtils {
+  val MONTHS_PER_YEAR: Int = 12
+  val MONTHS_PER_QUARTER: Byte = 3
+  val YEARS_PER_MILLENNIUM: Int = 1000
+  val YEARS_PER_CENTURY: Int = 100
+  val YEARS_PER_DECADE: Int = 10
+  val MICROS_PER_HOUR: Long = DateTimeUtils.MILLIS_PER_HOUR * DateTimeUtils.MICROS_PER_MILLIS
+  val MICROS_PER_MINUTE: Long = DateTimeUtils.MILLIS_PER_MINUTE * DateTimeUtils.MICROS_PER_MILLIS
+  val DAYS_PER_MONTH: Byte = 30
+  val MICROS_PER_MONTH: Long = DAYS_PER_MONTH * DateTimeUtils.SECONDS_PER_DAY
+  /* 365.25 days per year assumes leap year every four years */
+  val MICROS_PER_YEAR: Long = (36525L * DateTimeUtils.MICROS_PER_DAY) / 100
+
+  def getYears(interval: CalendarInterval): Int = {
+    interval.months / MONTHS_PER_YEAR
+  }
+
+  def getMillenniums(interval: CalendarInterval): Int = {
+    getYears(interval) / YEARS_PER_MILLENNIUM
+  }
+
+  def getCenturies(interval: CalendarInterval): Int = {
+    getYears(interval) / YEARS_PER_CENTURY
+  }
+
+  def getDecades(interval: CalendarInterval): Int = {
+    getYears(interval) / YEARS_PER_DECADE
+  }
+
+  def getMonths(interval: CalendarInterval): Byte = {
+    (interval.months % MONTHS_PER_YEAR).toByte
+  }
+
+  def getQuarters(interval: CalendarInterval): Byte = {
+    (getMonths(interval) / MONTHS_PER_QUARTER + 1).toByte
+  }
+
+  def getDays(interval: CalendarInterval): Long = {
+    interval.microseconds / DateTimeUtils.MICROS_PER_DAY
+  }
+
+  def getHours(interval: CalendarInterval): Byte = {
+    ((interval.microseconds % DateTimeUtils.MICROS_PER_DAY) / MICROS_PER_HOUR).toByte
+  }
+
+  def getMinutes(interval: CalendarInterval): Byte = {
+    ((interval.microseconds % MICROS_PER_HOUR) / MICROS_PER_MINUTE).toByte
+  }
+
+  def getMicroseconds(interval: CalendarInterval): Long = {
+    interval.microseconds % MICROS_PER_MINUTE
+  }
+
+  def getSeconds(interval: CalendarInterval): Decimal = {
+    Decimal(getMicroseconds(interval), 8, 6)
+  }
+
+  def getMilliseconds(interval: CalendarInterval): Decimal = {
+    Decimal(getMicroseconds(interval), 8, 3)
+  }
+
+  // Returns total number of seconds with microseconds fractional part in the given interval.
+  def getEpoch(interval: CalendarInterval): Decimal = {
+    var result = interval.microseconds
+    result += MICROS_PER_YEAR * (interval.months / MONTHS_PER_YEAR)
+    result += MICROS_PER_MONTH * (interval.months % MONTHS_PER_YEAR)
+    Decimal(result, 18, 6)
+  }
+}