Double for loop on pandas groupby - python-3.x

Question

I'm trying to have a pairwise comparison on two dataframes based on some key, but I'm having a hard time with pandas groupby in a double for loop since it is very slow. Is there any way I can optimize so that I don't have to recompute the groups every time I run the outer loop?
I tried using the same groupby variable but it doesn't seem to solve the recomputation problem.
mygroups = mydf.groupby('mykey')
for key1,subdf1 in mygroups:
for key2,subdf2 in mygroups:
if(key2 <= key1):
continue
do_some_work(subdf1,subdf2)
subdf2 seems to start recomputing from the first key rather than from the next key after key1. In my use-case scenario I expected that key2 will be the next in the iteration after key1 and so on. How can I have such behavior happen without the need to recompute?

Answer 1

Your observation is correct, the inner loop iterates over the whole dataframe, not just the records after key1.
The easiest way for smaller DataFrames
I would create a list with the groups first and then iterate over this list.
This is what I would do:
mygroups_list= [(key, subdf) for (key, subdf) mydf.groupby('mykey')]
for len(mygroups_list) > 0:
key1,subdf1= mygroups_list.pop(0)
for key2,subdf2 in mygroups_list:
do_some_work(subdf1,subdf2)
You just have to make sure, the groups are really sorted, but AFAIK this is done by the .groupby method anyways. If you are not sure, you can just add a mygroups_list.sort(key=lambda tup: tup[0]) outside your loop.
If size yet does matter
For larger dataframes you can avoid materializing the dataframes at once and just defer that until you actually need the data like this:
# create the groupby object as usual
group_by= mydf.groupby('mykey')
# now fetch the row indices from the groupby object
# and because this is actually a dictionary
# extract the keys from it and sort them
mygroups_dict= group_by.indices
mygroups_labels= list(mygroups_dict)
mygroups_labels.sort()
# now use a similar approach as above
while len(mygroups_labels) > 0:
key1= mygroups_labels.pop(0)
# but instead of creating the sub dataframes
# before you enter the loop, just do it
# within the loop and use the row indices
# the groupby object evaluated
subdf1= mydf.iloc[mygroups_dict[key1]]
for key2 in mygroups_labels:
subdf2= mydf.iloc[mygroups_dict[key2]]
do_some_work(subdf1, subdf2)
That should be much less memory extensive because you just need to store the row indices instead of the whole rows throughout the hole processing time.
For the following example setup:
import numpy as np
def do_some_work(subdf1, subdf2):
print('{} --> {} (len={}/{})'.format(subdf1['mykey'].iat[0], subdf2['mykey'].iat[0], len(subdf1), len(subdf2)))
mydf= pd.DataFrame(dict(mykey=np.random.randint(5, size=100), col=range(1, 101)))
This outputs something like (of course the len info will look different from run to run because of the randint). But note the group labels (left and right of the arrow). On the right side you have key2 which always is > key1:
0 --> 1 (len=21/16)
0 --> 2 (len=21/21)
0 --> 3 (len=21/20)
0 --> 4 (len=21/22)
1 --> 2 (len=16/21)
1 --> 3 (len=16/20)
1 --> 4 (len=16/22)
2 --> 3 (len=21/20)
2 --> 4 (len=21/22)
3 --> 4 (len=20/22)