可变块大小的 Pandas 数据框

Question

我有一个非常大的 DF，其中包含如下数据：

import pandas as pd
df = pd.DataFrame()
df['CODE'] = [1,2,3,1,2,4,2,2,4,5]
df["DATA"] = [ 'AA', 'BB', 'CC', 'DD', 'AA', 'BB', 'EE', 'FF','GG', 'HH']
df.sort_values('CODE')
df
  CODE DATA
 0     1   AA
 3     1   DD
 1     2   BB
 4     2   AA
 6     2   EE
 7     2   FF
 2     3   CC
 5     4   BB
 8     4   GG
 9     5   HH

因为我需要将它分成块并解析它。 但是，CODE 列中包含的 equals 元素不应以不同的块结束，而应将它们添加到前一个块中，即使超出大小也是如此。

基本上，如果我选择 4 行的块大小，第一个块可以增加到包括所有带有“2”的元素并且是：

chunk1:
      CODE DATA
 0     1   AA
 3     1   DD
 1     2   BB
 4     2   AA
 6     2   EE
 7     2   FF

我发现了一些关于分块和分组的帖子，如下所示：

split dataframe into multiple dataframes based on number of rows

但是，上面提供了相同大小的分块，我需要一个考虑 CODE 列中的值的智能分块。

有什么办法吗？

Answer 1

我可能想出了一个解决方案，（仍在测试所有案例），但不是很优雅。

我创建了一个递归函数，返回要采取的间隔：

def findrange(start,step):
        for i in range(start,len(df)+1, step):
            if i+step > len(df): return [i, len(df)]
            if df.CODE[i+step:i+step+1].values != df.CODE[i+step-1:i+step].values:
                return [i,i+step]
        else:
            return findrange(i,step+1)

然后我调用函数来获取范围并处理数据

interval = [0,0]
idx = 0
N=2
while interval[1] < len(df):
    if idx < interval[1]: idx = interval[1]
    interval = findrange(idx, N)
    idx+=N # this point became useless once interval[1] > idx 

我尝试使用 N > 0 的许多不同值发布的 DF，看起来不错。 如果你有更多像我这样的熊猫的方法。

Answer 2

我认为您可以通过cumcount 创建新列GROUPS，然后除以N - 获取每个CODE 值的块：

N = 2
df['GROUPS'] = df.groupby('CODE').cumcount() // N
print (df)
   CODE DATA  GROUPS
0     1   AA       0
3     1   DD       0
1     2   BB       0
4     2   AA       0
6     2   EE       1
7     2   FF       1
2     3   CC       0
5     4   BB       0
8     4   GG       0
9     5   HH       0

groups = df.groupby(['CODE','GROUPS'])
for (frameno, frame) in groups:
    print (frame.to_csv("%s.csv" % frameno))

---

您也可以创建新的Series 并将其用于groupby：

chunked_ser = df.groupby('CODE').cumcount() // N
print (chunked_ser)
0    0
3    0
1    0
4    0
6    1
7    1
2    0
5    0
8    0
9    0
dtype: int64

groups = df.groupby([df.CODE,chunked_ser])
for (frameno, frame) in groups:
    print (frame.to_csv("%s.csv" % frameno))