Browse Source

support width > 8

Jaguar Xiong 11 years ago
parent
commit
6001b5e777
1 changed files with 18 additions and 22 deletions
  1. 18 22
      parquet/encoding.py

+ 18 - 22
parquet/encoding.py

@@ -132,7 +132,6 @@ def read_bitpacked(fo, header, width):
 
     Currently only supports width <=8 (doesn't support crossing bytes).
     """
-    assert width <= 8
     num_groups = header >> 1
     logger.debug("Reading a bit-packed run with: %s groups", num_groups)
     count = num_groups * 8
@@ -140,31 +139,28 @@ def read_bitpacked(fo, header, width):
     current_byte = 0
     b = raw_bytes[current_byte]
     mask = _mask_for_bits(width)
-    bits_in_byte = 8
+    bits_wnd_l = 8
+    bits_wnd_r = 0
     res = []
-    while current_byte < len(raw_bytes):
+    total = len(raw_bytes)*8;
+    while (total >= width):
         # TODO zero-padding could produce extra zero-values
-        logger.debug("  read bitpacked: width=%s bits_in_byte=%s b=%s,"
+        logger.debug("  read bitpacked: width=%s window=(%s %s) b=%s,"
                      " current_byte=%s",
-                     width, bits_in_byte, bin(b), current_byte)
-        if bits_in_byte >= width:
-            res.append(b & mask)
-            b >>= width
-            bits_in_byte -= width
-        else:
-            if current_byte + 1 == len(raw_bytes):
-                break  # partial results / padding at the end.
-            next_b = raw_bytes[current_byte + 1]
-            borrowed_bits = next_b & _mask_for_bits(width - bits_in_byte)
-            logger.debug("    borrowing %d bits", width - bits_in_byte)
-            logger.debug("    next_b=%s, borrowed_bits=%s",
-                         bin(next_b), bin(borrowed_bits))
-            res.append((borrowed_bits << bits_in_byte) | b)
-            b = next_b >> (width - bits_in_byte)
-            logger.debug("  shifting away: %d", width - bits_in_byte)
-            bits_in_byte = 8 - (width - bits_in_byte)
+                     width, bits_wnd_l, bits_wnd_r, bin(b), current_byte)
+        if bits_wnd_r >= 8:
+            bits_wnd_r -= 8
+            bits_wnd_l -= 8
+            b >>= 8
+        elif bits_wnd_l - bits_wnd_r >= width:
+            res.append((b >> bits_wnd_r) & mask)
+            total -= width
+            bits_wnd_r += width
+            logger.debug("  read bitpackage: added: %s", res[-1])
+        elif current_byte + 1 < len(raw_bytes):
             current_byte += 1
-        logger.debug("  read bitpackage: added: %s", res[-1])
+            b |= (raw_bytes[current_byte] << bits_wnd_l)
+            bits_wnd_l += 8
     return res