Browse Source

cleanup

* docstring all the things.
* fix tons of pep8 errors.
* clean up tests and improve test coverage.
* move a bunch of prints to debug statements.
* add a readme, only contains a todo so far.
Joe Crobak 12 years ago
parent
commit
0d984bc6d6
7 changed files with 343 additions and 158 deletions
  1. 2 0
      .gitignore
  2. 10 0
      README.md
  3. 182 94
      parquet/__init__.py
  4. 5 1
      parquet/__main__.py
  5. 69 18
      parquet/encoding.py
  6. 23 20
      test/test_encoding.py
  7. 52 25
      test/test_read_support.py

+ 2 - 0
.gitignore

@@ -1 +1,3 @@
 *.pyc
 *.pyc
+.coverage
+cover

+ 10 - 0
README.md

@@ -0,0 +1,10 @@
+# Todos
+
+* Support the deprecated bitpacking
+* 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
+* Implement writing

+ 182 - 94
parquet/__init__.py

@@ -1,11 +1,11 @@
 import gzip
 import gzip
-import json
 import logging
 import logging
 import struct
 import struct
-import thrift
 import StringIO
 import StringIO
+import sys
 from collections import defaultdict
 from collections import defaultdict
-from ttypes import FileMetaData, CompressionCodec, Encoding, FieldRepetitionType, PageHeader, PageType, Type
+from ttypes import (FileMetaData, CompressionCodec, Encoding,
+                    FieldRepetitionType, PageHeader, PageType, Type)
 from thrift.protocol import TCompactProtocol
 from thrift.protocol import TCompactProtocol
 from thrift.transport import TTransport
 from thrift.transport import TTransport
 import encoding
 import encoding
@@ -17,33 +17,38 @@ logger = logging.getLogger("parquet")
 try:
 try:
     import snappy
     import snappy
 except ImportError:
 except ImportError:
-    logger.warn("Couldn't import snappy. Support for snappy compression disabled.")
+    logger.warn(
+        "Couldn't import snappy. Support for snappy compression disabled.")
+
 
 
 class ParquetFormatException(Exception):
 class ParquetFormatException(Exception):
     pass
     pass
 
 
+
 def _check_header_magic_bytes(fo):
 def _check_header_magic_bytes(fo):
-    """Returns true if the file-like obj has the PAR1 magic bytes at the header"""
+    "Returns true if the file-like obj has the PAR1 magic bytes at the header"
     fo.seek(0, 0)
     fo.seek(0, 0)
     magic = fo.read(4)
     magic = fo.read(4)
     return magic == 'PAR1'
     return magic == 'PAR1'
 
 
+
 def _check_footer_magic_bytes(fo):
 def _check_footer_magic_bytes(fo):
-    """Returns true if the file-like obj has the PAR1 magic bytes at the footer"""
+    "Returns true if the file-like obj has the PAR1 magic bytes at the footer"
     fo.seek(-4, 2)  # seek to four bytes from the end of the file
     fo.seek(-4, 2)  # seek to four bytes from the end of the file
     magic = fo.read(4)
     magic = fo.read(4)
     return magic == 'PAR1'
     return magic == 'PAR1'
 
 
 
 
 def _get_footer_size(fo):
 def _get_footer_size(fo):
-    """Readers the footer size in bytes, which is serialized as little endian"""
+    "Readers the footer size in bytes, which is serialized as little endian"
     fo.seek(-8, 2)
     fo.seek(-8, 2)
     tup = struct.unpack("<i", fo.read(4))
     tup = struct.unpack("<i", fo.read(4))
     return tup[0]
     return tup[0]
 
 
+
 def _read_footer(fo):
 def _read_footer(fo):
-    """Reads the footer from the given file object, returning a FileMetaData object. This method
-    assumes that the fo references a valid parquet file"""
+    """Reads the footer from the given file object, returning a FileMetaData
+    object. This method assumes that the fo references a valid parquet file"""
     footer_size = _get_footer_size(fo)
     footer_size = _get_footer_size(fo)
     logger.debug("Footer size in bytes: %s", footer_size)
     logger.debug("Footer size in bytes: %s", footer_size)
     fo.seek(-(8+footer_size), 2)  # seek to beginning of footer
     fo.seek(-(8+footer_size), 2)  # seek to beginning of footer
@@ -62,61 +67,98 @@ def _read_page_header(fo):
     ph.read(pin)
     ph.read(pin)
     return ph
     return ph
 
 
+
 def read_footer(filename):
 def read_footer(filename):
     """Reads and returns the FileMetaData object for the given file."""
     """Reads and returns the FileMetaData object for the given file."""
     with open(filename, 'rb') as fo:
     with open(filename, 'rb') as fo:
-        if not _check_header_magic_bytes(fo) or not _check_footer_magic_bytes(fo):
-            raise ParquetFormatException("{0} is not a valid parquet file (missing magic bytes)".format(filename))
+        if not _check_header_magic_bytes(fo) or \
+           not _check_footer_magic_bytes(fo):
+            raise ParquetFormatException("{0} is not a valid parquet file "
+                                         "(missing magic bytes)"
+                                         .format(filename))
         return _read_footer(fo)
         return _read_footer(fo)
 
 
 
 
-def dump_metadata(filename):
+def _get_name(type_, value):
+    """Returns the name for the given value of the given type_ unless value is
+    None, in which case it returns empty string"""
+    return type_._VALUES_TO_NAMES[value] if value else "None"
+
+
+def _get_offset(cmd):
+    """Returns the offset into the cmd based upon if it's a dictionary page or
+    a data page"""
+    dict_offset = cmd.dictionary_page_offset
+    data_offset = cmd.data_page_offset
+    if dict_offset is None or data_offset < dict_offset:
+        return data_offset
+    return dict_offset
+
+
+def dump_metadata(filename, out=sys.stdout):
+    def println(value):
+        out.write(value + "\n")
     footer = read_footer(filename)
     footer = read_footer(filename)
-    print("File: {0}".format(filename))
-    print("  version: {0}".format(footer.version))
-    print("  num rows: {0}".format(footer.num_rows))
-    print("  k/v metadata: ")
+    println("File Metadata: {0}".format(filename))
+    println("  Version: {0}".format(footer.version))
+    println("  Num Rows: {0}".format(footer.num_rows))
+    println("  k/v metadata: ")
     if footer.key_value_metadata and len(footer.key_value_metadata) > 0:
     if footer.key_value_metadata and len(footer.key_value_metadata) > 0:
         for kv in footer.key_value_metadata:
         for kv in footer.key_value_metadata:
-            print("    {0}={1}".format(kv.key, kv.value)) 
+            println("    {0}={1}".format(kv.key, kv.value))
     else:
     else:
-        print("    (none)")
-    print("  schema: ")
+        println("    (none)")
+    println("  schema: ")
     for se in footer.schema:
     for se in footer.schema:
-        print("    {name} ({type}): length={type_length}, repetition={repetition_type}, children={num_children}, converted_type={converted_type}".format(
-            name=se.name, type=Type._VALUES_TO_NAMES[se.type] if se.type else None, type_length=se.type_length,
-            repetition_type=FieldRepetitionType._VALUES_TO_NAMES[se.repetition_type] if se.repetition_type else None,
-            num_children=se.num_children, converted_type=se.converted_type))
-    print("  row groups: ")
+        println("    {name} ({type}): length={type_length}, "
+                "repetition={repetition_type}, "
+                "children={num_children}, "
+                "converted_type={converted_type}".format(
+                    name=se.name,
+                    type=Type._VALUES_TO_NAMES[se.type] if se.type else None,
+                    type_length=se.type_length,
+                    repetition_type=_get_name(FieldRepetitionType,
+                                              se.repetition_type),
+                    num_children=se.num_children,
+                    converted_type=se.converted_type))
+    println("  row groups: ")
     for rg in footer.row_groups:
     for rg in footer.row_groups:
         num_rows = rg.num_rows
         num_rows = rg.num_rows
         bytes = rg.total_byte_size
         bytes = rg.total_byte_size
-        print("  rows={num_rows}, bytes={bytes}".format(num_rows=num_rows, bytes=bytes))
-        print("    chunks:")
+        println("  rows={num_rows}, bytes={bytes}".format(num_rows=num_rows,
+                                                          bytes=bytes))
+        println("    chunks:")
         for cg in rg.columns:
         for cg in rg.columns:
             cmd = cg.meta_data
             cmd = cg.meta_data
-            print("      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=CompressionCodec._VALUES_TO_NAMES[cmd.codec],
-                    encodings=",".join([Encoding._VALUES_TO_NAMES[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
-                    ))
+            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:
             with open(filename, 'rb') as fo:
-                offset = cmd.data_page_offset if (cmd.dictionary_page_offset is None or cmd.data_page_offset < cmd.dictionary_page_offset) else cmd.dictionary_page_offset
+                offset = _get_offset(cmd)
                 fo.seek(offset, 0)
                 fo.seek(offset, 0)
                 values_read = 0
                 values_read = 0
-                print("      pages: ")
+                println("      pages: ")
                 while values_read < num_rows:
                 while values_read < num_rows:
                     ph = _read_page_header(fo)
                     ph = _read_page_header(fo)
-                    fo.seek(ph.compressed_page_size, 1) # seek past current page.
+                    # seek past current page.
+                    fo.seek(ph.compressed_page_size, 1)
                     daph = ph.data_page_header
                     daph = ph.data_page_header
-                    diph = ph.dictionary_page_header
-                    type_ = PageType._VALUES_TO_NAMES[ph.type] if ph.type else None
+                    type_ = _get_name(PageType, ph.type)
                     raw_bytes = ph.uncompressed_page_size
                     raw_bytes = ph.uncompressed_page_size
                     num_values = None
                     num_values = None
                     if ph.type == PageType.DATA_PAGE:
                     if ph.type == PageType.DATA_PAGE:
@@ -126,47 +168,59 @@ def dump_metadata(filename):
                         pass
                         pass
                         #num_values = diph.num_values
                         #num_values = diph.num_values
 
 
-                    encoding = None
+                    encoding_type = None
                     def_level_encoding = None
                     def_level_encoding = None
                     rep_level_encoding = None
                     rep_level_encoding = None
                     if daph:
                     if daph:
-                        encoding = Encoding._VALUES_TO_NAMES[daph.encoding]
-                        def_level_encoding = Encoding._VALUES_TO_NAMES[daph.definition_level_encoding]
-                        rep_level_encoding = Encoding._VALUES_TO_NAMES[daph.repetition_level_encoding]
-
-                    print("        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, def_level_encoding=def_level_encoding,
-                            rep_level_encoding=rep_level_encoding))
+                        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):
 def _read_page(fo, page_header, column_metadata):
-    """Reads the data page from the given file-object using the column metadata"""
+    """Internal function to read the data page from the given file-object
+    and convert it to raw, uncompressed bytes (if necessary)."""
     bytes_from_file = fo.read(page_header.compressed_page_size)
     bytes_from_file = fo.read(page_header.compressed_page_size)
-    if column_metadata.codec is not None and column_metadata.codec != CompressionCodec.UNCOMPRESSED:
+    codec = column_metadata.codec
+    if codec is not None and codec != CompressionCodec.UNCOMPRESSED:
         if column_metadata.codec == CompressionCodec.SNAPPY:
         if column_metadata.codec == CompressionCodec.SNAPPY:
             raw_bytes = snappy.decompress(bytes_from_file)
             raw_bytes = snappy.decompress(bytes_from_file)
         elif column_metadata.codec == CompressionCodec.GZIP:
         elif column_metadata.codec == CompressionCodec.GZIP:
             io_obj = StringIO.StringIO(bytes_from_file)
             io_obj = StringIO.StringIO(bytes_from_file)
             with gzip.GzipFile(fileobj=io_obj, mode='rb') as f:
             with gzip.GzipFile(fileobj=io_obj, mode='rb') as f:
                 raw_bytes = f.read()
                 raw_bytes = f.read()
+        else:
+            raise ParquetFormatException(
+                "Unsupported Codec: {0}".format(codec))
     else:
     else:
         raw_bytes = bytes_from_file
         raw_bytes = bytes_from_file
     return raw_bytes
     return raw_bytes
 
 
 
 
 def _read_data(fo, fo_encoding, value_count, bit_width):
 def _read_data(fo, fo_encoding, value_count, bit_width):
-    """Internal method to read data from the file-object using the given encoding. The data
-    could be definition levels, repetition levels, or actual values."""
+    """Internal method to read data from the file-object using the given
+    encoding. The data could be definition levels, repetition levels, or
+    actual values.
+    """
     vals = []
     vals = []
     if fo_encoding == Encoding.RLE:
     if fo_encoding == Encoding.RLE:
         seen = 0
         seen = 0
         while seen < value_count:
         while seen < value_count:
             values = encoding.read_rle_bit_packed_hybrid(fo, bit_width)
             values = encoding.read_rle_bit_packed_hybrid(fo, bit_width)
             if values is None:
             if values is None:
-                break  ## EOF was reached.
+                break  # EOF was reached.
             vals += values
             vals += values
             seen += len(values)
             seen += len(values)
     elif fo_encoding == Encoding.BIT_PACKED:
     elif fo_encoding == Encoding.BIT_PACKED:
@@ -175,52 +229,73 @@ def _read_data(fo, fo_encoding, value_count, bit_width):
     return vals
     return vals
 
 
 
 
-def read_data_page(fo, schema_helper, page_header, column_metadata, dictionary):
+def read_data_page(fo, schema_helper, page_header, column_metadata,
+                   dictionary):
+    """Reads the datapage from the given file-like object based upon the
+    metadata in the schema_helper, page_header, column_metadata, and
+    (optional) dictionary. Returns a list of values.
+    """
     daph = page_header.data_page_header
     daph = page_header.data_page_header
     raw_bytes = _read_page(fo, page_header, column_metadata)
     raw_bytes = _read_page(fo, page_header, column_metadata)
     io_obj = StringIO.StringIO(raw_bytes)
     io_obj = StringIO.StringIO(raw_bytes)
     vals = []
     vals = []
 
 
-    print("  definition_level_encoding: {0}".format(Encoding._VALUES_TO_NAMES[daph.definition_level_encoding]))
-    print("  repetition_level_encoding: {0}".format(Encoding._VALUES_TO_NAMES[daph.repetition_level_encoding]))
-    print("  encoding: {0}".format(Encoding._VALUES_TO_NAMES[daph.encoding]) )
+    logger.debug("  definition_level_encoding: %s",
+                 _get_name(Encoding, daph.definition_level_encoding))
+    logger.debug("  repetition_level_encoding: %s",
+                 _get_name(Encoding, daph.repetition_level_encoding))
+    logger.debug("  encoding: %s", _get_name(Encoding, daph.encoding))
 
 
     # definition levels are skipped if data is required.
     # definition levels are skipped if data is required.
     if not schema_helper.is_required(column_metadata.path_in_schema[-1]):
     if not schema_helper.is_required(column_metadata.path_in_schema[-1]):
-        max_definition_level = schema_helper.max_definition_level(column_metadata.path_in_schema)
-        bit_width = encoding.width_from_max_int(max_definition_level) # TODO Where does the -1 come from?
-        print "  max def level: {1}   bit_width: {0}".format(bit_width, max_definition_level)
+        max_definition_level = schema_helper.max_definition_level(
+            column_metadata.path_in_schema)
+        bit_width = encoding.width_from_max_int(max_definition_level)
+        logger.debug("  max def level: %s   bit_width: %s",
+                     max_definition_level, bit_width)
         if bit_width == 0:
         if bit_width == 0:
             definition_levels = [0] * daph.num_values
             definition_levels = [0] * daph.num_values
         else:
         else:
-            definition_levels = _read_data(io_obj, daph.definition_level_encoding, daph.num_values, bit_width)
+            definition_levels = _read_data(io_obj,
+                                           daph.definition_level_encoding,
+                                           daph.num_values,
+                                           bit_width)
+
+        logger.debug("  Definition levels: %s",
+                     ",".join([str(dl) for dl in definition_levels]))
 
 
-        print ("  Definition levels: {0}".format(",".join([str(dl) for dl in definition_levels])))
-    
     # repetition levels are skipped if data is at the first level.
     # repetition levels are skipped if data is at the first level.
     if len(column_metadata.path_in_schema) > 1:
     if len(column_metadata.path_in_schema) > 1:
-        max_repetition_level = schema_helper.max_repetition_level(column_metadata.path_in_schema)
+        max_repetition_level = schema_helper.max_repetition_level(
+            column_metadata.path_in_schema)
         bit_width = encoding.width_from_max_int(max_repetition_level)
         bit_width = encoding.width_from_max_int(max_repetition_level)
-        repetition_levels = _read_data(io_obj, daph.repetition_level_encoding, daph.num_values)
+        repetition_levels = _read_data(io_obj,
+                                       daph.repetition_level_encoding,
+                                       daph.num_values)
 
 
     # TODO Actually use the definition and repetition levels.
     # TODO Actually use the definition and repetition levels.
 
 
     if daph.encoding == Encoding.PLAIN:
     if daph.encoding == Encoding.PLAIN:
         for i in range(daph.num_values):
         for i in range(daph.num_values):
-            vals.append(encoding.read_plain(io_obj, column_metadata.type, None))
-        print "  Values: " + ",".join([str(x) for x in vals])
+            vals.append(
+                encoding.read_plain(io_obj, column_metadata.type, None))
+        logger.debug("  Values: ", ",".join([str(x) for x in vals]))
     elif daph.encoding == Encoding.PLAIN_DICTIONARY:
     elif daph.encoding == Encoding.PLAIN_DICTIONARY:
-        bit_width = struct.unpack("<B", io_obj.read(1))[0]  # bitwidth is stored as single byte.
-        print "bit_width: {0}".format(bit_width)
+        # bit_width is stored as single byte.
+        bit_width = struct.unpack("<B", io_obj.read(1))[0]
+        logger.debug("bit_width: %d", bit_width)
         total_seen = 0
         total_seen = 0
         dict_values_bytes = io_obj.read()
         dict_values_bytes = io_obj.read()
         dict_values_io_obj = StringIO.StringIO(dict_values_bytes)
         dict_values_io_obj = StringIO.StringIO(dict_values_bytes)
-        while total_seen < daph.num_values:  # TODO jcrobak -- not sure that this loop i sneeded?
-            values = encoding.read_rle_bit_packed_hybrid(dict_values_io_obj, bit_width, len(dict_values_bytes))
+        # TODO jcrobak -- not sure that this loop is needed?
+        while total_seen < daph.num_values:
+            values = encoding.read_rle_bit_packed_hybrid(
+                dict_values_io_obj, bit_width, len(dict_values_bytes))
             vals += [dictionary[v] for v in values]
             vals += [dictionary[v] for v in values]
             total_seen += len(values)
             total_seen += len(values)
     else:
     else:
-        raise ParquetFormatException("Unsupported encoding: " + Encoding._VALUES_TO_NAMES[daph.encoding])
+        raise ParquetFormatException("Unsupported encoding: %s",
+                                     _get_name(Encoding, daph.encoding))
     return vals
     return vals
 
 
 
 
@@ -229,43 +304,56 @@ def read_dictionary_page(fo, page_header, column_metadata):
     io_obj = StringIO.StringIO(raw_bytes)
     io_obj = StringIO.StringIO(raw_bytes)
     dict_items = []
     dict_items = []
     while io_obj.tell() < len(raw_bytes):
     while io_obj.tell() < len(raw_bytes):
-        dict_items.append(encoding.read_plain(io_obj, column_metadata.type, None))  # TODO - length for fixed byte array
+        # TODO - length for fixed byte array
+        dict_items.append(
+            encoding.read_plain(io_obj, column_metadata.type, None))
     return dict_items
     return dict_items
 
 
+## max_records=10
+
+
+def dump(filename, options, out=sys.stdout):
+    def println(value):
+        out.write(value + "\n")
 
 
-def dump(filename, max_records=10):
     footer = read_footer(filename)
     footer = read_footer(filename)
     schema_helper = schema.SchemaHelper(footer.schema)
     schema_helper = schema.SchemaHelper(footer.schema)
     for rg in footer.row_groups:
     for rg in footer.row_groups:
         res = defaultdict(list)
         res = defaultdict(list)
         row_group_rows = rg.num_rows
         row_group_rows = rg.num_rows
-        dict_items = []
         for idx, cg in enumerate(rg.columns):
         for idx, cg in enumerate(rg.columns):
+            dict_items = []
             cmd = cg.meta_data
             cmd = cg.meta_data
             with open(filename, 'rb') as fo:
             with open(filename, 'rb') as fo:
-                offset = cmd.data_page_offset if (cmd.dictionary_page_offset is None or cmd.data_page_offset < cmd.dictionary_page_offset) else cmd.dictionary_page_offset
+                offset = _get_offset(cmd)
                 fo.seek(offset, 0)
                 fo.seek(offset, 0)
                 values_seen = 0
                 values_seen = 0
-                print("reading column chunk of type: {0}".format(Type._VALUES_TO_NAMES[cmd.type]))
+                logger.debug("reading column chunk of type: %s",
+                             _get_name(Type, cmd.type))
                 while values_seen < row_group_rows:
                 while values_seen < row_group_rows:
                     ph = _read_page_header(fo)
                     ph = _read_page_header(fo)
-                    print("Reading page (type={2}, uncompressed={0} bytes, compressed={1} bytes)".format(
-                        ph.uncompressed_page_size, ph.compressed_page_size, PageType._VALUES_TO_NAMES[ph.type]))
-                    daph = ph.data_page_header
-                    diph = ph.dictionary_page_header
+                    logger.debug("Reading page (type=%s, "
+                                 "uncompressed=%s bytes, "
+                                 "compressed=%s bytes)",
+                                 _get_name(PageType, ph.type),
+                                 ph.uncompressed_page_size,
+                                 ph.compressed_page_size)
+
                     if ph.type == PageType.DATA_PAGE:
                     if ph.type == PageType.DATA_PAGE:
-                        values = read_data_page(fo, schema_helper, ph, cmd, dict_items)
+                        values = read_data_page(fo, schema_helper, ph, cmd,
+                                                dict_items)
                         res[".".join(cmd.path_in_schema)] += values
                         res[".".join(cmd.path_in_schema)] += values
                         values_seen += cmd.num_values
                         values_seen += cmd.num_values
                     elif ph.type == PageType.DICTIONARY_PAGE:
                     elif ph.type == PageType.DICTIONARY_PAGE:
-                        print ph
+                        logger.debug(ph)
                         assert dict_items == []
                         assert dict_items == []
                         dict_items = read_dictionary_page(fo, ph, cmd)
                         dict_items = read_dictionary_page(fo, ph, cmd)
-                        print("Dictionary: " + str(dict_items))
+                        logger.debug("Dictionary: %s", str(dict_items))
                     else:
                     else:
-                        logger.info("Skipping unknown page type={0}".format(ph.type))
-        print "Data for row group: "
+                        logger.warn("Skipping unknown page type={0}".format(
+                            _get_name(PageType, ph.type)))
+        println("Data for row group: ")
         keys = res.keys()
         keys = res.keys()
-        print "\t".join(keys)
+        println("\t".join(keys))
         for i in range(rg.num_rows):
         for i in range(rg.num_rows):
-            print "\t".join(str(res[k][i]) for k in keys)
+            println("\t".join(str(res[k][i]) for k in keys))

+ 5 - 1
parquet/__main__.py

@@ -1,4 +1,8 @@
 import parquet
 import parquet
 import sys
 import sys
 parquet.dump_metadata(sys.argv[1])
 parquet.dump_metadata(sys.argv[1])
-parquet.dump(sys.argv[1])
+print("")
+print("-"*80)
+print("")
+options = {}  # TODO(jcrobak)
+parquet.dump(sys.argv[1], options)

+ 69 - 18
parquet/encoding.py

@@ -2,36 +2,56 @@ import array
 import math
 import math
 import struct
 import struct
 import StringIO
 import StringIO
+import logging
 
 
 from ttypes import Type
 from ttypes import Type
 
 
+logger = logging.getLogger("parquet")
+
+
 def read_plain_boolean(fo):
 def read_plain_boolean(fo):
+    """Reads a boolean using the plain encoding"""
     raise NotImplemented
     raise NotImplemented
 
 
 
 
 def read_plain_int32(fo):
 def read_plain_int32(fo):
+    """Reads a 32-bit int using the plain encoding"""
     tup = struct.unpack("<i", fo.read(4))
     tup = struct.unpack("<i", fo.read(4))
     return tup[0]
     return tup[0]
 
 
+
 def read_plain_int64(fo):
 def read_plain_int64(fo):
+    """Reads a 64-bit int using the plain encoding"""
     tup = struct.unpack("<q", fo.read(8))
     tup = struct.unpack("<q", fo.read(8))
     return tup[0]
     return tup[0]
 
 
+
 def read_plain_int96(fo):
 def read_plain_int96(fo):
+    """Reads a 96-bit int using the plain encoding"""
     tup = struct.unpack("<q<i", fo.read(12))
     tup = struct.unpack("<q<i", fo.read(12))
     return tup[0] << 32 | tup[1]
     return tup[0] << 32 | tup[1]
 
 
+
 def read_plain_float(fo):
 def read_plain_float(fo):
+    """Reads a 32-bit float using the plain encoding"""
     tup = struct.unpack("<f", fo.read(4))
     tup = struct.unpack("<f", fo.read(4))
+    return tup[0]
+
 
 
 def read_plain_double(fo):
 def read_plain_double(fo):
+    """Reads a 64-bit float (double) using the plain encoding"""
     tup = struct.unpack("<d", fo.read(8))
     tup = struct.unpack("<d", fo.read(8))
+    return tup[0]
+
 
 
 def read_plain_byte_array(fo):
 def read_plain_byte_array(fo):
+    """Reads a byte array using the plain encoding"""
     length = read_plain_int32(fo)
     length = read_plain_int32(fo)
     return fo.read(length)
     return fo.read(length)
 
 
+
 def read_plain_byte_array_fixed(fo, fixed_length):
 def read_plain_byte_array_fixed(fo, fixed_length):
+    """Reads a byte array of the given fixed_length"""
     return fo.read(fixed_length)
     return fo.read(fixed_length)
 
 
 DECODE_PLAIN = {
 DECODE_PLAIN = {
@@ -45,6 +65,7 @@ DECODE_PLAIN = {
     Type.FIXED_LEN_BYTE_ARRAY: read_plain_byte_array_fixed
     Type.FIXED_LEN_BYTE_ARRAY: read_plain_byte_array_fixed
 }
 }
 
 
+
 def read_plain(fo, type_, type_length):
 def read_plain(fo, type_, type_length):
     conv = DECODE_PLAIN[type_]
     conv = DECODE_PLAIN[type_]
     if type_ == Type.FIXED_LEN_BYTE_ARRAY:
     if type_ == Type.FIXED_LEN_BYTE_ARRAY:
@@ -63,13 +84,19 @@ def read_unsigned_var_int(fo):
         shift += 7
         shift += 7
     return result
     return result
 
 
+
 def byte_width(bit_width):
 def byte_width(bit_width):
     "Returns the byte width for the given bit_width"
     "Returns the byte width for the given bit_width"
-    return (bit_width + 7) / 8;
+    return (bit_width + 7) / 8
+
 
 
 def read_rle(fo, header, bit_width):
 def read_rle(fo, header, bit_width):
-    """Grabs count from the header and uses width to grab the value that's
-    repeated. Returns an array with the value repeated count times."""
+    """Read a run-length encoded run from the given fo with the given header
+    and bit_width.
+
+    The count is determined from the header and the width is used to grab the
+    value that's repeated. Yields the value repeated count times.
+    """
     count = header >> 1
     count = header >> 1
     zero_data = "\x00\x00\x00\x00"
     zero_data = "\x00\x00\x00\x00"
     data = ""
     data = ""
@@ -79,47 +106,65 @@ def read_rle(fo, header, bit_width):
     elif width >= 2:
     elif width >= 2:
         data += fo.read(1)
         data += fo.read(1)
     elif width >= 3:
     elif width >= 3:
-        data +=  fo.read(1)
+        data += fo.read(1)
     elif width == 4:
     elif width == 4:
         data = fo.read(1)
         data = fo.read(1)
     data = data + zero_data[len(data):]
     data = data + zero_data[len(data):]
     value = struct.unpack("<i", data)[0]
     value = struct.unpack("<i", data)[0]
-
-    return [value]*count
+    logger.debug("Read RLE group with value %s of byte-width %s and count %s",
+                 value, width, count)
+    for i in range(count):
+        yield value
 
 
 
 
 def width_from_max_int(value):
 def width_from_max_int(value):
+    """Converts the value specified to a bit_width."""
     return int(math.ceil(math.log(value + 1, 2)))
     return int(math.ceil(math.log(value + 1, 2)))
 
 
-def mask_for_bits(i):
+
+def _mask_for_bits(i):
+    """Helper function for read_bitpacked to generage a mask to grab i bits."""
     return (1 << i) - 1
     return (1 << i) - 1
 
 
+
 def read_bitpacked(fo, header, width):
 def read_bitpacked(fo, header, width):
-    num_groups = header >> 1;
+    """Reads a bitpacked run of the rle/bitpack hybrid.
+
+    Currently only supports width <=8 (doesn't support crossing bytes).
+    """
+    assert width <= 8
+    num_groups = header >> 1
+    logger.debug("Reading a bit-packed run with: %s groups", num_groups)
     count = num_groups * 8
     count = num_groups * 8
     raw_bytes = array.array('B', fo.read(count)).tolist()
     raw_bytes = array.array('B', fo.read(count)).tolist()
     current_byte = 0
     current_byte = 0
     b = raw_bytes[current_byte]
     b = raw_bytes[current_byte]
-    mask = mask_for_bits(width)
+    mask = _mask_for_bits(width)
     bits_in_byte = 8
     bits_in_byte = 8
     res = []
     res = []
-    while current_byte < width and len(res) < (count / width):
-        print "width={0} bits_in_byte={1} b={2}".format(width, bits_in_byte, bin(b))
+    while current_byte < len(raw_bytes):
+        # TODO zero-padding could produce extra zero-values
+        logger.debug("  read bitpacked: width=%s bits_in_byte=%s b=%s,"
+                     " current_byte=%s",
+                     width, bits_in_byte, bin(b), current_byte)
         if bits_in_byte >= width:
         if bits_in_byte >= width:
             res.append(b & mask)
             res.append(b & mask)
             b >>= width
             b >>= width
             bits_in_byte -= width
             bits_in_byte -= width
         else:
         else:
+            if current_byte + 1 == len(raw_bytes):
+                break  # partial results / padding at the end.
             next_b = raw_bytes[current_byte + 1]
             next_b = raw_bytes[current_byte + 1]
-            borrowed_bits = next_b & mask_for_bits(width - bits_in_byte)
-            #print "  borrowing {0} bites".format(width - bits_in_byte)
-            #print "  next_b={0}, borrowed_bits={1}".format(bin(next_b), bin(borrowed_bits))
+            borrowed_bits = next_b & _mask_for_bits(width - bits_in_byte)
+            logger.debug("    borrowing %d bits", width - bits_in_byte)
+            logger.debug("    next_b=%s, borrowed_bits=%s",
+                         bin(next_b), bin(borrowed_bits))
             res.append((borrowed_bits << bits_in_byte) | b)
             res.append((borrowed_bits << bits_in_byte) | b)
             b = next_b >> (width - bits_in_byte)
             b = next_b >> (width - bits_in_byte)
-            #print "  shifting away: {0}".format(width - bits_in_byte)
+            logger.debug("  shifting away: %d", width - bits_in_byte)
             bits_in_byte = 8 - (width - bits_in_byte)
             bits_in_byte = 8 - (width - bits_in_byte)
             current_byte += 1
             current_byte += 1
-        print "  added: {0}".format(res[-1])
+        logger.debug("  read bitpackage: added: %s", res[-1])
     return res
     return res
 
 
 
 
@@ -128,12 +173,18 @@ def read_bitpacked_deprecated(fo, count, width):
     raw_bytes = array.array('B', fo.read(count)).tolist()
     raw_bytes = array.array('B', fo.read(count)).tolist()
     current_byte = 0
     current_byte = 0
     b = raw_bytes[current_byte]
     b = raw_bytes[current_byte]
-    mask = mask_for_bits(width)
+    mask = _mask_for_bits(width)
 
 
+    # TODO implement
+    return res
 
 
 
 
 def read_rle_bit_packed_hybrid(fo, width, length=None):
 def read_rle_bit_packed_hybrid(fo, width, length=None):
-#    import pdb; pdb.set_trace()
+    """Implemenation of a decoder for the rel/bit-packed hybrid encoding.
+
+    If length is not specified, then a 32-bit int is read first to grab the
+    length of the encoded data.
+    """
     io_obj = fo
     io_obj = fo
     if length is None:
     if length is None:
         length = read_plain_int32(fo)
         length = read_plain_int32(fo)

+ 23 - 20
test/test_encoding.py

@@ -4,10 +4,12 @@ import unittest
 
 
 import parquet.encoding
 import parquet.encoding
 
 
+
 class TestBitPacked(unittest.TestCase):
 class TestBitPacked(unittest.TestCase):
 
 
     def testFromExample(self):
     def testFromExample(self):
-        encoded_bitstring = array.array('B', [0b10001000, 0b11000110, 0b11111010]).tostring()
+        raw_data_in = [0b10001000, 0b11000110, 0b11111010]
+        encoded_bitstring = array.array('B', raw_data_in).tostring()
         fo = StringIO.StringIO(encoded_bitstring)
         fo = StringIO.StringIO(encoded_bitstring)
         count = 3 << 1
         count = 3 << 1
         res = parquet.encoding.read_bitpacked(fo, count, 3)
         res = parquet.encoding.read_bitpacked(fo, count, 3)
@@ -17,7 +19,8 @@ class TestBitPacked(unittest.TestCase):
 class TestBitPackedDeprecated(unittest.TestCase):
 class TestBitPackedDeprecated(unittest.TestCase):
 
 
     def testFromExample(self):
     def testFromExample(self):
-        encoded_bitstring = array.array('B', [0b00000101, 0b00111001, 0b01110111])
+        encoded_bitstring = array.array('B',
+                                        [0b00000101, 0b00111001, 0b01110111])
         fo = StringIO.StringIO(encoded_bitstring)
         fo = StringIO.StringIO(encoded_bitstring)
         res = parquet.encoding.read_bitpacked_deprecated(fo, 3, 3)
         res = parquet.encoding.read_bitpacked_deprecated(fo, 3, 3)
         self.assertEquals(range(8), res)
         self.assertEquals(range(8), res)
@@ -26,21 +29,21 @@ class TestBitPackedDeprecated(unittest.TestCase):
 class TestWidthFromMaxInt(unittest.TestCase):
 class TestWidthFromMaxInt(unittest.TestCase):
 
 
     def testWidths(self):
     def testWidths(self):
-        self.assertEquals(0, parquet.encoding.width_from_max_int(0));
-        self.assertEquals(1, parquet.encoding.width_from_max_int(1));
-        self.assertEquals(2, parquet.encoding.width_from_max_int(2));
-        self.assertEquals(2, parquet.encoding.width_from_max_int(3));
-        self.assertEquals(3, parquet.encoding.width_from_max_int(4));
-        self.assertEquals(3, parquet.encoding.width_from_max_int(5));
-        self.assertEquals(3, parquet.encoding.width_from_max_int(6));
-        self.assertEquals(3, parquet.encoding.width_from_max_int(7));
-        self.assertEquals(4, parquet.encoding.width_from_max_int(8));
-        self.assertEquals(4, parquet.encoding.width_from_max_int(15));
-        self.assertEquals(5, parquet.encoding.width_from_max_int(16));
-        self.assertEquals(5, parquet.encoding.width_from_max_int(31));
-        self.assertEquals(6, parquet.encoding.width_from_max_int(32));
-        self.assertEquals(6, parquet.encoding.width_from_max_int(63));
-        self.assertEquals(7, parquet.encoding.width_from_max_int(64));
-        self.assertEquals(7, parquet.encoding.width_from_max_int(127));
-        self.assertEquals(8, parquet.encoding.width_from_max_int(128));
-        self.assertEquals(8, parquet.encoding.width_from_max_int(255));
+        self.assertEquals(0, parquet.encoding.width_from_max_int(0))
+        self.assertEquals(1, parquet.encoding.width_from_max_int(1))
+        self.assertEquals(2, parquet.encoding.width_from_max_int(2))
+        self.assertEquals(2, parquet.encoding.width_from_max_int(3))
+        self.assertEquals(3, parquet.encoding.width_from_max_int(4))
+        self.assertEquals(3, parquet.encoding.width_from_max_int(5))
+        self.assertEquals(3, parquet.encoding.width_from_max_int(6))
+        self.assertEquals(3, parquet.encoding.width_from_max_int(7))
+        self.assertEquals(4, parquet.encoding.width_from_max_int(8))
+        self.assertEquals(4, parquet.encoding.width_from_max_int(15))
+        self.assertEquals(5, parquet.encoding.width_from_max_int(16))
+        self.assertEquals(5, parquet.encoding.width_from_max_int(31))
+        self.assertEquals(6, parquet.encoding.width_from_max_int(32))
+        self.assertEquals(6, parquet.encoding.width_from_max_int(63))
+        self.assertEquals(7, parquet.encoding.width_from_max_int(64))
+        self.assertEquals(7, parquet.encoding.width_from_max_int(127))
+        self.assertEquals(8, parquet.encoding.width_from_max_int(128))
+        self.assertEquals(8, parquet.encoding.width_from_max_int(255))

+ 52 - 25
test/test_read_support.py

@@ -1,35 +1,62 @@
+import StringIO
 import tempfile
 import tempfile
 import unittest
 import unittest
 
 
 import parquet
 import parquet
 
 
+
 class TestFileFormat(unittest.TestCase):
 class TestFileFormat(unittest.TestCase):
-	def test_header_magic_bytes(self):
-		with tempfile.NamedTemporaryFile() as t:
-			t.write("PAR1_some_bogus_data")
-			t.flush()
-			self.assertTrue(parquet._check_header_magic_bytes(t))
-
-	def test_footer_magic_bytes(self):
-		with tempfile.NamedTemporaryFile() as t:
-			t.write("PAR1_some_bogus_data_PAR1")
-			t.flush()
-			self.assertTrue(parquet._check_footer_magic_bytes(t))
-
-	def test_not_parquet_file(self):
-		with tempfile.NamedTemporaryFile() as t:
-			t.write("blah")
-			t.flush()
-			self.assertFalse(parquet._check_header_magic_bytes(t))
-			self.assertFalse(parquet._check_footer_magic_bytes(t))
+    def test_header_magic_bytes(self):
+        with tempfile.NamedTemporaryFile() as t:
+            t.write("PAR1_some_bogus_data")
+            t.flush()
+            self.assertTrue(parquet._check_header_magic_bytes(t))
+
+    def test_footer_magic_bytes(self):
+        with tempfile.NamedTemporaryFile() as t:
+            t.write("PAR1_some_bogus_data_PAR1")
+            t.flush()
+            self.assertTrue(parquet._check_footer_magic_bytes(t))
+
+    def test_not_parquet_file(self):
+        with tempfile.NamedTemporaryFile() as t:
+            t.write("blah")
+            t.flush()
+            self.assertFalse(parquet._check_header_magic_bytes(t))
+            self.assertFalse(parquet._check_footer_magic_bytes(t))
+
 
 
 class TestMetadata(unittest.TestCase):
 class TestMetadata(unittest.TestCase):
 
 
-	f = "/Users/joecrow/Code/parquet-compatibility/parquet-testdata/impala/1.1.1-SNAPPY/nation.impala.parquet"
-	
-	def testFooterBytes(self):
-		with open(self.f) as fo:
-			self.assertEquals(327, parquet._get_footer_size(fo))
+    f = "/Users/joecrow/Code/parquet-compatibility/parquet-testdata/impala/1.1.1-SNAPPY/nation.impala.parquet"
+
+    def test_footer_bytes(self):
+        with open(self.f) as fo:
+            self.assertEquals(327, parquet._get_footer_size(fo))
+
+    def test_read_footer(self):
+        footer = parquet.read_footer(self.f)
+        self.assertEquals(
+            set([s.name for s in footer.schema]),
+            set(["schema", "n_regionkey", "n_name", "n_nationkey",
+                 "n_comment"]))
+
+    def test_dump_metadata(self):
+        data = StringIO.StringIO()
+        parquet.dump_metadata(self.f, data)
+
+
+class TestCompatibility(unittest.TestCase):
+
+    files = []
+
+    def _test_file(self, parquet_file, csv_file):
+        """ Given the parquet_file and csv_file representation, converts the
+            parquet_file to a csv using the dump utility and then compares the
+            result to the csv_file using column agnostic ordering.
+        """
+        pass
 
 
-	def testReadFOoter(self):
-		parquet.read_footer(self.f)
+    def test_all_files(self):
+        for parquet_file, csv_file in self.files:
+            yield self._test_file, parquet_file, csv_file