<s> from datetime import timedelta <EOL> from operator import methodcaller <EOL> import itertools <EOL> import math <EOL> import pytest <EOL> sa = pytest . importorskip ( '<STR_LIT>' ) <EOL> pytest . importorskip ( '<STR_LIT>' ) <EOL> import os <EOL> import numpy as np <EOL> import pandas as pd <EOL> import pandas . util . testing as tm <EOL> from datashape import dshape <EOL> from odo import odo , drop , discover <EOL> from odo . utils import tmpfile <EOL> from blaze import ( <EOL> data , <EOL> atan2 , <EOL> by , <EOL> coalesce , <EOL> compute , <EOL> concat , <EOL> cos , <EOL> greatest , <EOL> join , <EOL> least , <EOL> radians , <EOL> sin , <EOL> sqrt , <EOL> symbol , <EOL> transform , <EOL> ) <EOL> from blaze . interactive import iscorescalar <EOL> from blaze . utils import example , normalize <EOL> names = ( '<STR_LIT>' % i for i in itertools . count ( ) ) <EOL> @ pytest . fixture ( scope = '<STR_LIT>' ) <EOL> def pg_ip ( ) : <EOL> return os . environ . get ( '<STR_LIT>' , '<STR_LIT:localhost>' ) <EOL> @ pytest . fixture <EOL> def url ( pg_ip ) : <EOL> return '<STR_LIT>' . format ( pg_ip ) <EOL> @ pytest . yield_fixture <EOL> def sql ( url ) : <EOL> ds = dshape ( '<STR_LIT>' ) <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = ds ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> assert t . dshape == ds <EOL> t = data ( odo ( [ ( '<STR_LIT:a>' , <NUM_LIT:1> ) , ( '<STR_LIT:b>' , <NUM_LIT:2> ) ] , t ) ) <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture ( scope = '<STR_LIT>' ) <EOL> def nyc ( pg_ip ) : <EOL> nyc_pd = odo ( example ( '<STR_LIT>' ) , pd . DataFrame ) <EOL> try : <EOL> t = odo ( nyc_pd , <EOL> '<STR_LIT>' . format ( pg_ip ) ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture <EOL> def big_sql ( url ) : <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = '<STR_LIT>' ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> t = odo ( zip ( list ( '<STR_LIT:a>' * <NUM_LIT:100> ) , list ( range ( <NUM_LIT:100> ) ) ) , t ) <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture <EOL> def sqla ( url ) : <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = '<STR_LIT>' ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> t = odo ( [ ( '<STR_LIT:a>' , <NUM_LIT:1> ) , ( None , <NUM_LIT:1> ) , ( '<STR_LIT:c>' , None ) ] , t ) <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture <EOL> def sqlb ( url ) : <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = '<STR_LIT>' ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> t = odo ( [ ( '<STR_LIT:a>' , <NUM_LIT:1> ) , ( '<STR_LIT:b>' , <NUM_LIT:2> ) ] , t ) <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture <EOL> def sql_with_dts ( url ) : <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = '<STR_LIT>' ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> t = odo ( [ ( d , ) for d in pd . date_range ( '<STR_LIT>' , '<STR_LIT>' ) ] , t ) <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture <EOL> def sql_with_timedeltas ( url ) : <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = '<STR_LIT>' ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> t = odo ( [ ( timedelta ( seconds = n ) , ) for n in range ( <NUM_LIT:10> ) ] , t ) <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture <EOL> def sql_two_tables ( url ) : <EOL> dshape = '<STR_LIT>' <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = dshape ) <EOL> u = data ( url % next ( names ) , dshape = dshape ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield u , t <EOL> finally : <EOL> drop ( t ) <EOL> drop ( u ) <EOL> @ pytest . yield_fixture <EOL> def products ( url ) : <EOL> try : <EOL> products = data ( url % '<STR_LIT>' , <EOL> dshape = """<STR_LIT>""" , <EOL> primary_key = [ '<STR_LIT>' ] ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield products <EOL> finally : <EOL> drop ( products ) <EOL> @ pytest . yield_fixture <EOL> def orders ( url , products ) : <EOL> try : <EOL> orders = data ( url % '<STR_LIT>' , <EOL> dshape = """<STR_LIT>""" , <EOL> foreign_keys = dict ( product_id = products . data . c . product_id ) , <EOL> primary_key = [ '<STR_LIT>' ] ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield orders <EOL> finally : <EOL> drop ( orders ) <EOL> @ pytest . yield_fixture <EOL> def main ( url ) : <EOL> try : <EOL> main = odo ( [ ( i , int ( np . random . randint ( <NUM_LIT:10> ) ) ) for i in range ( <NUM_LIT> ) ] , <EOL> url % '<STR_LIT>' , <EOL> dshape = dshape ( '<STR_LIT>' ) , <EOL> primary_key = [ '<STR_LIT:id>' ] ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield main <EOL> finally : <EOL> drop ( main ) <EOL> @ pytest . yield_fixture <EOL> def pkey ( url , main ) : <EOL> choices = [ u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , <EOL> u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , u'<STR_LIT>' , <EOL> u'<STR_LIT>' ] <EOL> n = <NUM_LIT:100> <EOL> data = list ( zip ( range ( n ) , <EOL> np . random . choice ( choices , size = n ) . tolist ( ) , <EOL> np . random . uniform ( <NUM_LIT> , <NUM_LIT> , size = n ) . tolist ( ) , <EOL> np . random . randint ( main . count ( ) . scalar ( ) , size = n ) . tolist ( ) ) ) <EOL> try : <EOL> pkey = odo ( data , url % '<STR_LIT>' , <EOL> dshape = dshape ( '<STR_LIT>' ) , <EOL> foreign_keys = dict ( main = main . c . id ) , <EOL> primary_key = [ '<STR_LIT:id>' ] ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield pkey <EOL> finally : <EOL> drop ( pkey ) <EOL> @ pytest . yield_fixture <EOL> def fkey ( url , pkey ) : <EOL> try : <EOL> fkey = odo ( [ ( i , <EOL> int ( np . random . randint ( pkey . count ( ) . scalar ( ) ) ) , <EOL> int ( np . random . randint ( <NUM_LIT> ) ) ) <EOL> for i in range ( <NUM_LIT:10> ) ] , <EOL> url % '<STR_LIT>' , <EOL> dshape = dshape ( '<STR_LIT>' ) , <EOL> foreign_keys = dict ( sym_id = pkey . c . id ) , <EOL> primary_key = [ '<STR_LIT:id>' ] ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield fkey <EOL> finally : <EOL> drop ( fkey ) <EOL> @ pytest . yield_fixture <EOL> def sql_with_float ( url ) : <EOL> try : <EOL> t = data ( url % next ( names ) , dshape = '<STR_LIT>' ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> @ pytest . yield_fixture ( scope = '<STR_LIT>' ) <EOL> def nyc_csv ( pg_ip ) : <EOL> try : <EOL> t = odo ( <EOL> example ( '<STR_LIT>' ) , <EOL> '<STR_LIT>' . format ( pg_ip ) , <EOL> ) <EOL> except sa . exc . OperationalError as e : <EOL> pytest . skip ( str ( e ) ) <EOL> else : <EOL> try : <EOL> yield t <EOL> finally : <EOL> drop ( t ) <EOL> def test_nyc_csv ( nyc_csv ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc_csv ) ) <EOL> assert compute ( t . nrows , nyc_csv , return_type = '<STR_LIT>' ) > <NUM_LIT:0> <EOL> def test_postgres_create ( sql ) : <EOL> assert odo ( sql , list ) == [ ( '<STR_LIT:a>' , <NUM_LIT:1> ) , ( '<STR_LIT:b>' , <NUM_LIT:2> ) ] <EOL> def test_postgres_isnan ( sql_with_float ) : <EOL> dta = ( <NUM_LIT:1.0> , ) , ( float ( '<STR_LIT>' ) , ) <EOL> table = odo ( dta , sql_with_float ) <EOL> sym = symbol ( '<STR_LIT:s>' , discover ( dta ) ) <EOL> assert compute ( sym . isnan ( ) , table , return_type = list ) == [ ( False , ) , ( True , ) ] <EOL> def test_insert_from_subselect ( sql_with_float ) : <EOL> data = pd . DataFrame ( [ { '<STR_LIT:c>' : <NUM_LIT> } , { '<STR_LIT:c>' : <NUM_LIT:1.0> } ] ) <EOL> tbl = odo ( data , sql_with_float ) <EOL> s = symbol ( '<STR_LIT:s>' , discover ( data ) ) <EOL> odo ( compute ( s [ s . c . isin ( ( <NUM_LIT:1.0> , <NUM_LIT> ) ) ] . sort ( ) , tbl , return_type = '<STR_LIT>' ) , sql_with_float ) , <EOL> tm . assert_frame_equal ( <EOL> odo ( sql_with_float , pd . DataFrame ) . iloc [ <NUM_LIT:2> : ] . reset_index ( drop = True ) , <EOL> pd . DataFrame ( [ { '<STR_LIT:c>' : <NUM_LIT:1.0> } , { '<STR_LIT:c>' : <NUM_LIT> } ] ) , <EOL> ) <EOL> def test_concat ( sql_two_tables ) : <EOL> t_table , u_table = sql_two_tables <EOL> t_data = pd . DataFrame ( np . arange ( <NUM_LIT:5> ) , columns = [ '<STR_LIT:a>' ] ) <EOL> u_data = pd . DataFrame ( np . arange ( <NUM_LIT:5> , <NUM_LIT:10> ) , columns = [ '<STR_LIT:a>' ] ) <EOL> odo ( t_data , t_table ) <EOL> odo ( u_data , u_table ) <EOL> t = symbol ( '<STR_LIT:t>' , discover ( t_data ) ) <EOL> u = symbol ( '<STR_LIT:u>' , discover ( u_data ) ) <EOL> tm . assert_frame_equal ( <EOL> compute ( concat ( t , u ) . sort ( '<STR_LIT:a>' ) , { t : t_table , u : u_table } , return_type = pd . DataFrame ) , <EOL> pd . DataFrame ( np . arange ( <NUM_LIT:10> ) , columns = [ '<STR_LIT:a>' ] ) , <EOL> ) <EOL> def test_concat_invalid_axis ( sql_two_tables ) : <EOL> t_table , u_table = sql_two_tables <EOL> t_data = pd . DataFrame ( np . arange ( <NUM_LIT:5> ) , columns = [ '<STR_LIT:a>' ] ) <EOL> u_data = pd . DataFrame ( np . arange ( <NUM_LIT:5> , <NUM_LIT:10> ) , columns = [ '<STR_LIT:a>' ] ) <EOL> odo ( t_data , t_table ) <EOL> odo ( u_data , u_table ) <EOL> t = symbol ( '<STR_LIT:t>' , '<STR_LIT>' ) <EOL> u = symbol ( '<STR_LIT:u>' , '<STR_LIT>' ) <EOL> with pytest . raises ( ValueError ) as e : <EOL> compute ( concat ( t , u , axis = <NUM_LIT:1> ) , { t : t_table , u : u_table } , return_type = '<STR_LIT>' ) <EOL> assert "<STR_LIT>" in str ( e . value ) <EOL> def test_timedelta_arith ( sql_with_dts ) : <EOL> delta = timedelta ( days = <NUM_LIT:1> ) <EOL> dates = pd . Series ( pd . date_range ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> sym = symbol ( '<STR_LIT:s>' , discover ( dates ) ) <EOL> assert ( <EOL> compute ( sym + delta , sql_with_dts , return_type = pd . Series ) == dates + delta <EOL> ) . all ( ) <EOL> assert ( <EOL> compute ( sym - delta , sql_with_dts , return_type = pd . Series ) == dates - delta <EOL> ) . all ( ) <EOL> assert ( <EOL> compute ( sym - ( sym - delta ) , sql_with_dts , return_type = pd . Series ) == <EOL> dates - ( dates - delta ) <EOL> ) . all ( ) <EOL> @ pytest . mark . parametrize ( '<STR_LIT>' , ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> def test_timedelta_stat_reduction ( sql_with_timedeltas , func ) : <EOL> sym = symbol ( '<STR_LIT:s>' , discover ( sql_with_timedeltas ) ) <EOL> expr = getattr ( sym . N , func ) ( ) <EOL> deltas = pd . Series ( [ timedelta ( seconds = n ) for n in range ( <NUM_LIT:10> ) ] ) <EOL> expected = timedelta ( <EOL> seconds = getattr ( deltas . astype ( '<STR_LIT>' ) / <NUM_LIT> , func ) ( ddof = expr . unbiased ) <EOL> ) <EOL> assert compute ( expr , sql_with_timedeltas , return_type = timedelta ) == expected <EOL> def test_coerce_bool_and_sum ( sql ) : <EOL> sql = sql . data <EOL> n = sql . name <EOL> t = symbol ( n , discover ( sql ) ) <EOL> expr = ( t . B > <NUM_LIT:1.0> ) . coerce ( to = '<STR_LIT>' ) . sum ( ) <EOL> result = compute ( expr , sql ) . scalar ( ) <EOL> expected = compute ( t . B , sql , return_type = pd . Series ) . gt ( <NUM_LIT:1> ) . sum ( ) <EOL> assert result == expected <EOL> def test_distinct_on ( sql ) : <EOL> sql = sql . data <EOL> t = symbol ( '<STR_LIT:t>' , discover ( sql ) ) <EOL> computation = compute ( t [ [ '<STR_LIT:A>' , '<STR_LIT:B>' ] ] . sort ( '<STR_LIT:A>' ) . distinct ( '<STR_LIT:A>' ) , sql , return_type = '<STR_LIT>' ) <EOL> assert normalize ( str ( computation ) ) == normalize ( """<STR_LIT>""" . format ( tbl = sql . name ) ) <EOL> assert odo ( computation , tuple ) == ( ( '<STR_LIT:a>' , <NUM_LIT:1> ) , ( '<STR_LIT:b>' , <NUM_LIT:2> ) ) <EOL> def test_relabel_columns_over_selection ( big_sql ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( big_sql ) ) <EOL> result = compute ( t [ t [ '<STR_LIT:B>' ] == <NUM_LIT:2> ] . relabel ( B = u'<STR_LIT:b>' ) , <EOL> big_sql , return_type = pd . DataFrame ) <EOL> expected = pd . DataFrame ( [ [ '<STR_LIT:a>' , <NUM_LIT:2> ] ] , columns = [ u'<STR_LIT:A>' , u'<STR_LIT:b>' ] ) <EOL> tm . assert_frame_equal ( result , expected ) <EOL> def test_auto_join_field ( orders ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( orders ) ) <EOL> expr = t . product_id . color <EOL> result = compute ( expr , orders , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> def test_auto_join_projection ( orders ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( orders ) ) <EOL> expr = t . product_id [ [ '<STR_LIT>' , '<STR_LIT>' ] ] <EOL> result = compute ( expr , orders , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> @ pytest . mark . xfail <EOL> @ pytest . mark . parametrize ( '<STR_LIT>' , [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> def test_foreign_key_reduction ( orders , products , func ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( orders ) ) <EOL> expr = methodcaller ( func ) ( t . product_id . price ) <EOL> result = compute ( expr , orders , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" . format ( func ) <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> def test_foreign_key_chain ( fkey ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( fkey ) ) <EOL> expr = t . sym_id . main . data <EOL> result = compute ( expr , fkey , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> @ pytest . mark . xfail ( raises = AssertionError , <EOL> reason = '<STR_LIT>' ) <EOL> @ pytest . mark . parametrize ( '<STR_LIT>' , [ '<STR_LIT>' , [ '<STR_LIT>' ] ] ) <EOL> def test_foreign_key_group_by ( fkey , grouper ) : <EOL> t = symbol ( '<STR_LIT>' , discover ( fkey ) ) <EOL> expr = by ( t . sym_id [ grouper ] , avg_price = t . sym_id . price . mean ( ) ) <EOL> result = compute ( expr , fkey , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> @ pytest . mark . parametrize ( '<STR_LIT>' , [ '<STR_LIT>' , [ '<STR_LIT>' ] ] ) <EOL> def test_group_by_map ( fkey , grouper ) : <EOL> t = symbol ( '<STR_LIT>' , discover ( fkey ) ) <EOL> expr = by ( t [ grouper ] , id_count = t . size . count ( ) ) <EOL> result = compute ( expr , fkey , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> def test_foreign_key_isin ( fkey ) : <EOL> t = symbol ( '<STR_LIT>' , discover ( fkey ) ) <EOL> expr = t . sym_id . isin ( [ <NUM_LIT:1> , <NUM_LIT:2> ] ) <EOL> result = compute ( expr , fkey , return_type = '<STR_LIT>' ) <EOL> expected = """<STR_LIT>""" <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> @ pytest . mark . xfail ( raises = AssertionError , reason = '<STR_LIT>' ) <EOL> def test_foreign_key_merge_expression ( fkey ) : <EOL> from blaze import merge <EOL> t = symbol ( '<STR_LIT>' , discover ( fkey ) ) <EOL> expr = merge ( t . sym_id . sym , t . sym_id . main . data ) <EOL> expected = """<STR_LIT>""" <EOL> result = compute ( expr , fkey , return_type = '<STR_LIT>' ) <EOL> assert normalize ( str ( result ) ) == normalize ( expected ) <EOL> def test_join_type_promotion ( sqla , sqlb ) : <EOL> t , s = symbol ( sqla . name , discover ( sqla ) ) , symbol ( sqlb . name , discover ( sqlb ) ) <EOL> expr = join ( t , s , '<STR_LIT:B>' , how = '<STR_LIT>' ) <EOL> result = set ( map ( tuple , compute ( expr , { t : sqla , s : sqlb } , return_type = '<STR_LIT>' ) . execute ( ) . fetchall ( ) ) ) <EOL> expected = set ( [ ( <NUM_LIT:1> , '<STR_LIT:a>' , '<STR_LIT:a>' ) , ( <NUM_LIT:1> , None , '<STR_LIT:a>' ) ] ) <EOL> assert result == expected <EOL> @ pytest . mark . parametrize ( [ '<STR_LIT:n>' , '<STR_LIT>' ] , <EOL> [ ( <NUM_LIT:1> , '<STR_LIT:A>' ) , ( - <NUM_LIT:1> , '<STR_LIT:A>' ) , <EOL> ( <NUM_LIT:1> , '<STR_LIT:B>' ) , ( - <NUM_LIT:1> , '<STR_LIT:B>' ) , <EOL> ( <NUM_LIT:0> , '<STR_LIT:A>' ) , ( <NUM_LIT:0> , '<STR_LIT:B>' ) ] ) <EOL> def test_shift_on_column ( n , column , sql ) : <EOL> sql = sql . data <EOL> t = symbol ( '<STR_LIT:t>' , discover ( sql ) ) <EOL> expr = t [ column ] . shift ( n ) <EOL> result = compute ( expr , sql , return_type = pd . Series ) <EOL> expected = odo ( sql , pd . DataFrame ) [ column ] . shift ( n ) <EOL> tm . assert_series_equal ( result , expected ) <EOL> @ pytest . mark . parametrize ( '<STR_LIT:n>' , [ - <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:1> ] ) <EOL> def test_shift_arithmetic ( sql , n ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( sql ) ) <EOL> expr = t . B - t . B . shift ( n ) <EOL> result = compute ( expr , sql , return_type = pd . Series ) <EOL> df = odo ( sql , pd . DataFrame ) <EOL> expected = df . B - df . B . shift ( n ) <EOL> tm . assert_series_equal ( result , expected ) <EOL> def test_dist ( nyc ) : <EOL> def distance ( lat1 , lon1 , lat2 , lon2 , R = <NUM_LIT> ) : <EOL> dlon = radians ( lon2 - lon1 ) <EOL> dlat = radians ( lat2 - lat1 ) <EOL> a = sin ( dlat / <NUM_LIT> ) ** <NUM_LIT:2> + cos ( lat1 ) * cos ( lat2 ) * sin ( dlon / <NUM_LIT> ) ** <NUM_LIT:2> <EOL> return R * <NUM_LIT:2> * atan2 ( sqrt ( a ) , sqrt ( <NUM_LIT:1> - a ) ) <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> filtered = t [ <EOL> ( t . pickup_latitude >= <NUM_LIT> ) & <EOL> ( t . pickup_latitude <= <NUM_LIT> ) & <EOL> ( t . dropoff_latitude >= <NUM_LIT> ) & <EOL> ( t . dropoff_latitude <= <NUM_LIT> ) & <EOL> ( t . pickup_longitude >= - <NUM_LIT> ) & <EOL> ( t . pickup_longitude <= - <NUM_LIT> ) & <EOL> ( t . dropoff_longitude >= - <NUM_LIT> ) & <EOL> ( t . dropoff_longitude <= - <NUM_LIT> ) & <EOL> ( t . passenger_count < <NUM_LIT:6> ) <EOL> ] <EOL> dist = distance ( filtered . pickup_latitude , filtered . pickup_longitude , <EOL> filtered . dropoff_latitude , filtered . dropoff_longitude ) <EOL> transformed = transform ( filtered , dist = dist ) <EOL> assert ( <EOL> compute ( transformed . dist . max ( ) , nyc , return_type = float ) == <EOL> compute ( transformed . dist , nyc , return_type = pd . Series ) . max ( ) <EOL> ) <EOL> def test_multiple_columns_in_transform ( nyc ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> t = t [ <EOL> ( t . pickup_latitude >= <NUM_LIT> ) & <EOL> ( t . pickup_latitude <= <NUM_LIT> ) & <EOL> ( t . dropoff_latitude >= <NUM_LIT> ) & <EOL> ( t . dropoff_latitude <= <NUM_LIT> ) & <EOL> ( t . pickup_longitude >= - <NUM_LIT> ) & <EOL> ( t . pickup_longitude <= - <NUM_LIT> ) & <EOL> ( t . dropoff_longitude >= - <NUM_LIT> ) & <EOL> ( t . dropoff_longitude <= - <NUM_LIT> ) & <EOL> ( t . passenger_count < <NUM_LIT:6> ) <EOL> ] <EOL> hours = t . trip_time_in_secs . coerce ( '<STR_LIT>' ) / <NUM_LIT> <EOL> avg_speed_in_mph = t . trip_distance / hours <EOL> d = transform ( t , avg_speed_in_mph = avg_speed_in_mph , mycol = avg_speed_in_mph + <NUM_LIT:1> ) <EOL> df = compute ( d [ d . avg_speed_in_mph <= <NUM_LIT:200> ] , nyc , return_type = pd . DataFrame ) <EOL> assert not df . empty <EOL> def test_coerce_on_select ( nyc ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> t = t [ <EOL> ( t . pickup_latitude >= <NUM_LIT> ) & <EOL> ( t . pickup_latitude <= <NUM_LIT> ) & <EOL> ( t . dropoff_latitude >= <NUM_LIT> ) & <EOL> ( t . dropoff_latitude <= <NUM_LIT> ) & <EOL> ( t . pickup_longitude >= - <NUM_LIT> ) & <EOL> ( t . pickup_longitude <= - <NUM_LIT> ) & <EOL> ( t . dropoff_longitude >= - <NUM_LIT> ) & <EOL> ( t . dropoff_longitude <= - <NUM_LIT> ) & <EOL> ( t . passenger_count < <NUM_LIT:6> ) <EOL> ] <EOL> t = transform ( t , pass_count = t . passenger_count + <NUM_LIT:1> ) <EOL> result = compute ( t . pass_count . coerce ( '<STR_LIT>' ) , nyc , return_type = '<STR_LIT>' ) <EOL> s = odo ( result , pd . Series ) <EOL> expected = compute ( t , nyc , return_type = pd . DataFrame ) . passenger_count . astype ( '<STR_LIT>' ) + <NUM_LIT:1.0> <EOL> assert list ( s ) == list ( expected ) <EOL> def test_interactive_len ( sql ) : <EOL> t = data ( sql ) <EOL> assert len ( t ) == int ( t . count ( ) ) <EOL> def test_sample_n ( nyc ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> result = compute ( t . sample ( n = <NUM_LIT> ) , nyc , return_type = pd . DataFrame ) <EOL> assert len ( result ) == <NUM_LIT> <EOL> def test_sample_bounded ( nyc ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> nrows = compute ( t . nrows , nyc , return_type = int ) <EOL> result = compute ( t . sample ( n = <NUM_LIT:2> * nrows ) , nyc , return_type = pd . DataFrame ) <EOL> assert len ( result ) == nrows <EOL> def test_sample_frac ( nyc ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> result = compute ( t . sample ( frac = <NUM_LIT:0.5> ) , nyc , return_type = pd . DataFrame ) <EOL> num_rows = compute ( t . nrows , nyc , return_type = int ) <EOL> fractional , integral = math . modf ( num_rows * <NUM_LIT:0.5> ) <EOL> assert int ( integral + ( <NUM_LIT:0> if fractional < <NUM_LIT:0.5> else <NUM_LIT:1> ) ) == len ( result ) <EOL> def test_sample ( big_sql ) : <EOL> nn = symbol ( '<STR_LIT>' , discover ( big_sql ) ) <EOL> nrows = odo ( compute ( nn . nrows , big_sql ) , int ) <EOL> result = compute ( nn . sample ( n = nrows // <NUM_LIT:2> ) , big_sql , return_type = pd . DataFrame ) <EOL> assert len ( result ) == nrows // <NUM_LIT:2> <EOL> result2 = compute ( nn . sample ( frac = <NUM_LIT:0.5> ) , big_sql , return_type = pd . DataFrame ) <EOL> assert len ( result ) == len ( result2 ) <EOL> def test_core_compute ( nyc ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( nyc ) ) <EOL> assert isinstance ( compute ( t , nyc , return_type = '<STR_LIT>' ) , pd . DataFrame ) <EOL> assert isinstance ( compute ( t . passenger_count , nyc , return_type = '<STR_LIT>' ) , pd . Series ) <EOL> assert iscorescalar ( compute ( t . passenger_count . mean ( ) , nyc , return_type = '<STR_LIT>' ) ) <EOL> assert isinstance ( compute ( t , nyc , return_type = list ) , list ) <EOL> @ pytest . fixture <EOL> def gl_data ( sql_two_tables ) : <EOL> u_data , t_data = sql_two_tables <EOL> return data ( odo ( [ ( <NUM_LIT:1> , ) ] , u_data ) ) , data ( odo ( [ ( <NUM_LIT:2> , ) ] , t_data ) ) <EOL> def test_greatest ( gl_data ) : <EOL> u , t = gl_data <EOL> assert odo ( greatest ( u . a . max ( ) , t . a . max ( ) ) , int ) == <NUM_LIT:2> <EOL> def test_least ( gl_data ) : <EOL> u , t = gl_data <EOL> assert odo ( least ( u . a . max ( ) , t . a . max ( ) ) , int ) == <NUM_LIT:1> <EOL> def test_coalesce ( sqla ) : <EOL> t = symbol ( '<STR_LIT:t>' , discover ( sqla ) ) <EOL> assert ( <EOL> compute ( coalesce ( t . B , - <NUM_LIT:1> ) , { t : sqla } , return_type = list ) == <EOL> [ ( <NUM_LIT:1> , ) , ( <NUM_LIT:1> , ) , ( - <NUM_LIT:1> , ) ] <EOL> ) <EOL> assert ( <EOL> compute ( coalesce ( t . A , '<STR_LIT:z>' ) , { t : sqla } , return_type = list ) == <EOL> [ ( '<STR_LIT:a>' , ) , ( '<STR_LIT:z>' , ) , ( '<STR_LIT:c>' , ) ] <EOL> ) </s>
