PySpark Pandas如何使用Lambda函数实现分组聚合？

问题描述 #

原标准Pandas分组聚合代码：

df1 = df1.groupby(['FISCAL_YEAR', 'FISCAL_MONTH', 'FISCAL_WEEK','ORDER_NUMBER', 'LINE_TYPE'], as_index=False).agg({
    'DEPARTMENT': lambda x: ' | '.join(sorted(x.unique()))
    ,'Dept_Subdept': lambda x: ' | '.join(sorted(x.unique()))
    ,'Demand': 'sum'
    ,'COGS': 'sum'
    ,'Units':'sum'
})

切换到PySpark Pandas DataFrame后报错：

ValueError: aggs must be a dict mapping from column name to aggregate functions (string or list of strings).

解决方案 #

1. 适配PySpark Pandas的写法 #

PySpark Pandas不支持自定义lambda作为聚合函数，需改用Spark原生函数组合实现去重排序拼接逻辑：

import pyspark.pandas as ps
from pyspark.sql import functions as F

# 转换为原生Spark DataFrame执行高效聚合，再转回PySpark Pandas
spark_df = df1.to_spark()
agg_spark_df = spark_df.groupBy(
    'FISCAL_YEAR', 'FISCAL_MONTH', 'FISCAL_WEEK', 'ORDER_NUMBER', 'LINE_TYPE'
).agg(
    F.concat_ws(' | ', F.sort_array(F.collect_set('DEPARTMENT'))).alias('DEPARTMENT'),
    F.concat_ws(' | ', F.sort_array(F.collect_set('Dept_Subdept'))).alias('Dept_Subdept'),
    F.sum('Demand').alias('Demand'),
    F.sum('COGS').alias('COGS'),
    F.sum('Units').alias('Units')
)
df1 = agg_spark_df.to_pandas_on_spark()

2. 原生Spark DataFrame直接实现 #

如果无需保留PySpark Pandas格式，直接用原生Spark代码更高效：

from pyspark.sql import functions as F

# 假设spark_df是你的输入Spark DataFrame
aggregated_df = spark_df.groupBy(
    'FISCAL_YEAR', 'FISCAL_MONTH', 'FISCAL_WEEK', 'ORDER_NUMBER', 'LINE_TYPE'
).agg(
    # 去重→排序→拼接字符串
    F.concat_ws(' | ', F.sort_array(F.collect_set('DEPARTMENT'))).alias('DEPARTMENT'),
    F.concat_ws(' | ', F.sort_array(F.collect_set('Dept_Subdept'))).alias('Dept_Subdept'),
    # 数值求和
    F.sum('Demand').alias('Demand'),
    F.sum('COGS').alias('COGS'),
    F.sum('Units').alias('Units')
)

# 查看结果
aggregated_df.show()

核心说明 #

• PySpark Pandas的agg仅支持字符串型聚合函数名（如'sum'），不支持自定义lambda，这是和标准Pandas的核心差异。
• Spark原生函数组合collect_set(去重)+sort_array(排序)+concat_ws(拼接)完全等价于原lambda逻辑，且分布式执行效率更高，适配大数据量场景。

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