PySpark Pandas API中groupby.agg结合@pandas_udf使用方法问询

你在使用PySpark的Pandas API执行groupby.agg操作时遇到格式错误，代码因聚合函数参数格式不符合要求抛出错误：

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

你的样本数据：

定义的Pandas UDF：

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import pandas_udf
import pyspark.pandas as ps

@pandas_udf(IntegerType())
def g1(x):
    return ((x == "1-URGENT") | (x == "2-HIGH")).sum()

@pandas_udf(IntegerType())
def g2(x):
    return ((x != "1-URGENT") & (x != "2-HIGH")).sum()

错误的调用方式：

total = jn.groupby("L_SHIPMODE", as_index=False)["O_ORDERPRIORITY"].agg({"O_ORDERPRIORITY": [g1, g2]})

正确调用方式 #

PySpark Pandas API的groupby.agg不支持嵌套字典格式，需直接将列名映射到函数列表，或通过命名参数指定聚合结果列名，以下是两种可行写法：

写法一：命名参数指定聚合列

total = jn.groupby("L_SHIPMODE", as_index=False).agg(
    high_priority=("O_ORDERPRIORITY", g1),
    low_priority=("O_ORDERPRIORITY", g2)
)

写法二：字典映射列到函数列表

total = jn.groupby("L_SHIPMODE", as_index=False).agg({
    "O_ORDERPRIORITY": [g1, g2]
})
# 可选：重命名结果列
total.columns = ["L_SHIPMODE", "high_priority", "low_priority"]

完整示例代码 #

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import pandas_udf
import pyspark.pandas as ps

# 创建样本数据框
data = [("MAIL", "2 -HIGH"), ("SHIP", "1 -URGENT"), ("MAIL", "3 -MEDIUM")]
jn = ps.DataFrame(data, columns=["L_SHIPMODE", "O_ORDERPRIORITY"])

# 定义Pandas UDF（匹配样本数据的空格格式）
@pandas_udf(IntegerType())
def g1(x):
    return ((x == "1 -URGENT") | (x == "2 -HIGH")).sum()

@pandas_udf(IntegerType())
def g2(x):
    return ((x != "1 -URGENT") & (x != "2 -HIGH")).sum()

# 执行聚合
total = jn.groupby("L_SHIPMODE", as_index=False).agg(
    high_priority=("O_ORDERPRIORITY", g1),
    low_priority=("O_ORDERPRIORITY", g2)
)

# 打印结果
print(total)

输出结果 #

L_SHIPMODE  high_priority  low_priority
0        MAIL              1             1
1        SHIP              1             0

关键说明 #

• 注意样本数据中O_ORDERPRIORITY的值带有空格（如"2 -HIGH"），UDF中的判断条件需与实际数据格式匹配，否则会导致计数错误。
• PySpark Pandas API的groupby.agg格式与原生Pandas有差异，禁止使用{"列名": {"列名": [函数]}}这类嵌套结构。

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