有条件地在pandas中读取csv

Question

我有大型CSV，我只对这些行的子集感兴趣。特别是，我想读取在满足特定条件之前发生的所有行。

例如，如果read_csv将产生数据框：

     A    B      C
1   34   3.20   'b'
2   24   9.21   'b'
3   34   3.32   'c'
4   24   24.3   'c'
5   35   1.12   'a'
... 
1e9 42   2.15   'd'

有什么方法可以读取csv中的所有行，直到col B超过10。在上面的示例中，我想读入：

     A    B      C
1   34   3.20   'b'
2   24   9.21   'b'
3   34   3.32   'c'
4   24   24.3   'c'

我知道在读入数据帧后如何将这些行扔掉，但是到现在为止，我已经花了所有的计算来读入它们。在读取csv之前，我无法访问最后一行的索引请不要跳过脚）

Answer 1

您可以分批读取csv。由于pd.read_csv在chunksize指定参数时将返回迭代器，因此您可以使用它itertools.takewhile来读取所需数量的块，而无需读取整个文件。

import itertools as IT
import pandas as pd

chunksize = 10 ** 5
chunks = pd.read_csv(filename, chunksize=chunksize, header=None)
chunks = IT.takewhile(lambda chunk: chunk['B'].iloc[-1] < 10, chunks)
df = pd.concat(chunks)
mask = df['B'] < 10
df = df.loc[mask]

或者，为了避免不得不df.loc[mask]从最后一个块中删除不需要的行，也许更干净的解决方案是定义一个自定义生成器：

import itertools as IT
import pandas as pd

def valid(chunks):
    for chunk in chunks:
        mask = chunk['B'] < 10
        if mask.all():
            yield chunk
        else:
            yield chunk.loc[mask]
            break

chunksize = 10 ** 5
chunks = pd.read_csv(filename, chunksize=chunksize, header=None)
df = pd.concat(valid(chunks))