| #!/usr/bin/env python3
|
| """Minimal pure-Python Parquet reader (flat schemas, PLAIN/dictionary pages,
|
| SNAPPY/GZIP/uncompressed). Enough to extract tick columns from public datasets.
|
| """
|
| import struct
|
| import sys
|
| import zlib
|
| from array import array
|
|
|
| # ---------------------------------------------------------------- thrift compact
|
| class TReader:
|
| def __init__(self, buf, pos=0):
|
| self.b = buf
|
| self.p = pos
|
|
|
| def byte(self):
|
| v = self.b[self.p]
|
| self.p += 1
|
| return v
|
|
|
| def varint(self):
|
| shift = 0
|
| result = 0
|
| while True:
|
| c = self.b[self.p]
|
| self.p += 1
|
| result |= (c & 0x7F) << shift
|
| if not (c & 0x80):
|
| return result
|
| shift += 7
|
|
|
| def zigzag(self):
|
| n = self.varint()
|
| return (n >> 1) ^ -(n & 1)
|
|
|
| def read_value(self, t):
|
| if t == 1:
|
| return True
|
| if t == 2:
|
| return False
|
| if t == 3:
|
| v = self.byte()
|
| return v - 256 if v > 127 else v
|
| if t in (4, 5, 6):
|
| return self.zigzag()
|
| if t == 7:
|
| v = struct.unpack_from('<d', self.b, self.p)[0]
|
| self.p += 8
|
| return v
|
| if t == 8:
|
| n = self.varint()
|
| v = bytes(self.b[self.p:self.p + n])
|
| self.p += n
|
| return v
|
| if t in (9, 10):
|
| h = self.byte()
|
| size = h >> 4
|
| et = h & 15
|
| if size == 15:
|
| size = self.varint()
|
| out = []
|
| for _ in range(size):
|
| if et in (1, 2):
|
| out.append(self.byte() == 1)
|
| else:
|
| out.append(self.read_value(et))
|
| return out
|
| if t == 11:
|
| size = self.varint()
|
| out = {}
|
| if size:
|
| kv = self.byte()
|
| for _ in range(size):
|
| k = self.read_value(kv >> 4)
|
| out[k] = self.read_value(kv & 15)
|
| return out
|
| if t == 12:
|
| return self.read_struct()
|
| raise ValueError('bad thrift type %d' % t)
|
|
|
| def read_struct(self):
|
| out = {}
|
| fid = 0
|
| while True:
|
| h = self.byte()
|
| if h == 0:
|
| return out
|
| t = h & 15
|
| d = h >> 4
|
| if d == 0:
|
| fid = self.zigzag()
|
| else:
|
| fid += d
|
| out[fid] = self.read_value(t)
|
|
|
|
|
| # ---------------------------------------------------------------- snappy
|
| def snappy_decompress(src):
|
| r = TReader(src)
|
| n = r.varint()
|
| out = bytearray()
|
| p = r.p
|
| L = len(src)
|
| while p < L:
|
| tag = src[p]
|
| p += 1
|
| kind = tag & 3
|
| if kind == 0:
|
| ln = tag >> 2
|
| if ln >= 60:
|
| nb = ln - 59
|
| ln = int.from_bytes(src[p:p + nb], 'little')
|
| p += nb
|
| ln += 1
|
| out += src[p:p + ln]
|
| p += ln
|
| continue
|
| if kind == 1:
|
| ln = ((tag >> 2) & 7) + 4
|
| off = ((tag >> 5) << 8) | src[p]
|
| p += 1
|
| elif kind == 2:
|
| ln = (tag >> 2) + 1
|
| off = src[p] | (src[p + 1] << 8)
|
| p += 2
|
| else:
|
| ln = (tag >> 2) + 1
|
| off = int.from_bytes(src[p:p + 4], 'little')
|
| p += 4
|
| start = len(out) - off
|
| if off >= ln:
|
| out += out[start:start + ln]
|
| else:
|
| while ln > 0:
|
| chunk = out[start:start + min(off, ln)]
|
| out += chunk
|
| ln -= len(chunk)
|
| if len(out) != n:
|
| raise ValueError('snappy length mismatch')
|
| return bytes(out)
|
|
|
|
|
| _zstd = None
|
|
|
|
|
| def zstd_decompress(data, usize):
|
| """ZSTD pages are decoded with the system libzstd (or $LIBZSTD) via ctypes."""
|
| global _zstd
|
| import ctypes
|
| import ctypes.util
|
| import os
|
| if _zstd is None:
|
| lib = os.environ.get('LIBZSTD') or ctypes.util.find_library('zstd')
|
| if not lib:
|
| raise RuntimeError('libzstd not found: install it or set LIBZSTD=/path/to/libzstd.so')
|
| _zstd = ctypes.CDLL(lib)
|
| _zstd.ZSTD_decompress.restype = ctypes.c_size_t
|
| _zstd.ZSTD_decompress.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t]
|
| _zstd.ZSTD_isError.restype = ctypes.c_uint
|
| _zstd.ZSTD_isError.argtypes = [ctypes.c_size_t]
|
| dst = ctypes.create_string_buffer(usize)
|
| src = bytes(data)
|
| r = _zstd.ZSTD_decompress(dst, usize, src, len(src))
|
| if _zstd.ZSTD_isError(r) or r != usize:
|
| raise ValueError('zstd failure')
|
| return dst.raw
|
|
|
|
|
| def decompress(codec, data, usize):
|
| if codec == 6:
|
| return zstd_decompress(data, usize)
|
| if codec == 0:
|
| return data
|
| if codec == 1:
|
| return snappy_decompress(data)
|
| if codec == 2:
|
| return zlib.decompress(data, 16 + zlib.MAX_WBITS)
|
| raise ValueError('unsupported codec %d' % codec)
|
|
|
|
|
| # ---------------------------------------------------------------- rle hybrid
|
| def rle_hybrid(buf, pos, end, width, count):
|
| out = []
|
| r = TReader(buf, pos)
|
| mask = (1 << width) - 1
|
| vbytes = (width + 7) // 8
|
| while len(out) < count and r.p < end:
|
| h = r.varint()
|
| if h & 1:
|
| groups = h >> 1
|
| nb = groups * width
|
| chunk = int.from_bytes(buf[r.p:r.p + nb], 'little')
|
| r.p += nb
|
| if width == 0:
|
| out.extend([0] * (groups * 8))
|
| else:
|
| for i in range(groups * 8):
|
| out.append((chunk >> (i * width)) & mask)
|
| else:
|
| run = h >> 1
|
| v = int.from_bytes(buf[r.p:r.p + vbytes], 'little') if vbytes else 0
|
| r.p += vbytes
|
| out.extend([v] * run)
|
| return out[:count]
|
|
|
|
|
| PTYPE_FMT = {1: ('i', 4), 2: ('q', 8), 4: ('f', 4), 5: ('d', 8)}
|
|
|
|
|
| def plain_values(ptype, buf, pos, count):
|
| code, size = PTYPE_FMT[ptype]
|
| a = array(code)
|
| a.frombytes(bytes(buf[pos:pos + count * size]))
|
| if sys.byteorder != 'little':
|
| a.byteswap()
|
| return a
|
|
|
|
|
| def read_parquet(path, wanted=None):
|
| with open(path, 'rb') as f:
|
| data = f.read()
|
| if data[:4] != b'PAR1' or data[-4:] != b'PAR1':
|
| raise ValueError('not parquet')
|
| mlen = struct.unpack_from('<I', data, len(data) - 8)[0]
|
| meta = TReader(data, len(data) - 8 - mlen).read_struct()
|
| schema = meta[2]
|
| leaves = []
|
| for el in schema[1:]:
|
| if el.get(5):
|
| raise ValueError('nested schema not supported')
|
| leaves.append((el[4].decode(), el.get(1), el.get(3, 0), el.get(6), el.get(10)))
|
| cols = {name: [] for name, *_ in leaves}
|
| for rg in meta[4]:
|
| for ci, cc in enumerate(rg[1]):
|
| name, ptype, rep, conv, logical = leaves[ci]
|
| if wanted and name not in wanted:
|
| continue
|
| md = cc[3]
|
| codec = md[4]
|
| nvals = md[5]
|
| start = md.get(11) or md[9]
|
| if md.get(11) is not None and md[11] < md[9]:
|
| start = md[11]
|
| pos = start
|
| dictionary = None
|
| got = 0
|
| while got < nvals:
|
| r = TReader(data, pos)
|
| ph = r.read_struct()
|
| pos = r.p
|
| ptype_page = ph[1]
|
| usize = ph[2]
|
| csize = ph[3]
|
| raw = data[pos:pos + csize]
|
| pos += csize
|
| if ptype_page == 2: # dictionary
|
| page = decompress(codec, raw, usize)
|
| dictionary = plain_values(ptype, page, 0, ph[7][1])
|
| continue
|
| if ptype_page == 0:
|
| page = decompress(codec, raw, usize)
|
| dh = ph[5]
|
| n = dh[1]
|
| enc = dh[2]
|
| p = 0
|
| ndef = n
|
| if rep == 1: # optional -> def levels
|
| ln = struct.unpack_from('<I', page, 0)[0]
|
| lv = rle_hybrid(page, 4, 4 + ln, 1, n)
|
| p = 4 + ln
|
| ndef = sum(lv)
|
| if ndef != n:
|
| raise ValueError('nulls not supported')
|
| elif ptype_page == 3:
|
| dh = ph[8]
|
| n = dh[1]
|
| if dh.get(2, 0):
|
| raise ValueError('nulls not supported')
|
| enc = dh[4]
|
| dl = dh.get(5, 0)
|
| rl = dh.get(6, 0)
|
| lev = raw[:dl + rl]
|
| body = raw[dl + rl:]
|
| if dh.get(7, True):
|
| body = decompress(codec, body, usize - dl - rl)
|
| page = body
|
| p = 0
|
| ndef = n
|
| else:
|
| continue
|
| if enc == 0:
|
| vals = plain_values(ptype, page, p, ndef)
|
| elif enc in (2, 8):
|
| width = page[p]
|
| idx = rle_hybrid(page, p + 1, len(page), width, ndef)
|
| d = dictionary
|
| vals = [d[i] for i in idx]
|
| else:
|
| raise ValueError('unsupported encoding %d' % enc)
|
| cols[name].extend(vals)
|
| got += n
|
| return cols, leaves, meta[3]
|
|
|
|
|
| if __name__ == '__main__':
|
| cols, leaves, nrows = read_parquet(sys.argv[1])
|
| print('rows', nrows)
|
| for name, ptype, rep, conv, logical in leaves:
|
| c = cols[name]
|
| print(name, ptype, rep, conv, logical, len(c), list(c[:5]))
|
541e992bab