3 Simple Ways To Handle Bigger Datasets In Pandas
The types we choose to represent our data have an impact on the underlying size of its representation. There exists a multitude of data types, each with a different precisions, and memory consumption. To be as memory-efficient as possible you must understand your data, and choose the most fine-grained types possible to represent it.
Find the code on GitHub.
Concept
Pandas is amazing at processing data on a single machine, on a single core, and when it fits into memory. If you deal with big datasets, you can easily exceed these settings and might have to look into other tools to handle your data. Before moving to distributed computing, with technologies like Spark, there are some simple optimizations that will help you scale to bigger datasets.
The main limiting factor is the amount of memory needed to load the data into memory. If we reduce the size of the data, we can load more data, process it faster … Consider a simple CSV file with two columns: age and job.
When you load data into memory, you have to define the types for each column of
your dataset, and it is often done for you by default. Pandas use numpy arrays
to represent each column in your data frame. These numpy arrays have a dtype,
which is the underlying type of all elements in the array. When using
pd.read_csv, pandas will guess the data types of each column and assign numpy
arrays with default types. However the default types are core-grained, using
more space than needed to represent the different columns.
In our case we would get an age column with an int64 type and a job column
with an object type. This is sub-optimal since an age has a range [0, 100],
needing only 8 bits instead of 64. Also a job is not “any string” but is
categorical, we most likely have a limited list of possible job and each job is
an element of this list.
Use the proper dtype for numerical columns
Pandas does not infer the best type for our numerical columns. In our case we
could represent an age with a uint8 type, because an age is unsigned (always
positive) and its range is [0, 100]. However, by default, pandas use an int64
which consumes 8x more memory.
The first step is to enforce optimized types for numerical columns.
Use the proper dtype for categorical columns
When a column is categorical with a low cardinality (the number of unique values is lower than the total number of rows -> lots of repetitions) the string representation is the worst. The string is a costly type, depending on the format (utf-8, ascii, …) we use between 1 and 4 bytes per character. It quickly consumes a huge amount of memory when repeated over thousands of rows.
But there is a better way. It’s more efficient to have a dictionary mapping
every unique value to an index with a small type, for instance 1 byte. Then we
repeat only 1 byte thousands of times instead of the complete string which is
far bigger. Fortunately, pandas has a dtype to handle it automatically for you,
you just have to use the category dtype.
It usually brings the biggest improvement in memory consumption. Handling
categorical data with simple strings is a huge waste of resources. The str
type is heavy instead of a mapping with only indices repeated with a proper
categorical type.
In the calculation above, we assume that the smallest chunk of memory we can allocate is 1 byte (8 bits). Indeed, we only have two different values for the vocabulary, Data Scientist and Student. To represent two different values we only need 1 bit: either 0 or 1. But since we assumed that we can allocate 1 byte integer at minimum, we get 1 byte for each entry.
Load only relevant columns
It’s likely that you are not using all the columns of your dataset to conduct a specific analysis. Using the usecols keyword, you can load only the column of interests. It is a simple trick, but it helps you deal with huge datasets easily.
Pandas dtypes contain all numpy dtypes and specific pandas types. You can find the detailed list of all types: for pandas and for numpy.
In Practice
Let's get our hands dirty on a real dataset, to demonstrate the importance of data types. We will use a public dataset from Kaggle entitled: Resale Flat Prices in Sigapore (between 1990 and 1999).
First we import the needed libraries, and build the path to the dataset.
from pathlib import Path
import numpy as np
import pandas as pd
# Data source:
# <https://www.kaggle.com/sveneschlbeck/resale-flat-prices-in-singapore>
data_path = Path(__file__).parents[0] / "flat-prices.zip"
Then we define two helper functions:
We use df.memory_usage(deep=True) to get the memory size per column and sum
it to get the total size. The argument deep=True is important to get real
amount of memory consumed.
def to_mb(bytes: float) -> float:
"""Function to convert bytes to mega bytes (MB)"""
return bytes / 1024 ** 2
def analyze(df: pd.DataFrame, title: str) -> float:
"""Print the memory used by the dataframe in MB and return the
bytes"""
bytes = df.memory_usage(deep=True).sum()
mb = to_mb(bytes)
print(f"{title}: {mb:.2f} MB")
print(df.dtypes)
return mb
Then we load the data using pd.read_csv without any optimizations:
def no_optimization() -> pd.DataFrame:
"""Load the data without any optimization"""
df = pd.read_csv(data_path)
return df
...
if __name__ == "__main__":
# First we start with no optimizations
df_no_opti = no_optimization()
mb_no_opti = analyze(df_no_opti, "No optimizations")
No optimizations: 131.36 MB
month object
town object
flat_type object
block object
street_name object
storey_range object
floor_area_sqm float64
flat_model object
lease_commence_date int64
resale_price int64
dtype: object
The data frame takes 131.36 MB in memory. We also take a
look at column types and notice that some numerical types can be improved. Let's
take a look at the columns floor_area_sqm and resale_price:
floor_area_sqm resale_price
31.0 9000
31.0 6000
31.0 8000
31.0 6000
73.0 47200
... ...
142.0 456000
142.0 408000
146.0 469000
146.0 440000
145.0 484000
The floor_area_sqm can fit within a 16 bits float (with max value 6.55 ×
10e4) column and can be represented with a float16. resale_price can fit
into a 32 bits unisgned int (with max value 4.29e+09) and can be represented
with a unin32.
You can inspect the count, min, max and some other statistics for every column in your dataset using df.describe().
It is particularly useful to find the range of numerical columns, and the cardinality (count of unique values) of a categorical columns.
- Numerical columns with small range can be optimized with a smaller type
- The
categorytype for a low cardinality column provides a huge memory gain
Let's apply our numrical type optimizations. We use the dtypes argument in
pd.read_csv to specify a type per column.
def with_numerical_types() -> pd.DataFrame:
"""Load the data with proper numerical types"""
df = pd.read_csv(
data_path,
dtype={
"floor_area_sqm": np.float16,
"resale_price": np.uint32,
},
)
return df
The data frame size in memory is now 128.62 MB which is a slight reduction of 2.09%.
Now let's deal with the categorical columns. flat_model, flat_type,
storey_range, block and town are good candidates for categorical colunms.
They have a low cardinality which indicates a great memory improvement
potential.
Let's apply these optimizations:
def with_numerical_and_categorical_types() -> pd.DataFrame:
"""Load the data with proper numerical and categorical types"""
df = pd.read_csv(
data_path,
dtype={
"floor_area_sqm": np.float16,
"resale_price": np.uint32,
"flat_model": "category",
"flat_type": "category",
"storey_range": "category",
"block": "category",
"town": "category",
},
)
return df
The data frame size in memory is now 42.54 MB which is a huge reduction of 67.61% compared to the case without optimizations!
def with_numerical_and_categorical_types_and_without_unused_columns() -> pd.DataFrame:
"""Load the data with proper numerical and categorical types
and without unused columns"""
dtype = {
"floor_area_sqm": np.float16,
"resale_price": np.uint32,
"flat_model": "category",
"flat_type": "category",
"storey_range": "category",
"block": "category",
"town": "category",
}
df = pd.read_csv(data_path, dtype=dtype, usecols=list(dtype.keys()))
return df
The data frame size in memory is now 3.39 MB which is a huge reduction of 97.42% compared to the case without optimizations!
Sometimes all you need is to be cautious about the types of your data and clarity about the columns you really need.
from pathlib import Path
import numpy as np
import pandas as pd
# Data source:
# <https://www.kaggle.com/sveneschlbeck/resale-flat-prices-in-singapore>
data_path = Path(__file__).parents[0] / "flat-prices.zip"
def to_mb(bytes: float) -> float:
"""Function to convert bytes to mega bytes (MB)"""
return bytes / 1024 ** 2
def no_optimization() -> pd.DataFrame:
"""Load the data without any optimization"""
df = pd.read_csv(data_path)
return df
def with_numerical_types() -> pd.DataFrame:
"""Load the data with proper numerical types"""
df = pd.read_csv(
data_path,
dtype={
"floor_area_sqm": np.float16,
"resale_price": np.uint32,
},
)
return df
def with_numerical_and_categorical_types() -> pd.DataFrame:
"""Load the data with proper numerical and categorical types"""
df = pd.read_csv(
data_path,
dtype={
"floor_area_sqm": np.float16,
"resale_price": np.uint32,
"flat_model": "category",
"flat_type": "category",
"storey_range": "category",
"block": "category",
"town": "category",
},
)
return df
def with_numerical_and_categorical_types_and_without_unused_columns() -> pd.DataFrame:
"""Load the data with proper numerical and categorical types
and without unused columns"""
dtype = {
"floor_area_sqm": np.float16,
"resale_price": np.uint32,
"flat_model": "category",
"flat_type": "category",
"storey_range": "category",
"block": "category",
"town": "category",
}
df = pd.read_csv(data_path, dtype=dtype, usecols=list(dtype.keys()))
return df
def analyze(df: pd.DataFrame, title: str) -> float:
"""Print the memory used by the dataframe in MB and return the
bytes"""
bytes = df.memory_usage(deep=True).sum()
mb = to_mb(bytes)
print(f"{title}: {mb:.2f} MB")
print(df.dtypes)
return mb
if __name__ == "__main__":
# First we start with no optimizations
df_no_opti = no_optimization()
mb_no_opti = analyze(df_no_opti, "No optimizations")
print()
# Then we use proper numerical types
df_proper_types = with_numerical_types()
mb_proper_types = analyze(df_proper_types, "Proper numerical types")
reduction = (mb_no_opti - mb_proper_types) / mb_no_opti * 100
print(f"Size reduced by: {reduction:.2f}%")
print()
# Then we use proper numerical and categorical types
df_categorical = with_numerical_and_categorical_types()
mb_categorical = analyze(df_categorical, "Proper numerical and categorical types")
reduction = (mb_no_opti - mb_categorical) / mb_no_opti * 100
print(f"Size reduced by: {reduction:.2f}%")
print()
# No need to load unused columns. We often use a subset of the
# columns, releveant to our analysis.
# We assume that we are only interessed in the columns where we
# defined the types (numerical and categrical)
df_without_unused_columns = (
with_numerical_and_categorical_types_and_without_unused_columns()
)
mb_without_unused_columns = analyze(
df_without_unused_columns,
"Proper numerical and categorical types, without unused columns",
)
reduction = (mb_no_opti - mb_without_unused_columns) / mb_no_opti * 100
print(f"Size reduced by: {reduction:.2f}%")
