<s> import pytest <EOL> from operator import getitem <EOL> from toolz import merge <EOL> from dask . dataframe . optimize import dataframe_from_ctable <EOL> import dask . dataframe as dd <EOL> import pandas as pd <EOL> dsk = { ( '<STR_LIT:x>' , <NUM_LIT:0> ) : pd . DataFrame ( { '<STR_LIT:a>' : [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] , '<STR_LIT:b>' : [ <NUM_LIT:4> , <NUM_LIT:5> , <NUM_LIT:6> ] } , <EOL> index = [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:3> ] ) , <EOL> ( '<STR_LIT:x>' , <NUM_LIT:1> ) : pd . DataFrame ( { '<STR_LIT:a>' : [ <NUM_LIT:4> , <NUM_LIT:5> , <NUM_LIT:6> ] , '<STR_LIT:b>' : [ <NUM_LIT:3> , <NUM_LIT:2> , <NUM_LIT:1> ] } , <EOL> index = [ <NUM_LIT:5> , <NUM_LIT:6> , <NUM_LIT:8> ] ) , <EOL> ( '<STR_LIT:x>' , <NUM_LIT:2> ) : pd . DataFrame ( { '<STR_LIT:a>' : [ <NUM_LIT:7> , <NUM_LIT:8> , <NUM_LIT:9> ] , '<STR_LIT:b>' : [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] } , <EOL> index = [ <NUM_LIT:9> , <NUM_LIT:9> , <NUM_LIT:9> ] ) } <EOL> dfs = list ( dsk . values ( ) ) <EOL> def test_column_optimizations_with_bcolz_and_rewrite ( ) : <EOL> try : <EOL> import bcolz <EOL> except ImportError : <EOL> return <EOL> bc = bcolz . ctable ( [ [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] , [ <NUM_LIT:10> , <NUM_LIT:20> , <NUM_LIT:30> ] ] , names = [ '<STR_LIT:a>' , '<STR_LIT:b>' ] ) <EOL> func = lambda x : x <EOL> for cols in [ None , '<STR_LIT:abc>' , [ '<STR_LIT:abc>' ] ] : <EOL> dsk2 = merge ( dict ( ( ( '<STR_LIT:x>' , i ) , <EOL> ( dataframe_from_ctable , bc , slice ( <NUM_LIT:0> , <NUM_LIT:2> ) , cols , { } ) ) <EOL> for i in [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ) , <EOL> dict ( ( ( '<STR_LIT:y>' , i ) , <EOL> ( getitem , ( '<STR_LIT:x>' , i ) , ( list , [ '<STR_LIT:a>' , '<STR_LIT:b>' ] ) ) ) <EOL> for i in [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ) ) <EOL> expected = dict ( ( ( '<STR_LIT:y>' , i ) , ( dataframe_from_ctable , <EOL> bc , slice ( <NUM_LIT:0> , <NUM_LIT:2> ) , ( list , [ '<STR_LIT:a>' , '<STR_LIT:b>' ] ) , { } ) ) <EOL> for i in [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ) <EOL> result = dd . optimize ( dsk2 , [ ( '<STR_LIT:y>' , i ) for i in [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ] ) <EOL> assert result == expected <EOL> @ pytest . mark . xfail ( reason = "<STR_LIT>" ) <EOL> def test_castra_column_store ( ) : <EOL> try : <EOL> from castra import Castra <EOL> except ImportError : <EOL> return <EOL> df = pd . DataFrame ( { '<STR_LIT:x>' : [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] , '<STR_LIT:y>' : [ <NUM_LIT:4> , <NUM_LIT:5> , <NUM_LIT:6> ] } ) <EOL> with Castra ( template = df ) as c : <EOL> c . extend ( df ) <EOL> df = c . to_dask ( ) <EOL> df2 = df [ [ '<STR_LIT:x>' ] ] <EOL> dsk = dd . optimize ( df2 . dask , df2 . _keys ( ) ) <EOL> assert dsk == { ( df2 . _name , <NUM_LIT:0> ) : ( Castra . load_partition , c , '<STR_LIT>' , <EOL> ( list , [ '<STR_LIT:x>' ] ) ) } <EOL> df3 = df . index <EOL> dsk = dd . optimize ( df3 . dask , df3 . _keys ( ) ) <EOL> assert dsk == { ( df3 . _name , <NUM_LIT:0> ) : ( Castra . load_index , c , '<STR_LIT>' ) } </s>
