How to optimize a function which contains for loops and 20 millions rows in dataframe

user_8034652 0
Viewed 52

I have a pandas dataframe df that looks like the following:

student_id   category_id  count
1              111        10
2              111        5
3              222        8
4              333        5
5              111        6

Likewise I have 20 million rows.

I want to compute the rating for each student_id. For example, let's consider a category_id “111”. We have 3 student_ids 1, 2 and 5 in this category. student_id 1 has 10 counts, student_id 2 has 5 counts and student_id 5 has 6 counts. The rating of each student_id for category_id is calculated by the formula:

(count per student_id / total number of counts per category_id) * 5

For student_id 1 -> 10/ 21 * 5 = 2.38

For student_id 2 -> 5/ 21 *5 = 1.19

For student_id 5 -> 6/ 21 * 5 = 1.43

Below is the function I already have to compute this:

countPerStudentID = datasetPandas.groupby('student_id').agg(list)
countPerCategoryID = datasetPandas.groupby('category_id').agg(list)

studentIDMap = dict()
def func1(student_id):
    if student_id in studentIDMap:
        return studentIDMap[student_id]
    runningSum = 0
    countList = countPerStudentID.loc[student_id, 'count']
    for count in countList:
        runningSum += count
    studentIDMap[student_id] = runningSum
    return studentIDMap[student_id]

#Similar to the above function
categoryIDMap = dict()
def func2(category_id):
    if category_id in categoryIDMap:
        return categoryIDMap[category_id]
    runningSum = 0
    countList = countPerCategoryID.loc[category_id, 'count']
    for count in countList:
        runningSum += count
    categoryIDMap[category_id] = runningSum
    return categoryIDMap[category_id]

Finally I call these two function from below:

#Calculating rating category-wise
rating = []
for index, row in df.iterrows():

    totalCountPerCategoryID = func1(row['category_id'])
    totalCountPerStudentID = func2(row['student_id'])

    rating.append((totalCountPerStudentID / totalCountPerCategoryID) * 5)

df['rating'] = rating

Required output:

student_id   category_id  count   rating
1              111        10       2.38
2              111        5        1.19
3              222        8         5
4              333        5         5 
5              111        6        1.43

Since the data is huge, it is taking lot of time to run this. I would like to know how to optimize this calculation

Thanks in advance

bigdata pandas dataframe optimization python
2 Answers

You don't need to loop, this is a groupby case:

df['rating'] = df['count']/df.groupby('category_id')['count'].transform('sum') * 5

Output:

   student_id  category_id  count    rating
0           1          111     10  2.380952
1           2          111      5  1.190476
2           3          222      8  5.000000
3           4          333      5  5.000000
4           5          111      6  1.428571
user_42384080

Good god, don't use iterrows and append, even less so together. No wonder your performance is crawling. With pandas, iterrows should be a last resort.

You should be able to achieve this with the vectorized methods:

>>> df['rating'] = df['count'].div(df.groupby('category_id')['count'].transform(sum)).mul(5)
>>> df
   student_id  category_id  count    rating
0           1          111     10  2.380952
1           2          111      5  1.190476
2           3          222      8  5.000000
3           4          333      5  5.000000
4           5          111      6  1.428571
user_91833440
Related