<s> from __future__ import absolute_import , division , print_function <EOL> from multipledispatch import MDNotImplementedError <EOL> from odo import Chunks , convert , into <EOL> from collections import Iterator , Iterable <EOL> from toolz import curry , concat <EOL> from datashape . dispatch import dispatch <EOL> import pandas as pd <EOL> import numpy as np <EOL> from . . expr import Head , ElemWise , Distinct , Symbol , Expr , path <EOL> from . . expr . split import split <EOL> from . core import compute <EOL> from . pmap import get_default_pmap <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> Cheap = ( Head , ElemWise , Distinct , Symbol ) <EOL> @ dispatch ( Head , Chunks ) <EOL> def pre_compute ( expr , data , ** kwargs ) : <EOL> leaf = expr . _leaves ( ) [ <NUM_LIT:0> ] <EOL> if all ( isinstance ( e , Cheap ) for e in path ( expr , leaf ) ) : <EOL> return convert ( Iterator , data ) <EOL> else : <EOL> raise MDNotImplementedError ( ) <EOL> def compute_chunk ( chunk , chunk_expr , part ) : <EOL> return compute ( chunk_expr , { chunk : part } ) <EOL> @ dispatch ( Expr , Chunks ) <EOL> def compute_down ( expr , data , map = None , ** kwargs ) : <EOL> if map is None : <EOL> map = get_default_pmap ( ) <EOL> leaf = expr . _leaves ( ) [ <NUM_LIT:0> ] <EOL> ( chunk , chunk_expr ) , ( agg , agg_expr ) = split ( leaf , expr ) <EOL> parts = list ( map ( curry ( compute_chunk , chunk , chunk_expr ) , data ) ) <EOL> if isinstance ( parts [ <NUM_LIT:0> ] , np . ndarray ) : <EOL> intermediate = np . concatenate ( parts ) <EOL> elif isinstance ( parts [ <NUM_LIT:0> ] , pd . DataFrame ) : <EOL> intermediate = pd . concat ( parts ) <EOL> elif isinstance ( parts [ <NUM_LIT:0> ] , ( Iterable , Iterator ) ) : <EOL> intermediate = list ( concat ( parts ) ) <EOL> return compute ( agg_expr , { agg : intermediate } ) <EOL> Cheap = ( Head , ElemWise , Distinct , Symbol ) <EOL> @ dispatch ( Head , Chunks ) <EOL> def compute_down ( expr , data , ** kwargs ) : <EOL> leaf = expr . _leaves ( ) [ <NUM_LIT:0> ] <EOL> if all ( isinstance ( e , Cheap ) for e in path ( expr , leaf ) ) : <EOL> return compute ( expr , { leaf : into ( Iterator , data ) } , ** kwargs ) <EOL> else : <EOL> raise MDNotImplementedError ( ) </s>
