<s> """<STR_LIT>""" <EOL> import random <EOL> from google . appengine . api import datastore <EOL> from google . appengine . ext import key_range <EOL> from google . appengine . ext . mapreduce import datastore_range_iterators as db_iters <EOL> from google . appengine . ext . mapreduce import errors <EOL> from google . appengine . ext . mapreduce import key_ranges <EOL> from google . appengine . ext . mapreduce import model <EOL> from google . appengine . ext . mapreduce import namespace_range <EOL> from google . appengine . ext . mapreduce . api . map_job import input_reader <EOL> class AbstractDatastoreInputReader ( input_reader . InputReader ) : <EOL> """<STR_LIT>""" <EOL> _BATCH_SIZE = <NUM_LIT:50> <EOL> _MAX_SHARD_COUNT = <NUM_LIT> <EOL> MAX_NAMESPACES_FOR_KEY_SHARD = <NUM_LIT:10> <EOL> _APP_PARAM = "<STR_LIT>" <EOL> NAMESPACE_PARAM = "<STR_LIT>" <EOL> ENTITY_KIND_PARAM = "<STR_LIT>" <EOL> KEYS_ONLY_PARAM = "<STR_LIT>" <EOL> BATCH_SIZE_PARAM = "<STR_LIT>" <EOL> KEY_RANGE_PARAM = "<STR_LIT>" <EOL> FILTERS_PARAM = "<STR_LIT>" <EOL> _KEY_RANGE_ITER_CLS = db_iters . AbstractKeyRangeIterator <EOL> def __init__ ( self , iterator ) : <EOL> """<STR_LIT>""" <EOL> self . _iter = iterator <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> for o in self . _iter : <EOL> yield o <EOL> def __str__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return repr ( self . _iter ) <EOL> def to_json ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _iter . to_json ( ) <EOL> @ classmethod <EOL> def from_json ( cls , state ) : <EOL> """<STR_LIT>""" <EOL> return cls ( db_iters . RangeIteratorFactory . from_json ( state ) ) <EOL> @ classmethod <EOL> def _get_query_spec ( cls , params ) : <EOL> """<STR_LIT>""" <EOL> entity_kind = params [ cls . ENTITY_KIND_PARAM ] <EOL> filters = params . get ( cls . FILTERS_PARAM ) <EOL> app = params . get ( cls . _APP_PARAM ) <EOL> ns = params . get ( cls . NAMESPACE_PARAM ) <EOL> return model . QuerySpec ( <EOL> entity_kind = cls . _get_raw_entity_kind ( entity_kind ) , <EOL> keys_only = bool ( params . get ( cls . KEYS_ONLY_PARAM , False ) ) , <EOL> filters = filters , <EOL> batch_size = int ( params . get ( cls . BATCH_SIZE_PARAM , cls . _BATCH_SIZE ) ) , <EOL> model_class_path = entity_kind , <EOL> app = app , <EOL> ns = ns ) <EOL> @ classmethod <EOL> def split_input ( cls , job_config ) : <EOL> """<STR_LIT>""" <EOL> shard_count = job_config . shard_count <EOL> params = job_config . input_reader_params <EOL> query_spec = cls . _get_query_spec ( params ) <EOL> namespaces = None <EOL> if query_spec . ns is not None : <EOL> k_ranges = cls . _to_key_ranges_by_shard ( <EOL> query_spec . app , [ query_spec . ns ] , shard_count , query_spec ) <EOL> else : <EOL> ns_keys = namespace_range . get_namespace_keys ( <EOL> query_spec . app , cls . MAX_NAMESPACES_FOR_KEY_SHARD + <NUM_LIT:1> ) <EOL> if not ns_keys : <EOL> return <EOL> elif len ( ns_keys ) <= cls . MAX_NAMESPACES_FOR_KEY_SHARD : <EOL> namespaces = [ ns_key . name ( ) or "<STR_LIT>" for ns_key in ns_keys ] <EOL> k_ranges = cls . _to_key_ranges_by_shard ( <EOL> query_spec . app , namespaces , shard_count , query_spec ) <EOL> else : <EOL> ns_ranges = namespace_range . NamespaceRange . split ( n = shard_count , <EOL> contiguous = False , <EOL> can_query = lambda : True , <EOL> _app = query_spec . app ) <EOL> k_ranges = [ key_ranges . KeyRangesFactory . create_from_ns_range ( ns_range ) <EOL> for ns_range in ns_ranges ] <EOL> iters = [ db_iters . RangeIteratorFactory . create_key_ranges_iterator ( <EOL> r , query_spec , cls . _KEY_RANGE_ITER_CLS ) for r in k_ranges ] <EOL> return [ cls ( i ) for i in iters ] <EOL> @ classmethod <EOL> def _to_key_ranges_by_shard ( cls , app , namespaces , shard_count , query_spec ) : <EOL> """<STR_LIT>""" <EOL> key_ranges_by_ns = [ ] <EOL> for namespace in namespaces : <EOL> ranges = cls . _split_ns_by_scatter ( <EOL> shard_count , <EOL> namespace , <EOL> query_spec . entity_kind , <EOL> app ) <EOL> random . shuffle ( ranges ) <EOL> key_ranges_by_ns . append ( ranges ) <EOL> ranges_by_shard = [ [ ] for _ in range ( shard_count ) ] <EOL> for ranges in key_ranges_by_ns : <EOL> for i , k_range in enumerate ( ranges ) : <EOL> if k_range : <EOL> ranges_by_shard [ i ] . append ( k_range ) <EOL> key_ranges_by_shard = [ ] <EOL> for ranges in ranges_by_shard : <EOL> if ranges : <EOL> key_ranges_by_shard . append ( key_ranges . KeyRangesFactory . create_from_list ( <EOL> ranges ) ) <EOL> return key_ranges_by_shard <EOL> @ classmethod <EOL> def _split_ns_by_scatter ( cls , <EOL> shard_count , <EOL> namespace , <EOL> raw_entity_kind , <EOL> app ) : <EOL> """<STR_LIT>""" <EOL> if shard_count == <NUM_LIT:1> : <EOL> return [ key_range . KeyRange ( namespace = namespace , _app = app ) ] <EOL> ds_query = datastore . Query ( kind = raw_entity_kind , <EOL> namespace = namespace , <EOL> _app = app , <EOL> keys_only = True ) <EOL> ds_query . Order ( "<STR_LIT>" ) <EOL> oversampling_factor = <NUM_LIT:32> <EOL> random_keys = ds_query . Get ( shard_count * oversampling_factor ) <EOL> if not random_keys : <EOL> return ( [ key_range . KeyRange ( namespace = namespace , _app = app ) ] + <EOL> [ None ] * ( shard_count - <NUM_LIT:1> ) ) <EOL> random_keys . sort ( ) <EOL> if len ( random_keys ) >= shard_count : <EOL> random_keys = cls . _choose_split_points ( random_keys , shard_count ) <EOL> k_ranges = [ ] <EOL> k_ranges . append ( key_range . KeyRange ( <EOL> key_start = None , <EOL> key_end = random_keys [ <NUM_LIT:0> ] , <EOL> direction = key_range . KeyRange . ASC , <EOL> include_start = False , <EOL> include_end = False , <EOL> namespace = namespace , <EOL> _app = app ) ) <EOL> for i in range ( <NUM_LIT:0> , len ( random_keys ) - <NUM_LIT:1> ) : <EOL> k_ranges . append ( key_range . KeyRange ( <EOL> key_start = random_keys [ i ] , <EOL> key_end = random_keys [ i + <NUM_LIT:1> ] , <EOL> direction = key_range . KeyRange . ASC , <EOL> include_start = True , <EOL> include_end = False , <EOL> namespace = namespace , <EOL> _app = app ) ) <EOL> k_ranges . append ( key_range . KeyRange ( <EOL> key_start = random_keys [ - <NUM_LIT:1> ] , <EOL> key_end = None , <EOL> direction = key_range . KeyRange . ASC , <EOL> include_start = True , <EOL> include_end = False , <EOL> namespace = namespace , <EOL> _app = app ) ) <EOL> if len ( k_ranges ) < shard_count : <EOL> k_ranges += [ None ] * ( shard_count - len ( k_ranges ) ) <EOL> return k_ranges <EOL> @ classmethod <EOL> def _choose_split_points ( cls , sorted_keys , shard_count ) : <EOL> """<STR_LIT>""" <EOL> assert len ( sorted_keys ) >= shard_count <EOL> index_stride = len ( sorted_keys ) / float ( shard_count ) <EOL> return [ sorted_keys [ int ( round ( index_stride * i ) ) ] <EOL> for i in range ( <NUM_LIT:1> , shard_count ) ] <EOL> @ classmethod <EOL> def validate ( cls , job_config ) : <EOL> """<STR_LIT>""" <EOL> super ( AbstractDatastoreInputReader , cls ) . validate ( job_config ) <EOL> params = job_config . input_reader_params <EOL> if cls . ENTITY_KIND_PARAM not in params : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> if cls . BATCH_SIZE_PARAM in params : <EOL> try : <EOL> batch_size = int ( params [ cls . BATCH_SIZE_PARAM ] ) <EOL> if batch_size < <NUM_LIT:1> : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" % batch_size ) <EOL> except ValueError , e : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" % e ) <EOL> try : <EOL> bool ( params . get ( cls . KEYS_ONLY_PARAM , False ) ) <EOL> except : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" <EOL> "<STR_LIT>" , <EOL> params [ cls . KEYS_ONLY_PARAM ] ) <EOL> if cls . NAMESPACE_PARAM in params : <EOL> if not isinstance ( params [ cls . NAMESPACE_PARAM ] , <EOL> ( str , unicode , type ( None ) ) ) : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" ) <EOL> if cls . FILTERS_PARAM in params : <EOL> filters = params [ cls . FILTERS_PARAM ] <EOL> if not isinstance ( filters , list ) : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" ) <EOL> for f in filters : <EOL> if not isinstance ( f , ( tuple , list ) ) : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" <EOL> "<STR_LIT:%s>" , f ) <EOL> if len ( f ) != <NUM_LIT:3> : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" , f ) <EOL> prop , op , _ = f <EOL> if not isinstance ( prop , basestring ) : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" , <EOL> prop ) <EOL> if not isinstance ( op , basestring ) : <EOL> raise errors . BadReaderParamsError ( "<STR_LIT>" , op ) <EOL> @ classmethod <EOL> def _get_raw_entity_kind ( cls , entity_kind_or_model_classpath ) : <EOL> """<STR_LIT>""" <EOL> return entity_kind_or_model_classpath </s>
