海量数据上的 pandas groupby 操作优化

在处理海量数据时，使用 pandas 进行数据分组（groupby）操作可能遇到性能瓶颈。例如，一个包含数亿条记录的 CSV 文件，其中包含用户 ID、时间戳和点击的类别。我们需要对用户 ID 进行分组，并计算每个用户对不同类别的喜爱程度。

用户 ID | 时间戳 | 点击的类别
20140512081646222000004-927168801 | 20140722 | 7
20140512081714121000004-383009763 | 20140727 | 4
201405011348508050000041009490586 | 20140728 | 1
20140512081646222000004-927168801 | 20140724 | 1
20140501135024818000004-1623130763 | 20140728 | 3

使用常规的 pandas groupby 操作可能非常耗时，尤其是在数据量很大的情况下。而且，随着数据量的增加，内存使用量也会随之增加，导致内存溢出问题。

2、解决方案

2.1 使用 HDFStore 分块存储

一种解决方法是将数据存储在 HDFStore 中，并以分块的方式进行读取和处理。HDFStore 是一个高效的二进制数据存储格式，它可以将数据存储在多个文件中，并支持快速的数据访问和检索。

clean_input_reader = read_csv(work_path + '/input/input.csv', chunksize=500000)
with get_store(work_path+'/input/input.h5') as store:
    for chunk in clean_input_reader:
        store.append('clean_input', chunk,
                     data_columns=['user_id','timestamp','category_clicked'],
                     min_itemsize=15)

2.2 使用哈希函数对数据进行分组

为了进一步提高性能，我们可以使用哈希函数对数据进行分组。哈希函数可以将数据映射到一个较小的范围，从而减少需要处理的组数。

def sub_group_hash(x):
    return x['user id'].str[-2:]

clean_input_reader = read_csv('input.csv', chunksize=500000)
with get_store('output.h5') as store:
    for chunk in clean_input_reader:
        g = chunk.groupby(sub_group_hash)
        for grp, grouped in g:
            store.append('group_%s' % grp, grouped,
                         data_columns=['user_id','timestamp','category_clicked'],
                         min_itemsize=15)

2.3 并行处理分组数据

为了进一步提高处理速度，我们可以并行处理分组数据。我们可以使用多进程或多线程技术来实现并行处理。

from multiprocessing import Pool

def process_group(grp):
    grouped = store.select(grp)
    grouped.groupby(......).apply(your_cool_function)

with get_store('output.h5') as store:
    pool = Pool(8)
    pool.map(process_group, store.keys())

代码示例

import pandas as pd
import h5py

# 读取数据并存储在 HDFStore 中
clean_input_reader = pd.read_csv('input.csv', chunksize=500000)
with h5py.File('input.h5', 'w') as store:
    for chunk in clean_input_reader:
        store.create_dataset('clean_input', data=chunk,
                             columns=['user_id','timestamp','category_clicked'],
                             dtype='S15')

# 使用哈希函数对数据进行分组
def sub_group_hash(x):
    return x['user id'].str[-2:]

# 将分组数据存储在 HDFStore 中
with h5py.File('output.h5', 'w') as store:
    clean_input_reader = pd.read_hdf('input.h5', 'clean_input', chunksize=500000)
    for chunk in clean_input_reader:
        g = chunk.groupby(sub_group_hash)
        for grp, grouped in g:
            store.create_dataset('group_%s' % grp, data=grouped,
                                 columns=['user_id','timestamp','category_clicked'],
                                 dtype='S15')

# 并行处理分组数据
from multiprocessing import Pool

def process_group(grp):
    grouped = pd.read_hdf('output.h5', 'group_%s' % grp)
    grouped.groupby(......).apply(your_cool_function)

with h5py.File('output.h5', 'r') as store:
    pool = Pool(8)
    pool.map(process_group, store.keys())