Memória: os tipos certos

Onde a memória vai

Cada coluna ocupa linhas × tamanho do tipo. Um inteiro de 64 bits ocupa 8 bytes mesmo quando só guarda valores de 1 a 10. Uma coluna de texto repete a mesma palavra milhares de vezes. Medir com deep=True mostra o custo real, inclusive o conteúdo dos textos:

import numpy as np
import pandas as pd

pd.set_option("display.width", 170)
pd.set_option("display.max_columns", 20)

pedidos = pd.read_csv("dados/pedidos.csv", parse_dates=["data_pedido"])
grande = pd.concat([pedidos] * 250, ignore_index=True)
grande["id_pedido"] = np.arange(len(grande))
print(len(grande))
print((grande.memory_usage(deep=True) / 1024 / 1024).round(2).to_dict())

200000
{'Index': 0.0, 'id_pedido': 1.53, 'id_cliente': 1.53, 'id_produto': 1.53, 'quantidade': 1.53, 'data_pedido': 1.53, 'desconto': 1.53, 'canal': 2.25}

Reduzir com cuidado

Uma função de redução aplica três ideias: inteiros menores quando os valores cabem, category para textos com poucos valores distintos, e nada nos decimais que guardam dinheiro (reduzir float64 para float32 perde precisão, e em valores monetários isso é um erro):

def reduzir_memoria(tabela: pd.DataFrame) -> pd.DataFrame:
    nova = tabela.copy()
    for coluna in nova.select_dtypes("integer").columns:
        nova[coluna] = pd.to_numeric(nova[coluna], downcast="integer")
    for coluna in nova.select_dtypes(include=["object", "str"]).columns:
        if nova[coluna].nunique() / len(nova) < 0.5:
            nova[coluna] = nova[coluna].astype("category")
    return nova


menor = reduzir_memoria(grande)
antes = grande.memory_usage(deep=True).sum()
depois = menor.memory_usage(deep=True).sum()
print(menor.dtypes.astype(str).to_dict())
print("usa menos da metade da memória:", depois < 0.5 * antes)
restaurada = menor.astype(grande.dtypes.to_dict())
pd.testing.assert_frame_equal(grande, restaurada)
print("os valores são os mesmos")

{'id_pedido': 'int32', 'id_cliente': 'int8', 'id_produto': 'int8', 'quantidade': 'int8', 'data_pedido': 'datetime64[us]', 'desconto': 'float64', 'canal': 'category'}
usa menos da metade da memória: True
os valores são os mesmos

Para provar que nenhum valor mudou, eu converto a tabela reduzida de volta aos tipos originais e comparo com assert_frame_equal. Não basta passar check_dtype=False: o assert_frame_equal ainda compara o tipo do contêiner, e uma coluna category contra uma de texto falha, mesmo com os mesmos valores. O id_pedido caiu para int32, o id_cliente, o id_produto e a quantidade para int8, e o canal virou category.

Cuidado com a redução automática

O downcast escolhe o menor tipo que cabe nos valores de agora. Se amanhã chegar um id_cliente maior que o limite, uma conta com ele pode estourar (no NumPy, o inteiro dá a volta sem avisar, como vimos no capítulo 7 do curso de NumPy). Eu reduzo os tipos no fim da preparação, para guardar e analisar, e não no meio de um fluxo que ainda recebe dados novos.

Ler já com os tipos certos

O melhor lugar para economizar é na leitura: com usecols você lê só as colunas que vai usar, e com dtype você dá o tipo certo desde o início, sem passar antes pelo caro:

from pathlib import Path

Path("saida").mkdir(exist_ok=True)
grande.to_csv("saida/pedidos_grande.csv", index=False)

padrao = pd.read_csv("saida/pedidos_grande.csv", usecols=["id_cliente", "quantidade", "canal"])
enxuto = pd.read_csv(
    "saida/pedidos_grande.csv",
    usecols=["id_cliente", "quantidade", "canal"],
    dtype={"id_cliente": "int16", "quantidade": "int8", "canal": "category"},
)
print(padrao.dtypes.astype(str).to_dict())
print(enxuto.dtypes.astype(str).to_dict())
print("a leitura enxuta usa menos de um quinto:", enxuto.memory_usage(deep=True).sum() < padrao.memory_usage(deep=True).sum() / 5)

{'id_cliente': 'int64', 'quantidade': 'int64', 'canal': 'str'}
{'id_cliente': 'int16', 'quantidade': 'int8', 'canal': 'category'}
a leitura enxuta usa menos de um quinto: True

Tipos do Arrow

Para uma leitura que já devolve colunas no formato do Arrow (que representa os ausentes de forma nativa e comprime bem), existe o dtype_backend="pyarrow":

arrow = pd.read_csv("saida/pedidos_grande.csv", nrows=1000, dtype_backend="pyarrow")
print(arrow.dtypes.astype(str).to_dict())

{'id_pedido': 'int64[pyarrow]', 'id_cliente': 'int64[pyarrow]', 'id_produto': 'int64[pyarrow]', 'quantidade': 'int64[pyarrow]', 'data_pedido': 'string[pyarrow]', 'desconto': 'double[pyarrow]', 'canal': 'string[pyarrow]'}

Os tipos terminam em [pyarrow]. Eles são uma boa opção para dados com muitos ausentes e textos, mas algumas operações menos comuns ainda são mais rápidas nos tipos clássicos. Meça com os seus dados antes de adotar.