How to make your Pandas operation 100x faster | by Yifei Huang | Towards Data Science

Pandas is a great tool for exploring and working with data. As such, it is deliberately optimized for versatility and ease of use, instead of performance. There are often many different ways to do the exact same operation, some of which are far more performant than others. This is not a problem if your data is relatively small or your transformation is relatively simple, but it can quickly becomes a bottleneck as the scale and complexity increase. In this article, I will walk through a few different approaches to optimizing the performance ranging from vectorization to parallelization.
Let’s start by making this issue more concrete with an example.
A common operation when using Pandas is to create a new derived value based on one or more existing columns. The examples below (inspired by [1]) illustrates the various different ways of doing this using Pandas, and measure the resulting performance using the%timeit magic command in Jupyter notebooks.
The first approach
[sum_square(row[0], row[1]) for _, row in df.iterrows()]
uses list comprehension along with the method iterrows , and is the slowest by a long shot. This is because it is effectively using a simple for loop and incurring the heavy overhead of using the pandas series object in each iteration. It is rarely necessary to use series object in your transformation function, so you should almost never use the iterrows method.
The second approach
[sum_square(a, b) for a, b in df[[0, 1]].itertuples(index=False)]
replaces the iterrows method of the first approach with itertuples method. The improvement in performance is a whopping 60x (575ms -> 9.53ms). This is because the itertuples method bypasses the overhead associated with the Pandas series object and uses the simple tuples instead. If you need to loop through a Pandas dataframe, you should almost always itertuples instead of the iterrows
The third approach
df.apply(lambda row: sum_square(row[0], row[1]), axis=1 )
uses the apply method instead of list comprehension to apply the function to each row of the dataframe. It still incurs the overhead of the Pandas series object when working with each row, but is 6x faster than the first approach because it is able to leverage some more efficient pre-compiled looping methods underneath the hood. It is, however, still effectively a for loop, just more efficient way to do so.
The fourth approach
df.apply(lambda row: sum_square(row[0], row[1]), raw=True, axis=1 )
is able to achieve a 4x speed up relative to the third approach, with a very simple parameter tweak in adding raw=True . This is telling the apply method to bypass the overhead associated with the Pandas series object and use simple map objects instead. It is interesting note that this is still slower than the second approach with itertuples
The fifth approach
np.vectorize(sum_square)(df[0], df[1])
uses np.vectorize to attempt to make the function provided more “vectorized”. This basically means it tries to leverage pre-compiled and native methods that are SIMD (Single Instruction Multiple Data) and significantly faster [2], as much as possible. Depending on the function, it may or may not actually yield a speed up, but it never hurts to try given the relative simplicity. In this particularly example, it worked quite well and yielded a 171x speed up relative to the first approach. It is even 3x faster than the itertuple method in the second approach.
The last approach
np.power(df[0] + df[1], 2)
fully takes advantage of the vectorized native Numpy methods and is 1888x faster than the first approach we started with. Just like the np.vectorize method, it may not always be possible (or clear as to how) to vectorize the given function, but when you can you should definitely try. A good rule of thumb for assessing whether you can vectorize the function is if you can represent it as linear algebra operations on matrices and scalar values. If you can, you can likely find native Numpy methods to vectorize those operations.
What if your function cannot be vectorized or you don’t know how to?
There are transformations of the data that simply cannot be vectorized, and therefore cannot benefit from the dramatic improvements that vectorization provided above (but you should still never use iterrows). One example of this type of transformation is regex operations on strings like extracting the top level domain from URLs. In these cases, there are a few different approaches like compiling your function using Cython or run-time compilers like Numba, or using parallel processing. I prefer parallel processing because it requires the least amount of rewrite of your existing code. You simply have to add a few lines wrapper code to make it work. The example below illustrates how you can do this
import concurrent.futures
import numpy as np
import pandas as pd
from functools import partial
import time

df = pd.DataFrame(np.random.binomial(n=1000, p=0.2, size=(10000,2)))

# a no op function that simulates a transformation that takes some time to complete
def do_work(row, duration):
    time.sleep(duration)
    return True

FUNCTION_RUN_TIME = 0.0001

# simple use of apply to execute the function against the pd dataframe
def serial_calc(df, duration):
    apply_partial = partial(do_work, duration=duration)
    df['result'] = df.apply(apply_partial, axis=1)
    return df

%timeit serial_calc(df, FUNCTION_RUN_TIME)

# simple wrapper code around serial_calc to parallelize the work
def parallel_calc(df, func, n_core, duration):
    futs = []
    df_split = np.array_split(df, n_core)
#     pool = concurrent.futures.ThreadPoolExecutor(max_workers = n_core)
    pool = concurrent.futures.ProcessPoolExecutor(max_workers = n_core)
    apply_partial = partial(func, duration=duration)
    return pd.concat(pool.map(apply_partial, df_split))

%timeit parallel_calc(df, serial_calc, 32, FUNCTION_RUN_TIME)

1.42 s ± 9.58 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
184 ms ± 776 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

function_time = []
non_parallel_time = []
parallel_time = []

for d in range(11):
#     print(d)
    d_ms = d/20000
    function_time.append(d_ms)
    start = time.time()
    serial_calc(df, d_ms)
    end = time.time()
    non_parallel_time.append(end-start)
    
    start = time.time()
    parallel_calc(df, serial_calc, 32, d_ms)
    end = time.time()
    parallel_time.append(end-start)
    
import matplotlib.pyplot as plt

font = {'family' : 'Arial',
        'weight': 'normal',
        'size'   : 20}

plt.rc('font', **font)

fig, ax = plt.subplots(figsize=(14,8))
plt.xlabel('Function run time (ms)')
plt.ylabel('Overall processing time (s)')

x = [t*1000 for t in function_time]

ax.plot(x, non_parallel_time, '-', linewidth=1)
ax.plot(x, parallel_time, '-', linewidth=1)
ax.legend(['Serial Apply', 'Parallelized Apply'])

plt.show()

As you can see with relatively little effort and no modification to the existing code, we were able to achieve a 7x speed up to the code by leveraging parallel processing. It is worth noting that this improvement increases with data scale and function run time, as can be seen in the graph below, where I varied simulated function run time from 0 to 0.5ms. The serial implementation processing time scales at a much higher rate than parallel implementation.
A few important consideration to keep in mind when using parallel processing
Overhead — parallel processing incurs quite a bit of overhead, and as a result it is not guaranteed to be faster than serial processing if you function is relatively fast and your data is not super large. For the example above (10k rows), we can see that parallel processing overtakes serial processing, at around 0.01ms function run time.CPU bound vs I/O bound — If your function is CPU bound, meaning that the transformation logic just takes a long time and it is not waiting for data, then you should use multiprocess (or process pool). If your function is I/O bound, meaning that it is spending a lot of time waiting for data (e.g. making api requests over the internet), then multithreading (or thread pool) will be the better and faster option. This article [3] provides a really good discussion on differences between the two.
Overhead — parallel processing incurs quite a bit of overhead, and as a result it is not guaranteed to be faster than serial processing if you function is relatively fast and your data is not super large. For the example above (10k rows), we can see that parallel processing overtakes serial processing, at around 0.01ms function run time.
CPU bound vs I/O bound — If your function is CPU bound, meaning that the transformation logic just takes a long time and it is not waiting for data, then you should use multiprocess (or process pool). If your function is I/O bound, meaning that it is spending a lot of time waiting for data (e.g. making api requests over the internet), then multithreading (or thread pool) will be the better and faster option. This article [3] provides a really good discussion on differences between the two.
Whenever possible write your transformation functions using native pre-compiled SIMD methods (most Numpy and many Pandas native methods are)If you are not sure how to vectorize yourself, try np.vectorize, but your mileage may varyUse df.apply or df.itertuples for looping, never use df.iterrows unless you know exactly why you must use itIf your data is large and transformation is slow and cannot be vectorized easily, parallel processing is an easy way to boost performance
Whenever possible write your transformation functions using native pre-compiled SIMD methods (most Numpy and many Pandas native methods are)
If you are not sure how to vectorize yourself, try np.vectorize, but your mileage may vary
Use df.apply or df.itertuples for looping, never use df.iterrows unless you know exactly why you must use it
If your data is large and transformation is slow and cannot be vectorized easily, parallel processing is an easy way to boost performance
Hope this was a useful discussion. Feel free to reach out if you have comments or questions. Twitter | Linkedin | Medium