test_read_support.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. import csv
  2. import os
  3. import StringIO
  4. import tempfile
  5. import unittest
  6. import parquet
  7. class TestFileFormat(unittest.TestCase):
  8. def test_header_magic_bytes(self):
  9. with tempfile.NamedTemporaryFile() as t:
  10. t.write("PAR1_some_bogus_data")
  11. t.flush()
  12. self.assertTrue(parquet._check_header_magic_bytes(t))
  13. def test_footer_magic_bytes(self):
  14. with tempfile.NamedTemporaryFile() as t:
  15. t.write("PAR1_some_bogus_data_PAR1")
  16. t.flush()
  17. self.assertTrue(parquet._check_footer_magic_bytes(t))
  18. def test_not_parquet_file(self):
  19. with tempfile.NamedTemporaryFile() as t:
  20. t.write("blah")
  21. t.flush()
  22. self.assertFalse(parquet._check_header_magic_bytes(t))
  23. self.assertFalse(parquet._check_footer_magic_bytes(t))
  24. class TestMetadata(unittest.TestCase):
  25. f = "test-data/nation.impala.parquet"
  26. def test_footer_bytes(self):
  27. with open(self.f) as fo:
  28. self.assertEquals(327, parquet._get_footer_size(fo))
  29. def test_read_footer(self):
  30. footer = parquet.read_footer(self.f)
  31. self.assertEquals(
  32. set([s.name for s in footer.schema]),
  33. set(["schema", "n_regionkey", "n_name", "n_nationkey",
  34. "n_comment"]))
  35. def test_dump_metadata(self):
  36. data = StringIO.StringIO()
  37. parquet.dump_metadata(self.f, data)
  38. class Options():
  39. col = None
  40. format = 'csv'
  41. no_headers = True
  42. limit = -1
  43. class TestCompatibility(object):
  44. td = "test-data"
  45. files = [(os.path.join(td, p), os.path.join(td, "nation.csv")) for p in
  46. ["gzip-nation.impala.parquet", "nation.dict.parquet",
  47. "nation.impala.parquet", "nation.plain.parquet",
  48. "snappy-nation.impala.parquet"]]
  49. def _test_file_csv(self, parquet_file, csv_file):
  50. """ Given the parquet_file and csv_file representation, converts the
  51. parquet_file to a csv using the dump utility and then compares the
  52. result to the csv_file using column agnostic ordering.
  53. """
  54. expected_data = []
  55. with open(csv_file, 'rb') as f:
  56. expected_data = list(csv.reader(f, delimiter='|'))
  57. actual_raw_data = StringIO.StringIO()
  58. parquet.dump(parquet_file, Options(), out=actual_raw_data)
  59. actual_raw_data.seek(0, 0)
  60. actual_data = list(csv.reader(actual_raw_data, delimiter='\t'))
  61. assert expected_data == actual_data, "{0} != {1}".format(
  62. str(expected_data), str(actual_data))
  63. def test_all_files(self):
  64. for parquet_file, csv_file in self.files:
  65. yield self._test_file_csv, parquet_file, csv_file