Optimize Python Code in Jupyter Notebook | by Winson Waisakurnia | Towards Data Science

Want to make a slow Python code fast in Jupyter Notebook? By using sum aggregate as an example, we will measure the runtime with %%time and find the bottleneck using %lprun line profiler.
Assume that Pandas doesn’t have group-by and aggregate functions and we need to implement it ourselves. We want to aggregate the table above by name and calculate the sum of value.
Perhaps, in our first iteration, we come out with code like below. We aggregate the value with the help of a dictionary (name ⇨ value). By iterating each row, we update the value of the dictionary. At the end of the iteration, we convert the dictionary back into Pandas’ DataFrame.
Before we optimize the code, we need to ask: “Do we really need to optimize the code? Why?” It is not worth it if we take more time to optimize compared to waiting a little bit longer for the job to finish, isn’t it?
“Premature optimization is the root of all evil.” — Donal Knuth
For now, let’s assume that the sum aggregate code is too slow and we need to optimize it. How can we do that?
“If you can’t measure it, you can’t improve it.” — Peter Drucker
We need the ability to measure code runtime to compare the before and after optimization. There is a chance that we make things worse especially if we use an external library like Numpy and Pandas. We don’t know how the library implemented.
Code with %% prefix is IPython cell magic-function. We can measure the runtime of a cell by adding %%time to the cell. We don’t need to install any new library for that. We just need to import time which is part of the Python standard library.
%%timesum_aggregate_1(df, 'name', 'value').sort_values('value', ascending=False).head()
The measured runtime is about 38 s. Now, which part do we need to optimize? If I am the one who is taking a guess, I will choose to optimize the time.sleep() part. It seems obvious. Optimizing based on gut feeling could end up with no significant improvement and just wasting time. We need to find the bottleneck.
“Any improvements made anywhere besides the bottleneck are an illusion.” — Gene Kim
Optimizing part of the code that is not the bottleneck will not reduce the total runtime significantly. That’s why Gene Kim said that the improvement is just an illusion and useless. How can we find the bottleneck?
We can use line-profiler library to get the execution time for each line. We need to install line-profiler. After that, we need to import the magic function with %load_ext.
! pip install line-profiler%load_ext line_profiler
We need to provide function to profile and execution trigger arguments to the line profiler. In our case, the function we want to profile is sum_aggregate_1 and we trigger the execution with sum_aggregate_1(random_df, 'name', 'value).
# %lprun -f [function to profile] [execution trigger]# [function to profile]: sum_aggregate_1# [execution trigger]: sum_aggregate_1(random_df, 'name', 'value')%lprun -f sum_aggregate_1 sum_aggregate_1(random_df, 'name', 'value')
From the result, we can see that the main bottleneck of the function is at line #4. Now, we know that DataFrame.iterrows() is very slow for iteration. It took 66% of execution time.
After we found out the bottleneck, we could proceed with optimizing the code. We could just simply find alternatives for DataFrame.iterrows() from a search engine.
One of the alternatives is to get the key and value columns as Numpy arrays. After that, zip both Numpy arrays and iterate over it. We could observe that the bottleneck is now moving from the iteration part to time.sleep() part. The total runtime decreased from 38 s to 7 s.
By repeatedly finding and optimizing the bottleneck, we do the most impactful optimization first. It is like using a “greedy algorithm” mindset to prioritize optimization. We can stop optimizing when the code runtime is fast enough, although it is tempting to do further optimization.
Three things that are very important to remember when doing optimization:
Do optimization when you really need itDon’t optimize the code if it doesn’t worth your time because instead of making it fast, there is a chance you will introduce new bugs.We can’t improve what we can’t measureFirst, we must choose which metric we want to optimize. After that, we need to find how to measure it. Currently, we tried to optimize runtime. We can also choose to optimize memory usage in other cases.Optimize the bottleneckOptimizing part of the code that is not the bottleneck is useless because it will not reduce the runtime significantly. We can use line-profiler library to find the bottleneck.
Do optimization when you really need itDon’t optimize the code if it doesn’t worth your time because instead of making it fast, there is a chance you will introduce new bugs.
We can’t improve what we can’t measureFirst, we must choose which metric we want to optimize. After that, we need to find how to measure it. Currently, we tried to optimize runtime. We can also choose to optimize memory usage in other cases.
Optimize the bottleneckOptimizing part of the code that is not the bottleneck is useless because it will not reduce the runtime significantly. We can use line-profiler library to find the bottleneck.
You can get and run all of the code above from my Google Colab notebook. From the notebook, you can see the examples of how we could profile a class method. There is an example of how you could profile another function called by the triggering function as well.
Thanks for reading. Hopefully, this could help you optimizing your code when it is too slow. Please leave a comment or note if you have any suggestions or feedback.