Add calculated column to pandas dataframe on the fly while iterating over the lines of a csv file?

Question

I have a large space separated input file input.csv , which I can't hold in memory:

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

If I use the iterator=True argument for pandas.read_csv , then it returns a TextFileReader / TextParser object. This allows filtering the file on the fly and only selecting rows for which column A is greater than 2.

But how do I add a third column to the dataframe on the fly without having to loop over all of the data once more?

Specifically I want column C to be equal to column A multiplied by the value in a dictionary d , which has the value of column B as its key; ie C = A*d[B] .

Currently I have this code:

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)

Which prints this output:

   A  B
1  3  4

How do I get it to print this output ( C = A*d[B] ):

   A  B  C
1  3  4  9

Answer 1

You can use a generator to work on the chunks one at a time:

Code:

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]

Test Code:

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)

Results:

   A  B     C
1  3  4   9.0
2  4  4  12.0