Efficiently filter a large (100gb+) csv file (v3)

Question

I'm in the process of filtering some very(!) large files (100gb+): I can't download files with a lower granularity. This is a followup from this question.

The problem is as follows: I need to filter large files that look like the following (3b+ rows).

TIC, Date, Time, Bid, Offer
AAPL, 20090901, 09:45, 145, 145.5 
AAPL, 20090902, 09:45, 145, 145.5
AAPL, 20090903, 09:45, 145, 145.5

I filter based on TICKER+DATE combinations found in an external file. I have, on average, ~ 1200 dates of interest per firm for ~ 700 firms. The large file contains all dates for the firms of interest, for which I want to extract only a few dates of interest. The big files are split by month (2013-01, 2013-02 etc.).

# Example firm_dates_of_interest file.
AAPL, 20090902
AAPL, 20090903

A few changes were made since the previous post:

• Extract the filtering process into a generator function
• Extract chunks from this iterator

I'm currently at 4 minutes of processing time for 30 million rows (1% of the file); I tested a few files and it works properly. However, with about 3 billion rows per file, that puts it at ~8 hours for one 120gb file. Seeing as I have about twelve files, I'm very curious whether I can get significant performance improvements by doing things differently. Example files can be found here (sourcefile, 3gb zipped) and here (firmdates).

Any tips are greatly appreciated.

import os
import datetime
import csv
import cProfile
import re
ROOT_DIR = "H:/ROOT_DIR/"
SOURCE_FILES = os.path.join(ROOT_DIR, '15. Speedtest')
EXPORT_DIR = ROOT_DIR
DATES_FILE = os.path.join(ROOT_DIR, "10. Dates of interest/firm_date_of_interest_withSP.csv")

# Build the original date dict
# For example:
#    d['AAPL'] is a list with ['20140901', '20140902', '20140901']
with open(DATES_FILE, "r") as csvfile:
    d = {}
    reader = csv.reader(csvfile)
    reader.next()
    for line in reader:
        firm = line[1]
        date = line[2]
        if firm in d.keys():
            d[firm].append(date)
        else:
            d[firm] = [date]


def filter_lines(filename, d):
    """ Given a dictionary with key, value a Ticker and dates_of_interest, yield
    only the filtered rows with those ticker / dates pairs. """
    with open(filename, "rb") as csvfile:
        datareader = csv.reader(csvfile)
        for row in datareader:
            try:
                if row[1] in d[row[0]]:
                    yield row
            except KeyError:
                continue
            except IndexError:
                continue


def get_chunk(iterable, chunk_size):
    """ Given an iterable and chunk_size, return chunks of chunk_size"""
    result = []
    for item in iterable:
        result.append(item)
        if len(result) == chunk_size:
            yield tuple(result)
            result = []
    if len(result) > 0:
        yield tuple(result)


def main():
    start = datetime.datetime.now()
    for root, dir, files in os.walk(SOURCE_FILES):
        for i, file in enumerate(files):
            basename = os.path.splitext(file)[0]
            source_filepath = os.path.join(root, file)

            # Annotate files with 'DONE' after succesful processing: skip those
            if re.search("DONE", basename):
                continue

            startfile = datetime.datetime.now()
            iterable = filter_lines(source_filepath, d)
            for num_saved, chunk in enumerate(get_chunk(iterable, 3000000)):
                output_filename = os.path.join(EXPORT_DIR, basename+' EXTRACT' + str(num_saved) + '.csv')
                with open(output_filename, 'wb') as csvfile:
                    writer = csv.writer(csvfile, quoting=csv.QUOTE_NONNUMERIC)
                    for line in chunk:
                        writer.writerow(line)
                file_elapsed = datetime.datetime.now() - startfile
                print("Took me %d seconds for 3000000 extracted rows.." % file_elapsed.seconds)

            new_filename = os.path.join(root, os.path.splitext(file)[0]+'- DONE.csv')
            os.rename(source_filepath, new_filename)

    elapsed = datetime.datetime.now() - start
    print("Took me %d seconds..." % elapsed.seconds)

if __name__ == "__main__":
    main()

Answer 1

I'm sorry to have to say so, but reading back over the answers to your two previous questions, it seems to me as if you've been given some bad advice. In particular, this answer advised you to switch to using the csv module. But that was a mistake, because:

1. The csv module has to handle all the details of the CSV format, which can be quite complicated (quoted fields, choice of field separator etc).

2. The csv module splits all the fields of each line, but here you are only interested in the first two fields.

So there is a lot of wasted effort. To demonstrate this, I made a file with 10 million records in your format:

>>> with open('data.csv', 'w') as f:
...     for i in range(10**7):
...         _ = f.write('AAPL, {:08d}, 09:45, 145, 145.5\n'.format(i//1000))

This is about a third of a gigabyte in size. To read (and discard) all the lines from this file takes about 7.5 seconds:

>>> from collections import deque
>>> from timeit import timeit
>>> with open('data.csv') as f:
...     timeit(lambda:deque(f, maxlen=0), number=1)
7.537129107047804

Which is a rate of 1.3 million lines a second. (Using collections.deque with maxlen=0 to consume an iterable with no Python interpreter overhead is a useful trick to know about.)

Now, using the csv module, reading all the lines in the file takes about three times as long:

>>> import csv
>>> with open('data.csv') as f:
...    timeit(lambda:deque(csv.reader(f), maxlen=0), number=1)
22.454868661938235

So switching to csv was a mistake.

What to do instead? Well, if I understand your question correctly, the records are quite constrained in format: they start with a stock ticker symbol, a comma, a space, a date in ISO 8601 format, and a comma. Moreover, you want to select lines based on just these two fields. So you might try putting the fields you are looking for into a set of strings:

keys = set(line.strip() for line in open('dates-of-interest.csv'))

and then finding the field boundaries yourself:

def filter_lines(in_filename, out_filename, keys):
    """Read records from in_filename and write records to out_filename if
    the string consisting of the first two comma-separated fields is
    found in the set keys.

    """
    with open(in_filename) as in_f, open(out_filename, 'w') as out_f:
        for line in in_f:
            ticker_end = line.find(',')
            date_end = line.find(',', ticker_end + 1)
            if line[:date_end] in keys:
                out_f.write(line)

In this test I filter out one in every five lines:

>>> keys = {'AAPL, {:08d}'.format(i) for i in range(0, 10000, 5)}
>>> timeit(lambda:filter_lines('data.csv', 'out.csv', keys), number=1)
14.926778070977889

This is about 670,000 lines a second.

Now, it might be the case that your records are even more constrained than I described above. For example, it might be the case that all the stock ticker symbols that you are looking for fall into a small range of lengths. For example maybe they are all between one and five letters long. Then the second comma must appear between positions 11 and 16 of the string:

01234567890123456789
A, 20150102, 09:45, 145, 145.5
AB, 20150102, 09:45, 145, 145.5
ABC, 20150102, 09:45, 145, 145.5
ABCD, 20150102, 09:45, 145, 145.5
ABCDE, 20150102, 09:45, 145, 145.5

So we can restrict our search to these positions:

def filter_lines(in_filename, out_filename, keys):
    """Read records from in_filename and write records to out_filename if
    the string up to the first comma between positions 11 and 16 of
    line is found in the set keys.

    """
    with open(in_filename) as in_f, open(out_filename, 'w') as out_f:
        for line in in_f:
            date_end = line.find(',', 11, 16)
            if line[:date_end] in keys:
                out_f.write(line)

and this is substantially faster, about 850,000 lines a second on my computer:

>>> timeit(lambda:filter_lines('data.csv', 'out.csv', keys), number=1)
11.833066276973113

There are some more optimizations we can do:

1. open the files in binary mode, to avoid dealing with character set encodings;
2. avoid storing date_end in a local variable;
3. avoid passing the end postion 16 to the find (this is superfluous);
4. (suggested in comments by d33tah) run it under PyPy.

This results in:

def filter_lines(in_filename, out_filename, keys):
    """Read records from in_filename and write records to out_filename if
    the beginning of the line (taken up to the first comma at or after
    position 11) is found in keys (which must be a set of byte strings).

    """
    with open(in_filename, 'rb') as in_f, open(out_filename, 'wb') as out_f:
        for line in in_f:
            if line[:line.find(b',', 11)] in keys:
                out_f.write(line)

Note that we now require keys to be a set of byte strings, so we better encode them:

>>>> keys = {'AAPL, {:08d}'.format(i).encode() for i in range(0, 10000, 5)}

This processes about 1.8 million lines per second:

>>>> timeit(lambda:filter_lines('data.csv', 'out.csv', keys), number=1)
5.53329086304

which suggests that a 100 GiB file could be filtered in about 30 minutes. Of course, this is all on my computer, which might be faster or slower than yours. But the general approach almost certainly applies.