Impose erforderliche Spalten Einschränkung in Pandas read_csv

Question

Ich suche eine große CSV in einem Datenrahmen mit der zusätzlichen Einschränkung, die ich fehlschlagen möchte früh wenn bestimmte Spalten fehlen (da die Eingabe nicht wie erwartet ist), aber ich wollen alle die Spalten, nicht nur die erforderlichen Spalten, in den Dataframe aufgenommen werden. In pandas.read_csv, es scheint, dass ich das usecols Argument verwenden kann, wenn ich die Teilmenge der Spalten zum Einlesen angeben möchte, aber die einzige offensichtliche Möglichkeit zu überprüfen, welche Spalten in einem Dataframe ich bin zu lesen ist zu lese eigentlich die Datei.

ich eine Arbeits First-Pass-Version vorgenommen haben, die den Datenrahmen als Iterator liest, erhält die erste Zeile überprüft, dass die Spalten vorhanden sind, dann die Datei mit den normalen Argumente lautet:

import pandas as pd 
from io import StringIO 

class MissingColumnsError(ValueError): 
    pass 

def cols_enforced_reader(*args, cols_must_exist=None, **kwargs): 
    if cols_must_exist is not None: 
     # Read the first line of the DataFrame and check the columns 
     new_kwargs = kwargs.copy() 
     new_kwargs['iterator'] = True 
     new_kwargs['chunksize'] = 1 

     if len(args): 
      filepath_or_buffer = args[0] 
      args = args[1:] 
     else: 
      filepath_or_buffer = new_kwargs.get('filepath_or_buffer', None) 

     df_iterator = pd.read_csv(filepath_or_buffer, *args, **new_kwargs) 

     c = next(df_iterator) 
     if not all(col in c.columns for col in cols_must_exist): 
      raise MissingColumnsError('Some required columns were missing!') 

     seek = getattr(filepath_or_buffer, 'seek', None) 
     if seek is not None: 
      if filepath_or_buffer.seekable(): 
       filepath_or_buffer.seek(0) 

    return pd.read_csv(filepath_or_buffer, *args, **kwargs) 

in_csv = """col1,col2,col3\n0,1,2\n3,4,5\n6,7,8""" 

# Should succeed 
df = cols_enforced_reader(StringIO(in_csv), cols_must_exist=['col1']) 
print('First call succeeded as expected.') 

# Should fail 
try: 
    df = cols_enforced_reader(StringIO(in_csv), cols_must_exist=['col7']) 
except MissingColumnsError: 
    print('Second call failed as expected.') 

Diese fühlt sich ein bisschen chaotisch an und behandelt nicht wirklich alle möglichen Eingaben für filepath_or_buffer (nicht suchbare Ströme zum Beispiel, oder Puffer, wo ich nicht bei 0 anfangen soll). Offensichtlich kann ich für den Moment das, was ich hier habe, an meinem spezifischen Anwendungsfall anpassen und bin damit fertig, aber ich frage mich, ob es eine elegantere Art gibt, dies zu tun (vorzugsweise nur Standard-Pandas-Funktionen), die in allgemein funktioniert .

Answer 1

Sie könnten nur eine Zeile lesen und testen, ob alle erforderlichen Spalten vorhanden sind? Zum Beispiel:

import pandas as pd 

required_cols = ['col1', 'col2'] 
cols = pd.read_csv('input.csv', nrows=1).columns 

if all(req in cols for req in required_cols): 
    print pd.read_csv('input.csv') 
else: 
    print "Columns missing" 

---

Um dies zu tun über einen Bach, ein alternativer Ansatz es über ein csv.reader() zu lesen wäre, dann ist dies kompatibel mit itertools.tee():

import pandas as pd 
from itertools import tee 
import csv 

required_cols = ['col1', 'col2'] 

with open('input.csv') as f_input: 
    csv_input = csv.reader(f_input) 
    csv_stream1, csv_stream2 = tee(csv_input, 2) 
    header = next(csv_stream1) 

    if all(req in header for req in required_cols): 
     df = pd.DataFrame(list(csv_stream2)[1:], columns=header) 
     print(df) 
    else: 
     print("Columns missing")