在pandas数据框中读取前几行的方法

Question

有没有一种简单的方法可以用 read_csv 只读取文件的前 n 行，而不需要提前知道文件的总行数？我有一个很大的文件，读取起来很慢，有时候我只想用前面，比如说20行，来做个样本（我更希望不加载整个文件，只取前面部分）。

如果我知道总行数，我可以这样做 footer_lines = total_lines - n，然后把这个值传给 skipfooter 这个参数。现在我用的方法是手动用Python抓取前 n 行，然后用 StringIO 转给 pandas：

import pandas as pd
from StringIO import StringIO

n = 20
with open('big_file.csv', 'r') as f:
    head = ''.join(f.readlines(n))

df = pd.read_csv(StringIO(head))

这样做还不错，但有没有更简洁的、符合 pandas 风格的方法可以用参数来实现呢？

Answer 1

我会在使用read_csv的时候加上'skiprows'这个参数，比如说：

df = pd.read_csv(filename, skiprows=range(2, 20000), nrows=10000)

Answer 2

我觉得你可以使用 nrows 这个参数。根据文档的说明：

nrows : int, default None

    Number of rows of file to read. Useful for reading pieces of large files

这个方法似乎有效。使用一个标准的大测试文件（大小为988504479字节，包含5344499行）：

In [1]: import pandas as pd

In [2]: time z = pd.read_csv("P00000001-ALL.csv", nrows=20)
CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
Wall time: 0.00 s

In [3]: len(z)
Out[3]: 20

In [4]: time z = pd.read_csv("P00000001-ALL.csv")
CPU times: user 27.63 s, sys: 1.92 s, total: 29.55 s
Wall time: 30.23 s