Pandas API on Spark中groupby遍历报错KeyError: (0,)问题咨询

问题场景 #

用户先创建如下Pandas DataFrame：

data = {'Team': ['Riders', 'Riders', 'Devils', 'Devils', 'Kings',
   'kings', 'Kings', 'Kings', 'Riders', 'Royals', 'Royals', 'Riders'],
   'Rank': [1, 2, 2, 3, 3,4 ,1 ,1,2 , 4,1,2],
   'Year': [2014,2015,2014,2015,2014,2015,2016,2017,2016,2014,2015,2017],
   'Points':[876,789,863,673,741,812,756,788,694,701,804,690]}
df = pd.DataFrame(data)

随后通过import pyspark.pandas as ps将其转换为Pandas API on Spark的DataFrame（类型为<class 'pyspark.pandas.frame.DataFrame'>），尝试执行如下遍历groupby结果的代码时：

for i,j in pdf.groupby("Team"):
    print(i)
    print(j)

出现KeyError: (0,)错误，咨询该遍历groupby的功能是否在Pandas API on Spark中可用。

解答 #

Pandas API on Spark的groupby对象不支持直接像原生Pandas那样迭代遍历，这是由Spark的分布式计算特性决定的：Spark的分组结果是分布式存储的，并非本地内存中的可迭代集合，直接触发迭代会引发内部索引错误。

替代方案 #

1. 小数据量场景：转本地Pandas对象遍历
   如果分组后的数据量较小，可以将groupby结果转换为本地Pandas对象再遍历，注意此方法不适用于大数据集（会导致内存溢出）：

for i, j in pdf.groupby("Team").to_pandas():
    print(i)
    print(j)

2. 大数据量场景：使用apply处理分组
   针对分布式大数据，建议用apply方法对每个分组执行处理逻辑，避免将数据拉取到本地：

def process_group(group):
    print(group.name)  # 输出分组键
    print(group)       # 输出分组数据
    return group

pdf.groupby("Team").apply(process_group)

额外提示 #

原数据中Team列存在大小写差异（如Kings和kings），如果需要统一分组逻辑，可先对列值做标准化处理：

pdf['Team'] = pdf['Team'].str.lower()

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