在pandas中強加所需的列約束read_csv

Question

我正在尋找將大CSV讀入數據框，並附加了一個我想失敗的約束條件如果缺少某些列（因爲輸入不像預期的那樣），我會提前，但是我確實需要全部列，而不僅僅是要包含在Dataframe中的所需列。在pandas.read_csv中，如果我想指定要讀入的列的子集，我似乎可以使用usecols參數，但是我可以看到檢查哪些列將在我要讀取的數據框中的唯一顯而易見的方式是實際上閱讀文件。

我做了讀取數據幀作爲一個迭代工作第一通版本，獲取第一線，檢查該列是否存在，然後讀取與正常參數的文件：

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.') 

這對我來說感覺有點混亂，並沒有真正處理​​（例如不可搜索的流，或者我不應該從0開始的緩衝區）的所有可能的輸入。顯然，我可以調整我在這裏對我的具體用例的時刻，並完成它，但我想知道是否有一個更優雅的方式來做到這一點（最好只是使用標準的熊貓功能），在一般 。

Answer 1

你可以只讀一行，並測試是否所有必需的列都存在？例如：

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" 

---

要通過流做到這一點，另一種方法是通過csv.reader()讀它，這是兼容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")