test_encoding.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. import array
  2. import struct
  3. import StringIO
  4. import unittest
  5. import parquet.encoding
  6. from parquet.ttypes import Type
  7. class TestPlain(unittest.TestCase):
  8. def test_int32(self):
  9. self.assertEquals(
  10. 999,
  11. parquet.encoding.read_plain_int32(
  12. StringIO.StringIO(struct.pack("<i", 999))))
  13. def test_int64(self):
  14. self.assertEquals(
  15. 999,
  16. parquet.encoding.read_plain_int64(
  17. StringIO.StringIO(struct.pack("<q", 999))))
  18. def test_int96(self):
  19. self.assertEquals(
  20. 999,
  21. parquet.encoding.read_plain_int96(
  22. StringIO.StringIO(struct.pack("<qi", 0, 999))))
  23. def test_float(self):
  24. self.assertAlmostEquals(
  25. 9.99,
  26. parquet.encoding.read_plain_float(
  27. StringIO.StringIO(struct.pack("<f", 9.99))),
  28. 2)
  29. def test_double(self):
  30. self.assertEquals(
  31. 9.99,
  32. parquet.encoding.read_plain_double(
  33. StringIO.StringIO(struct.pack("<d", 9.99))))
  34. def test_fixed(self):
  35. data = "foobar"
  36. fo = StringIO.StringIO(data)
  37. self.assertEquals(
  38. data[:3],
  39. parquet.encoding.read_plain_byte_array_fixed(
  40. fo, 3))
  41. self.assertEquals(
  42. data[3:],
  43. parquet.encoding.read_plain_byte_array_fixed(
  44. fo, 3))
  45. def test_fixed_read_plain(self):
  46. data = "foobar"
  47. fo = StringIO.StringIO(data)
  48. self.assertEquals(
  49. data[:3],
  50. parquet.encoding.read_plain(
  51. fo, Type.FIXED_LEN_BYTE_ARRAY, 3))
  52. class TestRle(unittest.TestCase):
  53. def testFourByteValue(self):
  54. fo = StringIO.StringIO(struct.pack("<i", 1 << 30))
  55. out = parquet.encoding.read_rle(fo, 2 << 1, 30)
  56. self.assertEquals([1 << 30] * 2, list(out))
  57. class TestVarInt(unittest.TestCase):
  58. def testSingleByte(self):
  59. fo = StringIO.StringIO(struct.pack("<B", 0x7F))
  60. out = parquet.encoding.read_unsigned_var_int(fo)
  61. self.assertEquals(0x7F, out)
  62. def testFourByte(self):
  63. fo = StringIO.StringIO(struct.pack("<BBBB", 0xFF, 0xFF, 0xFF, 0x7F))
  64. out = parquet.encoding.read_unsigned_var_int(fo)
  65. self.assertEquals(0x0FFFFFFF, out)
  66. class TestBitPacked(unittest.TestCase):
  67. def testFromExample(self):
  68. raw_data_in = [0b10001000, 0b11000110, 0b11111010]
  69. encoded_bitstring = array.array('B', raw_data_in).tostring()
  70. fo = StringIO.StringIO(encoded_bitstring)
  71. count = 3 << 1
  72. res = parquet.encoding.read_bitpacked(fo, count, 3)
  73. self.assertEquals(range(8), res)
  74. class TestBitPackedDeprecated(unittest.TestCase):
  75. def testFromExample(self):
  76. encoded_bitstring = array.array('B',
  77. [0b00000101, 0b00111001, 0b01110111])
  78. fo = StringIO.StringIO(encoded_bitstring)
  79. res = parquet.encoding.read_bitpacked_deprecated(fo, 3, 3)
  80. self.assertEquals(range(8), res)
  81. class TestWidthFromMaxInt(unittest.TestCase):
  82. def testWidths(self):
  83. self.assertEquals(0, parquet.encoding.width_from_max_int(0))
  84. self.assertEquals(1, parquet.encoding.width_from_max_int(1))
  85. self.assertEquals(2, parquet.encoding.width_from_max_int(2))
  86. self.assertEquals(2, parquet.encoding.width_from_max_int(3))
  87. self.assertEquals(3, parquet.encoding.width_from_max_int(4))
  88. self.assertEquals(3, parquet.encoding.width_from_max_int(5))
  89. self.assertEquals(3, parquet.encoding.width_from_max_int(6))
  90. self.assertEquals(3, parquet.encoding.width_from_max_int(7))
  91. self.assertEquals(4, parquet.encoding.width_from_max_int(8))
  92. self.assertEquals(4, parquet.encoding.width_from_max_int(15))
  93. self.assertEquals(5, parquet.encoding.width_from_max_int(16))
  94. self.assertEquals(5, parquet.encoding.width_from_max_int(31))
  95. self.assertEquals(6, parquet.encoding.width_from_max_int(32))
  96. self.assertEquals(6, parquet.encoding.width_from_max_int(63))
  97. self.assertEquals(7, parquet.encoding.width_from_max_int(64))
  98. self.assertEquals(7, parquet.encoding.width_from_max_int(127))
  99. self.assertEquals(8, parquet.encoding.width_from_max_int(128))
  100. self.assertEquals(8, parquet.encoding.width_from_max_int(255))