inferencekit-cli 0.1.0__py3-none-any.whl

inferencekit/quality/metrics.py

@@ -0,0 +1,235 @@
+"""Quality statistics: KL and TV on truncated top-k lists, position buckets, CIs, Wilson, McNemar.
+
+KL(reference || engine) from top-k lists
+----------------------------------------
+Servers return only the k most likely tokens per position (k <= 20 for vLLM's default
+``--max-logprobs``), so neither distribution is known in full. ``kl_topk`` builds two
+comparable distributions over the same finite support:
+
+1. Support U = the union of both top-k lists.
+2. A token in U that one side did not list gets, on that side, an imputed probability
+   ``min(p_min, tail / (m + 1))``: ``tail`` is that side's unlisted mass
+   (1 - sum of its listed probabilities), ``m`` is how many tokens of U it is missing,
+   and ``p_min`` is its smallest listed probability (an unlisted token cannot be more
+   likely than the k-th listed one). The "+1" leaves a share of the tail for the rest of
+   the vocabulary.
+3. Everything outside U is merged into one OTHER bucket on each side: 1 - sum over U.
+4. Probabilities are floored at 1e-10 and each side is renormalised, then
+   KL = sum_u p(u) log(p(u) / q(u)) over U and OTHER.
+
+Properties: identical inputs give exactly 0; when the lists cover the whole vocabulary
+the result is the exact KL; merging the tail into OTHER can only lower KL (data
+processing inequality), while step 2 is an approximation in either direction. Treat
+the value as a well-defined, comparable divergence between engines scored the same
+way, not as the exact full-vocabulary KL.
+
+Both lists must have the same length. A top-20 reference list against a top-5 engine
+list puts 15 tokens on the engine side at imputed tail values, which reads as divergence
+even for identical distributions (0.03 nats on a peaked 200-token toy model, more on
+real vocabularies). ``quality score`` therefore cuts both lists to the engine's k (plus
+the token actually generated) before calling ``kl_topk``, and pairs the noise floor at
+that same k. A smaller k makes the value coarser (less sensitive), never biased upward.
+
+Total variation from top-k lists
+--------------------------------
+``tv_topk`` builds the same two distributions (union support, imputed missing entries,
+one OTHER bucket) and returns TV = 1/2 * sum |p(u) - q(u)| over U and OTHER: the share of
+probability mass the engine moved, between 0 and 1. Identical inputs give exactly 0, full
+lists give the exact TV, and merging the tail into OTHER can only lower it. Unlike KL it
+stays bounded when one side gives a token almost no mass, so a single bad imputation
+cannot dominate the mean. ``tv_counts`` is the same distance between two sample
+histograms (``quality sampling`` tests draws against a known distribution instead:
+``quality/sampling.py``).
+
+Confidence intervals
+--------------------
+Tokens inside one generation are correlated, so every CI resamples whole prompts
+(cluster bootstrap) and pools the per-token values of the drawn prompts. Paired CIs
+(engine minus noise floor) draw the same prompts for both arms.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from fractions import Fraction
+from typing import Hashable, Mapping, Sequence
+
+from inferencekit.stats import Z95, _draw, percentile_ci
+
+FLOOR = 1e-10
+OTHER = "__other__"
+# Position buckets (token index inside the generated text, 0-based).
+BUCKETS = ((0, 15), (16, 63), (64, 255), (256, None))
+
+
+def bucket_label(lo: int, hi: int | None) -> str:
+    return f"{lo}+" if hi is None else f"{lo}-{hi}"
+
+
+BUCKET_LABELS = tuple(bucket_label(lo, hi) for lo, hi in BUCKETS)
+
+
+def bucket_of(position: int) -> str:
+    for lo, hi in BUCKETS:
+        if position >= lo and (hi is None or position <= hi):
+            return bucket_label(lo, hi)
+    raise ValueError(position)
+
+
+# --------------------------------------------------------------------------- KL
+
+def _complete(known: Mapping[Hashable, float], support: set) -> dict:
+    """Probabilities on ``support`` plus OTHER, from one side's listed logprobs."""
+    probs = {k: math.exp(lp) for k, lp in known.items() if lp is not None}
+    listed = sum(probs.values())
+    tail = max(0.0, 1.0 - listed)
+    missing = [k for k in support if k not in probs]
+    if missing:
+        cap = min(probs.values()) if probs else 1.0
+        each = max(min(cap, tail / (len(missing) + 1)), FLOOR)
+        for k in missing:
+            probs[k] = each
+    probs = {k: max(v, FLOOR) for k, v in probs.items() if k in support}
+    probs[OTHER] = max(1.0 - sum(probs.values()), FLOOR)
+    total = sum(probs.values())
+    return {k: v / total for k, v in probs.items()}
+
+
+def kl_topk(reference: Mapping[Hashable, float], engine: Mapping[Hashable, float]) -> float:
+    """KL(reference || engine) in nats from two {token: logprob} top-k lists (see module doc)."""
+    support = {k for k, v in reference.items() if v is not None} | {k for k, v in engine.items() if v is not None}
+    p = _complete(reference, support)
+    q = _complete(engine, support)
+    kl = 0.0
+    for k, pk in p.items():
+        if pk > 0:
+            kl += pk * math.log(pk / q[k])
+    return max(kl, 0.0)
+
+
+def kl_full(p_logprobs: Sequence[float], q_logprobs: Sequence[float]) -> float:
+    """Exact KL(p || q) for two full log-distributions over the same vocabulary."""
+    return sum(math.exp(a) * (a - b) for a, b in zip(p_logprobs, q_logprobs))
+
+
+# --------------------------------------------------------------------------- TV
+
+def tv_topk(reference: Mapping[Hashable, float], engine: Mapping[Hashable, float]) -> float:
+    """Total variation between two {token: logprob} top-k lists, on the support ``kl_topk`` uses."""
+    support = {k for k, v in reference.items() if v is not None} | {k for k, v in engine.items() if v is not None}
+    p = _complete(reference, support)
+    q = _complete(engine, support)
+    return min(1.0, 0.5 * sum(abs(p[k] - q[k]) for k in p))
+
+
+def tv_full(p_logprobs: Sequence[float], q_logprobs: Sequence[float]) -> float:
+    """Exact TV for two full log-distributions over the same vocabulary."""
+    return 0.5 * sum(abs(math.exp(a) - math.exp(b)) for a, b in zip(p_logprobs, q_logprobs))
+
+
+def tv_counts(a: Mapping[Hashable, int], b: Mapping[Hashable, int]) -> float | None:
+    """TV between the empirical distributions of two sample histograms ({outcome: count}); None if one is empty."""
+    na, nb = sum(a.values()), sum(b.values())
+    if na <= 0 or nb <= 0:
+        return None
+    return 0.5 * sum(abs(a.get(k, 0) / na - b.get(k, 0) / nb) for k in set(a) | set(b))
+
+
+# --------------------------------------------------------------------------- CIs over prompts
+
+def pooled(units: Sequence[tuple[float, int]]) -> float | None:
+    """Pooled mean sum(sums) / sum(counts) of per-prompt (sum, count) pairs."""
+    s = sum(u[0] for u in units)
+    n = sum(u[1] for u in units)
+    return s / n if n else None
+
+
+def cluster_ci(units: Sequence[tuple[float, int]], resamples: int = 2000, seed: int = 0) -> tuple:
+    """95% CI of the pooled mean, resampling prompts (each unit = one prompt's (sum, count))."""
+    units = [u for u in units if u[1] > 0]
+    k = len(units)
+    if k < 2 or resamples <= 0:
+        return (None, None)
+    rng = random.Random(seed)
+    reps = []
+    for _ in range(resamples):
+        s = n = 0.0
+        for _ in range(k):
+            a, b = units[_draw(rng, k)]
+            s += a
+            n += b
+        reps.append(s / n if n else float("nan"))
+    return percentile_ci(reps)
+
+
+def paired_delta_ci(engine: Sequence[tuple[float, int]], floor: Sequence[tuple[float, int]],
+                    resamples: int = 2000, seed: int = 0) -> dict:
+    """pooled(engine) - pooled(floor) with a paired bootstrap 95% CI.
+
+    ``engine[i]`` and ``floor[i]`` must be the same prompt; prompts where either arm has no
+    tokens are dropped. Each resample draws prompt indices once and uses them for both arms.
+    """
+    pairs = [(e, f) for e, f in zip(engine, floor) if e[1] > 0 and f[1] > 0]
+    k = len(pairs)
+    out = {"n_prompts": k, "delta": None, "ci95": (None, None)}
+    if k == 0:
+        return out
+    out["delta"] = pooled([p[0] for p in pairs]) - pooled([p[1] for p in pairs])
+    if k < 2 or resamples <= 0:
+        return out
+    rng = random.Random(seed)
+    reps = []
+    for _ in range(resamples):
+        es = en = fs = fn = 0.0
+        for _ in range(k):
+            (a, b), (c, d) = pairs[_draw(rng, k)]
+            es += a
+            en += b
+            fs += c
+            fn += d
+        reps.append(es / en - fs / fn)
+    out["ci95"] = percentile_ci(reps)
+    return out
+
+
+def margin_verdict(worse_lo: float | None, worse_hi: float | None, margin: float) -> tuple[str, str]:
+    """PASS when the whole CI of "how much worse than the floor" is <= margin, FAIL when all of it is above.
+
+    ``worse`` is oriented so that positive = the engine is worse than the noise floor.
+    """
+    if worse_lo is None or worse_hi is None:
+        return "UNPROVEN", "no confidence interval (need >= 2 prompts scored for both the engine and the floor)"
+    if worse_hi <= margin + 1e-12:
+        return "PASS", f"CI [{worse_lo:.4g}, {worse_hi:.4g}] <= margin {margin:g}"
+    if worse_lo > margin + 1e-12:
+        return "FAIL", f"CI [{worse_lo:.4g}, {worse_hi:.4g}] > margin {margin:g}"
+    return "UNPROVEN", f"CI [{worse_lo:.4g}, {worse_hi:.4g}] straddles margin {margin:g}: score more prompts"
+
+
+# --------------------------------------------------------------------------- task accuracy
+
+def wilson_ci(successes: int, n: int, z: float = Z95) -> tuple[float | None, float | None]:
+    """Wilson score interval for a binomial proportion (good coverage even near 0 or 1)."""
+    if n <= 0:
+        return (None, None)
+    p = successes / n
+    denom = 1 + z * z / n
+    centre = (p + z * z / (2 * n)) / denom
+    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / denom
+    return (max(0.0, centre - half), min(1.0, centre + half))
+
+
+def mcnemar_exact(b: int, c: int) -> float:
+    """Two-sided exact McNemar p-value from the discordant counts.
+
+    b = items only arm A got right, c = items only arm B got right. Under H0 (same
+    accuracy) each discordant item is a fair coin, so min(b, c) ~ Binomial(b + c, 1/2).
+    Computed with integers (2**-n underflows a float for n > 1074).
+    """
+    n = b + c
+    if n == 0:
+        return 1.0
+    k = min(b, c)
+    tail = sum(math.comb(n, i) for i in range(k + 1))
+    return float(min(Fraction(1), Fraction(2 * tail, 2 ** n)))

inferencekit/quality/niah.py

@@ -0,0 +1,63 @@
+"""Needle in a haystack: can the engine still retrieve one fact from a long context?
+
+Long-context paths (chunked prefill, paged/sliding attention, quantised KV cache, RoPE
+scaling) can break retrieval while short-prompt quality looks fine. This builds, for a
+target context length and a depth (0 = start, 1 = end), a haystack of synthetic filler
+sentences with one "needle" sentence holding a random 7-digit code, then asks for the
+code. Correct = the code's digits appear in the reply.
+
+Everything is generated from a seed (no third-party text, nothing to license): the
+filler is grammatical nonsense about rooms, boxes and colours, so it never contains a
+competing code. Lengths are hit approximately: characters per token are measured once
+with the server's ``/tokenize`` when it has one (else 4.0 is assumed), and the actual
+prompt length reported by the server (``usage.prompt_tokens``) is recorded per item.
+"""
+
+from __future__ import annotations
+
+import random
+import re
+
+_ADJ = ("quiet", "narrow", "amber", "distant", "hollow", "crowded", "frozen", "silver", "patient", "brittle")
+_NOUN = ("archive", "harbour", "orchard", "workshop", "station", "library", "market", "garden", "tower", "bridge")
+_PLACE = ("the northern valley", "the old quarter", "the river delta", "the high plateau", "the east coast",
+          "the lower town", "the island chain", "the western ridge")
+_VERB = ("stores", "counts", "repairs", "ships", "sorts", "paints", "weighs", "labels")
+_THING = ("wooden crates", "paper maps", "glass jars", "copper wires", "woollen blankets", "clay tiles",
+          "brass keys", "linen sacks")
+_COLOR = ("red", "green", "blue", "violet", "orange", "teal", "golden", "grey")
+_ANIMAL = ("falcon", "otter", "lynx", "heron", "badger", "salmon", "gecko", "raven")
+
+QUESTION = "What is the access code for the {color} {animal} vault? Reply with the number only."
+
+
+def filler_sentence(rng: random.Random) -> str:
+    # Small numbers in words only, so the filler never contains a digit string like the code.
+    return (f"Every morning the {rng.choice(_ADJ)} {rng.choice(_NOUN)} of {rng.choice(_PLACE)} "
+            f"{rng.choice(_VERB)} {rng.choice(('two', 'five', 'nine', 'twelve', 'forty'))} "
+            f"{rng.choice(_THING)}.")
+
+
+def build(target_tokens: int, depth: float, seed: int, chars_per_token: float = 4.0) -> dict:
+    """{messages, code, depth, target_tokens, needle} for one haystack item."""
+    rng = random.Random(f"niah:{seed}:{target_tokens}:{depth}")
+    color, animal = rng.choice(_COLOR), rng.choice(_ANIMAL)
+    code = str(rng.randrange(1_000_000, 10_000_000))
+    needle = f"The access code for the {color} {animal} vault is {code}."
+    question = QUESTION.format(color=color, animal=animal)
+    budget = max(0, int(target_tokens * chars_per_token) - len(needle) - len(question) - 200)
+    sentences, size = [], 0
+    while size < budget:
+        s = filler_sentence(rng)
+        sentences.append(s)
+        size += len(s) + 1
+    pos = round(min(max(depth, 0.0), 1.0) * len(sentences))
+    sentences.insert(pos, needle)
+    context = " ".join(sentences)
+    content = (f"Read the following document.\n\n<document>\n{context}\n</document>\n\n{question}")
+    return {"messages": [{"role": "user", "content": content}], "code": code, "depth": depth,
+            "target_tokens": target_tokens, "needle_sentence_index": pos, "sentences": len(sentences)}
+
+
+def found(reply: str | None, code: str) -> bool:
+    return bool(reply) and code in re.sub(r"[\s,._'-]", "", reply)

inferencekit/quality/parquet.py

@@ -0,0 +1,296 @@
+"""A small, dependency-free Parquet reader: just enough to read a flat table of string columns.
+
+It exists so ``quality gsm8k`` can read the Hugging Face copy of GSM8K (a Parquet file)
+without pulling in pyarrow. Supported: flat schemas (no nested or repeated columns),
+BYTE_ARRAY / INT32 / INT64 columns, REQUIRED or OPTIONAL, PLAIN and dictionary encodings
+(PLAIN_DICTIONARY / RLE_DICTIONARY), data pages v1 and v2, and the UNCOMPRESSED, SNAPPY
+and GZIP codecs. Anything else raises ``ParquetError`` instead of returning wrong data.
+
+Format reference: github.com/apache/parquet-format (parquet.thrift, Encodings.md); the
+footer and page headers use the Thrift compact protocol; Snappy is the raw block format.
+"""
+
+from __future__ import annotations
+
+import struct
+import zlib
+from pathlib import Path
+
+
+class ParquetError(ValueError):
+    pass
+
+
+# ----------------------------------------------------------------------------- thrift compact protocol
+
+class _Thrift:
+    """Reads Thrift compact-protocol structs as {field_id: value} dicts."""
+
+    def __init__(self, buf: bytes, pos: int = 0):
+        self.buf, self.pos = buf, pos
+
+    def byte(self) -> int:
+        if self.pos >= len(self.buf):
+            raise ParquetError("truncated thrift data")
+        b = self.buf[self.pos]
+        self.pos += 1
+        return b
+
+    def varint(self) -> int:
+        shift = result = 0
+        while True:
+            b = self.byte()
+            result |= (b & 0x7F) << shift
+            if not b & 0x80:
+                return result
+            shift += 7
+
+    def zigzag(self) -> int:
+        n = self.varint()
+        return (n >> 1) ^ -(n & 1)
+
+    def binary(self) -> bytes:
+        n = self.varint()
+        out = self.buf[self.pos:self.pos + n]
+        if len(out) != n:
+            raise ParquetError("truncated thrift binary")
+        self.pos += n
+        return out
+
+    def value(self, t: int):
+        if t in (1, 2):  # bool inside a container: one byte, 1 = true
+            return self.byte() == 1
+        if t == 3:
+            return struct.unpack("b", bytes([self.byte()]))[0]
+        if t in (4, 5, 6):
+            return self.zigzag()
+        if t == 7:
+            out = struct.unpack_from("<d", self.buf, self.pos)[0]
+            self.pos += 8
+            return out
+        if t == 8:
+            return self.binary()
+        if t in (9, 10):
+            head = self.byte()
+            size, et = head >> 4, head & 0x0F
+            if size == 15:
+                size = self.varint()
+            return [self.value(et) for _ in range(size)]
+        if t == 11:
+            size = self.varint()
+            if not size:
+                return {}
+            kv = self.byte()
+            return {self.value(kv >> 4): self.value(kv & 0x0F) for _ in range(size)}
+        if t == 12:
+            return self.struct()
+        raise ParquetError(f"unknown thrift type {t}")
+
+    def struct(self) -> dict:
+        out, last = {}, 0
+        while True:
+            head = self.byte()
+            if head == 0:
+                return out
+            delta, t = head >> 4, head & 0x0F
+            fid = last + delta if delta else self.zigzag()
+            last = fid
+            out[fid] = True if t == 1 else False if t == 2 else self.value(t)
+
+
+# ----------------------------------------------------------------------------- codecs and encodings
+
+def snappy_decompress(data: bytes) -> bytes:
+    """Raw (unframed) Snappy block decompression."""
+    t = _Thrift(data)
+    size = t.varint()
+    out = bytearray()
+    pos, n = t.pos, len(data)
+    while pos < n:
+        tag = data[pos]
+        pos += 1
+        kind = tag & 3
+        if kind == 0:  # literal
+            ln = tag >> 2
+            if ln >= 60:
+                nb = ln - 59
+                ln = int.from_bytes(data[pos:pos + nb], "little")
+                pos += nb
+            ln += 1
+            out += data[pos:pos + ln]
+            pos += ln
+            continue
+        if kind == 1:
+            ln = ((tag >> 2) & 7) + 4
+            off = ((tag >> 5) << 8) | data[pos]
+            pos += 1
+        elif kind == 2:
+            ln = (tag >> 2) + 1
+            off = int.from_bytes(data[pos:pos + 2], "little")
+            pos += 2
+        else:
+            ln = (tag >> 2) + 1
+            off = int.from_bytes(data[pos:pos + 4], "little")
+            pos += 4
+        if off <= 0 or off > len(out):
+            raise ParquetError("corrupt snappy data (bad copy offset)")
+        start = len(out) - off
+        if off >= ln:
+            out += out[start:start + ln]
+        else:  # overlapping copy: repeat the last `off` bytes
+            for i in range(ln):
+                out.append(out[start + i])
+    if len(out) != size:
+        raise ParquetError(f"corrupt snappy data ({len(out)} bytes, header says {size})")
+    return bytes(out)
+
+
+def _decompress(codec: int, data: bytes, size: int) -> bytes:
+    if codec == 0:
+        return data
+    if codec == 1:
+        return snappy_decompress(data)
+    if codec == 2:
+        return zlib.decompress(data, 47)  # gzip or zlib header
+    raise ParquetError(f"unsupported compression codec {codec} (supported: none, snappy, gzip)")
+
+
+def _hybrid(buf: bytes, pos: int, end: int, width: int, count: int) -> list[int]:
+    """RLE / bit-packed hybrid decoding of ``count`` values of ``width`` bits from buf[pos:end]."""
+    out: list[int] = []
+    if width == 0:
+        return [0] * count
+    nbytes = (width + 7) // 8
+    mask = (1 << width) - 1
+    t = _Thrift(buf, pos)
+    while len(out) < count and t.pos < end:
+        head = t.varint()
+        if head & 1:  # bit-packed: (head >> 1) groups of 8 values
+            ngroups = head >> 1
+            raw = int.from_bytes(buf[t.pos:t.pos + ngroups * width], "little")
+            t.pos += ngroups * width
+            out.extend((raw >> (i * width)) & mask for i in range(ngroups * 8))
+        else:  # run of one value
+            value = int.from_bytes(buf[t.pos:t.pos + nbytes], "little")
+            t.pos += nbytes
+            out.extend([value] * (head >> 1))
+    if len(out) < count:
+        raise ParquetError("RLE data ended early")
+    return out[:count]
+
+
+def _plain(ptype: int, buf: bytes, pos: int, count: int) -> list:
+    if ptype == 6:  # BYTE_ARRAY
+        out = []
+        for _ in range(count):
+            (ln,) = struct.unpack_from("<I", buf, pos)
+            out.append(bytes(buf[pos + 4:pos + 4 + ln]))
+            pos += 4 + ln
+        return out
+    if ptype in (1, 2):  # INT32, INT64
+        fmt = "<i" if ptype == 1 else "<q"
+        return list(struct.unpack_from(f"<{count}{fmt[1]}", buf, pos))
+    raise ParquetError(f"unsupported physical type {ptype}")
+
+
+# ----------------------------------------------------------------------------- reading
+
+def _column(data: bytes, meta: dict, ptype: int, max_def: int) -> list:
+    codec, total = meta.get(4, 0), meta.get(5, 0)
+    first_data, first_dict = meta.get(9), meta.get(11)
+    start = first_dict if first_dict and (first_data is None or first_dict < first_data) else first_data
+    dictionary = None
+    values: list = []
+    pos = start
+    while len(values) < total:
+        t = _Thrift(data, pos)
+        header = t.struct()
+        pos = t.pos
+        ptype_page, usize, csize = header.get(1), header.get(2, 0), header.get(3, 0)
+        body = data[pos:pos + csize]
+        pos += csize
+        if ptype_page == 2:  # dictionary page
+            dh = header.get(7) or {}
+            dictionary = _plain(ptype, _decompress(codec, body, usize), 0, dh.get(1, 0))
+            continue
+        if ptype_page == 0:  # data page v1: [def levels][values], all compressed together
+            dh = header.get(5) or {}
+            n, enc = dh.get(1, 0), dh.get(2, 0)
+            page = _decompress(codec, body, usize)
+            p = 0
+            if max_def:
+                (ln,) = struct.unpack_from("<I", page, 0)
+                defs = _hybrid(page, 4, 4 + ln, max_def.bit_length(), n)
+                p = 4 + ln
+            else:
+                defs = [0] * n
+        elif ptype_page == 3:  # data page v2: [rep levels][def levels] uncompressed, then values
+            dh = header.get(8) or {}
+            n, enc = dh.get(1, 0), dh.get(4, 0)
+            dlen, rlen = dh.get(5, 0), dh.get(6, 0)
+            if rlen:
+                raise ParquetError("repeated columns are not supported")
+            defs = _hybrid(body, 0, dlen, max_def.bit_length(), n) if max_def else [0] * n
+            rest = body[dlen:]
+            page = _decompress(codec, rest, usize - dlen) if dh.get(7, True) else rest
+            p = 0
+        else:
+            continue  # index pages and anything else carry no values
+        present = sum(1 for d in defs if d == max_def)
+        if enc == 0:
+            vals = _plain(ptype, page, p, present)
+        elif enc in (2, 8):
+            if dictionary is None:
+                raise ParquetError("dictionary-encoded page without a dictionary")
+            width = page[p]
+            idx = _hybrid(page, p + 1, len(page), width, present)
+            vals = [dictionary[i] for i in idx]
+        else:
+            raise ParquetError(f"unsupported value encoding {enc}")
+        it = iter(vals)
+        values.extend(next(it) if d == max_def else None for d in defs)
+    return values
+
+
+def read_table(source: str | Path | bytes) -> dict[str, list]:
+    """{column name: values} of a flat Parquet file (BYTE_ARRAY values as bytes; None for nulls).
+
+    A malformed file raises ``ParquetError`` (a ``ValueError``), never a low-level decoding error."""
+    data = source if isinstance(source, (bytes, bytearray)) else Path(source).read_bytes()
+    try:
+        return _read_table(data)
+    except ParquetError:
+        raise
+    except (struct.error, zlib.error, IndexError, KeyError, TypeError, AttributeError, ValueError,
+            OverflowError, RecursionError) as e:  # offsets or lengths that point outside the file, bad codecs
+        raise ParquetError(f"malformed Parquet file ({type(e).__name__}: {e})") from None
+
+
+def _read_table(data: bytes) -> dict[str, list]:
+    if len(data) < 12 or data[:4] != b"PAR1" or data[-4:] != b"PAR1":
+        raise ParquetError("not a Parquet file (missing PAR1 magic)")
+    (meta_len,) = struct.unpack_from("<I", data, len(data) - 8)
+    meta = _Thrift(data, len(data) - 8 - meta_len).struct()
+    schema = meta.get(2) or []
+    if not schema:
+        raise ParquetError("Parquet file without a schema")
+    leaves = schema[1:]
+    if any(el.get(5) for el in leaves):
+        raise ParquetError("nested columns are not supported")
+    info = {}
+    for el in leaves:
+        name = el[4].decode("utf-8")
+        rep = el.get(3, 0)
+        if rep == 2:
+            raise ParquetError(f"column {name!r} is repeated; not supported")
+        info[name] = (el.get(1), 1 if rep == 1 else 0)
+    out: dict[str, list] = {name: [] for name in info}
+    for rg in meta.get(4) or []:
+        for chunk in rg.get(1) or []:
+            cm = chunk.get(3) or {}
+            name = b".".join(cm.get(3) or []).decode("utf-8")
+            if name not in info:
+                raise ParquetError(f"column {name!r} is not in the schema")
+            ptype, max_def = info[name]
+            out[name].extend(_column(data, cm, ptype, max_def))
+    return out