Almacenar archivos CSV 10 veces más rápido en Python con PyArrow

Hace un tiempo publiqué una entrada en la que comparaba las ventajas de usar el formato de archivo Apache Arrow frente al tradicional CSV. Suponiendo grandes ahorros tanto en el tamaño de los archivos como el tiempo de carga. Usando para ello la librería PyArrow. Es posible que por temas de compatibilidad entre sistemas en algunas situaciones no podamos usar Apache Arrow, aunque hemos visto que este formato se puede leer y escribir tanto en R como en Julia, por lo que sea necesario seguir usando archivos CSV. En estos puede ser una buena idea recurrir a PyArrow en lugar de Pandas para escribir y leer archivos CSV. Ya que, como veremos a continuación, PyArrow almacena los archivos CSV 10 veces más rápido en Python que Pandas.

Escribir un archivo CSV con PyArrow

En primer lugar, es necesario saber cómo se puede guardar un archivo CSV con PyArrow. Para lo que es necesario importar tanto la librería PyArrow como el complemento para trabajar con CSV. Una vez hecho esto, si lo que tenemos en memoria es un objeto DataFrame es necesario convertirlo antes en un tabla mediante la función Table.from_pandas() de PyArrow. Una vez hecho esto se puede volcar la tabla a una archivos mediante la función csv.write_csv(), indicando en el primer parámetro la tabla y en el segundo el nombre del archivo. Lo que se muestra en el siguiente pedazo de código.

import pyarrow as pa
import pyarrow.csv

tbl = pa.Table.from_pandas(df)
pa.csv.write_csv(tbl, 'pyarrow.csv')

Asumiendo que el DataFrame que se desea guardar el el archivo se encuentra en la variable df.

Leer un archivo CSV con PyArrow

Para leer archivos CSV con PyArrow se tiene que recurrir a la función csv.read_csv(), una función que solamente requiere el nombre del archivo. Como resultado se obtendrá un objeto Table, por lo que si se desea trabajar con un DataFrame de Pandas será necesario convertir el resultado con el método to_pandas() del que disponen los objetos Table. Así para leer el archivo se puede hacer con el siguiente código.

import pyarrow as pa
import pyarrow.csv

csv = pa.csv.read_csv('pyarrow.csv').to_pandas()

Comparar el rendimiento de PyArrow y Pandas leyendo archivos CSV

La mejor manera de comparar el tiempo que necesita tanto Pandas como PyArrow para escribir y leer un archivo es usar un conjunto de datos y comprobar cuánto tarda cada uno en completar las mismas tareas en las mismas condiciones. Para ello, crearemos un DataFrame aleatorio con cinco columnas y un millón de registros. Una vez hecho esto mediremos cuánto tarda cada una de las librerías en guardar y leer los mismos datos. Comparativa que se puede hacer con el siguiente código.

import pandas as pd
import pyarrow as pa
import pyarrow.csv
import numpy as np
import time

cols = 5
rows = 1000000

np.random.seed(0)

df = pd.DataFrame(np.random.rand(rows, cols),
                  columns=["A", "B", "C", "D", "E"])

# CSV con Pandas
tic = time.process_time()
df.to_csv('pandas.csv')
pandas_write = time.process_time() - tic

tic = time.process_time()
csv = pd.read_csv('pandas.csv')
pandas_read = time.process_time() - tic

# CSV con PyArrow
tic = time.process_time()
tbl = pa.Table.from_pandas(df)
pa.csv.write_csv(tbl, 'pyarrow.csv')
pyarrow_write = time.process_time() - tic

tic = time.process_time()
csv = pa.csv.read_csv('pyarrow.csv').to_pandas()
pyarrow_read = time.process_time() - tic

# Compara tiempos
times = pd.DataFrame({'write': [pandas_write, pyarrow_write],
                      'read': [pandas_read, pyarrow_read]},
                     index=['Pandas', 'PyArrow'])

Obteniéndose como resultado la siguiente tabla, en la que se puede ver cómo PyArror puede almacenar archivos CSV 10 veces más rápido en Python que Pandas. Una mejora considerable cuando es necesario leer grandes conjuntos de datos. Aunque en la lectura de datos no existe una mejora apreciable entre las dos librerías.

Conclusiones

A partir de los resultados obtenidos es aconsejable reemplazar Pandas por PyArrow cuando necesitemos generar archivos CSV. El tiempo de escritura en las pruebas se ha reducido un orden de magnitud, algo que puede ser considerable especialmente al trabajar con grandes conjuntos de datos. Aunque, para leer los mismos archivos, no es necesario abandonar Pandas, ya que la diferencia de tiempo no es tan importante.

Imagen de Steve Sewell en Pixabay