[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.safeFromYearMonthString(s).orNull)
+  }
+
   // 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, ctx)
     case BooleanType => castToBooleanCode(from)
     case ByteType => castToByteCode(from, ctx)
     case ShortType => castToShortCode(from, ctx)
@@ -1353,6 +1361,23 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
 
   }
 
+  private[this] def castToYearMonthIntervalCode(
+      from: DataType,
+      ctx: CodegenContext): CastFunction = from match {
+    case StringType =>
+      val util = IntervalUtils.getClass.getCanonicalName.stripSuffix("$")
+      val intOpt = ctx.freshVariable("intOpt", classOf[Option[Int]])
+      (c, evPrim, evNull) =>
+        code"""
+           scala.Option<Integer> $intOpt = $util.safeFromYearMonthString($c);
+           if ($intOpt.isDefined()) {
+              $evPrim = ((Integer) $intOpt.get()).intValue();
+            } else {
+              $evNull = true;
+            }
+         """.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 +1937,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,24 @@ object IntervalUtils {
 
   private val yearMonthPattern = "^([+|-])?(\\d+)-(\\d+)$".r
 
+  def safeFromYearMonthString(input: UTF8String): Option[Int] = {
+    try {
+      if (input == null || input.toString == null) {
+        throw new IllegalArgumentException("Interval year-month string must be not null")
+      } else {
+        val regex = "INTERVAL '([-|+]?[0-9]+-[-|+]?[0-9]+)' YEAR TO MONTH".r
+        // scalastyle:off caselocale .toLowerCase
+        val intervalString = input.trimAll().toUpperCase.toString
+        // scalastyle:on
+        val interval = regex.findFirstMatchIn(intervalString)
+          .map(_.group(1)).getOrElse(intervalString)
+        Some(fromYearMonthString(interval).months)
+      }
+    } catch {
+      case _: IllegalArgumentException => None
+    }
+  }
+
   /**
    * Parse YearMonth string in form: [+|-]YYYY-MM
    *

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

@@ -1774,6 +1774,24 @@ class CastSuite extends CastSuiteBase {
       assert(e3.contains("Casting 2147483648 to int causes overflow"))
     }
   }
+
+  test("SPARK-35111: Cast string to year-month interval") {
+    val regex = "INTERVAL '([-|+]?[0-9]+-[-|+]?[0-9]+)' YEAR TO MONTH".r
+    val m = regex.findFirstMatchIn("INTERVAL '0-0' YEAR TO MONTH")
+    println(m)
+    println(m.map(_.group(1)))
+
+    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("null"), YearMonthIntervalType), null)
+  }
 }
 
 /**