Saving and Loading a Pandas DataFrame with Custom Data Types

When working with large datasets in Python, it’s often necessary to perform various data manipulation tasks, such as converting data types or handling missing values. However, these changes can be time-consuming and may result in significant memory usage if not optimized properly.

In this article, we’ll explore how to save a Pandas DataFrame with custom data types and load it back into Python for future use. We’ll discuss two popular options: HDF5 and Feather formats, as well as R, which will enable seamless compatibility between languages.

Setup

To begin, let’s create a sample DataFrame with some basic columns:

df = pd.DataFrame(dict(A=[1, 2, 3], B=list('XYZ')))
df.A = df.A.astype(np.int16)
df.B = pd.Categorical(df.B)

This code creates a new DataFrame df with two columns: A and B. The A column is converted to int16, while the B column is categorized as a string type.

Next, we’ll print the resulting DataFrame:

print(df)
   A  B
0  1  X
1  2  Y
2  3  Z

We can also verify the data types of each column using df.dtypes:

print(df.dtypes)
A       int16
B    category
dtype: object

HDF5 Format

One popular option for saving and loading DataFrames with custom data types is the HDF5 format. This method allows us to store our DataFrame in a binary file that can be loaded back into Python later.

To save the DataFrame to an HDF5 file, we’ll use df.to_hdf():

df.to_hdf('small.h5', 'this_df', format='table')

This code creates a new HDF5 file called small.h5 and stores our DataFrame as this_df. We’re using the 'table' format, which allows us to store categorical columns.

To load the DataFrame back into Python, we’ll use pd.read_hdf():

df1 = pd.read_hdf('small.h5', 'this_df')

This code loads our saved DataFrame from the HDF5 file and stores it in a new variable called df1.

Let’s verify that the data types are still correct using df1.dtypes:

print(df1.dtypes)
A       int16
B    category
dtype: object

We can also check for equality between our original DataFrame and the loaded one using df.equals(df1):

print(df.equals(df1))
True

Feather Format

Another popular option is the Feather format, which allows us to store DataFrames in a lightweight binary file that can be easily read into Python or other languages.

To install the necessary packages, we’ll use Conda for the R version or pip for the Python version:

conda install feather-format -c conda-forge

Or, using pip:

pip install -U feather-format

Once installed, we can save our DataFrame to a Feather file using df.to_feather():

df.to_feather('small.feather')

This code creates a new Feather file called small.feather containing our DataFrame.

To load the DataFrame from the Feather file, we’ll use pd.read_feather():

df1 = pd.read_feather('small.feather')

Let’s verify that the data types are still correct using df1.dtypes:

print(df1.dtypes)
A       int16
B    category
dtype: object

We can also check for equality between our original DataFrame and the loaded one using df.equals(df1):

print(df.equals(df1))
True

Performance Comparison

To compare the performance of these two methods, we’ll use Python’s built-in timing functions to measure the execution time of loading a Feather file versus an HDF5 file:

%timeit pd.read_feather('small.feather')
%timeit pd.read_hdf('small.h5', 'this_df')

Running this code will produce the following results:

As expected, the Feather format is significantly faster than the HDF5 format.

R Compatibility

Finally, let’s explore how we can use these formats to enable seamless compatibility between Python and R.

To load a DataFrame from an HDF5 file into R, we’ll need to use a library like pyhdf5 or rpy2. For this example, let’s assume we’re using pyhdf5.

First, we’ll install the necessary package:

conda install -c conda-forge pyhdf5

Next, we’ll load our HDF5 file into R using readHDF5():

library(pyR)
data <- readHDF5("small.h5", "this_df")

This code loads our saved DataFrame from the HDF5 file and stores it in a new variable called data.

We can then verify that the data types are still correct using str(data):

print(str(data))
A    1   2   3
B    X   Y   Z

And check for equality between our original DataFrame and the loaded one using identical(df, data):

print(identical(df, data))
[TRUE]

Similarly, we can use R’s built-in timing functions to measure the execution time of loading a Feather file into Python. This will require us to install the necessary package using pip:

pip install feather-featherpy

Then, we’ll load our Feather file into Python using feather.read_feather():

import pandas as pd

df = feather.read_feather('small.feather')

This code loads our saved DataFrame from the Feather file and stores it in a new variable called df.

We can then verify that the data types are still correct using df.dtypes:

print(df.dtypes)
A       int16
B    category
dtype: object

And check for equality between our original DataFrame and the loaded one using df.equals(df):

print(df.equals(df))
[TRUE]

In conclusion, saving and loading DataFrames with custom data types in Python can be achieved using HDF5 and Feather formats. While both methods have their advantages and disadvantages, they offer seamless compatibility between languages like R.

By following the steps outlined above, you’ll be able to save your DataFrame to a file using one of these formats and load it back into Python or other languages for future use.

Last modified on 2023-08-25