How to read chunks of multiple large CSV files from google cloud storage using Dask without overloading the memory all at once

Viewed 835

I'm trying to read a bunch of large csv files (multiple files) from google storage. I use the Dask distribution library for parallel computation, but the problem I'm facing here is, though I mention the blocksize (100mb), I'm not sure how to read partition by partition and save it to my postgres database so that I don't want overload my memory.

    from dask.distributed import Client
    from dask.diagnostics import ProgressBar
    client = Client(processes=False)
    import dask.dataframe as dd

    def read_csv_gcs():
      with ProgressBar():
        df = dd.read_csv('gs://mybucket/renish/*.csv', blocksize=100e6)
        pd = df.compute(scheduler='threads')
        return pd

    def write_df_to_db(df):
      try:
        from sqlalchemy import create_engine
        engine = create_engine('postgresql://usr:pass@localhost:5432/sampledb')
        df.to_sql('sampletable', engine, if_exists='replace',index=False)
      except Exception as e:
        print(e)
        pass

    pd = read_csv_gcs()
    write_df_to_db(pd)

The above code is my basic implementation, but as said I would like to read it in chunk and update the db. Something like

    df = dd.read_csv('gs://mybucket/renish/*.csv', blocksize=100e6)
    for chunk in df:
       write_it_to_db(chunk)

Is it possible to do it in Dask? or should I go for pandas's chunksize and iterate, then save it to DB (But I miss parallel computation here)?

Can someone shed some light?

1 Answers

This line

df.compute(scheduler='threads')

says: load the data in chunks in worker threads, and concatenate them all into a single in-memory dataframe, df. This is not what you wanted. You wanted to insert the chunks as they come and then drop them from memory.

You probably wanted to use map_partitions

df = dd.read_csv('gs://mybucket/renish/*.csv', blocksize=100e6)
df.map_partitions(write_it_to_db).compute()

or use df.to_delayed().

Note that, depending on your SQL driver, you might not be able to get parallelism this way, and if not, the pandas iter-chunk method would have worked just as well.

Related