310 라인
9.6 KiB
Python
 | #!/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]))
|