test_read_support.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. import csv
  2. import json
  3. import os
  4. import StringIO
  5. import tempfile
  6. import unittest
  7. import parquet
  8. class TestFileFormat(unittest.TestCase):
  9. def test_header_magic_bytes(self):
  10. with tempfile.NamedTemporaryFile() as t:
  11. t.write("PAR1_some_bogus_data")
  12. t.flush()
  13. self.assertTrue(parquet._check_header_magic_bytes(t))
  14. def test_footer_magic_bytes(self):
  15. with tempfile.NamedTemporaryFile() as t:
  16. t.write("PAR1_some_bogus_data_PAR1")
  17. t.flush()
  18. self.assertTrue(parquet._check_footer_magic_bytes(t))
  19. def test_not_parquet_file(self):
  20. with tempfile.NamedTemporaryFile() as t:
  21. t.write("blah")
  22. t.flush()
  23. self.assertFalse(parquet._check_header_magic_bytes(t))
  24. self.assertFalse(parquet._check_footer_magic_bytes(t))
  25. class TestMetadata(unittest.TestCase):
  26. f = "test-data/nation.impala.parquet"
  27. def test_footer_bytes(self):
  28. with open(self.f) as fo:
  29. self.assertEquals(327, parquet._get_footer_size(fo))
  30. def test_read_footer(self):
  31. footer = parquet.read_footer(self.f)
  32. self.assertEquals(
  33. set([s.name for s in footer.schema]),
  34. set(["schema", "n_regionkey", "n_name", "n_nationkey",
  35. "n_comment"]))
  36. def test_dump_metadata(self):
  37. data = StringIO.StringIO()
  38. parquet.dump_metadata(self.f, data)
  39. class Options(object):
  40. def __init__(self, col=None, format='csv', no_headers=True, limit=-1):
  41. self.col = col
  42. self.format = format
  43. self.no_headers = no_headers
  44. self.limit = limit
  45. class TestReadApi(unittest.TestCase):
  46. def test_projection(self):
  47. pass
  48. def test_limit(self):
  49. pass
  50. class TestCompatibility(object):
  51. td = "test-data"
  52. files = [(os.path.join(td, p), os.path.join(td, "nation.csv")) for p in
  53. ["gzip-nation.impala.parquet", "nation.dict.parquet",
  54. "nation.impala.parquet", "nation.plain.parquet",
  55. "snappy-nation.impala.parquet"]]
  56. def _test_file_csv(self, parquet_file, csv_file):
  57. """ Given the parquet_file and csv_file representation, converts the
  58. parquet_file to a csv using the dump utility and then compares the
  59. result to the csv_file.
  60. """
  61. expected_data = []
  62. with open(csv_file, 'rb') as f:
  63. expected_data = list(csv.reader(f, delimiter='|'))
  64. actual_raw_data = StringIO.StringIO()
  65. parquet.dump(parquet_file, Options(), out=actual_raw_data)
  66. actual_raw_data.seek(0, 0)
  67. actual_data = list(csv.reader(actual_raw_data, delimiter='\t'))
  68. assert expected_data == actual_data, "{0} != {1}".format(
  69. str(expected_data), str(actual_data))
  70. actual_raw_data = StringIO.StringIO()
  71. parquet.dump(parquet_file, Options(no_headers=False),
  72. out=actual_raw_data)
  73. actual_raw_data.seek(0, 0)
  74. actual_data = list(csv.reader(actual_raw_data, delimiter='\t'))[1:]
  75. assert expected_data == actual_data, "{0} != {1}".format(
  76. str(expected_data), str(actual_data))
  77. def _test_file_json(self, parquet_file, csv_file):
  78. """ Given the parquet_file and csv_file representation, converts the
  79. parquet_file to json using the dump utility and then compares the
  80. result to the csv_file using column agnostic ordering.
  81. """
  82. expected_data = []
  83. with open(csv_file, 'rb') as f:
  84. expected_data = list(csv.reader(f, delimiter='|'))
  85. actual_raw_data = StringIO.StringIO()
  86. parquet.dump(parquet_file, Options(format='json'),
  87. out=actual_raw_data)
  88. actual_raw_data.seek(0, 0)
  89. actual_data = [json.loads(x.rstrip()) for x in
  90. actual_raw_data.read().split("\n") if len(x) > 0]
  91. assert len(expected_data) == len(actual_data)
  92. footer = parquet.read_footer(parquet_file)
  93. cols = [s.name for s in footer.schema]
  94. for expected, actual in zip(expected_data, actual_raw_data):
  95. assert len(expected) == len(actual)
  96. for i, c in enumerate(cols):
  97. if c in actual:
  98. assert expected[i] == actual[c]
  99. def test_all_files(self):
  100. for parquet_file, csv_file in self.files:
  101. yield self._test_file_csv, parquet_file, csv_file
  102. yield self._test_file_json, parquet_file, csv_file