improve code

apache · Jun 30, 2020 · 3e28461 · 3e28461
1 parent 6802459
commit 3e28461
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 14 deletions.
diff --git a/...alyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/...alyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
@@ -2612,8 +2612,10 @@ object Sequence {
       val stepDays = step.days
       val stepMicros = step.microseconds
 
-      require(scale != MICROS_PER_DAY || stepMonths != 0 || stepDays != 0,
-        "sequence step must be a day interval if start and end values are dates")
+      if(scale == MICROS_PER_DAY && stepMonths == 0 && stepDays == 0) {
+        throw new IllegalArgumentException(
+          "sequence step must be a day interval if start and end values are dates")
+      }
 
       if (stepMonths == 0 && stepMicros == 0 && scale == MICROS_PER_DAY) {
         backedSequenceImpl.eval(start, stop, fromLong(stepDays))
@@ -2677,15 +2679,23 @@ object Sequence {
            |${genSequenceLengthCode(ctx, startMicros, stopMicros, intervalInMicros, arrLength)}
           """.stripMargin
 
+      val check = if (scale == MICROS_PER_DAY) {
+        s"""
+           if ($stepMonths == 0 && $stepDays == 0) {
+            throw new IllegalArgumentException(
+               "sequence step must be a day interval if start and end values are dates");
+           }
+         """
+      } else {
+        ""
+      }
+
       s"""
          |final int $stepMonths = $step.months;
          |final int $stepDays = $step.days;
          |final long $stepMicros = $step.microseconds;
          |
-         |if (${scale}L == ${MICROS_PER_DAY}L && $stepMonths == 0 && $stepDays == 0) {
-         |  throw new IllegalArgumentException(
-         |    "sequence step must be a day interval if start and end values are dates");
-         |}
+         |$check
          |
          |if ($stepMonths == 0 && $stepMicros == 0 && ${scale}L == ${MICROS_PER_DAY}L) {
          |  ${backedSequenceImpl.genCode(ctx, start, stop, stepDays, arr, elemType)};

diff --git a/...src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala b/...src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
@@ -1857,13 +1857,15 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
 
   test("SPARK-32133: Sequence step must be a day interval " +
     "if start and end values are dates") {
-    val e = intercept[Exception](
-      checkEvaluation(Sequence(
-        Cast(Literal("2011-03-01"), DateType),
-        Cast(Literal("2011-04-01"), DateType),
-        Option(Literal(stringToInterval("interval 1 hour")))),
-        Seq(Date.valueOf("2011-03-01"), Date.valueOf("2011-04-01"))))
-    assert(e.getCause.getMessage.contains(
-      "sequence step must be a day interval if start and end values are dates"))
+    checkExceptionInExpression[IllegalArgumentException](Sequence(
+      Cast(Literal("2011-03-01"), DateType),
+      Cast(Literal("2011-04-01"), DateType),
+      Option(Literal(stringToInterval("interval 1 hour")))), null,
+      "sequence step must be a day interval if start and end values are dates")
+    checkEvaluation(Sequence(
+      Cast(Literal("2011-03-01"), DateType),
+      Cast(Literal("2011-03-02"), DateType),
+      Option(Literal(stringToInterval("interval 1 day")))),
+      Seq(Date.valueOf("2011-03-01"), Date.valueOf("2011-03-02")))
   }
 }