rows CSV lazy import and export

brazilian-cities.csv

test.py

import time

import rows


def routes(cities):
    'Generate (fake) routes based on all cities'

    for city_1 in cities:
        for city_2 in cities:
            if city_1 == city_2:
                continue
            yield {
                    'city_from': city_1.city,
                    'state_from': city_1.state,
                    'city_to': city_2.city,
                    'state_to': city_2.state, }


def export_routes(routes, filename, batch_size):
    'Export routes to a CSV file'

    # The `import_from_dicts` will generate a `LazyTable` (`lazy=True`) and
    # will use the first 500 rows (`samples=500`) to determine the field types.
    # Only the first 500 rows will be stored in the memory (as a buffer), so
    # when we iterate over the `table` object they'll be returned and after
    # this the generator will be consumed until its end.
    table = rows.import_from_dicts(routes, lazy=True, samples=500)

    # Exporting a `LazyTable` is lazy by default: the exporter will consume it
    # while writing to the destination. The `batch_size` parameter is the
    # number of consumed rows (from the generator) per time.
    start = time.time()
    rows.export_to_csv(table, filename, batch_size=batch_size)
    end = time.time()
    print('batch_size={:06d}, time={:7.2f}'.format(batch_size, end - start))


if __name__ == '__main__':
    # Importing only the fields we need (`import_fields`) will speed up the
    # process.
    # Setting `lazy=False` will store cities in memory and allow us to iterate
    # over the `Table` object as many times as we want (if it was a `LazyTable`
    # we could only iterate once).
    cities = rows.import_from_csv('brazilian-cities.csv',
                                  import_fields=['city', 'state'],
                                  lazy=False)

    export_routes(routes(cities), '/tmp/ram/routes-000001.csv', 1)
    export_routes(routes(cities), '/tmp/ram/routes-000010.csv', 10)
    export_routes(routes(cities), '/tmp/ram/routes-000100.csv', 100)
    export_routes(routes(cities), '/tmp/ram/routes-001000.csv', 1000)
    export_routes(routes(cities), '/tmp/ram/routes-010000.csv', 10000)
    export_routes(routes(cities), '/tmp/ram/routes-100000.csv', 100000)