如何加快大型 pandas dataframe 的数据标记速度？

Question

I have a large pandas data frame which roughly looks this

  Identity  periods      one        two       three     Label
0   one      1       -0.462407    0.022811  -0.277357
1   one      1       -0.617588    1.667191  -0.370436
2   one      2       -0.604699    0.635473  -0.556088
3   one      2       -0.852943    1.087415  -0.784377
4   two      3        0.421453    2.390097   0.176333
5   two      3       -0.447321   -1.215280  -0.187156
6   two      4        0.398953   -0.334095  -1.194132
7   two      4       -0.324348   -0.842357   0.970825

I need to be able to categorise the data according to groupings in the various columns, for example one of my categorisation criteria is to label each of the groups in the identity column with a label if there is between x and y periods in the periods column.

The code I have to categorise this looks like this, generating a final column:

for i in df['Identity'].unique():
    if (2 <= df[df['Identity']==i]['periods'].max() <= 5) :
        df.loc[df['Identity']==i,'label']='label 1'

I have also tried a version using

df.groupby('Identity').apply().

But this is no quicker.

My data is approximately 2.8m rows at the moment, and there are about 900 unique identities. The code takes about 5 minutes to run, which to me suggests it's the code within the loop that is slow, rather than the looping making it slow.

Answer 1

Let's try to enhance the system performance by using all vectorized Pandas operations instead of using loops or .apply() function which is also just commonly using the relatively slow Python loops internally.

Use .groupby() and .transform() to broadcast max() of periods within group to get a series for making mask. Then use .loc[] with the mask of the condition 2 <= max <=5 and setup label for such rows fulfulling the mask.

Assumed same label for all rows of same Identity group whenever the max period within the group is within 2 <= max <=5.

m = df.groupby('Identity')['periods'].transform('max')
df.loc[(m >=2) & (m <=5), 'Label'] = 'label 1'


print(df)

  Identity  periods       one       two     three    Label
0      one        1 -0.462407  0.022811 -0.277357  label 1
1      one        1 -0.617588  1.667191 -0.370436  label 1
2      one        2 -0.604699  0.635473 -0.556088  label 1
3      one        2 -0.852943  1.087415 -0.784377  label 1
4      two        3  0.421453  2.390097  0.176333  label 1
5      two        3 -0.447321 -1.215280 -0.187156  label 1
6      two        4  0.398953 -0.334095 -1.194132  label 1
7      two        4 -0.324348 -0.842357  0.970825  label 1