python - 在迭代 csv 文件的行时，动态地将计算列添加到 pandas 数据框？

Question

我有一个很大的空格分隔输入文件input.csv，我无法将其保存在内存中:

## Header
# More header here
A   B
1   2
3   4

如果我对 pandas.read_csv 使用 iterator=True 参数，然后它返回一个 TextFileReader/TextParser 对象。这允许动态过滤文件并仅选择列 A 大于 2 的行。

但是如何动态地将第三列添加到数据框中，而不必再次循环所有数据？

具体来说，我希望列C等于列A乘以字典d中的值，该值具有列B 作为其键；即C = A*d[B]。

目前我有这个代码:

import pandas
d = {2: 2, 4: 3}
TextParser = pandas.read_csv('input.csv', sep=' ', iterator=True, comment='#')
df = pandas.concat([chunk[chunk['A'] > 2] for chunk in TextParser])
print(df)

打印此输出:

   A  B
1  3  4

如何让它打印此输出(C = A*d[B]):

   A  B  C
1  3  4  9

Answer 1

您可以使用生成器一次处理一个 block :

代码:

def on_the_fly(the_csv):
    d = {2: 2, 4: 3}
    chunked_csv = pd.read_csv(
        the_csv, sep='\s+', iterator=True, comment='#')

    for chunk in chunked_csv:
        rows_idx = chunk['A'] > 2
        chunk.loc[rows_idx, 'C'] = chunk[rows_idx].apply(
            lambda x: x.A * d[x.B], axis=1)
        yield chunk[rows_idx]

测试代码:

from io import StringIO
data = StringIO(u"""#
    A   B
    1   2
    3   4
    4   4
""")

import pandas as pd
df = pd.concat([c for c in on_the_fly(data)])
print(df)

结果:

   A  B     C
1  3  4   9.0
2  4  4  12.0