def __init__(self):

# Python 2.x

super(SparkCSVToAvro, self).__init__()

# Python 3.x

# super().__init__()

# logging.config.fileConfig(os.path.join(libdir, 'resources', 'logging.conf'))

# log = logging.getLogger(self.__class__.__name__)

self.verbose_default = 2

self.timeout_default = 86400

self.schema = None

self.types_mapping = {}

# dynamically generate types mapping from available types in PySpark

from pyspark.sql import types # pylint: disable=wrong-import-position,import-error

for _ in dir(types):

if _.endswith('Type'):

self.types_mapping[_[:-4].lower()] = _

if 'integer' in self.types_mapping:

# because I like typing int better than integer

self.types_mapping['int'] = self.types_mapping['integer']

# @override

def add_options(self):

self.add_opt('-c', '--csv', metavar='<file/dir>',

help='CSV input file/dir ($CSV)',

default=getenv('CSV'))

self.add_opt('-a', '--avro-dir', metavar='<dir>',

help='Avro output dir ($AVRODIR)',

default=getenv('AVRODIR'))

self.add_opt('-e', '--has-header', action='store_true',

help='CSV has header. Infers schema if --schema is not given in which case all ' +

"types are assumed to be 'string'. Must specify --schema to override this")

self.add_opt('-s', '--schema', metavar='name:type,name2:type2,...',

help="Schema for CSV. Types default to 'string'. Possible types are: %s" \

% ', '.join(sorted(self.types_mapping)))

def parse_args(self):

self.no_args()

if not self.get_opt('csv'):

self.usage('--csv not defined')

if not self.get_opt('avro_dir'):

self.usage('--avro-dir not defined')

if not (self.get_opt('has_header') or self.get_opt('schema')):

self.usage('must specify either --has-header or --schema')

# no longer mutually exclusive now this support schema override

# if self.get_opt('has_header') and self.get_opt('schema'):

# self.usage('--has-header and --schema are mutually exclusive')

def run(self):

csv_file = self.get_opt('csv')

avro_dir = self.get_opt('avro_dir')

has_header = self.get_opt('has_header')

# I don't know why the Spark guys made this a string instead of a bool

header_str = 'false'

if has_header:

header_str = 'true'

schema = self.get_opt('schema')

# let Spark fail if csv/avro dir aren't available

# can't check paths exist as want to remain generically portable

# to HDFS, local filesystm or any other uri scheme Spark supports

log.info("CSV Source: %s" % csv_file)

log.info("Avro Destination: %s" % avro_dir)

if schema:

def get_type(arg):

arg = str(arg).lower()

if arg not in self.types_mapping:

self.usage("invalid type '%s' defined in --schema, must be one of: %s"

% (arg, ', '.join(sorted(self.types_mapping.keys()))))

# return self.types_mapping[arg]

module = __import__('pyspark.sql.types', globals(), locals(), ['types'], -1)

class_ = getattr(module, self.types_mapping[arg])

_ = class_()

return _

def create_struct(arg):

name = str(arg).strip()

data_type = 'string'

if ':' in arg:

(name, data_type) = arg.split(':', 1)

data_class = get_type(data_type)

return StructField(name, data_class, True)

# see https://github.com/databricks/spark-csv#python-api

self.schema = StructType([create_struct(_) for _ in schema.split(',')])

log.info('generated CSV => Spark schema')

conf = SparkConf().setAppName('HS PySpark CSV => Avro')

sc = SparkContext(conf=conf) # pylint: disable=invalid-name

if self.verbose < 3 and 'setLogLevel' in dir(sc):

sc.setLogLevel('WARN')

sqlContext = SQLContext(sc) # pylint: disable=invalid-name

spark_version = sc.version

log.info('Spark version detected as %s' % spark_version)

if not isVersionLax(spark_version):

die("Spark version couldn't be determined. " + support_msg('pytools'))

# pylint: disable=invalid-name

df = None

if isMinVersion(spark_version, 1.4):

if has_header and not schema:

log.info('inferring schema from CSV headers')

df = sqlContext.read.format('com.databricks.spark.csv')\

.options(header=header_str, inferschema='true')\

.load(csv_file)

else:

log.info('using explicitly defined schema')

schema = self.schema

df = sqlContext.read\

.format('com.databricks.spark.csv')\

.options(header=header_str)\

.load(csv_file, schema=schema)

else:

die('Spark <= 1.3 is not supported due to avro dependency, sorry! ' + \

'I may change this on request but prefer people just upgrade')

# log.warn('running legacy code for Spark <= 1.3')

# if has_header and not schema:

# log.info('inferring schema from CSV headers')

# df = sqlContext.load(source="com.databricks.spark.csv", path=csv_file,

# header=header_str, inferSchema='true')

# elif self.schema:

# log.info('using explicitly defined schema')

# df = sqlContext.load(source="com.databricks.spark.csv", path=csv_file,

# header=header_str, schema=self.schema)

# else:

# die('no header and no schema, caught late')

# this doesn't work in Spark <= 1.3 and the github docs don't mention the older methods for writing avro using

# the databricks avro driver