PySpark大DataFrame逐行转JSON并发送至Kafka的问题排查

我太懂你这个痛点了——用collect()直接把全量DataFrame拉到Driver端，别说超大数据集了，稍微大一点就直接OOM。而且用UDF的时候踩了「传入列而非整行Row」的坑，确实容易懵。下面给你两个实用的解决方案，完美匹配你的需求：

方法一：用foreachPartition批量处理（最推荐） #

这个方法完全避开了把全量数据拉到Driver端的问题，因为foreachPartition是让每个Executor单独处理自己负责的分区数据，内存压力分散到各个节点，还能批量发送Kafka消息提升效率。

代码示例：

from kafka import KafkaProducer

def send_partition_to_kafka(partition):
    # 每个分区创建一个Kafka客户端（避免重复创建连接，减少开销）
    producer = KafkaProducer(bootstrap_servers='your_kafka_broker:9092')
    topic = 'your_target_topic'
    
    for row_json in partition:
        # 直接用Spark原生toJSON生成的字符串发送
        producer.send(topic, value=row_json.encode('utf-8'))
    
    # 确保分区内所有消息都发送完成
    producer.flush()

# 关键操作：用toJSON()把每行转成JSON字符串，再用foreachPartition分布式处理
df.toJSON().foreachPartition(send_partition_to_kafka)

为什么这个方法好用？ #

• 不会把全量数据拉到Driver，完全分布式处理，从根源上避免内存溢出
• 每个分区复用一个Kafka连接，减少连接建立的开销
• 用Spark原生的toJSON()生成JSON，格式完全符合你要的{"Col1":"A","Col2":1}样式

方法二：自定义UDF处理整行数据（适合需要自定义JSON格式的场景） #

如果你需要对JSON格式做定制化调整（比如修改字段名、过滤某些字段），可以用UDF，但要注意必须把整行打包成一个结构体传入UDF，而不是单独传列。

代码示例：

from pyspark.sql.functions import udf, struct
from pyspark.sql.types import StringType
import json

# 定义接收Row对象的UDF，把Row转成JSON字符串
def row_to_json(row):
    # 把Row转成Python字典再序列化，确保格式正确
    return json.dumps(row.asDict())

# 注册UDF
row_to_json_udf = udf(row_to_json, StringType())

# 把所有列打包成一个struct，这样UDF才能拿到完整的Row对象
df_with_json = df.withColumn("json_str", row_to_json_udf(struct(*df.columns)))

# 之后同样用foreachPartition发送Kafka，和方法一逻辑一致
def send_json_to_kafka(partition):
    producer = KafkaProducer(bootstrap_servers='your_kafka_broker:9092')
    topic = 'your_target_topic'
    
    for row in partition:
        producer.send(topic, value=row.json_str.encode('utf-8'))
    
    producer.flush()

df_with_json.select("json_str").foreachPartition(send_json_to_kafka)

注意点： #

• 必须用struct(*df.columns)把所有列打包成一个结构体，否则UDF只能拿到单个列的值，无法获取整行数据
• 用row.asDict()把Row转成Python字典，再用json.dumps()序列化，能保证生成的JSON格式和你预期一致

测试验证 #

用你给出的示例DataFrame测试：

df = spark.createDataFrame([("A", 1), ("B", 2), ("D", 3)],["Col1", "Col2"])

两种方法生成的JSON字符串都会是：
{"Col1":"A","Col2":1}、{"Col1":"B","Col2":2}、{"Col1":"D","Col2":3}，完全符合你的要求。

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