如何使用psycopg2将PostgreSQL查询结果直接写入PySpark DataFrame

核心结论 #

psycopg2本身无法直接将超大规模PostgreSQL数据集并行加载到PySpark DataFrame，因为psycopg2的查询结果默认全部拉取到运行Python代码的驱动节点，天然不适合Spark分布式计算场景。下面分场景给出解决方案：

场景1：小数据集，坚持使用psycopg2读取后转Spark DataFrame

你拿到的元组列表可以搭配查询返回的字段名直接构造Spark DataFrame，代码如下：

import psycopg2
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType # 可按需替换为对应字段的实际数据类型

spark = SparkSession.builder.appName("pg_test").getOrCreate()

conn = psycopg2.connect(
    database="databasename", 
    user='user', 
    password='pass', 
    host='postgres.host', 
    port= '5432'
)
cur = conn.cursor()
cur.execute("select * from database.table limit 10")
# 获取返回结果的字段名
col_names = [desc[0] for desc in cur.description]
data = cur.fetchall()
# 构造schema，也可以直接用spark.createDataFrame(data, schema=col_names)自动推断类型，大数据量下推荐显式指定类型
schema = StructType([StructField(name, StringType(), True) for name in col_names])
df = spark.createDataFrame(data, schema=schema)

cur.close()
conn.close()

注意：该方案仅适合小批量数据测试，全量拉取大数据集还是会出现驱动节点内存溢出问题，生产环境大数据场景请用下面的方案。

场景2：生产环境大数据集，使用Spark原生JDBC连接器并行读取

这是官方推荐的标准方案，不需要依赖psycopg2，Spark会自动将读取任务分发到多个Executor节点并行加载，不会把全量数据压到驱动节点。首先需要确保你的Spark环境已经包含PostgreSQL JDBC驱动包（可直接下载对应版本的jar包，启动时用--jars参数指定，或者在maven依赖中引入），代码示例：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("pg_jdbc_read").getOrCreate()

# 基础读取配置
pg_url = "jdbc:postgresql://postgres.host:5432/databasename"
pg_properties = {
    "user": "user",
    "password": "pass",
    "driver": "org.postgresql.Driver"
}

# 方式1：小表全量读取
df = spark.read.jdbc(
    url=pg_url,
    table="database.table", # 也可以传入子查询作为表，比如 "(select id, name from database.table where create_time > '2024-01-01') as t"
    properties=pg_properties
)

# 方式2：大表并行读取，指定分片字段、上下限、分片数，实现多Executor并行拉取
df = spark.read.jdbc(
    url=pg_url,
    table="database.table",
    column="id", # 用于分片的整数/时间类型字段，优先选索引字段
    lowerBound=1, # 分片字段的最小值
    upperBound=1000000, # 分片字段的最大值
    numPartitions=10, # 分片数，对应同时拉取的并发任务数
    properties=pg_properties
)

# 后续可直接对df做分布式计算，不需要全量拉到驱动节点
df.show()

并行读取建议选择分布均匀、有索引的字段作为分片键，避免数据倾斜和全表扫描

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