Ver código fonte

HUE-5009 [core] Backport parquet-python Add support for converted_types to dict-encoding.

Commit https://github.com/jcrobak/parquet-python/commit/45165f3159505524d708894337e68120fcd844e7

Jenny Kim 9 anos atrás
pai
commit
9ff637d

+ 7 - 4
desktop/core/ext-py/parquet-1.1/parquet/__init__.py

@@ -350,6 +350,7 @@ def read_data_page(file_obj, schema_helper, page_header, column_metadata,
             vals.extend(read_values)
         if debug_logging:
             logger.debug("  Values: %s, nulls: %s", len(vals), num_nulls)
+
     elif daph.encoding == parquet_thrift.Encoding.PLAIN_DICTIONARY:
         # bit_width is stored as single byte.
         bit_width = struct.unpack("<B", io_obj.read(1))[0]
@@ -387,18 +388,20 @@ def read_data_page(file_obj, schema_helper, page_header, column_metadata,
     return vals
 
 
-def _read_dictionary_page(file_obj, page_header, column_metadata):
+def _read_dictionary_page(file_obj, schema_helper, page_header, column_metadata):
     """Read a page containing dictionary data.
-
     Consumes data using the plain encoding and returns an array of values.
     """
     raw_bytes = _read_page(file_obj, page_header, column_metadata)
     io_obj = io.BytesIO(raw_bytes)
-    return encoding.read_plain(
+    values = encoding.read_plain(
         io_obj,
         column_metadata.type,
         page_header.dictionary_page_header.num_values
     )
+    # convert the values once, if the dictionary is associated with a converted_type.
+    schema_element = schema_helper.schema_element(column_metadata.path_in_schema[-1])
+    return convert_column(values, schema_element) if schema_element.converted_type is not None else values
 
 
 def DictReader(file_obj, columns=None):  # pylint: disable=invalid-name
@@ -475,7 +478,7 @@ def reader(file_obj, columns=None):
                     if debug_logging:
                         logger.debug(page_header)
                     assert dict_items == []
-                    dict_items = _read_dictionary_page(file_obj, page_header, cmd)
+                    dict_items = _read_dictionary_page(file_obj, schema_helper, page_header, cmd)
                     if debug_logging:
                         logger.debug("Dictionary: %s", str(dict_items))
                 else:

+ 12 - 0
desktop/core/ext-py/parquet-1.1/test/test_read_support.py

@@ -228,3 +228,15 @@ class TestDefinitionLevel(unittest.TestCase):
             [{"foo": "bar"}, {"foo": None}],
             actual_data
         )
+
+    def test_null_plain_dictionary(self):
+        """Test reading a file that contains null records for a plain dictionary column."""
+        with open(os.path.join(TEST_DATA, "test-null-dictionary.parquet"), "rb") as parquet_fo:
+            actual_data = list(parquet.DictReader(parquet_fo))
+
+        self.assertListEqual(
+            # this is the contents of test-null-dictionary.parquet. 7 records.
+            # The first record is null, and the rest alternate between values of 'bar' and 'baz.'
+            [{"foo": None}] + [{"foo": "bar"}, {"foo": "baz"}] * 3,
+            actual_data
+        )