encoding.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237
  1. """encoding.py - methods for reading parquet encoded data blocks."""
  2. from __future__ import absolute_import
  3. from __future__ import division
  4. from __future__ import print_function
  5. from __future__ import unicode_literals
  6. import array
  7. import io
  8. import logging
  9. import math
  10. import os
  11. import struct
  12. import sys
  13. import thriftpy2 as thriftpy
  14. THRIFT_FILE = os.path.join(os.path.dirname(__file__), "parquet.thrift")
  15. parquet_thrift = thriftpy.load(THRIFT_FILE, module_name=str("parquet_thrift")) # pylint: disable=invalid-name
  16. logger = logging.getLogger("parquet") # pylint: disable=invalid-name
  17. PY3 = sys.version_info.major > 2
  18. ARRAY_BYTE_STR = u'B' if PY3 else b'B'
  19. def read_plain_boolean(file_obj, count):
  20. """Read `count` booleans using the plain encoding."""
  21. # for bit packed, the count is stored shifted up. But we want to pass in a count,
  22. # so we shift up.
  23. # bit width is 1 for a single-bit boolean.
  24. return read_bitpacked(file_obj, count << 1, 1, logger.isEnabledFor(logging.DEBUG))
  25. def read_plain_int32(file_obj, count):
  26. """Read `count` 32-bit ints using the plain encoding."""
  27. length = 4 * count
  28. data = file_obj.read(length)
  29. if len(data) != length:
  30. raise EOFError("Expected {} bytes but got {} bytes".format(length, len(data)))
  31. res = struct.unpack("<{}i".format(count).encode("utf-8"), data)
  32. return res
  33. def read_plain_int64(file_obj, count):
  34. """Read `count` 64-bit ints using the plain encoding."""
  35. return struct.unpack("<{}q".format(count).encode("utf-8"), file_obj.read(8 * count))
  36. def read_plain_int96(file_obj, count):
  37. """Read `count` 96-bit ints using the plain encoding."""
  38. items = struct.unpack(b"<" + b"qi" * count, file_obj.read(12 * count))
  39. return [q << 32 | i for (q, i) in zip(items[0::2], items[1::2])]
  40. def read_plain_float(file_obj, count):
  41. """Read `count` 32-bit floats using the plain encoding."""
  42. return struct.unpack("<{}f".format(count).encode("utf-8"), file_obj.read(4 * count))
  43. def read_plain_double(file_obj, count):
  44. """Read `count` 64-bit float (double) using the plain encoding."""
  45. return struct.unpack("<{}d".format(count).encode("utf-8"), file_obj.read(8 * count))
  46. def read_plain_byte_array(file_obj, count):
  47. """Read `count` byte arrays using the plain encoding."""
  48. return [file_obj.read(struct.unpack(b"<i", file_obj.read(4))[0]) for i in range(count)]
  49. def read_plain_byte_array_fixed(file_obj, fixed_length):
  50. """Read a byte array of the given fixed_length."""
  51. return file_obj.read(fixed_length)
  52. DECODE_PLAIN = {
  53. parquet_thrift.Type.BOOLEAN: read_plain_boolean,
  54. parquet_thrift.Type.INT32: read_plain_int32,
  55. parquet_thrift.Type.INT64: read_plain_int64,
  56. parquet_thrift.Type.INT96: read_plain_int96,
  57. parquet_thrift.Type.FLOAT: read_plain_float,
  58. parquet_thrift.Type.DOUBLE: read_plain_double,
  59. parquet_thrift.Type.BYTE_ARRAY: read_plain_byte_array,
  60. parquet_thrift.Type.FIXED_LEN_BYTE_ARRAY: read_plain_byte_array_fixed
  61. }
  62. def read_plain(file_obj, type_, count):
  63. """Read `count` items `type` from the fo using the plain encoding."""
  64. if count == 0:
  65. return []
  66. conv = DECODE_PLAIN[type_]
  67. return conv(file_obj, count)
  68. def read_unsigned_var_int(file_obj):
  69. """Read a value using the unsigned, variable int encoding."""
  70. result = 0
  71. shift = 0
  72. while True:
  73. byte = struct.unpack(b"<B", file_obj.read(1))[0]
  74. result |= ((byte & 0x7F) << shift)
  75. if (byte & 0x80) == 0:
  76. break
  77. shift += 7
  78. return result
  79. def read_rle(file_obj, header, bit_width, debug_logging):
  80. """Read a run-length encoded run from the given fo with the given header and bit_width.
  81. The count is determined from the header and the width is used to grab the
  82. value that's repeated. Yields the value repeated count times.
  83. """
  84. count = header >> 1
  85. zero_data = b"\x00\x00\x00\x00"
  86. width = (bit_width + 7) // 8
  87. data = file_obj.read(width)
  88. data = data + zero_data[len(data):]
  89. value = struct.unpack(b"<i", data)[0]
  90. if debug_logging:
  91. logger.debug("Read RLE group with value %s of byte-width %s and count %s",
  92. value, width, count)
  93. for _ in range(count):
  94. yield value
  95. def width_from_max_int(value):
  96. """Convert the value specified to a bit_width."""
  97. return int(math.ceil(math.log(value + 1, 2)))
  98. def _mask_for_bits(i):
  99. """Generate a mask to grab `i` bits from an int value."""
  100. return (1 << i) - 1
  101. def read_bitpacked(file_obj, header, width, debug_logging):
  102. """Read a bitpacked run of the rle/bitpack hybrid.
  103. Supports width >8 (crossing bytes).
  104. """
  105. num_groups = header >> 1
  106. count = num_groups * 8
  107. byte_count = (width * count) // 8
  108. if debug_logging:
  109. logger.debug("Reading a bit-packed run with: %s groups, count %s, bytes %s",
  110. num_groups, count, byte_count)
  111. if width == 0:
  112. return [0 for _ in range(count)]
  113. raw_bytes = array.array(ARRAY_BYTE_STR, file_obj.read(byte_count)).tolist()
  114. current_byte = 0
  115. data = raw_bytes[current_byte]
  116. mask = _mask_for_bits(width)
  117. bits_wnd_l = 8
  118. bits_wnd_r = 0
  119. res = []
  120. total = len(raw_bytes) * 8
  121. while total >= width:
  122. # NOTE zero-padding could produce extra zero-values
  123. if debug_logging:
  124. logger.debug(" read bitpacked: width=%s window=(%s %s) b=%s,"
  125. " current_byte=%s",
  126. width, bits_wnd_l, bits_wnd_r, bin(data), current_byte)
  127. if bits_wnd_r >= 8:
  128. bits_wnd_r -= 8
  129. bits_wnd_l -= 8
  130. data >>= 8
  131. elif bits_wnd_l - bits_wnd_r >= width:
  132. res.append((data >> bits_wnd_r) & mask)
  133. total -= width
  134. bits_wnd_r += width
  135. if debug_logging:
  136. logger.debug(" read bitpackage: added: %s", res[-1])
  137. elif current_byte + 1 < len(raw_bytes):
  138. current_byte += 1
  139. data |= (raw_bytes[current_byte] << bits_wnd_l)
  140. bits_wnd_l += 8
  141. return res
  142. def read_bitpacked_deprecated(file_obj, byte_count, count, width, debug_logging):
  143. """Read `count` values from `fo` using the deprecated bitpacking encoding."""
  144. raw_bytes = array.array(ARRAY_BYTE_STR, file_obj.read(byte_count)).tolist()
  145. mask = _mask_for_bits(width)
  146. index = 0
  147. res = []
  148. word = 0
  149. bits_in_word = 0
  150. while len(res) < count and index <= len(raw_bytes):
  151. if debug_logging:
  152. logger.debug("index = %d", index)
  153. logger.debug("bits in word = %d", bits_in_word)
  154. logger.debug("word = %s", bin(word))
  155. if bits_in_word >= width:
  156. # how many bits over the value is stored
  157. offset = (bits_in_word - width)
  158. # figure out the value
  159. value = (word & (mask << offset)) >> offset
  160. if debug_logging:
  161. logger.debug("offset = %d", offset)
  162. logger.debug("value = %d (%s)", value, bin(value))
  163. res.append(value)
  164. bits_in_word -= width
  165. else:
  166. word = (word << 8) | raw_bytes[index]
  167. index += 1
  168. bits_in_word += 8
  169. return res
  170. def read_rle_bit_packed_hybrid(file_obj, width, length=None):
  171. """Read values from `fo` using the rel/bit-packed hybrid encoding.
  172. If length is not specified, then a 32-bit int is read first to grab the
  173. length of the encoded data.
  174. """
  175. debug_logging = logger.isEnabledFor(logging.DEBUG)
  176. io_obj = file_obj
  177. if length is None:
  178. length = read_plain_int32(file_obj, 1)[0]
  179. raw_bytes = file_obj.read(length)
  180. if raw_bytes == b'':
  181. return None
  182. io_obj = io.BytesIO(raw_bytes)
  183. res = []
  184. while io_obj.tell() < length:
  185. header = read_unsigned_var_int(io_obj)
  186. if header & 1 == 0:
  187. res += read_rle(io_obj, header, width, debug_logging)
  188. else:
  189. res += read_bitpacked(io_obj, header, width, debug_logging)
  190. return res