# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure the top-1% share of spend.

On your warehouse, run the query below against its query history. Without
one, this script simulates query histories whose per-query cost is
log-normal with a range of tail weights (sigma), and reports the share of
spend carried by the most expensive 1% of queries and what a cap at the 99th
percentile would block.
"""
import argparse, json
import numpy as np

SQL = """-- the share of spend carried by the most expensive 1% of queries
WITH q AS (
  SELECT query_id, bytes_scanned AS cost,
         NTILE(100) OVER (ORDER BY bytes_scanned DESC) AS pct
  FROM query_history
  WHERE start_time >= current_date - INTERVAL 30 DAY
)
SELECT SUM(CASE WHEN pct = 1 THEN cost END) / SUM(cost) AS top_1pct_share FROM q;"""

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--queries", type=int, default=100_000, help="Queries in each simulated history.")
p.add_argument("--sigmas", type=lambda s: [float(v) for v in s.split(",")], default=[1.0, 1.5, 2.0, 2.5, 3.0],
               help="Log-normal sigma of per-query cost (tail weight), comma-separated.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)
print(SQL, "\n")

rows = []
for s in args.sigmas:
    cost = np.sort(rng.lognormal(0, s, args.queries))[::-1]
    top = cost[: args.queries // 100].sum() / cost.sum()
    p99 = np.percentile(cost, 99)
    rows.append({"sigma": s, "top1_share_pct": top * 100, "p99_over_median": float(p99 / np.median(cost)),
                 "blocked_by_p99_cap": int((cost > p99).sum())})
    print(f"sigma {s:.1f}: top 1% of queries carry {top:6.1%} of spend; p99 is {rows[-1]['p99_over_median']:6.0f}x the median; "
          f"a p99 cap blocks {rows[-1]['blocked_by_p99_cap']:,} queries")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"sigmas": rows}}, f)