<s> from itertools import chain <EOL> import os <EOL> import time <EOL> from multiprocessing import cpu_count <EOL> import sys <EOL> import resource <EOL> import gc <EOL> import pandas <EOL> import numpy as np <EOL> import joblib <EOL> from chatto_transform . transforms . transform_base import Transform <EOL> from chatto_transform . lib . chunks import from_chunks , to_chunks , to_group_chunks , CHUNK_SIZE <EOL> from chatto_transform . lib import temp_file <EOL> from chatto_transform . datastores . hdf_datastore import HdfDataStore <EOL> """<STR_LIT>""" <EOL> class ParallelTransform ( Transform ) : <EOL> def __init__ ( self , transform , group_index = None , chunksize = None , n_jobs = - <NUM_LIT:1> ) : <EOL> self . transform_obj = transform <EOL> self . group_index = group_index <EOL> if n_jobs == - <NUM_LIT:1> : <EOL> n_jobs = cpu_count ( ) <EOL> self . n_jobs = n_jobs <EOL> if chunksize is None : <EOL> chunksize = CHUNK_SIZE <EOL> self . chunksize = chunksize <EOL> def output_schema ( self ) : <EOL> return self . transform_obj . output_schema ( ) <EOL> def input_schema ( self ) : <EOL> return self . transform_obj . input_schema ( ) <EOL> def _transform ( self , data ) : <EOL> start = time . time ( ) <EOL> print ( '<STR_LIT>' , data . memory_usage ( index = True ) . sum ( ) , '<STR_LIT>' ) <EOL> store_chunks_jobs = [ ] <EOL> transform_jobs = [ ] <EOL> hdf_stores = [ ] <EOL> try : <EOL> print ( '<STR_LIT>' ) <EOL> group_iter = self . _group_iter ( data , len ( data ) // self . n_jobs or self . chunksize ) <EOL> for group_data in group_iter : <EOL> if group_data . empty : <EOL> continue <EOL> f = temp_file . make_temporary_file ( ) <EOL> hdf_store = HdfDataStore ( self . input_schema ( ) , f ) <EOL> hdf_stores . append ( hdf_store ) <EOL> hdf_store . store ( group_data ) <EOL> store_chunks_jobs . append ( joblib . delayed ( self . store_chunks_job ) ( hdf_store ) ) <EOL> print ( '<STR_LIT>' ) <EOL> joblib . Parallel ( n_jobs = self . n_jobs ) ( store_chunks_jobs ) <EOL> chunk_stores = chain . from_iterable ( store . chunk_stores ( ) for store in hdf_stores ) <EOL> transform_jobs = [ joblib . delayed ( self . transform_job ) ( chunk_store ) for chunk_store in chunk_stores ] <EOL> print ( '<STR_LIT>' , len ( transform_jobs ) , '<STR_LIT>' ) <EOL> result_hdf_stores = joblib . Parallel ( n_jobs = self . n_jobs ) ( transform_jobs ) <EOL> print ( '<STR_LIT>' ) <EOL> results = from_chunks ( r_hdf_store . load ( ) for r_hdf_store in result_hdf_stores ) <EOL> print ( '<STR_LIT>' ) <EOL> finally : <EOL> for hdf_store in hdf_stores : <EOL> hdf_store . delete_chunks ( ) <EOL> hdf_store . delete ( ) <EOL> end = time . time ( ) <EOL> print ( '<STR_LIT>' , end - start , '<STR_LIT>' ) <EOL> return results <EOL> def _group_iter ( self , data , chunksize ) : <EOL> if self . group_index is not None : <EOL> group_iter = to_group_chunks ( data , self . group_index , chunksize ) <EOL> else : <EOL> group_iter = to_chunks ( data , chunksize ) <EOL> return group_iter <EOL> def store_chunks_job ( self , hdf_store ) : <EOL> data = hdf_store . load ( ) <EOL> chunks = self . _group_iter ( data , self . chunksize ) <EOL> def gc_chunks ( ) : <EOL> for chunk in chunks : <EOL> yield chunk <EOL> del chunk <EOL> gc . collect ( ) <EOL> hdf_store . store_chunks ( gc_chunks ( ) ) <EOL> def transform_job ( self , chunk_store ) : <EOL> gc . collect ( ) <EOL> start = resource . getrusage ( resource . RUSAGE_SELF ) . ru_maxrss / <NUM_LIT:10> ** <NUM_LIT:6> <EOL> data = chunk_store . load ( ) <EOL> gc . collect ( ) <EOL> finished_load = resource . getrusage ( resource . RUSAGE_SELF ) . ru_maxrss / <NUM_LIT:10> ** <NUM_LIT:6> <EOL> result = self . transform_chunk ( data ) <EOL> gc . collect ( ) <EOL> finished_transform = resource . getrusage ( resource . RUSAGE_SELF ) . ru_maxrss / <NUM_LIT:10> ** <NUM_LIT:6> <EOL> t_hdf_store = HdfDataStore ( self . output_schema ( ) , chunk_store . hdf_file ) <EOL> t_hdf_store . store ( result ) <EOL> gc . collect ( ) <EOL> finished_store = resource . getrusage ( resource . RUSAGE_SELF ) . ru_maxrss / <NUM_LIT:10> ** <NUM_LIT:6> <EOL> print ( '<STR_LIT>' , start , '<STR_LIT>' , finished_store , '<STR_LIT>' , finished_store - start ) <EOL> print ( '<STR_LIT>' , finished_load - start , '<STR_LIT>' ) <EOL> print ( '<STR_LIT>' , finished_transform - finished_load , '<STR_LIT>' ) <EOL> print ( '<STR_LIT>' , finished_store - finished_transform , '<STR_LIT>' ) <EOL> return t_hdf_store <EOL> def transform_chunk ( self , data ) : <EOL> print ( '<STR_LIT>' , len ( data ) , '<STR_LIT>' , data . memory_usage ( index = True ) . sum ( ) / <NUM_LIT:10> ** <NUM_LIT:6> , '<STR_LIT>' ) <EOL> transformed_groups = [ ] <EOL> if self . group_index is None : <EOL> return self . transform_obj . transform ( data ) <EOL> for group_id , group_data in data . groupby ( self . group_index , as_index = False , sort = False ) : <EOL> if group_data . empty : <EOL> continue <EOL> transformed_group = self . transform_obj . transform ( group_data ) <EOL> transformed_groups . append ( transformed_group ) <EOL> transformed_data = pandas . concat ( transformed_groups ) <EOL> return transformed_data </s>
