Databricks Spark 3.5.0升级后自定义JDBCDialect的Limit子句失效

在Databricks环境将Spark从3.1.2升级到3.5.0后，自定义SQL Server JDBC方言出现异常：

• 数据类型转换（如TimestampType转datetime2）功能正常
• 但当查询包含Limit子句时，Spark生成LIMIT 10001而非SQL Server支持的TOP 10001

自定义方言Scala代码 #

# custom_dialect.scala
package de.company.test.utils

...imports...

class Datetime2MsSqlServerDialect extends JdbcDialect {
  override def canHandle(url: String): Boolean =
    url.toLowerCase(Locale.ROOT).startsWith("jdbc:sqlserver")

  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
    case TimestampType => Some(JdbcType("DATETIME2", java.sql.Types.TIMESTAMP))
    case TimestampNTZType =>
      Some(JdbcType("DATETIME2", java.sql.Types.TIMESTAMP))
    case StringType  => Some(JdbcType("NVARCHAR(MAX)", java.sql.Types.NVARCHAR))
    case BooleanType => Some(JdbcType("BIT", java.sql.Types.BIT))
    case BinaryType =>
      Some(JdbcType("VARBINARY(MAX)", java.sql.Types.VARBINARY))
    case ShortType if !SQLConf.get.legacyMsSqlServerNumericMappingEnabled =>
      Some(JdbcType("SMALLINT", java.sql.Types.SMALLINT))
    case _ => None
  }

  override def getLimitClause(limit: Integer): String = {
    if (limit > 0) s"TOP ($limit)" else ""
  }

  ...
}

PySpark注册方言代码 #

spark_builder = SparkSession.builder.appName(self.app_name).config("spark.sql.caseSensitive", "true")
spark_session = spark_builder.getOrCreate()

if DatabricksUtils.is_running_on_databricks_cluster():
    self.__register_custom_jdbc_dialect(spark_session.sparkContext)

def __register_custom_jdbc_dialect(self, spark_context):
    from py4j.java_gateway import java_import

    gw = spark_context._gateway
    java_import(gw.jvm, "de.company.test.utils.Datetime2MsSqlServerDialect")
    gw.jvm.org.apache.spark.sql.jdbc.JdbcDialects.registerDialect(
        gw.jvm.de.company.test.utils.Datetime2MsSqlServerDialect()
    )

验证结果 #

> spark._jvm.org.apache.spark.sql.jdbc.JdbcDialects.registerDialect(spark._jvm.de.company.test.utils.Datetime2MsSqlServerDialect())  
> spark._jvm.de.company.test.utils.Datetime2MsSqlServerDialect().getLimitClause(10)
TOP (10)

> spark._jvm.org.apache.spark.sql.jdbc.JdbcDialects.get("jdbc:sqlserver").getClass().getName()
'org.apache.spark.sql.jdbc.AggregatedDialect'

> spark._jvm.org.apache.spark.sql.jdbc.JdbcDialects.get("jdbc:sqlserver").getLimitClause(10)  
'LIMIT 10'

解决方案 #

原因分析 #

Spark 3.2+引入AggregatedDialect机制，当多个方言匹配同一JDBC URL时会自动聚合。默认SQL Server方言（MsSqlServerDialect）优先级高于自定义方言，导致getLimitClause调用的是默认方言的实现，返回LIMIT而非自定义的TOP。

解决步骤 #

1. 提升自定义方言优先级：在Scala方言类中重写priority方法，设置比默认方言更高的优先级（默认MsSqlServerDialect优先级为1，自定义方言设为2即可）：

override def priority: Int = 2

2. 调整注册逻辑确保优先级：修改PySpark注册代码，先移除已匹配的默认方言，再注册自定义方言，避免聚合冲突：

def __register_custom_jdbc_dialect(self, spark_context):
    from py4j.java_gateway import java_import

    gw = spark_context._gateway
    java_import(gw.jvm, "de.company.test.utils.Datetime2MsSqlServerDialect")
    
    # 移除已存在的默认SQL Server方言
    current_dialect = gw.jvm.org.apache.spark.sql.jdbc.JdbcDialects.get("jdbc:sqlserver")
    if current_dialect.getClass().getName() == "org.apache.spark.sql.jdbc.AggregatedDialect":
        for dialect in current_dialect.dialects():
            if dialect.getClass().getName() == "org.apache.spark.sql.jdbc.MsSqlServerDialect":
                gw.jvm.org.apache.spark.sql.jdbc.JdbcDialects.unregisterDialect(dialect)
    
    # 注册自定义方言
    gw.jvm.org.apache.spark.sql.jdbc.JdbcDialects.registerDialect(
        gw.jvm.de.company.test.utils.Datetime2MsSqlServerDialect()
    )

3. 验证修复效果：重新注册后执行验证命令，确认返回结果为TOP (10)：

> spark._jvm.org.apache.spark.sql.jdbc.JdbcDialects.get("jdbc:sqlserver").getLimitClause(10)  
'TOP (10)'

额外说明 #

• 若只需修改部分默认行为，可直接继承MsSqlServerDialect而非JdbcDialect，仅重写需要调整的方法，避免聚合冲突。
• Databricks环境中需确保自定义方言的JAR包已正确部署到集群节点，或通过Workspace库加载。

内容的提问来源于stack exchange，提问作者user3657850