Data Partitioning skill for Claude Code

#!/usr/bin/env python3
"""Simulate how a candidate shard key spreads load across shards.

Input: a CSV or text file of observed keys, one per line as `key` or `key,weight`
(weight = rows, requests, or bytes for that key; default 1). A header line starting
with "key" is skipped. Use "-" to read stdin.

Usage:
    python3 scripts/shard_skew.py keys.csv --shards 8 [--vnodes 128] [--slots 1024]
                                  [--max-skew 1.25] [--hot-key-share 0.5]

For each strategy (hash mod N, consistent-hash ring with virtual nodes, fixed logical
slots mapped to shards) it prints per-shard load, skew (max / mean), and the share of
load that moves when growing from N to N+1 shards. It also lists the hottest keys.

Exit codes: 0 = key distributes acceptably, 1 = skew or hot keys exceed thresholds,
2 = bad input.
"""
import argparse
import bisect
import hashlib
import statistics
import sys
from collections import Counter


def h64(text: str) -> int:
    return int.from_bytes(hashlib.md5(text.encode("utf-8")).digest()[:8], "big")


def read_keys(path: str) -> Counter:
    fh = sys.stdin if path == "-" else open(path, encoding="utf-8")
    weights: Counter = Counter()
    with fh:
        for n, line in enumerate(fh, 1):
            line = line.strip()
            if not line or (n == 1 and line.lower().startswith("key")):
                continue
            key, _, w = line.partition(",")
            try:
                weights[key.strip()] += float(w) if w.strip() else 1.0
            except ValueError:
                sys.exit(f"error: line {n}: weight '{w}' is not a number")
    if not weights:
        print("error: no keys read", file=sys.stderr)
        sys.exit(2)
    return weights


class ModN:
    name = "hash mod N"

    def __init__(self, shards: int, **_):
        self.n = shards

    def shard(self, key: str) -> int:
        return h64(key) % self.n


class Ring:
    name = "consistent-hash ring"

    def __init__(self, shards: int, vnodes: int = 128, **_):
        points = sorted((h64(f"shard-{s}#vn-{v}"), s) for s in range(shards) for v in range(vnodes))
        self.hashes = [p[0] for p in points]
        self.owners = [p[1] for p in points]

    def shard(self, key: str) -> int:
        i = bisect.bisect(self.hashes, h64(key)) % len(self.hashes)
        return self.owners[i]


class Slots:
    name = "logical slots"

    def __init__(self, shards: int, slots: int = 1024, assignment=None, **_):
        self.slots = slots
        self.assignment = assignment or [s % shards for s in range(slots)]

    def shard(self, key: str) -> int:
        return self.assignment[h64(key) % self.slots]

    def grown(self, new_shards: int) -> "Slots":
        """Move the minimum number of slots to a new shard, taking evenly from the others."""
        assignment = list(self.assignment)
        target = self.slots // new_shards
        owned = Counter(assignment)
        new_id = new_shards - 1
        moved = 0
        for slot, owner in enumerate(assignment):
            if moved >= target:
                break
            if owned[owner] > target:
                assignment[slot] = new_id
                owned[owner] -= 1
                moved += 1
        return Slots(new_shards, self.slots, assignment)


def loads(strategy, weights: Counter, shards: int):
    per = [0.0] * shards
    for key, w in weights.items():
        per[strategy.shard(key)] += w
    return per


def moved_share(old, new, weights: Counter) -> float:
    total = sum(weights.values())
    moved = sum(w for k, w in weights.items() if old.shard(k) != new.shard(k))
    return moved / total


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("file")
    ap.add_argument("--shards", type=int, required=True)
    ap.add_argument("--vnodes", type=int, default=128)
    ap.add_argument("--slots", type=int, default=1024)
    ap.add_argument("--max-skew", type=float, default=1.25, help="max/mean shard load allowed (default 1.25)")
    ap.add_argument("--hot-key-share", type=float, default=0.5,
                    help="flag keys heavier than this fraction of a mean shard (default 0.5)")
    ap.add_argument("--top", type=int, default=5)
    args = ap.parse_args()
    if args.shards < 1 or args.slots < args.shards:
        print("error: need shards >= 1 and slots >= shards", file=sys.stderr)
        return 2

    weights = read_keys(args.file)
    total = sum(weights.values())
    mean = total / args.shards
    n = args.shards
    print(f"{len(weights)} distinct keys, total weight {total:,.0f}, {n} shards, mean load {mean:,.1f}\n")

    strategies = [
        (ModN(n), ModN(n + 1)),
        (Ring(n, vnodes=args.vnodes), Ring(n + 1, vnodes=args.vnodes)),
    ]
    slots = Slots(n, slots=args.slots)
    strategies.append((slots, slots.grown(n + 1)))

    print(f"{'strategy':<22} {'skew max/mean':>13} {'stdev %':>8} {'moved on N+1':>13}")
    results = {}
    for cur, nxt in strategies:
        per = loads(cur, weights, n)
        skew = max(per) / mean if mean else 0
        stdev = (statistics.pstdev(per) / mean * 100) if mean else 0
        mv = moved_share(cur, nxt, weights)
        results[cur.name] = (skew, per)
        print(f"{cur.name:<22} {skew:>13.3f} {stdev:>7.1f}% {mv * 100:>12.1f}%")

    best = min(results, key=lambda k: results[k][0])
    print(f"\nper-shard load ({best}):")
    for i, v in enumerate(results[best][1]):
        bar = "#" * int(40 * v / max(results[best][1]))
        print(f"  shard {i:>3}: {v:>12,.1f}  {bar}")

    print(f"\ntop {args.top} keys:")
    hot = []
    for key, w in weights.most_common(args.top):
        share_of_shard = w / mean if mean else 0
        flag = "  <-- hot key" if share_of_shard > args.hot_key_share else ""
        if flag:
            hot.append(key)
        print(f"  {key:<30} {w:>12,.1f}  {w / total * 100:5.1f}% of total, {share_of_shard:5.2f}x mean shard{flag}")

    problems = []
    if results[best][0] > args.max_skew:
        problems.append(f"best skew {results[best][0]:.3f} exceeds --max-skew {args.max_skew}")
    if hot:
        problems.append(f"{len(hot)} hot key(s) exceed {args.hot_key_share}x a mean shard; no hash spreads a single key")
    print()
    if problems:
        for p in problems:
            print(f"FAIL: {p}")
        print("Consider a higher-cardinality key, a compound key, dedicated shards, or salting hot keys.")
        return 1
    print("OK: key distributes within thresholds.")
    return 0


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