#!/usr/bin/env python3
"""DDIA Ch. 7 (Sharding) -- hashing the partition key destroys range queries.

Put 1,000,000 keys across 16 shards two ways and run the SAME contiguous range
scan (keys 500000..500999) against each:

  * Key-range sharding  -- shard = key // (1_000_000 // 16) -- keeps adjacent
    keys adjacent, so a contiguous range lives inside ONE shard.
  * Hash sharding        -- shard = md5(key) % 16 -- is designed to scatter
    adjacent keys, so the same 1000-key range is smeared across ALL 16 shards
    and the query must fan out to every one of them.

Then the controls: a single-key POINT lookup touches exactly 1 shard under BOTH
schemes (hashing only hurts ranges, not points); and a COMPOSITE key
(f"{sensor}#{ts}", hashed by the sensor prefix only) keeps a range over ts
within one sensor on a single shard even under hashing.

Deterministic: placement is decided by md5, never Python's per-process
randomized hash(), so every number below reproduces byte-for-byte. Stdlib only.
"""

import hashlib

SHARDS = 16
N_KEYS = 1_000_000
RANGE_PER_SHARD = N_KEYS // SHARDS  # 62_500 contiguous keys per range-shard

# Zero-padded so string order == numeric order (what a range scan relies on).
def kstr(i):
    return f"{i:06d}"


def h(s):
    """Stable, process-independent 128-bit hash of a string as an int."""
    return int(hashlib.md5(s.encode()).hexdigest(), 16)


# --- Two placement schemes ---------------------------------------------------
def range_shard(key_int):
    """Contiguous slices: shard = key // 62500."""
    return key_int // RANGE_PER_SHARD


def hash_shard(key_int):
    """Scatter: shard = md5(key) % 16."""
    return h(kstr(key_int)) % SHARDS


def shards_touched(key_ints, shard_fn):
    """Set of distinct shards holding at least one key in key_ints."""
    return {shard_fn(k) for k in key_ints}


def main():
    LO, HI = 500_000, 501_000  # 1000 consecutive keys: [500000, 501000)
    rng = list(range(LO, HI))

    print(f"{N_KEYS:,} keys (000000..999999) spread over {SHARDS} shards, two ways.")
    print(f"    key-range sharding : shard = key // {RANGE_PER_SHARD:,}  (contiguous slices)")
    print(f"    hash sharding      : shard = md5(key) % {SHARDS}       (scattered)")
    print()

    # --- Step 1: the range query -------------------------------------------
    print(f"RANGE QUERY -- scan the {len(rng):,} consecutive keys "
          f"[{LO}, {HI}):")
    print()

    r_hit = shards_touched(rng, range_shard)
    print("  key-range sharding:")
    print(f"    shards touched : {len(r_hit)}   -> {sorted(r_hit)}")
    print(f"    the whole range lives in shard {sorted(r_hit)[0]} "
          f"(keys {sorted(r_hit)[0]*RANGE_PER_SHARD}"
          f"..{(sorted(r_hit)[0]+1)*RANGE_PER_SHARD-1}); "
          f"query hits 1 shard.")
    print()

    h_hit = shards_touched(rng, hash_shard)
    # how the 1000 keys spread across the shards they hit
    spread = {}
    for k in rng:
        s = hash_shard(k)
        spread[s] = spread.get(s, 0) + 1
    lo_n, hi_n = min(spread.values()), max(spread.values())
    print("  hash sharding:")
    print(f"    shards touched : {len(h_hit)}  -> {sorted(h_hit)}")
    print(f"    the {len(rng):,} keys spread {lo_n}-{hi_n} per shard "
          f"(~{len(rng)/SHARDS:.1f} each); query must fan out to ALL "
          f"{len(h_hit)} shards.")
    print()
    print(f"  SAME query, SAME data: {len(r_hit)} shard vs {len(h_hit)} shards.")
    print()

    # --- Step 2: the point query control -----------------------------------
    point = 500_000
    print(f"POINT QUERY -- look up the single key {point}:")
    print(f"    key-range sharding : shard {range_shard(point)}   (1 shard)")
    print(f"    hash sharding      : shard {hash_shard(point)}   (1 shard)")
    print(f"    both touch exactly 1 shard -- hashing is only bad for RANGES.")
    print()

    # --- Step 3: composite-key redemption ----------------------------------
    # Keys are f"{sensor}#{ts}". Hash the SENSOR PREFIX only (the partition
    # key); ts is the sort key within a sensor. A range over ts for ONE sensor
    # therefore shares one hash and stays on one shard.
    def composite_shard(sensor, ts):
        return h(sensor) % SHARDS

    sensor = "sensor-0042"
    ts_lo, ts_hi = 1_000, 2_000  # 1000 consecutive readings for this sensor
    comp_range = [(sensor, ts) for ts in range(ts_lo, ts_hi)]
    c_hit = {composite_shard(s, t) for (s, t) in comp_range}

    print("COMPOSITE KEY -- key = f\"{sensor}#{ts}\", hashed by sensor prefix only:")
    print(f"    range over ts in [{ts_lo}, {ts_hi}) for {sensor!r} "
          f"({len(comp_range):,} readings)")
    print(f"    shards touched : {len(c_hit)}   -> {sorted(c_hit)}")
    print(f"    all readings share the sensor prefix -> one hash "
          f"-> shard {sorted(c_hit)[0]}; range stays on 1 shard even under hashing.")


if __name__ == "__main__":
    main()