如何得到一个精确的python迭代器？

Question

我正在用pandas读取一个超大csv文件（10G），read_csv（filename，chunksize=chunksize）返回一个迭代器（assum it names‘reader’）。现在我想得到一个精确的块，因为我只需要一些行（例如，我读的csv文件有100000000行，我想得到50000000行，后面有1000行），除了tranverse迭代器，我应该怎么做，直到它到达我想要的块？你知道吗

这是我以前的密码：

def get_lines_by_chunk(file_name, line_beg, line_end, chunk_size=-1):
func_name = 'get_lines_by_chunk'
line_no = get_file_line_no(file_name)

if chunk_size < 0:
    chunk_size = get_chunk_size(line_no, line_beg, line_end)

reader = pd.read_csv(file_name, chunksize=chunk_size)
data = pd.DataFrame({})

flag = 0

for chunk in reader:
    line_before = flag * chunk_size
    flag = flag + 1
    line_after = flag * chunk_size
    if line_beg >= line_before and line_beg <= line_after:
        if line_end >= line_after:
            temp = chunk[line_beg - line_before : chunk_size]
            data = pd.concat([data, temp], ignore_index=True)
        else:
            temp = chunk[line_beg - line_before : line_end - line_before]
            data = pd.concat([data, temp], ignore_index=True)
            return data
    elif line_end <= line_after and line_end >= line_before:
        temp = chunk[0 : line_end - line_before]
        data = pd.concat([data, temp], ignore_index=True)
        return data
    elif line_beg < line_before and line_end > line_after:
        temp = chunk[0 : chunk_size]
        data = pd.concat([data, temp], ignore_index=True)

return data

Answer 1

如果需要用不同大小的块读取CSV文件，可以使用iterator=True：

假设我们有一个1000行的DF（关于它是如何生成的，请参见设置部分）

In [103]: reader = pd.read_csv(fn, iterator=True)

In [104]: reader.get_chunk(5)
Out[104]:
   a   b
0  1   8
1  2  28
2  3  85
3  4  56
4  5  29

In [105]: reader.get_chunk(3)
Out[105]:
   a   b
5  6  55
6  7  16
7  8  96

注意：get_chunk不能跳过数据，它将连续读取具有指定块大小的数据

如果只想读取第100-110行：

In [106]: cols = pd.read_csv(fn, nrows=1).columns.tolist()

In [107]: cols
Out[107]: ['a', 'b']

In [109]: pd.read_csv(fn, header=None, skiprows=100, nrows=10, names=cols)
Out[109]:
     a   b
0  100  52
1  101  15
2  102  74
3  103  10
4  104  35
5  105  73
6  106  48
7  107  49
8  108   1
9  109  56

但是，如果您可以使用HDF5格式，它将更容易和更快：

先将其另存为HDF5：

In [110]: df.to_hdf('c:/temp/test.h5', 'mydf', format='t', data_columns=True, compression='blosc', complevel=9)

现在我们可以通过索引位置来读取它，如下所示：

In [113]: pd.read_hdf('c:/temp/test.h5', 'mydf', start=99, stop=109)
Out[113]:
       a   b
99   100  52
100  101  15
101  102  74
102  103  10
103  104  35
104  105  73
105  106  48
106  107  49
107  108   1
108  109  56

或查询（类似SQL）：

In [115]: pd.read_hdf('c:/temp/test.h5', 'mydf', where="a >= 100 and a <= 110")
Out[115]:
       a   b
99   100  52
100  101  15
101  102  74
102  103  10
103  104  35
104  105  73
105  106  48
106  107  49
107  108   1
108  109  56
109  110  23

设置：

In [99]: df = pd.DataFrame({'a':np.arange(1, 1001), 'b':np.random.randint(0, 100, 1000)})

In [100]: fn = r'C:\Temp\test.csv'

In [101]: df.to_csv(fn, index=False)

In [102]: df.shape
Out[102]: (1000, 2)