Bladeren bron

HUE-5009 [core] Backport parquet-python Logs error when a fileobj isn't in binary mode

Commit https://github.com/jcrobak/parquet-python/commit/d18e5d111521bb216de75271467eb2fa22803d3e
Jenny Kim 9 jaren geleden
bovenliggende
commit
9cc9b8c

+ 6 - 1
desktop/core/ext-py/parquet-1.1/parquet/__init__.py

@@ -20,6 +20,7 @@ except ImportError:
 import thriftpy
 from thriftpy.protocol.compact import TCompactProtocolFactory
 
+from .converted_types import convert_column
 from .thrift_filetransport import TFileTransport
 from . import encoding
 from . import schema
@@ -395,6 +396,8 @@ def reader(fo, columns=None):
     :param columns: the columns to include. If None (default), all columns
                     are included. Nested values are referenced with "." notation
     """
+    if hasattr(fo, 'mode') and 'b' not in fo.mode:
+        logger.error("parquet.reader requires the fileobj to be opened in binary mode!")
     footer = _read_footer(fo)
     schema_helper = schema.SchemaHelper(footer.schema)
     keys = columns if columns else [s.name for s in
@@ -429,7 +432,9 @@ def reader(fo, columns=None):
                 if ph.type == parquet_thrift.PageType.DATA_PAGE:
                     values = read_data_page(fo, schema_helper, ph, cmd,
                                             dict_items)
-                    res[".".join(cmd.path_in_schema)] += values
+                    schema_element = schema_helper.schema_element(cmd.path_in_schema[-1])
+                    res[".".join(cmd.path_in_schema)] += convert_column(values,
+                                                                        schema_element) if schema_element.converted_type else values
                     values_seen += ph.data_page_header.num_values
                 elif ph.type == parquet_thrift.PageType.DICTIONARY_PAGE:
                     if debug_logging:

+ 87 - 0
desktop/core/ext-py/parquet-1.1/parquet/converted_types.py

@@ -0,0 +1,87 @@
+# -#- coding: utf-8 -#-
+"""
+Deal with parquet logical types (aka converted types), higher-order
+things built from primitive types.
+
+The implementations in this class are pure python for the widest compatibility,
+but they're not necessarily the most performant.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import codecs
+import datetime
+import logging
+import json
+import os
+import struct
+import sys
+from decimal import Decimal
+
+import thriftpy
+
+THRIFT_FILE = os.path.join(os.path.dirname(__file__), "parquet.thrift")
+parquet_thrift = thriftpy.load(THRIFT_FILE, module_name=str("parquet_thrift"))
+
+logger = logging.getLogger('parquet')
+
+bson = None
+try:
+    import bson
+except ImportError:
+    pass
+
+PY3 = sys.version_info.major > 2
+
+# define bytes->int for non 2, 4, 8 byte ints
+if PY3:
+    intbig = lambda x: int.from_bytes(x, 'big', signed=True)
+else:
+    intbig = lambda x: int(codecs.encode(x, 'hex'), 16)
+
+DAYS_TO_MILLIS = 86400000000000
+"""Number of millis in a day. Used to convert a Date to a date"""
+
+
+def convert_unsigned(data, fmt):
+    num = len(data)
+    return struct.unpack(
+        "{}{}".format(num, fmt.upper()),
+        struct.pack("{}{}".format(num, fmt), *data)
+        )
+
+def convert_column(data, schemae):
+    """Convert known types from primitive to rich."""
+    ctype = schemae.converted_type
+    if ctype == parquet_thrift.ConvertedType.DECIMAL:
+        scale_factor = Decimal("10e-{}".format(schemae.scale))
+        if schemae.type == parquet_thrift.Type.INT32 or schemae.type == parquet_thrift.Type.INT64:
+            return [Decimal(unscaled) * scale_factor for unscaled in data]
+        return [Decimal(intbig(unscaled)) * scale_factor for unscaled in data]
+    elif ctype == parquet_thrift.ConvertedType.DATE:
+        return [datetime.date.fromordinal(d) for d in data]
+    elif ctype == parquet_thrift.ConvertedType.TIME_MILLIS:
+        return [datetime.timedelta(milliseconds=d) for d in data]
+    elif ctype == parquet_thrift.ConvertedType.TIMESTAMP_MILLIS:
+        return [datetime.datetime.utcfromtimestamp(d/1000.0) for d in data]
+    elif ctype == parquet_thrift.ConvertedType.UTF8:
+        return list(codecs.iterdecode(data, "utf-8"))
+    elif ctype == parquet_thrift.ConvertedType.UINT_8:
+        return convert_unsigned(data, 'b')
+    elif ctype == parquet_thrift.ConvertedType.UINT_16:
+        return convert_unsigned(data, 'h')
+    elif ctype == parquet_thrift.ConvertedType.UINT_32:
+        return convert_unsigned(data, 'i')
+    elif ctype == parquet_thrift.ConvertedType.UINT_64:
+        return convert_unsigned(data, 'q')
+    elif ctype == parquet_thrift.ConvertedType.JSON:
+        return [json.loads(s) for s in codecs.iterdecode(data, "utf-8")]
+    elif ctype == parquet_thrift.ConvertedType.BSON and bson:
+        return [bson.BSON(s).decode() for s in data]
+    else:
+        logger.warn("Converted type '{}'' not handled".format(
+            parquet_thrift.ConvertedType._VALUES_TO_NAMES[ctype]))
+    return data

+ 189 - 0
desktop/core/ext-py/parquet-1.1/test/test_converted_types.py

@@ -0,0 +1,189 @@
+# -*- coding: UTF-8 -*-
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import array
+import datetime
+import struct
+import io
+import unittest
+from decimal import Decimal
+
+from parquet.converted_types import convert_column
+from parquet import parquet_thrift as pt
+from bson import Binary
+
+class TestDecimal(unittest.TestCase):
+
+	def test_int32(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.INT32,
+			name="test",
+			converted_type=pt.ConvertedType.DECIMAL,
+			scale=10,
+			precision=9
+		)
+
+		self.assertEquals(
+			convert_column([9876543210], schema)[0],
+			Decimal('9.87654321')
+		)
+
+	def test_int64(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.INT64,
+			name="test",
+			converted_type=pt.ConvertedType.DECIMAL,
+			scale=3,
+			precision=13
+		)
+
+		self.assertEquals(
+			convert_column([1099511627776], schema)[0],
+			Decimal('10995116277.76')
+		)
+
+	def test_fixedlength(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.FIXED_LEN_BYTE_ARRAY,
+			type_length=3,
+			name="test",
+			converted_type=pt.ConvertedType.DECIMAL,
+			scale=3,
+			precision=13
+		)
+
+		self.assertEquals(
+			convert_column([b'\x02\x00\x01'], schema)[0],
+			Decimal('1310.73')
+		)
+
+	def test_binary(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.BYTE_ARRAY,
+			name="test",
+			converted_type=pt.ConvertedType.DECIMAL,
+			scale=3,
+			precision=13
+		)
+
+		self.assertEquals(
+			convert_column([b'\x02\x00\x00\x00\x00\x00\x00\x00\x00\x01'], schema)[0],
+			Decimal('94447329657392904273.93')
+		)
+
+
+class TestDateTypes(unittest.TestCase):
+
+	def test_date(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.INT32,
+			name="test",
+			converted_type=pt.ConvertedType.DATE,
+		)
+		self.assertEquals(
+			convert_column([731888], schema)[0],
+			datetime.date(2004, 11, 3)
+		)
+
+	def test_time_millis(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.INT32,
+			name="test",
+			converted_type=pt.ConvertedType.TIME_MILLIS,
+		)
+		self.assertEquals(
+			convert_column([731888], schema)[0],
+			datetime.timedelta(milliseconds=731888)
+		)
+
+	def test_timestamp_millis(self):
+		schema = pt.SchemaElement(
+			type=pt.Type.INT64,
+			name="test",
+			converted_type=pt.ConvertedType.TIMESTAMP_MILLIS,
+		)
+		self.assertEquals(
+			convert_column([1099511625014], schema)[0],
+			datetime.datetime(2004, 11, 3, 19, 53, 45, 14*1000)
+		)
+
+	def test_utf8(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.BYTE_ARRAY,
+			name="test",
+			converted_type=pt.ConvertedType.UTF8
+		)
+		data = b'foo\xf0\x9f\x91\xbe'
+		self.assertEquals(
+			convert_column([data], schema)[0],
+			'foo👾'
+		)
+
+	def test_uint8(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.INT32,
+			name="test",
+			converted_type=pt.ConvertedType.UINT_8
+		)
+		self.assertEquals(
+			convert_column([-3], schema)[0],
+			253
+		)
+
+	def test_uint16(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.INT32,
+			name="test",
+			converted_type=pt.ConvertedType.UINT_16
+		)
+		self.assertEquals(
+			convert_column([-3], schema)[0],
+			65533
+		)
+
+	def test_uint32(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.INT32,
+			name="test",
+			converted_type=pt.ConvertedType.UINT_32
+		)
+		self.assertEquals(
+			convert_column([-6884376], schema)[0],
+			4288082920
+		)
+
+	def test_uint64(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.INT64,
+			name="test",
+			converted_type=pt.ConvertedType.UINT_64
+		)
+		self.assertEquals(
+			convert_column([-6884376], schema)[0],
+			18446744073702667240
+		)
+
+	def test_json(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.BYTE_ARRAY,
+			name="test",
+			converted_type=pt.ConvertedType.JSON
+		)
+		self.assertEquals(
+			convert_column([b'{"foo": ["bar", "\\ud83d\\udc7e"]}'], schema)[0],
+			{'foo': ['bar','👾']}
+		)
+
+	def test_bson(self):
+		schema = pt.SchemaElement(
+			type = pt.Type.BYTE_ARRAY,
+			name="test",
+			converted_type=pt.ConvertedType.BSON
+		)
+		self.assertEquals(
+			convert_column([b'&\x00\x00\x00\x04foo\x00\x1c\x00\x00\x00\x020\x00\x04\x00\x00\x00bar\x00\x021\x00\x05\x00\x00\x00\xf0\x9f\x91\xbe\x00\x00\x00'], schema)[0],
+			{'foo': ['bar','👾']}
+		)

+ 8 - 0
desktop/core/ext-py/parquet-1.1/tox.ini

@@ -0,0 +1,8 @@
+[tox]
+envlist = py27, py34, py35, pypy
+[testenv]
+deps=
+   nose
+   python-snappy
+   pymongo
+commands=nosetests