Browse Source

Argument handling for parquet main

* enable debug output
* allow a limit
* support json and csv output

and more.
Joe Crobak 12 years ago
parent
commit
abab814365
3 changed files with 131 additions and 80 deletions
  1. 1 2
      README.md
  2. 83 71
      parquet/__init__.py
  3. 47 7
      parquet/__main__.py

+ 1 - 2
README.md

@@ -4,7 +4,6 @@
 * Support for bitwidths > 8
 * Support projection on read.
 * Fix handling of repetition-levels and definition-levels
-* Argument handlign for main (enable debug output, allow a limit,
-  support json and csv).
 * Tests for nested schemas, null data
+* Support reading of data from HDFS via snakebite
 * Implement writing

+ 83 - 71
parquet/__init__.py

@@ -1,4 +1,5 @@
 import gzip
+import json
 import logging
 import struct
 import StringIO
@@ -95,7 +96,7 @@ def _get_offset(cmd):
     return dict_offset
 
 
-def dump_metadata(filename, out=sys.stdout):
+def dump_metadata(filename, show_row_group_metadata, out=sys.stdout):
     def println(value):
         out.write(value + "\n")
     footer = read_footer(filename)
@@ -121,72 +122,73 @@ def dump_metadata(filename, out=sys.stdout):
                                               se.repetition_type),
                     num_children=se.num_children,
                     converted_type=se.converted_type))
-    println("  row groups: ")
-    for rg in footer.row_groups:
-        num_rows = rg.num_rows
-        bytes = rg.total_byte_size
-        println("  rows={num_rows}, bytes={bytes}".format(num_rows=num_rows,
-                                                          bytes=bytes))
-        println("    chunks:")
-        for cg in rg.columns:
-            cmd = cg.meta_data
-            println("      type={type} file_offset={offset} "
-                    "compression={codec} "
-                    "encodings={encodings} path_in_schema={path_in_schema} "
-                    "num_values={num_values} uncompressed_bytes={raw_bytes} "
-                    "compressed_bytes={compressed_bytes} "
-                    "data_page_offset={data_page_offset} "
-                    "dictionary_page_offset={dictionary_page_offset}".format(
-                        type=cmd.type,
-                        offset=cg.file_offset,
-                        codec=_get_name(CompressionCodec, cmd.codec),
-                        encodings=",".join(
-                            [_get_name(Encoding, s) for s in cmd.encodings]),
-                        path_in_schema=cmd.path_in_schema,
-                        num_values=cmd.num_values,
-                        raw_bytes=cmd.total_uncompressed_size,
-                        compressed_bytes=cmd.total_compressed_size,
-                        data_page_offset=cmd.data_page_offset,
-                        dictionary_page_offset=cmd.dictionary_page_offset))
-            with open(filename, 'rb') as fo:
-                offset = _get_offset(cmd)
-                fo.seek(offset, 0)
-                values_read = 0
-                println("      pages: ")
-                while values_read < num_rows:
-                    ph = _read_page_header(fo)
-                    # seek past current page.
-                    fo.seek(ph.compressed_page_size, 1)
-                    daph = ph.data_page_header
-                    type_ = _get_name(PageType, ph.type)
-                    raw_bytes = ph.uncompressed_page_size
-                    num_values = None
-                    if ph.type == PageType.DATA_PAGE:
-                        num_values = daph.num_values
-                        values_read += num_values
-                    if ph.type == PageType.DICTIONARY_PAGE:
-                        pass
-                        #num_values = diph.num_values
-
-                    encoding_type = None
-                    def_level_encoding = None
-                    rep_level_encoding = None
-                    if daph:
-                        encoding_type = _get_name(Encoding, daph.encoding)
-                        def_level_encoding = _get_name(Encoding, daph.definition_level_encoding)
-                        rep_level_encoding = _get_name(Encoding, daph.repetition_level_encoding)
-
-                    println("        page header: type={type} "
-                            "uncompressed_size={raw_bytes} "
-                            "num_values={num_values} encoding={encoding} "
-                            "def_level_encoding={def_level_encoding} "
-                            "rep_level_encoding={rep_level_encoding}".format(
-                                type=type_,
-                                raw_bytes=raw_bytes,
-                                num_values=num_values,
-                                encoding=encoding_type,
-                                def_level_encoding=def_level_encoding,
-                                rep_level_encoding=rep_level_encoding))
+    if show_row_group_metadata:
+        println("  row groups: ")
+        for rg in footer.row_groups:
+            num_rows = rg.num_rows
+            bytes = rg.total_byte_size
+            println("  rows={num_rows}, bytes={bytes}".format(num_rows=num_rows,
+                                                              bytes=bytes))
+            println("    chunks:")
+            for cg in rg.columns:
+                cmd = cg.meta_data
+                println("      type={type} file_offset={offset} "
+                        "compression={codec} "
+                        "encodings={encodings} path_in_schema={path_in_schema} "
+                        "num_values={num_values} uncompressed_bytes={raw_bytes} "
+                        "compressed_bytes={compressed_bytes} "
+                        "data_page_offset={data_page_offset} "
+                        "dictionary_page_offset={dictionary_page_offset}".format(
+                            type=cmd.type,
+                            offset=cg.file_offset,
+                            codec=_get_name(CompressionCodec, cmd.codec),
+                            encodings=",".join(
+                                [_get_name(Encoding, s) for s in cmd.encodings]),
+                            path_in_schema=cmd.path_in_schema,
+                            num_values=cmd.num_values,
+                            raw_bytes=cmd.total_uncompressed_size,
+                            compressed_bytes=cmd.total_compressed_size,
+                            data_page_offset=cmd.data_page_offset,
+                            dictionary_page_offset=cmd.dictionary_page_offset))
+                with open(filename, 'rb') as fo:
+                    offset = _get_offset(cmd)
+                    fo.seek(offset, 0)
+                    values_read = 0
+                    println("      pages: ")
+                    while values_read < num_rows:
+                        ph = _read_page_header(fo)
+                        # seek past current page.
+                        fo.seek(ph.compressed_page_size, 1)
+                        daph = ph.data_page_header
+                        type_ = _get_name(PageType, ph.type)
+                        raw_bytes = ph.uncompressed_page_size
+                        num_values = None
+                        if ph.type == PageType.DATA_PAGE:
+                            num_values = daph.num_values
+                            values_read += num_values
+                        if ph.type == PageType.DICTIONARY_PAGE:
+                            pass
+                            #num_values = diph.num_values
+
+                        encoding_type = None
+                        def_level_encoding = None
+                        rep_level_encoding = None
+                        if daph:
+                            encoding_type = _get_name(Encoding, daph.encoding)
+                            def_level_encoding = _get_name(Encoding, daph.definition_level_encoding)
+                            rep_level_encoding = _get_name(Encoding, daph.repetition_level_encoding)
+
+                        println("        page header: type={type} "
+                                "uncompressed_size={raw_bytes} "
+                                "num_values={num_values} encoding={encoding} "
+                                "def_level_encoding={def_level_encoding} "
+                                "rep_level_encoding={rep_level_encoding}".format(
+                                    type=type_,
+                                    raw_bytes=raw_bytes,
+                                    num_values=num_values,
+                                    encoding=encoding_type,
+                                    def_level_encoding=def_level_encoding,
+                                    rep_level_encoding=rep_level_encoding))
 
 
 def _read_page(fo, page_header, column_metadata):
@@ -279,7 +281,7 @@ def read_data_page(fo, schema_helper, page_header, column_metadata,
         for i in range(daph.num_values):
             vals.append(
                 encoding.read_plain(io_obj, column_metadata.type, None))
-        logger.debug("  Values: ", ",".join([str(x) for x in vals]))
+        logger.debug("  Values: %s", ",".join([str(x) for x in vals]))
     elif daph.encoding == Encoding.PLAIN_DICTIONARY:
         # bit_width is stored as single byte.
         bit_width = struct.unpack("<B", io_obj.read(1))[0]
@@ -318,12 +320,16 @@ def dump(filename, options, out=sys.stdout):
 
     footer = read_footer(filename)
     schema_helper = schema.SchemaHelper(footer.schema)
+    total_count = 0
     for rg in footer.row_groups:
         res = defaultdict(list)
         row_group_rows = rg.num_rows
         for idx, cg in enumerate(rg.columns):
             dict_items = []
             cmd = cg.meta_data
+            # skip if the list of columns is specified and this isn't in it
+            if options.col and not ".".join(cmd.path_in_schema) in options.col:
+                continue
             with open(filename, 'rb') as fo:
                 offset = _get_offset(cmd)
                 fo.seek(offset, 0)
@@ -352,8 +358,14 @@ def dump(filename, options, out=sys.stdout):
                     else:
                         logger.warn("Skipping unknown page type={0}".format(
                             _get_name(PageType, ph.type)))
-        println("Data for row group: ")
         keys = res.keys()
-        println("\t".join(keys))
+        if options.format == "csv" and not options.no_headers:
+            println("\t".join(keys))
         for i in range(rg.num_rows):
-            println("\t".join(str(res[k][i]) for k in keys))
+            if options.limit != -1 and i + total_count >= options.limit:
+                return
+            if options.format == "csv":
+                println("\t".join(str(res[k][i]) for k in keys))
+            elif options.format == "json":
+                println(json.dumps(dict([(k, res[k][i]) for k in keys])))
+        total_count += rg.num_rows

+ 47 - 7
parquet/__main__.py

@@ -1,8 +1,48 @@
-import parquet
+import argparse
+import logging
 import sys
-parquet.dump_metadata(sys.argv[1])
-print("")
-print("-"*80)
-print("")
-options = {}  # TODO(jcrobak)
-parquet.dump(sys.argv[1], options)
+
+import parquet
+
+
+def setup_logging(options):
+    level = logging.DEBUG if options.debug else logging.WARNING
+    logging.basicConfig(level=level)
+
+
+def main(argv=None):
+    argv = argv or sys.argv[1:]
+
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    parser.add_argument('--metadata', action='store_true',
+                        help='show metadata on file')
+    parser.add_argument('--row-group-metadata', action='store_true',
+                        help="show per row group metadata")
+    parser.add_argument('--no-data', action='store_true',
+                        help="don't dump any data from the file")
+    parser.add_argument('--limit', action='store', type=int, default=-1,
+                        help='max records to output')
+    parser.add_argument('--col', action='append', type=str,
+                        help='only include this column (can be '
+                             'specified multiple times)')
+    parser.add_argument('--no-headers', action='store_true',
+                        help='skip headers in output (only applies if '
+                             'format=csv)')
+    parser.add_argument('--format', action='store', type=str, default='csv',
+                        help='format for the output data. can be csv or json.')
+    parser.add_argument('--debug', action='store_true',
+                        help='log debug info to stderr')
+    parser.add_argument('file',
+                        help='path to the file to parse')
+
+    args = parser.parse_args(argv)
+
+    setup_logging(args)
+
+    if args.metadata:
+        parquet.dump_metadata(args.file, args.row_group_metadata)
+    if not args.no_data:
+        parquet.dump(args.file, args)
+
+if __name__ == '__main__':
+    main()