La forma más rápida de escribir CSV grandes con Python

Quiero escribir algunos datos de muestra aleatorios en un archivo csv hasta que tenga un tamaño de 1 GB. El siguiente código está funcionando:

import numpy as np import uuid import csv import os outfile = 'data.csv' outsize = 1024 # MB with open(outfile, 'ab') as csvfile: wtr = csv.writer(csvfile) while (os.path.getsize(outfile)//1024**2) < outsize: wtr.writerow(['%s,%.6f,%.6f,%i' % (uuid.uuid4(), np.random.random()*50, np.random.random()*50, np.random.randint(1000))]) 

¿Cómo conseguirlo más rápido?

Eliminar todas las cosas innecesarias, y por lo tanto debería ser más rápido y más fácil de entender:

 import random import uuid outfile = 'data.csv' outsize = 1024 * 1024 * 1024 # 1GB with open(outfile, 'ab') as csvfile: size = 0 while size < outsize: txt = '%s,%.6f,%.6f,%i\n' % (uuid.uuid4(), random.random()*50, random.random()*50, random.randrange(1000)) size += len(txt) csvfile.write(txt) 

El problema parece estar principalmente ligado a IO. Puede mejorar un poco la E / S escribiendo en el archivo en trozos más grandes en lugar de escribir una línea a la vez:

 import numpy as np import uuid import csv import os outfile = 'data-alt.csv' outsize = 10 # MB chunksize = 1000 with open(outfile, 'ab') as csvfile: while (os.path.getsize(outfile)//1024**2) < outsize: data = [[uuid.uuid4() for i in range(chunksize)], np.random.random(chunksize)*50, np.random.random(chunksize)*50, np.random.randint(1000, size=(chunksize,))] csvfile.writelines(['%s,%.6f,%.6f,%i\n' % row for row in zip(*data)]) 

Puede experimentar con el tamaño de trozo (el número de filas escritas por trozo) para ver qué funciona mejor en su máquina.


Aquí hay un punto de referencia, que compara el código anterior con su código original, con un outsize establecido en 10 MB:

 % time original.py real 0m5.379s user 0m4.839s sys 0m0.538s % time write_in_chunks.py real 0m4.205s user 0m3.850s sys 0m0.351s 

Así que esto es un 25% más rápido que el código original.


PD. Intenté reemplazar las llamadas a os.path.getsize con una estimación del número de líneas totales necesarias. Lamentablemente, no mejoró la velocidad. Como el número de bytes necesarios para representar el int final varía, la estimación también es inexacta, es decir, no reproduce perfectamente el comportamiento de su código original. Así que dejé el os.path.getsize en su lugar.