encoding.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. import array
  2. import math
  3. import struct
  4. import StringIO
  5. import logging
  6. from ttypes import Type
  7. logger = logging.getLogger("parquet")
  8. def read_plain_boolean(fo):
  9. """Reads a boolean using the plain encoding"""
  10. raise NotImplemented
  11. def read_plain_int32(fo):
  12. """Reads a 32-bit int using the plain encoding"""
  13. tup = struct.unpack("<i", fo.read(4))
  14. return tup[0]
  15. def read_plain_int64(fo):
  16. """Reads a 64-bit int using the plain encoding"""
  17. tup = struct.unpack("<q", fo.read(8))
  18. return tup[0]
  19. def read_plain_int96(fo):
  20. """Reads a 96-bit int using the plain encoding"""
  21. tup = struct.unpack("<qi", fo.read(12))
  22. return tup[0] << 32 | tup[1]
  23. def read_plain_float(fo):
  24. """Reads a 32-bit float using the plain encoding"""
  25. tup = struct.unpack("<f", fo.read(4))
  26. return tup[0]
  27. def read_plain_double(fo):
  28. """Reads a 64-bit float (double) using the plain encoding"""
  29. tup = struct.unpack("<d", fo.read(8))
  30. return tup[0]
  31. def read_plain_byte_array(fo):
  32. """Reads a byte array using the plain encoding"""
  33. length = read_plain_int32(fo)
  34. return fo.read(length)
  35. def read_plain_byte_array_fixed(fo, fixed_length):
  36. """Reads a byte array of the given fixed_length"""
  37. return fo.read(fixed_length)
  38. DECODE_PLAIN = {
  39. Type.BOOLEAN: read_plain_boolean,
  40. Type.INT32: read_plain_int32,
  41. Type.INT64: read_plain_int64,
  42. Type.INT96: read_plain_int96,
  43. Type.FLOAT: read_plain_float,
  44. Type.DOUBLE: read_plain_double,
  45. Type.BYTE_ARRAY: read_plain_byte_array,
  46. Type.FIXED_LEN_BYTE_ARRAY: read_plain_byte_array_fixed
  47. }
  48. def read_plain(fo, type_, type_length):
  49. conv = DECODE_PLAIN[type_]
  50. if type_ == Type.FIXED_LEN_BYTE_ARRAY:
  51. return conv(fo, type_length)
  52. return conv(fo)
  53. def read_unsigned_var_int(fo):
  54. result = 0
  55. shift = 0
  56. while True:
  57. byte = struct.unpack("<B", fo.read(1))[0]
  58. result |= ((byte & 0x7F) << shift)
  59. if (byte & 0x80) == 0:
  60. break
  61. shift += 7
  62. return result
  63. def byte_width(bit_width):
  64. "Returns the byte width for the given bit_width"
  65. return (bit_width + 7) / 8
  66. def read_rle(fo, header, bit_width):
  67. """Read a run-length encoded run from the given fo with the given header
  68. and bit_width.
  69. The count is determined from the header and the width is used to grab the
  70. value that's repeated. Yields the value repeated count times.
  71. """
  72. count = header >> 1
  73. zero_data = "\x00\x00\x00\x00"
  74. data = ""
  75. width = byte_width(bit_width)
  76. if width >= 1:
  77. data += fo.read(1)
  78. if width >= 2:
  79. data += fo.read(1)
  80. if width >= 3:
  81. data += fo.read(1)
  82. if width == 4:
  83. data += fo.read(1)
  84. data = data + zero_data[len(data):]
  85. value = struct.unpack("<i", data)[0]
  86. logger.debug("Read RLE group with value %s of byte-width %s and count %s",
  87. value, width, count)
  88. for i in range(count):
  89. yield value
  90. def width_from_max_int(value):
  91. """Converts the value specified to a bit_width."""
  92. return int(math.ceil(math.log(value + 1, 2)))
  93. def _mask_for_bits(i):
  94. """Helper function for read_bitpacked to generage a mask to grab i bits."""
  95. return (1 << i) - 1
  96. def read_bitpacked(fo, header, width):
  97. """Reads a bitpacked run of the rle/bitpack hybrid.
  98. Currently only supports width <=8 (doesn't support crossing bytes).
  99. """
  100. assert width <= 8
  101. num_groups = header >> 1
  102. logger.debug("Reading a bit-packed run with: %s groups", num_groups)
  103. count = num_groups * 8
  104. raw_bytes = array.array('B', fo.read(count)).tolist()
  105. current_byte = 0
  106. b = raw_bytes[current_byte]
  107. mask = _mask_for_bits(width)
  108. bits_in_byte = 8
  109. res = []
  110. while current_byte < len(raw_bytes):
  111. # TODO zero-padding could produce extra zero-values
  112. logger.debug(" read bitpacked: width=%s bits_in_byte=%s b=%s,"
  113. " current_byte=%s",
  114. width, bits_in_byte, bin(b), current_byte)
  115. if bits_in_byte >= width:
  116. res.append(b & mask)
  117. b >>= width
  118. bits_in_byte -= width
  119. else:
  120. if current_byte + 1 == len(raw_bytes):
  121. break # partial results / padding at the end.
  122. next_b = raw_bytes[current_byte + 1]
  123. borrowed_bits = next_b & _mask_for_bits(width - bits_in_byte)
  124. logger.debug(" borrowing %d bits", width - bits_in_byte)
  125. logger.debug(" next_b=%s, borrowed_bits=%s",
  126. bin(next_b), bin(borrowed_bits))
  127. res.append((borrowed_bits << bits_in_byte) | b)
  128. b = next_b >> (width - bits_in_byte)
  129. logger.debug(" shifting away: %d", width - bits_in_byte)
  130. bits_in_byte = 8 - (width - bits_in_byte)
  131. current_byte += 1
  132. logger.debug(" read bitpackage: added: %s", res[-1])
  133. return res
  134. def read_bitpacked_deprecated(fo, byte_count, count, width):
  135. raw_bytes = array.array('B', fo.read(byte_count)).tolist()
  136. mask = _mask_for_bits(width)
  137. index = 0
  138. res = []
  139. word = 0
  140. bits_in_word = 0
  141. while len(res) < count and index <= len(raw_bytes):
  142. logger.debug("index = %d", index)
  143. logger.debug("bits in word = %d", bits_in_word)
  144. logger.debug("word = %s", bin(word))
  145. if bits_in_word >= width:
  146. # how many bits over the value is stored
  147. offset = (bits_in_word - width)
  148. logger.debug("offset = %d", offset)
  149. # figure out the value
  150. value = (word & (mask << offset)) >> offset
  151. logger.debug("value = %d (%s)", value, bin(value))
  152. res.append(value)
  153. bits_in_word -= width
  154. else:
  155. word = (word << 8) | raw_bytes[index]
  156. index += 1
  157. bits_in_word += 8
  158. return res
  159. def read_rle_bit_packed_hybrid(fo, width, length=None):
  160. """Implemenation of a decoder for the rel/bit-packed hybrid encoding.
  161. If length is not specified, then a 32-bit int is read first to grab the
  162. length of the encoded data.
  163. """
  164. io_obj = fo
  165. if length is None:
  166. length = read_plain_int32(fo)
  167. raw_bytes = fo.read(length)
  168. if raw_bytes == '':
  169. return None
  170. io_obj = StringIO.StringIO(raw_bytes)
  171. res = []
  172. while io_obj.tell() < length:
  173. header = read_unsigned_var_int(io_obj)
  174. if header & 1 == 0:
  175. res += read_rle(io_obj, header, width)
  176. else:
  177. res += read_bitpacked(io_obj, header, width)
  178. return res