[SPARK-59051][SQL] Compile JDBC partition bounds with the dialect

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala

@@ -26,13 +26,13 @@ import org.apache.spark.internal.LogKeys.{CLAUSES, LOWER_BOUND, NEW_VALUE, NUM_P
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession, SQLContext}
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.util.{DateFormatter, DateTimeUtils, TimestampFormatter}
+import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.{getZoneId, stringToDate, stringToTimestamp, stringToTimestampWithoutTimeZone}
 import org.apache.spark.sql.connector.expressions.filter.Predicate
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.jdbc.JdbcDialects
+import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcDialects}
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types.{DataType, DateType, NumericType, StructType, TimestampNTZType, TimestampType}
 import org.apache.spark.unsafe.types.UTF8String
@@ -111,8 +111,9 @@ private[sql] object JDBCRelation extends Logging {
       "Operation not allowed: the lower bound of partitioning column is larger than the upper " +
       s"bound. Lower bound: $lowerBound; Upper bound: $upperBound")
 
+    val dialect = JdbcDialects.get(jdbcOptions.url)
     val boundValueToString: Long => String =
-      toBoundValueInWhereClause(_, partitioning.columnType, timeZoneId)
+      toBoundValueInWhereClause(_, partitioning.columnType, timeZoneId, dialect)
     val numPartitions =
       if ((upperBound - lowerBound) >= partitioning.numPartitions || /* check for overflow */
           (upperBound - lowerBound) < 0) {
@@ -216,26 +217,22 @@ private[sql] object JDBCRelation extends Logging {
   private def toBoundValueInWhereClause(
       value: Long,
       columnType: DataType,
-      timeZoneId: String): String = {
-    def dateTimeToString(): String = {
-      val dateTimeStr = columnType match {
+      timeZoneId: String,
+      dialect: JdbcDialect): String = {
+    def compileDateTimeValue(): String = {
+      val dateTimeValue = columnType match {
         case DateType =>
-          DateFormatter().format(value.toInt)
+          java.sql.Date.valueOf(DateTimeUtils.daysToLocalDate(value.toInt))
         case TimestampType =>
-          val timestampFormatter = TimestampFormatter.getFractionFormatter(
-            DateTimeUtils.getZoneId(timeZoneId))
-          timestampFormatter.format(value)
+          DateTimeUtils.microsToInstant(value).atZone(getZoneId(timeZoneId)).toLocalDateTime
         case TimestampNTZType =>
-          // NTZ micros are zoneless wall-clock values; format in UTC so no zone shift is applied.
-          val timestampFormatter = TimestampFormatter.getFractionFormatter(
-            DateTimeUtils.getZoneId("UTC"))
-          timestampFormatter.format(value)
+          DateTimeUtils.microsToLocalDateTime(value)
       }
-      s"'$dateTimeStr'"
+      dialect.compileValue(dateTimeValue).toString
     }
     columnType match {
       case _: NumericType => value.toString
-      case DateType | TimestampType | TimestampNTZType => dateTimeToString()
+      case DateType | TimestampType | TimestampNTZType => compileDateTimeValue()
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala

@@ -541,6 +541,41 @@ class JDBCSuite extends SharedSparkSession {
     assert(e.getMessage.contains("2018-07-06 10:00:00+05:00"))
   }
 
+  test("SPARK-59051: compile temporal partition bounds with the JDBC dialect") {
+    val testUrl = "jdbc:spark-test:"
+    val quotedColumn = "\"PartitionColumn\""
+    val cases = Seq(
+      (DateType, "2018-07-06", "2018-07-08", "{d '2018-07-07'}"),
+      (TimestampType, "2018-07-06 10:00:00", "2018-07-06 14:00:00",
+        "{ts '2018-07-06 12:00:00.0'}"),
+      (TimestampNTZType, "2018-07-06 10:00:00", "2018-07-06 14:00:00",
+        "{ts '2018-07-06 12:00:00.0'}"))
+
+    JdbcDialects.registerDialectForUrlPrefix(testUrl, OracleDialect())
+    try {
+      cases.foreach { case (dataType, lowerBound, upperBound, compiledMidpoint) =>
+        val schema = StructType(Seq(StructField("PartitionColumn", dataType)))
+        val partitions = JDBCRelation.columnPartition(
+          schema,
+          analysis.caseInsensitiveResolution,
+          "UTC",
+          new JDBCOptions(testUrl, "table", Map(
+            "driver" -> "org.h2.Driver",
+            "lowerBound" -> lowerBound,
+            "upperBound" -> upperBound,
+            "numPartitions" -> "2",
+            "partitionColumn" -> "PartitionColumn")))
+
+        val clauses = partitions.map(_.asInstanceOf[JDBCPartition].whereClause)
+        assert(clauses === Array(
+          s"$quotedColumn < $compiledMidpoint or $quotedColumn is null",
+          s"$quotedColumn >= $compiledMidpoint"))
+      }
+    } finally {
+      JdbcDialects.unregisterDialectForUrlPrefix(testUrl)
+    }
+  }
+
   test("overflow of partition bound difference does not give negative stride") {
     val df = sql("SELECT * FROM partsoverflow")
     checkNumPartitions(df, expectedNumPartitions = 3)