Pandas分组后将产品名聚合为列表/集合至新列报错解决

以下是初始实现代码：

import pandas as pd  # 2.0.3

df = pd.DataFrame(
    {
        "customer_id": [1, 2, 3, 2, 1],
        "order_id": [1, 2, 3, 4, 1],
        "products": ["foo", "bar", "baz", "foo", "bar"],
        "amount": [1, 1, 1, 1, 1]
    }
)

print(df)
grouped = df.groupby(["customer_id", "order_id"])
df["product_order_count"] = grouped["amount"].transform("sum")
df["all_products"] = grouped["products"].agg(list).reset_index()
print(df)

运行后抛出异常：

Traceback (most recent call last):
  File "C:\temp\tt.py", line 15, in <module>
    df["all_orders"] = grouped["products"].agg(list).reset_index()
  File "c:\Users\foo\.venvs\kapa_monitor-38\lib\site-packages\pandas\core\frame.py", line 3940, in __setitem__
    self._set_item_frame_value(key, value)
  File "c:\Users\foo\.venvs\kapa_monitor-38\lib\site-packages\pandas\core\frame.py", line 4094, in _set_item_frame_value
    raise ValueError(
ValueError: Cannot set a DataFrame with multiple columns to the single column all_products

期望输出（all_products为列表或集合形式）：

customer_id  order_id products  amount  product_order_count all_products
0            1         1      foo       1                    2 'foo', 'bar'
1            2         2      bar       1                    1 'bar'
2            3         3      baz       1                    1 'baz'
3            2         4      foo       1                    1 'foo'
4            1         1      bar       1                    2 'foo', 'bar'

错误原因 #

grouped["products"].agg(list).reset_index()返回的是包含customer_id、order_id和聚合后列表的多列DataFrame，你试图把这个多列结构赋值给原DataFrame的单一列all_products，因此触发ValueError。

解决方案 #

使用transform方法替代agg+reset_index，transform会返回和原DataFrame长度一致的结果，自动将聚合值匹配到对应分组的每一行。

1. 聚合为列表 #

import pandas as pd  # 2.0.3

df = pd.DataFrame(
    {
        "customer_id": [1, 2, 3, 2, 1],
        "order_id": [1, 2, 3, 4, 1],
        "products": ["foo", "bar", "baz", "foo", "bar"],
        "amount": [1, 1, 1, 1, 1]
    }
)

grouped = df.groupby(["customer_id", "order_id"])
df["product_order_count"] = grouped["amount"].transform("sum")
# 用transform聚合为列表
df["all_products"] = grouped["products"].transform(list)

print(df)

输出结果：

customer_id  order_id products  amount  product_order_count all_products
0            1         1      foo       1                    2  [foo, bar]
1            2         2      bar       1                    1        [bar]
2            3         3      baz       1                    1        [baz]
3            2         4      foo       1                    1        [foo]
4            1         1      bar       1                    2  [foo, bar]

2. 聚合为集合（自动去重） #

如果需要去重的集合形式，只需把list换成set：

import pandas as pd  # 2.0.3

df = pd.DataFrame(
    {
        "customer_id": [1, 2, 3, 2, 1],
        "order_id": [1, 2, 3, 4, 1],
        "products": ["foo", "bar", "baz", "foo", "bar"],
        "amount": [1, 1, 1, 1, 1]
    }
)

grouped = df.groupby(["customer_id", "order_id"])
df["product_order_count"] = grouped["amount"].transform("sum")
# 用transform聚合为集合
df["all_products"] = grouped["products"].transform(set)

print(df)

输出结果：

customer_id  order_id products  amount  product_order_count all_products
0            1         1      foo       1                    2  {bar, foo}
1            2         2      bar       1                    1        {bar}
2            3         3      baz       1                    1        {baz}
3            2         4      foo       1                    1        {foo}
4            1         1      bar       1                    2  {bar, foo}

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