Otimização de Arquivos Grandes com Pandas

Eu estava lendo algumas notícias no TabNews e achei bem interessante esse post sobre otimização de arquivos com Python: Como otimizei um Arquivo de 1.6 GB para 9 MB usando Python.

Eu entendi a ideia, mas eu queria tentar simular a ideia aplicada pelo autor e verificar a solução proposta por ele.

Todo o código de teste está disponível nesse Notebook.

Construção dos Dataset (CSV)

Para construir o dataset primeiro separei as variáveis comentadas pelo autor. De acordo com ele um trecho dele:

"Após a criação da coluna com as horas, procedi com a agregação da quantidade de passagens por data, hora, tipo de veículo e sentido:"

Desse trecho eu sei que o que ele tinha antes eram dados temporais (coluna com data e hora), tipo de veículo, sentido e quantidade de passagens. Nesse caso, basta eu criar um exemplo tenha todas essas variáveis.

Antes de tudo eu criei algumas variáveis que definem os dados, como tipos de veículo (vehicle_types), direções (directions), arquivos de entrada e saída e quandidade de dados (n).

big_csv_file_name = 'big_out.csv'
small_csv_file_name = 'small_out.csv'
vehicle_types = ['Carro', 'Moto', 'Caminhão', 'Ônibus']
directions = ['Norte', 'Sul', 'Leste', 'Oeste']
n = 22000000

Os códigos abaixo são apenas para construção da base. Nesse caso foi considerado um ano entre o começo e o fim e valores aleatórios de quantidade, entre 1 e 30.

from datetime import datetime, timedelta
import pandas as pd
import numpy as np

start_date = datetime(2023, 1, 1)
end_date = datetime(2023, 12, 31)
date_range = pd.date_range(start_date, end_date, freq='T')

dates = np.random.choice(date_range, size=n)
vehicles = np.random.choice(vehicle_types, size=n)
directions_data = np.random.choice(directions, size=n)
quantity = np.random.randint(1, 30, size=n)

df = pd.DataFrame({
  'datetime': dates,
  'vehicle_type': vehicles,
  'direction': directions_data,
  'quantity': quantity
})

Ao fim da criação do dataframe é necessário transformar em csv para simular o arquivo de input de dados.

df.to_csv(big_csv_file_name, index=False)

Nessa parte eu só vejo o tamanho do arquivo de entrada antes da agragação.

import os

file_size_big = os.path.getsize(big_csv_file_name)

print(f'O arquivo {big_csv_file_name} tem um tamanho de {file_size_big / (1024 * 1024):.2f} MB.')

Aplicação do Agrupamento

Uma vez tenho criado o arquivo de entrada, agora é necessário ler o arquivo de entrada.

import pandas as pd

csv_df = pd.read_csv(big_csv_file_name)

Analisando as informações do dataframe eu notei que da coluna datetime estava como object.

csv_df.info()

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 22000000 entries, 0 to 21999999
Data columns (total 4 columns):
 #   Column        Dtype 
---  ------        ----- 
 0   datetime      object
 1   vehicle_type  object
 2   direction     object
 3   quantity      int64 
dtypes: int64(1), object(3)
memory usage: 671.4+ MB

Isso não é bom de trabalhar, então transformei ela em datetime. Alterei os outros tipos também.

csv_df['datetime'] = pd.to_datetime(csv_df['datetime'])
csv_df['vehicle_type'] = csv_df['vehicle_type'].astype('category')
csv_df['direction'] = csv_df['direction'].astype('category')

Tendo transformado a coluna datetime em datetime, então fica fácil ter a data e a hora.

csv_df['hour'] = csv_df['datetime'].dt.hour
csv_df['date'] = csv_df['datetime'].dt.date

Depois é só agregar. Eu tentei agragar de acordo com o que o autor colocou.

aggregated_df = csv_df.groupby(['date', 'hour', 'vehicle_type', 'direction', 'quantity']).size().reset_index(name='count')

Depois de agragado eu transformei para csv para ter o arquivo de saída reduzido.

aggregated_df.to_csv(small_csv_file_name, index=False)

Aqui eu só faço a mesma coisa que fiz no arquivo de entrada, vendo o tamanho dele.

import os

file_size_small = os.path.getsize(small_csv_file_name)

print(f'O arquivo {big_csv_file_name} tem um tamanho de {file_size_small / (1024 * 1024):.2f} MB.')

Comparando os resultados, dá para notar que houve uma redução significativa. Para 22.000.000 de dados, eu tive uma redução de 84% do arquivo original. Acredito que a media que a base aumenta de tamanho esse resultado deve diminuir, mas já ajuda.

percentage = (file_size_small / file_size_big) * 100
print(f'O arquivo {small_csv_file_name} é {percentage:.2f}% do tamanho do arquivo {big_csv_file_name}.')

Conclusão

Deu para ver que a recomendação do agrupamento foi muito útil e reduziu para um valor bem melhor. Legal que essa estratégia parece ser muito interessante de aplicar para séries temporais, uma vez que eles, geralmente, são dados com muitos ruídos. Outro ponto é que talvez os meus valores não tenha abaixado mais o tamanho do arquivo final devido ao meu range de quantidade de 1 até 30. Isso piora o agrupamento, mas para valores menores de quantidades isso diminui ainda mais o arquivo final.