forked from HariSekhon/DevOps-Python-tools
-
Notifications
You must be signed in to change notification settings - Fork 0
/
spark_json_to_avro.py
executable file
·119 lines (102 loc) · 4.62 KB
/
spark_json_to_avro.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#!/usr/bin/env python
# vim:ts=4:sts=4:sw=4:et
#
# Author: Hari Sekhon
# Date: 2015-11-03 21:38:52 +0000 (Tue, 03 Nov 2015)
#
# https://github.com/harisekhon/pytools
#
# License: see accompanying Hari Sekhon LICENSE file
#
# If you're using my code you're welcome to connect with me on LinkedIn and optionally send me feedback
# to help improve or steer this or other code I publish
#
# http://www.linkedin.com/in/harisekhon
#
"""
PySpark program to convert JSON file(s) to Avro
Written to work across Python 2.x, supports Spark 1.4+
Tested on Spark 1.4.0, 1.6.0
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# from __future__ import unicode_literals
import os
import sys
libdir = os.path.join(os.path.dirname(__file__), 'pylib')
sys.path.append(libdir)
try:
# pylint: disable=wrong-import-position
from harisekhon.utils import log, isMinVersion, support_msg, isVersionLax, die, getenv, pyspark_path
from harisekhon import CLI
except ImportError as _:
print('module import failed: %s' % _, file=sys.stderr)
print("Did you remember to build the project by running 'make'?", file=sys.stderr)
print("Alternatively perhaps you tried to copy this program out without it's adjacent libraries?", file=sys.stderr)
sys.exit(4)
# com.databricks:spark-avro_2.10:2.0.1 - 2.0.1 is for Spark 1.4+
# you can edit this bit if you need to run it on Spark 1.3:
# https://github.com/databricks/spark-avro#linking
os.environ['PYSPARK_SUBMIT_ARGS'] = '--packages com.databricks:spark-avro_2.10:2.0.1 %s' \
% os.getenv('PYSPARK_SUBMIT_ARGS', '')
pyspark_path()
from pyspark import SparkContext # pylint: disable=wrong-import-position,import-error
from pyspark import SparkConf # pylint: disable=wrong-import-position,import-error
from pyspark.sql import SQLContext # pylint: disable=wrong-import-position,import-error
__author__ = 'Hari Sekhon'
__version__ = '0.7.0'
class SparkJSONToAvro(CLI):
def __init__(self):
# Python 2.x
super(SparkJSONToAvro, self).__init__()
# Python 3.x
# super().__init__()
# logging.config.fileConfig(os.path.join(libdir, 'resources', 'logging.conf'))
# log = logging.getLogger(self.__class__.__name__)
# @override
def add_options(self):
self.set_verbose_default(2)
self.set_timeout_default(86400)
self.parser.add_option('-j', '--json', metavar='<file/dir>',
help='JSON input file/dir ($JSON)',
default=getenv('JSON'))
self.parser.add_option('-a', '--avro-dir', metavar='<dir>',
help='Avro output dir ($AVRODIR)',
default=getenv('AVRODIR'))
def parse_args(self):
self.no_args()
if not self.options.json:
self.usage('--json not defined')
if not self.options.avro_dir:
self.usage('--avro-dir not defined')
def run(self):
self.no_args()
json_file = self.options.json
avro_dir = self.options.avro_dir
# let Spark fail if json/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("Json Source: %s" % json_file)
log.info("Avro Destination: %s" % avro_dir)
conf = SparkConf().setAppName('HS PySpark Json => Avro')
sc = SparkContext(conf=conf) # pylint: disable=invalid-name
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):
df = sqlContext.read.json(json_file)
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')
#json = sqlContext.jsonFile(json_file)
# 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
df.write.format('com.databricks.spark.avro').save(avro_dir)
if __name__ == '__main__':
SparkJSONToAvro().main()