#!/usr/bin/env python3 """ Simple test of read and write times for columnar data formats: python filetimes.py [pandas|dask [hdf5base [xcolumn [ycolumn] [categories...]]]] Test files may be generated starting from any file format supported by Pandas: python -c "import filetimes ; filetimes.base='' ; filetimes.categories=['','']; filetimes.timed_write('')" """ from __future__ import annotations import time global_start = time.time() import os, os.path, sys, glob, argparse, resource, multiprocessing import pandas as pd import dask.dataframe as dd import numpy as np import datashader as ds import feather import fastparquet as fp from datashader.utils import export_image from datashader import transfer_functions as tf #from multiprocessing.pool import ThreadPool #dask.set_options(pool=ThreadPool(3)) # select a specific number of threads from dask import distributed # Toggled by command-line arguments DEBUG = False DD_FORCE_LOAD = False DASK_CLIENT = None class Parameters: base,x,y='data','x','y' dftype='pandas' categories=[] chunksize=76668751 cat_width=1 # Size of fixed-width string for representing categories columns=None cachesize=9e9 parq_opts=dict(file_scheme='hive', has_nulls=False, write_index=False) n_workers=multiprocessing.cpu_count() p=Parameters() filetypes_storing_categories = {'parq'} class Kwargs(dict): """Used to distinguish between dictionary argument values, and keyword-arguments. """ pass def benchmark(fn, args, filetype=None): """Benchmark when "fn" function gets called on "args" tuple. "args" may have a Kwargs instance at the end. If "filetype" is provided, it may be used to convert columns to categorical dtypes after reading (the "loading" is assumed). """ posargs = list(args) kwargs = {} # Remove Kwargs instance at end of posargs list, if one exists if posargs and isinstance(posargs[-1], Kwargs): lastarg = posargs.pop() kwargs.update(lastarg) if DEBUG: printable_posargs = ', '.join([str(posarg.head()) if hasattr(posarg, 'head') else str(posarg) for posarg in posargs]) printable_kwargs = ', '.join(['{}={}'.format(k, v) for k,v in kwargs.items()]) print('DEBUG: {}({}{})'.format(fn.__name__, printable_posargs, ', '+printable_kwargs if printable_kwargs else ''), flush=True) # Benchmark fn when run on posargs and kwargs start = time.time() res = fn(*posargs, **kwargs) # If we're loading data if filetype is not None: if filetype not in filetypes_storing_categories: opts = {} if p.dftype == 'pandas': opts['copy']=False for c in p.categories: res[c]=res[c].astype('category',**opts) # Force loading (--cache=persist was provided) if p.dftype == 'dask' and DD_FORCE_LOAD: if DASK_CLIENT is not None: # 2017-04-28: This combination leads to a large drop in # aggregation performance (both --distributed and # --cache=persist were provided) res = DASK_CLIENT.persist(res) distributed.wait(res) else: if DEBUG: print("DEBUG: Force-loading Dask dataframe", flush=True) res = res.persist() end = time.time() return end-start, res read = dict([(f, dict()) for f in ["parq","snappy.parq","gz.parq","feather","h5","csv"]]) def read_csv_dask(filepath, usecols=None): # Pandas writes CSV files out as a single file if os.path.isfile(filepath): return dd.read_csv(filepath, usecols=usecols) # Dask may have written out CSV files in partitions filepath_expr = filepath.replace('.csv', '*.csv') return dd.read_csv(filepath_expr, usecols=usecols) read["csv"] ["dask"] = lambda filepath,p,filetype: benchmark(read_csv_dask, (filepath, Kwargs(usecols=p.columns)), filetype) read["h5"] ["dask"] = lambda filepath,p,filetype: benchmark(dd.read_hdf, (filepath, p.base, Kwargs(chunksize=p.chunksize, columns=p.columns)), filetype) def read_feather_dask(filepath): df = feather.read_dataframe(filepath, columns=p.columns) return dd.from_pandas(df, npartitions=p.n_workers) read["feather"] ["dask"] = lambda filepath,p,filetype: benchmark(read_feather_dask, (filepath,), filetype) read["parq"] ["dask"] = lambda filepath,p,filetype: benchmark(dd.read_parquet, (filepath, Kwargs(index=False, columns=p.columns)), filetype) read["gz.parq"] ["dask"] = lambda filepath,p,filetype: benchmark(dd.read_parquet, (filepath, Kwargs(index=False, columns=p.columns)), filetype) read["snappy.parq"] ["dask"] = lambda filepath,p,filetype: benchmark(dd.read_parquet, (filepath, Kwargs(index=False, columns=p.columns)), filetype) def read_csv_pandas(filepath, usecols=None): # Pandas writes CSV files out as a single file if os.path.isfile(filepath): return pd.read_csv(filepath, usecols=usecols) # Dask may have written out CSV files in partitions filepath_expr = filepath.replace('.csv', '*.csv') filepaths = glob.glob(filepath_expr) return pd.concat((pd.read_csv(f, usecols=usecols) for f in filepaths)) read["csv"] ["pandas"] = lambda filepath,p,filetype: benchmark(read_csv_pandas, (filepath, Kwargs(usecols=p.columns)), filetype) read["h5"] ["pandas"] = lambda filepath,p,filetype: benchmark(pd.read_hdf, (filepath, p.base, Kwargs(columns=p.columns)), filetype) read["feather"] ["pandas"] = lambda filepath,p,filetype: benchmark(feather.read_dataframe, (filepath,), filetype) def read_parq_pandas(filepath): return fp.ParquetFile(filepath).to_pandas() read["parq"] ["pandas"] = lambda filepath,p,filetype: benchmark(read_parq_pandas, (filepath,), filetype) read["gz.parq"] ["pandas"] = lambda filepath,p,filetype: benchmark(read_parq_pandas, (filepath,), filetype) read["snappy.parq"] ["pandas"] = lambda filepath,p,filetype: benchmark(read_parq_pandas, (filepath,), filetype) write = dict([(f, dict()) for f in ["parq","snappy.parq","gz.parq","feather","h5","csv"]]) write["csv"] ["dask"] = lambda df,filepath,p: benchmark(df.to_csv, (filepath.replace(".csv","*.csv"), Kwargs(index=False))) write["h5"] ["dask"] = lambda df,filepath,p: benchmark(df.to_hdf, (filepath, p.base)) def write_feather_dask(filepath, df): return feather.write_dataframe(df.compute(), filepath) write["feather"] ["dask"] = lambda df,filepath,p: benchmark(write_feather_dask, (filepath, df)) write["parq"] ["dask"] = lambda df,filepath,p: benchmark(dd.to_parquet, (filepath, df)) # **p.parq_opts write["snappy.parq"] ["dask"] = lambda df,filepath,p: benchmark(dd.to_parquet, (filepath, df, Kwargs(compression='SNAPPY'))) ## **p.parq_opts write["gz.parq"] ["dask"] = lambda df,filepath,p: benchmark(dd.to_parquet, (filepath, df, Kwargs(compression='GZIP'))) write["csv"] ["pandas"] = lambda df,filepath,p: benchmark(df.to_csv, (filepath, Kwargs(index=False))) write["h5"] ["pandas"] = lambda df,filepath,p: benchmark(df.to_hdf, (filepath, Kwargs(key=p.base, format='table'))) write["feather"] ["pandas"] = lambda df,filepath,p: benchmark(feather.write_dataframe, (df, filepath)) write["parq"] ["pandas"] = lambda df,filepath,p: benchmark(fp.write, (filepath, df, Kwargs(**p.parq_opts))) write["gz.parq"] ["pandas"] = lambda df,filepath,p: benchmark(fp.write, (filepath, df, Kwargs(compression='GZIP', **p.parq_opts))) write["snappy.parq"] ["pandas"] = lambda df,filepath,p: benchmark(fp.write, (filepath, df, Kwargs(compression='SNAPPY', **p.parq_opts))) def timed_write(filepath,dftype,fsize='double',output_directory="times"): """Accepts any file with a dataframe readable by the given dataframe type, and writes it out as a variety of file types""" assert fsize in ('single', 'double') p.dftype = dftype # This function may get called from outside main() df,duration=timed_read(filepath,dftype) for ext in write.keys(): directory,filename = os.path.split(filepath) basename, extension = os.path.splitext(filename) fname = output_directory+os.path.sep+basename+"."+ext if os.path.exists(fname): print("{:28} (keeping existing)".format(fname), flush=True) else: filetype=ext.split(".")[-1] if not filetype in filetypes_storing_categories: for c in p.categories: if filetype == 'parq' and df[c].dtype == 'object': df[c]=df[c].str.encode('utf8') else: df[c]=df[c].astype(str) # Convert doubles to floats when writing out datasets if fsize == 'single': for colname in df.columns: if df[colname].dtype == 'float64': df[colname] = df[colname].astype(np.float32) code = write[ext].get(dftype,None) if code is None: print("{:28} {:7} Operation not supported".format(fname,dftype), flush=True) else: duration, res = code(df,fname,p) print("{:28} {:7} {:05.2f}".format(fname,dftype,duration), flush=True) if not filetype in filetypes_storing_categories: for c in p.categories: df[c]=df[c].astype('category') def timed_read(filepath,dftype): basename, extension = os.path.splitext(filepath) extension = extension[1:] filetype=extension.split(".")[-1] code = read[extension].get(dftype,None) if code is None: return (None, -1) p.columns=[p.x]+[p.y]+p.categories duration, df = code(filepath,p,filetype) return df, duration CACHED_RANGES = (None, None) def timed_agg(df, filepath, plot_width=int(900), plot_height=int(900*7.0/12), cache_ranges=True): global CACHED_RANGES start = time.time() cvs = ds.Canvas(plot_width, plot_height, x_range=CACHED_RANGES[0], y_range=CACHED_RANGES[1]) agg = cvs.points(df, p.x, p.y) end = time.time() if cache_ranges: CACHED_RANGES = (cvs.x_range, cvs.y_range) img = export_image(tf.shade(agg),filepath,export_path=".") return img, end-start def get_size(path): total = 0 # CSV files are broken up by dask when they're written out if os.path.isfile(path): return os.path.getsize(path) elif path.endswith('csv'): for csv_fpath in glob.glob(path.replace('.csv', '*.csv')): total += os.path.getsize(csv_fpath) return total # If path is a directory (such as parquet), sum all files in directory for dirpath, dirnames, filenames in os.walk(path): for f in filenames: fp = os.path.join(dirpath, f) total += os.path.getsize(fp) return total def get_proc_mem(): return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 def main(argv): global DEBUG, DD_FORCE_LOAD, DASK_CLIENT parser = argparse.ArgumentParser(epilog=__doc__, formatter_class=argparse.RawTextHelpFormatter) parser.add_argument('filepath') parser.add_argument('dftype') parser.add_argument('base') parser.add_argument('x') parser.add_argument('y') parser.add_argument('categories', nargs='+') parser.add_argument('--debug', action='store_true', help='Enable increased verbosity and DEBUG messages') parser.add_argument('--cache', choices=('persist', 'cachey'), default=None, help='Enable caching: "persist" causes Dask dataframes to force loading into memory; "cachey" uses dask.cache.Cache with a cachesize of {}. Caching is disabled by default'.format(int(p.cachesize))) parser.add_argument('--distributed', action='store_true', help='Enable the distributed scheduler instead of the threaded, which is the default.') parser.add_argument('--recalc-ranges', action='store_true', help='Tell datashader to recalculate the ranges on each aggregation, instead of caching them (by default).') args = parser.parse_args(argv[1:]) if args.cache is None: if args.debug: print("DEBUG: Cache disabled", flush=True) else: if args.cache == 'cachey': from dask.cache import Cache cache = Cache(p.cachesize) cache.register() elif args.cache == 'persist': DD_FORCE_LOAD = True if args.debug: print('DEBUG: Cache "{}" mode enabled'.format(args.cache), flush=True) if args.dftype == 'dask' and args.distributed: local_cluster = distributed.LocalCluster(n_workers=p.n_workers, threads_per_worker=1) DASK_CLIENT = distributed.Client(local_cluster) if args.debug: print('DEBUG: "distributed" scheduler is enabled') else: if args.dftype != 'dask' and args.distributed: raise ValueError('--distributed argument is only available with the dask dataframe type (not pandas)') if args.debug: print('DEBUG: "threaded" scheduler is enabled') filepath = args.filepath basename, extension = os.path.splitext(filepath) p.dftype = args.dftype p.base = args.base p.x = args.x p.y = args.y p.categories = args.categories DEBUG = args.debug if DEBUG: print('DEBUG: Memory usage (before read):\t{} MB'.format(get_proc_mem()), flush=True) df,loadtime = timed_read(filepath, p.dftype) if df is None: if loadtime == -1: print("{:28} {:6} Operation not supported".format(filepath, p.dftype), flush=True) return 1 if DEBUG: print('DEBUG: Memory usage (after read):\t{} MB'.format(get_proc_mem()), flush=True) img,aggtime1 = timed_agg(df,filepath,5,5,cache_ranges=(not args.recalc_ranges)) if DEBUG: mem_usage = df.memory_usage(deep=True) if p.dftype == 'dask': mem_usage = mem_usage.compute() print('DEBUG:', mem_usage, flush=True) mem_usage_total = mem_usage.sum() print('DEBUG: DataFrame size:\t\t\t{} MB'.format(mem_usage_total / 1e6), flush=True) for colname in df.columns: print('DEBUG: column "{}" dtype: {}'.format(colname, df[colname].dtype)) print('DEBUG: Memory usage (after agg1):\t{} MB'.format(get_proc_mem()), flush=True) img,aggtime2 = timed_agg(df,filepath,cache_ranges=(not args.recalc_ranges)) if DEBUG: print('DEBUG: Memory usage (after agg2):\t{} MB'.format(get_proc_mem()), flush=True) in_size = get_size(filepath) out_size = get_size(filepath+".png") global_end = time.time() print("{:28} {:6} Aggregate1:{:06.2f} ({:06.2f}+{:06.2f}) Aggregate2:{:06.2f} In:{:011d} Out:{:011d} Total:{:06.2f}"\ .format(filepath, p.dftype, loadtime+aggtime1, loadtime, aggtime1, aggtime2, in_size, out_size, global_end-global_start), flush=True) return 0 if __name__ == '__main__': sys.exit(main(sys.argv))