benchmark.py

#!/usr/bin/env python3
import gc
import json
import sys
import time
import tracemalloc

import jsonfold


REPEATS = 3
MEM_FRACTION = 0.10


class NullWriter:
    def __init__(self, t0=None):
        self.t0 = t0 if t0 else time.perf_counter ;
        self.first_write = None
        self.bytes = 0
        self.writes = 0

    def write(self, s):
        if self.first_write is None:
            self.first_write = time.perf_counter()
        self.bytes += len(s)
        self.writes += 1
        return len(s)

    def ttfb_ms(self):
        if self.first_write is None:
            return ""
        return round((self.first_write - self.t0) * 1000, 1)


def make_data(rows):
    return {
        "meta": {"version": 1, "ok": True, "name": "jsonfold benchmark"},
        "long_ids": list(range(100)),
        "long_obj": {f"k{i}": i for i in range(50)},
        "rows": [
            {
                "id": i,
                "name": f"name_{i}",
                "active": i % 3 == 0,
                "score": i * 1.25,
                "tags": ["alpha", "beta", "gamma", "delta"],
                "pos": {"x": i, "y": i + 1, "z": i + 2},
                "values": [i, i + 1, i + 2, i + 3, i + 4],
                "pairs": [ [ i, i+1, [i+2, i+3], [i+4, i+5]]],
            }
            for i in range(rows)
        ],
    }


def mem_label():
    return "kb"

def mem_units(n):
    return round(n / 1024, 1)


def run_case(data, name):
    if name == "base.dumps.plain":
        return lambda t0: write_string(t0, json.dumps(data))
    if name == "base.dumps.pretty":
        return lambda t0: write_string(t0, json.dumps(data, indent=2))
    if name == "base.dump.pretty":
        return lambda t0: run_json_dump(data, t0)
    if name == "base.dump.plain":
        return lambda t0: run_json_dump_plain(data, t0)

    kind, func, compact = name.split(".")

    if kind == "jf":
        if func == "dumps":
            return lambda t0: write_string(t0, jsonfold.dumps(data, compact=compact, indent=2))
        if func == "dump":
            return lambda t0: run_jsonfold_dump(data, t0, compact)

    raise ValueError(name)


def write_string(t0, s):
    w = NullWriter(t0)
    w.write(s)
    return w


def run_json_dump(data, t0):
    w = NullWriter(t0)
    json.dump(data, w, indent=2)
    return w


def run_jsonfold_dump(data, t0, compact):
    w = NullWriter(t0)
    jsonfold.dump(data, w, compact=compact, indent=2)
    return w

def run_json_dump_plain(data, t0):
    w = NullWriter(t0)
    json.dump(data, w)
    return w

def time_one(name, data):
    best = None
    best_dt = 0

    for _ in range(REPEATS):
        gc.collect()
        t0 = time.perf_counter()
        p0 = time.process_time()
        w = run_case(data, name)(t0)
        p1 = time.process_time()
        t1 = time.perf_counter()
        dt = t1-t0

        row = {
            "time(ms)": round(dt * 1000, 1),
            "CPU(ms)": round((p1 - p0) * 1000, 1),
            "ttfb(ms)": w.ttfb_ms(),
            f"out({mem_label()})": mem_units(w.bytes),
            "writes": w.writes,
        }

        if best is None or dt < best_dt:
            best = row
            best_dt = dt

    return best_dt, best


def memory_one(name, data):
    gc.collect()
    tracemalloc.start()

    t0 = time.perf_counter()
    run_case(data, name)(t0)

    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()

    return mem_units(peak)


def print_table(rows):
    cols = list(rows[0].keys())
    widths = {c: max(len(c), *(len(str(r[c])) for r in rows)) for c in cols}

    def is_num(v):
        return isinstance(v, (int, float)) and not isinstance(v, bool)

    numeric = {
        c: all(is_num(r[c]) or r[c] == "" for r in rows)
        for c in cols
    }

    def cell(c, v):
        s = str(v)
        return " " + (s.rjust(widths[c]) if numeric[c] else s.ljust(widths[c])) + " "

    line = "+" + "+".join("-" * (widths[c] + 2) for c in cols) + "+"

    print(line)
    print("|" + "|".join(cell(c, c) for c in cols) + "|")
    print(line)
    for r in rows:
        print("|" + "|".join(cell(c, r[c]) for c in cols) + "|")
    print(line)

def run_one_size(rows, tests):
    data = make_data(rows)

    if not tests:
        tests = [
            "base.dump.plain",
            "base.dump.pretty",
            "jf.dump.off",
            "jf.dump.none",
            "jf.dump.default",
            "jf.dump.low",
            "jf.dump.med",
            "jf.dump.classic",
            "jf.dump.high",
            "jf.dump.max",
            "jf.dump.pack",
            "jf.dump.fold",
            "jf.dump.grid",
            "jf.dump.join",
            "base.dumps.plain",
            "base.dumps.pretty",
            "jf.dumps.none",
            "jf.dumps.default",
            "jf.dumps.high",
            "jf.dumps.max",
            ]
    results = []

    for name in tests:
        print(f"{name} ({rows})... ", end="", file=sys.stderr, flush=True)

        t0 = time.perf_counter() ;
        best_dt, speed = time_one(name, data)
        peak_mb = memory_one(name, data)
        t1 = time.perf_counter() ;
        dt = t1-t0 ;

        print(f"{round(dt*1000,0)} ms", file=sys.stderr)

        results.append({
            "rows": rows,
            "name": name,
            **speed,
            f"peak({mem_label()})": peak_mb,
        })

    return results

def show_data(rows):
    data = make_data(rows)
    json.dump(make_data(rows), fp=sys.stdout, indent=2)


def main(argv):
    import argparse

    p = argparse.ArgumentParser(
        description="Read JSON from stdin; write folded JSON to stdout.")
    p.add_argument("--show", type=int, default=None)
    args, rest = p.parse_known_args(argv)

    t0 = time.perf_counter() ;
    filter = []
    last_sz = None
    results = []
    if (args.show):
        show_data(args.show)
        return

    for arg in rest[1:]:
        if arg == "-":
            filter = []
            continue

        try:
            last_sz = int(arg)
        except ValueError:
            filter.append(arg)
            continue
        results.extend(run_one_size(last_sz, filter))

    if last_sz is None:
        results.extend(run_one_size(1_000, filter))

    t1 = time.perf_counter() ;

    print_table(results)
    print(f"completed in: {round(t1-t0, 1)}", file=sys.stderr)

if __name__ == "__main__":
    main(sys.argv)