avro 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. #!/usr/bin/env python
  2. # Licensed to the Apache Software Foundation (ASF) under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. The ASF licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing, software
  13. # distributed under the License is distributed on an "AS IS" BASIS,
  14. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15. # See the License for the specific language governing permissions and
  16. # limitations under the License.
  17. """Command line utlity for reading and writing Avro files."""
  18. from avro.io import DatumReader, DatumWriter
  19. from avro.datafile import DataFileReader, DataFileWriter
  20. import avro.schema
  21. try:
  22. import json
  23. except ImportError:
  24. import simplejson as json
  25. import csv
  26. from sys import stdout, stdin
  27. from itertools import ifilter, imap
  28. from functools import partial
  29. from os.path import splitext
  30. class AvroError(Exception):
  31. pass
  32. def print_json(row):
  33. print(json.dumps(row))
  34. def print_json_pretty(row):
  35. print(json.dumps(row, indent=4))
  36. _write_row = csv.writer(stdout).writerow
  37. _encoding = stdout.encoding or "UTF-8"
  38. def _encode(v, encoding=_encoding):
  39. if not isinstance(v, basestring):
  40. return v
  41. return v.encode(_encoding)
  42. def print_csv(row):
  43. # We sort the keys to the fields will be in the same place
  44. # FIXME: Do we want to do it in schema order?
  45. _write_row([_encode(row[key]) for key in sorted(row)])
  46. def select_printer(format):
  47. return {
  48. "json" : print_json,
  49. "json-pretty" : print_json_pretty,
  50. "csv" : print_csv
  51. }[format]
  52. def record_match(expr, record):
  53. return eval(expr, None, {"r" : record})
  54. def parse_fields(fields):
  55. fields = fields or ''
  56. if not fields.strip():
  57. return None
  58. return [field.strip() for field in fields.split(',') if field.strip()]
  59. def field_selector(fields):
  60. fields = set(fields)
  61. def keys_filter(obj):
  62. return dict((k, obj[k]) for k in (set(obj) & fields))
  63. return keys_filter
  64. def print_avro(avro, opts):
  65. if opts.header and (opts.format != "csv"):
  66. raise AvroError("--header applies only to CSV format")
  67. # Apply filter first
  68. if opts.filter:
  69. avro = ifilter(partial(record_match, opts.filter), avro)
  70. for i in xrange(opts.skip):
  71. try:
  72. next(avro)
  73. except StopIteration:
  74. return
  75. fields = parse_fields(opts.fields)
  76. if fields:
  77. avro = imap(field_selector(fields), avro)
  78. printer = select_printer(opts.format)
  79. for i, record in enumerate(avro):
  80. if i == 0 and opts.header:
  81. _write_row(sorted(record.keys()))
  82. if i >= opts.count:
  83. break
  84. printer(record)
  85. def print_schema(avro):
  86. schema = avro.meta["avro.schema"]
  87. # Pretty print
  88. print json.dumps(json.loads(schema), indent=4)
  89. def cat(opts, args):
  90. if not args:
  91. raise AvroError("No files to show")
  92. for filename in args:
  93. try:
  94. fo = open(filename, "rb")
  95. except (OSError, IOError), e:
  96. raise AvroError("Can't open %s - %s" % (filename, e))
  97. avro = DataFileReader(fo, DatumReader())
  98. if opts.print_schema:
  99. print_schema(avro)
  100. continue
  101. print_avro(avro, opts)
  102. def _open(filename, mode):
  103. if filename == "-":
  104. return {
  105. "rb" : stdin,
  106. "wb" : stdout
  107. }[mode]
  108. return open(filename, mode)
  109. def iter_json(info, _):
  110. return imap(json.loads, info)
  111. def convert(value, field):
  112. type = field.type.type
  113. if type == "union":
  114. return convert_union(value, field)
  115. return {
  116. "int" : int,
  117. "long" : long,
  118. "float" : float,
  119. "double" : float,
  120. "string" : str,
  121. "bytes" : str,
  122. "boolean" : bool,
  123. "null" : lambda _: None,
  124. "union" : lambda v: convert_union(v, field),
  125. }[type](value)
  126. def convert_union(value, field):
  127. for name in [s.name for s in field.type.schemas]:
  128. try:
  129. return convert(name)(value)
  130. except ValueError:
  131. continue
  132. def iter_csv(info, schema):
  133. header = [field.name for field in schema.fields]
  134. for row in csv.reader(info):
  135. values = [convert(v, f) for v, f in zip(row, schema.fields)]
  136. yield dict(zip(header, values))
  137. def guess_input_type(files):
  138. if not files:
  139. return None
  140. ext = splitext(files[0])[1].lower()
  141. if ext in (".json", ".js"):
  142. return "json"
  143. elif ext in (".csv",):
  144. return "csv"
  145. return None
  146. def write(opts, files):
  147. if not opts.schema:
  148. raise AvroError("No schema specified")
  149. input_type = opts.input_type or guess_input_type(files)
  150. if not input_type:
  151. raise AvroError("Can't guess input file type (not .json or .csv)")
  152. try:
  153. schema = avro.schema.parse(open(opts.schema, "rb").read())
  154. out = _open(opts.output, "wb")
  155. except (IOError, OSError), e:
  156. raise AvroError("Can't open file - %s" % e)
  157. writer = DataFileWriter(out, DatumWriter(), schema)
  158. iter_records = {"json" : iter_json, "csv" : iter_csv}[input_type]
  159. for filename in (files or ["-"]):
  160. info = _open(filename, "rb")
  161. for record in iter_records(info, schema):
  162. writer.append(record)
  163. writer.close()
  164. def main(argv=None):
  165. import sys
  166. from optparse import OptionParser, OptionGroup
  167. argv = argv or sys.argv
  168. parser = OptionParser(description="Display/write for Avro files",
  169. version="1.8.2",
  170. usage="usage: %prog cat|write [options] FILE [FILE...]")
  171. # cat options
  172. cat_options = OptionGroup(parser, "cat options")
  173. cat_options.add_option("-n", "--count", default=float("Infinity"),
  174. help="number of records to print", type=int)
  175. cat_options.add_option("-s", "--skip", help="number of records to skip",
  176. type=int, default=0)
  177. cat_options.add_option("-f", "--format", help="record format",
  178. default="json",
  179. choices=["json", "csv", "json-pretty"])
  180. cat_options.add_option("--header", help="print CSV header", default=False,
  181. action="store_true")
  182. cat_options.add_option("--filter", help="filter records (e.g. r['age']>1)",
  183. default=None)
  184. cat_options.add_option("--print-schema", help="print schema",
  185. action="store_true", default=False)
  186. cat_options.add_option('--fields', default=None,
  187. help='fields to show, comma separated (show all by default)')
  188. parser.add_option_group(cat_options)
  189. # write options
  190. write_options = OptionGroup(parser, "write options")
  191. write_options.add_option("--schema", help="schema file (required)")
  192. write_options.add_option("--input-type",
  193. help="input file(s) type (json or csv)",
  194. choices=["json", "csv"], default=None)
  195. write_options.add_option("-o", "--output", help="output file", default="-")
  196. parser.add_option_group(write_options)
  197. opts, args = parser.parse_args(argv[1:])
  198. if len(args) < 1:
  199. parser.error("You much specify `cat` or `write`") # Will exit
  200. command = args.pop(0)
  201. try:
  202. if command == "cat":
  203. cat(opts, args)
  204. elif command == "write":
  205. write(opts, args)
  206. else:
  207. raise AvroError("Unknown command - %s" % command)
  208. except AvroError, e:
  209. parser.error("%s" % e) # Will exit
  210. except Exception, e:
  211. raise SystemExit("panic: %s" % e)
  212. if __name__ == "__main__":
  213. main()