<s> """<STR_LIT>""" <EOL> from __future__ import with_statement <EOL> import gc <EOL> import heapq <EOL> import logging <EOL> import time <EOL> from appengine_pipeline . src import pipeline <EOL> from appengine_pipeline . src . pipeline import common as pipeline_common <EOL> from google . appengine . api import files <EOL> from google . appengine . api . files import file_service_pb <EOL> from google . appengine . api . files import records <EOL> from google . appengine . ext import db <EOL> from google . appengine . ext . mapreduce import base_handler <EOL> from google . appengine . ext . mapreduce import context <EOL> from google . appengine . ext . mapreduce import errors <EOL> from google . appengine . ext . mapreduce import input_readers <EOL> from google . appengine . ext . mapreduce import mapper_pipeline <EOL> from google . appengine . ext . mapreduce import operation <EOL> from google . appengine . ext . mapreduce import output_writers <EOL> class _OutputFile ( db . Model ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def kind ( cls ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" <EOL> @ classmethod <EOL> def get_root_key ( cls , job_id ) : <EOL> """<STR_LIT>""" <EOL> return db . Key . from_path ( cls . kind ( ) , job_id ) <EOL> def _compare_keys ( key_record1 , key_record2 ) : <EOL> """<STR_LIT>""" <EOL> return cmp ( key_record1 [ <NUM_LIT:0> ] , key_record2 [ <NUM_LIT:0> ] ) <EOL> class _BatchRecordsReader ( input_readers . RecordsReader ) : <EOL> """<STR_LIT>""" <EOL> BATCH_SIZE = <NUM_LIT> * <NUM_LIT> * <NUM_LIT:3> <EOL> def __iter__ ( self ) : <EOL> records = [ ] <EOL> size = <NUM_LIT:0> <EOL> for record in input_readers . RecordsReader . __iter__ ( self ) : <EOL> records . append ( record ) <EOL> size += len ( record ) <EOL> if size > self . BATCH_SIZE : <EOL> yield records <EOL> size = <NUM_LIT:0> <EOL> records = [ ] <EOL> gc . collect ( ) <EOL> if records : <EOL> yield records <EOL> records = [ ] <EOL> gc . collect ( ) <EOL> def _sort_records_map ( records ) : <EOL> """<STR_LIT>""" <EOL> ctx = context . get ( ) <EOL> l = len ( records ) <EOL> key_records = [ None ] * l <EOL> logging . debug ( "<STR_LIT>" ) <EOL> for i in range ( l ) : <EOL> proto = file_service_pb . KeyValue ( ) <EOL> proto . ParseFromString ( records [ i ] ) <EOL> key_records [ i ] = ( proto . key ( ) , records [ i ] ) <EOL> logging . debug ( "<STR_LIT>" ) <EOL> key_records . sort ( cmp = _compare_keys ) <EOL> logging . debug ( "<STR_LIT>" ) <EOL> blob_file_name = ( ctx . mapreduce_spec . name + "<STR_LIT:->" + <EOL> ctx . mapreduce_id + "<STR_LIT>" ) <EOL> output_path = files . blobstore . create ( <EOL> _blobinfo_uploaded_filename = blob_file_name ) <EOL> with output_writers . RecordsPool ( output_path , ctx = ctx ) as pool : <EOL> for key_record in key_records : <EOL> pool . append ( key_record [ <NUM_LIT:1> ] ) <EOL> logging . debug ( "<STR_LIT>" ) <EOL> files . finalize ( output_path ) <EOL> output_path = files . blobstore . get_file_name ( <EOL> files . blobstore . get_blob_key ( output_path ) ) <EOL> entity = _OutputFile ( key_name = output_path , <EOL> parent = _OutputFile . get_root_key ( ctx . mapreduce_id ) ) <EOL> entity . put ( ) <EOL> class _SortChunksPipeline ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , job_name , filenames ) : <EOL> sort_mappers = [ ] <EOL> for i in range ( len ( filenames ) ) : <EOL> filename = filenames [ i ] <EOL> sort_mapper = yield mapper_pipeline . MapperPipeline ( <EOL> "<STR_LIT>" % ( job_name , str ( i ) ) , <EOL> __name__ + "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> None , <EOL> { <EOL> "<STR_LIT>" : [ filename ] , <EOL> "<STR_LIT>" : <NUM_LIT> , <EOL> } , <EOL> shards = <NUM_LIT:1> ) <EOL> sort_mappers . append ( sort_mapper ) <EOL> with pipeline . After ( * sort_mappers ) : <EOL> job_ids = yield pipeline_common . Append ( * [ mapper . job_id for mapper in <EOL> sort_mappers ] ) <EOL> result = yield _CollectOutputFiles ( job_ids ) <EOL> with pipeline . After ( result ) : <EOL> yield _CleanupOutputFiles ( job_ids ) <EOL> yield pipeline_common . Return ( result ) <EOL> class _CollectOutputFiles ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , job_ids ) : <EOL> result = [ ] <EOL> for job_id in job_ids : <EOL> entities = _OutputFile . all ( ) . ancestor ( _OutputFile . get_root_key ( job_id ) ) <EOL> result . append ( [ entity . key ( ) . name ( ) for entity in entities ] ) <EOL> return result <EOL> class _CleanupOutputFiles ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , job_ids ) : <EOL> result = [ ] <EOL> for job_id in job_ids : <EOL> db . delete ( _OutputFile . all ( ) . ancestor ( _OutputFile . get_root_key ( job_id ) ) ) <EOL> class _MergingReader ( input_readers . InputReader ) : <EOL> """<STR_LIT>""" <EOL> expand_parameters = True <EOL> FILES_PARAM = "<STR_LIT>" <EOL> def __init__ ( self , offsets ) : <EOL> """<STR_LIT>""" <EOL> self . _offsets = offsets <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> ctx = context . get ( ) <EOL> mapper_spec = ctx . mapreduce_spec . mapper <EOL> shard_number = ctx . shard_state . shard_number <EOL> filenames = mapper_spec . params [ self . FILES_PARAM ] [ shard_number ] <EOL> if len ( filenames ) != len ( self . _offsets ) : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> readers = [ ] <EOL> for ( i , filename ) in enumerate ( filenames ) : <EOL> offset = self . _offsets [ i ] <EOL> reader = records . RecordsReader ( files . BufferedFile ( filename ) ) <EOL> reader . seek ( offset ) <EOL> readers . append ( ( None , None , i , reader ) ) <EOL> current_result = None <EOL> while readers : <EOL> ( key , value , index , reader ) = readers [ <NUM_LIT:0> ] <EOL> if key is not None : <EOL> if current_result and key != current_result [ <NUM_LIT:0> ] : <EOL> yield current_result <EOL> if not current_result or key != current_result [ <NUM_LIT:0> ] : <EOL> current_result = ( key , [ ] ) <EOL> current_result [ <NUM_LIT:1> ] . append ( value ) <EOL> try : <EOL> self . _offsets [ index ] = reader . tell ( ) <EOL> start_time = time . time ( ) <EOL> binary_record = reader . read ( ) <EOL> if context . get ( ) : <EOL> operation . counters . Increment ( <EOL> input_readers . COUNTER_IO_READ_BYTES , <EOL> len ( binary_record ) ) ( context . get ( ) ) <EOL> operation . counters . Increment ( <EOL> input_readers . COUNTER_IO_READ_MSEC , <EOL> int ( ( time . time ( ) - start_time ) * <NUM_LIT:1000> ) ) ( context . get ( ) ) <EOL> proto = file_service_pb . KeyValue ( ) <EOL> proto . ParseFromString ( binary_record ) <EOL> heapq . heapreplace ( readers , <EOL> ( proto . key ( ) , proto . value ( ) , index , reader ) ) <EOL> except EOFError : <EOL> heapq . heappop ( readers ) <EOL> if current_result : <EOL> yield current_result <EOL> @ classmethod <EOL> def from_json ( cls , json ) : <EOL> """<STR_LIT>""" <EOL> return cls ( json [ "<STR_LIT>" ] ) <EOL> def to_json ( self ) : <EOL> """<STR_LIT>""" <EOL> return { "<STR_LIT>" : self . _offsets } <EOL> @ classmethod <EOL> def split_input ( cls , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> filelists = mapper_spec . params [ cls . FILES_PARAM ] <EOL> return [ cls ( [ <NUM_LIT:0> ] * len ( files ) ) for files in filelists ] <EOL> @ classmethod <EOL> def validate ( cls , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> if mapper_spec . input_reader_class ( ) != cls : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" ) <EOL> params = mapper_spec . params <EOL> if not cls . FILES_PARAM in params : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" ) <EOL> class _HashingBlobstoreOutputWriter ( output_writers . BlobstoreOutputWriterBase ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , filenames ) : <EOL> """<STR_LIT>""" <EOL> self . _filenames = filenames <EOL> @ classmethod <EOL> def validate ( cls , mapper_spec ) : <EOL> """<STR_LIT>""" <EOL> if mapper_spec . output_writer_class ( ) != cls : <EOL> raise errors . BadWriterParamsError ( "<STR_LIT>" ) <EOL> @ classmethod <EOL> def init_job ( cls , mapreduce_state ) : <EOL> """<STR_LIT>""" <EOL> shards = mapreduce_state . mapreduce_spec . mapper . shard_count <EOL> filenames = [ ] <EOL> for i in range ( shards ) : <EOL> blob_file_name = ( mapreduce_state . mapreduce_spec . name + <EOL> "<STR_LIT:->" + mapreduce_state . mapreduce_spec . mapreduce_id + <EOL> "<STR_LIT>" + str ( i ) ) <EOL> filenames . append ( <EOL> files . blobstore . create ( <EOL> _blobinfo_uploaded_filename = blob_file_name ) ) <EOL> mapreduce_state . writer_state = { "<STR_LIT>" : filenames } <EOL> @ classmethod <EOL> def finalize_job ( cls , mapreduce_state ) : <EOL> """<STR_LIT>""" <EOL> finalized_filenames = [ ] <EOL> for filename in mapreduce_state . writer_state [ "<STR_LIT>" ] : <EOL> files . finalize ( filename ) <EOL> finalized_filenames . append ( <EOL> files . blobstore . get_file_name ( <EOL> files . blobstore . get_blob_key ( filename ) ) ) <EOL> mapreduce_state . writer_state = { "<STR_LIT>" : finalized_filenames } <EOL> @ classmethod <EOL> def from_json ( cls , json ) : <EOL> """<STR_LIT>""" <EOL> return cls ( json [ "<STR_LIT>" ] ) <EOL> def to_json ( self ) : <EOL> """<STR_LIT>""" <EOL> return { "<STR_LIT>" : self . _filenames } <EOL> @ classmethod <EOL> def create ( cls , mapreduce_state , shard_number ) : <EOL> """<STR_LIT>""" <EOL> return cls ( mapreduce_state . writer_state [ "<STR_LIT>" ] ) <EOL> @ classmethod <EOL> def get_filenames ( cls , mapreduce_state ) : <EOL> """<STR_LIT>""" <EOL> return mapreduce_state . writer_state [ "<STR_LIT>" ] <EOL> def write ( self , data , ctx ) : <EOL> """<STR_LIT>""" <EOL> if len ( data ) != <NUM_LIT:2> : <EOL> logging . error ( "<STR_LIT>" , <EOL> len ( data ) , data ) <EOL> try : <EOL> key = str ( data [ <NUM_LIT:0> ] ) <EOL> value = str ( data [ <NUM_LIT:1> ] ) <EOL> except TypeError : <EOL> logging . error ( "<STR_LIT>" , <EOL> data . __class__ . __name__ , data ) <EOL> file_index = key . __hash__ ( ) % len ( self . _filenames ) <EOL> pool_name = "<STR_LIT>" % file_index <EOL> filename = self . _filenames [ file_index ] <EOL> if ctx . get_pool ( pool_name ) is None : <EOL> ctx . register_pool ( pool_name , <EOL> output_writers . RecordsPool ( filename = filename , ctx = ctx ) ) <EOL> proto = file_service_pb . KeyValue ( ) <EOL> proto . set_key ( key ) <EOL> proto . set_value ( value ) <EOL> ctx . get_pool ( pool_name ) . append ( proto . Encode ( ) ) <EOL> def _merge_map ( k , values ) : <EOL> """<STR_LIT>""" <EOL> proto = file_service_pb . KeyValues ( ) <EOL> proto . set_key ( k ) <EOL> proto . value_list ( ) . extend ( values ) <EOL> yield proto . Encode ( ) <EOL> class _MergePipeline ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , job_name , filenames ) : <EOL> yield mapper_pipeline . MapperPipeline ( <EOL> job_name + "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> output_writer_spec = <EOL> output_writers . __name__ + "<STR_LIT>" , <EOL> params = { '<STR_LIT>' : filenames } , <EOL> shards = len ( filenames ) ) <EOL> def _hashing_map ( binary_record ) : <EOL> """<STR_LIT>""" <EOL> proto = file_service_pb . KeyValue ( ) <EOL> proto . ParseFromString ( binary_record ) <EOL> yield ( proto . key ( ) , proto . value ( ) ) <EOL> class _HashPipeline ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , job_name , filenames ) : <EOL> yield mapper_pipeline . MapperPipeline ( <EOL> job_name + "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> input_readers . __name__ + "<STR_LIT>" , <EOL> output_writer_spec = __name__ + "<STR_LIT>" , <EOL> params = { '<STR_LIT>' : filenames } , <EOL> shards = len ( filenames ) ) <EOL> class ShufflePipeline ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , job_name , filenames ) : <EOL> hashed_files = yield _HashPipeline ( job_name , filenames ) <EOL> sorted_files = yield _SortChunksPipeline ( job_name , hashed_files ) <EOL> merged_files = yield _MergePipeline ( job_name , sorted_files ) <EOL> with pipeline . After ( merged_files ) : <EOL> all_temp_files = yield pipeline_common . Extend ( <EOL> hashed_files , sorted_files ) <EOL> yield mapper_pipeline . _CleanupPipeline ( all_temp_files ) <EOL> yield pipeline_common . Return ( merged_files ) </s>
