如何遍历PySpark Pandas分组DataFrame？报错求解决方案

问题场景 #

尝试用原生Pandas的遍历方式处理PySpark Pandas的分组DataFrame，代码如下：

import pyspark.pandas as ps

dataframe = ps.read_excel("data.xlsx")
groups = dataframe.groupby(['col1', 'col2'])
for name, group in groups:
    print(name)
    ...

运行后触发KeyError: (0,)报错。

报错详情 #

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[29], line 1
----> 1 for name, group in groups:
      2     print(name)

File /opt/spark/python/pyspark/pandas/groupby.py:2806, in DataFrameGroupBy.__getitem__(self, item)
   2803 def __getitem__(self, item: Any) -> GroupBy:
   2804     if self._as_index and is_name_like_value(item):
   2805         return SeriesGroupBy(
-> 2806             self._psdf._psser_for(item if is_name_like_tuple(item) else (item,)),
   2807             self._groupkeys,
   2808             dropna=self._dropna,
   2809         )
   2810     else:
   2811         if is_name_like_tuple(item):

File /opt/spark/python/pyspark/pandas/frame.py:699, in DataFrame._psser_for(self, label)
    672 def _psser_for(self, label: Label) -> "Series":
    673     """
    674     Create Series with a proper column label.
    675 
   (...)
    697     Name: id, dtype: int64
    698     """
--> 699     return self._pssers[label]

KeyError: (0,)

可行解决方案 #

方案1：使用apply处理分组逻辑 #

PySpark Pandas的GroupBy推荐用apply方法执行自定义分组逻辑，这是适配分布式计算的标准方式：

def process_group(group):
    print(f"分组键: {group.name}")
    # 在此编写分组处理逻辑，如统计、转换等
    return group

# 将函数应用到所有分组
processed_df = groups.apply(process_group)

group.name会返回当前分组的键值（对应原代码中的name），处理后的结果会自动合并为DataFrame。

方案2：通过分组键逐个筛选数据 #

若需单独获取每个分组的DataFrame，可先提取所有分组键，再逐一筛选：

# 获取所有分组键的列表
group_keys = groups.groups.keys()

for key in group_keys:
    # 根据分组键筛选对应数据
    group = dataframe[(dataframe['col1'] == key[0]) & (dataframe['col2'] == key[1])]
    print(f"分组键: {key}")
    # 处理当前分组数据

注意：此方法适合分组数量较少的场景，过多分组会频繁触发Spark查询，影响性能。

报错原因 #

PySpark Pandas的GroupBy对象未实现原生Pandas的迭代器协议，直接用for循环遍历会触发内部的__getitem__方法，而该方法无法识别迭代时传入的(0,)键，因此抛出KeyError。

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