Pandas Tips：切分DataFrame

当我们面临一个数据量很大的 DataFrame 时，分而治之是一个很有效的选择，这就涉及到将 DataFrame 进行切分。比如我最近需要对 20 万个样本（保存在一个 DataFrame 里）计算 CAR，我的做法是将这个 DataFrame 切分成 20 个子 DataFrame，然后并行计算。

使用 np.array_split 可以实现这个需求。np.array_split 函数签名如下：

np.array_split(ary, indices_or_sections, axis=0)

indices_or_sections 可以是一个整数，也可以是一个列表。当传入整数，若整数不能整除数组长度，即使数组不能均匀分割，np.array_split 也会通过将部分子数组稍微大一些来处理这种情况。比如：

import numpy as np
import pandas as pd

arr = np.arange(10)
# 将数组分割成 3 份，即使不能整除
np.array_split(arr, 3)

[array([0, 1, 2, 3]), array([4, 5, 6]), array([7, 8, 9])]

以上我们传入一个数组，长度为 10，np.array_split 自动帮我们切分为 3 个 array。如果我们传入一个二维的 array，np.array_split 同样可以切分为多个子 array。因此对于一个较大的 DataFrame，我们可以先转化为 array，然后利用 np.array_split 切分，再转化为 DataFrame。比如，定义这样一个函数：

def get_batches(df: pd.DataFrame, num_splits: int) -> list:
    """将一个 DataFrame 切分为给定数量的子 DataFrame
    Args:
        df (pd.DataFrame): 待切分 DataFrame
        num_splits (int): 切分成多少个子 DataFrame 
        
    Returns:
        list: 子 DataFrame 列表
    """ 
    batch_array_list = np.array_split(df.values, num_splits)
    
    return [pd.DataFrame(batch_array, columns=df.columns) for batch_array in batch_array_list]