<s> from __future__ import absolute_import , division , print_function <EOL> from toolz import memoize , merge , partition_all <EOL> from multipledispatch import MDNotImplementedError <EOL> import re <EOL> import os <EOL> from . csv import CSV <EOL> from . json import JSON , JSONLines <EOL> from . text import TextFile <EOL> import pandas as pd <EOL> import uuid <EOL> import datashape <EOL> import sqlalchemy as sa <EOL> from datashape import discover <EOL> from datashape import coretypes as ct <EOL> from collections import namedtuple , Iterator <EOL> from contextlib import contextmanager <EOL> from . ssh import SSH <EOL> from . sql import metadata_of_engine <EOL> from . . utils import tmpfile , sample , ignoring , raises <EOL> from . . temp import Temp <EOL> from . . append import append <EOL> from . . convert import convert <EOL> from . . chunks import chunks <EOL> from . . resource import resource <EOL> from . . directory import Directory <EOL> from . . compatibility import unicode <EOL> with ignoring ( ImportError ) : <EOL> from pywebhdfs . webhdfs import PyWebHdfsClient <EOL> from pywebhdfs . errors import FileNotFound <EOL> class _HDFS ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> hdfs = kwargs . get ( '<STR_LIT>' , None ) <EOL> host = kwargs . get ( '<STR_LIT:host>' , None ) <EOL> user = ( kwargs . get ( '<STR_LIT>' ) or <EOL> kwargs . get ( '<STR_LIT:user>' ) or <EOL> kwargs . get ( '<STR_LIT:username>' ) or None ) <EOL> port = str ( kwargs . get ( '<STR_LIT:port>' , <NUM_LIT> ) ) <EOL> if not hdfs and ( host and port and user ) : <EOL> hdfs = PyWebHdfsClient ( host = host , port = str ( port ) , <EOL> user_name = user ) <EOL> if hdfs is None : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> self . hdfs = hdfs <EOL> self . subtype . __init__ ( self , * args , ** kwargs ) <EOL> @ memoize <EOL> def HDFS ( cls ) : <EOL> return type ( '<STR_LIT>' % cls . __name__ , ( _HDFS , cls ) , { '<STR_LIT>' : cls } ) <EOL> @ sample . register ( _HDFS ) <EOL> @ contextmanager <EOL> def sample_hdfs_csv ( data , length = <NUM_LIT> ) : <EOL> sample = data . hdfs . read_file ( data . path . lstrip ( '<STR_LIT:/>' ) , length = length ) <EOL> with tmpfile ( data . canonical_extension ) as fn : <EOL> with open ( fn , '<STR_LIT:w>' ) as f : <EOL> f . write ( sample ) <EOL> yield fn <EOL> @ discover . register ( HDFS ( JSON ) ) <EOL> @ discover . register ( HDFS ( JSONLines ) ) <EOL> @ discover . register ( HDFS ( TextFile ) ) <EOL> def discover_hdfs_file ( data , ** kwargs ) : <EOL> with sample ( data ) as fn : <EOL> result = discover ( data . subtype ( fn , ** kwargs ) ) <EOL> return result <EOL> @ discover . register ( HDFS ( CSV ) ) <EOL> def discover_hdfs_csv ( data , length = <NUM_LIT> , ** kwargs ) : <EOL> with sample ( data , length = length ) as fn : <EOL> result = discover ( CSV ( fn , encoding = data . encoding , <EOL> dialect = data . dialect , <EOL> has_header = data . has_header ) ) <EOL> return result <EOL> @ sample . register ( HDFS ( Directory ( CSV ) ) ) <EOL> @ contextmanager <EOL> def sample_hdfs_directory_csv ( data , ** kwargs ) : <EOL> files = data . hdfs . list_dir ( data . path . lstrip ( '<STR_LIT:/>' ) ) <EOL> one_file = data . path + '<STR_LIT:/>' + files [ '<STR_LIT>' ] [ '<STR_LIT>' ] [ <NUM_LIT:0> ] [ '<STR_LIT>' ] <EOL> csv = HDFS ( CSV ) ( one_file , hdfs = data . hdfs ) <EOL> with sample ( csv , ** kwargs ) as fn : <EOL> yield fn <EOL> @ discover . register ( HDFS ( Directory ( CSV ) ) ) <EOL> def discover_hdfs_directory ( data , length = <NUM_LIT> , ** kwargs ) : <EOL> with sample ( data , length = length ) as fn : <EOL> o = data . container ( fn , ** data . kwargs ) <EOL> result = discover ( o ) <EOL> return result <EOL> """<STR_LIT>""" <EOL> TableProxy = namedtuple ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> """<STR_LIT>""" <EOL> @ resource . register ( '<STR_LIT>' , priority = <NUM_LIT:16> ) <EOL> def resource_hive_table ( uri , stored_as = '<STR_LIT>' , external = True , dshape = None , ** kwargs ) : <EOL> if dshape : <EOL> dshape = datashape . dshape ( dshape ) <EOL> uri , table = uri . split ( '<STR_LIT>' ) <EOL> engine = resource ( uri ) <EOL> metadata = metadata_of_engine ( engine ) <EOL> with ignoring ( sa . exc . NoSuchTableError ) : <EOL> return sa . Table ( table , metadata , autoload = True , <EOL> autoload_with = engine ) <EOL> if dshape and ( not external or external and kwargs . get ( '<STR_LIT:path>' ) ) : <EOL> table_type = '<STR_LIT>' if external else '<STR_LIT>' <EOL> statement = create_hive_statement ( table , dshape , <EOL> db_name = engine . url . database , stored_as = stored_as , <EOL> table_type = table_type , ** kwargs ) <EOL> with engine . connect ( ) as conn : <EOL> conn . execute ( statement ) <EOL> return sa . Table ( table , metadata , autoload = True , <EOL> autoload_with = engine ) <EOL> else : <EOL> return TableProxy ( engine , table , stored_as ) <EOL> hive_types = { <EOL> ct . int8 : '<STR_LIT>' , <EOL> ct . int16 : '<STR_LIT>' , <EOL> ct . int32 : '<STR_LIT>' , <EOL> ct . int64 : '<STR_LIT>' , <EOL> ct . float32 : '<STR_LIT>' , <EOL> ct . float64 : '<STR_LIT>' , <EOL> ct . date_ : '<STR_LIT>' , <EOL> ct . datetime_ : '<STR_LIT>' , <EOL> ct . string : '<STR_LIT>' , <EOL> ct . bool_ : '<STR_LIT>' } <EOL> def dshape_to_hive ( ds ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( ds , ( str , unicode ) ) : <EOL> ds = datashape . dshape ( ds ) <EOL> if isinstance ( ds , ct . DataShape ) : <EOL> ds = ds . measure <EOL> if isinstance ( ds , ct . Record ) : <EOL> columns = [ ( name , dshape_to_hive ( typ ) ) <EOL> for name , typ in zip ( ds . measure . names , ds . measure . types ) ] <EOL> return [ '<STR_LIT>' % col for col in columns ] <EOL> if isinstance ( ds , ct . Option ) : <EOL> ds = ds . ty <EOL> if isinstance ( ds , ct . String ) : <EOL> if ds . fixlen : <EOL> return '<STR_LIT>' % ds . fixlen <EOL> else : <EOL> return '<STR_LIT>' <EOL> if ds in hive_types : <EOL> return hive_types [ ds ] <EOL> raise NotImplementedError ( "<STR_LIT>" % ds ) <EOL> def create_hive_statement ( tbl_name , dshape , path = None , table_type = '<STR_LIT>' , <EOL> db_name = '<STR_LIT:default>' , stored_as = '<STR_LIT>' , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if db_name : <EOL> db_name = db_name + '<STR_LIT:.>' <EOL> if not table_type : <EOL> table_type = '<STR_LIT>' <EOL> assert isinstance ( dshape . measure , ct . Record ) <EOL> columns = dshape_to_hive ( dshape ) <EOL> column_text = '<STR_LIT>' . join ( '<STR_LIT>' % tuple ( col . split ( ) ) <EOL> for col in columns ) . lstrip ( ) <EOL> statement = """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> statement += """<STR_LIT>""" <EOL> statement += """<STR_LIT>""" <EOL> if path : <EOL> statement = statement + """<STR_LIT>""" <EOL> if kwargs . get ( '<STR_LIT>' ) : <EOL> statement = statement + """<STR_LIT>""" <EOL> return statement . format ( ** merge ( kwargs , locals ( ) ) ) . strip ( '<STR_LIT:\n>' ) <EOL> """<STR_LIT>""" <EOL> @ append . register ( TableProxy , HDFS ( Directory ( CSV ) ) ) <EOL> def create_new_hive_table_from_csv ( tbl , data , dshape = None , path = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> table_type = '<STR_LIT>' <EOL> if not dshape : <EOL> dshape = discover ( data ) <EOL> if tbl . engine . dialect . name == '<STR_LIT>' : <EOL> statement = create_hive_statement ( tbl . name , dshape , <EOL> path = data . path , <EOL> db_name = str ( tbl . engine . url ) . split ( '<STR_LIT:/>' ) [ - <NUM_LIT:1> ] , <EOL> table_type = table_type , <EOL> ** dialect_of ( data ) ) <EOL> else : <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> "<STR_LIT>" % tbl . engine . dialect . name ) <EOL> with tbl . engine . connect ( ) as conn : <EOL> conn . execute ( statement ) <EOL> metadata = metadata_of_engine ( tbl . engine ) <EOL> tbl2 = sa . Table ( tbl . name , metadata , autoload = True , <EOL> autoload_with = tbl . engine ) <EOL> return tbl2 <EOL> @ append . register ( TableProxy , ( HDFS ( CSV ) , Temp ( HDFS ( CSV ) ) ) ) <EOL> def create_new_hive_table_from_csv_file ( tbl , data , dshape = None , path = None , ** kwargs ) : <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> @ append . register ( TableProxy , ( SSH ( CSV ) , SSH ( Directory ( CSV ) ) ) ) <EOL> def append_remote_csv_to_new_table ( tbl , data , dshape = None , ** kwargs ) : <EOL> if not dshape : <EOL> dshape = discover ( data ) <EOL> if tbl . engine . dialect . name == '<STR_LIT>' : <EOL> statement = create_hive_statement ( tbl . name , dshape , <EOL> db_name = str ( tbl . engine . url ) . split ( '<STR_LIT:/>' ) [ - <NUM_LIT:1> ] , <EOL> ** dialect_of ( data ) ) <EOL> else : <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> "<STR_LIT>" % tbl . engine . dialect . name ) <EOL> with tbl . engine . connect ( ) as conn : <EOL> conn . execute ( statement ) <EOL> metadata = metadata_of_engine ( tbl . engine ) <EOL> tbl2 = sa . Table ( tbl . name , metadata , autoload = True , <EOL> autoload_with = tbl . engine ) <EOL> return append ( tbl2 , data , ** kwargs ) <EOL> @ append . register ( TableProxy , object ) <EOL> def append_anything_to_tableproxy ( tbl , data , ** kwargs ) : <EOL> return append ( tbl , convert ( Temp ( SSH ( CSV ) ) , data , ** kwargs ) , ** kwargs ) <EOL> @ append . register ( sa . Table , ( SSH ( CSV ) , SSH ( Directory ( CSV ) ) ) ) <EOL> def append_remote_csv_to_table ( tbl , csv , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> path = csv . path <EOL> if path [ <NUM_LIT:0> ] != '<STR_LIT:/>' : <EOL> path = '<STR_LIT>' % ( csv . auth [ '<STR_LIT:username>' ] , csv . path ) <EOL> if tbl . bind . dialect . name == '<STR_LIT>' : <EOL> statement = ( '<STR_LIT>' <EOL> % { '<STR_LIT:path>' : path , '<STR_LIT>' : tbl . name } ) <EOL> else : <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> "<STR_LIT>" % tbl . engine . dialect . name ) <EOL> with tbl . bind . connect ( ) as conn : <EOL> conn . execute ( statement ) <EOL> return tbl <EOL> @ append . register ( sa . Table , ( HDFS ( CSV ) , HDFS ( Directory ( CSV ) ) , Temp ( HDFS ( CSV ) ) ) ) <EOL> def append_hdfs_csv_to_table ( tbl , csv , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if tbl . bind . dialect . name != '<STR_LIT>' : <EOL> raise NotImplementedError ( "<STR_LIT>" <EOL> "<STR_LIT>" % tbl . engine . dialect . name ) <EOL> statement = ( '<STR_LIT>' <EOL> % { '<STR_LIT:path>' : csv . path , '<STR_LIT>' : tbl . name } ) <EOL> with tbl . bind . connect ( ) as conn : <EOL> conn . execute ( statement ) <EOL> return tbl <EOL> @ append . register ( TextFile , HDFS ( TextFile ) ) <EOL> @ append . register ( JSONLines , HDFS ( JSONLines ) ) <EOL> @ append . register ( JSON , HDFS ( JSON ) ) <EOL> @ append . register ( CSV , HDFS ( CSV ) ) <EOL> def append_hdfs_file_to_local ( target , source , ** kwargs ) : <EOL> text = source . hdfs . read_file ( source . path . lstrip ( '<STR_LIT:/>' ) , ** kwargs ) <EOL> with open ( target . path , '<STR_LIT:w>' ) as f : <EOL> f . write ( text ) <EOL> return target <EOL> @ convert . register ( Temp ( TextFile ) , HDFS ( TextFile ) ) <EOL> @ convert . register ( Temp ( JSONLines ) , HDFS ( JSONLines ) ) <EOL> @ convert . register ( Temp ( JSON ) , HDFS ( JSON ) ) <EOL> @ convert . register ( Temp ( CSV ) , HDFS ( CSV ) ) <EOL> def convert_hdfs_file_to_temp_local ( source , ** kwargs ) : <EOL> ext = os . path . splitext ( source . path ) [ <NUM_LIT:1> ] . strip ( '<STR_LIT:.>' ) <EOL> fn = '<STR_LIT>' % ( str ( uuid . uuid1 ( ) ) , ext ) <EOL> tmp = Temp ( source . subtype ) ( fn ) <EOL> return append ( tmp , source , ** kwargs ) <EOL> @ append . register ( HDFS ( TextFile ) , TextFile ) <EOL> @ append . register ( HDFS ( JSONLines ) , JSONLines ) <EOL> @ append . register ( HDFS ( JSON ) , JSON ) <EOL> @ append . register ( HDFS ( CSV ) , ( CSV , Temp ( CSV ) ) ) <EOL> def append_local_file_to_hdfs ( target , source , blocksize = <NUM_LIT> , ** kwargs ) : <EOL> if raises ( FileNotFound , <EOL> lambda : target . hdfs . list_dir ( target . path . lstrip ( '<STR_LIT:/>' ) ) ) : <EOL> target . hdfs . create_file ( target . path . lstrip ( '<STR_LIT:/>' ) , '<STR_LIT>' ) <EOL> with open ( source . path , '<STR_LIT:r>' ) as f : <EOL> blocks = partition_all ( blocksize , f ) <EOL> for block in blocks : <EOL> target . hdfs . append_file ( target . path . lstrip ( '<STR_LIT:/>' ) , '<STR_LIT>' . join ( block ) ) <EOL> return target <EOL> import csv <EOL> sniffer = csv . Sniffer ( ) <EOL> def dialect_of ( data , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> keys = set ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:strict>' , '<STR_LIT>' ] ) <EOL> if isinstance ( data , ( HDFS ( CSV ) , SSH ( CSV ) ) ) : <EOL> with sample ( data ) as fn : <EOL> d = dialect_of ( CSV ( fn , ** data . dialect ) ) <EOL> elif isinstance ( data , ( HDFS ( Directory ( CSV ) ) , SSH ( Directory ( CSV ) ) ) ) : <EOL> with sample ( data ) as fn : <EOL> d = dialect_of ( CSV ( fn , ** data . kwargs ) ) <EOL> elif isinstance ( data , Directory ( CSV ) ) : <EOL> d = dialect_of ( next ( data ) ) <EOL> else : <EOL> assert isinstance ( data , CSV ) <EOL> with open ( data . path , '<STR_LIT:r>' ) as f : <EOL> text = f . read ( ) <EOL> result = dict ( ) <EOL> d = sniffer . sniff ( text ) <EOL> d = dict ( ( k , getattr ( d , k ) ) for k in keys if hasattr ( d , k ) ) <EOL> if data . has_header is None : <EOL> d [ '<STR_LIT>' ] = sniffer . has_header ( text ) <EOL> else : <EOL> d [ '<STR_LIT>' ] = data . has_header <EOL> d . update ( data . dialect ) <EOL> d . update ( kwargs ) <EOL> d = dict ( ( k , v ) for k , v in d . items ( ) if k in keys ) <EOL> return d <EOL> types_by_extension = { '<STR_LIT>' : CSV , '<STR_LIT>' : JSONLines , '<STR_LIT>' : TextFile , <EOL> '<STR_LIT>' : TextFile } <EOL> hdfs_pattern = '<STR_LIT>' <EOL> @ resource . register ( '<STR_LIT>' , priority = <NUM_LIT:16> ) <EOL> def resource_hdfs ( uri , ** kwargs ) : <EOL> if '<STR_LIT>' in uri : <EOL> uri = uri [ len ( '<STR_LIT>' ) : ] <EOL> d = re . match ( hdfs_pattern , uri ) . groupdict ( ) <EOL> d = dict ( ( k , v ) for k , v in d . items ( ) if v is not None ) <EOL> path = d . pop ( '<STR_LIT:path>' ) <EOL> kwargs . update ( d ) <EOL> try : <EOL> subtype = types_by_extension [ path . split ( '<STR_LIT:.>' ) [ - <NUM_LIT:1> ] ] <EOL> if '<STR_LIT:*>' in path : <EOL> subtype = Directory ( subtype ) <EOL> path = path . rsplit ( '<STR_LIT:/>' , <NUM_LIT:1> ) [ <NUM_LIT:0> ] + '<STR_LIT:/>' <EOL> except KeyError : <EOL> subtype = type ( resource ( path ) ) <EOL> return HDFS ( subtype ) ( path , ** kwargs ) <EOL> @ append . register ( HDFS ( TextFile ) , ( Iterator , object ) ) <EOL> @ append . register ( HDFS ( JSON ) , ( list , object ) ) <EOL> @ append . register ( HDFS ( CSV ) , ( chunks ( pd . DataFrame ) , pd . DataFrame , object ) ) <EOL> def append_object_to_hdfs ( target , source , ** kwargs ) : <EOL> tmp = convert ( Temp ( target . subtype ) , source , ** kwargs ) <EOL> return append ( target , tmp , ** kwargs ) <EOL> @ append . register ( HDFS ( TextFile ) , SSH ( TextFile ) ) <EOL> @ append . register ( HDFS ( JSONLines ) , SSH ( JSONLines ) ) <EOL> @ append . register ( HDFS ( JSON ) , SSH ( JSON ) ) <EOL> @ append . register ( HDFS ( CSV ) , SSH ( CSV ) ) <EOL> def append_remote_file_to_hdfs ( target , source , ** kwargs ) : <EOL> raise MDNotImplementedError ( ) <EOL> @ append . register ( HDFS ( TextFile ) , HDFS ( TextFile ) ) <EOL> @ append . register ( HDFS ( JSONLines ) , HDFS ( JSONLines ) ) <EOL> @ append . register ( HDFS ( JSON ) , HDFS ( JSON ) ) <EOL> @ append . register ( HDFS ( CSV ) , HDFS ( CSV ) ) <EOL> def append_hdfs_file_to_hdfs_file ( target , source , ** kwargs ) : <EOL> raise MDNotImplementedError ( ) <EOL> @ append . register ( SSH ( TextFile ) , HDFS ( TextFile ) ) <EOL> @ append . register ( SSH ( JSONLines ) , HDFS ( JSONLines ) ) <EOL> @ append . register ( SSH ( JSON ) , HDFS ( JSON ) ) <EOL> @ append . register ( SSH ( CSV ) , HDFS ( CSV ) ) <EOL> def append_hdfs_file_to_remote ( target , source , ** kwargs ) : <EOL> raise MDNotImplementedError ( ) </s>
