Browse Source

This adds the ability to output arbitrary datatypes.

Luke Carmichael 10 years ago
parent
commit
042567d155
2 changed files with 36 additions and 0 deletions
  1. 3 0
      parquet/__init__.py
  2. 33 0
      test/test_read_support.py

+ 3 - 0
parquet/__init__.py

@@ -372,6 +372,9 @@ def _dump(fo, options, out=sys.stdout):
                         _get_name(PageType, ph.type)))
         keys = options.col if options.col else [s.name for s in
                                                 footer.schema if s.name in res]
+        if options.format == 'custom':
+            custom_datatype = out(res, keys)
+            return custom_datatype
         if options.format == "csv" and not options.no_headers:
             println("\t".join(keys))
         for i in range(rg.num_rows):

+ 33 - 0
test/test_read_support.py

@@ -125,9 +125,42 @@ class TestCompatibility(object):
                 if c in actual:
                     assert expected[i] == actual[c]
 
+    def _test_file_custom(self, parquet_file, csv_file):
+        """ Given the parquet_file and csv_file representation, converts the
+            parquet_file to json using the dump utility and then compares the
+            result to the csv_file using column agnostic ordering.
+        """
+        expected_data = []
+        with open(csv_file, 'rb') as f:
+            expected_data = list(csv.reader(f, delimiter='|'))
+
+        def _custom_datatype(in_dict, keys):
+            '''
+            return rows like the csv outputter
+
+            Could convert to a dataframe like this:
+                import pandas
+                df = pandas.DataFrame(in_dict)
+                return df
+            '''
+            columns = [in_dict[key] for key in keys]
+            rows = zip(*columns)
+            return rows
 
+        actual_data = parquet.dump(parquet_file, Options(format='custom'), out=_custom_datatype)
+
+        assert len(expected_data) == len(actual_data)
+        footer = parquet.read_footer(parquet_file)
+        cols = [s.name for s in footer.schema]
+
+        for expected, actual in zip(expected_data, actual_data):
+            assert len(expected) == len(actual)
+            for i, c in enumerate(cols):
+                if c in actual:
+                    assert expected[i] == actual[c]
 
     def test_all_files(self):
         for parquet_file, csv_file in self.files:
             yield self._test_file_csv, parquet_file, csv_file
             yield self._test_file_json, parquet_file, csv_file
+            yield self._test_file_custom, parquet_file, csv_file