基于a0分组并保留a1、a2唯一信息的DataFrame处理需求

原始数据 #

a0 a1    a2
1  1980  Y04
1  1983  Y08
1  1990  Y08
1  1980  Y08
1  1998  Y10
2  2003  Y02
3  1970  Y04
3  2009  Y10
3  1995  Y34

需求 #

按a0字段分组，将a1、a2列的所有唯一值分别整理为列表，得到如下结果：

a0  new_a1                        new_a2
1  [1980, 1983, 1990, 1998]  [Y04, Y08, Y10]
2  [2003]                          [Y02]
3  [1970,2009, 1995]               [Y04]

Pandas实现方案 #

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    'a0': [1,1,1,1,1,2,3,3,3],
    'a1': [1980,1983,1990,1980,1998,2003,1970,2009,1995],
    'a2': ['Y04','Y08','Y08','Y08','Y10','Y02','Y04','Y10','Y34']
})

# 分组聚合：提取唯一值并转为列表
result = df.groupby('a0').agg(
    new_a1=('a1', lambda x: list(pd.unique(x))),
    new_a2=('a2', lambda x: list(pd.unique(x)))
).reset_index()

print(result)

• 使用pd.unique()而非np.unique()，能保留值首次出现的顺序，和需求结果一致；
• 通过agg方法指定列的聚合逻辑，同时重命名为目标列名。

Dask实现方案 #

import dask.dataframe as dd
import pandas as pd

# 构造原始DataFrame并转为Dask DataFrame
df = pd.DataFrame({
    'a0': [1,1,1,1,1,2,3,3,3],
    'a1': [1980,1983,1990,1980,1998,2003,1970,2009,1995],
    'a2': ['Y04','Y08','Y08','Y08','Y10','Y02','Y04','Y10','Y34']
})
dask_df = dd.from_pandas(df, npartitions=2)  # 可根据数据规模调整分区数

# 分组聚合
dask_result = dask_df.groupby('a0').agg(
    new_a1=('a1', lambda x: x.unique().tolist()),
    new_a2=('a2', lambda x: x.unique().tolist())
).reset_index()

# 触发计算得到最终结果
final_result = dask_result.compute()
print(final_result)

• Dask采用延迟计算模式，需调用compute()方法执行运算并获取结果；
• 聚合逻辑与Pandas兼容，x.unique()会返回该分组下的唯一值集合，转为列表即可。

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