使用Psycopg2写入Spark DataFrame到Redshift出错：无法序列化游标对象

我来帮你搞定这个问题！你遇到的“无法序列化psycopg2.extensions.cursor对象”错误，本质原因是Spark的分布式特性和Psycopg2本地对象的冲突：Spark会把任务分发到多个worker节点执行，而psycopg2的cursor是本地进程内的对象，没办法被序列化后传递给远端worker，直接在全局使用cursor自然就会报错。

下面给你两种靠谱的解决方案，根据你的需求选就行：

方案1：使用Spark Redshift官方连接器（推荐） #

这是最适合Spark写Redshift的方式，专门做了分布式适配，不需要手动管理数据库连接，性能也更优。

首先确保你的环境里已经安装了Spark Redshift连接器（比如通过pip install spark-redshift或者在Spark配置里添加依赖），然后用以下代码写入：

from pyspark.sql import SparkSession

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

# 假设你的目标DataFrame是df
df.write \
  .format("io.github.spark_redshift_community.spark.redshift") \
  .option("url", f"jdbc:redshift://{__credential__.host_redshift}:{__credential__.port_redshift}/{__credential__.dbname_redshift}?user={__credential__.user_redshift}&password={__credential__.password_redshift}") \
  .option("dbtable", "your_target_table")  # 替换成你要写入的表名
  .option("tempdir", "s3://your-temp-bucket/temp-path/")  # 必须指定S3临时目录，Redshift通过它交换数据
  .mode("append")  # 可选模式：overwrite（覆盖）、ignore（忽略已存在）、append（追加）
  .save()

这个方法的优势：自动处理分布式数据写入，不需要关心连接序列化问题，支持批量写入，适合大数据量场景。

方案2：用foreachPartition自定义更新逻辑（适合特殊业务需求） #

如果你的需求不是单纯写入，而是需要像你之前那样执行自定义的UPDATE操作，那可以用foreachPartition方法——每个分区在worker节点上单独处理，我们在分区内部创建连接和cursor，这样就不会涉及到全局对象的序列化了。

代码示例：

def process_partition(partition_rows):
    # 每个分区内单独创建数据库连接和cursor
    conn = psycopg2.connect(
        host=__credential__.host_redshift,
        dbname=__credential__.dbname_redshift,
        user=__credential__.user_redshift,
        password=__credential__.password_redshift,
        port=__credential__.port_redshift
    )
    cur = conn.cursor()
    
    try:
        # 遍历分区内的每一行，执行自定义更新
        for row in partition_rows:
            # 这里替换成你的业务逻辑，比如用row的字段作为参数
            target_col1 = row["col1"]
            new_col2_val = row["col2"]
            cur.execute("UPDATE tb SET col2=%s WHERE col1=%s;", (new_col2_val, target_col1))
        conn.commit()
    except Exception as e:
        # 出错时回滚
        conn.rollback()
        raise e
    finally:
        # 无论成功失败，都关闭cursor和连接
        cur.close()
        conn.close()

# 对DataFrame执行分区级处理
df.foreachPartition(process_partition)

注意事项：

• 每个分区会建立一个数据库连接，所以要控制DataFrame的分区数，避免超过Redshift的最大连接数限制。
• 建议在代码里添加异常处理，避免单个分区出错导致整个任务失败（或者根据需求调整容错策略）。

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