如何将数据框拆分为固定大小的组？

Question

我正在处理大型数据框（>100 000 行和多列）。我需要对数据框进行排序，然后将其拆分为预定义大小的相等大小的组。如果有剩余的行（即如果行数不能被组的大小整除），则需要从数据框中删除任何较小的组。

例如1, 2, 3, 4, 5, 6, 7, 8, 9, 10，组大小为3 应该被拆分成[1, 2, 3]、[4, 5, 6]、[7, 8, 9]和10应该被丢弃。

我有一个解决方案，可以使用

创建一个新列

list(range(len(df.index) // group_size)) * group_size

然后使用sort() 和group_by() 将行分组在一起。之后我可以filter 删除任何小于group_size 的组。

示例工作代码：

import pandas as pd

df = pd.DataFrame([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])  # data frame has been sorted before this point and the rows are in the correct order
group_size = 3  

numbers = list(range(len(df.index) // group_size)) * group_size
numbers.sort()
numbers = pd.Series(numbers)
df = pd.concat([df, numbers], ignore_index=True, axis=1)
df.columns = ['value', 'group number']

groups = df.groupby('group number').filter(lambda x: len(x) == group_size)
print(groups)

这非常有效。不幸的是，我有很大的数据框，运行时间太长了。我的方法有替代方案吗？

Answer 1

这将为您提供数据帧列表：

lst = [df.iloc[i:i+group_size] for i in range(0,len(df)-group_size+1,group_size)]

它只是使用内置索引，所以应该很快。如果它太小，停止索引的烦躁会丢弃最后一帧 - 您也可以将其分解为

lst = [df.iloc[i:i+group_size] for i in range(0,len(df),group_size)]
if len(lst[-1]) < group_size:
   lst.pop()

Answer 2

用切片分隔，然后 ffill()。

df['group'] = df[::3]
df['group'].ffill(inplace=True)

现在您可以进行分组并丢弃太小的组。

# df has a RangeIndex, so we get to slice 
group_size = 3
df = pd.DataFrame({'a':[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})  # data frame has been sorted before this point and the rows are in the correct order
slices = df[::group_size]

# but you don't want the group number to be the ordinal at the slices
# so make a copy of the slice to assign good group numbers to it (or get a chained assignment warning)
slices=slices.copy()
slices['group'] = [i for i in range(len(slices))]
df['group'] = slices['group']

# ffill with the nice group numbers
df['group'].ffill(inplace=True)

#now trim the last group
last_group = df['group'].max()
if len(df[df['group']==last_group]) < group_size:
    df = df[df['group'] != last_group]

print(df)

次：

import pandas as pd
from datetime import datetime as dt
print(pd.__version__)


def test1():
    df = pd.DataFrame({'a':[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})  # data frame has been sorted before this point and the rows are in the correct order
    #print(df)
    group_size = 3
    numbers = list(range(len(df.index) // group_size)) * group_size
    numbers.sort()
    numbers = pd.Series(numbers)
    df = pd.concat([df, numbers], ignore_index=True, axis=1)
    df.columns = ['value', 'group number']
    groups = df.groupby('group number').filter(lambda x: len(x) == group_size)
    #print(groups)

def test2():
    # Won't work well because there is no easy way to calculate the remainder that should
    # not be grouped.  But cut() is good for discretizing continuous values
    df = pd.DataFrame({'a':[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})  # data frame has been sorted before this point and the rows are in the correct order
    num_groups = len(df.index)/3
    df['group'] = pd.cut(df['a'], num_groups, right=False)
    #print(df)

def test3():
    # df has a RangeIndex, so we get to slice 
    df = pd.DataFrame({'a':[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})  # data frame has been sorted before this point and the rows are in the correct order
    df['group'] = df[::3]
    df['group'].ffill(inplace=True)
    #print(df['group'])

def test4():
    # A mask can also be used
    df = pd.DataFrame({'a':[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})  # data frame has been sorted before this point and the rows are in the correct order
    df['group'] = df[df.index % 3 == 0]
    df['group'].ffill(inplace=True)
    #print(df)

def test5():
    # maybe go after grouping with iloc
    df = pd.DataFrame({'a':[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]})  # data frame has been sorted before this point and the rows are in the correct order
    group = 0
    for i in range(0,len(df), 3):
        df.loc[i:i+3, 'group'] = group
        group+=1
    #print(df)


funcs = [test1, test2, test3, test4, test5]
for func in funcs:
    print(func.__name__)
    a = dt.now()
    for i in range(1000):
        func()
    b = dt.now()
    print(b-a)

Answer 3

这是 Perigon 答案的变体。就我而言，我不想扔掉最后几个，所以这显示了如何将其余部分放入最终列表中。我正在阅读 CSV，并且想要进行多处理，因此我会将较小的数据帧传递给单独的进程，并且不会丢失 CSV 中的任何行。所以在我的情况下，desired_number_per_group 设置为我想要多进程的相同数量的进程。

    import pandas as pd
    
    test_dict = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]
    
    
    df = pd.DataFrame.from_dict(test_dict)
    
    print ('Size of dataFrame=', len(df.index))
    desired_number_of_groups = 4
    group_size = int(len(df.index) / (desired_number_of_groups))
    print("group_size=", group_size)
    remainder_size = len(df.index) % group_size
    print("remainder_size=", remainder_size)
    df_split_list = [df.iloc[i:i + group_size] for i in range(0, len(df) - group_size + 1, group_size)]
    print("Number of split_dataframes=", len(df_split_list))
    if remainder_size > 0:
        df_remainder = df.iloc[-remainder_size:len(df.index)]
        df_split_list.append(df_remainder)
    print("Revised Number of split_dataframes=", len(df_split_list))
    print("Splitting complete, verifying counts")
    
    count_all_rows_after_split = 0
    for index, split_df in enumerate(df_split_list):
        print("split_df:", index, " size=", len(split_df.index))
        count_all_rows_after_split += len(split_df.index)
    
    if count_all_rows_after_split != len(df.index):
        raise Exception('count_all_rows_after_split = ', count_all_rows_after_split,
                         " but original CSV DataFrame has count =", len(df.index)
                         )

Rich 的单元测试用例做得更好。我刚刚用 1:17、1:18、1:19、1:20、1:21 测试了 test_dict）