不使用collect()遍历PySpark DataFrame的其他方法咨询

除了你尝试过的foreach和map，还有以下几种分布式遍历处理PySpark DataFrame的方法，均不会将数据拉取到Driver端（避免collect()带来的内存压力）：

1. 使用foreachPartition #

针对每个数据分区进行批量遍历处理，能减少Driver与Executor之间的通信开销，适合批量操作场景：

def process_partition(partition):
    for row in partition:
        # 自定义单条数据处理逻辑
        print(f"姓名: {row.firstname} {row.lastName}, 性别: {row.gender}")

df.foreachPartition(process_partition)

2. 使用mapInPandas #

基于Pandas DataFrame进行分区级别的遍历处理，适合熟悉Pandas语法的场景，效率优于普通行处理：

import pandas as pd
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# 定义输出Schema
output_schema = StructType([
    StructField("full_name", StringType()),
    StructField("gender_lower", StringType()),
    StructField("double_salary", IntegerType())
])

def process_pandas_partition(pdf):
    # 对分区内的Pandas DataFrame做批量处理
    pdf["full_name"] = pdf["firstname"] + ", " + pdf["lastName"]
    pdf["gender_lower"] = pdf["gender"].str.lower()
    pdf["double_salary"] = pdf["salary"] * 2
    return pdf[["full_name", "gender_lower", "double_salary"]]

result_df = df.mapInPandas(process_pandas_partition, schema=output_schema)

3. 使用pandas_udf（矢量化UDF） #

利用Pandas矢量化能力实现高效的逐行处理，比普通UDF性能更高，本质是分布式遍历处理：

from pyspark.sql.functions import pandas_udf, col

@pandas_udf(output_schema)
def process_data(firstname, lastName, gender, salary):
    full_name = firstname + ", " + lastName
    gender_lower = gender.str.lower()
    double_salary = salary * 2
    return pd.DataFrame({
        "full_name": full_name,
        "gender_lower": gender_lower,
        "double_salary": double_salary
    })

result_df = df.select(
    process_data(col("firstname"), col("lastName"), col("gender"), col("salary")).alias("processed")
).select("processed.*")

4. 结合SQL与自定义UDF（间接遍历） #

通过注册临时视图，配合自定义UDF在SQL中实现逐行处理，底层也是分布式执行：

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, IntegerType

# 定义单行处理UDF
@udf(StringType())
def get_full_name(firstname, lastName):
    return f"{firstname}, {lastName}"

@udf(StringType())
def lower_gender(gender):
    return gender.lower()

@udf(IntegerType())
def double_salary(salary):
    return salary * 2

# 注册临时视图
df.createOrReplaceTempView("employees")

# 通过SQL完成处理
result_df = spark.sql("""
    SELECT 
        get_full_name(firstname, lastName) as full_name,
        lower_gender(gender) as gender_lower,
        double_salary(salary) as double_salary
    FROM employees
""")

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