Use Instant for CurrentDate, more test coverage

Author: olaky

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala

@@ -17,14 +17,14 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
-import java.time.{Instant, LocalDateTime}
+import java.time.{Instant, LocalDate, LocalDateTime}
 
 import org.apache.spark.sql.catalyst.CurrentUserContext.CURRENT_USER
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.catalyst.trees.TreePattern._
-import org.apache.spark.sql.catalyst.util.DateTimeUtils.{convertSpecialDate, convertSpecialTimestamp, convertSpecialTimestampNTZ, currentDate, instantToMicros, localDateTimeToMicros}
+import org.apache.spark.sql.catalyst.util.DateTimeUtils.{convertSpecialDate, convertSpecialTimestamp, convertSpecialTimestampNTZ, instantToMicros, localDateTimeToMicros, localDateToDays}
 import org.apache.spark.sql.connector.catalog.CatalogManager
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
@@ -82,7 +82,9 @@ object ComputeCurrentTime extends Rule[LogicalPlan] {
     plan.transformDownWithSubqueries {
       case subQuery =>
         subQuery.transformAllExpressionsWithPruning(_.containsPattern(CURRENT_LIKE)) {
-          case cd: CurrentDate => Literal.create(currentDate(cd.zoneId).asInstanceOf[Int], DateType)
+          case cd: CurrentDate =>
+            Literal.create(
+              localDateToDays(LocalDate.ofInstant(instant, cd.zoneId)).asInstanceOf[Int], DateType)
           case CurrentTimestamp() | Now() => currentTime
           case CurrentTimeZone() => timezone
           case localTimestamp: LocalTimestamp =>

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala

@@ -20,10 +20,10 @@ package org.apache.spark.sql.catalyst.optimizer
 import java.time.{LocalDateTime, ZoneId}
 
 import scala.collection.JavaConverters.mapAsScalaMap
-import scala.collection.mutable
+import scala.concurrent.duration._
 
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Alias, CurrentDate, CurrentTimeZone, CurrentTimestamp, InSubquery, ListQuery, Literal, LocalTimestamp}
+import org.apache.spark.sql.catalyst.expressions.{Alias, CurrentDate, CurrentTimestamp, CurrentTimeZone, InSubquery, ListQuery, Literal, LocalTimestamp, Now}
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
@@ -44,11 +44,7 @@ class ComputeCurrentTimeSuite extends PlanTest {
     val plan = Optimize.execute(in.analyze).asInstanceOf[Project]
     val max = (System.currentTimeMillis() + 1) * 1000
 
-    val lits = new scala.collection.mutable.ArrayBuffer[Long]
-    plan.transformAllExpressions { case e: Literal =>
-      lits += e.value.asInstanceOf[Long]
-      e
-    }
+    val lits = literals[Long](plan)
     assert(lits.size == 2)
     assert(lits(0) >= min && lits(0) <= max)
     assert(lits(1) >= min && lits(1) <= max)
@@ -62,11 +58,7 @@ class ComputeCurrentTimeSuite extends PlanTest {
     val plan = Optimize.execute(in.analyze).asInstanceOf[Project]
     val max = DateTimeUtils.currentDate(ZoneId.systemDefault())
 
-    val lits = new scala.collection.mutable.ArrayBuffer[Int]
-    plan.transformAllExpressions { case e: Literal =>
-      lits += e.value.asInstanceOf[Int]
-      e
-    }
+    val lits = literals[Int](plan)
     assert(lits.size == 2)
     assert(lits(0) >= min && lits(0) <= max)
     assert(lits(1) >= min && lits(1) <= max)
@@ -76,13 +68,9 @@ class ComputeCurrentTimeSuite extends PlanTest {
   test("SPARK-33469: Add current_timezone function") {
     val in = Project(Seq(Alias(CurrentTimeZone(), "c")()), LocalRelation())
     val plan = Optimize.execute(in.analyze).asInstanceOf[Project]
-    val lits = new scala.collection.mutable.ArrayBuffer[String]
-    plan.transformAllExpressions { case e: Literal =>
-      lits += e.value.asInstanceOf[UTF8String].toString
-      e
-    }
+    val lits = literals[UTF8String](plan)
     assert(lits.size == 1)
-    assert(lits.head == SQLConf.get.sessionLocalTimeZone)
+    assert(lits.head == UTF8String.fromString(SQLConf.get.sessionLocalTimeZone))
   }
 
   test("analyzer should replace localtimestamp with literals") {
@@ -95,11 +83,7 @@ class ComputeCurrentTimeSuite extends PlanTest {
     val plan = Optimize.execute(in.analyze).asInstanceOf[Project]
     val max = DateTimeUtils.localDateTimeToMicros(LocalDateTime.now(zoneId))
 
-    val lits = new scala.collection.mutable.ArrayBuffer[Long]
-    plan.transformAllExpressions { case e: Literal =>
-      lits += e.value.asInstanceOf[Long]
-      e
-    }
+    val lits = literals[Long](plan)
     assert(lits.size == 2)
     assert(lits(0) >= min && lits(0) <= max)
     assert(lits(1) >= min && lits(1) <= max)
@@ -115,15 +99,9 @@ class ComputeCurrentTimeSuite extends PlanTest {
 
     val plan = Optimize.execute(input.analyze).asInstanceOf[Project]
 
-    val literals = new scala.collection.mutable.ArrayBuffer[Long]
-    plan.transformDownWithSubqueries { case subQuery =>
-      subQuery.transformAllExpressions { case expression: Literal =>
-        literals += expression.value.asInstanceOf[Long]
-        expression
-      }
-    }
-    assert(literals.size == 3) // transformDownWithSubqueries covers the inner timestamp twice
-    assert(literals.toSet.size == 1)
+    val lits = literals[Long](plan)
+    assert(lits.size == 3) // transformDownWithSubqueries covers the inner timestamp twice
+    assert(lits.toSet.size == 1)
   }
 
   test("analyzer should use consistent timestamps for different timezones") {
@@ -133,12 +111,38 @@ class ComputeCurrentTimeSuite extends PlanTest {
 
     val plan = Optimize.execute(input).asInstanceOf[Project]
 
-    val literals = new scala.collection.mutable.ArrayBuffer[Long]
-    plan.transformAllExpressions { case e: Literal =>
-      literals += e.value.asInstanceOf[Long]
-      e
-    }
+    val lits = literals[Long](plan)
+    assert(lits.size === localTimestamps.size)
+    // there are timezones with a 30 or 45 minute offset
+    val offsetsFromQuarterHour = lits.map( _ % Duration(15, MINUTES).toMicros).toSet
+    assert(offsetsFromQuarterHour.size == 1)
+  }
 
-    assert(literals.size === localTimestamps.size)
+  test("analyzer should use consistent timestamps for different timestamp functions") {
+    val differentTimestamps = Seq(
+      Alias(CurrentTimestamp(), "currentTimestamp")(),
+      Alias(Now(), "now")(),
+      Alias(LocalTimestamp(Some("PLT")), "localTimestampWithTimezone")(),
+    )
+    val input = Project(differentTimestamps, LocalRelation())
+
+    val plan = Optimize.execute(input).asInstanceOf[Project]
+
+    val lits = literals[Long](plan)
+    assert(lits.size === differentTimestamps.size)
+    // there are timezones with a 30 or 45 minute offset
+    val offsetsFromQuarterHour = lits.map( _ % Duration(15, MINUTES).toMicros).toSet
+    assert(offsetsFromQuarterHour.size == 1)
+  }
+
+  private def literals[T](plan: LogicalPlan): Seq[T] = {
+    val literals = new scala.collection.mutable.ArrayBuffer[T]
+    plan.transformDownWithSubqueries { case subQuery =>
+      subQuery.transformAllExpressions { case expression: Literal =>
+        literals += expression.value.asInstanceOf[T]
+        expression
+      }
+    }
+    literals
   }
 }