Fractionner le fichier csv basé sur une colonne particulière en utilisant Python

Question

Je suis un débutant Python, et j'ai fait quelques scripts de base. Mon dernier défi est de prendre un très grand fichier csv (10gb +) et le diviser en un certain nombre de fichiers plus petits, en fonction de la valeur d'une variable particulière dans chaque ligne.

Par exemple, le fichier peut ressembler à ceci:

Category,Title,Sales 
"Books","Harry Potter",1441556 
"Books","Lord of the Rings",14251154 
"Series", "Breaking Bad",6246234 
"Books","The Alchemist",12562166 
"Movie","Inception",1573437 

Et je veux diviser le fichier en plusieurs fichiers séparés: Books.csv, Series.csv, Movie.csv

En la réalité il y aura des centaines de catégories, et ils ne seront pas triés. Dans ce cas, ils sont dans la première colonne mais dans le futur ils ne le seront peut-être pas.

J'ai trouvé quelques solutions en ligne mais rien en Python. Il existe une commande AWK très simple qui peut le faire en une ligne, mais je ne peux pas accéder à AWK au travail.

J'ai écrit le code suivant qui fonctionne, mais je pense que c'est probablement très inefficace. Quelqu'un peut-il suggérer comment l'accélérer?

import csv 

#Creates empty set - this will be used to store the values that have already been used 
filelist = set() 

#Opens the large csv file in "read" mode 
with open('//directory/largefile', 'r') as csvfile: 

    #Read the first row of the large file and store the whole row as a string (headerstring) 
    read_rows = csv.reader(csvfile) 
    headerrow = next(read_rows) 
    headerstring=','.join(headerrow) 

    for row in read_rows: 

     #Store the whole row as a string (rowstring) 
     rowstring=','.join(row) 

     #Defines filename as the first entry in the row - This could be made dynamic so that the user inputs a column name to use 
     filename = (row[0]) 

     #This basically makes sure it is not looking at the header row. 
     if filename != "Category": 

      #If the filename is not in the filelist set, add it to the list and create new csv file with header row. 
      if filename not in filelist:  
       filelist.add(filename) 
       with open('//directory/subfiles/' +str(filename)+'.csv','a') as f: 
        f.write(headerstring) 
        f.write("\n") 
        f.close()  
      #If the filename is in the filelist set, append the current row to the existing csv file.  
      else: 
       with open('//directory/subfiles/' +str(filename)+'.csv','a') as f: 
        f.write(rowstring) 
        f.write("\n") 
        f.close() 

Merci!

Answer 1

Une manière efficace de mémoire et qui évite de rouvrir des fichiers à ajouter ici (tant que vous n'allez pas générer d'énormes quantités de handles de fichiers ouverts) est d'utiliser un dict pour mapper la catégorie à un fichier . Lorsque ce fichier n'est pas encore ouvert, créez-le et écrivez l'en-tête, puis écrivez toujours toutes les lignes dans le fichier correspondant, par exemple:

import csv 

with open('somefile.csv') as fin:  
    csvin = csv.DictReader(fin) 
    # Category -> open file lookup 
    outputs = {} 
    for row in csvin: 
     cat = row['Category'] 
     # Open a new file and write the header 
     if cat not in outputs: 
      fout = open('{}.csv'.format(cat), 'w') 
      dw = csv.DictWriter(fout, fieldnames=csvin.fieldnames) 
      dw.writeheader() 
      outputs[cat] = fout, dw 
     # Always write the row 
     outputs[cat][1].writerow(row) 
    # Close all the files 
    for fout, _ in outputs.values(): 
     fout.close()