chart_results.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["matplotlib"]
# ///
"""Generate benchmark charts from a results JSON file.

Produces a 3-panel figure: Wall Time, CPU Time, and Peak Memory,
comparing all strategies across row counts.

Usage:
    python chart_results.py                           # defaults
    python chart_results.py --results FILE --chart FILE
"""

import argparse
import json

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


DEFAULT_RESULTS = "benchmark_results.json"
DEFAULT_CHART = "benchmark_chart.png"

STRATEGIES = [
    ("hash",         "Hash Join (py)",  "#d62728", "o"),
    ("merge",        "Merge Join (py)", "#2ca02c", "s"),
    ("sql",          "SQL Join (py)",   "#1f77b4", "D"),
    ("index_lookup", "Index Lookup (py)", "#9467bd", "v"),
    ("node_merge",   "Merge Join (js)", "#ff7f0e", "^"),
    ("rust_merge",   "Merge Join (rs)", "#8c564b", "P"),
]


def fmt_rows(x, _pos):
    """Format tick labels: 10000 → '10K', 1000000 → '1M'."""
    if x >= 1_000_000:
        return f"{x / 1_000_000:.0f}M"
    if x >= 1_000:
        return f"{x / 1_000:.0f}K"
    return str(int(x))


def fmt_mem(b):
    if b < 1024:
        return f"{b:.0f} B"
    elif b < 1024**2:
        return f"{b / 1024:.1f} KiB"
    elif b < 1024**3:
        return f"{b / 1024**2:.1f} MiB"
    else:
        return f"{b / 1024**3:.2f} GiB"


def main(results_file=DEFAULT_RESULTS, chart_file=DEFAULT_CHART, title=None):
    with open(results_file) as f:
        data = json.load(f)

    rows = data["row_counts"]

    # Filter to strategies present in the data
    active = [s for s in STRATEGIES if f"{s[0]}_wall" in data]

    if title is None:
        title = "Sync Strategy Benchmark"

    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    fig.suptitle(title, fontsize=14, fontweight="bold")

    panels = [
        (0, "wall",     "Wall Time vs Row Count",  "Wall Time (seconds)"),
        (1, "cpu",      "CPU Time vs Row Count",   "CPU Time (seconds)"),
        (2, "peak_rss", "Peak Memory vs Row Count", "Peak RSS (MiB)"),
    ]

    for ax_idx, metric, title, ylabel in panels:
        ax = axes[ax_idx]
        for key_prefix, label, color, marker in active:
            key = f"{key_prefix}_{metric}"
            values = data[key]
            if metric == "peak_rss":
                values = [b / 1024**2 for b in values]  # -> MiB
            ax.loglog(rows, values, color=color, marker=marker,
                      linewidth=2, markersize=5, label=label)
        ax.set_title(title)
        ax.set_xlabel("CSV Rows")
        ax.set_ylabel(ylabel)
        ax.xaxis.set_major_formatter(ticker.FuncFormatter(fmt_rows))
        ax.legend(fontsize=8)
        ax.grid(True, which="both", alpha=0.3)

    plt.tight_layout()
    plt.savefig(chart_file, dpi=150, bbox_inches="tight")
    print(f"Chart saved to {chart_file}")

    # Print summary table
    header = f"{'Rows':>12}"
    for _, label, _, _ in active:
        header += f"  {label + ' Wall':>18} {label + ' RSS':>12}"
    print(f"\n{header}")
    print("-" * len(header))
    for i, n in enumerate(rows):
        line = f"{n:>12,}"
        for key_prefix, _, _, _ in active:
            w = data[f"{key_prefix}_wall"][i]
            m = data[f"{key_prefix}_peak_rss"][i]
            line += f"  {w:>17.2f}s {fmt_mem(m):>12}"
        print(line)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate benchmark charts")
    parser.add_argument("--results", default=DEFAULT_RESULTS,
                        help="Path to results JSON file")
    parser.add_argument("--chart", default=DEFAULT_CHART,
                        help="Path to output chart PNG")
    parser.add_argument("--title", default=None,
                        help="Chart super title (default: Sync Strategy Benchmark)")
    args = parser.parse_args()
    main(results_file=args.results, chart_file=args.chart, title=args.title)