polypress 0.2.0__py3-none-any.whl

polypress/stream.py

@@ -0,0 +1,438 @@
+"""Bounded-memory compression for files larger than RAM.
+
+The single-shot codec in fast.py holds the whole table as Python strings.
+Measured expansion is about 8x -- a 291 MB CSV peaks at 2.4 GB -- so a 50 GB
+file would need hundreds of gigabytes and simply will not run.
+
+This splits the table into row blocks, compresses each independently, and
+concatenates them behind an index. Peak memory is one block, not one file,
+and it is settable. Blocks are independent, so this is also the thing that
+makes parallelism possible later.
+
+The cost is real and worth stating: the cross-column reordering only sees
+correlations *inside* a block, so smaller blocks compress slightly worse.
+Measure with `--rows` before choosing a small one.
+
+    python3 tzip.py stream-compress big.csv -o big.ppz --budget 1.0
+    python3 tzip.py stream-restore  big.ppz -o back.csv
+    python3 tzip.py stream-info     big.ppz
+
+Or as a module, which is the same code path:
+
+    python3 -m polypress.stream compress big.csv --budget 1.0
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import lzma
+import os
+import sys
+import time
+from typing import Iterator, List, Optional
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(
+    os.path.abspath(__file__))))
+
+from . import dtz
+from . import fast
+
+MAGIC = b"PPZS"
+DEFAULT_BUDGET_GB = 1.0
+# Python strings cost roughly this much more than the CSV bytes they came
+# from; measured at 8.1x on NHAMCS (291 MB file -> 2.36 GB resident).
+EXPANSION = 8.5
+# Verification decodes the block and compares, so the original block and the
+# decoded copy are both resident at once. Without this the budget lied by
+# roughly 2x: asking for 1 GB gave a 2.07 GB peak.
+VERIFY_FACTOR = 2.2
+MIN_ROWS = 1000
+MAX_ROWS = 2_000_000
+
+csv.field_size_limit(1 << 31)
+
+
+def _decode_guard(rows, path: str, enc: str):
+    """Turn a lazy decode failure into the same refusal dtz raises.
+
+    The single-shot reader decodes the whole file inside one call, so a bad
+    byte surfaces there. Here the file is read block by block, so it surfaces
+    hundreds of rows in -- possibly after several blocks have already been
+    written. Same error, same message, wherever it lands."""
+    try:
+        for row in rows:
+            yield row
+    except UnicodeDecodeError as exc:
+        raise dtz.EncodingRefused(
+            dtz.encoding_error(path, exc, enc)) from None
+
+
+def _reader(path: str, encoding: Optional[str] = None):
+    ext = os.path.splitext(path)[1].lower()
+    enc = encoding or dtz.sniff_encoding(path)
+    fh = dtz.open_text(path, enc)
+    try:
+        delimiter = dtz.EXT_DELIMITER.get(ext)
+        if delimiter is None:
+            delimiter = dtz._sniff_delimiter(fh.read(65536))
+            fh.seek(0)
+    except UnicodeDecodeError as exc:
+        fh.close()
+        raise dtz.EncodingRefused(
+            dtz.encoding_error(path, exc, enc)) from None
+    return fh, _decode_guard(csv.reader(fh, delimiter=delimiter), path, enc)
+
+
+def plan_rows(path: str, budget_gb: float, verify: bool = True,
+              encoding: Optional[str] = None) -> tuple:
+    """Rows per block that keep peak memory near the budget.
+
+    Estimated from the real average row width of the file rather than a
+    guess, because a 209-column survey row and a 5-column sensor row differ
+    by two orders of magnitude."""
+    size = os.path.getsize(path)
+    fh, rd = _reader(path, encoding)
+    try:
+        header = next(rd, [])
+        n, seen = 0, 0
+        for row in rd:
+            seen += sum(len(c) for c in row) + len(row)
+            n += 1
+            if n >= 2000:
+                break
+    finally:
+        fh.close()
+    per_row = max(1.0, seen / max(n, 1))
+    factor = EXPANSION * (VERIFY_FACTOR if verify else 1.0)
+    rows = int((budget_gb * (1 << 30)) / (per_row * factor))
+    rows = max(MIN_ROWS, min(MAX_ROWS, rows))
+    est_blocks = max(1, int(size / (per_row * rows)) + 1)
+    return rows, per_row, est_blocks, header
+
+
+def _blocks(path: str, rows_per_block: int,
+            encoding: Optional[str] = None) -> Iterator[dtz.Table]:
+    fh, rd = _reader(path, encoding)
+    try:
+        header = next(rd, None)
+        if header is None:
+            return
+        width = len(header)
+        batch: List[List[str]] = []
+        yielded = False
+        for row in rd:
+            if len(row) != width:            # ragged rows, same rule as dtz
+                row = (row + [""] * width)[:width]
+            batch.append(row)
+            if len(batch) >= rows_per_block:
+                yield dtz.Table(list(header), batch)
+                yielded = True
+                batch = []
+        # Only emit a trailing block if it holds rows. The exception is a file
+        # with a header and no rows at all: that still needs one empty block,
+        # or the archive would record no columns. Yielding unconditionally --
+        # as an earlier `if batch or True` did -- appended a wasted ~178-byte
+        # encoding of nothing whenever the row count divided evenly.
+        if batch or not yielded:
+            yield dtz.Table(list(header), batch)
+    finally:
+        fh.close()
+
+
+def compress(src: str, dst: str, budget_gb: float = DEFAULT_BUDGET_GB,
+             rows: Optional[int] = None, verify: bool = True,
+             progress=None, encoding: Optional[str] = None) -> dict:
+    if rows is None:
+        rows, per_row, est, _ = plan_rows(src, budget_gb, verify, encoding)
+    else:
+        per_row, est = 0.0, 0
+    sizes: List[int] = []
+    total_rows = 0
+    columns: List[str] = []
+
+    with open(dst, "wb") as out:
+        out.write(b"\0" * 8)                 # header length patched at the end
+        for i, table in enumerate(_blocks(src, rows, encoding)):
+            if not columns:
+                columns = table.columns
+            blob = fast.encode(table)
+            if verify:
+                back = fast.decode(blob)
+                if back.columns != table.columns or back.rows != table.rows:
+                    raise RuntimeError(
+                        "block {} failed verification; nothing usable written"
+                        .format(i))
+            out.write(blob)
+            sizes.append(len(blob))
+            total_rows += len(table.rows)
+            if progress:
+                progress(i + 1, total_rows, sum(sizes))
+
+        header = {"columns": columns, "nrows": total_rows,
+                  "rows_per_block": rows, "blocks": sizes}
+        hb = lzma.compress(json.dumps(header, separators=(",", ":")).encode(),
+                           **fast.XZ)
+        out.write(hb)
+        out.write(len(hb).to_bytes(8, "big"))
+
+    # Rewrite the magic now that the file is complete: a truncated run leaves
+    # eight zero bytes at the front rather than something that looks valid.
+    with open(dst, "r+b") as out:
+        out.seek(0)
+        out.write(MAGIC + b"\0\0\0\0")
+    return {"rows": total_rows, "blocks": len(sizes),
+            "rows_per_block": rows, "bytes": os.path.getsize(dst),
+            "per_row": per_row}
+
+
+def _read_header(fh):
+    fh.seek(0)
+    if fh.read(4) != MAGIC:
+        raise ValueError("not a Polypress stream archive")
+    fh.seek(-8, os.SEEK_END)
+    hlen = int.from_bytes(fh.read(8), "big")
+    fh.seek(-8 - hlen, os.SEEK_END)
+    return json.loads(lzma.decompress(fh.read(hlen), **fast.XZ))
+
+
+def iter_blocks(src: str) -> Iterator[dtz.Table]:
+    """Yield each block's table in order, holding one block at a time."""
+    with open(src, "rb") as fh:
+        header = _read_header(fh)
+        fh.seek(8)
+        for size in header["blocks"]:
+            yield fast.decode(fh.read(size))
+
+
+# Restoring has to honour the output extension the same way tzip.py does --
+# `restore -o out.parquet` must produce parquet, not a CSV wearing the name.
+# But it cannot call dtz.write_any, which needs the whole table in memory; the
+# entire point here is that the table does not fit. So each format gets a
+# writer that consumes one block at a time and holds nothing else.
+
+class _DelimitedOut:
+    def __init__(self, path, delim):
+        self.fh = open(path, "w", newline="", encoding="utf-8")
+        self.w = csv.writer(self.fh, delimiter=delim, lineterminator="\n")
+        self.first = True
+
+    def block(self, table):
+        if self.first:
+            self.w.writerow(table.columns)
+            self.first = False
+        self.w.writerows(table.rows)
+
+    def close(self):
+        self.fh.close()
+
+
+class _JsonlOut:
+    def __init__(self, path, columns):
+        dtz._require_unique_columns(dtz.Table(list(columns), []), "jsonl")
+        self.columns = columns
+        self.path = path
+        self.fh = open(path, "w", encoding="utf-8")
+        self.rows = 0
+
+    def block(self, table):
+        dump = json.dumps
+        for row in table.rows:
+            self.fh.write(dump(dict(zip(self.columns, row)),
+                               ensure_ascii=False, separators=(",", ":")))
+            self.fh.write("\n")
+            self.rows += 1
+
+    def close(self):
+        self.fh.close()
+        if not self.rows and self.columns:
+            # dtz's contract is that a FormatLimit leaves nothing behind. It
+            # can check up front; we only learn the table was empty after the
+            # last block, so remove the file we opened.
+            os.unlink(self.path)
+            raise dtz.FormatLimit(
+                "jsonl cannot carry column names for a zero-row table; "
+                "write .csv, .json or .parquet instead")
+
+
+class _JsonOut:
+    """A JSON array emitted incrementally, so the records never coexist."""
+
+    def __init__(self, path, columns):
+        dtz._require_unique_columns(dtz.Table(list(columns), []), "json")
+        self.columns = columns
+        self.path = path
+        self.fh = open(path, "w", encoding="utf-8")
+        self.rows = 0
+
+    def block(self, table):
+        dump = json.dumps
+        for row in table.rows:
+            self.fh.write("[\n " if not self.rows else ",\n ")
+            self.fh.write(dump(dict(zip(self.columns, row)),
+                               ensure_ascii=False))
+            self.rows += 1
+
+    def close(self):
+        if self.rows:
+            self.fh.write("\n]")
+            self.fh.close()
+            return
+        # zero rows: a record list cannot carry headers, so use the column
+        # form, exactly as dtz.write_json does
+        self.fh.close()
+        with open(self.path, "w", encoding="utf-8") as fh:
+            json.dump({c: [] for c in self.columns}, fh, ensure_ascii=False)
+
+
+class _ParquetOut:
+    """One row group per block -- the block structure maps straight onto it."""
+
+    def __init__(self, path, columns):
+        if dtz.pq is None:
+            raise RuntimeError("writing parquet needs pyarrow installed")
+        dtz._require_unique_columns(dtz.Table(list(columns), []), "parquet")
+        self.columns = columns
+        # Every cell is a string in a Table, so the schema is fixed up front
+        # and stays identical across blocks -- which is what lets the row
+        # groups append without buffering.
+        self.schema = dtz.pa.schema([(c, dtz.pa.string()) for c in columns])
+        self.writer = dtz.pq.ParquetWriter(path, self.schema,
+                                           compression="zstd")
+
+    def block(self, table):
+        if not table.rows:
+            return
+        arrays = [dtz.pa.array(table.column(i))
+                  for i in range(len(self.columns))]
+        self.writer.write_table(dtz.pa.table(arrays, schema=self.schema))
+
+    def close(self):
+        self.writer.close()
+
+
+def _open_writer(dst: str, columns: List[str]):
+    ext = os.path.splitext(dst)[1].lower()
+    if ext == ".tsv":
+        return _DelimitedOut(dst, "\t")
+    if ext == ".json":
+        return _JsonOut(dst, columns)
+    if ext in (".jsonl", ".ndjson"):
+        return _JsonlOut(dst, columns)
+    if ext in (".parquet", ".pq"):
+        return _ParquetOut(dst, columns)
+    return _DelimitedOut(dst, dtz.EXT_DELIMITER.get(ext, ","))
+
+
+def restore(src: str, dst: str, progress=None) -> dict:
+    with open(src, "rb") as fh:
+        header = _read_header(fh)
+    out = _open_writer(dst, header["columns"])
+    written = 0
+    try:
+        for table in iter_blocks(src):
+            out.block(table)
+            written += len(table.rows)
+            if progress:
+                progress(written)
+    finally:
+        out.close()
+    return {"rows": written, "blocks": len(header["blocks"]),
+            "bytes": os.path.getsize(dst)}
+
+
+def info(src: str) -> dict:
+    with open(src, "rb") as fh:
+        header = _read_header(fh)
+    header["size"] = os.path.getsize(src)
+    return header
+
+
+# ---------------------------------------------------------------------- cli
+
+def human(n: float) -> str:
+    for unit in ("B", "KB", "MB", "GB"):
+        if abs(n) < 1024 or unit == "GB":
+            return "{:,.0f} {}".format(n, unit) if unit == "B" \
+                else "{:,.1f} {}".format(n, unit)
+        n /= 1024.0
+    return str(n)
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(prog="polypress-stream",
+                                 description=__doc__.split("\n")[0])
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    c = sub.add_parser("compress")
+    c.add_argument("path")
+    c.add_argument("-o", "--output")
+    c.add_argument("--budget", type=float, default=DEFAULT_BUDGET_GB,
+                   help="approximate peak memory in GB (default 1.0)")
+    c.add_argument("--rows", type=int, help="rows per block, overrides budget")
+    c.add_argument("--no-verify", action="store_true")
+    c.add_argument("--encoding",
+                   help="text encoding of the input (default: utf-8, or "
+                        "whatever a byte-order mark says)")
+
+    r = sub.add_parser("restore")
+    r.add_argument("path")
+    r.add_argument("-o", "--output")
+
+    i = sub.add_parser("info")
+    i.add_argument("path")
+
+    a = ap.parse_args(argv)
+
+    if a.cmd == "compress":
+        dst = a.output or a.path + ".ppz"
+        raw = os.path.getsize(a.path)
+        t0 = time.time()
+
+        def prog(nblk, nrows, nbytes):
+            sys.stderr.write("\r  block {:>4}   {:>10,} rows   {:>10}"
+                             .format(nblk, nrows, human(nbytes)))
+            sys.stderr.flush()
+
+        st = compress(a.path, dst, a.budget, a.rows, not a.no_verify, prog,
+                      a.encoding)
+        sys.stderr.write("\r" + " " * 60 + "\r")
+        secs = time.time() - t0
+        print("{:,} rows in {} blocks of {:,}   {} -> {}   {:.2f}x   "
+              "{:.1f} MB/s".format(
+                  st["rows"], st["blocks"], st["rows_per_block"], human(raw),
+                  human(st["bytes"]), raw / max(st["bytes"], 1),
+                  raw / 1e6 / max(secs, 1e-9)))
+        print(dst)
+        return 0
+
+    if a.cmd == "restore":
+        dst = a.output or (a.path[:-4] if a.path.lower().endswith(".ppz")
+                           else a.path + ".csv")
+        t0 = time.time()
+        st = restore(a.path, dst)
+        secs = time.time() - t0
+        print("{:,} rows from {} blocks   {}   {:.1f} MB/s".format(
+            st["rows"], st["blocks"], human(st["bytes"]),
+            st["bytes"] / 1e6 / max(secs, 1e-9)))
+        print(dst)
+        return 0
+
+    h = info(a.path)
+    print("file          {}".format(a.path))
+    print("size          {}".format(human(h["size"])))
+    print("rows          {:,}".format(h["nrows"]))
+    print("columns       {}".format(len(h["columns"])))
+    print("blocks        {} of {:,} rows".format(
+        len(h["blocks"]), h["rows_per_block"]))
+    if h["blocks"]:
+        print("block sizes   min {}  median {}  max {}".format(
+            human(min(h["blocks"])),
+            human(sorted(h["blocks"])[len(h["blocks"]) // 2]),
+            human(max(h["blocks"]))))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

polypress/tcz.c

@@ -0,0 +1,207 @@
+/* Hot loops for the table codec, in C.
+ *
+ * Only four things live here, and each earned its place by showing up in a
+ * profile:
+ *
+ *   scan_decimals / parse_fixed / fmt_fixed
+ *       Converting between decimal text and scaled integers was 70% of decode
+ *       and a large share of encode. Both directions are here so the caller
+ *       can verify exactness by parsing, re-formatting, and comparing bytes.
+ *
+ *   pack_ints / unpack_ints
+ *       One-byte varints with an escape.
+ *
+ * Deliberately plain C with no Python headers: it builds with a bare `cc`
+ * and is loaded through ctypes, so there is no extension-module toolchain to
+ * keep working. Python falls back to the numpy path if the library is
+ * missing, so this is an accelerator, never a requirement.
+ */
+
+#include <stdint.h>
+#include <string.h>
+
+#define ESCAPE 255
+
+/* Largest magnitude we accept, matching INT_LIMIT on the Python side. */
+static const int64_t LIMIT = ((int64_t)1) << 62;
+
+/* ------------------------------------------------------------------ scan */
+
+/* Scan a newline-separated buffer of numbers.
+ * Returns the maximum number of decimal places, or -1 if any field is not a
+ * plain decimal number (which sends the column down the text path).
+ * `count` receives the number of fields. */
+int32_t scan_decimals(const char *buf, int64_t len, int64_t *count)
+{
+    int32_t max_dec = 0;
+    int64_t fields = 0;
+    int64_t i = 0;
+
+    while (i <= len) {
+        int64_t start = i;
+        while (i < len && buf[i] != '\n') i++;
+        int64_t end = i;
+        int64_t p = start;
+        int digits = 0, dec = 0;
+
+        if (p < end && buf[p] == '-') p++;
+        while (p < end && buf[p] >= '0' && buf[p] <= '9') { p++; digits++; }
+        if (p < end && buf[p] == '.') {
+            p++;
+            while (p < end && buf[p] >= '0' && buf[p] <= '9') { p++; dec++; }
+            if (dec == 0) return -1;            /* "12." is not accepted */
+        }
+        if (digits == 0 || p != end) return -1; /* empty or trailing junk */
+        if (dec > max_dec) max_dec = dec;
+        fields++;
+
+        if (i >= len) break;
+        i++;                                     /* step over the newline */
+    }
+    if (count) *count = fields;
+    return max_dec > 18 ? -1 : max_dec;
+}
+
+/* ----------------------------------------------------------------- parse */
+
+/* Parse `n` newline-separated decimals into integers scaled by 10^dec.
+ * Returns 0 on success, -1 on malformed input or overflow. */
+int32_t parse_fixed(const char *buf, int64_t len, int32_t dec,
+                    int64_t *out, int64_t n)
+{
+    static const int64_t POW10[19] = {
+        1LL, 10LL, 100LL, 1000LL, 10000LL, 100000LL, 1000000LL,
+        10000000LL, 100000000LL, 1000000000LL, 10000000000LL,
+        100000000000LL, 1000000000000LL, 10000000000000LL,
+        100000000000000LL, 1000000000000000LL, 10000000000000000LL,
+        100000000000000000LL, 1000000000000000000LL};
+
+    int64_t i = 0, k = 0;
+    while (k < n) {
+        int64_t start = i;
+        while (i < len && buf[i] != '\n') i++;
+        int64_t end = i;
+        int64_t p = start;
+        int neg = 0;
+        int64_t v = 0;
+        int seen = 0;
+
+        /* Overflow has to be caught BEFORE it happens. Checking `v >= LIMIT`
+         * after `v = v*10 + d` is too late: signed overflow is undefined, and
+         * in practice it wraps negative, so the guard never fires. That let
+         * "-9223372036854775808" through as a valid value even though the
+         * documented limit is 2^62 -- and the numpy fallback rejected it, so
+         * the accelerator and the reference disagreed about what a column
+         * even was. Rearranged, the test is exact and never overflows. */
+        #define ACC_DIGIT(d)                                                 \
+            do {                                                             \
+                int dgt_ = (d);                                              \
+                if (v > (LIMIT - 1 - dgt_) / 10) return -1;                  \
+                v = v * 10 + dgt_;                                           \
+            } while (0)
+
+        if (p < end && buf[p] == '-') { neg = 1; p++; }
+        while (p < end && buf[p] >= '0' && buf[p] <= '9') {
+            ACC_DIGIT(buf[p] - '0');
+            p++; seen = 1;
+        }
+        if (!seen) return -1;
+
+        int used = 0;
+        if (p < end && buf[p] == '.') {
+            p++;
+            while (p < end && buf[p] >= '0' && buf[p] <= '9') {
+                if (used < dec) {
+                    ACC_DIGIT(buf[p] - '0');
+                    used++;
+                }
+                p++;
+            }
+        }
+        if (p != end) return -1;
+        while (used < dec) {                     /* pad to the column scale */
+            ACC_DIGIT(0);
+            used++;
+        }
+        #undef ACC_DIGIT
+        out[k++] = neg ? -v : v;
+
+        if (i >= len) break;
+        i++;
+    }
+    return k == n ? 0 : -1;
+}
+
+/* ---------------------------------------------------------------- format */
+
+/* Write `n` scaled integers back as newline-separated decimal text.
+ * Returns bytes written, or -1 if `cap` was too small. */
+int64_t fmt_fixed(const int64_t *in, int64_t n, int32_t dec,
+                  char *out, int64_t cap)
+{
+    char digits[24];
+    int64_t at = 0;
+
+    for (int64_t k = 0; k < n; k++) {
+        int64_t v = in[k];
+        /* worst case: sign + 19 integer digits + point + 18 decimals + \n */
+        if (at + 40 > cap) return -1;
+
+        if (v < 0) { out[at++] = '-'; v = -v; }
+
+        int64_t scale = 1;
+        for (int32_t d = 0; d < dec; d++) scale *= 10;
+        int64_t q = dec ? v / scale : v;
+        int64_t r = dec ? v % scale : 0;
+
+        int nd = 0;
+        if (q == 0) digits[nd++] = '0';
+        while (q > 0) { digits[nd++] = (char)('0' + (q % 10)); q /= 10; }
+        while (nd > 0) out[at++] = digits[--nd];
+
+        if (dec) {
+            out[at++] = '.';
+            for (int32_t d = dec - 1; d >= 0; d--) {
+                int64_t p = 1;
+                for (int32_t e = 0; e < d; e++) p *= 10;
+                out[at++] = (char)('0' + ((r / p) % 10));
+            }
+        }
+        out[at++] = '\n';
+    }
+    return at ? at - 1 : 0;                      /* drop the final newline */
+}
+
+/* ------------------------------------------------------------- varints */
+
+/* Zigzag each value; small ones take a byte, the rest an escape plus a
+ * 64-bit tail entry. The tail must be 64 bits: a 32-bit tail silently
+ * truncated any value past 2^31, which is well inside the int64 range this
+ * codec accepts. */
+int64_t pack_ints(const int64_t *in, int64_t n,
+                  uint8_t *head, uint64_t *tail)
+{
+    int64_t nbig = 0;
+    for (int64_t i = 0; i < n; i++) {
+        int64_t v = in[i];
+        uint64_t u = v < 0 ? (uint64_t)(-v) * 2 - 1 : (uint64_t)v * 2;
+        if (u < ESCAPE) {
+            head[i] = (uint8_t)u;
+        } else {
+            head[i] = ESCAPE;
+            tail[nbig++] = u;
+        }
+    }
+    return nbig;
+}
+
+void unpack_ints(const uint8_t *head, int64_t n,
+                 const uint64_t *tail, int64_t *out)
+{
+    int64_t nbig = 0;
+    for (int64_t i = 0; i < n; i++) {
+        uint64_t u = head[i];
+        if (u == ESCAPE) u = tail[nbig++];
+        out[i] = (u & 1) ? -(int64_t)((u + 1) / 2) : (int64_t)(u / 2);
+    }
+}