Hey Pandas, why you no fast loop?!?
By: Tanmay Chimurkar
Tags: Numpy, Vectorization, Data Science, Pandas, Python

Speeding up your loop by X times.

A beginner’s guide to faster looping for DataFrame objects.

An everyday Data Scientist starts his day with a coffee and 2 kind-of mandatory lines of code:

import pandas as pd

import numpy as np

While we may be far too familiar with pandas for its DataFrame structure, we sometime simply ignore many of the key features it has to offer, along-with the conspicuous yet powerful numpy. This post is a small attempt to help build a familiarisation with some of the power that comes in with pandas and numpy (Obviously limited to looping)

Let’s Begin with a simple example of bucketing(grouping) based on column values.

Dang! No data. Let’s go ahead and create some dummy data and a dummy condition for this case, generating integers as the age of a person, on which we will be bucketing in a new column, say ‘age_bucket’

Dummy data for 1 lakh members, done!

A little better. Now, let’s assume we have some governing rules across the 1 lack ages of the people in our data-set as follows:

“age-range” -> “bucket-label” 0–5 -> “Infant” 6–20 -> “Teenage” 21–40 -> “Early Adulthood” 41–80 -> “Middle Adulthood” ≥81 -> “Late Adulthood”

Alright, now that we have our conditions, let’s loop it!

Naive Looping

If we use a very basic and naive type of looping,

Naive Looping on DataFrame

and calculate the total computation time to bucket (%%timeit), we get the following results:

Computation Time from Naive Looping

(Please make sure not to run this code if you have a small machine, since it took 5hrs to compute run-time. If you want to test the compute time for this, I suggest you change the overall records in the initial dataframe to 1000 or smaller)

So, not very good for just 1 lakh age grouping, since 1 lakh is actually really minimal data considered for a good machine learning model build(in most cases). Naive looping in Pandas does not take advantage of any built-in optimisations (which we’ll see below), making it extremely slow. Let’s move ahead and look at one of the built-in functions of Pandas.

Leveraging from pd.df.iterrows()

pd.df.itterows() iterates over DataFrame objects as a pair of index and the value of the objects that is associated with that index. For example, if we use itterows() to check one object, we get the following output:

checking objects with itterows()

person_age 40

buckets

Name: 0, dtype: object

Using itterows() to loop over our DataFrame to label buckets,

Using itterows() to loop over DataFrame

and timing it, we get the following results:

Results when using iterows to loop over DataFrame

We just boosted our loop ~149 times as compared to the normal loop. Let’s see if we can further reduce computation times!

Using apply() for looping

The apply() function loops over the DataFrame in a specific axis, i.e., it can either loop over columns(axis=1) or loop over rows(axis=0). apply() is better than iterrows() since it uses C extensions for Python in Cython.

We can use apply with lambda function along-with the axis argument by first making a custom function,

Custom bucketing function for apply()

and then applying it as follows, giving us the following results:

Results from df.apply() with lambda function

We are now in microseconds, making out loop faster by ~1900 times the naive loop in time. So far, we’ve had a good improvement in improving computation times for our code using pandas Vectorization.

Pandas Vectorization

It’s pretty simple actually. Vectorization operates on complete DataFrame at once, not on individual rows. In essence, vectorization is a process of executing operations on entire arrays. In the case of pandas, vectorization is applied on complete columns, passed as a series to the function. Here is what the vectorization part for our case looks like in pandas:

Pandas vectorization, where input is a Series from the DataFrame

In the function pd_vector , age is a Series from the dataframe(df). In passing a series as an input, pandas uses a lot of the built-in optimisations available for vectorized functions, performing operations on the entire series simultaneously.

Timing the computation time for this vectorized form of the version, we get the following results:

Results with Pandas Vectorization

We can see that using vectorization for pandas, we get about ~39 times faster loop from our apply() method and ~82608 times faster than our standard loop

Let’s see if we can fasten it up even further!

Numpy Vectorization

While pandas use series objects for vectorization, we can simply tweak the series object from series to an array, making it even faster. It becomes faster by removing all the extra overheads like indexing, data type, data formatting, etc. However, numpy vectorization should be carefully used, i.e., if series index is of no utmost essence, then we can surely go ahead and use numpy vectorization to even further reduce computation times!

Converting a series object to an array is pretty simple if your case favours such conversion to an array, and if done, we get to see the following results:

Numpy Vectorized loop

So, starting now from our initial naive loop on DataFrames, which took a total of ~38 minutes for just 1 lakh conditional looping, we are now at 14.7 milliseconds, improving the speed by about ~155102 times.

We can certainly call it a day till now, but if speed is what you really desire, then we can still go even further using just “numpy”

Cutting Down the Middle-Man

So far, all of our loops were built and implemented by transforming our pandas rows or series to some form which would then use the built-in optimisations. Now, we are ready to cut down our Middle-Man (Pandas, not a person 🙊). This is where we totally rely on just one module that we imported from the very beginning.

But to use just numpy in looping, you have to get your hands dirty (again, just dirty to code 😅). Here is a quote I had once read that I hope inspires you too to write such code at times:

Don’t be afraid to write complicated code, since at times the best written code is not optimised too well!

That said, our numpy code is going to look pretty horrible, but trust me, it is worth the loss of all the aesthetics.

Using np.where() to loop

np.where() is sometimes case specific, but if you have an if-else loop, it can almost always be used as the best alternative for faster code. Here is a snippet from the np.where() documentation explaining how it works:

np.where doc

So basically, it just returns elements from either “x” or “y” depending on what condition we provide. Let’s go ahead and look at how our code will look using just np.where() function, used as a loop:

Numpy loop using np.where

Testing compute time for this np.where loop, we get the following results:

So, we are at 11.6 ms now, strating from 38 minutes, improving code speed by ~196551.

Even so, this is not the end. You can still use Cython custom functions, np.digitize(), and many other optimizations to get runtimes in micro-seconds!

So, in conclusion, the order generally preferred (atleast by me) to loop over DataFrames is as follows:

Vectorization {numpy} Vectorization {pandas} pd.df.apply() {using lambda function} pd.df.iterrows() Basic Looping {Mostly Never}

I would be really glad if you use the looping methods mentioned above on test cases you may encounter some day! 😸

Sources:

https://docs.scipy.org/doc/numpy/reference/

https://pandas.pydata.org/pandas-docs/stable/

https://stackoverflow.com/questions/1422149/what-is-vectorization

https://en.wikipedia.org/wiki/Array_programming

Further, If you feel there are any mistakes in the post, I would be more than happy to be rectified of them. Otherwise, if you liked this article, and you learned some good optimization methods from this post, then my work here is done!