encoding.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. import array
  2. import math
  3. import struct
  4. import StringIO
  5. from ttypes import Type
  6. def read_plain_boolean(fo):
  7. raise NotImplemented
  8. def read_plain_int32(fo):
  9. tup = struct.unpack("<i", fo.read(4))
  10. return tup[0]
  11. def read_plain_int64(fo):
  12. tup = struct.unpack("<q", fo.read(8))
  13. return tup[0]
  14. def read_plain_int96(fo):
  15. tup = struct.unpack("<q<i", fo.read(12))
  16. return tup[0] << 32 | tup[1]
  17. def read_plain_float(fo):
  18. tup = struct.unpack("<f", fo.read(4))
  19. def read_plain_double(fo):
  20. tup = struct.unpack("<d", fo.read(8))
  21. def read_plain_byte_array(fo):
  22. length = read_plain_int32(fo)
  23. return fo.read(length)
  24. def read_plain_byte_array_fixed(fo, fixed_length):
  25. return fo.read(fixed_length)
  26. DECODE_PLAIN = {
  27. Type.BOOLEAN: read_plain_boolean,
  28. Type.INT32: read_plain_int32,
  29. Type.INT64: read_plain_int64,
  30. Type.INT96: read_plain_int96,
  31. Type.FLOAT: read_plain_float,
  32. Type.DOUBLE: read_plain_double,
  33. Type.BYTE_ARRAY: read_plain_byte_array,
  34. Type.FIXED_LEN_BYTE_ARRAY: read_plain_byte_array_fixed
  35. }
  36. def read_plain(fo, type_, type_length):
  37. conv = DECODE_PLAIN[type_]
  38. if type_ == Type.FIXED_LEN_BYTE_ARRAY:
  39. return conv(fo, type_length)
  40. return conv(fo)
  41. def read_unsigned_var_int(fo):
  42. result = 0
  43. shift = 0
  44. while True:
  45. byte = struct.unpack("<B", fo.read(1))[0]
  46. result |= ((byte & 0x7F) << shift)
  47. if (byte & 0x80) == 0:
  48. break
  49. shift += 7
  50. return result
  51. def byte_width(bit_width):
  52. "Returns the byte width for the given bit_width"
  53. return (bit_width + 7) / 8;
  54. def read_rle(fo, header, bit_width):
  55. """Grabs count from the header and uses width to grab the value that's
  56. repeated. Returns an array with the value repeated count times."""
  57. count = header >> 1
  58. zero_data = "\x00\x00\x00\x00"
  59. data = ""
  60. width = byte_width(bit_width)
  61. if width >= 1:
  62. data += fo.read(1)
  63. elif width >= 2:
  64. data += fo.read(1)
  65. elif width >= 3:
  66. data += fo.read(1)
  67. elif width == 4:
  68. data = fo.read(1)
  69. data = data + zero_data[len(data):]
  70. value = struct.unpack("<i", data)[0]
  71. return [value]*count
  72. def width_from_max_int(value):
  73. return int(math.ceil(math.log(value + 1, 2)))
  74. def mask_for_bits(i):
  75. return (1 << i) - 1
  76. def read_bitpacked(fo, header, width):
  77. num_groups = header >> 1;
  78. count = num_groups * 8
  79. raw_bytes = array.array('B', fo.read(count)).tolist()
  80. current_byte = 0
  81. b = raw_bytes[current_byte]
  82. mask = mask_for_bits(width)
  83. bits_in_byte = 8
  84. res = []
  85. while current_byte < width and len(res) < (count / width):
  86. print "width={0} bits_in_byte={1} b={2}".format(width, bits_in_byte, bin(b))
  87. if bits_in_byte >= width:
  88. res.append(b & mask)
  89. b >>= width
  90. bits_in_byte -= width
  91. else:
  92. next_b = raw_bytes[current_byte + 1]
  93. borrowed_bits = next_b & mask_for_bits(width - bits_in_byte)
  94. #print " borrowing {0} bites".format(width - bits_in_byte)
  95. #print " next_b={0}, borrowed_bits={1}".format(bin(next_b), bin(borrowed_bits))
  96. res.append((borrowed_bits << bits_in_byte) | b)
  97. b = next_b >> (width - bits_in_byte)
  98. #print " shifting away: {0}".format(width - bits_in_byte)
  99. bits_in_byte = 8 - (width - bits_in_byte)
  100. current_byte += 1
  101. print " added: {0}".format(res[-1])
  102. return res
  103. def read_bitpacked_deprecated(fo, count, width):
  104. res = []
  105. raw_bytes = array.array('B', fo.read(count)).tolist()
  106. current_byte = 0
  107. b = raw_bytes[current_byte]
  108. mask = mask_for_bits(width)
  109. def read_rle_bit_packed_hybrid(fo, width, length=None):
  110. # import pdb; pdb.set_trace()
  111. io_obj = fo
  112. if length is None:
  113. length = read_plain_int32(fo)
  114. raw_bytes = fo.read(length)
  115. if raw_bytes == '':
  116. return None
  117. io_obj = StringIO.StringIO(raw_bytes)
  118. res = []
  119. while io_obj.tell() < length:
  120. header = read_unsigned_var_int(io_obj)
  121. if header & 1 == 0:
  122. res += read_rle(io_obj, header, width)
  123. else:
  124. res += read_bitpacked(io_obj, header, width)
  125. return res