|
| 1 | +"""Benchmark the flat chicago-taxi query across arrow / pandas / polars / blosc2. |
| 2 | +
|
| 3 | +Each method's select script is run in a *fresh* subprocess under ``/usr/bin/time`` |
| 4 | +(``-l`` on macOS, ``-v`` on Linux -- auto-detected), so the captured numbers |
| 5 | +include interpreter + library import cost and are directly comparable to what |
| 6 | +you see running the scripts by hand. Three metrics are captured per run: |
| 7 | +
|
| 8 | +* ``script time`` -- whole-process wall-clock (``real`` on macOS, ``Elapsed |
| 9 | + (wall clock) time`` on Linux), dominated by interpreter + library import cost. |
| 10 | +* ``query time`` -- the ``total:`` line each script prints: the actual |
| 11 | + open+compute+print work, excluding import/startup. This is the fair |
| 12 | + engine-to-engine comparison. |
| 13 | +* ``peak mem`` -- ``peak memory footprint`` on macOS, ``Maximum resident set |
| 14 | + size`` on Linux (a decent stand-in); reported in MB either way. |
| 15 | +
|
| 16 | +It also records each method's input ``size`` on disk (summary table only) and |
| 17 | +the result ``rows`` count, asserts all methods agree on the row count (else they |
| 18 | +aren't computing the same query), and -- with ``--nruns > 1`` -- draws min->max |
| 19 | +error bars. Bars are annotated with a "(x best)" multiple. Plots: script-time, |
| 20 | +query-time and peak-mem (each tagged cold/warm), plus a cache-independent |
| 21 | +two-bar input-size chart (shared .parquet vs .b2z). |
| 22 | +
|
| 23 | +Only macOS and Linux are supported (Windows has no ``/usr/bin/time``). |
| 24 | +
|
| 25 | +With ``--nruns N`` each script is run N times and the *minimum* across runs is |
| 26 | +used for the summary and plots (best-of-N: the fastest/leanest run is the one |
| 27 | +least perturbed by background interference and warms the OS file cache). Time |
| 28 | +and memory minima are taken independently. |
| 29 | +
|
| 30 | +``--nruns 1`` (default) measures a cold OS file cache and tags the plots |
| 31 | +``-cold.png``; ``--nruns > 1`` measures a warm cache (steady state) and tags |
| 32 | +them ``-warm.png``. |
| 33 | +
|
| 34 | +For a *true* cold-cache measurement pass ``--purge``: it flushes the OS file |
| 35 | +cache before every run (``sudo purge`` on macOS, ``drop_caches`` on Linux) and |
| 36 | +then reads a few MB from the *other* input file so the first timed read does |
| 37 | +not pay the storage device's idle-state exit latency (which can be tens of ms |
| 38 | +on NVMe drives with power management). Run ``sudo -v`` beforehand so sudo |
| 39 | +does not prompt for a password mid-benchmark. |
| 40 | +""" |
| 41 | +import argparse |
| 42 | +import os |
| 43 | +import platform |
| 44 | +import re |
| 45 | +import subprocess |
| 46 | +import sys |
| 47 | + |
| 48 | +import matplotlib |
| 49 | +matplotlib.use("Agg") # no display needed; we only save PNGs |
| 50 | +import matplotlib.pyplot as plt # noqa: E402 |
| 51 | + |
| 52 | +# (label, script, input-file kind). Order here drives the plot/column order. |
| 53 | +METHODS = [ |
| 54 | + ("duckdb", "select-duckdb-flat.py", "parquet"), |
| 55 | + ("arrow", "select-arrow-flat.py", "parquet"), |
| 56 | + ("pandas", "select-pandas-flat.py", "parquet"), |
| 57 | + ("polars", "select-polars-flat.py", "parquet"), |
| 58 | + ("blosc2", "select-blosc2.py", "b2z"), |
| 59 | +] |
| 60 | + |
| 61 | +# The `total:` line each select script prints to stdout (always a '.' decimal). |
| 62 | +_TOTAL_RE = re.compile(r"total:\s+([\d.]+)\s*s") |
| 63 | +# pandas' "[67 rows x 5 columns]" footer -- used as a cross-method sanity check. |
| 64 | +_ROWS_RE = re.compile(r"\[(\d+) rows x \d+ columns\]") |
| 65 | + |
| 66 | +# --- macOS: BSD `/usr/bin/time -l` (uses the locale decimal separator) --- |
| 67 | +_MAC_WALL_RE = re.compile(r"([\d.,]+)\s+real") |
| 68 | +_MAC_MEM_RE = re.compile(r"([\d,]+)\s+peak memory footprint") |
| 69 | +_MAC_MAXRSS_RE = re.compile(r"([\d,]+)\s+maximum resident set size") # bytes on macOS |
| 70 | + |
| 71 | +# --- Linux: GNU `/usr/bin/time -v` --- |
| 72 | +# Greedy `.*` walks to the rightmost colon followed by whitespace -- that's the |
| 73 | +# label/value separator; the value's own "m:ss" colons are not space-followed. |
| 74 | +_LIN_WALL_RE = re.compile(r"Elapsed \(wall clock\) time.*:\s+([\d:.]+)") |
| 75 | +_LIN_MEM_RE = re.compile(r"Maximum resident set size \(kbytes\):\s*(\d+)") |
| 76 | + |
| 77 | + |
| 78 | +def _mac_parse(err): |
| 79 | + m_wall = _MAC_WALL_RE.search(err) |
| 80 | + m_mem = _MAC_MEM_RE.search(err) or _MAC_MAXRSS_RE.search(err) |
| 81 | + if not (m_wall and m_mem): |
| 82 | + return None |
| 83 | + wall = float(m_wall.group(1).replace(",", ".")) # comma decimal -> dot |
| 84 | + return wall, int(m_mem.group(1).replace(",", "")) # bytes |
| 85 | + |
| 86 | + |
| 87 | +def _lin_parse(err): |
| 88 | + m_wall = _LIN_WALL_RE.search(err) |
| 89 | + m_mem = _LIN_MEM_RE.search(err) |
| 90 | + if not (m_wall and m_mem): |
| 91 | + return None |
| 92 | + # Elapsed is "h:mm:ss(.ss)" or "m:ss(.ss)"; fold sexagesimal parts to seconds. |
| 93 | + wall = 0.0 |
| 94 | + for part in m_wall.group(1).split(":"): |
| 95 | + wall = wall * 60 + float(part) |
| 96 | + return wall, int(m_mem.group(1)) * 1024 # kbytes -> bytes |
| 97 | + |
| 98 | + |
| 99 | +_SYSTEM = platform.system() |
| 100 | +if _SYSTEM == "Darwin": |
| 101 | + _TIME_FLAG, _parse_time = "-l", _mac_parse |
| 102 | +elif _SYSTEM == "Linux": |
| 103 | + _TIME_FLAG, _parse_time = "-v", _lin_parse |
| 104 | +else: |
| 105 | + sys.exit(f"Unsupported platform {_SYSTEM!r}: only macOS and Linux are supported.") |
| 106 | + |
| 107 | + |
| 108 | +def run_once(script, infile): |
| 109 | + """Run one script under /usr/bin/time (-l on macOS, -v on Linux). |
| 110 | +
|
| 111 | + Returns (script_time_s, query_time_s, peak_bytes, rows): wall-clock of the |
| 112 | + whole process, the script's own ``total:`` query time, peak memory |
| 113 | + footprint, and the result row count (for a cross-method sanity check). |
| 114 | + """ |
| 115 | + proc = subprocess.run( |
| 116 | + ["/usr/bin/time", _TIME_FLAG, sys.executable, script, infile], |
| 117 | + capture_output=True, text=True, |
| 118 | + ) |
| 119 | + if proc.returncode != 0: |
| 120 | + sys.stderr.write(proc.stdout + proc.stderr) |
| 121 | + raise RuntimeError(f"{script} exited with {proc.returncode}") |
| 122 | + timed = _parse_time(proc.stderr) |
| 123 | + m_total = _TOTAL_RE.search(proc.stdout) |
| 124 | + m_rows = _ROWS_RE.search(proc.stdout) |
| 125 | + if not (timed and m_total and m_rows): |
| 126 | + sys.stderr.write(proc.stdout + proc.stderr) |
| 127 | + raise RuntimeError(f"could not parse timing/result output for {script}") |
| 128 | + wall, peak = timed |
| 129 | + return wall, float(m_total.group(1)), peak, int(m_rows.group(1)) |
| 130 | + |
| 131 | + |
| 132 | +def flush_os_cache(): |
| 133 | + """Flush the OS file cache (needs sudo; run `sudo -v` first to cache credentials).""" |
| 134 | + if sys.platform == "darwin": |
| 135 | + subprocess.run(["sudo", "purge"], check=True) |
| 136 | + elif sys.platform.startswith("linux"): |
| 137 | + subprocess.run(["sudo", "sh", "-c", "sync; echo 3 > /proc/sys/vm/drop_caches"], |
| 138 | + check=True) |
| 139 | + else: |
| 140 | + sys.exit("--purge is only supported on macOS and Linux") |
| 141 | + |
| 142 | + |
| 143 | +def wake_disk(path, nbytes=4 * 2 ** 20): |
| 144 | + """Read a few MB so the first timed read does not pay the device's |
| 145 | + idle-state exit latency (tens of ms on NVMe drives with power management).""" |
| 146 | + with open(path, "rb") as f: |
| 147 | + f.read(nbytes) |
| 148 | + |
| 149 | + |
| 150 | +def ensure_b2z(parquet, b2z): |
| 151 | + """Build *b2z* from *parquet* via parquet-to-blosc2 if it doesn't exist.""" |
| 152 | + if os.path.exists(b2z): |
| 153 | + return |
| 154 | + if not os.path.exists(parquet): |
| 155 | + sys.exit(f"Neither {b2z!r} nor its source {parquet!r} exist; " |
| 156 | + f"cannot build the blosc2 input.") |
| 157 | + print(f"{b2z} not found -> building from {parquet} ...") |
| 158 | + subprocess.run(["parquet-to-blosc2", parquet, b2z, "--overwrite"], check=True) |
| 159 | + |
| 160 | + |
| 161 | +def main(): |
| 162 | + p = argparse.ArgumentParser(description=__doc__, |
| 163 | + formatter_class=argparse.RawDescriptionHelpFormatter) |
| 164 | + p.add_argument("--parquet", default="chicago-taxi-flat.parquet", |
| 165 | + help="parquet file fed to arrow/pandas/polars") |
| 166 | + p.add_argument("--b2z", default="chicago-taxi-flat.b2z", |
| 167 | + help="blosc2 CTable file fed to the blosc2 script") |
| 168 | + p.add_argument("--nruns", type=int, default=1, |
| 169 | + help="runs per method; the best (min) run is kept. 1 = cold " |
| 170 | + "cache (plots tagged -cold), >1 = warm cache (-warm)") |
| 171 | + p.add_argument("--prefix", default="compare", |
| 172 | + help="output PNG filename stem (e.g. compare-script-time.png)") |
| 173 | + p.add_argument("--purge", action="store_true", |
| 174 | + help="flush the OS file cache before every run (true cold-cache " |
| 175 | + "timing; needs sudo -- run `sudo -v` first). The disk is " |
| 176 | + "woken with a small read of the other input file so the " |
| 177 | + "timed run does not pay the device idle-exit latency.") |
| 178 | + args = p.parse_args() |
| 179 | + |
| 180 | + # The blosc2 input is derived from the parquet; build it on first use. |
| 181 | + ensure_b2z(args.parquet, args.b2z) |
| 182 | + |
| 183 | + infile = {"parquet": args.parquet, "b2z": args.b2z} |
| 184 | + results = {} # label -> dict(scripts=[], queries=[], peaks=[], rows, size) |
| 185 | + for label, script, kind in METHODS: |
| 186 | + path = infile[kind] |
| 187 | + size = os.path.getsize(path) |
| 188 | + print(f"== {label:7s} ({script} {path}, {size / 1e6:.1f} MB) ==") |
| 189 | + # Wake the disk with the *other* input file: same device, but does not |
| 190 | + # warm any byte of the file about to be timed. |
| 191 | + wake_path = infile["b2z" if kind == "parquet" else "parquet"] |
| 192 | + scripts, queries, peaks, rows = [], [], [], None |
| 193 | + for i in range(args.nruns): |
| 194 | + if args.purge: |
| 195 | + flush_os_cache() |
| 196 | + wake_disk(wake_path) |
| 197 | + script_s, query_s, peak, rows = run_once(script, path) |
| 198 | + scripts.append(script_s) |
| 199 | + queries.append(query_s) |
| 200 | + peaks.append(peak) |
| 201 | + print(f" run {i + 1}/{args.nruns}: script {script_s:6.3f} s " |
| 202 | + f"query {query_s:6.3f} s {peak / 1e6:8.1f} MB") |
| 203 | + results[label] = {"scripts": scripts, "queries": queries, |
| 204 | + "peaks": peaks, "rows": rows, "size": size} |
| 205 | + |
| 206 | + # ---- cross-method correctness check ---- |
| 207 | + rowcounts = {l: results[l]["rows"] for l in results} |
| 208 | + if len(set(rowcounts.values())) > 1: |
| 209 | + print("\n!! WARNING: methods returned DIFFERENT row counts: " |
| 210 | + + ", ".join(f"{l}={r}" for l, r in rowcounts.items()) |
| 211 | + + " -- they may not be computing the same query!") |
| 212 | + else: |
| 213 | + print(f"\nrow-count check OK: all methods returned " |
| 214 | + f"{next(iter(rowcounts.values()))} rows") |
| 215 | + |
| 216 | + # ---- summary table ---- |
| 217 | + hdr = (f"{'method':8s} {'script (s)':>11s} {'query (s)':>11s} " |
| 218 | + f"{'peak mem (MB)':>15s} {'size (MB)':>11s} {'rows':>7s}") |
| 219 | + print(f"\n{hdr}\n{'-' * len(hdr)}") |
| 220 | + for label, _, _ in METHODS: |
| 221 | + r = results[label] |
| 222 | + print(f"{label:8s} {min(r['scripts']):11.3f} {min(r['queries']):11.3f} " |
| 223 | + f"{min(r['peaks']) / 1e6:15.1f} {r['size'] / 1e6:11.1f} {r['rows']:7d}") |
| 224 | + |
| 225 | + # ---- plots ---- |
| 226 | + labels = [m[0] for m in METHODS] |
| 227 | + colors = ["#8172B3", "#4C72B0", "#DD8452", "#55A868", "#C44E52"] |
| 228 | + |
| 229 | + def min_err(key, scale=1.0): |
| 230 | + """Best-of-N value per method plus an upward (min->max) error bar.""" |
| 231 | + mins = [min(results[l][key]) * scale for l in labels] |
| 232 | + ups = [(max(results[l][key]) - min(results[l][key])) * scale for l in labels] |
| 233 | + yerr = None if all(u == 0 for u in ups) else [[0] * len(mins), ups] |
| 234 | + return mins, yerr |
| 235 | + |
| 236 | + def barplot(values, ylabel, title, fname, fmt, yerr=None, ratio=True, cats=None): |
| 237 | + cats = cats if cats is not None else labels |
| 238 | + fig, ax = plt.subplots(figsize=(7, 4.5)) |
| 239 | + bars = ax.bar(cats, values, color=colors[:len(cats)], |
| 240 | + yerr=yerr, capsize=4 if yerr is not None else 0) |
| 241 | + ax.set_ylabel(ylabel) |
| 242 | + ax.set_title(title) |
| 243 | + # value, plus a "x best" multiple so the chart reads at a glance |
| 244 | + best = min(v for v in values if v > 0) if any(values) else 1.0 |
| 245 | + text = [f"{v:{fmt}}\n({v / best:.1f}×)" if ratio else f"{v:{fmt}}" |
| 246 | + for v in values] |
| 247 | + ax.bar_label(bars, labels=text, padding=3) |
| 248 | + ax.margins(y=0.20) |
| 249 | + ax.grid(axis="y", alpha=0.3) |
| 250 | + fig.tight_layout() |
| 251 | + fig.savefig(fname, dpi=120) |
| 252 | + print(f"wrote {fname}") |
| 253 | + |
| 254 | + n = args.nruns |
| 255 | + cache = "cold" if n == 1 else "warm" |
| 256 | + suffix = f"(best of {n} runs, {cache} cache)" if n > 1 else "(single run, cold cache)" |
| 257 | + |
| 258 | + s_vals, s_err = min_err("scripts") |
| 259 | + q_vals, q_err = min_err("queries") |
| 260 | + m_vals, m_err = min_err("peaks", 1 / 1e6) |
| 261 | + barplot(s_vals, "script time (s)", |
| 262 | + f"Whole-process wall-clock time {suffix}", |
| 263 | + f"{args.prefix}-script-time-{cache}.png", ".3f", yerr=s_err) |
| 264 | + barplot(q_vals, "query time (s)", |
| 265 | + f"In-script query time (excl. import) {suffix}", |
| 266 | + f"{args.prefix}-query-time-{cache}.png", ".3f", yerr=q_err) |
| 267 | + barplot(m_vals, "peak memory footprint (MB)", |
| 268 | + f"Query peak memory {suffix}", |
| 269 | + f"{args.prefix}-query-mem-{cache}.png", ".0f", yerr=m_err) |
| 270 | + # input size: cache-independent, two bars (shared parquet vs the .b2z) |
| 271 | + barplot([os.path.getsize(args.parquet) / 1e6, os.path.getsize(args.b2z) / 1e6], |
| 272 | + "file size on disk (MB)", "Input file size", |
| 273 | + f"{args.prefix}-size.png", ".1f", cats=[".parquet", ".b2z"], ratio=False) |
| 274 | + |
| 275 | + |
| 276 | +if __name__ == "__main__": |
| 277 | + main() |
0 commit comments