<s> from ibis . common import RelationError , ExpressionError <EOL> from ibis . expr . datatypes import HasSchema <EOL> from ibis . expr . window import window <EOL> import ibis . expr . types as ir <EOL> import ibis . expr . operations as ops <EOL> import ibis . util as util <EOL> def sub_for ( expr , substitutions ) : <EOL> helper = _Substitutor ( expr , substitutions ) <EOL> return helper . get_result ( ) <EOL> class _Substitutor ( object ) : <EOL> def __init__ ( self , expr , substitutions , sub_memo = None ) : <EOL> self . expr = expr <EOL> self . substitutions = substitutions <EOL> self . _id_to_expr = { } <EOL> for k , v in substitutions : <EOL> self . _id_to_expr [ self . _key ( k ) ] = v <EOL> self . sub_memo = sub_memo or { } <EOL> self . unchanged = True <EOL> def get_result ( self ) : <EOL> expr = self . expr <EOL> node = expr . op ( ) <EOL> if node . blocks ( ) : <EOL> return expr <EOL> subbed_args = [ ] <EOL> for arg in node . args : <EOL> if isinstance ( arg , ( tuple , list ) ) : <EOL> subbed_arg = [ self . _sub_arg ( x ) for x in arg ] <EOL> else : <EOL> subbed_arg = self . _sub_arg ( arg ) <EOL> subbed_args . append ( subbed_arg ) <EOL> if self . unchanged : <EOL> return expr <EOL> subbed_node = type ( node ) ( * subbed_args ) <EOL> if isinstance ( expr , ir . ValueExpr ) : <EOL> result = expr . _factory ( subbed_node , name = expr . _name ) <EOL> else : <EOL> result = expr . _factory ( subbed_node ) <EOL> return result <EOL> def _sub_arg ( self , arg ) : <EOL> if isinstance ( arg , ir . Expr ) : <EOL> subbed_arg = self . sub ( arg ) <EOL> if subbed_arg is not arg : <EOL> self . unchanged = False <EOL> else : <EOL> subbed_arg = arg <EOL> return subbed_arg <EOL> def _key ( self , expr ) : <EOL> return repr ( expr . op ( ) ) <EOL> def sub ( self , expr ) : <EOL> key = self . _key ( expr ) <EOL> if key in self . sub_memo : <EOL> return self . sub_memo [ key ] <EOL> if key in self . _id_to_expr : <EOL> return self . _id_to_expr [ key ] <EOL> result = self . _sub ( expr ) <EOL> self . sub_memo [ key ] = result <EOL> return result <EOL> def _sub ( self , expr ) : <EOL> helper = _Substitutor ( expr , self . substitutions , <EOL> sub_memo = self . sub_memo ) <EOL> return helper . get_result ( ) <EOL> class ScalarAggregate ( object ) : <EOL> def __init__ ( self , expr , memo = None , default_name = '<STR_LIT>' ) : <EOL> self . expr = expr <EOL> self . memo = memo or { } <EOL> self . tables = [ ] <EOL> self . default_name = default_name <EOL> def get_result ( self ) : <EOL> expr = self . expr <EOL> subbed_expr = self . _visit ( expr ) <EOL> try : <EOL> name = subbed_expr . get_name ( ) <EOL> named_expr = subbed_expr <EOL> except : <EOL> name = self . default_name <EOL> named_expr = subbed_expr . name ( self . default_name ) <EOL> tables = list ( self . memo . values ( ) ) <EOL> table = self . tables [ <NUM_LIT:0> ] <EOL> for other in self . tables [ <NUM_LIT:1> : ] : <EOL> table = table . cross_join ( other ) <EOL> return table . projection ( [ named_expr ] ) , name <EOL> def _visit ( self , expr ) : <EOL> if is_scalar_reduce ( expr ) and not has_multiple_bases ( expr ) : <EOL> key = self . _key ( expr ) <EOL> if key not in self . memo : <EOL> agg_expr , name = reduction_to_aggregation ( expr ) <EOL> self . memo [ key ] = agg_expr , name <EOL> self . tables . append ( agg_expr ) <EOL> else : <EOL> agg_expr , name = self . memo [ key ] <EOL> return agg_expr [ name ] <EOL> elif not isinstance ( expr , ir . Expr ) : <EOL> return expr <EOL> node = expr . op ( ) <EOL> subbed_args = [ ] <EOL> for arg in node . args : <EOL> if isinstance ( arg , ( tuple , list ) ) : <EOL> subbed_arg = [ self . _visit ( x ) for x in arg ] <EOL> else : <EOL> subbed_arg = self . _visit ( arg ) <EOL> subbed_args . append ( subbed_arg ) <EOL> subbed_node = type ( node ) ( * subbed_args ) <EOL> if isinstance ( expr , ir . ValueExpr ) : <EOL> result = expr . _factory ( subbed_node , name = expr . _name ) <EOL> else : <EOL> result = expr . _factory ( subbed_node ) <EOL> return result <EOL> def _key ( self , expr ) : <EOL> return repr ( expr . op ( ) ) <EOL> def has_multiple_bases ( expr ) : <EOL> return len ( find_all_tables ( expr ) ) > <NUM_LIT:1> <EOL> def reduction_to_aggregation ( expr , default_name = '<STR_LIT>' ) : <EOL> tables = find_all_tables ( expr ) <EOL> try : <EOL> name = expr . get_name ( ) <EOL> named_expr = expr <EOL> except : <EOL> name = default_name <EOL> named_expr = expr . name ( default_name ) <EOL> if len ( tables ) == <NUM_LIT:1> : <EOL> table = list ( tables . values ( ) ) [ <NUM_LIT:0> ] <EOL> return table . aggregate ( [ named_expr ] ) , name <EOL> else : <EOL> return ScalarAggregate ( expr , None , default_name ) . get_result ( ) <EOL> def find_all_tables ( expr , memo = None ) : <EOL> if memo is None : <EOL> memo = { } <EOL> node = expr . op ( ) <EOL> if isinstance ( expr , ir . TableExpr ) : <EOL> key = id ( node ) <EOL> if key not in memo : <EOL> memo [ key ] = expr <EOL> return memo <EOL> for arg in node . flat_args ( ) : <EOL> if isinstance ( arg , ir . Expr ) : <EOL> find_all_tables ( arg , memo ) <EOL> return memo <EOL> def is_scalar_reduce ( x ) : <EOL> return isinstance ( x , ir . ScalarExpr ) and ops . is_reduction ( x ) <EOL> def substitute_parents ( expr , lift_memo = None , past_projection = True ) : <EOL> rewriter = ExprSimplifier ( expr , lift_memo = lift_memo , <EOL> block_projection = not past_projection ) <EOL> return rewriter . get_result ( ) <EOL> class ExprSimplifier ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , expr , lift_memo = None , block_projection = False ) : <EOL> self . expr = expr <EOL> self . lift_memo = lift_memo or { } <EOL> self . block_projection = block_projection <EOL> def get_result ( self ) : <EOL> expr = self . expr <EOL> node = expr . op ( ) <EOL> if isinstance ( node , ir . Literal ) : <EOL> return expr <EOL> if isinstance ( node , ops . TableColumn ) : <EOL> result = self . _lift_TableColumn ( expr , block = self . block_projection ) <EOL> if result is not expr : <EOL> return result <EOL> elif isinstance ( node , ops . Selection ) : <EOL> return self . _lift_Selection ( expr , block = self . block_projection ) <EOL> elif isinstance ( node , ops . Aggregation ) : <EOL> return self . _lift_Aggregation ( expr , block = self . block_projection ) <EOL> unchanged = True <EOL> lifted_args = [ ] <EOL> for arg in node . args : <EOL> lifted_arg , unch_arg = self . _lift_arg ( <EOL> arg , block = self . block_projection ) <EOL> lifted_args . append ( lifted_arg ) <EOL> unchanged = unchanged and unch_arg <EOL> if unchanged : <EOL> return expr <EOL> lifted_node = type ( node ) ( * lifted_args ) <EOL> if isinstance ( expr , ir . ValueExpr ) : <EOL> result = expr . _factory ( lifted_node , name = expr . _name ) <EOL> else : <EOL> result = expr . _factory ( lifted_node ) <EOL> return result <EOL> def _lift_arg ( self , arg , block = None ) : <EOL> unchanged = [ True ] <EOL> def _lift ( x ) : <EOL> if isinstance ( x , ir . Expr ) : <EOL> lifted_arg = self . lift ( x , block = block ) <EOL> if lifted_arg is not x : <EOL> unchanged [ <NUM_LIT:0> ] = False <EOL> else : <EOL> lifted_arg = x <EOL> return lifted_arg <EOL> if arg is None : <EOL> return arg , True <EOL> if isinstance ( arg , ( tuple , list ) ) : <EOL> result = [ _lift ( x ) for x in arg ] <EOL> else : <EOL> result = _lift ( arg ) <EOL> return result , unchanged [ <NUM_LIT:0> ] <EOL> def lift ( self , expr , block = None ) : <EOL> key = id ( expr . op ( ) ) , block <EOL> if key in self . lift_memo : <EOL> return self . lift_memo [ key ] <EOL> op = expr . op ( ) <EOL> if isinstance ( op , ops . ValueNode ) : <EOL> return self . _sub ( expr , block = block ) <EOL> elif isinstance ( op , ops . Selection ) : <EOL> result = self . _lift_Selection ( expr , block = block ) <EOL> elif isinstance ( op , ops . Join ) : <EOL> result = self . _lift_Join ( expr , block = block ) <EOL> elif isinstance ( op , ( ops . TableNode , HasSchema ) ) : <EOL> return expr <EOL> else : <EOL> return self . _sub ( expr , block = block ) <EOL> self . lift_memo [ key ] = result <EOL> return result <EOL> def _lift_TableColumn ( self , expr , block = None ) : <EOL> node = expr . op ( ) <EOL> tnode = node . table . op ( ) <EOL> root = _base_table ( tnode ) <EOL> result = expr <EOL> if isinstance ( root , ops . Selection ) : <EOL> can_lift = False <EOL> for val in root . selections : <EOL> if ( isinstance ( val . op ( ) , ops . PhysicalTable ) and <EOL> node . name in val . schema ( ) ) : <EOL> can_lift = True <EOL> lifted_root = self . lift ( val ) <EOL> elif ( isinstance ( val . op ( ) , ops . TableColumn ) and <EOL> val . op ( ) . name == val . get_name ( ) and <EOL> node . name == val . get_name ( ) ) : <EOL> can_lift = True <EOL> lifted_root = self . lift ( val . op ( ) . table ) <EOL> if can_lift and not block : <EOL> lifted_node = ops . TableColumn ( node . name , lifted_root ) <EOL> result = expr . _factory ( lifted_node , name = expr . _name ) <EOL> return result <EOL> def _lift_Aggregation ( self , expr , block = None ) : <EOL> if block is None : <EOL> block = self . block_projection <EOL> op = expr . op ( ) <EOL> if block : <EOL> lifted_table = op . table <EOL> else : <EOL> lifted_table = self . lift ( op . table , block = True ) <EOL> unch = lifted_table is op . table <EOL> lifted_aggs , unch1 = self . _lift_arg ( op . agg_exprs , block = True ) <EOL> lifted_by , unch2 = self . _lift_arg ( op . by , block = True ) <EOL> lifted_having , unch3 = self . _lift_arg ( op . having , block = True ) <EOL> unchanged = unch and unch1 and unch2 and unch3 <EOL> if not unchanged : <EOL> lifted_op = ops . Aggregation ( lifted_table , lifted_aggs , <EOL> by = lifted_by , having = lifted_having ) <EOL> result = ir . TableExpr ( lifted_op ) <EOL> else : <EOL> result = expr <EOL> return result <EOL> def _lift_Selection ( self , expr , block = None ) : <EOL> if block is None : <EOL> block = self . block_projection <EOL> op = expr . op ( ) <EOL> if block and op . blocks ( ) : <EOL> return expr <EOL> else : <EOL> lifted_table , unch = self . _lift_arg ( op . table , block = True ) <EOL> lifted_selections , unch_sel = self . _lift_arg ( op . selections , block = True ) <EOL> unchanged = unch and unch_sel <EOL> lifted_predicates , unch_sel = self . _lift_arg ( op . predicates , block = True ) <EOL> unchanged = unch and unch_sel <EOL> lifted_sort_keys , unch_sel = self . _lift_arg ( op . sort_keys , block = True ) <EOL> unchanged = unch and unch_sel <EOL> if not unchanged : <EOL> lifted_projection = ops . Selection ( lifted_table , lifted_selections , <EOL> lifted_predicates , <EOL> lifted_sort_keys ) <EOL> result = ir . TableExpr ( lifted_projection ) <EOL> else : <EOL> result = expr <EOL> return result <EOL> def _lift_Join ( self , expr , block = None ) : <EOL> op = expr . op ( ) <EOL> left_lifted = self . lift ( op . left , block = block ) <EOL> right_lifted = self . lift ( op . right , block = block ) <EOL> unchanged = ( left_lifted is op . left and <EOL> right_lifted is op . right ) <EOL> lifted_preds = [ ] <EOL> for x in op . predicates : <EOL> subbed = self . _sub ( x , block = True ) <EOL> if subbed is not x : <EOL> unchanged = False <EOL> lifted_preds . append ( subbed ) <EOL> if not unchanged : <EOL> lifted_join = type ( op ) ( left_lifted , right_lifted , lifted_preds ) <EOL> result = ir . TableExpr ( lifted_join ) <EOL> else : <EOL> result = expr <EOL> return result <EOL> def _sub ( self , expr , block = None ) : <EOL> if block is None : <EOL> block = self . block_projection <EOL> helper = ExprSimplifier ( expr , lift_memo = self . lift_memo , <EOL> block_projection = block ) <EOL> return helper . get_result ( ) <EOL> def _base_table ( table_node ) : <EOL> if table_node . blocks ( ) : <EOL> return table_node <EOL> else : <EOL> return _base_table ( table_node . table . op ( ) ) <EOL> def apply_filter ( expr , predicates ) : <EOL> op = expr . op ( ) <EOL> if isinstance ( op , ops . Selection ) : <EOL> return _filter_selection ( expr , predicates ) <EOL> elif isinstance ( op , ops . Aggregation ) : <EOL> simplified_predicates = [ sub_for ( x , [ ( expr , op . table ) ] ) <EOL> for x in predicates ] <EOL> if op . table . _is_valid ( simplified_predicates ) : <EOL> result = ops . Aggregation ( <EOL> op . table , op . agg_exprs , by = op . by , having = op . having , <EOL> predicates = op . predicates + simplified_predicates , <EOL> sort_keys = op . sort_keys ) <EOL> return ir . TableExpr ( result ) <EOL> elif isinstance ( op , ops . Join ) : <EOL> expr = expr . materialize ( ) <EOL> result = ops . Selection ( expr , [ ] , predicates ) <EOL> return ir . TableExpr ( result ) <EOL> def _filter_selection ( expr , predicates ) : <EOL> op = expr . op ( ) <EOL> if not op . blocks ( ) : <EOL> simplified_predicates = [ sub_for ( x , [ ( expr , op . table ) ] ) <EOL> for x in predicates ] <EOL> if op . table . _is_valid ( simplified_predicates ) : <EOL> result = ops . Selection ( <EOL> op . table , [ ] , <EOL> predicates = op . predicates + simplified_predicates , <EOL> sort_keys = op . sort_keys ) <EOL> return ir . TableExpr ( result ) <EOL> can_pushdown = _can_pushdown ( op , predicates ) <EOL> if can_pushdown : <EOL> simplified_predicates = [ substitute_parents ( x ) for x in predicates ] <EOL> fused_predicates = op . predicates + simplified_predicates <EOL> result = ops . Selection ( op . table , <EOL> proj_exprs = op . selections , <EOL> predicates = fused_predicates , <EOL> sort_keys = op . sort_keys ) <EOL> else : <EOL> result = ops . Selection ( expr , proj_exprs = [ ] , <EOL> predicates = predicates ) <EOL> return ir . TableExpr ( result ) <EOL> def _can_pushdown ( op , predicates ) : <EOL> can_pushdown = True <EOL> for pred in predicates : <EOL> validator = _PushdownValidate ( op , pred ) <EOL> predicate_is_valid = validator . get_result ( ) <EOL> can_pushdown = can_pushdown and predicate_is_valid <EOL> return can_pushdown <EOL> class _PushdownValidate ( object ) : <EOL> def __init__ ( self , parent , predicate ) : <EOL> self . parent = parent <EOL> self . pred = predicate <EOL> self . validator = ExprValidator ( [ self . parent . table ] ) <EOL> self . valid = True <EOL> def get_result ( self ) : <EOL> self . _walk ( self . pred ) <EOL> return self . valid <EOL> def _walk ( self , expr ) : <EOL> node = expr . op ( ) <EOL> if isinstance ( node , ops . TableColumn ) : <EOL> is_valid = self . _validate_column ( expr ) <EOL> self . valid = self . valid and is_valid <EOL> for arg in node . flat_args ( ) : <EOL> if isinstance ( arg , ir . ValueExpr ) : <EOL> self . _walk ( arg ) <EOL> def _validate_column ( self , expr ) : <EOL> if isinstance ( self . parent , ops . Selection ) : <EOL> return self . _validate_projection ( expr ) <EOL> else : <EOL> validator = ExprValidator ( [ self . parent . table ] ) <EOL> return validator . validate ( expr ) <EOL> def _validate_projection ( self , expr ) : <EOL> is_valid = False <EOL> node = expr . op ( ) <EOL> if _is_aliased ( expr ) : <EOL> return False <EOL> for val in self . parent . selections : <EOL> if ( isinstance ( val . op ( ) , ops . PhysicalTable ) and <EOL> node . name in val . schema ( ) ) : <EOL> is_valid = True <EOL> elif ( isinstance ( val . op ( ) , ops . TableColumn ) and <EOL> node . name == val . get_name ( ) and <EOL> not _is_aliased ( val ) ) : <EOL> col_table = val . op ( ) . table . op ( ) <EOL> lifted_node = substitute_parents ( expr ) . op ( ) <EOL> is_valid = ( col_table . is_ancestor ( node . table ) or <EOL> col_table . is_ancestor ( lifted_node . table ) ) <EOL> return is_valid <EOL> def _is_aliased ( col_expr ) : <EOL> return col_expr . op ( ) . name != col_expr . get_name ( ) <EOL> def windowize_function ( expr , w = None ) : <EOL> def _check_window ( x ) : <EOL> arg , window = x . op ( ) . args <EOL> if isinstance ( arg . op ( ) , ops . RowNumber ) : <EOL> if len ( window . _order_by ) == <NUM_LIT:0> : <EOL> raise ExpressionError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> return x <EOL> def _windowize ( x , w ) : <EOL> if not isinstance ( x . op ( ) , ops . WindowOp ) : <EOL> walked = _walk ( x , w ) <EOL> else : <EOL> window_arg , window_w = x . op ( ) . args <EOL> walked_child = _walk ( window_arg , w ) <EOL> if walked_child is not window_arg : <EOL> walked = x . _factory ( ops . WindowOp ( walked_child , window_w ) , <EOL> name = x . _name ) <EOL> else : <EOL> walked = x <EOL> op = walked . op ( ) <EOL> if ( isinstance ( op , ops . AnalyticOp ) or <EOL> getattr ( op , '<STR_LIT>' , False ) ) : <EOL> if w is None : <EOL> w = window ( ) <EOL> return _check_window ( walked . over ( w ) ) <EOL> elif isinstance ( op , ops . WindowOp ) : <EOL> if w is not None : <EOL> return _check_window ( walked . over ( w ) ) <EOL> else : <EOL> return _check_window ( walked ) <EOL> else : <EOL> return walked <EOL> def _walk ( x , w ) : <EOL> op = x . op ( ) <EOL> unchanged = True <EOL> windowed_args = [ ] <EOL> for arg in op . args : <EOL> if not isinstance ( arg , ir . ValueExpr ) : <EOL> windowed_args . append ( arg ) <EOL> continue <EOL> new_arg = _windowize ( arg , w ) <EOL> unchanged = unchanged and arg is new_arg <EOL> windowed_args . append ( new_arg ) <EOL> if not unchanged : <EOL> new_op = type ( op ) ( * windowed_args ) <EOL> return x . _factory ( new_op , name = x . _name ) <EOL> else : <EOL> return x <EOL> return _windowize ( expr , w ) <EOL> class Projector ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , parent , proj_exprs ) : <EOL> self . parent = parent <EOL> self . input_exprs = proj_exprs <EOL> self . resolved_exprs = [ parent . _ensure_expr ( e ) for e in proj_exprs ] <EOL> node = self . parent . op ( ) <EOL> if isinstance ( node , ops . Selection ) : <EOL> roots = [ node ] <EOL> else : <EOL> roots = node . root_tables ( ) <EOL> self . parent_roots = roots <EOL> clean_exprs = [ ] <EOL> for expr in self . resolved_exprs : <EOL> expr = windowize_function ( expr ) <EOL> clean_exprs . append ( expr ) <EOL> self . clean_exprs = clean_exprs <EOL> def get_result ( self ) : <EOL> roots = self . parent_roots <EOL> if len ( roots ) == <NUM_LIT:1> and isinstance ( roots [ <NUM_LIT:0> ] , ops . Selection ) : <EOL> fused_op = self . _check_fusion ( roots [ <NUM_LIT:0> ] ) <EOL> if fused_op is not None : <EOL> return fused_op <EOL> return ops . Selection ( self . parent , self . clean_exprs ) <EOL> def _check_fusion ( self , root ) : <EOL> roots = root . table . _root_tables ( ) <EOL> validator = ExprValidator ( [ root . table ] ) <EOL> fused_exprs = [ ] <EOL> can_fuse = False <EOL> resolved = _maybe_resolve_exprs ( root . table , self . input_exprs ) <EOL> if not resolved : <EOL> return None <EOL> for val in resolved : <EOL> lifted_val = substitute_parents ( val ) <EOL> if ( isinstance ( val , ir . TableExpr ) and <EOL> ( self . parent . op ( ) . is_ancestor ( val ) or <EOL> len ( roots ) == <NUM_LIT:1> and val . _root_tables ( ) [ <NUM_LIT:0> ] is roots [ <NUM_LIT:0> ] ) ) : <EOL> can_fuse = True <EOL> have_root = False <EOL> for y in root . selections : <EOL> if y . equals ( root . table ) : <EOL> fused_exprs . append ( root . table ) <EOL> have_root = True <EOL> continue <EOL> fused_exprs . append ( y ) <EOL> if not have_root and len ( root . selections ) == <NUM_LIT:0> : <EOL> fused_exprs = [ root . table ] + fused_exprs <EOL> elif validator . validate ( lifted_val ) : <EOL> can_fuse = True <EOL> fused_exprs . append ( lifted_val ) <EOL> elif not validator . validate ( val ) : <EOL> can_fuse = False <EOL> break <EOL> else : <EOL> fused_exprs . append ( val ) <EOL> if can_fuse : <EOL> return ops . Selection ( root . table , fused_exprs , <EOL> predicates = root . predicates , <EOL> sort_keys = root . sort_keys ) <EOL> else : <EOL> return None <EOL> def _maybe_resolve_exprs ( table , exprs ) : <EOL> try : <EOL> return table . _resolve ( exprs ) <EOL> except : <EOL> return None <EOL> class ExprValidator ( object ) : <EOL> def __init__ ( self , exprs ) : <EOL> self . parent_exprs = exprs <EOL> self . roots = [ ] <EOL> for expr in self . parent_exprs : <EOL> self . roots . extend ( expr . _root_tables ( ) ) <EOL> def has_common_roots ( self , expr ) : <EOL> return self . validate ( expr ) <EOL> def validate ( self , expr ) : <EOL> op = expr . op ( ) <EOL> if isinstance ( op , ops . TableColumn ) : <EOL> if self . _among_roots ( op . table . op ( ) ) : <EOL> return True <EOL> elif isinstance ( op , ops . Selection ) : <EOL> if self . _among_roots ( op ) : <EOL> return True <EOL> expr_roots = expr . _root_tables ( ) <EOL> for root in expr_roots : <EOL> if not self . _among_roots ( root ) : <EOL> return False <EOL> return True <EOL> def _among_roots ( self , node ) : <EOL> return self . roots_shared ( node ) > <NUM_LIT:0> <EOL> def roots_shared ( self , node ) : <EOL> count = <NUM_LIT:0> <EOL> for root in self . roots : <EOL> if root . is_ancestor ( node ) : <EOL> count += <NUM_LIT:1> <EOL> return count <EOL> def shares_some_roots ( self , expr ) : <EOL> expr_roots = expr . _root_tables ( ) <EOL> return any ( self . _among_roots ( root ) <EOL> for root in expr_roots ) <EOL> def shares_one_root ( self , expr ) : <EOL> expr_roots = expr . _root_tables ( ) <EOL> total = sum ( self . roots_shared ( root ) <EOL> for root in expr_roots ) <EOL> return total == <NUM_LIT:1> <EOL> def shares_multiple_roots ( self , expr ) : <EOL> expr_roots = expr . _root_tables ( ) <EOL> total = sum ( self . roots_shared ( expr_roots ) <EOL> for root in expr_roots ) <EOL> return total > <NUM_LIT:1> <EOL> def validate_all ( self , exprs ) : <EOL> for expr in exprs : <EOL> self . assert_valid ( expr ) <EOL> def assert_valid ( self , expr ) : <EOL> if not self . validate ( expr ) : <EOL> msg = self . _error_message ( expr ) <EOL> raise RelationError ( msg ) <EOL> def _error_message ( self , expr ) : <EOL> return ( '<STR_LIT>' <EOL> '<STR_LIT>' % repr ( expr ) ) <EOL> class CommonSubexpr ( object ) : <EOL> def __init__ ( self , exprs ) : <EOL> self . parent_exprs = exprs <EOL> def validate ( self , expr ) : <EOL> if isinstance ( expr , ir . TableExpr ) : <EOL> if not self . _check ( expr ) : <EOL> return False <EOL> op = expr . op ( ) <EOL> for arg in op . flat_args ( ) : <EOL> if not isinstance ( arg , ir . Expr ) : <EOL> continue <EOL> elif not isinstance ( arg , ir . TableExpr ) : <EOL> if not self . validate ( arg ) : <EOL> return False <EOL> else : <EOL> if not self . _check ( arg ) : <EOL> return False <EOL> return True <EOL> def _check ( self , expr ) : <EOL> is_valid = False <EOL> for parent in self . parent_exprs : <EOL> is_valid = is_valid or self . _check_table ( parent , expr ) <EOL> return is_valid <EOL> def _check_table ( self , parent , needle ) : <EOL> def _matches ( expr ) : <EOL> op = expr . op ( ) <EOL> if expr . equals ( needle ) : <EOL> return True <EOL> if op . blocks ( ) : <EOL> return False <EOL> for arg in op . flat_args ( ) : <EOL> if not isinstance ( arg , ir . Expr ) : <EOL> continue <EOL> if _matches ( arg ) : <EOL> return True <EOL> return True <EOL> return _matches ( parent ) <EOL> def validate_all ( self , exprs ) : <EOL> for expr in exprs : <EOL> self . assert_valid ( expr ) <EOL> def assert_valid ( self , expr ) : <EOL> if not self . validate ( expr ) : <EOL> msg = self . _error_message ( expr ) <EOL> raise RelationError ( msg ) <EOL> def _error_message ( self , expr ) : <EOL> return ( '<STR_LIT>' <EOL> '<STR_LIT>' % repr ( expr ) ) <EOL> class FilterValidator ( ExprValidator ) : <EOL> """<STR_LIT>""" <EOL> def validate ( self , expr ) : <EOL> op = expr . op ( ) <EOL> is_valid = True <EOL> if isinstance ( op , ops . Contains ) : <EOL> value_valid = ExprValidator . validate ( self , op . value ) <EOL> is_valid = value_valid <EOL> else : <EOL> roots_valid = [ ] <EOL> for arg in op . flat_args ( ) : <EOL> if isinstance ( arg , ir . ScalarExpr ) : <EOL> pass <EOL> elif isinstance ( arg , ops . TopKExpr ) : <EOL> roots_valid . append ( True ) <EOL> elif isinstance ( arg , ( ir . ArrayExpr , ir . AnalyticExpr ) ) : <EOL> roots_valid . append ( self . shares_some_roots ( arg ) ) <EOL> elif isinstance ( arg , ir . Expr ) : <EOL> raise NotImplementedError <EOL> else : <EOL> pass <EOL> is_valid = any ( roots_valid ) <EOL> return is_valid <EOL> def find_source_table ( expr ) : <EOL> node = expr . op ( ) <EOL> first_tables = [ ] <EOL> def push_first ( arg ) : <EOL> if not isinstance ( arg , ir . Expr ) : <EOL> return <EOL> if isinstance ( arg , ir . TableExpr ) : <EOL> first_tables . append ( arg ) <EOL> else : <EOL> collect ( arg . op ( ) ) <EOL> def collect ( node ) : <EOL> for arg in node . flat_args ( ) : <EOL> push_first ( arg ) <EOL> collect ( node ) <EOL> options = util . unique_by_key ( first_tables , id ) <EOL> if len ( options ) > <NUM_LIT:1> : <EOL> raise NotImplementedError <EOL> return options [ <NUM_LIT:0> ] <EOL> def unwrap_ands ( expr ) : <EOL> out_exprs = [ ] <EOL> def walk ( expr ) : <EOL> op = expr . op ( ) <EOL> if isinstance ( op , ops . Comparison ) : <EOL> out_exprs . append ( expr ) <EOL> elif isinstance ( op , ops . And ) : <EOL> walk ( op . left ) <EOL> walk ( op . right ) <EOL> else : <EOL> raise Exception ( '<STR_LIT>' <EOL> . format ( expr . _repr ( ) ) ) <EOL> walk ( expr ) <EOL> return out_exprs </s>
