using Dask to load many CSV files with different columns

Question

I have many CSV files saved in AWS s3 with the same first set of columns and a lot of optional columns. I don't want to download them one by one and than use pd.concat to read it, since this takes a lot of time and has to fit in to the computer memory. Instead, I'm trying to use Dask to load and sum up all of these files, when optional columns should should be treated as zeros.

If all columns where the same I could use:

    import dask.dataframe as dd
    addr = "s3://SOME_BASE_ADDRESS/*.csv"
    df = dd.read_csv(addr)
    df.groupby(["index"]).sum().compute()

but it doesn't work with files that don't have same number of columns, since Dask assumes it can use the first columns for all files:

File ".../lib/python3.7/site-packages/pandas/core/internals/managers.py", line 155, in set_axis 'values have {new} elements'.format(old=old_len, new=new_len)) ValueError: Length mismatch: Expected axis has 64 elements, new values have 62 elements

According to this thread I can either read all headers in advanced (for example by writing them as I produce and save all of the small CSV's) or use something like this:

df = dd.concat([dd.read_csv(f) for f in filelist])

I wonder if this solution is actually faster/better than just directly use pandas? In general, I'd like to know what is the best (mainly fastest) way to tackle this issue?

user_10693596 · Answer

It might be a good idea to use delayed to standardize dataframes before converting them to a dask dataframe (whether this is optimal for your use case is difficult to judge).

import dask.dataframe as dd
from dask import delayed

list_files = [...] # create a list of files inside s3 bucket
list_cols_to_keep = ['col1', 'col2']

@delayed
def standard_csv(file_path):
    df = pd.read_csv(file_path)
    df = df[list_cols_to_keep]
    # add any other standardization routines, e.g. dtype conversion
    return df

ddf = dd.from_delayed([standard_csv(f) for f in list_files])

user_9156913 · Answer

I ended up giving up using Dask since it was too slow and used aws s3 sync to download the data and multiprocessing.Pool to read and concat them:

# download:
def sync_outputs(out_path):
    local_dir_path = f"/tmp/outputs/"
    safe_mkdir(os.path.dirname(local_dir_path))
    cmd = f'aws s3 sync {job_output_dir} {local_dir_path} > /tmp/null' # the last part is to avoid prints
    os.system(cmd)
    return local_dir_path

# concat:
def read_csv(path):
    return pd.read_csv(path,index_col=0)

def read_csvs_parallel(local_paths):
    from multiprocessing import Pool
    import os
    with Pool(os.cpu_count()) as p:
        csvs = list(tqdm(p.imap(read_csv, local_paths), desc='reading csvs', total=len(paths)))
    return csvs

# all togeter:
def concat_csvs_parallel(out_path):
    local_paths = sync_outputs(out_path)
    csvs = read_csvs_parallel(local_paths)
    df = pd.concat(csvs)
    return df

aws s3 sync dowloaded about 1000 files (~1KB each) in about 30 second, and reading than with multiproccesing (8 cores) took 3 seconds, this was much faster than also downloading the files using multiprocessing (almost 2 minutes for 1000 files)

using Dask to load many CSV files with different columns

2 Answers