kumar6591/data-quality-env

1from __future__ import annotations2 3import itertools4import re5from dataclasses import dataclass6from hashlib import sha17 8 9_ALGO_BANK: list["AlgorithmSpec"] | None = None10_BEST_SPEC_CACHE: dict[str, "AlgorithmSpec"] = {}11 12 13@dataclass(frozen=True)14class AlgorithmSpec:15    algorithm_id: int16    w_coverage: float17    w_stat: float18    w_risk: float19    w_novelty: float20    w_limit: float21    w_prior: float22    repeat_penalty: float23 24 25def generate_100k_algorithms() -> list[AlgorithmSpec]:26    """Generate exactly 100,000 deterministic algorithm specs."""27    global _ALGO_BANK28    if _ALGO_BANK is not None:29        return _ALGO_BANK30 31    out: list[AlgorithmSpec] = []32    # 10 * 10 * 10 * 10 * 5 * 2 = 100,00033    grids = [34        [i / 10 for i in range(10)],35        [i / 10 for i in range(10)],36        [i / 10 for i in range(10)],37        [i / 10 for i in range(10)],38        [i / 5 for i in range(5)],39        [0.0, 1.0],40    ]41 42    idx = 043    for a, b, c, d, e, f in itertools.product(*grids):44        out.append(45            AlgorithmSpec(46                algorithm_id=idx,47                w_coverage=a,48                w_stat=b,49                w_risk=c,50                w_novelty=d,51                w_limit=e,52                w_prior=(idx % 5) / 5,53                repeat_penalty=f * 0.03,54            )55        )56        idx += 157 58    _ALGO_BANK = out59    return _ALGO_BANK60 61 62def _query_features(sql: str) -> dict[str, float]:63    s = (sql or "").lower()64    return {65        "coverage": float(any(k in s for k in ["count(", "sum(", "avg(", "group by", "distinct"])),66        "stat": float(any(k in s for k in ["avg(", "stddev", "variance", "percentile", "try_cast", "strptime"])),67        "risk": float(any(k in s for k in ["drop", "truncate", "delete", "insert", "update", "alter", "create"])),68        "novelty": float(any(k in s for k in ["left join", "except", "not in", "having", "case when"])),69        "has_limit": float("limit" in s),70    }71 72 73def _task_relevance(task_id: int, sql: str) -> float:74    s = (sql or "").lower()75    if task_id == 1:76        keys = ["null", "email", "customer_id", "duplicate", "group by"]77    elif task_id == 2:78        keys = ["quantity", "amount", "n/a", "try_cast", "order_date"]79    else:80        keys = ["transactions_baseline", "transactions_current", "category", "user_id", "avg(amount)"]81    hits = sum(1 for k in keys if k in s)82    return hits / max(1, len(keys))83 84 85def algorithm_rule_check(spec: AlgorithmSpec, queries: list[str], max_steps: int = 10) -> bool:86    """87    Enforces constraints aligned with hackathon rules for this environment:88    - non-destructive SQL preference89    - bounded steps90    - deterministic finite parameters91    """92    if max_steps <= 0 or max_steps > 10:93        return False94    if spec.w_risk < 0.0 or spec.w_risk > 1.0:95        return False96    if spec.repeat_penalty < 0.0 or spec.repeat_penalty > 0.03:97        return False98 99    for q in queries:100        s = (q or "").strip()101        if not s:102            return False103        if re.search(r"\b(drop|truncate|delete|insert|update|alter|create)\b", s, flags=re.IGNORECASE):104            return False105        if not re.match(r"^\s*(select|with)\b", s, flags=re.IGNORECASE):106            return False107    return True108 109 110def rank_queries(task_id: int, queries: list[str], priors: list[float], spec: AlgorithmSpec) -> list[int]:111    scored: list[tuple[int, float]] = []112    for i, q in enumerate(queries):113        f = _query_features(q)114        prior = priors[i] if i < len(priors) else 0.0115        relevance = _task_relevance(task_id, q)116        score = (117            spec.w_coverage * f["coverage"]118            + spec.w_stat * f["stat"]119            + spec.w_novelty * f["novelty"]120            + spec.w_limit * f["has_limit"]121            + spec.w_prior * prior122            + 0.8 * relevance123            - spec.w_risk * f["risk"]124        )125        scored.append((i, score))126    scored.sort(key=lambda x: x[1], reverse=True)127    return [i for i, _ in scored]128 129 130def choose_best_algorithm(task_id: int, queries: list[str], priors: list[float], max_algorithms: int = 100_000) -> AlgorithmSpec:131    key_payload = f"t={task_id}|n={len(queries)}|m={max_algorithms}|q={'||'.join(queries)}|p={','.join(f'{x:.4f}' for x in priors)}"132    cache_key = sha1(key_payload.encode("utf-8")).hexdigest()133    if cache_key in _BEST_SPEC_CACHE:134        return _BEST_SPEC_CACHE[cache_key]135 136    algorithms = generate_100k_algorithms()137    n = min(max_algorithms, len(algorithms))138 139    best = algorithms[0]140    best_obj = -1e18141 142    for spec in algorithms[:n]:143        if not algorithm_rule_check(spec, queries, max_steps=10):144            continue145        ranking = rank_queries(task_id, queries, priors, spec)146        top = ranking[:2]147        obj = 0.0148        for pos, i in enumerate(top):149            base = 2.0 - pos150            rel = _task_relevance(task_id, queries[i])151            obj += base * rel152        # Prefer slight risk aversion153        obj -= 0.1 * spec.w_risk154        if obj > best_obj:155            best_obj = obj156            best = spec157 158    _BEST_SPEC_CACHE[cache_key] = best159    return best160 161 162def order_queries_with_100k_algorithms(task_id: int, queries: list[str], priors: list[float]) -> list[str]:163    spec = choose_best_algorithm(task_id, queries, priors, max_algorithms=100_000)164    ranked_idx = rank_queries(task_id, queries, priors, spec)165    return [queries[i] for i in ranked_idx]166