<s> from __future__ import absolute_import , division , print_function <EOL> import datetime <EOL> import numpy as np <EOL> from pandas import DataFrame , Series <EOL> from datashape import to_numpy , to_numpy_dtype <EOL> from numbers import Number <EOL> from . . expr import ( <EOL> Reduction , Field , Projection , Broadcast , Selection , ndim , <EOL> Distinct , Sort , Tail , Head , Label , ReLabel , Expr , Slice , Join , <EOL> std , var , count , nunique , Summary , IsIn , <EOL> BinOp , UnaryOp , USub , Not , nelements , Repeat , Concat , Interp , <EOL> UTCFromTimestamp , DateTimeTruncate , <EOL> Transpose , TensorDot , Coerce , isnan , <EOL> greatest , least , BinaryMath , atan2 , Coalesce , Cast <EOL> ) <EOL> from . . utils import keywords <EOL> from . core import base , compute <EOL> from . pandas import array_coalesce <EOL> from . . dispatch import dispatch <EOL> from odo import into <EOL> import pandas as pd <EOL> __all__ = [ '<STR_LIT>' ] <EOL> @ dispatch ( Field , np . ndarray ) <EOL> def compute_up ( c , x , ** kwargs ) : <EOL> if x . dtype . names and c . _name in x . dtype . names : <EOL> return x [ c . _name ] <EOL> if not x . dtype . names and x . shape [ <NUM_LIT:1> ] == len ( c . _child . fields ) : <EOL> return x [ : , c . _child . fields . index ( c . _name ) ] <EOL> raise NotImplementedError ( ) <EOL> @ dispatch ( Projection , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> if x . dtype . names and all ( col in x . dtype . names for col in t . fields ) : <EOL> return x [ t . fields ] <EOL> if not x . dtype . names and x . shape [ <NUM_LIT:1> ] == len ( t . _child . fields ) : <EOL> return x [ : , [ t . _child . fields . index ( col ) for col in t . fields ] ] <EOL> raise NotImplementedError ( ) <EOL> try : <EOL> from . numba import broadcast_numba as broadcast_ndarray <EOL> except ImportError : <EOL> def broadcast_ndarray ( t , * data , ** kwargs ) : <EOL> del kwargs [ '<STR_LIT>' ] <EOL> d = dict ( zip ( t . _scalar_expr . _leaves ( ) , data ) ) <EOL> return compute ( t . _scalar_expr , d , return_type = '<STR_LIT>' , ** kwargs ) <EOL> compute_up . register ( Broadcast , np . ndarray ) ( broadcast_ndarray ) <EOL> for i in range ( <NUM_LIT:2> , <NUM_LIT:6> ) : <EOL> compute_up . register ( Broadcast , * ( [ ( np . ndarray , Number ) ] * i ) ) ( broadcast_ndarray ) <EOL> @ dispatch ( Repeat , np . ndarray ) <EOL> def compute_up ( t , data , _char_mul = np . char . multiply , ** kwargs ) : <EOL> if isinstance ( t . lhs , Expr ) : <EOL> return _char_mul ( data , t . rhs ) <EOL> else : <EOL> return _char_mul ( t . lhs , data ) <EOL> @ compute_up . register ( Repeat , np . ndarray , ( np . ndarray , base ) ) <EOL> @ compute_up . register ( Repeat , base , np . ndarray ) <EOL> def compute_up_np_repeat ( t , lhs , rhs , _char_mul = np . char . multiply , ** kwargs ) : <EOL> return _char_mul ( lhs , rhs ) <EOL> def _interp ( arr , v , _Series = pd . Series , _charmod = np . char . mod ) : <EOL> """<STR_LIT>""" <EOL> if len ( arr ) >= <NUM_LIT> : <EOL> return _Series ( arr ) % v <EOL> return _charmod ( arr , v ) <EOL> @ dispatch ( Interp , np . ndarray ) <EOL> def compute_up ( t , data , ** kwargs ) : <EOL> if isinstance ( t . lhs , Expr ) : <EOL> return _interp ( data , t . rhs ) <EOL> else : <EOL> return _interp ( t . lhs , data ) <EOL> @ compute_up . register ( Interp , np . ndarray , ( np . ndarray , base ) ) <EOL> @ compute_up . register ( Interp , base , np . ndarray ) <EOL> def compute_up_np_interp ( t , lhs , rhs , ** kwargs ) : <EOL> return _interp ( lhs , rhs ) <EOL> @ compute_up . register ( greatest , np . ndarray , ( np . ndarray , base ) ) <EOL> @ compute_up . register ( greatest , base , np . ndarray ) <EOL> def compute_up_greatest ( expr , lhs , rhs , ** kwargs ) : <EOL> return np . maximum ( lhs , rhs ) <EOL> @ compute_up . register ( least , np . ndarray , ( np . ndarray , base ) ) <EOL> @ compute_up . register ( least , base , np . ndarray ) <EOL> def compute_up_least ( expr , lhs , rhs , ** kwargs ) : <EOL> return np . minimum ( lhs , rhs ) <EOL> @ dispatch ( BinOp , np . ndarray , ( np . ndarray , base ) ) <EOL> def compute_up ( t , lhs , rhs , ** kwargs ) : <EOL> return t . op ( lhs , rhs ) <EOL> @ dispatch ( BinOp , base , np . ndarray ) <EOL> def compute_up ( t , lhs , rhs , ** kwargs ) : <EOL> return t . op ( lhs , rhs ) <EOL> @ dispatch ( BinOp , np . ndarray ) <EOL> def compute_up ( t , data , ** kwargs ) : <EOL> if isinstance ( t . lhs , Expr ) : <EOL> return t . op ( data , t . rhs ) <EOL> else : <EOL> return t . op ( t . lhs , data ) <EOL> @ compute_up . register ( BinaryMath , np . ndarray , ( np . ndarray , base ) ) <EOL> @ compute_up . register ( BinaryMath , base , np . ndarray ) <EOL> def compute_up_binary_math ( t , lhs , rhs , ** kwargs ) : <EOL> return getattr ( np , type ( t ) . __name__ ) ( lhs , rhs ) <EOL> @ dispatch ( BinaryMath , np . ndarray ) <EOL> def compute_up ( t , data , ** kwargs ) : <EOL> func = getattr ( np , type ( t ) . __name__ ) <EOL> if isinstance ( t . lhs , Expr ) : <EOL> return func ( data , t . rhs ) <EOL> else : <EOL> return func ( t . lhs , data ) <EOL> @ compute_up . register ( atan2 , np . ndarray , ( np . ndarray , base ) ) <EOL> @ compute_up . register ( atan2 , base , np . ndarray ) <EOL> def compute_up_binary_math ( t , lhs , rhs , ** kwargs ) : <EOL> return np . arctan2 ( lhs , rhs ) <EOL> @ dispatch ( atan2 , np . ndarray ) <EOL> def compute_up ( t , data , ** kwargs ) : <EOL> if isinstance ( t . lhs , Expr ) : <EOL> return np . arctan2 ( data , t . rhs ) <EOL> else : <EOL> return np . arctan2 ( t . lhs , data ) <EOL> @ dispatch ( UnaryOp , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return getattr ( np , t . symbol ) ( x ) <EOL> @ dispatch ( Not , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return np . logical_not ( x ) <EOL> @ dispatch ( USub , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return np . negative ( x ) <EOL> inat = np . datetime64 ( '<STR_LIT>' ) . view ( '<STR_LIT>' ) <EOL> @ dispatch ( count , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> result_dtype = to_numpy_dtype ( t . dshape ) <EOL> if issubclass ( x . dtype . type , ( np . floating , np . object_ ) ) : <EOL> return pd . notnull ( x ) . sum ( keepdims = t . keepdims , axis = t . axis , <EOL> dtype = result_dtype ) <EOL> elif issubclass ( x . dtype . type , np . datetime64 ) : <EOL> return ( x . view ( '<STR_LIT>' ) != inat ) . sum ( keepdims = t . keepdims , axis = t . axis , <EOL> dtype = result_dtype ) <EOL> else : <EOL> return np . ones ( x . shape , dtype = result_dtype ) . sum ( keepdims = t . keepdims , <EOL> axis = t . axis , <EOL> dtype = result_dtype ) <EOL> @ dispatch ( nunique , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> assert t . axis == tuple ( range ( ndim ( t . _child ) ) ) <EOL> result = len ( np . unique ( x ) ) <EOL> if t . keepdims : <EOL> result = np . array ( [ result ] ) <EOL> return result <EOL> @ dispatch ( Reduction , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> reducer = getattr ( np , t . symbol ) <EOL> if '<STR_LIT>' in keywords ( reducer ) : <EOL> return reducer ( x , axis = t . axis , keepdims = t . keepdims , <EOL> dtype = to_numpy_dtype ( t . schema ) ) <EOL> return reducer ( x , axis = t . axis , keepdims = t . keepdims ) <EOL> def axify ( expr , axis , keepdims = False ) : <EOL> """<STR_LIT>""" <EOL> return type ( expr ) ( expr . _child , axis = axis , keepdims = keepdims ) <EOL> @ dispatch ( Summary , np . ndarray ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> shape , dtype = to_numpy ( expr . dshape ) <EOL> if shape : <EOL> result = np . empty ( shape = shape , dtype = dtype ) <EOL> for n , v in zip ( expr . names , expr . values ) : <EOL> result [ n ] = compute ( <EOL> axify ( v , expr . axis , expr . keepdims ) , <EOL> data , <EOL> return_type = '<STR_LIT>' , <EOL> ) <EOL> return result <EOL> else : <EOL> return tuple ( <EOL> compute ( axify ( v , expr . axis ) , data , return_type = '<STR_LIT>' ) <EOL> for v in expr . values <EOL> ) <EOL> @ dispatch ( ( std , var ) , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return getattr ( x , t . symbol ) ( ddof = t . unbiased , axis = t . axis , <EOL> keepdims = t . keepdims ) <EOL> @ compute_up . register ( Distinct , np . recarray ) <EOL> def recarray_distinct ( t , rec , ** kwargs ) : <EOL> return pd . DataFrame . from_records ( rec ) . drop_duplicates ( <EOL> subset = t . on or None ) . to_records ( index = False ) . astype ( rec . dtype ) <EOL> @ dispatch ( Distinct , np . ndarray ) <EOL> def compute_up ( t , arr , _recarray_distinct = recarray_distinct , ** kwargs ) : <EOL> if t . on : <EOL> if getattr ( arr . dtype , '<STR_LIT>' , None ) is not None : <EOL> return _recarray_distinct ( t , arr , ** kwargs ) . view ( np . ndarray ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return np . unique ( arr ) <EOL> @ dispatch ( Sort , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> if x . dtype . names is None : <EOL> result = np . sort ( x ) <EOL> elif ( t . key in x . dtype . names or <EOL> isinstance ( t . key , list ) and all ( k in x . dtype . names for k in t . key ) ) : <EOL> result = np . sort ( x , order = t . key ) <EOL> elif t . key : <EOL> raise NotImplementedError ( "<STR_LIT>" % t . key ) <EOL> if not t . ascending : <EOL> result = result [ : : - <NUM_LIT:1> ] <EOL> return result <EOL> @ dispatch ( Head , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return x [ : t . n ] <EOL> @ dispatch ( Tail , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return x [ - t . n : ] <EOL> @ dispatch ( Label , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return np . array ( x , dtype = [ ( t . label , x . dtype . type ) ] ) <EOL> @ dispatch ( ReLabel , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> types = [ x . dtype [ i ] for i in range ( len ( x . dtype ) ) ] <EOL> return np . array ( x , dtype = list ( zip ( t . fields , types ) ) ) <EOL> @ dispatch ( Selection , np . ndarray ) <EOL> def compute_up ( sel , x , ** kwargs ) : <EOL> predicate = compute ( sel . predicate , { sel . _child : x } , return_type = '<STR_LIT>' ) <EOL> cond = getattr ( predicate , '<STR_LIT>' , predicate ) <EOL> return x [ cond ] <EOL> @ dispatch ( Selection , np . ndarray , np . ndarray ) <EOL> def compute_up ( expr , arr , predicate , ** kwargs ) : <EOL> return arr [ predicate ] <EOL> @ dispatch ( Selection , np . ndarray , Series ) <EOL> def compute_up ( expr , arr , predicate , ** kwargs ) : <EOL> return arr [ predicate . values ] <EOL> @ dispatch ( UTCFromTimestamp , np . ndarray ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> return ( data * <NUM_LIT> ) . astype ( '<STR_LIT>' ) <EOL> @ dispatch ( Slice , np . ndarray ) <EOL> def compute_up ( expr , x , ** kwargs ) : <EOL> return x [ expr . index ] <EOL> @ dispatch ( Cast , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> return x <EOL> @ dispatch ( Expr , np . ndarray ) <EOL> def compute_up ( t , x , ** kwargs ) : <EOL> ds = t . _child . dshape <EOL> if x . ndim > <NUM_LIT:1> or isinstance ( x , np . recarray ) or x . dtype . fields is not None : <EOL> return compute_up ( t , into ( DataFrame , x , dshape = ds ) , ** kwargs ) <EOL> else : <EOL> return compute_up ( t , into ( Series , x , dshape = ds ) , ** kwargs ) <EOL> @ dispatch ( nelements , np . ndarray ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> axis = expr . axis <EOL> if expr . keepdims : <EOL> shape = tuple ( data . shape [ i ] if i not in axis else <NUM_LIT:1> <EOL> for i in range ( ndim ( expr . _child ) ) ) <EOL> else : <EOL> shape = tuple ( data . shape [ i ] for i in range ( ndim ( expr . _child ) ) <EOL> if i not in axis ) <EOL> value = np . prod ( [ data . shape [ i ] for i in axis ] ) <EOL> result = np . empty ( shape ) <EOL> result . fill ( value ) <EOL> result = result . astype ( '<STR_LIT>' ) <EOL> return result <EOL> precision_map = { '<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> '<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> epoch = datetime . datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> offsets = { <EOL> '<STR_LIT>' : epoch . isoweekday ( ) , <EOL> '<STR_LIT>' : epoch . toordinal ( ) <EOL> } <EOL> @ dispatch ( DateTimeTruncate , ( np . ndarray , np . datetime64 ) ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> np_dtype = precision_map [ expr . unit ] <EOL> offset = offsets . get ( expr . unit , <NUM_LIT:0> ) <EOL> measure = expr . measure * <NUM_LIT:7> if expr . unit == '<STR_LIT>' else expr . measure <EOL> result = ( ( ( data . astype ( np_dtype ) <EOL> . view ( '<STR_LIT>' ) <EOL> + offset ) <EOL> // measure <EOL> * measure <EOL> - offset ) <EOL> . astype ( np_dtype ) ) <EOL> return result <EOL> @ dispatch ( isnan , np . ndarray ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> return np . isnan ( data ) <EOL> @ dispatch ( np . ndarray ) <EOL> def chunks ( x , chunksize = <NUM_LIT> ) : <EOL> start = <NUM_LIT:0> <EOL> n = len ( x ) <EOL> while start < n : <EOL> yield x [ start : start + chunksize ] <EOL> start += chunksize <EOL> @ dispatch ( Transpose , np . ndarray ) <EOL> def compute_up ( expr , x , ** kwargs ) : <EOL> return np . transpose ( x , axes = expr . axes ) <EOL> @ dispatch ( TensorDot , np . ndarray , np . ndarray ) <EOL> def compute_up ( expr , lhs , rhs , ** kwargs ) : <EOL> return np . tensordot ( lhs , rhs , axes = [ expr . _left_axes , expr . _right_axes ] ) <EOL> @ dispatch ( IsIn , np . ndarray ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> return np . in1d ( data , tuple ( expr . _keys ) ) <EOL> @ compute_up . register ( Join , DataFrame , np . ndarray ) <EOL> @ compute_up . register ( Join , np . ndarray , DataFrame ) <EOL> @ compute_up . register ( Join , np . ndarray , np . ndarray ) <EOL> def join_ndarray ( expr , lhs , rhs , ** kwargs ) : <EOL> if isinstance ( lhs , np . ndarray ) : <EOL> lhs = DataFrame ( lhs ) <EOL> if isinstance ( rhs , np . ndarray ) : <EOL> rhs = DataFrame ( rhs ) <EOL> return compute_up ( expr , lhs , rhs , ** kwargs ) <EOL> @ dispatch ( Coerce , np . ndarray ) <EOL> def compute_up ( expr , data , ** kwargs ) : <EOL> return data . astype ( to_numpy_dtype ( expr . schema ) ) <EOL> @ dispatch ( Concat , np . ndarray , np . ndarray ) <EOL> def compute_up ( expr , lhs , rhs , _concat = np . concatenate , ** kwargs ) : <EOL> return _concat ( ( lhs , rhs ) , axis = expr . axis ) <EOL> compute_up . register ( Coalesce , np . ndarray , ( np . ndarray , base ) ) ( array_coalesce ) <EOL> compute_up . register ( Coalesce , base , np . ndarray ) ( array_coalesce ) <EOL> @ dispatch ( Coalesce , np . ndarray ) <EOL> def compute_up ( t , data , ** kwargs ) : <EOL> if isinstance ( t . lhs , Expr ) : <EOL> lhs = data <EOL> rhs = t . rhs <EOL> else : <EOL> lhs = t . lhs <EOL> rhs = data <EOL> return array_coalesce ( t , lhs , rhs ) </s>
