Category: Python

In this article you’ll find some tips to reduce the amount of RAM used when working with pandas, the fundamental Python library for data analysis and data manipulation.

When dealing with large(ish) datasets, reducing the memory usage is something you need to consider if you’re stretching to the limits of using a single machine. For example, when you try to load some data from a big CSV file, you want to avoid your program crashing with a MemoryError. These tips can also help speeding up some downstream analytical queries.

The overall strategy boils down to choosing the right data types and loading only what you need. In this article you’ll learn about:

• Finding out how much memory is used
• Saving memory using categories
• Saving memory using smaller number representations
• Saving memory using sparse data (when you have a lot of NaN)
• Choosing the right dtypes when loading the data
• Loading only the columns you need
• Loading only a subset of rows

Finding out how much memory is used

First, let’s look into some simple steps to observe how much memory is taken by a pandas DataFrame.

For the examples I’m using a dataset about Olympic history from Kaggle. The dataset is in CSV format and takes roughly 40Mb on disk.

>>> import pandas as pd
>>> athletes = pd.read_csv('athlete_events.csv')
>>> athletes.shape
(271116, 15)

There are ~271K records with 15 columns

For a breakdown of the memory usage, column by column, we can use memory_usage() on the whole DataFrame. The memory is reported in bytes:

>>> athletes.memory_usage(deep=True)
Index          128
ID         2168928
Name      20697535
Sex       15724728
Age        2168928
Height     2168928
Weight     2168928
Team      17734961
NOC       16266960
Games     18435888
Year       2168928
Season    17080308
City      17563109
Sport     18031019
Event     24146495
Medal      9882241
dtype: int64

The function also works for a single column:

>>> athletes['Name'].memory_usage(deep=True)
20697663

The difference between the two outputs is due to the memory taken by the index: when calling the function on the whole DataFrame, the Index has its own entry (128 bytes), while for a single column (i.e. a pandas Series) the memory used by the index is aggregated.

For an aggregated figure on the whole table, we can simply sum:

>>> athletes.memory_usage(deep=True).sum()
186408012  # roughly 178Mb

Why do we need deep=True? This flag will introspect the data deeply, reporting the actual system-level memory consumption. Without setting this flag, the function returns an estimate which could be quite far from the actual number, for example:

>>> athletes.memory_usage()
Index         128
ID        2168928
Name      2168928
Sex       2168928
Age       2168928
Height    2168928
Weight    2168928
Team      2168928
NOC       2168928
Games     2168928
Year      2168928
Season    2168928
City      2168928
Sport     2168928
Event     2168928
Medal     2168928
dtype: int64

>>> athletes['Name'].memory_usage()
2169056

Another way of getting the overall memory consumption is through the function info(), which is going to be useful because it also gives us information on the data types (dtype) used by the DataFrame. Notice again the use of deep introspection for the memory usage:

>>> athletes.info(memory_usage='deep')
<class 'pandas.core.frame.DataFrame'>
RangeIndex: 271116 entries, 0 to 271115
Data columns (total 15 columns):
 #   Column  Non-Null Count   Dtype  
---  ------  --------------   -----  
 0   ID      271116 non-null  int64  
 1   Name    271116 non-null  object 
 2   Sex     271116 non-null  object 
 3   Age     261642 non-null  float64
 4   Height  210945 non-null  float64
 5   Weight  208241 non-null  float64
 6   Team    271116 non-null  object 
 7   NOC     271116 non-null  object 
 8   Games   271116 non-null  object 
 9   Year    271116 non-null  int64  
 10  Season  271116 non-null  object 
 11  City    271116 non-null  object 
 12  Sport   271116 non-null  object 
 13  Event   271116 non-null  object 
 14  Medal   39783 non-null   object 
dtypes: float64(3), int64(2), object(10)
memory usage: 177.8 MB

Notice how all the string fields are loaded as object, while all the numerical fields use a 64-bit representation, due to the architecture of the local machine — it could be 32-bit with older hardware.

Saving memory using categories

Some of the variables in our dataset are categorical, meaning they only have a handful of possible values. Rather than using a generic object for these variables, when appropriate we can use the more relevant Categorical dtype in pandas. For example, good candidates for this data type include the variables Medal, Season, or Team, amongst others.

If you don’t have a full description of the data, in order to decide which columns should be treated as categorical, you can simply observe the number of unique values to confirm this is much smaller than the dataset size:

>>> athletes['Medal'].unique()
array([nan, 'Gold', 'Bronze', 'Silver'], dtype=object)

There are only three different values, plus the null value nan.

Observe the difference in memory consumption between using object and using categories:

>>> athletes['Medal'].memory_usage(deep=True)
9882369  # 9.4+ Mb
>>> athletes['Medal'].astype('category').memory_usage(deep=True)
271539  # 0.25 Mb

Besides saving memory, another advantage of using categorical data in pandas is that we can include a notion of logical order between the values, different from the lexical order.

Saving memory using smaller number representations.

Let’s look at some numerical variables, for example ID (int64), Height (float64) and Year (int64).

You can observe their range by checking the minimum and maximum values:

>>> athletes['ID'].min(), athletes['ID'].max()
(1, 135571)

The int64 dtype is able to hold numbers on a much broader range, at the price of a much bigger memory footprint:

>>> import numpy as np
>>> np.iinfo('int64')  # integer info
iinfo(min=-9223372036854775808, max=9223372036854775807, dtype=int64)

Using int32 for the column ID is enough to store its values and it will save us half of the memory space:

>>> athletes['ID'].memory_usage(deep=True)
2169056
>>> athletes['ID'].astype('int32').memory_usage(deep=True)
1084592

We can do the same with floats:

>>> athletes['Height'].min(), athletes['Height'].max()
(127.0, 226.0)

In this case, a float16 is enough, and costs a quarter of the memory price:

>>> athletes['Height'].memory_usage(deep=True)
2169056
>>> athletes['Height'].astype('float16').memory_usage(deep=True)
542360

Finally, let’s look at the variable Year:

>>> athletes['Year'].min(), athletes['Year'].max()
(1896, 2016)

In this case, it looks like an int16 would be enough. On a closer look though, we can consider the Year column to be categorical as well: there’s only a handful of possible values. Let’s check the difference:

>>> athletes['Year'].memory_usage(deep=True)
2169056
>>> athletes['Year'].astype('int16').memory_usage(deep=True)
542360
>>> athletes['Year'].astype('category').memory_usage(deep=True)
272596

For this particular situation, it makes more sense to use categories rather than numbers, unless we plan on performing arithmetic operations on this column (you cannot sum or multiply two categories).

Saving memory using sparse data (when you have a lot of NaN)

The sparse dtypes in pandas are useful wen dealing with columns that have a lot of null values. Depending on your variables, you may want to consider representing your data as sparse. The info() function used earlier tells us how many non-null records we have for each column, so if that number is much lower than the size of the dataset, it means we have a lot of null values.

This is exactly the case of the Medal column that we treated as categorical earlier:

>>> athletes['Medal'].memory_usage(deep=True)
9882369
>>> athletes['Medal'].astype('category').memory_usage(deep=True)
271539
>>> athletes['Medal'].astype('Sparse[category]').memory_usage(deep=True)
199067

Choosing the right dtypes when loading the data

So far we have looked at the memory usage of different dtypes, converting the columns after the dataset was loaded.

Once we have chosen the desired dtypes, we can make sure they are used when loading the data, by passing the schema as a dictionary to the read_csv() function:

>>> schema = {
... 'ID': 'int32',
... 'Height': 'float16',
... # add all your Column->dtype mappings
... }
>>> athletes = pd.read_csv('athlete_events.csv', dtype=schema)

Note: it’s not possible to use the Sparse dtype when loading the data in this way, we still need to convert the sparse columns after the dataset is loaded.

Loading only the columns you need

Depending on the application, we often don’t need the full set of columns in memory.

In our medal example, let’s say we simply want to compute the overall count of the medals per nation. For this specific use case, we only need to look at the columns Medal and NOC (National Olympic Committee).

We can pass the argument usecols to the read_csv() function:

>>> athletes = pd.read_csv('athlete_events.csv', usecols=['NOC', 'Medal'], dtype=schema)
>>> athletes.info(memory_usage='deep')
<class 'pandas.core.frame.DataFrame'>
RangeIndex: 271116 entries, 0 to 271115
Data columns (total 2 columns):
 #   Column  Non-Null Count   Dtype   
---  ------  --------------   -----   
 0   NOC     271116 non-null  category
 1   Medal   39783 non-null   category
dtypes: category(2)
memory usage: 816.3 KB

If we only want to count the medals per NOC, we can use a groupby operation on the above DataFrame:

>>> athletes.groupby('NOC')['Medal'].count()
# (long output omitted)

It’s also worth noting that for this groupby operation there is a significant speed-up when using the categorical data types.

Loading only a subset of rows

To complete the picture, the read_csv() function also offers options to limit the number of rows we’re loading. This can be useful in a few circumstances, for example when we just want to take a peek at the data without looking at the whole dataset, or when the dataset is big enough that we can answer our analytical questions with a sample.

The first case, just taking a peek at the data, is straightforward:

>>> athletes = pd.read_csv('athlete_events.csv', nrows=1000)
>>> len(athletes)
1000

Using the nrows argument, we’ll load the first N records (1000 in the example above) into the DataFrame. This can often be enough to have a first feeling for the data before digging into further analysis.

If we want to implement some random sampling instead, the read_csv() function also offers the skiprows argument. Specifically, if we pass a custom function to this argument, we can implement our sampling logic. The function takes one argument (the row number) and should return True if you want to skip that row.

In the example, we want to keep the first row because it has the column names, and we load only ~10% of the data, using the function random() which returns a random float in the [0, 1) range (if this number is greater than 0.1, we skip the row):

>>> from random import random
>>> def skip_record(row_number):
...     return random() > 0.1 and row_number > 0
... 
>>> athletes = pd.read_csv('athlete_events.csv', skiprows=skip_record)
>>> len(athletes)
27176  # 27K rows in the sample, ~271K in the full dataset

Summary

In this article we have discussed some options to save memory with pandas choosing the most appropriate data types and loading only the data that we need for our analysis.

Do you need to upskill your team in pandas? Marco runs public and private training courses on Effective pandas and other Python topics, please get in touch to know more.

Sign up to the newsletter Musings on Data to receive periodic updates and recommendations from Marco on Data Science.

Follow Marco on Twitter.