在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")