[SPARK-35111][SQL] Support Cast string to year-month interval

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

@@ -72,6 +72,7 @@ object Cast {
     case (TimestampType, DateType) => true
 
     case (StringType, CalendarIntervalType) => true
+    case (StringType, YearMonthIntervalType) => true
 
     case (StringType, _: NumericType) => true
     case (BooleanType, _: NumericType) => true
@@ -533,6 +534,11 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
       buildCast[UTF8String](_, s => IntervalUtils.safeStringToInterval(s))
   }
 
+  private[this] def castToYearMonthInterval(from: DataType): Any => Any = from match {
+    case StringType =>
+      buildCast[UTF8String](_, s => IntervalUtils.fromYearMonthString(s).months)
+  }
+
   // LongConverter
   private[this] def castToLong(from: DataType): Any => Any = from match {
     case StringType if ansiEnabled =>
@@ -837,6 +843,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
         case decimal: DecimalType => castToDecimal(from, decimal)
         case TimestampType => castToTimestamp(from)
         case CalendarIntervalType => castToInterval(from)
+        case YearMonthIntervalType => castToYearMonthInterval(from)
         case BooleanType => castToBoolean(from)
         case ByteType => castToByte(from)
         case ShortType => castToShort(from)
@@ -895,6 +902,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
     case decimal: DecimalType => castToDecimalCode(from, decimal, ctx)
     case TimestampType => castToTimestampCode(from, ctx)
     case CalendarIntervalType => castToIntervalCode(from)
+    case YearMonthIntervalType => castToYearMonthIntervalCode(from)
     case BooleanType => castToBooleanCode(from)
     case ByteType => castToByteCode(from, ctx)
     case ShortType => castToShortCode(from, ctx)
@@ -1353,6 +1361,13 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
 
   }
 
+  private[this] def castToYearMonthIntervalCode(from: DataType): CastFunction = from match {
+    case StringType =>
+      val util = IntervalUtils.getClass.getCanonicalName.stripSuffix("$")
+      (c, evPrim, evNull) =>
+        code"""$evPrim = (Integer)$util.fromYearMonthString($c).months;""".stripMargin
+  }
+
   private[this] def decimalToTimestampCode(d: ExprValue): Block = {
     val block = inline"new java.math.BigDecimal($MICROS_PER_SECOND)"
     code"($d.toBigDecimal().bigDecimal().multiply($block)).longValue()"
@@ -1912,6 +1927,7 @@ object AnsiCast {
     case (DateType, TimestampType) => true
 
     case (StringType, _: CalendarIntervalType) => true
+    case (StringType, YearMonthIntervalType) => true
 
     case (StringType, DateType) => true
     case (TimestampType, DateType) => true

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

@@ -93,6 +93,23 @@ object IntervalUtils {
 
   private val yearMonthPattern = "^([+|-])?(\\d+)-(\\d+)$".r
 
+  private val yearMonthFuzzyPattern = "[^-|+]*?([+|-]?[\\d]+-[\\d]+).*".r
+
+  def fromYearMonthString(input: UTF8String): CalendarInterval = {
+    if (input == null || input.toString == null) {
+      throw new IllegalArgumentException("Interval year-month string must be not null")
+    } else {
+      val intervalString = input.trimAll().toString
+      intervalString match {
+        case yearMonthFuzzyPattern(payLoad) =>
+          fromYearMonthString(payLoad)
+        case _ =>
+          throw new IllegalArgumentException(
+            s"Interval string does not match year-month format of 'y-m': $input")
+      }
+    }
+  }
+
   /**
    * Parse YearMonth string in form: [+|-]YYYY-MM
    *

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala

@@ -1774,6 +1774,45 @@ class CastSuite extends CastSuiteBase {
       assert(e3.contains("Casting 2147483648 to int causes overflow"))
     }
   }
+
+  test("SPARK-35111: Cast string to year-month interval") {
+    checkEvaluation(cast(Literal.create("INTERVAL '0-0' YEAR TO MONTH"),
+      YearMonthIntervalType), 0)
+    checkEvaluation(cast(Literal.create("0-0"), YearMonthIntervalType), 0)
+    checkEvaluation(cast(Literal.create("INTERVAL '-1-0' YEAR TO MONTH"),
+      YearMonthIntervalType), -12)
+    checkEvaluation(cast(Literal.create("-1-0"), YearMonthIntervalType), -12)
+    checkEvaluation(cast(Literal.create("INTERVAL '10-1' YEAR TO MONTH"),
+      YearMonthIntervalType), 121)
+    checkEvaluation(cast(Literal.create("10-1"), YearMonthIntervalType), 121)
+    checkEvaluation(cast(Literal.create("abc10-1"), YearMonthIntervalType), 121)
+    checkEvaluation(cast(Literal.create("abca10-1-da-"), YearMonthIntervalType), 121)
+    checkEvaluation(cast(Literal.create(null, StringType), YearMonthIntervalType), null)
+
+    Seq("INTERVAL '0-0' YEAR TO MONTH" -> "INTERVAL '0-0' YEAR TO MONTH",
+      "INTERVAL '10-1' YEAR TO MONTH" -> "INTERVAL '10-1' YEAR TO MONTH",
+      "INTERVAL '-178956970-7' YEAR TO MONTH" -> "INTERVAL '-178956970-7' YEAR TO MONTH",
+      "INTERVAL '178956970-7' YEAR TO MONTH" -> "INTERVAL '178956970-7' YEAR TO MONTH",
+      "INTERVAL '-178956970-8' YEAR TO MONTH" -> "INTERVAL '-178956970-8' YEAR TO MONTH"
+    ).foreach { case (interval, result) =>
+      checkEvaluation(
+        cast(cast(Literal.create(interval), YearMonthIntervalType), StringType), result)
+    }
+
+    Seq("INTERVAL '-178956970-9' YEAR TO MONTH", "INTERVAL '178956970-8' YEAR TO MONTH")
+      .foreach { interval =>
+        val e = intercept[IllegalArgumentException] {
+          cast(Literal.create(interval), YearMonthIntervalType).eval()
+        }.getMessage
+        assert(e.contains("Error parsing interval year-month string: integer overflow"))
+      }
+
+    Seq(Byte.MaxValue, Short.MaxValue, Int.MaxValue, Int.MinValue + 1, Int.MinValue)
+      .foreach { period =>
+        val interval = Literal.create(Period.ofMonths(period), YearMonthIntervalType)
+        checkEvaluation(cast(cast(interval, StringType), YearMonthIntervalType), period)
+      }
+  }
 }
 
 /**