python - 使用 .agg() 进行 pandas DataFrame 多项操作的进度条

Question

我想将 .agg pandas 操作应用于一个巨大的数据集

例如，我有这段代码:

from tqdm import tqdm
import pandas as pd
df = pd.DataFrame({"A":[1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0], 
                   "B":[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0], 
                   "C":[1.0, 1.5, 2.0, 2.0, 3.0, 4.0, 5.0, 6.0, 10.0],
                   "D":[2.0, 5.0, 3.0, 6.0, 4.0, 2.0, 5.0, 1.0, 2.0],
                   "E":['a', 'a', 'b', 'a', 'b', 'b', 'b', 'a', 'a']}) 

df2 = df.groupby('B').agg({
                           'C': 'mean',
                           'D': 'sum',
                           'E': lambda x: x.mode()
                         })
print(df2)

问题是我的原始数据集有 2.000.000 行。将其转换为 130.000 需要几分钟时间，我希望看到进度条

我试过使用 tqdm 但我不知道如何在这里应用它。有没有类似.progress_apply()但.agg()的函数？

Answer 1

这将在您进行时打印进度，其中进度是通过计算统计数据的组的分数来衡量的。但我不确定循环会减慢您的计算速度。

agger = {
   'C': 'mean',
   'D': 'sum',
   'E': lambda x: x.mode()}


gcols = ['B'] # columns defining the groups
groupby = df.groupby(gcols)

ngroups = len(groupby)
gfrac = 0.3 # fraction of groups for which you want to print progress
gfrac_size = max((1, int(ngroups*gfrac)))
groups = []
rows = []
for i,g in enumerate(groupby):

    if (i+1)%gfrac_size == 0:
        print('{:.0f}% complete'.format(100*(i/ngroups)))

    gstats = g[1].agg(agger)
    if i==0:
        if gstats.ndim==2:
            newcols = gstats.columns.tolist()
        else:
            newcols = gstats.index.tolist()

    groups.append(g[0])
    rows.append(gstats.values.flat)

df3 = pd.DataFrame(np.vstack(rows), columns=newcols)
if len(gcols) == 1:
    df3.index = groups
else:
    df3.index = pd.MultiIndex.from_tuples(groups, names=gcols)
df3 = df3.astype(df[newcols].dtypes)
df3
       C     D  E
1.0  1.5  10.0  a
2.0  3.0  12.0  b
3.0  7.0   8.0  a

另一种(虽然有点 hacky)方法是利用您使用自己的函数 lambda x: x.mode 这一事实。由于您已经在使用此函数时牺牲了速度，因此您可以编写一个类来存储有关进度的信息。例如，

import pandas as pd
import numpy as np
df = pd.DataFrame({"A":[1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0], 
                   "B":[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0], 
                   "C":[1.0, 1.5, 2.0, 2.0, 3.0, 4.0, 5.0, 6.0, 10.0],
                   "D":[2.0, 5.0, 3.0, 6.0, 4.0, 2.0, 5.0, 1.0, 2.0],
                   "E":['a', 'a', 'b', 'a', 'b', 'b', 'b', 'a', 'a']}) 

df2 = df.groupby('B').agg({
                           'C': 'mean',
                           'D': 'sum',
                           'E': lambda x: x.mode()
                         })
print(df2)

class ModeHack:

    def __init__(self, size=5, N=10):
        self.ix = 0
        self.K = 1 
        self.size = size
        self.N = N

    def mode(self, x):
        self.ix = self.ix + x.shape[0]
        if self.K*self.size <= self.ix:
            print('{:.0f}% complete'.format(100*self.ix/self.N))
            self.K += 1

        return x.mode()

    def reset(self):    
        self.ix = 0
        self.K = 1

mymode = ModeHack(size=int(.1*df.shape[0]), N=df.shape[0])
mymode.reset()

agger = {
   'C': 'mean',
   'D': 'sum',
   'E': lambda x: mymode.mode(x)}

df3 = df.groupby('B').agg(agger)