PySpark如何遍历大型DataFrame每行数据且避免使用collect()方法

首先明确两种适用场景，你可以根据自己的业务需求选择：

场景1：无需把转换后的自定义对象拉回Driver端，处理后直接落地 #

这类场景是大数据处理的最优选择，全程数据都在Executor端分布式运行，完全避免Driver侧OOM风险：

• 用foreach算子逐行处理
  直接在每个Executor节点上逐行转换对象并执行后续逻辑，示例代码：

def row_handler(row):
    # 替换为你自己的自定义对象转换逻辑
    custom_obj = YourCustomClass(row.string_field, row.other_field)
    # 后续业务逻辑，比如写入数据库、写文件等
    custom_obj.persist()

test_df.foreach(row_handler)

• 用foreachPartition按分区优化处理
  如果自定义对象初始化成本高，或者需要复用数据库连接这类公共资源，推荐用按分区处理的算子，每个分区仅初始化一次公共资源，大幅提升性能：

def partition_handler(partition_rows):
    # 分区级初始化，仅执行一次
    db_connection = get_db_connection()
    for row in partition_rows:
        custom_obj = YourCustomClass(row.string_field, row.other_field)
        custom_obj.save(conn=db_connection)
    # 处理完分区释放资源
    db_connection.close()

test_df.foreachPartition(partition_handler)

注意：分布式处理场景下，你的自定义类需要兼容pickle序列化规则，否则Executor端无法正确实例化对象会抛出序列化异常。

场景2：必须把所有自定义对象拉回Driver端生成本地列表 #

如果你的后续逻辑只能在Driver侧运行，50万条数据体量不算特别极端，可以采用分批拉取的方式降低内存峰值占用，避免全量collect导致的OOM：

• 方案1：用toLocalIterator按分区逐批拉取
  Spark会按顺序拉取单个分区的数据到Driver，处理完再释放内存拉取下一个分区，内存占用远低于全量collect：

result_list = []
for row in test_df.rdd.toLocalIterator():
    # 转换为自定义对象后加入列表
    result_list.append(YourCustomClass(row.string_field, row.other_field))

• 方案2：手动按分区拉取
  如果你需要更灵活的分批控制，可以自己遍历分区ID逐批拉取：

from pyspark.sql.functions import spark_partition_id

result_list = []
partition_count = test_df.rdd.getNumPartitions()
for partition_idx in range(partition_count):
    # 单次仅拉取对应分区的数据
    single_partition_data = test_df.where(spark_partition_id() == partition_idx).collect()
    # 批量转换后加入结果列表
    result_list.extend([YourCustomClass(row.string_field, row.other_field) for row in single_partition_data])

注意：使用本类方案前请确保Driver的内存配置足够容纳所有转换后的自定义对象，否则还是优先选择场景1的分布式处理方案。

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