基于配置驱动的Databricks SQL查询多线程并行方案咨询

云原生阿岚

2026-6-29

问题描述

我们公司正处于从本地/Azure MSSQL向Databricks迁移的阶段，目前采用Databricks SQL构建配置驱动的多流程处理模式：从指定表读取所有需处理的配置，将每行配置传入process_table函数，该函数执行Create Table、Merge Into、VACUUM等Spark SQL语句，并全程记录运行时、失败信息等日志。所有配置相互独立、可乱序执行，但当前流程为串行运行，需要实现并行处理。

已尝试将process_table转为UDF但失败（Worker节点无法调用spark.sql），使用ThreadPoolExecutor虽提升了处理速度，但感觉该方式不够规范。现咨询更优的并行化实现方案。

当前串行实现代码

sql_user = 'dbconfigreader'
sql_pass =  dbutils.secrets.get(scope = "REDACTED", key = "REDACTED")
jdbc_url = REDACTED

config_table = (spark.read
  .format("jdbc")
  .option("url",jdbc_url)
  .option("dbtable", "REDACTED")
  .option("user", sql_user)
  .option("password", sql_pass)
  .load()
)


x = "Test"
driver_manager = spark._sc._gateway.jvm.java.sql.DriverManager
con = driver_manager.getConnection(jdbc_url, sql_user, sql_pass)
configs = config_table.where(f"""ScheduleName = '{x}' and IsEnabled = 1""")
config_ls = configs.toJSON().collect()
for config in config_ls:
    failure = process_table(config)
    print(failure)
con.close()

process_table函数代码

import json
from datetime import datetime

def process_table(config):
  sql_statement = None
  columns = None
  failure_reason = None
  success = 1
  start_time = datetime.now()
  item = json.loads(config)
  full_sql = ''
  sql_statement = None
  try:
    full_sql += f"USE CATALOG {environment};"
    spark.sql(f"USE CATALOG {environment}")
    spark.sql(f"""CREATE TABLE IF NOT EXISTS silver.{item['SilverTable']} ({item['StartingTableStructure']}, Hash STRING NOT NULL, CreatedOnDate TIMESTAMP NOT NULL, LastUpdateDate TIMESTAMP, DeletedOnDate TIMESTAMP)""")

    # Run query for Quarantine table
    if 'QuarantineQuery' in item:
      if item['QuarantineQuery'] is not None:
        sql_statement = f"""CREATE OR REPLACE TABLE dq.{item['QuarantineTable']} AS {item['QuarantineQuery']} SELECT * FROM quarantine_final"""
        full_sql += sql_statement
        spark.sql(sql_statement)

    # Run query for stage table
    view = f"""CREATE OR REPLACE TABLE silver.stage_{item['SilverTable']} AS {item['StageQuery']} SELECT * FROM stage_final"""
    full_sql += view + ';'
    spark.sql(view)


    
    hash_columns = ",".join([f"IFNULL({column.strip()}, '')" for column in item['HashedColumns'].split(',')])
    update_columns = ",".join([f"`{column.replace(' ','')}` = Source.`{column.replace(' ','')}`" for column in item['MergeSourceColumns'].split(',')])
    source_columns = ",".join([f"{column.strip()}" for column in item['MergeSourceColumns'].split(',')])

    sql_statement = f"""MERGE INTO silver.{item['SilverTable']} as Target USING (
        SELECT
          {source_columns},
          sha2(
            CONCAT({ hash_columns }),
            256
          ) as Hash
        FROM
          silver.stage_{item['SilverTable']}
      ) as Source
      ON {item['MergeJoin']}
      WHEN MATCHED AND Target.`Hash` <> Source.`Hash` THEN 
      UPDATE SET {update_columns}, `Hash` = Source.`Hash`, LastUpdateDate = GETDATE()
      WHEN NOT MATCHED BY TARGET THEN 
      INSERT ({source_columns}, `Hash`, `CreatedOnDate`, `LastUpdateDate`, `DeletedOnDate`)
      VALUES ({source_columns}, `Hash`, GETDATE(), GETDATE(), NULL)
      WHEN NOT MATCHED BY Source THEN
      UPDATE SET Target.`DeletedOnDate` = GETDATE(); """
    
    full_sql += sql_statement
    spark.sql(sql_statement)

    sql_statement = (f"""DELETE FROM silver.{item['SilverTable']} 
              WHERE dateadd(Day, -5, GETDATE()) > DeletedOnDate""")
    full_sql += sql_statement + ";"
    spark.sql(sql_statement)


    # Clean up stage tables 
    sql_statement = f"""DROP Table silver.stage_{item['SilverTable']};""" 
    full_sql += sql_statement
    spark.sql(sql_statement)
    sql_statement = f"""VACUUM silver.{item['SilverTable']} RETAIN 168 HOURS"""
    full_sql += sql_statement + ";"
    spark.sql(sql_statement)
    if 'QuarantineQuery' in item:
      if spark.catalog.tableExists('silver', item['QuarantineTable']):
        sql_statement = f"""VACUUM dq.{item['QuarantineTable']} RETAIN 168 HOURS"""
        full_sql += sql_statement + ";"
        spark.sql(sql_statement)

  except Exception as e: 
    success = 0
    failure_reason = e

  finally: 
    end_time = datetime.now()
    sql_statement = None
    full_sql = full_sql.replace('\'', '\'\'')
    if failure_reason is None: 
      failure_reason = "NULL"
    sql_statement = f"""EXEC dbo.pPut_dbRunLog '{item['SilverTable']}', '{start_time.strftime('%Y-%m-%d %H:%M:%S')}', '{end_time.strftime('%Y-%m-%d %H:%M:%S')}', {success}, {failure_reason}, '{full_sql}'"""
    exec_statement  = con.prepareCall(sql_statement)
    exec_statement.execute()
    if success == 1:
      sql_statement = f"""EXEC dbo.pPut_dbTableWatermark '{item['SilverTable']}', '{start_time.strftime('%Y-%m-%d %H:%M:%S')}'"""
      exec_statement  = con.prepareCall(sql_statement)
      exec_statement.execute()
    exec_statement.close()
  return failure_reason

ThreadPoolExecutor尝试代码

import concurrent.futures as cf

...


x = "Test"
driver_manager = spark._sc._gateway.jvm.java.sql.DriverManager
con = driver_manager.getConnection(jdbc_url, sql_user, sql_pass)
configs = config_table.where(f"""ScheduleName = '{x}' and IsEnabled = 1""")
config_ls = configs.toJSON().collect()
results = []
with cf.ThreadPoolExecutor(max_workers=2) as executor:
    futures = {executor.submit(process_table, config): config for config in config_ls}
    for future in cf.as_completed(futures):
        results.append(future)
con.close()

优化方案

针对Databricks环境下的配置驱动并行处理，推荐以下三种规范且高效的实现方式：

1. Databricks Jobs API + 任务集群

将每个配置对应的process_table逻辑封装为独立作业，通过Jobs API批量提交并行任务：

核心优势：利用Databricks原生集群调度能力，任务隔离性强，支持自动重试、日志聚合、资源弹性伸缩
实现步骤：
- 将process_table拆分为独立脚本，接收配置参数作为输入
- 遍历配置列表，通过Jobs API为每个配置创建并提交作业，可指定共享任务集群或按需创建小集群
- 通过API统一监控所有任务执行状态，完成后汇总日志

2. Spark foreachPartition + 线程池（规范版）

在Driver端用foreachPartition拆分配置，每个分区内用线程池并行执行，同时解决原方案的线程安全问题：

关键修正：避免共享单个JDBC连接，每个线程独立创建连接；所有Spark SQL操作仍在Driver端线程执行
优化代码示例：

import json
import concurrent.futures as cf
from datetime import datetime
import pyodbc

# 线程安全的JDBC连接获取
def get_jdbc_connection():
    return pyodbc.connect(f"DRIVER={{ODBC Driver 17 for SQL Server}};SERVER={jdbc_url};UID={sql_user};PWD={sql_pass}")

def process_single_config(config):
    conn = get_jdbc_connection()
    failure_reason = None
    success = 1
    start_time = datetime.now()
    try:
        item = json.loads(config)
        # 原process_table中的Spark SQL执行逻辑（保持不变）
        spark.sql(f"USE CATALOG {environment}")
        # ... 其余SQL操作省略
    except Exception as e:
        success = 0
        failure_reason = str(e)
    finally:
        # 日志记录使用当前线程的独立连接
        end_time = datetime.now()
        # ... 原有的日志写入逻辑，替换con为当前conn
        conn.close()
    return failure_reason

# 处理单个分区的配置
def process_config_partition(partition):
    with cf.ThreadPoolExecutor(max_workers=4) as executor:
        executor.map(process_single_config, partition)

# 主逻辑
configs = config_table.where(f"""ScheduleName = '{x}' and IsEnabled = 1""")
# 根据配置数量调整分区数，避免分区过大
configs.repartition(4).rdd.foreachPartition(process_config_partition)

3. Delta Live Tables（DLT）适配方案

如果是批量ETL场景，可将配置对应的表处理逻辑转为DLT管道，通过配置驱动生成多个DLT任务：

核心优势：内置数据质量监控、自动重试、增量处理支持，完全适配Databricks湖仓架构
实现思路：将每个配置对应的Merge、清洗逻辑封装为DLT的Python/SQL模块，通过配置动态生成DLT任务列表

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

基于配置驱动的Databricks SQL查询多线程并行方案咨询

当前串行实现代码 #

process_table函数代码 #

ThreadPoolExecutor尝试代码 #

1. Databricks Jobs API + 任务集群 #

2. Spark foreachPartition + 线程池（规范版） #