Usar chunksize en pandas

La biblioteca pandas organiza los datos en DataFrames de filas y columnas, y funciones de lectura como read_csv normalmente cargan todo el archivo en un solo DataFrame. Cuando un archivo es demasiado grande para cargarlo cómodamente —por ejemplo, un conjunto de datos como la exportación de calificaciones de MovieLens con millones de filas—, el parámetro chunksize permite leerlo y procesarlo en lotes de filas.

Cómo funciona chunksize en pandas

Al pasar chunksize=n a una función de lectura de pandas que lo admita, el valor devuelto cambia. En lugar de un DataFrame, read_csv() devuelve un TextFileReader, un objeto similar a un iterador que produce un DataFrame por lote. Cada lote contiene como máximo n filas, y el último puede contener menos cuando el número total de filas no es un múltiplo exacto del tamaño del lote.

Hay dos detalles importantes al adaptar este patrón. Primero, chunksize cuenta filas, no bytes: cada lote conserva todas las columnas del archivo y solo varía el número de filas. Segundo, el parámetro no acelera el análisis del archivo por sí solo. Cambia cuándo se asigna la memoria, así que si resulta útil depende de lo que el programa conserve de cada lote.

Leer un archivo CSV en lotes

Para leer un archivo en lotes, pasa chunksize a read_csv e itera sobre el resultado. El ejemplo siguiente escribe un pequeño CSV de calificaciones en un archivo temporal y luego lo lee de vuelta de a tres filas, imprimiendo el atributo shape de cada DataFrame producido para que puedas ver la división en lotes.

Ejecuta este ejemplo:

El código y las etiquetas de salida se mantienen en inglés; la salida muestra el valor de shape de cada lote.

import os
import tempfile

import pandas as pd

rows = ["userId,movieId,rating"]
for i in range(7):
    rows.append(f"{i + 1},{100 + i},{1.0 + (i % 5) * 0.5}")

path = os.path.join(tempfile.mkdtemp(), "ratings.csv")
with open(path, "w") as f:
    f.write("\n".join(rows) + "\n")

reader = pd.read_csv(path, chunksize=3)
for number, chunk in enumerate(reader, start=1):
    print(f"chunk {number}: shape={chunk.shape}")

chunk 1: shape=(3, 3)
chunk 2: shape=(3, 3)
chunk 3: shape=(1, 3)

El archivo contiene siete filas de datos y chunksize=3, por lo que el TextFileReader produce dos lotes completos seguidos de un resto de una fila. El número de columnas se mantiene en tres en cada lote, lo que confirma que la división en lotes se aplica solo a las filas.

Procesar cada lote

El patrón útil consiste en acumular un resultado pequeño por lote —un conteo, una suma o un subconjunto filtrado— en lugar de ir guardando todos los lotes. Este ejemplo agrega el mismo archivo manteniendo un recuento acumulado de filas y una suma acumulada de calificaciones, de modo que el archivo completo nunca tenga que permanecer en memoria al mismo tiempo.

Ejecuta este ejemplo:

El código y las etiquetas de salida se mantienen en inglés; la salida muestra el total de filas leídas y la calificación media.

import os
import tempfile

import pandas as pd

rows = ["userId,movieId,rating"]
for i in range(7):
    rows.append(f"{i + 1},{100 + i},{1.0 + (i % 5) * 0.5}")

path = os.path.join(tempfile.mkdtemp(), "ratings.csv")
with open(path, "w") as f:
    f.write("\n".join(rows) + "\n")

total_rows = 0
rating_sum = 0.0
for chunk in pd.read_csv(path, chunksize=3):
    total_rows += len(chunk)
    rating_sum += chunk["rating"].sum()

print(f"rows read: {total_rows}")
print(f"average rating: {rating_sum / total_rows:.2f}")

rows read: 7
average rating: 1.79

Cualquier operación de DataFrame funciona dentro del bucle —filtrado, selección de columnas o agregaciones con groupby—, así que puedes transformar cada lote antes de incorporarlo al acumulador.

Comprender el tamaño de los lotes y las compensaciones de memoria

Leer en lotes puede reducir el consumo máximo de memoria, pero no es automático. El uso máximo es aproximadamente el tamaño del lote más grande más lo que el programa retenga. Si añades cada lote a una lista y luego los concatenas, vuelves a mantener todo el conjunto de datos en memoria y puedes acabar usando más memoria que con una sola lectura normal.

Elige el valor equilibrando dos factores. Un chunksize mayor implica menos iteraciones del bucle y menos sobrecarga por iteración, pero cada lote ocupa más memoria. Un valor menor reduce el pico a costa de realizar más pasadas por el bucle. Además, escribe el código del bucle sin suponer un número fijo de filas, ya que el último lote suele ser más corto.

Usar la lectura por lotes con otras fuentes de datos

chunksize no es exclusivo de read_csv. Otras funciones de lectura de pandas que lo aceptan incluyen pandas.read_json (con lines=True), pandas.read_sql y pandas.read_sql_table, pandas.read_stata y pandas.read_sas; algunas funciones, como pandas.read_hdf, ofrecen en su lugar un modo iterator equivalente. La compatibilidad y las opciones aceptadas varían según la función y la versión de pandas, así que consulta la documentación oficial de la función que vayas a usar.

El patrón se mantiene igual siempre que haya lectura por lotes disponible: solicita lotes, itera sobre ellos e incorpora cada lote a un resultado pequeño.