From b6302c07e728ec04aa6bf1c325b5951bb4c483a9 Mon Sep 17 00:00:00 2001 From: charliechen Date: Fri, 4 Sep 2026 16:32:16 -0700 Subject: [PATCH] [SPARK-59051][SQL] Compile JDBC partition bounds with the dialect --- .../datasources/jdbc/JDBCRelation.scala | 29 +++++++-------- .../org/apache/spark/sql/jdbc/JDBCSuite.scala | 35 +++++++++++++++++++ 2 files changed, 48 insertions(+), 16 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala index 1c529279744e0..f8881c205781c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala +++ b/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() } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala index b1114194f9c74..4716268d0dc09 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala +++ b/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)