<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import , division , print_function <EOL> import numbers <EOL> import fnmatch <EOL> from operator import attrgetter <EOL> from pymongo . collection import Collection <EOL> from pymongo . database import Database <EOL> from datashape . predicates import isscalar <EOL> from toolz import pluck , first , get , compose <EOL> import toolz <EOL> import datetime <EOL> from . . expr import ( <EOL> And , <EOL> Arithmetic , <EOL> Broadcast , <EOL> By , <EOL> Distinct , <EOL> Eq , <EOL> Expr , <EOL> Field , <EOL> FloorDiv , <EOL> Ge , <EOL> Gt , <EOL> Head , <EOL> Le , <EOL> Lt , <EOL> Ne , <EOL> Or , <EOL> Projection , <EOL> Reduction , <EOL> SimpleSelection , <EOL> Sort , <EOL> Summary , <EOL> Like , <EOL> Symbol , <EOL> ceil , <EOL> count , <EOL> floor , <EOL> math , <EOL> max , <EOL> mean , <EOL> min , <EOL> nelements , <EOL> nunique , <EOL> sum , <EOL> symbol , <EOL> ) <EOL> from . . expr . broadcast import broadcast_collect , Broadcastable <EOL> from . . expr . datetime import Day , Month , Year , Minute , Second , UTCFromTimestamp <EOL> from . . expr . optimize import simple_selections <EOL> from . . compatibility import _strtypes <EOL> from . . dispatch import dispatch <EOL> __all__ = [ '<STR_LIT>' ] <EOL> @ dispatch ( Expr , Collection ) <EOL> def pre_compute ( expr , data , scope = None , ** kwargs ) : <EOL> return MongoQuery ( data , [ ] ) <EOL> @ dispatch ( Expr , Database ) <EOL> def pre_compute ( expr , data , ** kwargs ) : <EOL> return data <EOL> class MongoQuery ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , coll , query ) : <EOL> self . coll = coll <EOL> self . query = tuple ( query ) <EOL> def append ( self , clause ) : <EOL> return MongoQuery ( self . coll , self . query + ( clause , ) ) <EOL> def info ( self ) : <EOL> return self . coll , self . query <EOL> def __eq__ ( self , other ) : <EOL> return type ( self ) == type ( other ) and self . info ( ) == other . info ( ) <EOL> def __hash__ ( self ) : <EOL> return hash ( ( type ( self ) , self . info ( ) ) ) <EOL> @ dispatch ( Expr , ( MongoQuery , Collection ) ) <EOL> def optimize ( expr , seq ) : <EOL> return broadcast_collect ( simple_selections ( expr ) , no_recurse = SimpleSelection ) <EOL> @ dispatch ( Head , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> return q . append ( { '<STR_LIT>' : t . n } ) <EOL> @ dispatch ( Broadcast , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> s = t . _scalars [ <NUM_LIT:0> ] <EOL> d = dict ( ( s [ c ] , symbol ( c , s [ c ] . dshape . measure ) ) for c in s . fields ) <EOL> expr = t . _scalar_expr . _subs ( d ) <EOL> name = expr . _name or '<STR_LIT>' % abs ( hash ( expr ) ) <EOL> return q . append ( { '<STR_LIT>' : { name : compute_sub ( expr ) } } ) <EOL> binops = { '<STR_LIT:+>' : '<STR_LIT>' , <EOL> '<STR_LIT:*>' : '<STR_LIT>' , <EOL> '<STR_LIT:/>' : '<STR_LIT>' , <EOL> '<STR_LIT:->' : '<STR_LIT>' , <EOL> '<STR_LIT:%>' : '<STR_LIT>' } <EOL> def compute_sub ( t ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( t , _strtypes + ( Symbol , ) ) : <EOL> return '<STR_LIT>' % t <EOL> elif isinstance ( t , numbers . Number ) : <EOL> return t <EOL> elif isinstance ( t , FloorDiv ) : <EOL> return compute_sub ( floor ( t . lhs / t . rhs ) ) <EOL> elif isinstance ( t , Arithmetic ) and hasattr ( t , '<STR_LIT>' ) and t . symbol in binops : <EOL> op = binops [ t . symbol ] <EOL> return { '<STR_LIT>' % op : [ compute_sub ( t . lhs ) , compute_sub ( t . rhs ) ] } <EOL> elif isinstance ( t , floor ) : <EOL> x = compute_sub ( t . _child ) <EOL> return { '<STR_LIT>' : [ x , { '<STR_LIT>' : [ x , <NUM_LIT:1> ] } ] } <EOL> elif isinstance ( t , math . abs ) : <EOL> x = compute_sub ( t . _child ) <EOL> return { '<STR_LIT>' : [ { '<STR_LIT>' : [ x , <NUM_LIT:0> ] } , <EOL> { '<STR_LIT>' : [ <NUM_LIT:0> , x ] } , <EOL> x ] } <EOL> elif isinstance ( t , ceil ) : <EOL> x = compute_sub ( t . _child ) <EOL> return { '<STR_LIT>' : [ x , { '<STR_LIT>' : [ <NUM_LIT:1> , { '<STR_LIT>' : [ x , <NUM_LIT:1> ] } ] } ] } <EOL> elif isinstance ( t , tuple ( datetime_terms ) ) : <EOL> op = datetime_terms [ type ( t ) ] <EOL> return { '<STR_LIT>' % op : compute_sub ( t . _child ) } <EOL> elif isinstance ( t , UTCFromTimestamp ) : <EOL> return { '<STR_LIT>' : [ datetime . datetime . utcfromtimestamp ( <NUM_LIT:0> ) , <EOL> { '<STR_LIT>' : [ <NUM_LIT:1000> , compute_sub ( t . _child ) ] } ] } <EOL> raise NotImplementedError ( '<STR_LIT>' % type ( t ) . __name__ ) <EOL> @ dispatch ( ( Projection , Field ) , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> return q . append ( { '<STR_LIT>' : dict ( ( col , <NUM_LIT:1> ) for col in t . fields ) } ) <EOL> @ dispatch ( SimpleSelection , MongoQuery ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> predicate = optimize ( expr . predicate , data ) <EOL> if not isinstance ( predicate , Broadcast ) : <EOL> raise TypeError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( type ( predicate ) . __name__ , <EOL> '<STR_LIT:U+002CU+0020>' . join ( map ( compose ( repr , attrgetter ( '<STR_LIT>' ) ) , <EOL> Broadcastable ) ) ) ) <EOL> s = predicate . _scalars [ <NUM_LIT:0> ] <EOL> d = dict ( ( s [ c ] , symbol ( c , s [ c ] . dshape . measure ) ) for c in s . fields ) <EOL> expr = predicate . _scalar_expr . _subs ( d ) <EOL> return data . append ( { '<STR_LIT>' : match ( expr ) } ) <EOL> @ dispatch ( By , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> if not isinstance ( t . grouper , ( Field , Projection , Symbol ) ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % t . grouper ) <EOL> apply = optimize ( t . apply , q ) <EOL> names = apply . fields <EOL> return MongoQuery ( q . coll , q . query + <EOL> ( { <EOL> '<STR_LIT>' : toolz . merge ( <EOL> { '<STR_LIT>' : dict ( ( col , '<STR_LIT:$>' + col ) for col in t . grouper . fields ) } , <EOL> group_apply ( apply ) <EOL> ) <EOL> } , <EOL> { <EOL> '<STR_LIT>' : toolz . merge ( dict ( ( col , '<STR_LIT>' + col ) for col in t . grouper . fields ) , <EOL> dict ( ( name , '<STR_LIT:$>' + name ) for name in names ) ) <EOL> } ) ) <EOL> @ dispatch ( nunique , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> return MongoQuery ( q . coll , q . query + <EOL> ( { '<STR_LIT>' : { '<STR_LIT>' : dict ( ( col , '<STR_LIT:$>' + col ) for col in t . fields ) , <EOL> t . _name : { '<STR_LIT>' : <NUM_LIT:1> } } } , <EOL> { '<STR_LIT>' : { '<STR_LIT>' : <NUM_LIT:0> , t . _name : <NUM_LIT:1> } } ) ) <EOL> @ dispatch ( Distinct , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> return MongoQuery ( q . coll , q . query + <EOL> ( { '<STR_LIT>' : { '<STR_LIT>' : dict ( ( col , '<STR_LIT:$>' + col ) for col in t . fields ) } } , <EOL> { '<STR_LIT>' : toolz . merge ( dict ( ( col , '<STR_LIT>' + col ) for col in t . fields ) , <EOL> { '<STR_LIT>' : <NUM_LIT:0> } ) } ) ) <EOL> @ dispatch ( Reduction ) <EOL> def group_apply ( expr ) : <EOL> """<STR_LIT>""" <EOL> key = expr . _name <EOL> col = '<STR_LIT:$>' + expr . _child . _name <EOL> if isinstance ( expr , count ) : <EOL> return { key : { '<STR_LIT>' : <NUM_LIT:1> } } <EOL> if isinstance ( expr , sum ) : <EOL> return { key : { '<STR_LIT>' : col } } <EOL> if isinstance ( expr , max ) : <EOL> return { key : { '<STR_LIT>' : col } } <EOL> if isinstance ( expr , min ) : <EOL> return { key : { '<STR_LIT>' : col } } <EOL> if isinstance ( expr , mean ) : <EOL> return { key : { '<STR_LIT>' : col } } <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> % type ( expr ) . __name__ ) <EOL> reductions = { mean : '<STR_LIT>' , count : '<STR_LIT>' , max : '<STR_LIT>' , min : '<STR_LIT>' , sum : '<STR_LIT>' } <EOL> def scalar_expr ( expr ) : <EOL> if isinstance ( expr , Broadcast ) : <EOL> s = expr . _scalars [ <NUM_LIT:0> ] <EOL> d = dict ( ( s [ c ] , symbol ( c , s [ c ] . dshape . measure ) ) for c in s . fields ) <EOL> return expr . _scalar_expr . _subs ( d ) <EOL> elif isinstance ( expr , Field ) : <EOL> return symbol ( expr . _name , expr . dshape . measure ) <EOL> else : <EOL> expr2 = broadcast_collect ( expr ) <EOL> if not expr2 . isidentical ( expr ) : <EOL> return scalar_expr ( expr2 ) <EOL> @ dispatch ( Summary ) <EOL> def group_apply ( expr ) : <EOL> reducs = expr . values <EOL> names = expr . fields <EOL> values = [ ( name , c , getattr ( c . _child , '<STR_LIT>' , None ) or name ) <EOL> for name , c in zip ( names , reducs ) ] <EOL> query = dict ( ( k , { reductions [ type ( v ) ] : <NUM_LIT:1> if isinstance ( v , count ) <EOL> else compute_sub ( scalar_expr ( v . _child ) ) } ) <EOL> for k , v , z in values ) <EOL> return query <EOL> @ dispatch ( ( count , nelements ) , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> name = t . _name <EOL> return q . append ( { '<STR_LIT>' : { '<STR_LIT>' : { } , name : { '<STR_LIT>' : <NUM_LIT:1> } } } ) <EOL> @ dispatch ( ( sum , min , max , mean ) , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> name = t . _name <EOL> reduction = { sum : '<STR_LIT>' , min : '<STR_LIT>' , max : '<STR_LIT>' , mean : '<STR_LIT>' } [ type ( t ) ] <EOL> column = '<STR_LIT:$>' + t . _child . _name <EOL> arg = { '<STR_LIT>' : { '<STR_LIT>' : { } , name : { reduction : column } } } <EOL> return q . append ( arg ) <EOL> @ dispatch ( Sort , MongoQuery ) <EOL> def compute_up ( t , q , ** kwargs ) : <EOL> return q . append ( { '<STR_LIT>' : { t . key : <NUM_LIT:1> if t . ascending else - <NUM_LIT:1> } } ) <EOL> datetime_terms = { Day : '<STR_LIT>' , <EOL> Month : '<STR_LIT>' , <EOL> Year : '<STR_LIT>' , <EOL> Minute : '<STR_LIT>' , <EOL> Second : '<STR_LIT>' } <EOL> @ dispatch ( Field , Database ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> return getattr ( data , expr . _name ) <EOL> @ dispatch ( Expr , Collection ) <EOL> def post_compute ( e , c , scope = None ) : <EOL> """<STR_LIT>""" <EOL> return post_compute ( e , MongoQuery ( c , ( ) ) , scope = scope ) <EOL> def get_result ( result ) : <EOL> try : <EOL> return result [ '<STR_LIT:result>' ] <EOL> except TypeError : <EOL> return list ( result ) <EOL> @ dispatch ( Expr , MongoQuery ) <EOL> def post_compute ( e , q , scope = None ) : <EOL> """<STR_LIT>""" <EOL> scope = { '<STR_LIT>' : toolz . merge ( { '<STR_LIT>' : <NUM_LIT:0> } , <EOL> dict ( ( col , <NUM_LIT:1> ) for col in e . fields ) ) } <EOL> q = q . append ( scope ) <EOL> if not e . dshape . shape : <EOL> result = get_result ( q . coll . aggregate ( list ( q . query ) ) ) [ <NUM_LIT:0> ] <EOL> if isscalar ( e . dshape . measure ) : <EOL> return result [ e . _name ] <EOL> else : <EOL> return get ( e . fields , result ) <EOL> dicts = get_result ( q . coll . aggregate ( list ( q . query ) ) ) <EOL> if isscalar ( e . dshape . measure ) : <EOL> return list ( pluck ( e . fields [ <NUM_LIT:0> ] , dicts , default = None ) ) <EOL> else : <EOL> return list ( pluck ( e . fields , dicts , default = None ) ) <EOL> @ dispatch ( Broadcast , MongoQuery ) <EOL> def post_compute ( e , q , scope = None ) : <EOL> """<STR_LIT>""" <EOL> columns = dict ( ( col , <NUM_LIT:1> ) for qry in q . query <EOL> for col in qry . get ( '<STR_LIT>' , [ ] ) ) <EOL> scope = { '<STR_LIT>' : toolz . merge ( { '<STR_LIT>' : <NUM_LIT:0> } , <EOL> dict ( ( col , <NUM_LIT:1> ) for col in columns ) ) } <EOL> q = q . append ( scope ) <EOL> dicts = get_result ( q . coll . aggregate ( list ( q . query ) ) ) <EOL> assert len ( columns ) == <NUM_LIT:1> <EOL> return list ( pluck ( first ( columns . keys ( ) ) , dicts ) ) <EOL> def name ( e ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( e , Symbol ) : <EOL> return e . _name <EOL> elif isinstance ( e , Expr ) : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> else : <EOL> return e <EOL> opnames = { <EOL> Lt : '<STR_LIT>' , <EOL> Le : '<STR_LIT>' , <EOL> Gt : '<STR_LIT>' , <EOL> Ge : '<STR_LIT>' , <EOL> Ne : '<STR_LIT>' , <EOL> And : '<STR_LIT>' , <EOL> Or : '<STR_LIT>' , <EOL> Like : '<STR_LIT>' , <EOL> } <EOL> @ dispatch ( ( Lt , Le , Gt , Ge , Ne ) ) <EOL> def match ( expr ) : <EOL> return { name ( expr . lhs ) : { opnames [ type ( expr ) ] : expr . rhs } } <EOL> @ dispatch ( ( And , Or ) ) <EOL> def match ( expr ) : <EOL> return { opnames [ type ( expr ) ] : [ match ( expr . lhs ) , match ( expr . rhs ) ] } <EOL> @ dispatch ( Eq ) <EOL> def match ( expr ) : <EOL> return { name ( expr . lhs ) : name ( expr . rhs ) } <EOL> @ dispatch ( Like ) <EOL> def match ( expr ) : <EOL> return { <EOL> name ( expr . _child ) : { <EOL> opnames [ type ( expr ) ] : fnmatch . translate ( expr . pattern ) <EOL> } <EOL> } <EOL> @ dispatch ( object ) <EOL> def match ( expr ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( <EOL> '<STR_LIT>' % type ( expr ) . __name__ <EOL> ) </s>
