PySpark写入数据库：非空约束预检查及BatchUpdateException异常处理问询

一、写入前预检查非空约束 #

可以直接从数据库元数据中获取目标表的非空约束规则，在Spark层提前过滤不符合要求的数据，从根源避免写入报错：

1. 获取目标表的非空约束列
   以PostgreSQL为例，通过Spark JDBC查询系统表提取非空列：

   # 数据库连接配置
   db_config = {
       "url": "jdbc:postgresql://your-host:5432/your-db",
       "dbtable": "(SELECT column_name FROM information_schema.columns WHERE table_schema='public' AND table_name='empl' AND is_nullable='NO') AS non_null_cols",
       "user": "your-user",
       "password": "your-password",
       "driver": "org.postgresql.Driver"
   }
   
   # 读取非空约束列列表
   non_null_cols_df = spark.read.format("jdbc").options(**db_config).load()
   non_null_cols = [row.column_name for row in non_null_cols_df.collect()]
   

2. 分离合规/违规数据
   基于非空列列表过滤DataFrame，将含null的违规行单独存储，只写入合规数据：

   from pyspark.sql.functions import col
   
   # 动态生成非空过滤条件
   filter_expr = " AND ".join([f"`{col_name}` IS NOT NULL" for col_name in non_null_cols])
   valid_df = df.filter(filter_expr)
   # 分离违规行
   invalid_df = df.exceptAll(valid_df)
   
   # 写入合规数据到目标表
   valid_df.write.format("jdbc").options(**db_config).mode("append").save()
   # 把违规行写入错误存储（比如本地文件、HDFS或错误日志表）
   invalid_df.write.mode("append").csv("/path/to/error-records")
   

二、捕获并处理BatchUpdateException #

如果预检查存在遗漏或约束动态变化，可通过捕获异常拆分失败批次，保证程序不终止：

from py4j.protocol import Py4JJavaError
import java.sql

try:
    df.write.format("jdbc").options(**db_config).mode("append").save()
except Py4JJavaError as e:
    cause = e.java_exception.getCause()
    if isinstance(cause, java.sql.BatchUpdateException):
        # 解析失败批次的索引（不同数据库返回规则有差异，PostgreSQL失败批次返回-2）
        update_counts = cause.getUpdateCounts()
        batch_size = 1000  # 需匹配你设置的JDBC批次大小
        failed_ranges = []
        
        for batch_idx, count in enumerate(update_counts):
            if count == -2:
                start = batch_idx * batch_size
                end = start + batch_size
                failed_ranges.append((start, end))
        
        # 标记行索引并分离合规/违规数据
        rdd_with_idx = df.rdd.zipWithIndex()
        
        def is_failed(idx):
            for start, end in failed_ranges:
                if start <= idx < end:
                    return True
            return False
        
        valid_rdd = rdd_with_idx.filter(lambda x: not is_failed(x[1])).map(lambda x: x[0])
        invalid_rdd = rdd_with_idx.filter(lambda x: is_failed(x[1])).map(lambda x: x[0])
        
        # 重新写入合规数据
        spark.createDataFrame(valid_rdd, df.schema).write.format("jdbc").options(**db_config).mode("append").save()
        # 保存违规行
        spark.createDataFrame(invalid_rdd, df.schema).write.mode("append").csv("/path/to/error-records")
    else:
        # 非约束类异常，重新抛出
        raise

注意：不同数据库的getUpdateCounts()返回值规则不同，需根据实际使用的数据库调整判断逻辑。

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