Python: Optimizar uso de memoria con Pandas

Si estás trabajando con dataframes de Pandas en Python te propongo varias formas de optimizar el consumo de memoria RAM. Para nuestro ejemplo vamos a usar el siguiente dataset:

import pandas as pd
data = [
        {'name' : 'John Connor', 'world' : 'Earth', 'age' : 16, 'survive' : 1},
        {'name' : 'Max Rockatansky', 'world' : 'Earth', 'age' : 31, 'survive' : 1},
        {'name' : 'Ender', 'world' : 'Albion' , 'age' : 6, 'survive' : 1},
        {'name' : 'Anakin Skywalker', 'world' : 'Tatooine', 'age' : 8, 'survive' : 0},
        {'name' : 'Ellen Ripley', 'world' : 'Earth', 'age' : 37, 'survive' : 0},
        {'name' : 'Willow Ufgood', 'world' : 'Earth', 'age' : 25, 'survive' : 1}
]
df = pd.DataFrame(data)
df.dtypes

En primer lugar, para analizar el consumo de memoria de cada columna de un dataframe usamos el método de Pandas memory_usage. Vamos a usar dos parámetros con este método:

• deep = True : estima el uso de memoria más preciso a nivel de fila y tipo de dato.
• index = False : si no indicamos nada, por defecto nos indica el consumo de RAM del índice del dataframe además de cada columna.

Como el resultado lo representa en bytes, podemos dividir dos veces por 1024 para tener el dato en MB.

df.memory_usage(deep = True, index = False)

Nos devuelve el peso en bytes del índice y cada una de las columnas.

Index      128
name       414
world      376
age         48
survive     48
dtype: int64

Si comprobamos el tipo de datos inferido del dataset:

df.dtypes

Vemos que no es lo más óptimo. Las columnas de tipo string las ha tipado como OBJECT y las numéricas INT64, veamos qué podemos hacer:

name       object
world      object
age         int64
survive     int64
dtype: object

Filtro de datos categóricos (CATEGORY)

Las columnas de datos categóricos de tipo STRING o DATE podemos convertirlos en CATEGORY. Es importante tener en cuenta la cardinalidad de los datos (cantidad de valores distintos). Vamos a conseguir reducir el consumo de memoria siempre que tengan baja o media cardinalidad, si convertimos a category una columna con una cardinalidad muy alta probablemente necesite más memoria que si no lo hiciéramos. Lo que hace el tipo CATEGORY es crear un diccionario de todos los valores distintos de una columna, sustituyéndolos por punteros al diccionario. Vamos a probar con nuestro ejemplo, primero vamos a observar el consumo de memoria de las columnas name y world según están definidas como tipo OBJECT con df.memory_usage(deep = True, index = False):

name       414
world      376
dtype: int64

Aplicamos la optimización cambiando el tipo a CATEGORY:

df["name"] = df["name"].astype("category")
df["world"] = df["world"].astype("category")

Y volvemos a observar el uso de memoria:

name       592
world      304
dtype: int64

La columna world logra reducir el peso un 19% (de 375 bytes a 304), sin embargo ¿qué ha pasado con name? ¡ha aumentado el consumo de memoria! este comportamiento se debe a que la cardinalidad de la columna es muy alta (un dataset con 5 registros y 5 valores distintos), por lo que al generar el diccionario de la categoría necesita más memoria.

Optimizar el tipado de los datos

En este paso vamos a cambiar el tipo de cada columna para intentar ahorrar costes. Las columnas de texto a STRING y las numéricas que por defecto asigna como INT64, podríamos convertirlas a INT8, INT16 o INT32. Siguiendo el ejemplo del post:

name       object
world      object
age         int64
survive     int64
dtype: object

Vamos a analizar cómo cambia el consumo de memoria al convertir las columnas OBJECT a STRING y las numéricas a INT8 (nos vale un entero de 8 bits porque abarca valores de -128 a 127). Primero observamos el consumo de memoria antes de hacer la conversión:

name       414
world      376
age         48
survive     48
dtype: int64

Y aplicamos las conversiones:

df["name"] = df["name"].astype("string")
df["world"] = df["world"].astype("string")
df["age"] = df["age"].astype("int8")
df["survive"] = df["survive"].astype("int8")

¿Cómo cambia?

name       414
world      376
age          6
survive      6
dtype: int64

Observamos que las columnas numéricas age y survive han reducido su peso 87,5% (de 48 bytes a 8). En las columnas de texto no apreciamos cambios. Lo ideal es asignar el tipado al crear el dataframe, pero depende de lo que estemos usando como origen. Por ejemplo, al cargar un fichero Parquet se carga el esquema definido en el propio fichero. Si se trata de un CSV o XML que cargamos con los métodos de Pandas read_xml y read_csv podemos especificar un tipo común a todas las columnas (dtype="string") o bien pasarle un diccionario con el tipo de cada columna (dtype={"col1" : "string"}) :

# Asignamos el tipo String a todas las columnas
df = pd.DataFrame(data, dtype = "string")

# Especificamos el tipo por columna gracias a un diccionario
dictTypes = {"name":"string", "world":"string", "age":"int8", "survive":"int8"}
df = pd.DataFrame(data, dtype = dictTypes )

Y observamos los tipos con dtypes:

name       string
world      string
age          int8
survive      int8
dtype: object

Reducir número de columnas que cargamos en un dataframe

Es la optimización más sencilla y lógica, la forma más fácil de reducir el consumo de recursos es utilizar sólo los datos que necesitamos. A la hora de cargar un dataframe seleccionamos sólo las columnas necesarias. Para este caso vamos a imaginar que cargamos el dataset desde un CSV con READ_CSV, podemos pasarle en el parámetro USECOLS un listado de las columna que deseamos cargar.

import pandas as pd

fields = ['name','world']
df = pd.read_csv('dataset.csv', usecols = fields)
print(df)