<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import , division , print_function <EOL> from toolz import concat as tconcat <EOL> import datashape <EOL> from datashape . predicates import isscalar , isrecord <EOL> from math import floor <EOL> from . core import * <EOL> from . expressions import * <EOL> from . strings import Like <EOL> from . expressions import shape <EOL> from . math import sqrt <EOL> from . reductions import * <EOL> from . split_apply_combine import * <EOL> from . collections import * <EOL> from . . dispatch import dispatch <EOL> from . . compatibility import builtins <EOL> good_to_split = ( Reduction , Summary , By , Distinct ) <EOL> can_split = good_to_split + ( Like , Selection , ElemWise , Apply ) <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> def path_split ( leaf , expr ) : <EOL> """<STR_LIT>""" <EOL> last = None <EOL> for node in list ( path ( expr , leaf ) ) [ : - <NUM_LIT:1> ] [ : : - <NUM_LIT:1> ] : <EOL> if isinstance ( node , good_to_split ) : <EOL> return node <EOL> elif not isinstance ( node , can_split ) : <EOL> return last <EOL> last = node <EOL> return node <EOL> def split ( leaf , expr , chunk = None , agg = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> center = path_split ( leaf , expr ) <EOL> if chunk is None : <EOL> if leaf . ndim > <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> else : <EOL> chunk = symbol ( '<STR_LIT>' , datashape . var * leaf . dshape . measure ) <EOL> chunk_expr = _split_chunk ( center , leaf = leaf , chunk = chunk , ** kwargs ) <EOL> chunk_expr_with_keepdims = _split_chunk ( center , leaf = leaf , chunk = chunk , <EOL> keepdims = True ) <EOL> if agg is None : <EOL> agg_shape = aggregate_shape ( leaf , expr , chunk , chunk_expr_with_keepdims ) <EOL> agg_dshape = DataShape ( * ( agg_shape + ( chunk_expr . dshape . measure , ) ) ) <EOL> agg = symbol ( '<STR_LIT>' , agg_dshape ) <EOL> agg_expr = _split_agg ( center , leaf = leaf , agg = agg ) <EOL> return ( ( chunk , chunk_expr ) , <EOL> ( agg , expr . _subs ( { center : agg } ) . _subs ( { agg : agg_expr } ) ) ) <EOL> reductions = { sum : ( sum , sum ) , count : ( count , sum ) , <EOL> min : ( min , min ) , max : ( max , max ) , <EOL> any : ( any , any ) , all : ( all , all ) , <EOL> nelements : ( nelements , sum ) } <EOL> @ dispatch ( Expr ) <EOL> def _split ( expr , leaf = None , chunk = None , agg = None , keepdims = True ) : <EOL> return ( ( chunk , _split_chunk ( expr , leaf = leaf , chunk = chunk , <EOL> keepdims = keepdims ) ) , <EOL> ( agg , _split_agg ( expr , leaf = leaf , agg = agg ) ) ) <EOL> @ dispatch ( tuple ( reductions ) ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , keepdims = True ) : <EOL> a , b = reductions [ type ( expr ) ] <EOL> return a ( expr . _subs ( { leaf : chunk } ) . _child , <EOL> keepdims = keepdims , <EOL> axis = expr . axis ) <EOL> @ dispatch ( tuple ( reductions ) ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> a , b = reductions [ type ( expr ) ] <EOL> return b ( agg , axis = expr . axis , keepdims = expr . keepdims ) <EOL> @ dispatch ( mean ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , keepdims = True ) : <EOL> child = expr . _subs ( { leaf : chunk } ) . _child <EOL> return summary ( total = child . sum ( ) , count = child . count ( ) , keepdims = keepdims , <EOL> axis = expr . axis ) <EOL> @ dispatch ( mean ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> total = agg . total . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> count = agg [ '<STR_LIT:count>' ] . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> return total / count <EOL> @ dispatch ( ( std , var ) ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , keepdims = True ) : <EOL> child = expr . _subs ( { leaf : chunk } ) . _child <EOL> return summary ( x = child . sum ( ) , x2 = ( child ** <NUM_LIT:2> ) . sum ( ) , n = child . count ( ) , <EOL> keepdims = keepdims , axis = expr . axis ) <EOL> @ dispatch ( var ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> x = agg . x . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> x2 = agg . x2 . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> n = agg . n . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> result = ( x2 / n ) - ( x / n ) ** <NUM_LIT:2> <EOL> if expr . unbiased : <EOL> result = result / ( n - <NUM_LIT:1> ) * n <EOL> return result <EOL> @ dispatch ( std ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> x = agg . x . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> x2 = agg . x2 . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> n = agg . n . sum ( axis = expr . axis , keepdims = expr . keepdims ) <EOL> result = ( x2 / n ) - ( x / n ) ** <NUM_LIT:2> <EOL> if expr . unbiased : <EOL> result = result / ( n - <NUM_LIT:1> ) * n <EOL> return sqrt ( result ) <EOL> @ dispatch ( Distinct ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , ** kwargs ) : <EOL> return expr . _subs ( { leaf : chunk } ) <EOL> @ dispatch ( Distinct ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> return agg . distinct ( ) <EOL> @ dispatch ( nunique ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , ** kwargs ) : <EOL> return ( expr . _child <EOL> . _subs ( { leaf : chunk } ) <EOL> . distinct ( ) ) <EOL> @ dispatch ( nunique ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> return agg . distinct ( ) . count ( keepdims = expr . keepdims ) <EOL> @ dispatch ( Summary ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , keepdims = True ) : <EOL> exprs = [ ( name , split ( leaf , val , chunk = chunk , <EOL> keepdims = False ) [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] ) <EOL> for name , val in zip ( expr . fields , expr . values ) ] <EOL> d = dict ( ) <EOL> for name , e in exprs : <EOL> if isinstance ( e , Reduction ) : <EOL> d [ name ] = e <EOL> elif isinstance ( e , Summary ) : <EOL> for n , v in zip ( e . names , e . values ) : <EOL> d [ name + '<STR_LIT:_>' + n ] = v <EOL> else : <EOL> raise NotImplementedError ( ) <EOL> return summary ( keepdims = keepdims , ** d ) <EOL> @ dispatch ( Summary ) <EOL> def _split_agg ( expr , leaf = None , chunk = None , agg = None , keepdims = True ) : <EOL> exprs = [ ( name , split ( leaf , val , keepdims = False ) [ <NUM_LIT:1> ] ) <EOL> for name , val in zip ( expr . fields , expr . values ) ] <EOL> d = dict ( ) <EOL> for name , ( a , ae ) in exprs : <EOL> if isscalar ( a . dshape . measure ) : <EOL> d [ name ] = ae . _subs ( { a : agg [ name ] } ) <EOL> elif isrecord ( a . dshape . measure ) : <EOL> names = [ '<STR_LIT>' % ( name , field ) for field in a . fields ] <EOL> namedict = dict ( zip ( a . fields , names ) ) <EOL> d [ name ] = ae . _subs ( toolz . merge ( { a : agg } , namedict ) ) <EOL> return summary ( ** d ) <EOL> @ dispatch ( By ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , ** kwargs ) : <EOL> chunk_apply = _split_chunk ( expr . apply , leaf = leaf , chunk = chunk , keepdims = False ) <EOL> chunk_grouper = expr . grouper . _subs ( { leaf : chunk } ) <EOL> return by ( chunk_grouper , chunk_apply ) <EOL> @ dispatch ( By ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> agg_apply = _split_agg ( expr . apply , leaf = leaf , agg = agg ) <EOL> agg_grouper = expr . grouper . _subs ( { leaf : agg } ) <EOL> ngroup = len ( expr . grouper . fields ) <EOL> if isscalar ( expr . grouper . dshape . measure ) : <EOL> agg_grouper = agg [ agg . fields [ <NUM_LIT:0> ] ] <EOL> else : <EOL> agg_grouper = agg [ list ( agg . fields [ : ngroup ] ) ] <EOL> return ( by ( agg_grouper , agg_apply ) <EOL> . relabel ( dict ( zip ( agg . fields [ : ngroup ] , <EOL> expr . fields [ : ngroup ] ) ) ) ) <EOL> @ dispatch ( ( ElemWise , Like , Selection ) ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , ** kwargs ) : <EOL> return expr . _subs ( { leaf : chunk } ) <EOL> @ dispatch ( ( ElemWise , Like , Selection ) ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> return agg <EOL> @ dispatch ( Apply ) <EOL> def _split_chunk ( expr , leaf = None , chunk = None , ** kwargs ) : <EOL> if expr . _splittable : <EOL> return expr . _subs ( { leaf : chunk } ) <EOL> else : <EOL> raise NotImplementedError ( ) <EOL> @ dispatch ( Apply ) <EOL> def _split_agg ( expr , leaf = None , agg = None ) : <EOL> return agg <EOL> from datashape import Fixed <EOL> from math import ceil <EOL> def dimension_div ( a , b ) : <EOL> """<STR_LIT>""" <EOL> if a == datashape . var or b == datashape . var : <EOL> return datashape . var <EOL> if isinstance ( a , Fixed ) : <EOL> a = int ( a ) <EOL> if isinstance ( b , Fixed ) : <EOL> b = int ( b ) <EOL> return int ( ceil ( a / b ) ) <EOL> def dimension_mul ( a , b ) : <EOL> """<STR_LIT>""" <EOL> if a == datashape . var or b == datashape . var : <EOL> return datashape . var <EOL> if isinstance ( a , Fixed ) : <EOL> a = int ( a ) <EOL> if isinstance ( b , Fixed ) : <EOL> b = int ( b ) <EOL> return int ( a * b ) <EOL> def aggregate_shape ( leaf , expr , chunk , chunk_expr ) : <EOL> """<STR_LIT>""" <EOL> if datashape . var in tconcat ( map ( shape , [ leaf , expr , chunk , chunk_expr ] ) ) : <EOL> return ( datashape . var , ) * leaf . ndim <EOL> numblocks = [ int ( floor ( l / c ) ) for l , c in zip ( leaf . shape , chunk . shape ) ] <EOL> last_chunk_shape = [ l % c for l , c in zip ( leaf . shape , chunk . shape ) ] <EOL> if builtins . sum ( last_chunk_shape ) != <NUM_LIT:0> : <EOL> old = last_chunk_shape <EOL> last_chunk = symbol ( chunk . _name , <EOL> DataShape ( * ( last_chunk_shape + [ chunk . dshape . measure ] ) ) ) <EOL> last_chunk_expr = chunk_expr . _subs ( { chunk : last_chunk } ) <EOL> last_chunk_shape = shape ( last_chunk_expr ) <EOL> last_chunk_shape = tuple ( a if b != <NUM_LIT:0> else <NUM_LIT:0> for a , b in <EOL> zip ( last_chunk_shape , old ) ) <EOL> return tuple ( int ( floor ( l / c ) ) * ce + lce <EOL> for l , c , ce , lce <EOL> in zip ( shape ( leaf ) , shape ( chunk ) , shape ( chunk_expr ) , last_chunk_shape ) ) </s>
