Python Pandas按需求实现explode操作 有没有更快方案及Vaex实现

原有方案性能差的核心原因是使用了嵌套Python循环+逐行修改DataFrame，属于Pandas性能最低的写法，百万级数据下会产生极大的时间和内存开销。可以直接使用Pandas内置的groupby聚合+向量化操作实现，性能提升至少两个数量级：

import pandas as pd 
import numpy as np
# 测试数据和原来一致
df = pd.DataFrame([('bird', 'Falconiformes', 2),
                      ('bird', 'Psittaciformes', 4),
                      ('mammal', 'Carnivora', 8),
                     ('mammal', 'Primates', np.nan),
                    ('mammal', 'Carnivora', 12)],
                     index=['falcon', 'parrot', 'lion', 'monkey', 'leopard'],
                    columns=('class', 'order', 'count'))

# 分组聚合：拼接order，收集count为列表
agg_df = df.groupby('class', as_index=False).agg(
    order=('order', '|'.join),
    count_list=('count', list)
)
# 拆分count列表为多列
count_df = agg_df['count_list'].apply(pd.Series)
count_df.columns = [f'count{i}' for i in count_df.columns]
# 合并得到最终结果
final = pd.concat([agg_df.drop('count_list', axis=1), count_df], axis=1)

运行后得到的结果和原有逻辑完全一致。 #

如果数据量超过内存容量，推荐使用Vaex实现，Vaex采用懒加载机制，全程不会将全量数据载入内存，支持TB级数据处理：

import vaex
import pandas as pd
import numpy as np

# 可以直接从文件读入，不需要转pandas：vdf = vaex.read_csv('your_file.csv')
# 这里用测试数据做示例，从pandas转换
df = pd.DataFrame([('bird', 'Falconiformes', 2),
                      ('bird', 'Psittaciformes', 4),
                      ('mammal', 'Carnivora', 8),
                     ('mammal', 'Primates', np.nan),
                    ('mammal', 'Carnivora', 12)],
                     index=['falcon', 'parrot', 'lion', 'monkey', 'leopard'],
                    columns=('class', 'order', 'count'))
vdf = vaex.from_pandas(df)

# 分组聚合order字段
agg_order = vdf.groupby('class', agg={
    'order': vaex.agg.str_join('order', sep='|')
})
# 分组收集count字段，先转字符串再拼接方便后续拆分
vdf['count_str'] = vdf['count'].astype('string')
agg_count = vdf.groupby('class', agg={
    'count_str_list': vaex.agg.str_join('count_str', sep=',')
})
# 合并两个聚合结果
final_vdf = agg_order.join(agg_count, on='class')

# 拆分count列
max_count = final_vdf['count_str_list'].str.split(',').len().max()
for i in range(max_count):
    final_vdf[f'count{i}'] = final_vdf['count_str_list'].str.split(',').get(i)
    # 转换回数值类型，处理空值
    final_vdf[f'count{i}'] = final_vdf[f'count{i}'].where(final_vdf[f'count{i}'] != 'nan', None).astype('float64')

# 清理多余字段
final_vdf = final_vdf.drop('count_str_list')
# 如果需要转成pandas格式，调用final_vdf.to_pandas()即可

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