将10亿行级SQL表加载到PySpark的正确方法

问题根源 #

你当前的方案虽然设置了numPartitions，但用子查询作为dbtable时，Spark无法自动识别拆分逻辑，只能全量拉取数据到内存，导致内存溢出。以下是针对这类超大规模数据的可行解决思路：

1. 用Spark JDBC原生分区实现并行分块读取 #

这是最高效的方案，核心是让Spark基于指定字段将查询拆分为多个并行任务，避免全量加载。需要满足两个前提：目标表有分布均匀的数值/日期类型字段（比如自增ID、时间戳），且能获取该字段的上下边界。

修改后的读取函数示例：

def read_large_jdbc_table(url, driver, user, password, table_name, select_cols, partition_col, lower_bound, upper_bound, num_partitions, fetch_size):
    return spark.read.format("jdbc") \
        .option("url", url) \
        .option("driver", driver) \
        .option("user", user) \
        .option("password", password) \
        .option("dbtable", table_name) \
        .option("select", select_cols)  # 指定要查询的列，替代子查询
        .option("partitionColumn", partition_col) \
        .option("lowerBound", lower_bound) \
        .option("upperBound", upper_bound) \
        .option("numPartitions", num_partitions) \
        .option("fetchsize", fetch_size) \
        .load()

# 使用示例（假设表有自增ID列，范围1-10亿）
df_result = read_large_jdbc_table(
    url=url,
    driver=driver,
    user=user,
    password=password,
    table_name="db.table",
    select_cols="COL1, COL2",
    partition_col="ID",
    lower_bound=1,
    upper_bound=1000000000,
    num_partitions=1000,
    fetch_size=10000
)

参数说明：

• partitionColumn: 用于拆分数据的列，必须是数值/日期类型
• lowerBound/upperBound: 该列的最小/最大值，Spark会自动拆分出numPartitions个区间并行查询
• fetchsize: 单批次从数据库拉取的行数，建议设为1万-10万，平衡网络和内存开销

2. 绝对禁止触发全量拉取的Action操作 #

你的代码中collect()是致命操作——它会把所有Executor上的数据拉到Driver节点内存，10亿行数据必然导致OOM。所有数据处理应保持延迟计算的Transformation操作，直到最后一步直接写入存储（比如Parquet、Hive表）。

3. 分块查询写入文件再导入的备选方案（针对JDBC分区受限场景） #

如果目标表没有合适的分区字段，可手动拆分查询条件，分批次将数据写入临时存储（比如Parquet），最后合并为DataFrame：

# 手动拆分ID范围，分批次查询
batch_size = 10000000  # 每批次1000万行
for i in range(0, 1000000000, batch_size):
    start = i + 1
    end = i + batch_size
    query = f"(SELECT COL1, COL2 FROM db.table WHERE ID BETWEEN {start} AND {end}) as batch"
    batch_df = spark.read.format("jdbc") \
        .option("url", url) \
        .option("driver", driver) \
        .option("user", user) \
        .option("password", password) \
        .option("dbtable", query) \
        .option("fetchsize", 10000) \
        .load()
    # 写入临时Parquet文件
    batch_df.write.format("parquet").mode("append").save("/tmp/large_table_temp")

# 合并所有临时文件为DataFrame
df_result = spark.read.parquet("/tmp/large_table_temp")

4. 小表关联的性能优化 #

因为另一张表是小表，用broadcast()将其广播到所有Executor节点，避免大规模Shuffle操作：

from pyspark.sql.functions import broadcast

# 加载小表
small_df = spark.read.format("jdbc").option(...).load()

# 关联时广播小表
final_df = df_result.join(broadcast(small_df), on="COL1", how="inner")

# 直接写入最终存储，避免全量加载到内存
final_df.write.format("parquet").mode("overwrite").save("/path/to/final_output")

关键注意事项 #

• 分区列必须分布均匀，否则会出现部分分区数据量过大的倾斜问题
• numPartitions不要设置超过集群Executor的总核心数，避免压垮数据库
• 测试时用limit(100).show()查看数据，绝对不要用show()或collect()查看全量数据

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